From 69224adb4825e66c0bdb9211860418399ae3b3b0 Mon Sep 17 00:00:00 2001 From: Xin Yao Date: Thu, 11 Jun 2026 00:41:10 -0700 Subject: [PATCH 1/3] initial commit Signed-off-by: Xin Yao --- qa/L0_pytorch_unittest/test.sh | 1 + tests/pytorch/test_strided_batched_gemm.py | 317 ++++++++++++++++++ .../common/gemm/cublaslt_gemm.cu | 276 +++++++++++++++ .../common/include/transformer_engine/gemm.h | 51 +++ .../pytorch/cpp_extensions/gemm.py | 80 +++++ transformer_engine/pytorch/csrc/extensions.h | 7 + .../pytorch/csrc/extensions/gemm.cpp | 54 +++ .../pytorch/csrc/extensions/pybind.cpp | 7 + 8 files changed, 793 insertions(+) create mode 100644 tests/pytorch/test_strided_batched_gemm.py diff --git a/qa/L0_pytorch_unittest/test.sh b/qa/L0_pytorch_unittest/test.sh index 7370e42559..a7f704d73d 100644 --- a/qa/L0_pytorch_unittest/test.sh +++ b/qa/L0_pytorch_unittest/test.sh @@ -39,6 +39,7 @@ python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_jit.xml $TE_PATH 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_strided_batched_gemm.xml $TE_PATH/tests/pytorch/test_strided_batched_gemm.py || test_fail "test_strided_batched_gemm.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_weight_swizzle_in_layers.xml $TE_PATH/tests/pytorch/test_weight_swizzle_in_layers.py || test_fail "test_weight_swizzle_in_layers.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_torch_compile.xml $TE_PATH/tests/pytorch/test_torch_compile.py || test_fail "test_torch_compile.py" diff --git a/tests/pytorch/test_strided_batched_gemm.py b/tests/pytorch/test_strided_batched_gemm.py new file mode 100644 index 0000000000..2442a14ab8 --- /dev/null +++ b/tests/pytorch/test_strided_batched_gemm.py @@ -0,0 +1,317 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Tests for experimental strided batched GEMM.""" + +import pytest +import torch + +import transformer_engine.pytorch as te +from transformer_engine.pytorch import ( + Float8BlockQuantizer, + Float8Quantizer, + MXFP8Quantizer, + NVFP4Quantizer, + get_device_compute_capability, + is_bf16_available, +) +from transformer_engine.pytorch.cpp_extensions import general_gemm, strided_batched_gemm +from transformer_engine.pytorch.tensor.storage.float8_blockwise_tensor_storage import ( + Float8BlockwiseQTensorStorage, +) +from transformer_engine.pytorch.tensor.storage.float8_tensor_storage import Float8TensorStorage +from transformer_engine.pytorch.tensor.storage.mxfp8_tensor_storage import MXFP8TensorStorage +from transformer_engine.pytorch.tensor.storage.nvfp4_tensor_storage import NVFP4TensorStorage +import transformer_engine_torch as tex + +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) + +_UNSUPPORTED_INPUT_ERROR = ( + "strided_batched_gemm supports only high-precision or MXFP8 A and B tensor pairs" +) + + +def _skip_if_unavailable(recipe): + if not is_bf16_available(): + pytest.skip("bfloat16 is not available.") + if recipe == "fp8_block": + if not te.is_fp8_block_scaling_available(): + pytest.skip("FP8 block scaling is not available.") + if get_device_compute_capability() >= (10, 0): + pytest.skip("FP8 block scaling GEMM is emulated on Blackwell; test on Hopper.") + if recipe == "mxfp8" and not mxfp8_available: + pytest.skip(reason_for_no_mxfp8) + if recipe == "mxfp8" and tex.get_cublasLt_version() < 120800: + pytest.skip("MXFP8 GEMM requires cuBLASLt 12.8+.") + if recipe == "nvfp4" and not nvfp4_available: + pytest.skip(reason_for_no_nvfp4) + + +def _group_quantize(tensors, quantizer): + first_dims = torch.tensor([x.shape[0] for x in tensors], dtype=torch.int64, device="cuda") + grouped = tex.group_quantize(torch.cat(tensors, dim=0), quantizer, len(tensors), first_dims) + assert grouped._with_gemm_swizzled_scales + return grouped.split_into_quantized_tensors() + + +def _quantize_operands(recipe, tensors): + if recipe == "fp8": + quantized = [] + for i, tensor in enumerate(tensors): + quantizer = Float8Quantizer( + scale=torch.tensor([1.0 + 0.25 * i], dtype=torch.float32, device="cuda"), + amax=torch.zeros(1, dtype=torch.float32, device="cuda"), + fp8_dtype=tex.DType.kFloat8E4M3, + rowwise=True, + columnwise=False, + ) + quantized.append(quantizer(tensor)) + return quantized + + if recipe == "fp8_block": + quantizer = Float8BlockQuantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + rowwise=True, + columnwise=False, + force_pow_2_scales=True, + amax_epsilon=0.0, + block_scaling_dim=1, + ) + return [quantizer(tensor) for tensor in tensors] + + if recipe == "mxfp8": + quantizer = MXFP8Quantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + rowwise=True, + columnwise=False, + ) + quantizer.optimize_for_gemm = True + return _group_quantize(tensors, quantizer) + + if recipe == "nvfp4": + quantizer = NVFP4Quantizer( + fp4_dtype=tex.DType.kFloat4E2M1, + rowwise=True, + columnwise=False, + with_rht=True, + with_post_rht_amax=True, + with_2d_quantization=False, + stochastic_rounding=False, + with_random_sign_mask=False, + ) + quantizer.optimize_for_gemm = True + return _group_quantize(tensors, quantizer) + + raise ValueError(f"Unknown quantized recipe: {recipe}") + + +def _rowwise_data(tensor, recipe): + if recipe == "fp8": + return tensor._data + return tensor._rowwise_data + + +def _packed_storage(recipe, data, tensors, dtype): + quantizer = tensors[0]._quantizer + if recipe == "fp8": + scales = torch.cat([tensor._scale_inv.reshape(-1) for tensor in tensors]) + return Float8TensorStorage( + data=data, + fp8_scale_inv=scales, + fp8_dtype=tex.DType.kFloat8E4M3, + quantizer=quantizer, + fake_dtype=dtype, + ) + + scales = torch.cat([tensor._rowwise_scale_inv.reshape(-1) for tensor in tensors]) + if recipe == "fp8_block": + return Float8BlockwiseQTensorStorage( + rowwise_data=data, + rowwise_scale_inv=scales, + columnwise_data=None, + columnwise_scale_inv=None, + fp8_dtype=tex.DType.kFloat8E4M3, + quantizer=quantizer, + is_2D_scaled=False, + fake_dtype=dtype, + ) + if recipe == "mxfp8": + return MXFP8TensorStorage( + rowwise_data=data, + rowwise_scale_inv=scales, + columnwise_data=None, + columnwise_scale_inv=None, + fp8_dtype=tex.DType.kFloat8E4M3, + quantizer=quantizer, + with_gemm_swizzled_scales=True, + fake_dtype=dtype, + ) + if recipe == "nvfp4": + amax = torch.cat([tensor._amax_rowwise.reshape(-1) for tensor in tensors]) + return NVFP4TensorStorage( + rowwise_data=data, + rowwise_scale_inv=scales, + columnwise_data=None, + columnwise_scale_inv=None, + amax_rowwise=amax, + amax_columnwise=None, + fp4_dtype=tex.DType.kFloat4E2M1, + quantizer=quantizer, + with_gemm_swizzled_scales=True, + fake_dtype=dtype, + ) + raise ValueError(f"Unknown packed recipe: {recipe}") + + +def _make_operands(recipe, x, w, dtype): + seq, micro_batch, groups, hidden = x.shape + _, out_features, _ = w.shape + rows = seq * micro_batch + x_mats = [x[:, :, g, :].reshape(rows, hidden).contiguous() for g in range(groups)] + w_mats = [w[g].contiguous() for g in range(groups)] + + if recipe == "bf16": + return w, x, w_mats, x_mats + + w_quantized = _quantize_operands(recipe, w_mats) + x_quantized = _quantize_operands(recipe, x_mats) + packed_hidden = hidden // 2 if recipe == "nvfp4" else hidden + w_data = torch.cat([_rowwise_data(tensor, recipe).reshape(-1) for tensor in w_quantized]).view( + groups, out_features, packed_hidden + ) + x_data = torch.empty(seq, micro_batch, groups, packed_hidden, dtype=torch.uint8, device="cuda") + for g, tensor in enumerate(x_quantized): + x_data[:, :, g, :].copy_( + _rowwise_data(tensor, recipe).view(seq, micro_batch, packed_hidden) + ) + return ( + _packed_storage(recipe, w_data, w_quantized, dtype), + _packed_storage(recipe, x_data, x_quantized, dtype), + w_quantized, + x_quantized, + ) + + +@pytest.mark.parametrize("accumulate", [False, True]) +@pytest.mark.parametrize("recipe", ["bf16", "mxfp8"]) +def test_strided_batched_gemm_interleaved_activation(recipe, accumulate): + _skip_if_unavailable(recipe) + + torch.manual_seed(1234) + dtype = torch.bfloat16 + seq, micro_batch, groups, hidden, out_features = 16, 8, 3, 128, 128 + rows = seq * micro_batch + + # Logical operation: + # X: [S, B, G, D], W: [G, R, D] -> Y: [S, B, G, R] + # The GEMM view for each group is: + # X_g: [S*B, D], W_g: [R, D], Y_g: [S*B, R]. + x = torch.randn(seq, micro_batch, groups, hidden, dtype=dtype, device="cuda") + w = torch.randn(groups, out_features, hidden, dtype=dtype, device="cuda") + A, B, ref_A, ref_B = _make_operands(recipe, x, w, dtype) + + out = torch.empty(seq, micro_batch, groups, out_features, dtype=dtype, device="cuda") + out_initial = None + if accumulate: + out_initial = torch.randn_like(out) + out.copy_(out_initial) + + strided_batched_gemm( + A, + B, + out, + m=out_features, + n=rows, + k=hidden, + batch_count=groups, + lda=hidden, + stridea=out_features * hidden, + ldb=groups * hidden, + strideb=hidden, + ldd=groups * out_features, + strided=out_features, + layout="TN", + accumulate=accumulate, + ) + + ref = out_initial.clone() if accumulate else torch.empty_like(out) + for g in range(groups): + if accumulate: + out_g = ref[:, :, g, :].reshape(rows, out_features).contiguous() + else: + out_g = torch.empty(rows, out_features, dtype=dtype, device="cuda") + general_gemm( + ref_A[g], + ref_B[g], + out_dtype=dtype, + out=out_g, + layout="TN", + accumulate=accumulate, + ) + ref[:, :, g, :].copy_(out_g.view(seq, micro_batch, out_features)) + + torch.testing.assert_close(out, ref, rtol=0.125, atol=0.0675) + + +@pytest.mark.parametrize("api", ["python", "cpp"]) +@pytest.mark.parametrize("recipe", ["fp8", "fp8_block", "nvfp4"]) +def test_strided_batched_gemm_rejects_unsupported_inputs(recipe, api): + _skip_if_unavailable(recipe) + + torch.manual_seed(1234) + dtype = torch.bfloat16 + seq, micro_batch, groups, hidden, out_features = 16, 8, 2, 128, 128 + rows = seq * micro_batch + x = torch.randn(seq, micro_batch, groups, hidden, dtype=dtype, device="cuda") + w = torch.randn(groups, out_features, hidden, dtype=dtype, device="cuda") + A, B, _, _ = _make_operands(recipe, x, w, dtype) + out = torch.empty(seq, micro_batch, groups, out_features, dtype=dtype, device="cuda") + + if api == "python": + with pytest.raises(AssertionError, match=_UNSUPPORTED_INPUT_ERROR): + strided_batched_gemm( + A, + B, + out, + m=out_features, + n=rows, + k=hidden, + batch_count=groups, + lda=hidden, + stridea=out_features * hidden, + ldb=groups * hidden, + strideb=hidden, + ldd=groups * out_features, + strided=out_features, + layout="TN", + ) + return + + workspace = torch.empty(1, dtype=torch.uint8, device="cuda") + with pytest.raises(RuntimeError, match=_UNSUPPORTED_INPUT_ERROR): + tex.strided_batched_gemm( + A, + True, + B, + False, + out, + workspace, + workspace.numel(), + out_features, + rows, + hidden, + groups, + hidden, + out_features * hidden, + groups * hidden, + hidden, + groups * out_features, + out_features, + False, + False, + 0, + 1.0, + 0.0, + ) diff --git a/transformer_engine/common/gemm/cublaslt_gemm.cu b/transformer_engine/common/gemm/cublaslt_gemm.cu index a0529c80c0..174cb6d4b7 100644 --- a/transformer_engine/common/gemm/cublaslt_gemm.cu +++ b/transformer_engine/common/gemm/cublaslt_gemm.cu @@ -14,6 +14,7 @@ #include #include +#include #include #include @@ -798,6 +799,227 @@ void cublas_gemm(const Tensor *inputA, const Tensor *inputB, Tensor *outputD, NVTE_CHECK_CUBLAS(cublasLtMatmulDescDestroy(operationDesc)); } +void cublas_gemm_strided_batched(const Tensor *inputA, const Tensor *inputB, const Tensor *inputC, + Tensor *outputD, cublasOperation_t transa, + cublasOperation_t transb, void *workspace, size_t workspaceSize, + const void *alpha, const void *beta, bool use_split_accumulator, + int math_sm_count, int64_t lda, int64_t stridea, int64_t ldb, + int64_t strideb, int64_t ldc, int64_t stridec, int64_t ldd, + int64_t strided, int64_t batch_count, int64_t m, int64_t n, + int64_t k, cudaStream_t stream) { + NVTE_CHECK(inputA != nullptr && inputB != nullptr && inputC != nullptr && outputD != nullptr, + "Strided batched GEMM requires A, B, C, and D tensors."); + NVTE_CHECK(batch_count >= 0, "Strided batched GEMM got negative batch_count=", batch_count); + NVTE_CHECK(m >= 0 && n >= 0 && k >= 0, "Strided batched GEMM got invalid dims m=", m, ", n=", n, + ", k=", k); + if (batch_count == 0 || m == 0 || n == 0) { + return; + } + NVTE_CHECK(k > 0); + + auto checked_int = [](int64_t value, const char *name) { + NVTE_CHECK(value <= static_cast(std::numeric_limits::max()), name, + " is too large for cuBLASLt: ", value); + NVTE_CHECK(value >= 0, name, " must be non-negative, got ", value); + return static_cast(value); + }; + auto checked_positive_int = [&checked_int](int64_t value, const char *name) { + NVTE_CHECK(value > 0, name, " must be positive, got ", value); + return checked_int(value, name); + }; + const int m_int = checked_positive_int(m, "m"); + const int n_int = checked_positive_int(n, "n"); + const int k_int = checked_positive_int(k, "k"); + const int batch_count_int = checked_int(batch_count, "batch_count"); + const int lda_int = checked_positive_int(lda, "lda"); + const int ldb_int = checked_positive_int(ldb, "ldb"); + const int ldc_int = checked_positive_int(ldc, "ldc"); + const int ldd_int = checked_positive_int(ldd, "ldd"); + NVTE_CHECK(stridea >= 0 && strideb >= 0 && stridec >= 0 && strided >= 0, + "Strided batched GEMM requires non-negative batch strides (got stridea=", stridea, + ", strideb=", strideb, ", stridec=", stridec, ", strided=", strided, ")."); + + GemmParam param = CanonicalizeGemmInput(*inputA, transa, *inputB, transb, m_int, n_int, k_int); + param.lda = lda_int; + param.ldb = ldb_int; + + void *C = inputC->data.dptr; + void *D = outputD->data.dptr; + NVTE_CHECK(C != nullptr && D != nullptr, "Strided batched GEMM requires allocated C and D."); + + const bool high_precision_inputs = + is_high_precision_dtype(param.Atype) && is_high_precision_dtype(param.Btype); + const bool use_mxfp8 = is_fp8_dtype(param.Atype) && is_fp8_dtype(param.Btype) && + is_mxfp8_scaling(inputA->scaling_mode) && + is_mxfp8_scaling(inputB->scaling_mode); + NVTE_CHECK(high_precision_inputs || use_mxfp8, + "Strided batched GEMM supports only high-precision or MXFP8 A and B tensor pairs " + "(got A dtype=", + to_string(param.Atype), ", A scaling mode=", to_string(inputA->scaling_mode), + ", B dtype=", to_string(param.Btype), + ", B scaling mode=", to_string(inputB->scaling_mode), ")."); + if (use_mxfp8) { + NVTE_CHECK(param.A_scale_inv != nullptr && param.B_scale_inv != nullptr, + "MXFP8 inputs to strided batched GEMM require inverse scales."); + } + NVTE_CHECK(is_high_precision_dtype(outputD->data.dtype), + "Strided batched GEMM currently supports high-precision output only."); + + const cudaDataType_t A_type = get_cuda_dtype(param.Atype); + const cudaDataType_t B_type = get_cuda_dtype(param.Btype); + const cudaDataType_t C_type = get_cuda_dtype(inputC->data.dtype); + const cudaDataType_t D_type = get_cuda_dtype(outputD->data.dtype); + NVTE_CHECK(C_type == D_type, + "Strided batched GEMM currently requires C and D to have the same dtype."); + cublasLtHandle_t handle = cublasHandleManager::Instance().GetHandle(); + + cublasLtMatmulDesc_t operationDesc = nullptr; + cublasLtMatrixLayout_t Adesc = nullptr, Bdesc = nullptr, Cdesc = nullptr, Ddesc = nullptr; + cublasLtMatmulPreference_t preference = nullptr; + int returnedResults = 0; + cublasLtMatmulHeuristicResult_t heuristicResult = {}; + + auto set_strided_batch = [&](cublasLtMatrixLayout_t desc, int64_t stride) { + NVTE_CHECK_CUBLAS(cublasLtMatrixLayoutSetAttribute(desc, CUBLASLT_MATRIX_LAYOUT_BATCH_COUNT, + &batch_count_int, sizeof(batch_count_int))); + NVTE_CHECK_CUBLAS(cublasLtMatrixLayoutSetAttribute( + desc, CUBLASLT_MATRIX_LAYOUT_STRIDED_BATCH_OFFSET, &stride, sizeof(stride))); + }; + + NVTE_CHECK_CUBLAS( + cublasLtMatrixLayoutCreate(&Adesc, A_type, param.transA == CUBLAS_OP_N ? m_int : k_int, + param.transA == CUBLAS_OP_N ? k_int : m_int, param.lda)); + set_strided_batch(Adesc, stridea); + + NVTE_CHECK_CUBLAS( + cublasLtMatrixLayoutCreate(&Bdesc, B_type, param.transB == CUBLAS_OP_N ? k_int : n_int, + param.transB == CUBLAS_OP_N ? n_int : k_int, param.ldb)); + set_strided_batch(Bdesc, strideb); + + NVTE_CHECK_CUBLAS(cublasLtMatrixLayoutCreate(&Cdesc, C_type, m_int, n_int, ldc_int)); + set_strided_batch(Cdesc, stridec); + NVTE_CHECK_CUBLAS(cublasLtMatrixLayoutCreate(&Ddesc, D_type, m_int, n_int, ldd_int)); + set_strided_batch(Ddesc, strided); + + cublasComputeType_t gemm_compute_type = CUBLAS_COMPUTE_32F; + if (A_type == CUDA_R_32F && B_type == CUDA_R_32F && D_type == CUDA_R_32F) { + gemm_compute_type = CUBLAS_COMPUTE_32F_FAST_TF32; + } + NVTE_CHECK_CUBLAS(cublasLtMatmulDescCreate(&operationDesc, gemm_compute_type, CUDA_R_32F)); + NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(operationDesc, CUBLASLT_MATMUL_DESC_TRANSA, + ¶m.transA, sizeof(param.transA))); + NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(operationDesc, CUBLASLT_MATMUL_DESC_TRANSB, + ¶m.transB, sizeof(param.transB))); + if (math_sm_count != 0) { + NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(operationDesc, + CUBLASLT_MATMUL_DESC_SM_COUNT_TARGET, + &math_sm_count, sizeof(math_sm_count))); + } + + if (use_mxfp8) { + const int8_t fastAccuMode = use_split_accumulator ? 0 : 1; + NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(operationDesc, CUBLASLT_MATMUL_DESC_FAST_ACCUM, + &fastAccuMode, sizeof(fastAccuMode))); + } + + if (use_mxfp8) { +#if CUBLAS_VERSION >= 120800 + NVTE_CHECK(transformer_engine::cuda::cublas_version() >= 120800, + "MXFP8 strided batched GEMM requires cuBLAS 12.8+, but run-time cuBLAS version is ", + transformer_engine::cuda::cublas_version()); + NVTE_CHECK( + inputA->with_gemm_swizzled_scales, + "MXFP8 A scales for strided batched GEMM must be packed by batch and GEMM-swizzled."); + NVTE_CHECK( + inputB->with_gemm_swizzled_scales, + "MXFP8 B scales for strided batched GEMM must be packed by batch and GEMM-swizzled."); + + 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, + CUBLASLT_MATMUL_DESC_A_SCALE_POINTER, + &A_scale_inverse, sizeof(A_scale_inverse))); + NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(operationDesc, + CUBLASLT_MATMUL_DESC_B_SCALE_POINTER, + &B_scale_inverse, sizeof(B_scale_inverse))); + const cublasLtMatmulMatrixScale_t scaling_mode_a = CUBLASLT_MATMUL_MATRIX_SCALE_VEC32_UE8M0; + const cublasLtMatmulMatrixScale_t scaling_mode_b = CUBLASLT_MATMUL_MATRIX_SCALE_VEC32_UE8M0; + + NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute( + operationDesc, CUBLASLT_MATMUL_DESC_A_SCALE_MODE, &scaling_mode_a, sizeof(scaling_mode_a))); + NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute( + operationDesc, CUBLASLT_MATMUL_DESC_B_SCALE_MODE, &scaling_mode_b, sizeof(scaling_mode_b))); + + // Workaround for the same heuristic-cache issue handled in the non-batched MXFP8 GEMM. + 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, + sizeof(dummy_a_vec_stride))); + } +#else + NVTE_ERROR( + "MXFP8 strided batched GEMM requires cuBLAS 12.8+, but compile-time cuBLAS version is ", + CUBLAS_VERSION); +#endif // CUBLAS_VERSION >= 120800 + } + + cublasLtEpilogue_t epilogue = CUBLASLT_EPILOGUE_DEFAULT; + NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(operationDesc, CUBLASLT_MATMUL_DESC_EPILOGUE, + &epilogue, sizeof(epilogue))); + + uint8_t *aligned_workspace_ptr = nullptr; + if (workspaceSize > 0) { + NVTE_CHECK(workspace != nullptr, + "Strided batched GEMM got non-zero workspace size with null workspace pointer."); + constexpr uintptr_t required_alignment = 256; + const uintptr_t workspace_addr = reinterpret_cast(workspace); + const uintptr_t aligned_addr = (workspace_addr + required_alignment - 1) & + ~(required_alignment - static_cast(1)); + const size_t padding = aligned_addr - workspace_addr; + NVTE_CHECK(workspaceSize >= padding, "Workspace is too small to align to 256 bytes."); + aligned_workspace_ptr = reinterpret_cast(aligned_addr); + workspaceSize -= padding; + } + + NVTE_CHECK_CUBLAS(cublasLtMatmulPreferenceCreate(&preference)); + NVTE_CHECK_CUBLAS(cublasLtMatmulPreferenceSetAttribute( + preference, CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, &workspaceSize, sizeof(workspaceSize))); + const auto A_alignment = _getAlignment(reinterpret_cast(param.A)); + const auto B_alignment = _getAlignment(reinterpret_cast(param.B)); + const auto C_alignment = _getAlignment(reinterpret_cast(C)); + const auto D_alignment = _getAlignment(reinterpret_cast(D)); + NVTE_CHECK_CUBLAS(cublasLtMatmulPreferenceSetAttribute( + preference, CUBLASLT_MATMUL_PREF_MIN_ALIGNMENT_A_BYTES, &A_alignment, sizeof(A_alignment))); + NVTE_CHECK_CUBLAS(cublasLtMatmulPreferenceSetAttribute( + preference, CUBLASLT_MATMUL_PREF_MIN_ALIGNMENT_B_BYTES, &B_alignment, sizeof(B_alignment))); + NVTE_CHECK_CUBLAS(cublasLtMatmulPreferenceSetAttribute( + preference, CUBLASLT_MATMUL_PREF_MIN_ALIGNMENT_C_BYTES, &C_alignment, sizeof(C_alignment))); + NVTE_CHECK_CUBLAS(cublasLtMatmulPreferenceSetAttribute( + preference, CUBLASLT_MATMUL_PREF_MIN_ALIGNMENT_D_BYTES, &D_alignment, sizeof(D_alignment))); + + const auto status = + cublasLtMatmulAlgoGetHeuristic(handle, operationDesc, Adesc, Bdesc, Cdesc, Ddesc, preference, + 1, &heuristicResult, &returnedResults); + NVTE_CHECK(status != CUBLAS_STATUS_NOT_SUPPORTED, + "Unable to find suitable cuBLAS strided batched GEMM algorithm"); + NVTE_CHECK_CUBLAS(status); + if (returnedResults == 0) { + NVTE_ERROR("Unable to find any suitable strided batched GEMM algorithms"); + } + + NVTE_CHECK_CUBLAS(cublasLtMatmul(handle, operationDesc, alpha, param.A, Adesc, param.B, Bdesc, + beta, C, Cdesc, D, Ddesc, &heuristicResult.algo, + aligned_workspace_ptr, workspaceSize, stream)); + + NVTE_CHECK_CUBLAS(cublasLtMatmulPreferenceDestroy(preference)); + NVTE_CHECK_CUBLAS(cublasLtMatrixLayoutDestroy(Ddesc)); + NVTE_CHECK_CUBLAS(cublasLtMatrixLayoutDestroy(Cdesc)); + NVTE_CHECK_CUBLAS(cublasLtMatrixLayoutDestroy(Bdesc)); + NVTE_CHECK_CUBLAS(cublasLtMatrixLayoutDestroy(Adesc)); + NVTE_CHECK_CUBLAS(cublasLtMatmulDescDestroy(operationDesc)); +} + } // namespace transformer_engine void nvte_cublas_gemm(const NVTETensor A, const NVTETensor B, NVTETensor D, const NVTETensor bias, @@ -892,6 +1114,60 @@ void nvte_cublas_gemm_v2(int transa, int transb, const float *alpha, const NVTET config_.use_split_accumulator, config_.sm_count, 0, 0, false, nullptr, stream); } +void nvte_cublas_gemm_strided_batched(int transa, int transb, const float *alpha, + const NVTETensor A, int64_t lda, int64_t stridea, + const NVTETensor B, int64_t ldb, int64_t strideb, + const float *beta, const NVTETensor C, int64_t ldc, + int64_t stridec, NVTETensor D, int64_t ldd, int64_t strided, + int64_t batch_count, int64_t m, int64_t n, int64_t k, + NVTETensor workspace, NVTEMatmulConfig config, + cudaStream_t stream) { + NVTE_API_CALL(nvte_cublas_gemm_strided_batched); + using namespace transformer_engine; + + const Tensor *A_tensor = convertNVTETensorCheck(A); + const Tensor *B_tensor = convertNVTETensorCheck(B); + const Tensor *C_tensor = convertNVTETensorCheck(C); + Tensor *D_tensor = convertNVTETensorCheck(D); + + const bool high_precision_inputs = + is_high_precision_dtype(A_tensor->dtype()) && is_high_precision_dtype(B_tensor->dtype()); + const bool mxfp8_inputs = is_mxfp8_scaling(A_tensor->scaling_mode) && + is_mxfp8_scaling(B_tensor->scaling_mode) && + is_fp8_dtype(A_tensor->dtype()) && is_fp8_dtype(B_tensor->dtype()); + NVTE_CHECK(high_precision_inputs || mxfp8_inputs, + "nvte_cublas_gemm_strided_batched supports only high-precision or MXFP8 A and B " + "tensor pairs."); + if (mxfp8_inputs) { + NVTE_CHECK(A_tensor->with_gemm_swizzled_scales && B_tensor->with_gemm_swizzled_scales, + "nvte_cublas_gemm_strided_batched expects packed, GEMM-swizzled MXFP8 scales."); + } + + void *workspace_ptr = nullptr; + size_t workspace_size = 0; + Tensor *workspace_tensor = convertNVTETensor(workspace); + if (workspace_tensor != nullptr) { + workspace_ptr = workspace_tensor->data.dptr; + workspace_size = + get_buffer_size_bytes(workspace_tensor->data.numel(), workspace_tensor->data.dtype); + } + + MatmulConfig config_; + if (config != nullptr) { + config_ = *reinterpret_cast(config); + } + NVTE_CHECK(config_.bias_tensor == nullptr && config_.dbias_tensor == nullptr && + !config_.with_gelu_epilogue && !config_.with_dgelu_epilogue && + config_.epilogue_aux_tensor == nullptr, + "Strided batched GEMM does not support fused epilogues yet."); + + cublas_gemm_strided_batched( + A_tensor, B_tensor, C_tensor, D_tensor, transa ? CUBLAS_OP_T : CUBLAS_OP_N, + transb ? CUBLAS_OP_T : CUBLAS_OP_N, workspace_ptr, workspace_size, alpha, beta, + config_.use_split_accumulator, config_.sm_count, lda, stridea, ldb, strideb, ldc, stridec, + ldd, strided, batch_count, m, n, k, stream); +} + void nvte_cublas_gemm_scaled(const NVTETensor A, const NVTETensor B, NVTETensor D, const NVTETensor bias, NVTETensor pre_gelu_out, bool transa, bool transb, bool grad, NVTETensor workspace, float alpha, float beta, diff --git a/transformer_engine/common/include/transformer_engine/gemm.h b/transformer_engine/common/include/transformer_engine/gemm.h index a99e0946ef..3869f6d7b4 100644 --- a/transformer_engine/common/include/transformer_engine/gemm.h +++ b/transformer_engine/common/include/transformer_engine/gemm.h @@ -201,6 +201,57 @@ void nvte_cublas_gemm_v2(int transa, int transb, const float *alpha, const NVTET const NVTETensor B, const float *beta, const NVTETensor C, NVTETensor D, NVTETensor workspace, NVTEMatmulConfig config, cudaStream_t stream); +/*! \brief Compute a strided batched matrix multiplication. + * + * This is an experimental low-level interface. The batch dimension is + * explicitly specified instead of inferred from tensor shape. Matrix data uses + * cuBLASLt strided-batch layout attributes. A and B must both be high-precision + * tensors or both be MXFP8 tensors. MXFP8 scale tensors are packed consecutively + * by batch, i.e. [batch0_scales][batch1_scales]..., and must already be in the + * GEMM-swizzled layout. Per-tensor FP8, FP8 block scaling, and NVFP4 are not + * supported by this interface. + * + * Computes for each batch: + * - `D_i = alpha * op(A_i) * op(B_i) + beta * C_i` + * + * The `m`, `n`, and `k` dimensions follow the cuBLAS column-major convention + * used internally by Transformer Engine's GEMM path. Row-major output tensors + * with logical shape [N, M] are represented by a cuBLAS matrix with dimensions + * M-by-N. + * + * \param[in] transa Whether to transpose A matrix. + * \param[in] transb Whether to transpose B matrix. + * \param[in] alpha Scaling factor applied to matmul output. + * \param[in] A A matrix batch base. + * \param[in] lda Leading dimension for each A matrix. + * \param[in] stridea Element stride between A matrices. + * \param[in] B B matrix batch base. + * \param[in] ldb Leading dimension for each B matrix. + * \param[in] strideb Element stride between B matrices. + * \param[in] beta Scaling factor applied to C matrix. + * \param[in] C C matrix batch base. + * \param[in] ldc Leading dimension for each C matrix. + * \param[in] stridec Element stride between C matrices. + * \param[out] D Output matrix batch base. + * \param[in] ldd Leading dimension for each D matrix. + * \param[in] strided Element stride between D matrices. + * \param[in] batch_count Number of matrices in the batch. + * \param[in] m Number of rows of op(A) and D in cuBLAS layout. + * \param[in] n Number of columns of op(B) and D in cuBLAS layout. + * \param[in] k Reduction dimension. + * \param[in] workspace Workspace tensor. + * \param[in] config Additional configuration. + * \param[in] stream CUDA stream used for the operation. + */ +void nvte_cublas_gemm_strided_batched(int transa, int transb, const float *alpha, + const NVTETensor A, int64_t lda, int64_t stridea, + const NVTETensor B, int64_t ldb, int64_t strideb, + const float *beta, const NVTETensor C, int64_t ldc, + int64_t stridec, NVTETensor D, int64_t ldd, int64_t strided, + int64_t batch_count, int64_t m, int64_t n, int64_t k, + NVTETensor workspace, NVTEMatmulConfig config, + cudaStream_t stream); + /*! \brief Compute matrix multiplication of 2 matrices, potentially fused with other operations, * allowing for using a scaling factor for the GEMM result and the accumulation input (deprecated) * diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index 18451976ab..eb4b9ec684 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -29,6 +29,7 @@ __all__ = [ "general_gemm", + "strided_batched_gemm", "general_grouped_gemm", "general_grouped_gemm_for_grouped_tensor", ] @@ -376,6 +377,85 @@ def general_gemm( return out, bias_grad, gelu_input, extra_output +def strided_batched_gemm( + A: torch.Tensor, + B: torch.Tensor, + out: torch.Tensor, + *, + m: int, + n: int, + k: int, + batch_count: int, + lda: int, + stridea: int, + ldb: int, + strideb: int, + ldd: int, + strided: int, + layout: str = "TN", + accumulate: bool = False, + use_split_accumulator: bool = False, + alpha: float = 1.0, + beta: Optional[float] = None, +) -> torch.Tensor: + """Experimental strided batched GEMM. + + The shape is intentionally explicit. ``m``, ``n``, and ``k`` use the same + cuBLAS/TE convention as the C++ GEMM backend. For the common row-major + ``layout="TN"`` case, this computes per batch: + ``out[n, m] = B[n, k] @ A[m, k].T``. + + Inputs must both be high-precision tensors or both be MXFP8 tensors. MXFP8 + scale tensors are packed consecutively by batch and must already use the + GEMM-swizzled layout. + """ + assert layout in ("TN", "NN", "NT"), f"GEMM layout {layout} not supported." + high_precision_dtypes = (torch.float32, torch.float16, torch.bfloat16) + high_precision_inputs = all( + isinstance(tensor, torch.Tensor) + and not isinstance(tensor, QuantizedTensorStorage) + and tensor.dtype in high_precision_dtypes + for tensor in (A, B) + ) + mxfp8_inputs = isinstance(A, MXFP8TensorStorage) and isinstance(B, MXFP8TensorStorage) + assert ( + high_precision_inputs or mxfp8_inputs + ), "strided_batched_gemm supports only high-precision or MXFP8 A and B tensor pairs." + if mxfp8_inputs: + assert ( + A._with_gemm_swizzled_scales and B._with_gemm_swizzled_scales + ), "strided_batched_gemm expects packed, GEMM-swizzled MXFP8 scales." + transa = layout[0] == "T" + transb = layout[1] == "T" + beta = validate_gemm_scale(beta, accumulate) + workspace = get_cublas_workspace(out.device.index, False, False) + sm_count = get_sm_count() + return tex.strided_batched_gemm( + A, + transa, + B, + transb, + out, + workspace, + workspace.shape[0], + m, + n, + k, + batch_count, + lda, + stridea, + ldb, + strideb, + ldd, + strided, + accumulate, + use_split_accumulator, + sm_count - int(os.getenv("NVTE_EXT_MARGIN_SM", str(sm_count))), + alpha, + beta, + ) + + def general_grouped_gemm( A: List[torch.Tensor], B: List[torch.Tensor], diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h index 51f8b636b2..ec286e12bc 100644 --- a/transformer_engine/pytorch/csrc/extensions.h +++ b/transformer_engine/pytorch/csrc/extensions.h @@ -161,6 +161,13 @@ std::vector gemm(py::handle A, bool transa, py::handle B, bool trans MaybeTensor extra_output = std::nullopt, bool bulk_overlap = false, float alpha = 1.0f, std::optional beta = std::nullopt); +at::Tensor strided_batched_gemm(py::handle A, bool transa, py::handle B, bool transb, at::Tensor D, + at::Tensor workspace, size_t workspaceSize, int64_t m, int64_t n, + int64_t k, int64_t batch_count, int64_t lda, int64_t stridea, + int64_t ldb, int64_t strideb, int64_t ldd, int64_t strided, + bool accumulate, bool use_split_accumulator, int math_sm_count, + float alpha = 1.0f, std::optional beta = std::nullopt); + void te_atomic_gemm(at::Tensor A, at::Tensor A_scale_inverse, DType A_type, std::vector A_scaling_mode, bool transa, at::Tensor B, at::Tensor B_scale_inverse, DType B_type, std::vector B_scaling_mode, diff --git a/transformer_engine/pytorch/csrc/extensions/gemm.cpp b/transformer_engine/pytorch/csrc/extensions/gemm.cpp index b1e552ec8b..590600858d 100644 --- a/transformer_engine/pytorch/csrc/extensions/gemm.cpp +++ b/transformer_engine/pytorch/csrc/extensions/gemm.cpp @@ -413,6 +413,60 @@ std::vector gemm(py::handle A, bool transa, py::handle B, bool trans return out; } +at::Tensor strided_batched_gemm(py::handle A, bool transa, py::handle B, bool transb, at::Tensor D, + at::Tensor workspace, size_t workspaceSize, int64_t m, int64_t n, + int64_t k, int64_t batch_count, int64_t lda, int64_t stridea, + int64_t ldb, int64_t strideb, int64_t ldd, int64_t strided, + bool accumulate, bool use_split_accumulator, int math_sm_count, + float alpha, std::optional beta) { + NVTE_CHECK(!A.is_none(), "Tensor A has not been provided"); + NVTE_CHECK(!B.is_none(), "Tensor B has not been provided"); + NVTE_CHECK(D.is_cuda(), "Output tensor D must be on CUDA."); + NVTE_CHECK(workspace.is_cuda(), "cuBLASLt workspace must be on CUDA."); + + at::cuda::CUDAGuard device_guard(D.device()); + + if (accumulate) { + beta = beta.value_or(1.0f); + } else { + beta = beta.value_or(0.0f); + NVTE_CHECK(beta.value() == 0.0f, + "Trying to use non-zero beta while not accumulating into D tensor."); + } + const auto none = py::none(); + TensorWrapper A_tensor = makeTransformerEngineTensor(A, none); + TensorWrapper B_tensor = makeTransformerEngineTensor(B, none); + TensorWrapper D_tensor = makeTransformerEngineTensor(D); + auto te_workspace = makeTransformerEngineTensor(workspace.data_ptr(), + std::vector{workspaceSize}, DType::kByte); + + const bool high_precision_inputs = + is_high_precision_dtype(A_tensor.dtype()) && is_high_precision_dtype(B_tensor.dtype()); + const bool mxfp8_inputs = A_tensor.scaling_mode() == NVTE_MXFP8_1D_SCALING && + B_tensor.scaling_mode() == NVTE_MXFP8_1D_SCALING && + is_fp8_dtype(A_tensor.dtype()) && is_fp8_dtype(B_tensor.dtype()); + NVTE_CHECK(high_precision_inputs || mxfp8_inputs, + "strided_batched_gemm supports only high-precision or MXFP8 A and B tensor pairs."); + if (mxfp8_inputs) { + NVTE_CHECK(A_tensor.get_with_gemm_swizzled_scales() && B_tensor.get_with_gemm_swizzled_scales(), + "strided_batched_gemm expects packed, GEMM-swizzled MXFP8 scales."); + } + NVTE_CHECK(is_high_precision_dtype(D_tensor.dtype()), + "strided_batched_gemm currently expects a high-precision output tensor."); + + transformer_engine::MatmulConfigWrapper config; + config.set_use_split_accumulator(use_split_accumulator); + config.set_sm_count(math_sm_count); + + NVTE_SCOPED_GIL_RELEASE({ + nvte_cublas_gemm_strided_batched( + transa, transb, &alpha, A_tensor.data(), lda, stridea, B_tensor.data(), ldb, strideb, + &beta.value(), D_tensor.data(), ldd, strided, D_tensor.data(), ldd, strided, batch_count, m, + n, k, te_workspace.data(), config, at::cuda::getCurrentCUDAStream()); + }); + return D; +} + void te_atomic_gemm(at::Tensor A, at::Tensor A_scale_inverse, DType A_type, std::vector A_scaling_mode, bool transa, at::Tensor B, at::Tensor B_scale_inverse, DType B_type, std::vector B_scaling_mode, diff --git a/transformer_engine/pytorch/csrc/extensions/pybind.cpp b/transformer_engine/pytorch/csrc/extensions/pybind.cpp index c24e95e951..58af79295f 100644 --- a/transformer_engine/pytorch/csrc/extensions/pybind.cpp +++ b/transformer_engine/pytorch/csrc/extensions/pybind.cpp @@ -251,6 +251,13 @@ 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); + m.def("strided_batched_gemm", transformer_engine::pytorch::strided_batched_gemm, + "Compute experimental strided batched GEMM", py::arg("A"), py::arg("transA"), py::arg("B"), + py::arg("transB"), py::arg("D"), py::arg("workspace"), py::arg("workspace_size"), + py::arg("m"), py::arg("n"), py::arg("k"), py::arg("batch_count"), py::arg("lda"), + py::arg("stridea"), py::arg("ldb"), py::arg("strideb"), py::arg("ldd"), py::arg("strided"), + py::arg("accumulate"), py::arg("use_split_accumulator"), py::arg("math_sm_count"), + 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")); From c09a325280bc8e37a5daab86da13b3498120d221 Mon Sep 17 00:00:00 2001 From: Xin Yao Date: Tue, 30 Jun 2026 02:54:37 -0700 Subject: [PATCH 2/3] add ops.batched_linear Signed-off-by: Xin Yao --- docs/api/pytorch.rst | 2 + qa/L0_pytorch_unittest/test.sh | 1 + tests/pytorch/test_batched_linear.py | 453 +++++++++++ tests/pytorch/test_strided_batched_gemm.py | 364 +++++---- .../common/gemm/cublaslt_gemm.cu | 198 ++++- .../include/transformer_engine/swizzle.h | 35 + transformer_engine/common/swizzle/swizzle.cu | 220 +++++ .../pytorch/cpp_extensions/gemm.py | 12 +- .../pytorch/csrc/extensions/gemm.cpp | 151 +++- .../pytorch/ops/basic/__init__.py | 1 + .../pytorch/ops/basic/batched_linear.py | 758 ++++++++++++++++++ 11 files changed, 2014 insertions(+), 181 deletions(-) create mode 100644 tests/pytorch/test_batched_linear.py create mode 100644 transformer_engine/pytorch/ops/basic/batched_linear.py diff --git a/docs/api/pytorch.rst b/docs/api/pytorch.rst index 5fac0a89a6..44f3d61e6c 100644 --- a/docs/api/pytorch.rst +++ b/docs/api/pytorch.rst @@ -189,6 +189,8 @@ Operation fuser .. autoapiclass:: transformer_engine.pytorch.ops.BasicLinear :members: _functional_forward, _functional_backward +.. autoapiclass:: transformer_engine.pytorch.ops.BatchedLinear + .. autoapiclass:: transformer_engine.pytorch.ops.Bias .. autoapiclass:: transformer_engine.pytorch.ops.ClampedSwiGLU diff --git a/qa/L0_pytorch_unittest/test.sh b/qa/L0_pytorch_unittest/test.sh index a7f704d73d..7cfaae025c 100644 --- a/qa/L0_pytorch_unittest/test.sh +++ b/qa/L0_pytorch_unittest/test.sh @@ -40,6 +40,7 @@ python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_fused_rope.xml $ 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_strided_batched_gemm.xml $TE_PATH/tests/pytorch/test_strided_batched_gemm.py || test_fail "test_strided_batched_gemm.py" +python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_batched_linear.xml $TE_PATH/tests/pytorch/test_batched_linear.py || test_fail "test_batched_linear.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_weight_swizzle_in_layers.xml $TE_PATH/tests/pytorch/test_weight_swizzle_in_layers.py || test_fail "test_weight_swizzle_in_layers.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_torch_compile.xml $TE_PATH/tests/pytorch/test_torch_compile.py || test_fail "test_torch_compile.py" diff --git a/tests/pytorch/test_batched_linear.py b/tests/pytorch/test_batched_linear.py new file mode 100644 index 0000000000..afd1d25252 --- /dev/null +++ b/tests/pytorch/test_batched_linear.py @@ -0,0 +1,453 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Tests for the PyTorch BatchedLinear operation.""" + +from __future__ import annotations + +import contextlib + +import pytest +import torch + +from utils import assert_close, dtype_tols, make_recipe, quantization_tols + +import transformer_engine.common.recipe +import transformer_engine.pytorch as te +import transformer_engine.pytorch.ops as te_ops +import transformer_engine.pytorch.ops.basic.batched_linear as batched_linear_op +from transformer_engine.pytorch.quantization import FP8GlobalStateManager +from transformer_engine.pytorch.quantized_tensor import QuantizedTensorStorage + +mxfp8_available, reason_for_no_mxfp8 = te.is_mxfp8_available(return_reason=True) + + +@pytest.fixture(autouse=True) +def reset_global_fp8_state(): + """Keep FP8 global state isolated between tests.""" + yield + FP8GlobalStateManager.reset() + + +def _to_reference(tensor: torch.Tensor, *, requires_grad: bool) -> torch.Tensor: + """Make an FP64 CPU reference tensor.""" + if isinstance(tensor, QuantizedTensorStorage): + tensor = tensor.dequantize() + return tensor.detach().to(dtype=torch.float64, device="cpu").requires_grad_(requires_grad) + + +def _reference_linear( + inp: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + batch_dim: int, +) -> torch.Tensor: + """Reference batched linear with the two supported input layouts.""" + if batch_dim == 0: + out = torch.einsum("g...d,grd->g...r", inp, weight) + bias_shape = [bias.size(0)] + [1] * (out.ndim - 2) + [bias.size(1)] + return out + bias.view(bias_shape) + return torch.einsum("...gd,grd->...gr", inp, weight) + bias + + +@pytest.mark.parametrize("batch_dim", (0, -2)) +@pytest.mark.parametrize("quantized_compute", (False, True)) +@pytest.mark.parametrize("quantized_weight", (False, True)) +@pytest.mark.parametrize( + "accumulate_into_main_grad,overwrite_main_grad", + ((False, False), (True, False), (True, True)), +) +def test_forward_backward( + monkeypatch: pytest.MonkeyPatch, + *, + batch_dim: int, + quantized_compute: bool, + quantized_weight: bool, + accumulate_into_main_grad: bool, + overwrite_main_grad: bool, +) -> None: + """Check numerics, layouts, bias, and main-grad accumulation.""" + if (quantized_compute or quantized_weight) and not mxfp8_available: + pytest.skip(reason_for_no_mxfp8) + + dtype = torch.bfloat16 + device = torch.device("cuda") + num_gemms, rows0, rows1 = 3, 3, 32 + in_features, out_features = 160, 96 + if batch_dim == 0: + input_shape = (num_gemms, rows0, rows1, in_features) + output_shape = (num_gemms, rows0, rows1, out_features) + else: + input_shape = (rows0, rows1, num_gemms, in_features) + output_shape = (rows0, rows1, num_gemms, out_features) + + inp = torch.rand(input_shape, dtype=dtype, device=device, requires_grad=True) + weight = torch.rand( + num_gemms, + out_features, + in_features, + dtype=dtype, + device=device, + ) + bias = torch.rand(num_gemms, out_features, dtype=dtype, device=device) + grad_output = torch.rand( + *output_shape[:-1], + 2 * output_shape[-1], + dtype=dtype, + device=device, + )[..., ::2] + assert not grad_output.is_contiguous() + + recipe = make_recipe("mxfp8") + with te.quantized_model_init(enabled=quantized_weight, recipe=recipe): + op = te_ops.BatchedLinear( + num_gemms, + in_features, + out_features, + batch_dim=batch_dim, + accumulate_into_main_grad=accumulate_into_main_grad, + dtype=dtype, + device=device, + ) + with torch.no_grad(): + op.weight.copy_(weight) + op.bias.copy_(bias) + op.weight.main_grad = torch.full( + op.weight.shape, + 0.5, + dtype=torch.float32, + device=device, + ) + op.weight.overwrite_main_grad = overwrite_main_grad + + inp_ref = _to_reference(inp, requires_grad=True) + weight_ref = _to_reference(op.weight, requires_grad=True) + bias_ref = _to_reference(op.bias, requires_grad=True) + grad_output_ref = _to_reference(grad_output, requires_grad=False) + output_ref = _reference_linear(inp_ref, weight_ref, bias_ref, batch_dim) + output_ref.backward(grad_output_ref) + + gemm_calls = [] + original_strided_batched_gemm = batched_linear_op.strided_batched_gemm + + def capture_strided_batched_gemm(*args, **kwargs): + gemm_calls.append((kwargs["layout"], args[0], args[1], kwargs.get("accumulate", False))) + return original_strided_batched_gemm(*args, **kwargs) + + monkeypatch.setattr( + batched_linear_op, + "strided_batched_gemm", + capture_strided_batched_gemm, + ) + with te.autocast(enabled=quantized_compute, recipe=recipe): + output = op(inp) + output.backward(grad_output) + + assert tuple(output.shape) == output_shape + assert output.is_contiguous() + tols = ( + quantization_tols("mxfp8") if quantized_compute or quantized_weight else dtype_tols(dtype) + ) + assert_close(output, output_ref, **tols) + assert_close(inp.grad, inp_ref.grad, **tols) + assert_close(op.bias.grad, bias_ref.grad, **tols) + + if accumulate_into_main_grad: + assert op.weight.grad is None + grad_weight = op.weight.main_grad + if not overwrite_main_grad: + grad_weight = grad_weight - 0.5 + assert_close(grad_weight, weight_ref.grad, **tols) + else: + assert_close(op.weight.grad, weight_ref.grad, **tols) + torch.testing.assert_close( + op.weight.main_grad, + torch.full_like(op.weight.main_grad, 0.5), + rtol=0, + atol=0, + ) + + assert [layout for layout, _, _, _ in gemm_calls] == ["TN", "NN", "NT"] + assert gemm_calls[-1][3] == (accumulate_into_main_grad and not overwrite_main_grad) + if quantized_compute: + + def assert_mxfp8_usage(tensor, *, rowwise: bool, columnwise: bool) -> None: + assert (tensor._rowwise_data is not None) == rowwise + assert (tensor._columnwise_data is not None) == columnwise + + _, fprop_weight, fprop_input, _ = gemm_calls[0] + _, dgrad_weight, dgrad_grad_output, _ = gemm_calls[1] + _, wgrad_input, wgrad_grad_output, _ = gemm_calls[2] + assert_mxfp8_usage(fprop_weight, rowwise=True, columnwise=True) + assert_mxfp8_usage(fprop_input, rowwise=True, columnwise=True) + assert_mxfp8_usage(dgrad_weight, rowwise=True, columnwise=True) + assert_mxfp8_usage(dgrad_grad_output, rowwise=True, columnwise=True) + assert_mxfp8_usage(wgrad_input, rowwise=True, columnwise=True) + assert_mxfp8_usage(wgrad_grad_output, rowwise=True, columnwise=True) + assert dgrad_weight is fprop_weight + assert wgrad_input is fprop_input + assert dgrad_grad_output is wgrad_grad_output + if quantized_weight: + assert fprop_weight is op.weight + assert not op.weight._with_gemm_swizzled_scales + + +@pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8) +@pytest.mark.parametrize("batch_dim", (0, -2)) +@pytest.mark.parametrize( + "input_requires_grad,weight_requires_grad", + ((True, False), (False, True)), +) +def test_mxfp8_selective_gradients( + *, + batch_dim: int, + input_requires_grad: bool, + weight_requires_grad: bool, +) -> None: + """MXFP8 backward supports independently frozen inputs and weights.""" + recipe = make_recipe("mxfp8") + op = te_ops.BatchedLinear( + 2, + 64, + 32, + batch_dim=batch_dim, + bias=False, + dtype=torch.bfloat16, + device="cuda", + ) + op.weight.requires_grad_(weight_requires_grad) + input_shape = (2, 32, 64) if batch_dim == 0 else (32, 2, 64) + inp = torch.rand( + input_shape, + dtype=torch.bfloat16, + device="cuda", + requires_grad=input_requires_grad, + ) + + with te.autocast(enabled=True, recipe=recipe): + output = op(inp) + output.sum().backward() + + assert (inp.grad is not None) == input_requires_grad + assert (op.weight.grad is not None) == weight_requires_grad + + +@pytest.mark.parametrize("deferred_init", (False, True)) +def test_init_method_and_rng_tracker(deferred_init: bool) -> None: + """Initialization uses the requested RNG tracker, including after meta init.""" + events = [] + + class TestRNGTracker: + """Minimal CUDA RNG tracker interface for initialization tests.""" + + @contextlib.contextmanager + def fork(self): + """Record entry and exit around the initialization context.""" + events.append("enter") + try: + yield + finally: + events.append("exit") + + tracker = TestRNGTracker() + + def get_rng_tracker(): + events.append("tracker") + return tracker + + def init_method(weight): + assert events[-1] == "enter" + events.append("init") + torch.nn.init.constant_(weight, 0.25) + + op = te_ops.BatchedLinear( + 2, + 32, + 64, + device="meta" if deferred_init else "cuda", + dtype=torch.bfloat16, + rng_state_tracker_function=get_rng_tracker, + init_method=init_method, + ) + if deferred_init: + assert not events + assert all(param.device.type == "meta" for param in op.parameters()) + + inp = torch.randn(32, 2, 32, dtype=torch.bfloat16, device="cuda", requires_grad=True) + output = op(inp) + output.sum().backward() + + assert events == ["tracker", "enter", "init", "exit"] + torch.testing.assert_close( + op.weight, + torch.full_like(op.weight, 0.25), + rtol=0, + atol=0, + ) + torch.testing.assert_close(op.bias, torch.zeros_like(op.bias), rtol=0, atol=0) + assert op.weight.device.type == "cuda" + assert op.bias.device.type == "cuda" + assert op.weight.grad is not None and op.weight.grad.device.type == "cuda" + assert op.bias.grad is not None and op.bias.grad.device.type == "cuda" + + +@pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8) +def test_quantized_weight_deferred_init() -> None: + """Meta initialization materializes a primary MXFP8 weight on first use.""" + recipe = make_recipe("mxfp8") + with te.quantized_model_init(enabled=True, recipe=recipe): + op = te_ops.BatchedLinear( + 2, + 32, + 32, + device="meta", + dtype=torch.bfloat16, + ) + assert all(param.device.type == "meta" for param in op.parameters()) + + inp = torch.randn(32, 2, 32, dtype=torch.bfloat16, device="cuda", requires_grad=True) + with te.autocast(enabled=True, recipe=recipe): + output = op(inp) + output.sum().backward() + + assert isinstance(op.weight, te.QuantizedTensor) + assert op.weight.device.type == "cuda" + assert op.bias.device.type == "cuda" + assert op.weight.grad is not None and op.weight.grad.device.type == "cuda" + assert op.bias.grad is not None and op.bias.grad.device.type == "cuda" + assert not op.weight._with_gemm_swizzled_scales + + +@pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8) +def test_quantized_weight_inference() -> None: + """Inference primary weight needs only row-wise compact MXFP8 data.""" + recipe = make_recipe("mxfp8") + weight = torch.rand(2, 96, 160, device="cuda", dtype=torch.bfloat16) + inp = torch.rand(96, 2, 160, device="cuda", dtype=torch.bfloat16) + + with torch.no_grad(), te.quantized_model_init(enabled=True, recipe=recipe): + op = te_ops.BatchedLinear( + 2, + 160, + 96, + bias=False, + dtype=torch.bfloat16, + device="cuda", + ) + op.requires_grad_(False) + with torch.no_grad(): + op.weight.copy_(weight) + reference = torch.einsum("...gd,grd->...gr", inp, op.weight.dequantize()) + + assert op.weight._rowwise_data is not None + assert op.weight._columnwise_data is None + with torch.no_grad(), te.autocast(enabled=True, recipe=recipe): + output = op(inp) + assert_close(output, reference, **quantization_tols("mxfp8")) + assert not op.weight._with_gemm_swizzled_scales + + +@pytest.mark.parametrize("batch_dim", (0, -2)) +@pytest.mark.parametrize("bias", (False, True)) +def test_return_bias(batch_dim: int, bias: bool) -> None: + """return_bias leaves the bias unapplied and returns it separately.""" + op = te_ops.BatchedLinear( + 2, + 32, + 32, + batch_dim=batch_dim, + bias=bias, + return_bias=True, + dtype=torch.bfloat16, + device="cuda", + ) + input_shape = (2, 32, 32) if batch_dim == 0 else (32, 2, 32) + inp = torch.randn(input_shape, dtype=torch.bfloat16, device="cuda") + output, returned_bias = op(inp) + weight_ref = op.weight.detach() + if batch_dim == 0: + output_ref = torch.einsum("g...d,grd->g...r", inp, weight_ref) + else: + output_ref = torch.einsum("...gd,grd->...gr", inp, weight_ref) + assert_close(output, output_ref, **dtype_tols(torch.bfloat16)) + if bias: + assert_close(returned_bias, op.bias, **dtype_tols(torch.bfloat16)) + grad_bias = torch.randn_like(returned_bias) + torch.autograd.backward((output, returned_bias), (torch.ones_like(output), grad_bias)) + torch.testing.assert_close(op.bias.grad, grad_bias, rtol=0, atol=0) + else: + assert returned_bias is None + + +def test_return_bias_in_sequential() -> None: + """The bias is exposed through the operation fuser's extra-output path.""" + op = te_ops.BatchedLinear( + 2, + 32, + 32, + bias=True, + return_bias=True, + dtype=torch.bfloat16, + device="cuda", + ) + model = te_ops.Sequential(op) + inp = torch.randn(32, 2, 32, dtype=torch.bfloat16, device="cuda") + output, returned_bias = model(inp) + assert_close(returned_bias, op.bias, **dtype_tols(torch.bfloat16)) + (output.sum() + returned_bias.sum()).backward() + torch.testing.assert_close(op.bias.grad, torch.ones_like(op.bias), rtol=0, atol=0) + + +@pytest.mark.parametrize( + "recipe", + ( + transformer_engine.common.recipe.DelayedScaling(), + transformer_engine.common.recipe.Float8CurrentScaling(), + transformer_engine.common.recipe.Float8BlockScaling(), + transformer_engine.common.recipe.NVFP4BlockScaling(), + ), +) +def test_rejects_non_mxfp8_recipe(recipe) -> None: + """Only high precision and MXFP8 compute are accepted.""" + with pytest.raises(ValueError, match="only high-precision compute or the MXFP8 recipe"): + with te.quantized_model_init(enabled=True, recipe=recipe): + te_ops.BatchedLinear( + 2, + 32, + 32, + bias=False, + dtype=torch.bfloat16, + device="cuda", + ) + + +@pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8) +@pytest.mark.parametrize("backward_override", ("high_precision", "dequantized")) +def test_rejects_backward_override(backward_override: str) -> None: + """MXFP8 backward overrides are not implemented.""" + op = te_ops.BatchedLinear( + 2, + 32, + 32, + bias=False, + dtype=torch.bfloat16, + device="cuda", + ) + inp = torch.randn(32, 2, 32, dtype=torch.bfloat16, device="cuda") + recipe = transformer_engine.common.recipe.MXFP8BlockScaling(backward_override=backward_override) + with pytest.raises(ValueError, match="does not support MXFP8 backward_override"): + with te.autocast(enabled=True, recipe=recipe): + op(inp) + + +def test_rejects_unsupported_layout() -> None: + """Reject ambiguous batch dimensions and non-contiguous input storage.""" + with pytest.raises(ValueError, match="batch_dim=0 or batch_dim=-2"): + te_ops.BatchedLinear(2, 32, 32, batch_dim=1, device="cuda") + + op = te_ops.BatchedLinear(2, 32, 32, batch_dim=-2, device="cuda") + inp = torch.empty(32, 2, 64, device="cuda")[:, :, ::2] + assert inp.shape == (32, 2, 32) and not inp.is_contiguous() + with pytest.raises(ValueError, match="contiguous input"): + op(inp) diff --git a/tests/pytorch/test_strided_batched_gemm.py b/tests/pytorch/test_strided_batched_gemm.py index 2442a14ab8..0e735737c8 100644 --- a/tests/pytorch/test_strided_batched_gemm.py +++ b/tests/pytorch/test_strided_batched_gemm.py @@ -17,12 +17,6 @@ is_bf16_available, ) from transformer_engine.pytorch.cpp_extensions import general_gemm, strided_batched_gemm -from transformer_engine.pytorch.tensor.storage.float8_blockwise_tensor_storage import ( - Float8BlockwiseQTensorStorage, -) -from transformer_engine.pytorch.tensor.storage.float8_tensor_storage import Float8TensorStorage -from transformer_engine.pytorch.tensor.storage.mxfp8_tensor_storage import MXFP8TensorStorage -from transformer_engine.pytorch.tensor.storage.nvfp4_tensor_storage import NVFP4TensorStorage import transformer_engine_torch as tex mxfp8_available, reason_for_no_mxfp8 = te.is_mxfp8_available(return_reason=True) @@ -31,6 +25,7 @@ _UNSUPPORTED_INPUT_ERROR = ( "strided_batched_gemm supports only high-precision or MXFP8 A and B tensor pairs" ) +_SWIZZLED_MXFP8_ERROR = "strided_batched_gemm expects compact MXFP8 scales" def _skip_if_unavailable(recipe): @@ -49,48 +44,52 @@ def _skip_if_unavailable(recipe): pytest.skip(reason_for_no_nvfp4) -def _group_quantize(tensors, quantizer): - first_dims = torch.tensor([x.shape[0] for x in tensors], dtype=torch.int64, device="cuda") - grouped = tex.group_quantize(torch.cat(tensors, dim=0), quantizer, len(tensors), first_dims) - assert grouped._with_gemm_swizzled_scales - return grouped.split_into_quantized_tensors() +def _quantize_mxfp8_batched(tensor, batch_dim, *, rowwise=True, columnwise=False): + quantizer = MXFP8Quantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + rowwise=rowwise, + columnwise=columnwise, + ) + quantizer.internal = True + quantizer.optimize_for_gemm = False + quantizer_input = tensor + if columnwise and batch_dim == -2: + groups = tensor.size(0 if batch_dim == 0 else tensor.ndim - 2) + hidden = tensor.size(-1) + rows = tensor.numel() // (groups * hidden) + quantizer_input = tensor.view(rows, groups * hidden) + return quantizer(quantizer_input) + + +def _quantize_mxfp8_reference(tensors, *, rowwise=True, columnwise=False): + quantizer = MXFP8Quantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + rowwise=rowwise, + columnwise=columnwise, + ) + quantizer.optimize_for_gemm = True + return [quantizer(tensor) for tensor in tensors] -def _quantize_operands(recipe, tensors): +def _quantize_unsupported_operands(recipe, x, w): if recipe == "fp8": - quantized = [] - for i, tensor in enumerate(tensors): - quantizer = Float8Quantizer( - scale=torch.tensor([1.0 + 0.25 * i], dtype=torch.float32, device="cuda"), - amax=torch.zeros(1, dtype=torch.float32, device="cuda"), - fp8_dtype=tex.DType.kFloat8E4M3, - rowwise=True, - columnwise=False, - ) - quantized.append(quantizer(tensor)) - return quantized - - if recipe == "fp8_block": - quantizer = Float8BlockQuantizer( + 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, rowwise=True, columnwise=False, - force_pow_2_scales=True, - amax_epsilon=0.0, - block_scaling_dim=1, ) - return [quantizer(tensor) for tensor in tensors] - - if recipe == "mxfp8": - quantizer = MXFP8Quantizer( + elif recipe == "fp8_block": + quantizer = Float8BlockQuantizer( fp8_dtype=tex.DType.kFloat8E4M3, rowwise=True, columnwise=False, + force_pow_2_scales=True, + amax_epsilon=0.0, + block_scaling_dim=1, ) - quantizer.optimize_for_gemm = True - return _group_quantize(tensors, quantizer) - - if recipe == "nvfp4": + elif recipe == "nvfp4": quantizer = NVFP4Quantizer( fp4_dtype=tex.DType.kFloat4E2M1, rowwise=True, @@ -102,95 +101,51 @@ def _quantize_operands(recipe, tensors): with_random_sign_mask=False, ) quantizer.optimize_for_gemm = True - return _group_quantize(tensors, quantizer) - - raise ValueError(f"Unknown quantized recipe: {recipe}") - - -def _rowwise_data(tensor, recipe): - if recipe == "fp8": - return tensor._data - return tensor._rowwise_data - - -def _packed_storage(recipe, data, tensors, dtype): - quantizer = tensors[0]._quantizer - if recipe == "fp8": - scales = torch.cat([tensor._scale_inv.reshape(-1) for tensor in tensors]) - return Float8TensorStorage( - data=data, - fp8_scale_inv=scales, - fp8_dtype=tex.DType.kFloat8E4M3, - quantizer=quantizer, - fake_dtype=dtype, - ) - - scales = torch.cat([tensor._rowwise_scale_inv.reshape(-1) for tensor in tensors]) - if recipe == "fp8_block": - return Float8BlockwiseQTensorStorage( - rowwise_data=data, - rowwise_scale_inv=scales, - columnwise_data=None, - columnwise_scale_inv=None, - fp8_dtype=tex.DType.kFloat8E4M3, - quantizer=quantizer, - is_2D_scaled=False, - fake_dtype=dtype, - ) - if recipe == "mxfp8": - return MXFP8TensorStorage( - rowwise_data=data, - rowwise_scale_inv=scales, - columnwise_data=None, - columnwise_scale_inv=None, - fp8_dtype=tex.DType.kFloat8E4M3, - quantizer=quantizer, - with_gemm_swizzled_scales=True, - fake_dtype=dtype, - ) - if recipe == "nvfp4": - amax = torch.cat([tensor._amax_rowwise.reshape(-1) for tensor in tensors]) - return NVFP4TensorStorage( - rowwise_data=data, - rowwise_scale_inv=scales, - columnwise_data=None, - columnwise_scale_inv=None, - amax_rowwise=amax, - amax_columnwise=None, - fp4_dtype=tex.DType.kFloat4E2M1, - quantizer=quantizer, - with_gemm_swizzled_scales=True, - fake_dtype=dtype, - ) - raise ValueError(f"Unknown packed recipe: {recipe}") - - -def _make_operands(recipe, x, w, dtype): - seq, micro_batch, groups, hidden = x.shape - _, out_features, _ = w.shape - rows = seq * micro_batch - x_mats = [x[:, :, g, :].reshape(rows, hidden).contiguous() for g in range(groups)] - w_mats = [w[g].contiguous() for g in range(groups)] - - if recipe == "bf16": - return w, x, w_mats, x_mats - - w_quantized = _quantize_operands(recipe, w_mats) - x_quantized = _quantize_operands(recipe, x_mats) - packed_hidden = hidden // 2 if recipe == "nvfp4" else hidden - w_data = torch.cat([_rowwise_data(tensor, recipe).reshape(-1) for tensor in w_quantized]).view( - groups, out_features, packed_hidden - ) - x_data = torch.empty(seq, micro_batch, groups, packed_hidden, dtype=torch.uint8, device="cuda") - for g, tensor in enumerate(x_quantized): - x_data[:, :, g, :].copy_( - _rowwise_data(tensor, recipe).view(seq, micro_batch, packed_hidden) - ) - return ( - _packed_storage(recipe, w_data, w_quantized, dtype), - _packed_storage(recipe, x_data, x_quantized, dtype), - w_quantized, - x_quantized, + else: + raise ValueError(f"Unknown quantized recipe: {recipe}") + return quantizer(w), quantizer(x) + + +def _cpp_strided_batched_gemm( + A, + B, + out, + *, + m, + n, + k, + batch_count, + lda, + stridea, + ldb, + strideb, + ldd, + strided, +): + workspace = torch.empty(1, dtype=torch.uint8, device="cuda") + return tex.strided_batched_gemm( + A, + True, + B, + False, + out, + workspace, + workspace.numel(), + m, + n, + k, + batch_count, + lda, + stridea, + ldb, + strideb, + ldd, + strided, + False, + False, + 0, + 1.0, + 0.0, ) @@ -201,7 +156,7 @@ def test_strided_batched_gemm_interleaved_activation(recipe, accumulate): torch.manual_seed(1234) dtype = torch.bfloat16 - seq, micro_batch, groups, hidden, out_features = 16, 8, 3, 128, 128 + seq, micro_batch, groups, hidden, out_features = 12, 8, 3, 160, 96 rows = seq * micro_batch # Logical operation: @@ -210,7 +165,19 @@ def test_strided_batched_gemm_interleaved_activation(recipe, accumulate): # X_g: [S*B, D], W_g: [R, D], Y_g: [S*B, R]. x = torch.randn(seq, micro_batch, groups, hidden, dtype=dtype, device="cuda") w = torch.randn(groups, out_features, hidden, dtype=dtype, device="cuda") - A, B, ref_A, ref_B = _make_operands(recipe, x, w, dtype) + x_mats = [x[:, :, g, :].reshape(rows, hidden).contiguous() for g in range(groups)] + w_mats = list(w.unbind(dim=0)) + if recipe == "bf16": + A, B = w, x + ref_A, ref_B = w_mats, x_mats + else: + A = _quantize_mxfp8_batched(w, batch_dim=0) + B = _quantize_mxfp8_batched(x, batch_dim=-2) + scale_ptrs = (A._rowwise_scale_inv.data_ptr(), B._rowwise_scale_inv.data_ptr()) + assert not A._with_gemm_swizzled_scales + assert not B._with_gemm_swizzled_scales + ref_A = _quantize_mxfp8_reference(w_mats) + ref_B = _quantize_mxfp8_reference(x_mats) out = torch.empty(seq, micro_batch, groups, out_features, dtype=dtype, device="cuda") out_initial = None @@ -236,6 +203,11 @@ def test_strided_batched_gemm_interleaved_activation(recipe, accumulate): accumulate=accumulate, ) + if recipe == "mxfp8": + assert (A._rowwise_scale_inv.data_ptr(), B._rowwise_scale_inv.data_ptr()) == scale_ptrs + assert not A._with_gemm_swizzled_scales + assert not B._with_gemm_swizzled_scales + ref = out_initial.clone() if accumulate else torch.empty_like(out) for g in range(groups): if accumulate: @@ -255,6 +227,40 @@ def test_strided_batched_gemm_interleaved_activation(recipe, accumulate): torch.testing.assert_close(out, ref, rtol=0.125, atol=0.0675) +@pytest.mark.parametrize("api", ["python", "cpp"]) +def test_strided_batched_gemm_rejects_swizzled_mxfp8_scales(api): + _skip_if_unavailable("mxfp8") + + torch.manual_seed(1234) + dtype = torch.bfloat16 + seq, micro_batch, groups, hidden, out_features = 12, 8, 3, 160, 96 + rows = seq * micro_batch + x = torch.randn(seq, micro_batch, groups, hidden, dtype=dtype, device="cuda") + w = torch.randn(groups, out_features, hidden, dtype=dtype, device="cuda") + A, B = _quantize_mxfp8_reference([w, x]) + assert A._with_gemm_swizzled_scales and B._with_gemm_swizzled_scales + out = torch.empty(seq, micro_batch, groups, out_features, dtype=dtype, device="cuda") + + kwargs = { + "m": out_features, + "n": rows, + "k": hidden, + "batch_count": groups, + "lda": hidden, + "stridea": out_features * hidden, + "ldb": groups * hidden, + "strideb": hidden, + "ldd": groups * out_features, + "strided": out_features, + } + if api == "python": + with pytest.raises(AssertionError, match=_SWIZZLED_MXFP8_ERROR): + strided_batched_gemm(A, B, out, layout="TN", **kwargs) + else: + with pytest.raises(RuntimeError, match=_SWIZZLED_MXFP8_ERROR): + _cpp_strided_batched_gemm(A, B, out, **kwargs) + + @pytest.mark.parametrize("api", ["python", "cpp"]) @pytest.mark.parametrize("recipe", ["fp8", "fp8_block", "nvfp4"]) def test_strided_batched_gemm_rejects_unsupported_inputs(recipe, api): @@ -266,7 +272,7 @@ def test_strided_batched_gemm_rejects_unsupported_inputs(recipe, api): rows = seq * micro_batch x = torch.randn(seq, micro_batch, groups, hidden, dtype=dtype, device="cuda") w = torch.randn(groups, out_features, hidden, dtype=dtype, device="cuda") - A, B, _, _ = _make_operands(recipe, x, w, dtype) + A, B = _quantize_unsupported_operands(recipe, x, w) out = torch.empty(seq, micro_batch, groups, out_features, dtype=dtype, device="cuda") if api == "python": @@ -289,29 +295,87 @@ def test_strided_batched_gemm_rejects_unsupported_inputs(recipe, api): ) return - workspace = torch.empty(1, dtype=torch.uint8, device="cuda") with pytest.raises(RuntimeError, match=_UNSUPPORTED_INPUT_ERROR): - tex.strided_batched_gemm( + _cpp_strided_batched_gemm( A, - True, B, - False, out, - workspace, - workspace.numel(), - out_features, - rows, - hidden, - groups, - hidden, - out_features * hidden, - groups * hidden, - hidden, - groups * out_features, - out_features, - False, - False, - 0, - 1.0, - 0.0, + m=out_features, + n=rows, + k=hidden, + batch_count=groups, + lda=hidden, + stridea=out_features * hidden, + ldb=groups * hidden, + strideb=hidden, + ldd=groups * out_features, + strided=out_features, ) + + +@pytest.mark.parametrize("api", ["python", "cpp"]) +@pytest.mark.parametrize( + "buffer_name,stride_name", + [("A", "stridea"), ("B", "strideb"), ("D", "strided")], +) +def test_strided_batched_gemm_rejects_out_of_bounds_layout(api, buffer_name, stride_name): + _skip_if_unavailable("bf16") + + groups, m, n, k = 2, 32, 64, 96 + A = torch.randn(groups, m, k, dtype=torch.bfloat16, device="cuda") + B = torch.randn(groups, n, k, dtype=torch.bfloat16, device="cuda") + out = torch.empty(groups, n, m, dtype=torch.bfloat16, device="cuda") + kwargs = { + "m": m, + "n": n, + "k": k, + "batch_count": groups, + "lda": k, + "stridea": m * k, + "ldb": k, + "strideb": n * k, + "ldd": m, + "strided": n * m, + } + kwargs[stride_name] += 1 + + call = strided_batched_gemm if api == "python" else _cpp_strided_batched_gemm + call_kwargs = {**kwargs, "layout": "TN"} if api == "python" else kwargs + with pytest.raises(RuntimeError, match=rf"{buffer_name} data buffer is too small"): + call(A, B, out, **call_kwargs) + + +@pytest.mark.parametrize("api", ["python", "cpp"]) +@pytest.mark.parametrize("buffer_name", ["A", "B"]) +def test_strided_batched_gemm_rejects_invalid_compact_scale_shape(api, buffer_name): + _skip_if_unavailable("mxfp8") + + groups, m, n, k = 2, 128, 128, 128 + A = _quantize_mxfp8_batched( + torch.randn(groups, m, k, dtype=torch.bfloat16, device="cuda"), + batch_dim=0, + ) + B = _quantize_mxfp8_batched( + torch.randn(groups, n, k, dtype=torch.bfloat16, device="cuda"), + batch_dim=0, + ) + operand = A if buffer_name == "A" else B + operand._rowwise_scale_inv = operand._rowwise_scale_inv.reshape(-1)[:-1] + out = torch.empty(groups, n, m, dtype=torch.bfloat16, device="cuda") + kwargs = { + "m": m, + "n": n, + "k": k, + "batch_count": groups, + "lda": k, + "stridea": m * k, + "ldb": k, + "strideb": n * k, + "ldd": m, + "strided": n * m, + } + + call = strided_batched_gemm if api == "python" else _cpp_strided_batched_gemm + call_kwargs = {**kwargs, "layout": "TN"} if api == "python" else kwargs + with pytest.raises(RuntimeError, match="expects 2D compact scaling factors"): + call(A, B, out, **call_kwargs) diff --git a/transformer_engine/common/gemm/cublaslt_gemm.cu b/transformer_engine/common/gemm/cublaslt_gemm.cu index 174cb6d4b7..4d7da2dd0e 100644 --- a/transformer_engine/common/gemm/cublaslt_gemm.cu +++ b/transformer_engine/common/gemm/cublaslt_gemm.cu @@ -93,6 +93,128 @@ struct GemmParam { int ldb = 0; // B column strides }; +struct CublasLtDescriptorCleanup { + cublasLtMatmulDesc_t operation_desc = nullptr; + cublasLtMatrixLayout_t a_desc = nullptr; + cublasLtMatrixLayout_t b_desc = nullptr; + cublasLtMatrixLayout_t c_desc = nullptr; + cublasLtMatrixLayout_t d_desc = nullptr; + cublasLtMatmulPreference_t preference = nullptr; + + ~CublasLtDescriptorCleanup() noexcept { + // Cleanup must not throw while unwinding an earlier cuBLASLt error. + if (preference != nullptr) { + static_cast(cublasLtMatmulPreferenceDestroy(preference)); + } + if (d_desc != nullptr) { + static_cast(cublasLtMatrixLayoutDestroy(d_desc)); + } + if (c_desc != nullptr) { + static_cast(cublasLtMatrixLayoutDestroy(c_desc)); + } + if (b_desc != nullptr) { + static_cast(cublasLtMatrixLayoutDestroy(b_desc)); + } + if (a_desc != nullptr) { + static_cast(cublasLtMatrixLayoutDestroy(a_desc)); + } + if (operation_desc != nullptr) { + static_cast(cublasLtMatmulDescDestroy(operation_desc)); + } + } +}; + +size_t CheckedSizeMul(size_t lhs, size_t rhs, const char *description) { + NVTE_CHECK(lhs == 0 || rhs <= std::numeric_limits::max() / lhs, + "Integer overflow while calculating ", description, "."); + return lhs * rhs; +} + +size_t CheckedSizeAdd(size_t lhs, size_t rhs, const char *description) { + NVTE_CHECK(rhs <= std::numeric_limits::max() - lhs, "Integer overflow while calculating ", + description, "."); + return lhs + rhs; +} + +size_t RoundUpToMultiple(size_t value, size_t multiple, const char *description) { + const size_t quotient = value / multiple + static_cast(value % multiple != 0); + return CheckedSizeMul(quotient, multiple, description); +} + +size_t GetStridedBatchSpan(int64_t rows, int64_t cols, int64_t leading_dim, int64_t batch_stride, + int64_t batch_count, const char *name) { + NVTE_CHECK(leading_dim >= rows, "Strided batched GEMM ", name, " leading dimension ", leading_dim, + " is smaller than the matrix row count ", rows, "."); + const auto matrix_span = CheckedSizeAdd( + CheckedSizeMul(static_cast(cols - 1), static_cast(leading_dim), name), + static_cast(rows), name); + return CheckedSizeAdd( + CheckedSizeMul(static_cast(batch_count - 1), static_cast(batch_stride), name), + matrix_span, name); +} + +const transformer_engine::SimpleTensor &GetGemmDataBuffer(const transformer_engine::Tensor &tensor, + const void *data_ptr, const char *name) { + if (data_ptr == tensor.data.dptr) { + return tensor.data; + } + if (data_ptr == tensor.columnwise_data.dptr) { + return tensor.columnwise_data; + } + NVTE_ERROR("Strided batched GEMM could not identify the selected ", name, " data buffer."); +} + +const transformer_engine::SimpleTensor &GetGemmScaleBuffer(const transformer_engine::Tensor &tensor, + const void *scale_ptr, bool *rowwise, + const char *name) { + if (scale_ptr == tensor.scale_inv.dptr) { + *rowwise = true; + return tensor.scale_inv; + } + if (scale_ptr == tensor.columnwise_scale_inv.dptr) { + *rowwise = false; + return tensor.columnwise_scale_inv; + } + NVTE_ERROR("Strided batched GEMM could not identify the selected ", name, " scale buffer."); +} + +void CheckStridedBatchBuffer(const transformer_engine::SimpleTensor &buffer, int64_t rows, + int64_t cols, int64_t leading_dim, int64_t batch_stride, + int64_t batch_count, const char *name) { + NVTE_CHECK(buffer.dptr != nullptr, "Strided batched GEMM ", name, " buffer is not allocated."); + const size_t required_span = + GetStridedBatchSpan(rows, cols, leading_dim, batch_stride, batch_count, name); + NVTE_CHECK(required_span <= buffer.numel(), "Strided batched GEMM ", name, + " buffer is too small for the requested layout (required span=", required_span, + " elements, available=", buffer.numel(), ")."); +} + +void CheckMXFP8BatchScaleBuffer(const transformer_engine::SimpleTensor &buffer, bool rowwise, + size_t rows_per_batch, size_t features, size_t batch_count, + const char *name) { + constexpr size_t block_size = 32; + size_t scale_rows = 0; + size_t scale_cols = 0; + if (rowwise) { + NVTE_CHECK(features % block_size == 0, "Strided batched GEMM ", name, + " row-wise scale feature count must be divisible by ", block_size, " (got ", + features, ")."); + scale_rows = RoundUpToMultiple(rows_per_batch, 128, name); + scale_cols = RoundUpToMultiple(features / block_size, 4, name); + } else { + NVTE_CHECK(rows_per_batch % block_size == 0, "Strided batched GEMM ", name, + " column-wise scale row count must be divisible by ", block_size, " (got ", + rows_per_batch, ")."); + scale_rows = RoundUpToMultiple(rows_per_batch / block_size, 4, name); + scale_cols = RoundUpToMultiple(features, 128, name); + } + const size_t expected_numel = + CheckedSizeMul(batch_count, CheckedSizeMul(scale_rows, scale_cols, name), name); + NVTE_CHECK(buffer.numel() == expected_numel, "Strided batched GEMM ", name, + " packed scale buffer has invalid size (expected ", expected_numel, ", got ", + buffer.numel(), ")."); +} + /* Populate parameters for cuBLAS GEMM * * cuBLAS follows the BLAS convention of column-major ordering. This @@ -861,6 +983,21 @@ void cublas_gemm_strided_batched(const Tensor *inputA, const Tensor *inputB, con if (use_mxfp8) { NVTE_CHECK(param.A_scale_inv != nullptr && param.B_scale_inv != nullptr, "MXFP8 inputs to strided batched GEMM require inverse scales."); +#if CUBLAS_VERSION >= 120800 + NVTE_CHECK(transformer_engine::cuda::cublas_version() >= 120800, + "MXFP8 strided batched GEMM requires cuBLAS 12.8+, but run-time cuBLAS version is ", + transformer_engine::cuda::cublas_version()); +#else + NVTE_ERROR( + "MXFP8 strided batched GEMM requires cuBLAS 12.8+, but compile-time cuBLAS version is ", + CUBLAS_VERSION); +#endif // CUBLAS_VERSION >= 120800 + NVTE_CHECK( + inputA->with_gemm_swizzled_scales, + "MXFP8 A scales for strided batched GEMM must be packed by batch and GEMM-swizzled."); + NVTE_CHECK( + inputB->with_gemm_swizzled_scales, + "MXFP8 B scales for strided batched GEMM must be packed by batch and GEMM-swizzled."); } NVTE_CHECK(is_high_precision_dtype(outputD->data.dtype), "Strided batched GEMM currently supports high-precision output only."); @@ -871,11 +1008,45 @@ void cublas_gemm_strided_batched(const Tensor *inputA, const Tensor *inputB, con const cudaDataType_t D_type = get_cuda_dtype(outputD->data.dtype); NVTE_CHECK(C_type == D_type, "Strided batched GEMM currently requires C and D to have the same dtype."); + + const int64_t a_rows = param.transA == CUBLAS_OP_N ? m : k; + const int64_t a_cols = param.transA == CUBLAS_OP_N ? k : m; + const int64_t b_rows = param.transB == CUBLAS_OP_N ? k : n; + const int64_t b_cols = param.transB == CUBLAS_OP_N ? n : k; + CheckStridedBatchBuffer(GetGemmDataBuffer(*inputA, param.A, "A"), a_rows, a_cols, param.lda, + stridea, batch_count, "A data"); + CheckStridedBatchBuffer(GetGemmDataBuffer(*inputB, param.B, "B"), b_rows, b_cols, param.ldb, + strideb, batch_count, "B data"); + CheckStridedBatchBuffer(outputD->data, m, n, ldd, strided, batch_count, "D data"); + CheckStridedBatchBuffer(inputC->data, m, n, ldc, stridec, batch_count, "C data"); + + if (use_mxfp8) { + bool a_scales_are_rowwise = false; + const auto &a_scale_buffer = + GetGemmScaleBuffer(*inputA, param.A_scale_inv, &a_scales_are_rowwise, "A"); + const size_t a_scale_rows = static_cast(a_scales_are_rowwise ? m : k); + const size_t a_scale_features = static_cast(a_scales_are_rowwise ? k : m); + CheckMXFP8BatchScaleBuffer(a_scale_buffer, a_scales_are_rowwise, a_scale_rows, a_scale_features, + static_cast(batch_count), "A"); + + bool b_scales_are_rowwise = false; + const auto &b_scale_buffer = + GetGemmScaleBuffer(*inputB, param.B_scale_inv, &b_scales_are_rowwise, "B"); + const size_t b_scale_rows = static_cast(b_scales_are_rowwise ? n : k); + const size_t b_scale_features = static_cast(b_scales_are_rowwise ? k : n); + CheckMXFP8BatchScaleBuffer(b_scale_buffer, b_scales_are_rowwise, b_scale_rows, b_scale_features, + static_cast(batch_count), "B"); + } + cublasLtHandle_t handle = cublasHandleManager::Instance().GetHandle(); - cublasLtMatmulDesc_t operationDesc = nullptr; - cublasLtMatrixLayout_t Adesc = nullptr, Bdesc = nullptr, Cdesc = nullptr, Ddesc = nullptr; - cublasLtMatmulPreference_t preference = nullptr; + CublasLtDescriptorCleanup descriptors; + auto &operationDesc = descriptors.operation_desc; + auto &Adesc = descriptors.a_desc; + auto &Bdesc = descriptors.b_desc; + auto &Cdesc = descriptors.c_desc; + auto &Ddesc = descriptors.d_desc; + auto &preference = descriptors.preference; int returnedResults = 0; cublasLtMatmulHeuristicResult_t heuristicResult = {}; @@ -924,16 +1095,6 @@ void cublas_gemm_strided_batched(const Tensor *inputA, const Tensor *inputB, con if (use_mxfp8) { #if CUBLAS_VERSION >= 120800 - NVTE_CHECK(transformer_engine::cuda::cublas_version() >= 120800, - "MXFP8 strided batched GEMM requires cuBLAS 12.8+, but run-time cuBLAS version is ", - transformer_engine::cuda::cublas_version()); - NVTE_CHECK( - inputA->with_gemm_swizzled_scales, - "MXFP8 A scales for strided batched GEMM must be packed by batch and GEMM-swizzled."); - NVTE_CHECK( - inputB->with_gemm_swizzled_scales, - "MXFP8 B scales for strided batched GEMM must be packed by batch and GEMM-swizzled."); - 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, @@ -957,10 +1118,6 @@ void cublas_gemm_strided_batched(const Tensor *inputA, const Tensor *inputB, con operationDesc, CUBLASLT_MATMUL_DESC_ALPHA_VECTOR_BATCH_STRIDE, &dummy_a_vec_stride, sizeof(dummy_a_vec_stride))); } -#else - NVTE_ERROR( - "MXFP8 strided batched GEMM requires cuBLAS 12.8+, but compile-time cuBLAS version is ", - CUBLAS_VERSION); #endif // CUBLAS_VERSION >= 120800 } @@ -1011,13 +1168,6 @@ void cublas_gemm_strided_batched(const Tensor *inputA, const Tensor *inputB, con NVTE_CHECK_CUBLAS(cublasLtMatmul(handle, operationDesc, alpha, param.A, Adesc, param.B, Bdesc, beta, C, Cdesc, D, Ddesc, &heuristicResult.algo, aligned_workspace_ptr, workspaceSize, stream)); - - NVTE_CHECK_CUBLAS(cublasLtMatmulPreferenceDestroy(preference)); - NVTE_CHECK_CUBLAS(cublasLtMatrixLayoutDestroy(Ddesc)); - NVTE_CHECK_CUBLAS(cublasLtMatrixLayoutDestroy(Cdesc)); - NVTE_CHECK_CUBLAS(cublasLtMatrixLayoutDestroy(Bdesc)); - NVTE_CHECK_CUBLAS(cublasLtMatrixLayoutDestroy(Adesc)); - NVTE_CHECK_CUBLAS(cublasLtMatmulDescDestroy(operationDesc)); } } // namespace transformer_engine diff --git a/transformer_engine/common/include/transformer_engine/swizzle.h b/transformer_engine/common/include/transformer_engine/swizzle.h index 396093b543..716174e491 100644 --- a/transformer_engine/common/include/transformer_engine/swizzle.h +++ b/transformer_engine/common/include/transformer_engine/swizzle.h @@ -30,6 +30,41 @@ extern "C" { */ void nvte_swizzle_scaling_factors(const NVTETensor input, NVTETensor output, cudaStream_t stream); +/*! \brief Pack and swizzle MXFP8 scaling factors for strided batched GEMM. + * + * \param[in] input Input MXFP8 tensor with compact row-wise scale_inv. + * \param[in,out] output Output MXFP8 tensor with packed, GEMM-swizzled row-wise scale_inv. + * It must be marked as having GEMM-swizzled scales. + * \param[in] batch_dim GEMM batch dimension. Only the first or penultimate dimension is + * supported. + * \param[in] batch_in_features + * Whether compact scales were produced from a 2D view with batches + * concatenated in the feature dimension. This is supported only for + * a penultimate batch dimension. + * \param[in] stream CUDA stream used for the operation. + * + * The quantized data buffer is not modified. Output scales are laid out as + * [batch0_swizzled_scales][batch1_swizzled_scales]... . + */ +void nvte_pack_mxfp8_scales_for_batched_gemm(const NVTETensor input, NVTETensor output, + int64_t batch_dim, int batch_in_features, + cudaStream_t stream); + +/*! \brief Pack and swizzle column-wise MXFP8 scaling factors for strided batched GEMM. + * + * \param[in] input Input MXFP8 tensor with compact column-wise scale_inv. + * \param[in,out] output Output MXFP8 tensor with packed, GEMM-swizzled column-wise + * scale_inv. It must be marked as having GEMM-swizzled scales. + * \param[in] batch_dim GEMM batch dimension. Only the first or penultimate dimension is + * supported. + * \param[in] stream CUDA stream used for the operation. + * + * The quantized data buffer is not modified. Output scales are laid out as + * [batch0_swizzled_scales][batch1_swizzled_scales]... . + */ +void nvte_pack_mxfp8_columnwise_scales_for_batched_gemm(const NVTETensor input, NVTETensor output, + int64_t batch_dim, cudaStream_t stream); + /*! \brief Swizzling scaling factors into the required interleaved layout for GEMM * * \param[in] inputs Input tensors with non-swizzled scale_inv. diff --git a/transformer_engine/common/swizzle/swizzle.cu b/transformer_engine/common/swizzle/swizzle.cu index 38b526360e..f2d2cf8691 100644 --- a/transformer_engine/common/swizzle/swizzle.cu +++ b/transformer_engine/common/swizzle/swizzle.cu @@ -24,6 +24,62 @@ namespace { constexpr int MXFP8_BLOCK_SIZE = 32; constexpr int NVFP4_BLOCK_SIZE = 16; +__device__ __forceinline__ size_t mxfp8_gemm_swizzled_scale_idx(size_t row, size_t col, + size_t num_tiles_x) { + constexpr size_t tile_dim_x = 4; + constexpr size_t tile_dim_y = 128; + constexpr size_t tile_size = tile_dim_x * tile_dim_y; + const size_t tile_x = col / tile_dim_x; + const size_t tile_y = row / tile_dim_y; + const size_t in_tile_x = col % tile_dim_x; + const size_t in_tile_y = row % tile_dim_y; + size_t idx = (tile_y * num_tiles_x + tile_x) * tile_size; + idx += (in_tile_y % 32) * 16 + (in_tile_y / 32) * 4 + in_tile_x; + return idx; +} + +__global__ void pack_mxfp8_rowwise_scales_for_batched_gemm_kernel( + const uint8_t* input, uint8_t* output, int64_t rows_per_batch, int64_t batch_count, + int64_t scale_cols, int64_t input_row_stride, int64_t output_batch_stride, int64_t num_tiles_x, + bool batch_major, bool batch_in_features) { + const int64_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + const int64_t numel = batch_count * rows_per_batch * scale_cols; + if (idx >= numel) return; + + const int64_t col = idx % scale_cols; + const int64_t row_and_batch = idx / scale_cols; + const int64_t row = row_and_batch % rows_per_batch; + const int64_t batch = row_and_batch / rows_per_batch; + const int64_t input_row = + batch_in_features ? row + : (batch_major ? batch * rows_per_batch + row : row * batch_count + batch); + const int64_t input_col = batch_in_features ? batch * scale_cols + col : col; + const size_t output_idx = + static_cast(batch) * output_batch_stride + + mxfp8_gemm_swizzled_scale_idx(row, col, static_cast(num_tiles_x)); + output[output_idx] = input[input_row * input_row_stride + input_col]; +} + +__global__ void pack_mxfp8_columnwise_scales_for_batched_gemm_kernel( + const uint8_t* input, uint8_t* output, int64_t scale_rows_per_batch, int64_t batch_count, + int64_t features, int64_t input_row_stride, int64_t output_batch_stride, int64_t num_tiles_x, + bool batch_major) { + const int64_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + const int64_t numel = batch_count * scale_rows_per_batch * features; + if (idx >= numel) return; + + const int64_t feature = idx % features; + const int64_t scale_row_and_batch = idx / features; + const int64_t scale_row = scale_row_and_batch % scale_rows_per_batch; + const int64_t batch = scale_row_and_batch / scale_rows_per_batch; + const int64_t input_row = batch_major ? batch * scale_rows_per_batch + scale_row : scale_row; + const int64_t input_col = batch_major ? feature : batch * features + feature; + const size_t output_idx = + static_cast(batch) * output_batch_stride + + mxfp8_gemm_swizzled_scale_idx(feature, scale_row, static_cast(num_tiles_x)); + output[output_idx] = input[input_row * input_row_stride + input_col]; +} + int get_max_dynamic_smem(int device_id = -1) { static std::vector cache(cuda::num_devices(), -1); static std::vector flags(cuda::num_devices()); @@ -1947,6 +2003,153 @@ void multi_tensor_unswizzle_scaling_factors(const std::vector& input, kernel_args, vec_load_size, false, stream); } } + +void pack_mxfp8_scales_for_batched_gemm(const Tensor* input, Tensor* output, int64_t batch_dim, + bool batch_in_features, cudaStream_t stream) { + NVTE_CHECK( + input->scaling_mode == NVTE_MXFP8_1D_SCALING && output->scaling_mode == NVTE_MXFP8_1D_SCALING, + "Batched scale packing supports only MXFP8 tensors."); + NVTE_CHECK(input->has_data(), "Batched scale packing requires row-wise input data."); + NVTE_CHECK(input->scale_inv.has_data(), + "Batched scale packing requires row-wise input scaling factors."); + NVTE_CHECK(output->scale_inv.has_data(), + "Batched scale packing requires row-wise output scaling factors."); + NVTE_CHECK(!input->with_gemm_swizzled_scales, + "Batched scale packing expects compact input scaling factors."); + NVTE_CHECK(output->with_gemm_swizzled_scales, + "Batched scale packing expects GEMM-swizzled output scaling factors."); + NVTE_CHECK(is_fp8_dtype(input->data.dtype), "Batched scale packing expects FP8 input data."); + NVTE_CHECK( + input->scale_inv.dtype == DType::kFloat8E8M0 && output->scale_inv.dtype == DType::kFloat8E8M0, + "Batched scale packing expects E8M0 scaling factors."); + + const auto& data_shape = input->data.shape; + const int64_t ndim = static_cast(data_shape.size()); + NVTE_CHECK(ndim >= 2, "Batched scale packing expects input with at least two dimensions."); + if (batch_dim < 0) batch_dim += ndim; + NVTE_CHECK(batch_dim == 0 || batch_dim == ndim - 2, + "Batched scale packing supports only batch_dim=0 or batch_dim=-2."); + NVTE_CHECK(!batch_in_features || batch_dim == ndim - 2, + "Feature-batched compact scales require batch_dim=-2."); + + const int64_t batch_count = static_cast(data_shape[batch_dim]); + const int64_t features = static_cast(data_shape.back()); + NVTE_CHECK(batch_count > 0 && features > 0, + "Batched scale packing requires non-empty batch and feature dimensions."); + NVTE_CHECK(features % MXFP8_BLOCK_SIZE == 0, "MXFP8 feature dimension must be divisible by ", + MXFP8_BLOCK_SIZE, "."); + const int64_t total_rows = static_cast(input->data.numel()) / features; + NVTE_CHECK(total_rows > 0, "Batched scale packing does not support empty matrices."); + NVTE_CHECK(total_rows % batch_count == 0, + "Input rows must be divisible by the GEMM batch count."); + const int64_t rows_per_batch = total_rows / batch_count; + const int64_t scale_cols = features / MXFP8_BLOCK_SIZE; + const int64_t padded_rows = static_cast(round_up_to_multiple(rows_per_batch, 128)); + const int64_t padded_scale_cols = static_cast(round_up_to_multiple(scale_cols, 4)); + const int64_t output_batch_stride = padded_rows * padded_scale_cols; + const int64_t output_numel = batch_count * output_batch_stride; + + NVTE_CHECK(input->scale_inv.shape.size() == 2, + "Batched scale packing expects 2D compact scaling factors."); + const int64_t input_scale_rows = static_cast(input->scale_inv.shape[0]); + const int64_t input_row_stride = static_cast(input->scale_inv.shape[1]); + if (batch_in_features) { + NVTE_CHECK(input_scale_rows >= rows_per_batch && input_row_stride >= batch_count * scale_cols, + "Compact scaling factor buffer is too small for feature-batched scales."); + } else { + NVTE_CHECK(input_scale_rows >= total_rows && input_row_stride >= scale_cols, + "Compact scaling factor buffer is too small for the input tensor."); + } + NVTE_CHECK(static_cast(output->scale_inv.numel()) == output_numel, + "Packed scaling factor buffer has invalid size (expected ", output_numel, ", got ", + output->scale_inv.numel(), ")."); + + auto* output_ptr = static_cast(output->scale_inv.dptr); + NVTE_CHECK_CUDA(cudaMemsetAsync(output_ptr, 0, output_numel, stream)); + const int64_t numel = batch_count * rows_per_batch * scale_cols; + constexpr int threads = 256; + const int blocks = static_cast((numel + threads - 1) / threads); + pack_mxfp8_rowwise_scales_for_batched_gemm_kernel<<>>( + static_cast(input->scale_inv.dptr), output_ptr, rows_per_batch, batch_count, + scale_cols, input_row_stride, output_batch_stride, padded_scale_cols / 4, batch_dim == 0, + batch_in_features); + NVTE_CHECK_CUDA(cudaGetLastError()); +} + +void pack_mxfp8_columnwise_scales_for_batched_gemm(const Tensor* input, Tensor* output, + int64_t batch_dim, cudaStream_t stream) { + NVTE_CHECK( + input->scaling_mode == NVTE_MXFP8_1D_SCALING && output->scaling_mode == NVTE_MXFP8_1D_SCALING, + "Batched scale packing supports only MXFP8 tensors."); + NVTE_CHECK(input->has_columnwise_data(), + "Batched scale packing requires column-wise input data."); + NVTE_CHECK(input->columnwise_scale_inv.has_data(), + "Batched scale packing requires column-wise input scaling factors."); + NVTE_CHECK(output->columnwise_scale_inv.has_data(), + "Batched scale packing requires column-wise output scaling factors."); + NVTE_CHECK(!input->with_gemm_swizzled_scales, + "Batched scale packing expects compact input scaling factors."); + NVTE_CHECK(output->with_gemm_swizzled_scales, + "Batched scale packing expects GEMM-swizzled output scaling factors."); + NVTE_CHECK(is_fp8_dtype(input->columnwise_data.dtype), + "Batched scale packing expects FP8 input data."); + NVTE_CHECK(input->columnwise_scale_inv.dtype == DType::kFloat8E8M0 && + output->columnwise_scale_inv.dtype == DType::kFloat8E8M0, + "Batched scale packing expects E8M0 scaling factors."); + + const auto& data_shape = input->columnwise_data.shape; + const int64_t ndim = static_cast(data_shape.size()); + NVTE_CHECK(ndim >= 2, "Batched scale packing expects input with at least two dimensions."); + if (batch_dim < 0) batch_dim += ndim; + NVTE_CHECK(batch_dim == 0 || batch_dim == ndim - 2, + "Batched scale packing supports only batch_dim=0 or batch_dim=-2."); + + const int64_t batch_count = static_cast(data_shape[batch_dim]); + const int64_t features = static_cast(data_shape.back()); + NVTE_CHECK(batch_count > 0 && features > 0, + "Batched scale packing requires non-empty batch and feature dimensions."); + const int64_t total_rows = static_cast(input->columnwise_data.numel()) / features; + NVTE_CHECK(total_rows > 0, "Batched scale packing does not support empty matrices."); + NVTE_CHECK(total_rows % batch_count == 0, + "Input rows must be divisible by the GEMM batch count."); + const int64_t rows_per_batch = total_rows / batch_count; + NVTE_CHECK(rows_per_batch % MXFP8_BLOCK_SIZE == 0, "MXFP8 rows per GEMM must be divisible by ", + MXFP8_BLOCK_SIZE, "."); + const int64_t scale_rows_per_batch = rows_per_batch / MXFP8_BLOCK_SIZE; + const int64_t padded_scale_rows = + static_cast(round_up_to_multiple(scale_rows_per_batch, 4)); + const int64_t padded_features = static_cast(round_up_to_multiple(features, 128)); + const int64_t output_batch_stride = padded_scale_rows * padded_features; + const int64_t output_numel = batch_count * output_batch_stride; + + NVTE_CHECK(input->columnwise_scale_inv.shape.size() == 2, + "Batched scale packing expects 2D compact scaling factors."); + const int64_t input_scale_rows = static_cast(input->columnwise_scale_inv.shape[0]); + const int64_t input_row_stride = static_cast(input->columnwise_scale_inv.shape[1]); + if (batch_dim == 0) { + NVTE_CHECK( + input_scale_rows >= batch_count * scale_rows_per_batch && input_row_stride >= features, + "Compact scaling factor buffer is too small for the batch-major input tensor."); + } else { + NVTE_CHECK( + input_scale_rows >= scale_rows_per_batch && input_row_stride >= batch_count * features, + "Compact scaling factor buffer is too small for the interleaved input tensor."); + } + NVTE_CHECK(static_cast(output->columnwise_scale_inv.numel()) == output_numel, + "Packed scaling factor buffer has invalid size (expected ", output_numel, ", got ", + output->columnwise_scale_inv.numel(), ")."); + + auto* output_ptr = static_cast(output->columnwise_scale_inv.dptr); + NVTE_CHECK_CUDA(cudaMemsetAsync(output_ptr, 0, output_numel, stream)); + const int64_t numel = batch_count * scale_rows_per_batch * features; + constexpr int threads = 256; + const int blocks = static_cast((numel + threads - 1) / threads); + pack_mxfp8_columnwise_scales_for_batched_gemm_kernel<<>>( + static_cast(input->columnwise_scale_inv.dptr), output_ptr, + scale_rows_per_batch, batch_count, features, input_row_stride, output_batch_stride, + padded_scale_rows / 4, batch_dim == 0); + NVTE_CHECK_CUDA(cudaGetLastError()); +} } // namespace transformer_engine /* @@ -1960,6 +2163,23 @@ void nvte_swizzle_scaling_factors(const NVTETensor input, NVTETensor output, cud swizzle_scaling_factors(convertNVTETensorCheck(input), convertNVTETensorCheck(output), stream); } +void nvte_pack_mxfp8_scales_for_batched_gemm(const NVTETensor input, NVTETensor output, + int64_t batch_dim, int batch_in_features, + cudaStream_t stream) { + NVTE_API_CALL(nvte_pack_mxfp8_scales_for_batched_gemm); + using namespace transformer_engine; + pack_mxfp8_scales_for_batched_gemm(convertNVTETensorCheck(input), convertNVTETensorCheck(output), + batch_dim, static_cast(batch_in_features), stream); +} + +void nvte_pack_mxfp8_columnwise_scales_for_batched_gemm(const NVTETensor input, NVTETensor output, + int64_t batch_dim, cudaStream_t stream) { + NVTE_API_CALL(nvte_pack_mxfp8_columnwise_scales_for_batched_gemm); + using namespace transformer_engine; + pack_mxfp8_columnwise_scales_for_batched_gemm(convertNVTETensorCheck(input), + convertNVTETensorCheck(output), batch_dim, stream); +} + void nvte_multi_tensor_swizzle_scaling_factors(const NVTETensor* inputs, NVTETensor* outputs, const size_t num_tensors, cudaStream_t stream) { NVTE_API_CALL(nvte_multi_tensor_swizzle_scaling_factors); diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index eb4b9ec684..a0be39de03 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -404,10 +404,14 @@ def strided_batched_gemm( cuBLAS/TE convention as the C++ GEMM backend. For the common row-major ``layout="TN"`` case, this computes per batch: ``out[n, m] = B[n, k] @ A[m, k].T``. + ``out`` serves as both the C and D matrices, so accumulation applies + ``beta`` to the existing contents of ``out`` in place. Inputs must both be high-precision tensors or both be MXFP8 tensors. MXFP8 - scale tensors are packed consecutively by batch and must already use the - GEMM-swizzled layout. + inputs must use compact scales; the extension packs and GEMM-swizzles the + required scaling direction for each operand before launching cuBLASLt. + MXFP8 operand strides must describe a contiguous ``[G, ..., D]`` or + ``[..., G, D]`` layout. """ assert layout in ("TN", "NN", "NT"), f"GEMM layout {layout} not supported." high_precision_dtypes = (torch.float32, torch.float16, torch.bfloat16) @@ -423,8 +427,8 @@ def strided_batched_gemm( ), "strided_batched_gemm supports only high-precision or MXFP8 A and B tensor pairs." if mxfp8_inputs: assert ( - A._with_gemm_swizzled_scales and B._with_gemm_swizzled_scales - ), "strided_batched_gemm expects packed, GEMM-swizzled MXFP8 scales." + not A._with_gemm_swizzled_scales and not B._with_gemm_swizzled_scales + ), "strided_batched_gemm expects compact MXFP8 scales." transa = layout[0] == "T" transb = layout[1] == "T" beta = validate_gemm_scale(beta, accumulate) diff --git a/transformer_engine/pytorch/csrc/extensions/gemm.cpp b/transformer_engine/pytorch/csrc/extensions/gemm.cpp index 590600858d..fa2b6b91d1 100644 --- a/transformer_engine/pytorch/csrc/extensions/gemm.cpp +++ b/transformer_engine/pytorch/csrc/extensions/gemm.cpp @@ -6,6 +6,7 @@ #include +#include #include #include @@ -13,6 +14,7 @@ #include "common/util/cuda_runtime.h" #include "common/util/system.h" #include "pybind.h" +#include "transformer_engine/swizzle.h" #include "transformer_engine/transformer_engine.h" #include "util.h" @@ -413,6 +415,141 @@ std::vector gemm(py::handle A, bool transa, py::handle B, bool trans return out; } +namespace { + +struct PackedMXFP8Operand { + at::Tensor scale_inv; + TensorWrapper tensor{NVTE_MXFP8_1D_SCALING}; +}; + +int64_t checked_mul(int64_t lhs, int64_t rhs, const char* description) { + NVTE_CHECK(lhs >= 0 && rhs >= 0, "Negative value while calculating ", description, "."); + NVTE_CHECK(lhs == 0 || rhs <= std::numeric_limits::max() / lhs, + "Integer overflow while calculating ", description, "."); + return lhs * rhs; +} + +int64_t round_up(int64_t value, int64_t multiple, const char* description) { + NVTE_CHECK(value >= 0, "Negative value while calculating ", description, "."); + const int64_t quotient = value / multiple + static_cast(value % multiple != 0); + return checked_mul(quotient, multiple, description); +} + +PackedMXFP8Operand pack_mxfp8_operand_scales(py::handle operand, + const TensorWrapper& operand_tensor, bool rowwise, + int64_t rows_per_batch, int64_t features, + int64_t batch_count, int64_t leading_dim, + int64_t batch_stride, const char* name) { + NVTE_CHECK(rows_per_batch > 0 && features > 0 && batch_count > 0, + "MXFP8 strided batched GEMM requires positive dimensions for ", name, "."); + NVTE_CHECK(!operand_tensor.get_with_gemm_swizzled_scales(), + "strided_batched_gemm expects compact MXFP8 scales for ", name, "."); + + const char* data_attr = rowwise ? "_rowwise_data" : "_columnwise_data"; + const char* scale_attr = rowwise ? "_rowwise_scale_inv" : "_columnwise_scale_inv"; + const py::object data_py = operand.attr(data_attr); + const py::object scale_py = operand.attr(scale_attr); + NVTE_CHECK(!data_py.is_none() && !scale_py.is_none(), "MXFP8 strided_batched_gemm requires ", + rowwise ? "row-wise" : "column-wise", " data and scales for ", name, "."); + const at::Tensor data = data_py.cast(); + const at::Tensor scale_inv = scale_py.cast(); + NVTE_CHECK(data.is_cuda() && scale_inv.is_cuda(), "MXFP8 strided_batched_gemm requires CUDA ", + name, " data and scales."); + NVTE_CHECK(data.device() == scale_inv.device(), "MXFP8 strided_batched_gemm requires ", name, + " data and scales on the same device."); + NVTE_CHECK(data.scalar_type() == torch::kUInt8 && scale_inv.scalar_type() == torch::kUInt8, + "MXFP8 strided_batched_gemm expects byte data and E8M0 scales for ", name, "."); + NVTE_CHECK(data.is_contiguous() && scale_inv.is_contiguous(), + "MXFP8 strided_batched_gemm requires contiguous ", name, " data and scales."); + NVTE_CHECK(data.dim() >= 2, "MXFP8 strided_batched_gemm requires at least 2D data for ", name, + "."); + + const int64_t matrix_numel = checked_mul(rows_per_batch, features, name); + const int64_t expected_data_numel = checked_mul(batch_count, matrix_numel, name); + NVTE_CHECK(data.numel() == expected_data_numel, "MXFP8 strided_batched_gemm ", name, + " data buffer has invalid size (expected ", expected_data_numel, ", got ", + data.numel(), ")."); + const int64_t interleaved_leading_dim = checked_mul(batch_count, features, name); + const bool batch_major = leading_dim == features && batch_stride == matrix_numel; + const bool interleaved = leading_dim == interleaved_leading_dim && batch_stride == features; + NVTE_CHECK(batch_major || interleaved, "MXFP8 strided_batched_gemm supports only contiguous ", + "[G, ..., D] or [..., G, D] layouts for ", name, " (leading_dim=", leading_dim, + ", batch_stride=", batch_stride, ")."); + + const int64_t compact_features = data.size(-1); + const bool batch_in_features = interleaved && compact_features == interleaved_leading_dim; + if (batch_major) { + NVTE_CHECK(compact_features == features, "Batch-major MXFP8 ", name, + " must be quantized with the per-GEMM feature dimension last."); + } else { + NVTE_CHECK(compact_features == features || batch_in_features, "Interleaved MXFP8 ", name, + " has incompatible compact scale layout."); + NVTE_CHECK(rowwise || batch_in_features, "Interleaved column-wise MXFP8 ", name, + " must be quantized with batches concatenated in the feature dimension."); + } + + const std::vector logical_shape = + batch_major + ? std::vector{static_cast(batch_count), + static_cast(rows_per_batch), static_cast(features)} + : std::vector{static_cast(rows_per_batch), + static_cast(batch_count), static_cast(features)}; + const int64_t batch_dim = batch_major ? 0 : 1; + const auto input_data = + rowwise ? operand_tensor.get_rowwise_data() : operand_tensor.get_columnwise_data(); + + TensorWrapper compact_tensor(NVTE_MXFP8_1D_SCALING); + PackedMXFP8Operand packed; + if (rowwise) { + NVTE_CHECK(features % 32 == 0, "Row-wise MXFP8 ", name, + " feature dimension must be divisible by 32."); + const int64_t scale_rows = round_up(rows_per_batch, 128, name); + const int64_t scale_cols = round_up(features / 32, 4, name); + const int64_t output_numel = + checked_mul(batch_count, checked_mul(scale_rows, scale_cols, name), name); + packed.scale_inv = at::empty({output_numel}, scale_inv.options()); + compact_tensor.set_rowwise_data(input_data.data_ptr, static_cast(input_data.dtype), + logical_shape); + compact_tensor.set_rowwise_scale_inv(scale_inv.data_ptr(), DType::kFloat8E8M0, + getTensorShape(scale_inv)); + packed.tensor.set_rowwise_data(input_data.data_ptr, static_cast(input_data.dtype), + logical_shape); + packed.tensor.set_rowwise_scale_inv(packed.scale_inv.data_ptr(), DType::kFloat8E8M0, + getTensorShape(packed.scale_inv)); + } else { + NVTE_CHECK(rows_per_batch % 32 == 0, "Column-wise MXFP8 ", name, + " row count must be divisible by 32."); + const int64_t scale_rows = round_up(rows_per_batch / 32, 4, name); + const int64_t scale_cols = round_up(features, 128, name); + const int64_t output_numel = + checked_mul(batch_count, checked_mul(scale_rows, scale_cols, name), name); + packed.scale_inv = at::empty({output_numel}, scale_inv.options()); + compact_tensor.set_columnwise_data(input_data.data_ptr, static_cast(input_data.dtype), + logical_shape); + compact_tensor.set_columnwise_scale_inv(scale_inv.data_ptr(), DType::kFloat8E8M0, + getTensorShape(scale_inv)); + packed.tensor.set_columnwise_data(input_data.data_ptr, static_cast(input_data.dtype), + logical_shape); + packed.tensor.set_columnwise_scale_inv(packed.scale_inv.data_ptr(), DType::kFloat8E8M0, + getTensorShape(packed.scale_inv)); + } + packed.tensor.set_with_gemm_swizzled_scales(true); + + NVTE_SCOPED_GIL_RELEASE({ + if (rowwise) { + nvte_pack_mxfp8_scales_for_batched_gemm(compact_tensor.data(), packed.tensor.data(), + batch_dim, static_cast(batch_in_features), + at::cuda::getCurrentCUDAStream()); + } else { + nvte_pack_mxfp8_columnwise_scales_for_batched_gemm( + compact_tensor.data(), packed.tensor.data(), batch_dim, at::cuda::getCurrentCUDAStream()); + } + }); + return packed; +} + +} // namespace + at::Tensor strided_batched_gemm(py::handle A, bool transa, py::handle B, bool transb, at::Tensor D, at::Tensor workspace, size_t workspaceSize, int64_t m, int64_t n, int64_t k, int64_t batch_count, int64_t lda, int64_t stridea, @@ -447,9 +584,17 @@ at::Tensor strided_batched_gemm(py::handle A, bool transa, py::handle B, bool tr is_fp8_dtype(A_tensor.dtype()) && is_fp8_dtype(B_tensor.dtype()); NVTE_CHECK(high_precision_inputs || mxfp8_inputs, "strided_batched_gemm supports only high-precision or MXFP8 A and B tensor pairs."); + const TensorWrapper* gemm_A = &A_tensor; + const TensorWrapper* gemm_B = &B_tensor; + std::optional packed_A; + std::optional packed_B; if (mxfp8_inputs) { - NVTE_CHECK(A_tensor.get_with_gemm_swizzled_scales() && B_tensor.get_with_gemm_swizzled_scales(), - "strided_batched_gemm expects packed, GEMM-swizzled MXFP8 scales."); + packed_A.emplace(pack_mxfp8_operand_scales(A, A_tensor, transa, transa ? m : k, transa ? k : m, + batch_count, lda, stridea, "A")); + packed_B.emplace(pack_mxfp8_operand_scales(B, B_tensor, !transb, transb ? k : n, transb ? n : k, + batch_count, ldb, strideb, "B")); + gemm_A = &packed_A->tensor; + gemm_B = &packed_B->tensor; } NVTE_CHECK(is_high_precision_dtype(D_tensor.dtype()), "strided_batched_gemm currently expects a high-precision output tensor."); @@ -460,7 +605,7 @@ at::Tensor strided_batched_gemm(py::handle A, bool transa, py::handle B, bool tr NVTE_SCOPED_GIL_RELEASE({ nvte_cublas_gemm_strided_batched( - transa, transb, &alpha, A_tensor.data(), lda, stridea, B_tensor.data(), ldb, strideb, + transa, transb, &alpha, gemm_A->data(), lda, stridea, gemm_B->data(), ldb, strideb, &beta.value(), D_tensor.data(), ldd, strided, D_tensor.data(), ldd, strided, batch_count, m, n, k, te_workspace.data(), config, at::cuda::getCurrentCUDAStream()); }); diff --git a/transformer_engine/pytorch/ops/basic/__init__.py b/transformer_engine/pytorch/ops/basic/__init__.py index 4eb2796b2c..f52374d977 100644 --- a/transformer_engine/pytorch/ops/basic/__init__.py +++ b/transformer_engine/pytorch/ops/basic/__init__.py @@ -20,6 +20,7 @@ from .add_extra_input import AddExtraInput from .all_gather import AllGather from .all_reduce import AllReduce +from .batched_linear import BatchedLinear from .basic_linear import BasicLinear from .bias import Bias from .constant_scale import ConstantScale diff --git a/transformer_engine/pytorch/ops/basic/batched_linear.py b/transformer_engine/pytorch/ops/basic/batched_linear.py new file mode 100644 index 0000000000..d40291311b --- /dev/null +++ b/transformer_engine/pytorch/ops/basic/batched_linear.py @@ -0,0 +1,758 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Fusible operation for strided batched linear transformations.""" + +from __future__ import annotations + +from collections.abc import Callable, Iterable, Sequence +import contextlib +from typing import Any, Optional +import warnings + +import torch + +from transformer_engine.common.recipe import Recipe + +from ...cpp_extensions import strided_batched_gemm +from ...distributed import CudaRNGStatesTracker +from ...module.base import _2X_ACC_DGRAD, _2X_ACC_FPROP, _2X_ACC_WGRAD +from ...quantization import FP8GlobalStateManager, QuantizerRole +from ...tensor import MXFP8Quantizer, Quantizer +from ...tensor.storage.mxfp8_tensor_storage import MXFP8TensorStorage +from ...utils import ( + canonicalize_device, + canonicalize_dtype, + devices_match, + get_default_init_method, +) +from .._common import ( + get_accumulate_flag_in_param, + get_dummy_wgrads_for_params, + get_main_grad_from_param, + is_quantized_tensor, + maybe_autocast_dtype, + maybe_dequantize, + view_main_grad_as_grouped_buffer, +) +from ..op import BasicOperation, OperationContext + + +_HIGH_PRECISION_DTYPES = (torch.float32, torch.float16, torch.bfloat16) + + +class BatchedLinear(BasicOperation): + """Apply one linear transformation per batch entry. + + The weight has shape ``[G, R, D]``. Inputs may use either contiguous + ``[..., G, D]`` storage with ``batch_dim=-2`` or contiguous + ``[G, ..., D]`` storage with ``batch_dim=0``. The output replaces ``D`` + with ``R``. High-precision and MXFP8 forward and backward computation are + supported. Tensor parallelism is not supported. + + Parameters + ---------- + num_gemms : int + Number of independent linear transformations (``G``). + in_features : int + Input feature dimension (``D``). + out_features : int + Output feature dimension (``R``). + batch_dim : {0, -2}, default = -2 + Position of the GEMM batch dimension in the input. + bias : bool, default = True + Add one learned bias of shape ``[G, R]``. + return_bias : bool, default = False + Return the bias separately instead of applying it. + device : torch.device, default = default CUDA device + Parameter device. + dtype : torch.dtype, default = default dtype + Parameter datatype. + rng_state_tracker_function : callable, optional + Function returning a ``CudaRNGStatesTracker`` used during parameter + initialization. + accumulate_into_main_grad : bool, default = False + Write weight gradients directly into the externally allocated + ``weight.main_grad`` buffer. Setting ``weight.overwrite_main_grad`` to + ``True`` overwrites that buffer instead of accumulating into it. + init_method : callable, optional + Weight initialization method. The default is TE's normal initializer. + name : str, optional + Name used by quantizer-role dispatch and debugging. + + Notes + ----- + Constructing this operation under ``quantized_model_init`` supports only + an MXFP8 recipe. Meta-device parameter materialization relies on the + operation fuser's deferred-initialization support. + """ + + num_extra_outputs: int = 0 + + def __init__( + self, + num_gemms: int, + in_features: int, + out_features: int, + *, + batch_dim: int = -2, + bias: bool = True, + return_bias: bool = False, + 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, + init_method: Optional[Callable[[torch.Tensor], torch.Tensor]] = None, + name: Optional[str] = None, + ) -> None: + # The fuser allocates extra-output routing state in BasicOperation.__init__. + self.num_extra_outputs = int(bias and return_bias) + super().__init__() + + for arg_name, value in ( + ("num_gemms", num_gemms), + ("in_features", in_features), + ("out_features", out_features), + ): + if not isinstance(value, int) or isinstance(value, bool) or value <= 0: + raise ValueError(f"{arg_name} must be a positive integer (got {value!r})") + if batch_dim not in (0, -2): + raise ValueError( + f"BatchedLinear supports only batch_dim=0 or batch_dim=-2 (got {batch_dim})" + ) + + device = canonicalize_device(device) + dtype = canonicalize_dtype(dtype) + if dtype not in _HIGH_PRECISION_DTYPES: + raise ValueError( + f"BatchedLinear parameters must use float32, float16, or bfloat16 (got {dtype})" + ) + + self.num_gemms = num_gemms + self.in_features = in_features + self.out_features = out_features + self.batch_dim = batch_dim + self.use_bias = bias + self.return_bias = return_bias + self.apply_bias = bias and not return_bias + self.name = name + self._rng_state_tracker_function = rng_state_tracker_function + self._accumulate_into_main_grad = accumulate_into_main_grad + self._init_method = get_default_init_method() if init_method is None else init_method + + # Initialize recipe state if the weight itself is stored in MXFP8. + self._with_quantized_weight = FP8GlobalStateManager.with_fp8_parameters() + if self._with_quantized_weight: + self.reset_recipe_state(recipe=FP8GlobalStateManager.get_fp8_recipe()) + + weight = torch.empty( + num_gemms, + out_features, + in_features, + dtype=dtype, + device=device, + ) + self.weight: torch.nn.Parameter + self.register_parameter("weight", torch.nn.Parameter(weight)) + + bias_tensor = None + if bias: + bias_tensor = torch.empty( + num_gemms, + out_features, + dtype=dtype, + device=device, + ) + bias_tensor = torch.nn.Parameter(bias_tensor) + self.bias: Optional[torch.nn.Parameter] + self.register_parameter("bias", bias_tensor) + + if device.type != "meta": + self.reset_parameters() + + def reset_parameters(self) -> None: + """Allocate and initialize parameter values.""" + old_weight = self.weight + device = old_weight.device + if device.type == "meta": + device = canonicalize_device(None) + + if is_quantized_tensor(old_weight): + weight = torch.empty(old_weight.size(), dtype=old_weight.dtype, device=device) + elif not devices_match(old_weight.device, device): + weight = torch.empty_like(old_weight, device=device) + else: + weight = old_weight + + init_context = contextlib.nullcontext() + if self._rng_state_tracker_function is not None: + init_context = self._rng_state_tracker_function().fork() + with torch.no_grad(), init_context: + self._init_method(weight) + + if self._with_quantized_weight: + quantizer = self.get_quantizer("forward", 1) + if quantizer is None: + raise RuntimeError( + "Tried to quantize BatchedLinear weight after deferred initialization, " + "but no quantizer was available. The forward pass must run under " + "MXFP8 autocast." + ) + self._configure_quantizer(quantizer, internal=False) + quantizer.set_usage(rowwise=True, columnwise=torch.is_grad_enabled()) + with torch.no_grad(): + weight = quantizer(weight) + + if not isinstance(weight, torch.nn.Parameter): + weight = torch.nn.Parameter(weight, requires_grad=old_weight.requires_grad) + self.weight = weight + + if self.use_bias: + old_bias = self.bias + if old_bias is None: + raise RuntimeError("BatchedLinear bias parameter is missing") + if devices_match(old_bias.device, device): + bias = old_bias + else: + bias = torch.empty_like(old_bias, device=device) + with torch.no_grad(): + bias.zero_() + if not isinstance(bias, torch.nn.Parameter): + bias = torch.nn.Parameter(bias, requires_grad=old_bias.requires_grad) + self.bias = bias + + def pre_first_fuser_forward(self) -> None: + super().pre_first_fuser_forward() + if self.weight.device.type == "meta": + self.reset_parameters() + + def num_quantizers(self, mode: str) -> int: + if mode == "forward": + return 2 + if mode == "backward": + return 1 + return 0 + + def get_quantizer_roles(self, mode: str) -> Optional[list[QuantizerRole]]: + name = self.name or "" + if mode == "forward": + return [ + QuantizerRole(module_type="batched_linear", tensor_type="input", name=name), + QuantizerRole(module_type="batched_linear", tensor_type="weight", name=name), + ] + if mode == "backward": + return [ + QuantizerRole( + module_type="batched_linear", + tensor_type="grad_output", + name=name, + ) + ] + return None + + def get_input_quantizer(self) -> None: + # Input scales must be packed with knowledge of this op's batch dimension. + return None + + def get_grad_output_quantizer(self) -> None: + # Grad-output scales also require batch-aware packing. + return None + + @staticmethod + def _configure_quantizer(quantizer: Quantizer, *, internal: bool = True) -> None: + if not isinstance(quantizer, MXFP8Quantizer): + raise RuntimeError("BatchedLinear expected an MXFP8 quantizer") + quantizer.set_usage(rowwise=True, columnwise=False) + quantizer.internal = internal + quantizer.optimize_for_gemm = False + + @staticmethod + def _validate_recipe(recipe: Recipe) -> None: + if not recipe.mxfp8(): + raise ValueError( + "BatchedLinear supports only high-precision compute or the MXFP8 recipe " + f"(got {recipe.__class__.__name__})" + ) + if recipe.backward_override is not None: + raise ValueError( + "BatchedLinear does not support MXFP8 backward_override " + f"(got {recipe.backward_override!r})" + ) + + def reset_recipe_state(self, *, recipe: Optional[Recipe]) -> None: + if recipe is not None: + self._validate_recipe(recipe) + super().reset_recipe_state(recipe=recipe) + if recipe is None: + return + + self._configure_quantizer(self.get_quantizer("forward", 0)) + weight = getattr(self, "weight", None) + weight_is_quantized = is_quantized_tensor(weight) + weight_quantizer = self.get_quantizer("forward", 1) + self._configure_quantizer( + weight_quantizer, + internal=not ( + FP8GlobalStateManager.with_fp8_parameters() + or getattr(self, "_with_quantized_weight", False) + or weight_is_quantized + ), + ) + self._configure_quantizer(self.get_quantizer("backward", 0)) + + if isinstance(weight, MXFP8TensorStorage): + if weight._quantizer is not None: + weight_quantizer.set_usage( + rowwise=weight._quantizer.rowwise_usage, + columnwise=weight._quantizer.columnwise_usage, + ) + weight.update_quantizer(weight_quantizer.copy()) + + def pre_fuser_forward(self, *, requires_grad: bool) -> None: + super().pre_fuser_forward(requires_grad=requires_grad) + if not FP8GlobalStateManager.is_fp8_enabled(): + return + self._validate_recipe(FP8GlobalStateManager.get_fp8_recipe()) + self._configure_quantizer(self.get_quantizer("forward", 0)) + self._configure_quantizer( + self.get_quantizer("forward", 1), + internal=not ( + getattr(self, "_with_quantized_weight", False) + or is_quantized_tensor(getattr(self, "weight", None)) + ), + ) + self._configure_quantizer(self.get_quantizer("backward", 0)) + + def _validate_input(self, input_: torch.Tensor) -> int: + if not isinstance(input_, torch.Tensor): + raise TypeError( + f"BatchedLinear expects a torch.Tensor input (got {type(input_).__name__})" + ) + if is_quantized_tensor(input_): + raise ValueError("BatchedLinear expects a high-precision input tensor") + if input_.device.type != "cuda": + raise ValueError(f"BatchedLinear requires a CUDA input tensor (got {input_.device})") + if not devices_match(input_.device, self.weight.device): + raise ValueError( + "BatchedLinear input and weight must be on the same device " + f"(got {input_.device} and {self.weight.device})" + ) + if input_.dtype not in _HIGH_PRECISION_DTYPES: + raise ValueError(f"BatchedLinear input has unsupported dtype {input_.dtype}") + if not input_.is_contiguous(): + raise ValueError("BatchedLinear requires a contiguous input tensor") + if input_.ndim < 2: + raise ValueError( + f"BatchedLinear input must have at least two dimensions (got {input_.ndim})" + ) + + batch_axis = 0 if self.batch_dim == 0 else input_.ndim - 2 + if input_.size(batch_axis) != self.num_gemms: + raise ValueError( + "BatchedLinear input batch dimension has invalid size " + f"(expected {self.num_gemms}, got {input_.size(batch_axis)})" + ) + if input_.size(-1) != self.in_features: + raise ValueError( + "BatchedLinear input feature dimension has invalid size " + f"(expected {self.in_features}, got {input_.size(-1)})" + ) + rows = input_.numel() // (self.num_gemms * self.in_features) + if rows <= 0: + raise ValueError("BatchedLinear does not support empty input matrices") + return rows + + def _validate_parameters(self) -> None: + expected_weight_shape = (self.num_gemms, self.out_features, self.in_features) + if tuple(self.weight.shape) != expected_weight_shape: + raise ValueError( + "BatchedLinear weight has invalid shape " + f"(expected {expected_weight_shape}, got {tuple(self.weight.shape)})" + ) + if self.weight.device.type != "cuda": + raise ValueError(f"BatchedLinear requires a CUDA weight (got {self.weight.device})") + if not self.weight.is_contiguous(): + raise ValueError("BatchedLinear requires a contiguous weight") + if is_quantized_tensor(self.weight) and not isinstance(self.weight, MXFP8TensorStorage): + raise ValueError("BatchedLinear supports only MXFP8 quantized weights") + if isinstance(self.weight, MXFP8TensorStorage) and self.weight._with_gemm_swizzled_scales: + raise ValueError("BatchedLinear quantized weights must use compact MXFP8 scales") + + if self.use_bias: + if self.bias is None: + raise ValueError("BatchedLinear bias parameter is missing") + expected_bias_shape = (self.num_gemms, self.out_features) + if tuple(self.bias.shape) != expected_bias_shape: + raise ValueError( + "BatchedLinear bias has invalid shape " + f"(expected {expected_bias_shape}, got {tuple(self.bias.shape)})" + ) + if not devices_match(self.bias.device, self.weight.device): + raise ValueError( + "BatchedLinear bias and weight must be on the same device " + f"(got {self.bias.device} and {self.weight.device})" + ) + if self.bias.dtype not in _HIGH_PRECISION_DTYPES: + raise ValueError(f"BatchedLinear bias has unsupported dtype {self.bias.dtype}") + if not self.bias.is_contiguous(): + raise ValueError("BatchedLinear requires a contiguous bias") + + def _matrix_strides(self, rows: int, features: int) -> tuple[int, int]: + if self.batch_dim == 0: + return features, rows * features + return self.num_gemms * features, features + + def _validate_mxfp8_dimensions(self, rows: int) -> None: + for name, size in ( + ("rows per GEMM", rows), + ("in_features", self.in_features), + ("out_features", self.out_features), + ): + if size % 32 != 0: + raise ValueError( + f"MXFP8 BatchedLinear requires {name} divisible by 32 (got {size})" + ) + + @staticmethod + def _quantize_for_batched_gemm( + tensor: torch.Tensor, + quantizer: Quantizer, + batch_dim: int, + num_gemms: int, + *, + rowwise: bool, + columnwise: bool, + ) -> MXFP8TensorStorage: + if not isinstance(quantizer, MXFP8Quantizer): + raise RuntimeError("BatchedLinear expected an MXFP8 quantizer") + if not rowwise and not columnwise: + raise RuntimeError("BatchedLinear quantization requires at least one scaling direction") + features = tensor.size(-1) + rows = tensor.numel() // (num_gemms * features) + quantizer_input = tensor + if columnwise and batch_dim == -2: + quantizer_input = tensor.view(rows, num_gemms * features) + quantizer.set_usage(rowwise=rowwise, columnwise=columnwise) + quantizer.internal = True + quantizer.optimize_for_gemm = False + return quantizer(quantizer_input) + + @staticmethod + def _validate_quantized_weight_usage( + weight: MXFP8TensorStorage, + *, + columnwise: bool, + ) -> None: + if weight._rowwise_data is None or weight._rowwise_scale_inv is None: + raise RuntimeError("BatchedLinear MXFP8 weight is missing row-wise data") + if columnwise and (weight._columnwise_data is None or weight._columnwise_scale_inv is None): + raise RuntimeError( + "BatchedLinear MXFP8 weight is missing column-wise data required for backward" + ) + + def _apply_bias(self, output: torch.Tensor, dtype: torch.dtype) -> torch.Tensor: + if not self.apply_bias: + return output + if self.bias is None: + raise RuntimeError("BatchedLinear bias parameter is missing") + bias = maybe_dequantize(self.bias, dtype) + if self.batch_dim == 0: + bias_shape = [self.num_gemms] + [1] * (output.ndim - 2) + [self.out_features] + bias = bias.view(bias_shape) + return output + bias + + def _reduce_bias_gradient(self, grad_output: torch.Tensor) -> torch.Tensor: + if self.batch_dim == 0: + reduce_dims = tuple(range(1, grad_output.ndim - 1)) + else: + reduce_dims = tuple(range(grad_output.ndim - 2)) + if reduce_dims: + grad_output = grad_output.sum(dim=reduce_dims) + if self.bias is not None and grad_output.dtype != self.bias.dtype: + grad_output = grad_output.to(dtype=self.bias.dtype) + return grad_output + + def op_forward( + self, + ctx: OperationContext, + input_: torch.Tensor, + *, + prev_op_grad_output_quantizer: Optional[Quantizer], + next_op_input_quantizer: Optional[Quantizer], + ) -> torch.Tensor: + del prev_op_grad_output_quantizer, next_op_input_quantizer + + rows = self._validate_input(input_) + self._validate_parameters() + dtype = maybe_autocast_dtype(default_dtype=self.weight.dtype) + x = maybe_dequantize(input_, dtype) + input_requires_grad = ctx.requires_grad + weight_requires_grad = ctx.requires_grad and self.weight.requires_grad + bias_requires_grad = ctx.requires_grad and self.bias is not None and self.bias.requires_grad + with_mxfp8_compute = FP8GlobalStateManager.is_fp8_enabled() + + if not with_mxfp8_compute and is_quantized_tensor(self.weight): + warnings.warn( + "BatchedLinear is using an MXFP8 weight without MXFP8 compute. " + "The weight will be dequantized.", + stacklevel=2, + ) + if with_mxfp8_compute and isinstance(self.weight, MXFP8TensorStorage): + w = self.weight + else: + w = maybe_dequantize(self.weight, dtype) + + gemm_x: torch.Tensor = x + gemm_w: torch.Tensor = w + if with_mxfp8_compute: + self._validate_mxfp8_dimensions(rows) + gemm_x = self._quantize_for_batched_gemm( + x, + self.get_quantizer("forward", 0), + self.batch_dim, + self.num_gemms, + rowwise=True, + columnwise=weight_requires_grad, + ) + if isinstance(w, MXFP8TensorStorage): + self._validate_quantized_weight_usage(w, columnwise=input_requires_grad) + gemm_w = w + else: + gemm_w = self._quantize_for_batched_gemm( + w, + self.get_quantizer("forward", 1), + 0, + self.num_gemms, + rowwise=True, + columnwise=input_requires_grad, + ) + + output_shape = list(input_.shape) + output_shape[-1] = self.out_features + output = torch.empty(output_shape, dtype=dtype, device=input_.device) + ldb, strideb = self._matrix_strides(rows, self.in_features) + ldd, strided = self._matrix_strides(rows, self.out_features) + strided_batched_gemm( + gemm_w, + gemm_x, + output, + m=self.out_features, + n=rows, + k=self.in_features, + batch_count=self.num_gemms, + lda=self.in_features, + stridea=self.out_features * self.in_features, + ldb=ldb, + strideb=strideb, + ldd=ldd, + strided=strided, + layout="TN", + use_split_accumulator=_2X_ACC_FPROP, + ) + output = self._apply_bias(output, dtype) + + if ctx.requires_grad: + ctx.save_for_backward( + gemm_x if weight_requires_grad else None, + gemm_w if input_requires_grad else None, + ) + ctx.input_requires_grad = input_requires_grad + ctx.weight_requires_grad = weight_requires_grad + ctx.bias_requires_grad = bias_requires_grad + ctx.input_shape = tuple(input_.shape) + ctx.output_shape = tuple(output_shape) + ctx.rows = rows + ctx.dtype = dtype + ctx.with_mxfp8_compute = with_mxfp8_compute + ctx.grad_output_quantizer = self.get_quantizer("backward", 0) + ctx.apply_bias = self.apply_bias + + return output + + def _op_backward( + self, + ctx: OperationContext, + grad_output: torch.Tensor, + grad_returned_bias: Optional[torch.Tensor], + ) -> tuple[torch.Tensor, Iterable[Optional[torch.Tensor]]]: + x, w = ctx.saved_tensors + dy = maybe_dequantize(grad_output, ctx.dtype).contiguous() + if tuple(dy.shape) != ctx.output_shape: + raise ValueError( + "BatchedLinear grad output has invalid shape " + f"(expected {ctx.output_shape}, got {tuple(dy.shape)})" + ) + + gemm_dy = dy + if ctx.with_mxfp8_compute: + gemm_dy = self._quantize_for_batched_gemm( + dy, + ctx.grad_output_quantizer, + self.batch_dim, + self.num_gemms, + rowwise=ctx.input_requires_grad, + columnwise=ctx.weight_requires_grad, + ) + + grad_input = None + if ctx.input_requires_grad: + if w is None: + raise RuntimeError("BatchedLinear weight was not saved for input gradient") + grad_input = torch.empty(ctx.input_shape, dtype=ctx.dtype, device=dy.device) + ldb, strideb = self._matrix_strides(ctx.rows, self.out_features) + ldd, strided = self._matrix_strides(ctx.rows, self.in_features) + strided_batched_gemm( + w, + gemm_dy, + grad_input, + m=self.in_features, + n=ctx.rows, + k=self.out_features, + batch_count=self.num_gemms, + lda=self.in_features, + stridea=self.out_features * self.in_features, + ldb=ldb, + strideb=strideb, + ldd=ldd, + strided=strided, + layout="NN", + use_split_accumulator=_2X_ACC_DGRAD, + ) + + grad_weight = None + if ctx.weight_requires_grad: + if x is None: + raise RuntimeError("BatchedLinear input was not saved for weight gradient") + accumulate_wgrad = False + if self._accumulate_into_main_grad: + main_grad = get_main_grad_from_param( + self.weight, + op_label="BatchedLinear", + ).detach() + grad_weight = view_main_grad_as_grouped_buffer( + main_grad, + self.num_gemms, + (self.out_features, self.in_features), + label="BatchedLinear weight", + ) + if not grad_weight.is_contiguous(): + raise RuntimeError("BatchedLinear weight main_grad must be contiguous") + if not devices_match(grad_weight.device, dy.device): + raise RuntimeError( + "BatchedLinear weight main_grad must be on the grad output device " + f"(got {grad_weight.device} and {dy.device})" + ) + if grad_weight.dtype not in _HIGH_PRECISION_DTYPES: + raise RuntimeError( + "BatchedLinear weight main_grad must have a high-precision dtype " + f"(got {grad_weight.dtype})" + ) + accumulate_wgrad = get_accumulate_flag_in_param(self.weight) + else: + grad_weight = torch.empty( + self.num_gemms, + self.out_features, + self.in_features, + dtype=ctx.dtype, + device=dy.device, + ) + + lda, stridea = self._matrix_strides(ctx.rows, self.in_features) + ldb, strideb = self._matrix_strides(ctx.rows, self.out_features) + strided_batched_gemm( + x, + gemm_dy, + grad_weight, + m=self.in_features, + n=self.out_features, + k=ctx.rows, + batch_count=self.num_gemms, + lda=lda, + stridea=stridea, + ldb=ldb, + strideb=strideb, + ldd=self.in_features, + strided=self.out_features * self.in_features, + layout="NT", + accumulate=accumulate_wgrad, + use_split_accumulator=_2X_ACC_WGRAD, + ) + + if self._accumulate_into_main_grad: + grad_weight = get_dummy_wgrads_for_params([self.weight])[0] + + grad_bias = None + if ctx.bias_requires_grad: + if ctx.apply_bias: + grad_bias = self._reduce_bias_gradient(dy) + elif grad_returned_bias is not None: + grad_bias = maybe_dequantize(grad_returned_bias, self.bias.dtype) + + grad_params = [grad_weight] + if self.use_bias: + grad_params.append(grad_bias) + return grad_input, grad_params + + def op_backward( + self, + ctx: OperationContext, + grad_output: torch.Tensor, + ) -> tuple[torch.Tensor, Iterable[Optional[torch.Tensor]]]: + return self._op_backward(ctx, grad_output, None) + + def fuser_forward( + self, + basic_op_ctxs: list[OperationContext], + input_: torch.Tensor, + *, + basic_op_extra_inputs: Sequence[Sequence[Optional[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, Sequence[Sequence[Optional[torch.Tensor]]]]: + del basic_op_extra_inputs + output = self.op_forward( + basic_op_ctxs[0], + input_, + prev_op_grad_output_quantizer=prev_op_grad_output_quantizer, + next_op_input_quantizer=next_op_input_quantizer, + **basic_op_kwargs[0], + ) + if self.num_extra_outputs == 0: + return output, [()] + if self.bias is None: + raise RuntimeError("BatchedLinear bias parameter is missing") + return output, [(maybe_dequantize(self.bias, output.dtype),)] + + def fuser_backward( + self, + basic_op_ctxs: list[OperationContext], + grad_output: torch.Tensor, + *, + basic_op_grad_extra_outputs: Sequence[Sequence[Optional[torch.Tensor]]], + ) -> tuple[ + torch.Tensor, + Sequence[Sequence[Optional[torch.Tensor]]], + Sequence[Sequence[Optional[torch.Tensor]]], + ]: + grad_returned_bias = None + if self.num_extra_outputs > 0: + grad_returned_bias = basic_op_grad_extra_outputs[0][0] + grad_input, grad_params = self._op_backward( + basic_op_ctxs[0], + grad_output, + grad_returned_bias, + ) + return grad_input, [grad_params], [()] + + def forward( + self, + input: torch.Tensor, # pylint: disable=redefined-builtin + *extra_inputs: torch.Tensor, + **kwargs: Any, + ) -> torch.Tensor | tuple[torch.Tensor, Optional[torch.Tensor]]: + output = super().forward(input, *extra_inputs, **kwargs) + if self.return_bias and not self.use_bias: + return output, None + return output From cf6b6dd214e691f7bb57bc5b3228961f402bcddf Mon Sep 17 00:00:00 2001 From: Xin Yao Date: Wed, 2 Sep 2026 02:56:40 -0700 Subject: [PATCH 3/3] add save_original_input Signed-off-by: Xin Yao --- tests/pytorch/test_batched_linear.py | 78 +++++++++++++++++++ .../pytorch/ops/basic/batched_linear.py | 28 ++++++- 2 files changed, 104 insertions(+), 2 deletions(-) diff --git a/tests/pytorch/test_batched_linear.py b/tests/pytorch/test_batched_linear.py index afd1d25252..9810a0342a 100644 --- a/tests/pytorch/test_batched_linear.py +++ b/tests/pytorch/test_batched_linear.py @@ -233,6 +233,84 @@ def test_mxfp8_selective_gradients( assert (op.weight.grad is not None) == weight_requires_grad +@pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8) +@pytest.mark.parametrize("batch_dim", (0, -2)) +def test_save_original_input(monkeypatch: pytest.MonkeyPatch, batch_dim: int) -> None: + """MXFP8 backward reconstructs its column-wise operand from the original input.""" + dtype = torch.bfloat16 + num_gemms, rows = 2, 32 + in_features, out_features = 64, 32 + input_shape = ( + (num_gemms, rows, in_features) if batch_dim == 0 else (rows, num_gemms, in_features) + ) + output_shape = list(input_shape) + output_shape[-1] = out_features + + op = te_ops.BatchedLinear( + num_gemms, + in_features, + out_features, + batch_dim=batch_dim, + bias=False, + save_original_input=True, + dtype=dtype, + device="cuda", + ) + inp = torch.rand(input_shape, dtype=dtype, device="cuda", requires_grad=True) + grad_output = torch.rand(output_shape, dtype=dtype, device="cuda") + + inp_ref = _to_reference(inp, requires_grad=True) + weight_ref = _to_reference(op.weight, requires_grad=True) + output_ref = ( + torch.einsum("g...d,grd->g...r", inp_ref, weight_ref) + if batch_dim == 0 + else torch.einsum("...gd,grd->...gr", inp_ref, weight_ref) + ) + output_ref.backward(_to_reference(grad_output, requires_grad=False)) + + input_quantizations = [] + original_quantize = batched_linear_op.BatchedLinear._quantize_for_batched_gemm + + def capture_quantize(tensor, quantizer, batch_dim_, num_gemms_, *, rowwise, columnwise): + quantized = original_quantize( + tensor, + quantizer, + batch_dim_, + num_gemms_, + rowwise=rowwise, + columnwise=columnwise, + ) + if tensor.data_ptr() == inp.data_ptr(): + input_quantizations.append((quantized, rowwise, columnwise)) + return quantized + + monkeypatch.setattr( + batched_linear_op.BatchedLinear, + "_quantize_for_batched_gemm", + staticmethod(capture_quantize), + ) + + recipe = make_recipe("mxfp8") + with te.autocast(enabled=True, recipe=recipe): + output = op(inp) + output.backward(grad_output) + + assert [(rowwise, columnwise) for _, rowwise, columnwise in input_quantizations] == [ + (True, False), + (False, True), + ] + forward_input, backward_input = (entry[0] for entry in input_quantizations) + assert forward_input._rowwise_data is not None + assert forward_input._columnwise_data is None + assert backward_input._rowwise_data is None + assert backward_input._columnwise_data is not None + + tols = quantization_tols("mxfp8") + assert_close(output, output_ref, **tols) + assert_close(inp.grad, inp_ref.grad, **tols) + assert_close(op.weight.grad, weight_ref.grad, **tols) + + @pytest.mark.parametrize("deferred_init", (False, True)) def test_init_method_and_rng_tracker(deferred_init: bool) -> None: """Initialization uses the requested RNG tracker, including after meta init.""" diff --git a/transformer_engine/pytorch/ops/basic/batched_linear.py b/transformer_engine/pytorch/ops/basic/batched_linear.py index d40291311b..eb754961ae 100644 --- a/transformer_engine/pytorch/ops/basic/batched_linear.py +++ b/transformer_engine/pytorch/ops/basic/batched_linear.py @@ -76,6 +76,10 @@ class BatchedLinear(BasicOperation): Write weight gradients directly into the externally allocated ``weight.main_grad`` buffer. Setting ``weight.overwrite_main_grad`` to ``True`` overwrites that buffer instead of accumulating into it. + save_original_input : bool, default = False + Save the original high-precision input for the weight-gradient + computation instead of the cast or quantized forward operand. With + MXFP8 compute, the column-wise input is quantized during backward. init_method : callable, optional Weight initialization method. The default is TE's normal initializer. name : str, optional @@ -103,6 +107,7 @@ def __init__( dtype: Optional[torch.dtype] = None, rng_state_tracker_function: Optional[Callable[[], CudaRNGStatesTracker]] = None, accumulate_into_main_grad: bool = False, + save_original_input: bool = False, init_method: Optional[Callable[[torch.Tensor], torch.Tensor]] = None, name: Optional[str] = None, ) -> None: @@ -139,6 +144,7 @@ def __init__( self.name = name self._rng_state_tracker_function = rng_state_tracker_function self._accumulate_into_main_grad = accumulate_into_main_grad + self.save_original_input = save_original_input self._init_method = get_default_init_method() if init_method is None else init_method # Initialize recipe state if the weight itself is stored in MXFP8. @@ -490,6 +496,7 @@ def op_forward( input_requires_grad = ctx.requires_grad weight_requires_grad = ctx.requires_grad and self.weight.requires_grad bias_requires_grad = ctx.requires_grad and self.bias is not None and self.bias.requires_grad + save_original_input = self.save_original_input and weight_requires_grad with_mxfp8_compute = FP8GlobalStateManager.is_fp8_enabled() if not with_mxfp8_compute and is_quantized_tensor(self.weight): @@ -513,7 +520,7 @@ def op_forward( self.batch_dim, self.num_gemms, rowwise=True, - columnwise=weight_requires_grad, + columnwise=weight_requires_grad and not save_original_input, ) if isinstance(w, MXFP8TensorStorage): self._validate_quantized_weight_usage(w, columnwise=input_requires_grad) @@ -554,7 +561,7 @@ def op_forward( if ctx.requires_grad: ctx.save_for_backward( - gemm_x if weight_requires_grad else None, + (input_ if save_original_input else gemm_x) if weight_requires_grad else None, gemm_w if input_requires_grad else None, ) ctx.input_requires_grad = input_requires_grad @@ -565,6 +572,12 @@ def op_forward( ctx.rows = rows ctx.dtype = dtype ctx.with_mxfp8_compute = with_mxfp8_compute + ctx.save_original_input = save_original_input + ctx.input_quantizer = ( + self.get_quantizer("forward", 0) + if save_original_input and with_mxfp8_compute + else None + ) ctx.grad_output_quantizer = self.get_quantizer("backward", 0) ctx.apply_bias = self.apply_bias @@ -624,6 +637,17 @@ def _op_backward( if ctx.weight_requires_grad: if x is None: raise RuntimeError("BatchedLinear input was not saved for weight gradient") + if ctx.save_original_input: + x = maybe_dequantize(x, ctx.dtype) + if ctx.with_mxfp8_compute: + x = self._quantize_for_batched_gemm( + x, + ctx.input_quantizer, + self.batch_dim, + self.num_gemms, + rowwise=False, + columnwise=True, + ) accumulate_wgrad = False if self._accumulate_into_main_grad: main_grad = get_main_grad_from_param(