diff --git a/ci/pytorch.sh b/ci/pytorch.sh index 1cccdff1b..c3a8d4f8e 100755 --- a/ci/pytorch.sh +++ b/ci/pytorch.sh @@ -58,6 +58,8 @@ run_test_config(){ run_default_fa 1 test_fused_router.py run_default_fa 1 test_fusible_ops.py run_default_fa 1 test_gemm_autotune.py + NVTE_USE_GEMM_TRITON=1 run_default_fa_lbl "triton" 1 triton_kernels/test_gemm.py + NVTE_USE_GEMM_TRITON=1 run_default_fa_lbl "triton" 1 triton_kernels/test_gemm_kernel.py run 1 test_gqa.py run 1 test_jit.py NVTE_ROCM_ENABLE_MXFP8=1 run_default_fa 1 test_multi_tensor.py @@ -85,6 +87,9 @@ run_test_config(){ NVTE_USE_DEQUANTIZE_TRITON=1 NVTE_USE_CAST_TRANSPOSE_TRITON=1 NVTE_USE_RMSNORM_TRITON=1 NVTE_USE_LAYERNORM_TRITON=1 run_default_fa_lbl "triton" 3 test_numerics.py NVTE_USE_CAST_TRANSPOSE_TRITON=1 NVTE_USE_RMSNORM_TRITON=1 run_default_fa_lbl "triton" 1 test_fusible_ops.py NVTE_USE_CAST_TRANSPOSE_TRITON=1 run_default_fa_lbl "triton" 1 test_float8_current_scaling_exact.py + NVTE_USE_GEMM_TRITON=1 run_default_fa_lbl "gemm-triton" 3 test_numerics.py + NVTE_USE_GEMM_TRITON=1 run_default_fa_lbl "gemm-triton" 1 test_fusible_ops.py + NVTE_USE_GEMM_TRITON=1 run_default_fa_lbl "gemm-triton" 1 test_float8_current_scaling_exact.py NVTE_USE_ATOMIC_AMAX=1 run_default_fa_lbl "amax" 3 test_numerics.py NVTE_USE_ATOMIC_AMAX=1 run_default_fa_lbl "amax" 3 test_fusible_ops.py NVTE_USE_ATOMIC_AMAX=1 NVTE_USE_CAST_TRANSPOSE_TRITON=1 run_default_fa_lbl "amax+triton" 3 test_numerics.py diff --git a/tests/pytorch/conftest.py b/tests/pytorch/conftest.py new file mode 100644 index 000000000..4549c743e --- /dev/null +++ b/tests/pytorch/conftest.py @@ -0,0 +1,42 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. + +"""Convert Triton GEMM backend refusals (HYBRID, mixed FP8, NVFP4, ...) +into pytest.skip so the CI sweep under NVTE_USE_GEMM_TRITON=1 does not +flag intentionally-unsupported combinations as failures. + +The gates raise ValueError from quantization.py and +triton_kernels/gemm/gemm_wrapper.py; when they are relaxed the marker +text disappears and this hook stops firing. +""" + +import pytest + + +# Substrings identifying our Triton GEMM backend refusals. Kept short so +# they are easy to grep. +_TRITON_GEMM_GATE_MARKERS = ( + # Mixed FP8 (e4m3 x e5m2) refused at the low-level matmul entry. + "Mixed FP8 types", + # Covers both quantization.py::check_recipe_support (HYBRID) and + # gemm_wrapper._classify_input's refusal of NVFP4 / other + # QuantizedTensorStorage subclasses. + "The Triton GEMM backend (NVTE_USE_GEMM_TRITON=1) does not support", +) + + +@pytest.hookimpl(hookwrapper=True) +def pytest_runtest_call(item): + """Convert known Triton GEMM backend gate errors into skips.""" + outcome = yield + if outcome.excinfo is None: + return + exc = outcome.excinfo[1] + if not isinstance(exc, ValueError): + return + msg = str(exc) + if any(marker in msg for marker in _TRITON_GEMM_GATE_MARKERS): + outcome.force_exception( + pytest.skip.Exception(f"Triton GEMM backend gate: {msg}") + ) diff --git a/tests/pytorch/test_numerics.py b/tests/pytorch/test_numerics.py index 5c9686f15..d35198d0f 100644 --- a/tests/pytorch/test_numerics.py +++ b/tests/pytorch/test_numerics.py @@ -196,6 +196,27 @@ def get_nvfp4_inp_supported_dtypes(recipe: recipe.Recipe, dtype: torch.dtype) -> use_cutlass_grouped_gemm.append(True) +# Marker for tests that compare a sequence of individual `Linear` layers +# against a multi-stream / grouped GEMM (or that otherwise assume a +# bit-exact `rtol=0, atol=0` match). The original invariant these tests +# check is that the same GEMM kernel, used both sequentially and inside +# the multi-stream GEMM, reproduces the same result -- see the upstream +# `# cuBLAS implementation should be bit-wise match` comment. +# +# Under NVTE_USE_GEMM_TRITON=1 the Triton GEMM kernel is NOT used inside +# the multi-stream GEMM path, so the invariant no longer applies to these +# tests. Skip them under the override rather than loosen tolerances just +# to make them pass vacuously. +_skip_grouped_under_gemm_triton = pytest.mark.skipif( + bool(int(os.environ.get("NVTE_USE_GEMM_TRITON", "0"))), + reason=( + "The Triton GEMM kernel is not used inside the multi-stream GEMM " + "path, so the sequential-vs-multi-stream equivalence these tests " + "check does not apply under NVTE_USE_GEMM_TRITON=1." + ), +) + + def get_causal_attn_mask(sq: int) -> torch.Tensor: return torch.triu(torch.ones(sq, sq, device="cuda"), diagonal=1).bool() @@ -2103,6 +2124,7 @@ def _test_grouped_linear_accuracy( return outputs +@_skip_grouped_under_gemm_triton @pytest.mark.parametrize("dtype", param_types, ids=str) @pytest.mark.parametrize("num_gemms", [3, 6]) @pytest.mark.parametrize("bs", batch_sizes) @@ -2231,6 +2253,7 @@ def test_grouped_linear_accuracy( torch.testing.assert_close(o, o_ref, rtol=rtol, atol=atol) +@_skip_grouped_under_gemm_triton @pytest.mark.skipif( torch.cuda.get_device_capability() != (9, 0) and not IS_HIP_EXTENSION, reason="Only enable CUTLASS grouped gemm on Hopper", @@ -2277,6 +2300,7 @@ def test_grouped_linear_accuracy_cutlass( os.environ.pop("NVTE_USE_CUTLASS_GROUPED_GEMM", None) +@_skip_grouped_under_gemm_triton @pytest.mark.parametrize("dtype", param_types, ids=str) @pytest.mark.parametrize("num_gemms", [3]) @pytest.mark.parametrize("bs", [1]) @@ -2383,6 +2407,7 @@ def test_grouped_linear_accuracy_save_original_input( torch.testing.assert_close(o, o_ref, rtol=0, atol=0) +@_skip_grouped_under_gemm_triton @pytest.mark.parametrize("recipe", fp8_recipes + [None]) def test_grouped_linear_accuracy_single_gemm(recipe): """Split the tests to save CI time""" @@ -2493,6 +2518,7 @@ def _generate_random_numbers(n, total_sum): return outputs +@_skip_grouped_under_gemm_triton @pytest.mark.parametrize("dtype", param_types) @pytest.mark.parametrize("num_gemms", [3, 6]) @pytest.mark.parametrize("bs", batch_sizes) @@ -2568,6 +2594,7 @@ def test_padding_grouped_linear_accuracy( torch.testing.assert_close(o, o_ref, rtol=0, atol=0) +@_skip_grouped_under_gemm_triton @pytest.mark.parametrize("dtype", param_types) @pytest.mark.parametrize("num_gemms", [3]) @pytest.mark.parametrize("bs", [1]) @@ -2988,6 +3015,7 @@ def test_transformer_layer_hidden_states_format(dtype, bs, model): ) +@_skip_grouped_under_gemm_triton @pytest.mark.parametrize( "shape", [ @@ -3740,6 +3768,7 @@ def test_fp8gemm_with_unfused_quantization(N, datatype, input_quantizer, out_qua torch.testing.assert_close(expected_quantized_out.dequantize(), quantized_out.dequantize()) +@_skip_grouped_under_gemm_triton @pytest.mark.parametrize( "shape", [ diff --git a/tests/pytorch/triton_kernels/test_gemm.py b/tests/pytorch/triton_kernels/test_gemm.py new file mode 100644 index 000000000..2e1feba7c --- /dev/null +++ b/tests/pytorch/triton_kernels/test_gemm.py @@ -0,0 +1,506 @@ +# Copyright (c) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. +# +# License for AMD contributions = MIT. See LICENSE for more information + +"""User-facing Triton GEMM tests -- ``general_gemm()`` under ``NVTE_USE_GEMM_TRITON=1``. + +Exercises the same public entry point that TE ``Linear`` / ``LayerNormLinear`` +use. Covers the full precision surface (fp32 / fp16 / bf16 / same-format FP8 / +mixed FP8 / MXFP8) and all three layouts (TN / NN / NT), plus bias / bias-grad +epilogues and a batched-fp8 multidim case. + +Each test compares the Triton path against two independent references: + + 1. ``torch.matmul`` on dequantized inputs -- catches functional bugs + independent of any hipBLASLt behavior. + 2. The C++ ``tex.generic_gemm`` backend under the same TE ``general_gemm`` + surface -- catches divergence from the production path. + +Complementary file: + +- ``test_gemm_kernel.py`` -- low-level ``te_gemm_triton()`` kernel-direct + correctness (bypasses ``general_gemm`` and the wrappers; ~1000 + parametrizations of a single ``test_correctness`` function). + +MXFP8 end-to-end numerical correctness lives here, in +``test_triton_vs_pytorch_mxfp8`` / ``test_triton_vs_cpp_mxfp8``. +""" + +import os +import pytest +import torch + +from transformer_engine.pytorch.cpp_extensions.gemm import general_gemm +from transformer_engine.pytorch import Float8Tensor +from transformer_engine.pytorch.tensor.float8_tensor import Float8Quantizer +from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Quantizer, MXFP8Tensor +import transformer_engine_torch as tex + +# --- Feature detection -------------------------------------------------------- + +major, minor = torch.cuda.get_device_capability() +is_gfx950 = (major == 9 and minor >= 5) + +from transformer_engine.pytorch import torch_version +_torch_ver = torch_version() + +requires_gfx950 = pytest.mark.skipif( + not is_gfx950, + reason="MXFP8 requires gfx950 (compute capability >= 9.5)", +) + +# --- Test parameters ---------------------------------------------------------- + +REGULAR_FP8_SHAPES = [ + (2304, 768, 4096), + (768, 768, 4096), + (768, 3072, 4096), + (229, 541, 541), + (71, 71, 3571), + (29, 29, 17389), +] + +MXFP8_SHAPES = [ + (128, 256, 512), + (768, 768, 4096), +] + +LAYOUTS = ["TN", "NN", "NT"] + +FP8_FORMAT_COMBOS = [ + (tex.DType.kFloat8E4M3, tex.DType.kFloat8E4M3), + (tex.DType.kFloat8E5M2, tex.DType.kFloat8E5M2), +] + +# Mixed FP8 formats are disabled due to a Triton compiler bug on gfx950: +# when the MFMA layout is transposed, operand B is packed using A's element type, +# and the instruction format encoding doesn't account for the operand swap. +# This affects both v_mfma_f32_32x32x16_{fp8|bf8} and v_mfma_f32_32x32x64_f8f6f4. +# Fixed upstream in triton-lang/triton PR #9567 (commit eaaa75cf5, 2026-02-27). +# Not yet in any pytorch-triton-rocm release as of PyTorch 2.11. +# TODO: Re-enable once pytorch-triton-rocm includes the fix (expected PyTorch 2.12+). +FP8_MIXED_FORMAT_COMBOS = [ + (tex.DType.kFloat8E4M3, tex.DType.kFloat8E5M2), + (tex.DType.kFloat8E5M2, tex.DType.kFloat8E4M3), +] + +REGULAR_DTYPES = [torch.float32, torch.float16, torch.bfloat16] + +# --- Fixtures ----------------------------------------------------------------- + +@pytest.fixture(autouse=True) +def cleanup_env(): + """Save/restore NVTE_USE_GEMM_TRITON env var between tests.""" + old_val = os.environ.get('NVTE_USE_GEMM_TRITON', None) + yield + if old_val is None: + os.environ.pop('NVTE_USE_GEMM_TRITON', None) + else: + os.environ['NVTE_USE_GEMM_TRITON'] = old_val + + +# --- Helpers ------------------------------------------------------------------ + +def get_shapes(layout, M, K, N): + """Returns (A_shape, B_shape) based on layout.""" + if layout == "TN": + return (M, K), (N, K) + elif layout == "NN": + return (M, K), (K, M) + elif layout == "NT": + return (M, K), (M, K) + else: + raise ValueError(f"Unsupported layout: {layout}") + + +def compute_pytorch_reference(A_ref, B_ref, layout): + """torch.matmul with correct transpose for layout.""" + if layout == "TN": + return torch.matmul(B_ref, A_ref.T) + elif layout == "NN": + return torch.matmul(B_ref, A_ref) + elif layout == "NT": + return torch.matmul(B_ref.T, A_ref) + else: + raise ValueError(f"Unsupported layout: {layout}") + + +def create_fp8_tensors(M, K, N, layout, fp8_dtype_a, fp8_dtype_b): + """Create Float8Tensor inputs and dequantized references.""" + A_shape, B_shape = get_shapes(layout, M, K, N) + A_f32 = torch.randn(A_shape, dtype=torch.float32, device='cuda') * 0.5 + B_f32 = torch.randn(B_shape, dtype=torch.float32, device='cuda') * 0.5 + + A_fp8 = Float8Quantizer( + scale=torch.full([1], 1.0, dtype=torch.float32, device='cuda'), + amax=torch.empty([1], dtype=torch.float32, device='cuda'), + fp8_dtype=fp8_dtype_a, + )(A_f32) + B_fp8 = Float8Quantizer( + scale=torch.full([1], 1.0, dtype=torch.float32, device='cuda'), + amax=torch.empty([1], dtype=torch.float32, device='cuda'), + fp8_dtype=fp8_dtype_b, + )(B_f32) + + return A_fp8, B_fp8, A_fp8.dequantize(), B_fp8.dequantize() + + +def create_mxfp8_tensors(M, K, N, layout): + """Create MXFP8Tensor inputs and dequantized references.""" + A_shape, B_shape = get_shapes(layout, M, K, N) + A_f32 = torch.randn(A_shape, dtype=torch.float32, device='cuda') * 0.5 + B_f32 = torch.randn(B_shape, dtype=torch.float32, device='cuda') * 0.5 + + quantizer = MXFP8Quantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + rowwise=True, + columnwise=True, + ) + A_mxfp8 = quantizer(A_f32) + B_mxfp8 = quantizer(B_f32) + + return A_mxfp8, B_mxfp8, A_mxfp8.dequantize(), B_mxfp8.dequantize() + + +def call_gemm(A, B, layout, out_dtype, use_triton=True): + """Call general_gemm() with appropriate env var setting.""" + os.environ['NVTE_USE_GEMM_TRITON'] = '1' if use_triton else '0' + output, _, _, _ = general_gemm( + A=A, + B=B, + out_dtype=out_dtype, + layout=layout, + ) + return output + + +def call_gemm_with_bias(A, B, layout, out_dtype, bias, grad, use_triton=True): + """Call general_gemm() with a bias argument. + + Returns (output, bias_grad). When grad=True the GEMM uses the BGRADB + epilogue and bias_grad contains the reduced bias gradient; otherwise + it uses the BIAS epilogue and bias is fused into the output. + """ + os.environ['NVTE_USE_GEMM_TRITON'] = '1' if use_triton else '0' + output, bias_grad, _, _ = general_gemm( + A=A, + B=B, + out_dtype=out_dtype, + layout=layout, + bias=bias, + grad=grad, + ) + return output, bias_grad + + +# ============================================================================== +# Approach 1: Triton vs PyTorch torch.matmul reference +# ============================================================================== + +@pytest.mark.parametrize("M, K, N", REGULAR_FP8_SHAPES) +@pytest.mark.parametrize("layout", LAYOUTS) +@pytest.mark.parametrize("dtype", REGULAR_DTYPES, ids=["fp32", "fp16", "bf16"]) +def test_triton_vs_pytorch_regular(M, K, N, layout, dtype): + """Test Triton GEMM vs torch.matmul for regular tensors.""" + torch.manual_seed(42) + A_shape, B_shape = get_shapes(layout, M, K, N) + A = torch.randn(A_shape, dtype=dtype, device='cuda') * 0.5 + B = torch.randn(B_shape, dtype=dtype, device='cuda') * 0.5 + + # Triton result + output = call_gemm(A, B, layout, out_dtype=dtype, use_triton=True) + + # PyTorch reference on fp32 copies + expected = compute_pytorch_reference(A.float(), B.float(), layout) + + torch.testing.assert_close( + output.float(), expected.float(), + atol=1e-3, rtol=1e-2, + ) + + +@pytest.mark.parametrize("M, K, N", REGULAR_FP8_SHAPES) +@pytest.mark.parametrize("layout", LAYOUTS) +@pytest.mark.parametrize("fp8_format", FP8_FORMAT_COMBOS, + ids=["e4m3_e4m3", "e5m2_e5m2"]) +def test_triton_vs_pytorch_fp8(M, K, N, layout, fp8_format): + """Test Triton GEMM vs torch.matmul for Float8Tensor inputs.""" + torch.manual_seed(42) + fp8_dtype_a, fp8_dtype_b = fp8_format + A_fp8, B_fp8, A_deq, B_deq = create_fp8_tensors(M, K, N, layout, fp8_dtype_a, fp8_dtype_b) + + output = call_gemm(A_fp8, B_fp8, layout, out_dtype=torch.float32, use_triton=True) + expected = compute_pytorch_reference(A_deq.float(), B_deq.float(), layout) + + torch.testing.assert_close( + output.float(), expected.float(), + atol=5e-3, rtol=1e-2, + ) + + +@pytest.mark.skip(reason="Triton compiler bug with mixed FP8 formats (triton-lang/triton#9567)") +@pytest.mark.parametrize("M, K, N", REGULAR_FP8_SHAPES) +@pytest.mark.parametrize("layout", LAYOUTS) +@pytest.mark.parametrize("fp8_format", FP8_MIXED_FORMAT_COMBOS, + ids=["e4m3_e5m2", "e5m2_e4m3"]) +def test_triton_vs_pytorch_fp8_mixed(M, K, N, layout, fp8_format): + """Test Triton GEMM vs torch.matmul for mixed Float8Tensor formats.""" + torch.manual_seed(42) + fp8_dtype_a, fp8_dtype_b = fp8_format + A_fp8, B_fp8, A_deq, B_deq = create_fp8_tensors(M, K, N, layout, fp8_dtype_a, fp8_dtype_b) + + output = call_gemm(A_fp8, B_fp8, layout, out_dtype=torch.float32, use_triton=True) + expected = compute_pytorch_reference(A_deq.float(), B_deq.float(), layout) + + torch.testing.assert_close( + output.float(), expected.float(), + atol=5e-3, rtol=1e-2, + ) + + +@requires_gfx950 +@pytest.mark.skipif( + _torch_ver < (2, 10), + reason=( + "Triton tl.dot_scaled() RHS scale bug fixed in PyTorch 2.10 " + f"(found {_torch_ver}). The TE kernel uses the new dot_scaled API " + "(rhs_scale in [N, K//32] layout) which requires PyTorch >= 2.10." + ), +) +@pytest.mark.parametrize("M, K, N", MXFP8_SHAPES) +@pytest.mark.parametrize("layout", LAYOUTS) +def test_triton_vs_pytorch_mxfp8(M, K, N, layout): + """Test Triton GEMM vs torch.matmul for MXFP8Tensor inputs.""" + torch.manual_seed(42) + A_mxfp8, B_mxfp8, A_deq, B_deq = create_mxfp8_tensors(M, K, N, layout) + + output = call_gemm(A_mxfp8, B_mxfp8, layout, out_dtype=torch.bfloat16, use_triton=True) + expected = compute_pytorch_reference(A_deq.float(), B_deq.float(), layout) + + torch.testing.assert_close( + output.float(), expected.float(), + atol=5e-3, rtol=1e-2, + ) + + +# ============================================================================== +# Approach 2: Triton vs C++ tex.generic_gemm reference +# ============================================================================== + +@pytest.mark.parametrize("M, K, N", REGULAR_FP8_SHAPES) +@pytest.mark.parametrize("layout", LAYOUTS) +@pytest.mark.parametrize("dtype", REGULAR_DTYPES, ids=["fp32", "fp16", "bf16"]) +def test_triton_vs_cpp_regular(M, K, N, layout, dtype): + """Test Triton GEMM vs C++ generic_gemm for regular tensors.""" + torch.manual_seed(42) + A_shape, B_shape = get_shapes(layout, M, K, N) + A = torch.randn(A_shape, dtype=dtype, device='cuda') * 0.5 + B = torch.randn(B_shape, dtype=dtype, device='cuda') * 0.5 + + triton_out = call_gemm(A, B, layout, out_dtype=dtype, use_triton=True) + cpp_out = call_gemm(A, B, layout, out_dtype=dtype, use_triton=False) + + torch.testing.assert_close( + triton_out.float(), cpp_out.float(), + atol=1e-3, rtol=1e-2, + ) + + +@pytest.mark.parametrize("M, K, N", REGULAR_FP8_SHAPES) +@pytest.mark.parametrize("layout", LAYOUTS) +@pytest.mark.parametrize("fp8_format", FP8_FORMAT_COMBOS, + ids=["e4m3_e4m3", "e5m2_e5m2"]) +def test_triton_vs_cpp_fp8(M, K, N, layout, fp8_format): + """Test Triton GEMM vs C++ generic_gemm for Float8Tensor inputs.""" + torch.manual_seed(42) + fp8_dtype_a, fp8_dtype_b = fp8_format + A_fp8, B_fp8, _, _ = create_fp8_tensors(M, K, N, layout, fp8_dtype_a, fp8_dtype_b) + + triton_out = call_gemm(A_fp8, B_fp8, layout, out_dtype=torch.float32, use_triton=True) + cpp_out = call_gemm(A_fp8, B_fp8, layout, out_dtype=torch.float32, use_triton=False) + + torch.testing.assert_close( + triton_out.float(), cpp_out.float(), + atol=5e-3, rtol=1e-2, + ) + + +@pytest.mark.skip(reason="Triton compiler bug with mixed FP8 formats (triton-lang/triton#9567)") +@pytest.mark.parametrize("M, K, N", REGULAR_FP8_SHAPES) +@pytest.mark.parametrize("layout", LAYOUTS) +@pytest.mark.parametrize("fp8_format", FP8_MIXED_FORMAT_COMBOS, + ids=["e4m3_e5m2", "e5m2_e4m3"]) +def test_triton_vs_cpp_fp8_mixed(M, K, N, layout, fp8_format): + """Test Triton GEMM vs C++ generic_gemm for mixed Float8Tensor formats.""" + torch.manual_seed(42) + fp8_dtype_a, fp8_dtype_b = fp8_format + A_fp8, B_fp8, _, _ = create_fp8_tensors(M, K, N, layout, fp8_dtype_a, fp8_dtype_b) + + triton_out = call_gemm(A_fp8, B_fp8, layout, out_dtype=torch.float32, use_triton=True) + cpp_out = call_gemm(A_fp8, B_fp8, layout, out_dtype=torch.float32, use_triton=False) + + torch.testing.assert_close( + triton_out.float(), cpp_out.float(), + atol=5e-3, rtol=1e-2, + ) + + +@requires_gfx950 +@pytest.mark.skipif( + _torch_ver < (2, 10), + reason=( + "Triton tl.dot_scaled() RHS scale bug fixed in PyTorch 2.10 " + f"(found {_torch_ver}). The TE kernel uses the new dot_scaled API " + "(rhs_scale in [N, K//32] layout) which requires PyTorch >= 2.10." + ), +) +@pytest.mark.parametrize("M, K, N", MXFP8_SHAPES) +@pytest.mark.parametrize("layout", LAYOUTS) +def test_triton_vs_cpp_mxfp8(M, K, N, layout): + """Test Triton GEMM vs C++ generic_gemm for MXFP8Tensor inputs.""" + torch.manual_seed(42) + A_mxfp8, B_mxfp8, _, _ = create_mxfp8_tensors(M, K, N, layout) + + triton_out = call_gemm(A_mxfp8, B_mxfp8, layout, out_dtype=torch.bfloat16, use_triton=True) + cpp_out = call_gemm(A_mxfp8, B_mxfp8, layout, out_dtype=torch.bfloat16, use_triton=False) + + torch.testing.assert_close( + triton_out.float(), cpp_out.float(), + atol=5e-3, rtol=1e-2, + ) + + +# ============================================================================== +# Bias epilogue coverage (regression guard for gemm_triton.py bias wiring) +# +# The Triton wrapper must honor the `bias` + `grad` arguments to general_gemm: +# - grad=False + bias present → BIAS epilogue, bias added to output +# - grad=True + bias present → BGRADB epilogue, bias gradient returned as +# the second element of general_gemm's tuple +# Layout TN matches TE Linear's forward convention: A=weight[M,K], +# B=input[N,K], output[N,M]; BIAS reads bias[M], BGRADB reduces to shape [N]. +# ============================================================================== + +BIAS_SHAPES = [(128, 256, 512), (229, 541, 541), (71, 71, 3571)] + + +@pytest.mark.parametrize("M, K, N", BIAS_SHAPES) +@pytest.mark.parametrize("dtype", REGULAR_DTYPES, ids=["fp32", "fp16", "bf16"]) +def test_triton_vs_cpp_bias_forward(M, K, N, dtype): + """Forward with BIAS epilogue: Triton must match C++ when bias is fused.""" + torch.manual_seed(42) + A_shape, B_shape = get_shapes("TN", M, K, N) + A = torch.randn(A_shape, dtype=dtype, device='cuda') * 0.5 + B = torch.randn(B_shape, dtype=dtype, device='cuda') * 0.5 + bias = torch.randn((M,), dtype=dtype, device='cuda') + + triton_out, _ = call_gemm_with_bias(A, B, "TN", dtype, bias, grad=False, use_triton=True) + cpp_out, _ = call_gemm_with_bias(A, B, "TN", dtype, bias, grad=False, use_triton=False) + + # Bias must actually change the result vs. no-bias path; otherwise BIAS + # silently reverted to DEFAULT would pass a simple Triton-vs-C++ check. + no_bias_out = call_gemm(A, B, "TN", out_dtype=dtype, use_triton=True) + assert not torch.allclose(triton_out.float(), no_bias_out.float(), atol=1e-4), ( + "Triton output matches no-bias output; BIAS epilogue appears inactive." + ) + + torch.testing.assert_close( + triton_out.float(), cpp_out.float(), + atol=5e-3, rtol=1e-2, + ) + + +WGRAD_SHAPES = [ + # (batch*seq, in_features, out_features) — TE Linear wgrad pattern + (256, 128, 512), + (512, 541, 229), + (128, 3571, 71), +] + + +@pytest.mark.parametrize("batch, in_features, out_features", WGRAD_SHAPES) +@pytest.mark.parametrize("dtype", REGULAR_DTYPES, ids=["fp32", "fp16", "bf16"]) +def test_triton_vs_cpp_bias_grad(batch, in_features, out_features, dtype): + """Backward with BGRADB epilogue: Triton must produce the correct bias gradient. + + Exercises the same call shape TE Linear uses for weight-grad: + general_gemm(x, dy, layout="NT", bias=, grad=True) + A=x[batch, in_features], B=dy[batch, out_features]. + The reduced bias gradient is expected to equal dy.sum(dim=0). + + Regression guard for the wrapper bug where the epilogue was hardcoded to + DEFAULT, which silently zeroed the returned bias gradient. + """ + torch.manual_seed(42) + A = torch.randn((batch, in_features), dtype=dtype, device='cuda') * 0.5 + B = torch.randn((batch, out_features), dtype=dtype, device='cuda') * 0.5 + bias = torch.zeros((out_features,), dtype=dtype, device='cuda') + + _, triton_bias_grad = call_gemm_with_bias(A, B, "NT", dtype, bias, grad=True, use_triton=True) + _, cpp_bias_grad = call_gemm_with_bias(A, B, "NT", dtype, bias, grad=True, use_triton=False) + + assert triton_bias_grad is not None, "Triton did not return a bias gradient tensor." + assert cpp_bias_grad is not None, "C++ did not return a bias gradient tensor." + # A correct BGRADB must not produce an all-zero gradient for non-trivial B. + assert triton_bias_grad.abs().sum().item() > 0, ( + "Triton bias gradient is all zeros — BGRADB epilogue appears inactive." + ) + + # Cross-check against the analytical reduction. + expected = B.float().sum(dim=0) + torch.testing.assert_close( + triton_bias_grad.float(), expected, + atol=5e-2, rtol=1e-2, + ) + torch.testing.assert_close( + triton_bias_grad.float(), cpp_bias_grad.float(), + atol=5e-3, rtol=1e-2, + ) + + +# Batched (multi-dim) FP8 coverage. The 2D fp8 case is already covered by +# test_triton_vs_pytorch_fp8 above; this exercises the backend's +# flatten-leading-dims semantics for tensors with ndim > 2. +@pytest.mark.parametrize("batch_size, M, K, N", [ + (2, 128, 256, 512), + (4, 64, 128, 256), +]) +def test_triton_vs_pytorch_fp8_multidim(batch_size, M, K, N): + os.environ['NVTE_USE_GEMM_TRITON'] = '1' + torch.manual_seed(42) + # TN layout: A=[batch, M, K], B=[batch, N, K]. Leading dims flatten. + A_f32 = torch.randn(batch_size, M, K, dtype=torch.float32, device='cuda') * 0.5 + B_f32 = torch.randn(batch_size, N, K, dtype=torch.float32, device='cuda') * 0.5 + + A_fp8 = Float8Quantizer( + scale=torch.full([1], 1.0, dtype=torch.float32, device='cuda'), + amax=torch.empty([1], dtype=torch.float32, device='cuda'), + fp8_dtype=tex.DType.kFloat8E4M3, + )(A_f32) + B_fp8 = Float8Quantizer( + scale=torch.full([1], 1.0, dtype=torch.float32, device='cuda'), + amax=torch.empty([1], dtype=torch.float32, device='cuda'), + fp8_dtype=tex.DType.kFloat8E4M3, + )(B_f32) + + output = call_gemm(A_fp8, B_fp8, layout="TN", out_dtype=torch.float32) + + # Reference: flatten leading dims, then B @ A.T (TN semantics). + A_flat = A_fp8.dequantize().reshape(-1, K) # [batch*M, K] + B_flat = B_fp8.dequantize().reshape(-1, K) # [batch*N, K] + expected = torch.matmul(B_flat, A_flat.T).reshape(batch_size, N, batch_size * M) + + torch.testing.assert_close( + output.to(torch.float32), expected.to(torch.float32), + atol=5e-3, rtol=1e-2, + ) + + +if __name__ == "__main__": + # Quick smoke test + os.environ['NVTE_USE_GEMM_TRITON'] = '1' + test_triton_vs_pytorch_regular(128, 256, 512, "TN", torch.float16) + test_triton_vs_pytorch_fp8(128, 256, 512, "TN", + (tex.DType.kFloat8E4M3, tex.DType.kFloat8E4M3)) + test_triton_vs_cpp_regular(128, 256, 512, "TN", torch.float16) + print("All smoke tests passed!") diff --git a/tests/pytorch/triton_kernels/test_gemm_kernel.py b/tests/pytorch/triton_kernels/test_gemm_kernel.py new file mode 100644 index 000000000..1829b39f7 --- /dev/null +++ b/tests/pytorch/triton_kernels/test_gemm_kernel.py @@ -0,0 +1,272 @@ +# Copyright (c) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. +# +# License for AMD contributions = MIT. See LICENSE for more information + +"""Low-level ``te_gemm_triton()`` kernel-direct correctness tests. + +Bypasses ``general_gemm`` and the tensor wrappers. Drives +``te_gemm_triton()`` with raw tensors and a Triton reference implementation +across a broad matrix of ``(M, K, N) x layout x in_dtype x out_dtype x bias +x grad`` parametrizations (~1000 collected). Isolates kernel-level bugs from +wrapper-layer issues. + +Complementary file: + +- ``test_gemm.py`` -- user-facing ``general_gemm()`` tests with real + ``Float8Tensor`` / ``MXFP8Tensor`` and equivalence vs. both PyTorch and + the C++ backend (covers MXFP8 and mxfp8 wrapper paths there too). +""" + +import pytest +import torch +import triton +import triton.language as tl + +from transformer_engine.pytorch.triton_kernels.common import ( + torch_dtype_to_te_dtype as torch_to_te_dtype, +) +from transformer_engine.pytorch.triton_kernels.gemm import te_gemm_triton +from test_common import str_to_torch_dtype + + +# String -> Triton kernel dtype. Kept local because `gen_input` needs the +# Triton dtype for the reference `copy_kernel` and to decide the fp8 +# codepath. String -> torch.dtype comes from `str_to_torch_dtype` (shared). +name_to_tl_types = { + 'int8': tl.int8, + 'int32': tl.int32, + 'fp16': tl.float16, + 'fp32': tl.float32, + 'bf16': tl.bfloat16, + 'fp8e4': tl.float8e4b8, + 'fp8e5': tl.float8e5b16, +} + + +def gen_input(M, N, ty_name, needTrans, seed, device='cuda'): + d_type = name_to_tl_types[ty_name] + torch.manual_seed(seed) + torch.cuda.manual_seed(seed) + + @triton.jit + def copy_kernel(input_ptr, output_ptr, n_elements, BLOCK_SIZE: tl.constexpr): + offsets = tl.program_id(axis=0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = offsets < n_elements + input = tl.load(input_ptr + offsets, mask=mask) + output = input + tl.store(output_ptr + offsets, output, mask=mask) + + if needTrans: + raw_data = torch.randn((N, M), dtype=torch.float32, device='cuda').T + else: + raw_data = torch.randn((M, N), dtype=torch.float32, device='cuda') + # avoid type conversion rounding errors of subnormal values + raw_data += 0.1 + if d_type == tl.float8e4b8: + raw_data += torch.sign(raw_data) + + if ty_name in ('fp16', 'bf16', 'fp32', 'fp8e4', 'fp8e5'): + input = raw_data.to(str_to_torch_dtype(ty_name)) + input_f16 = input.to(torch.float16) + else: + f8_tensor = raw_data.to(torch.int8) + # keep only two bits of exponent to avoid overflow + f8_tensor = f8_tensor & 0b00111111 + input = triton.reinterpret(f8_tensor, d_type) + input_f16 = torch.empty_like(f8_tensor, dtype=torch.float16) + grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']),) + n_elements = raw_data.numel() + copy_kernel[grid](input, input_f16, n_elements, BLOCK_SIZE=1024) + + return input, input_f16 + +def get_in_dtypes(in_type): + types = in_type.split('-') + if len(types) == 2: + return types[0], types[1] + else: + return types[0], types[0] + +def is_fp8(type_name): + a_type, b_type = get_in_dtypes(type_name) + return ( a_type in ('fp8e4', 'fp8e5') ) and ( b_type in ('fp8e4', 'fp8e5') ) + +def is_mixed_fp8(type_name): + a_type, b_type = get_in_dtypes(type_name) + return is_fp8(type_name) and a_type != b_type + +# Mixed FP8 (e4m3 + e5m2) on gfx950 hits a Triton compiler bug in mixed-type +# MFMA instruction selection. Fixed in triton-lang/triton PR #9567 +# (commit eaaa75cf5, 2026-02-27). The fix lives only on Triton main — it is +# NOT on release/3.6.x, release/3.7.x, or release/3.8.x, so no released +# PyTorch through 2.13 carries it. First lands in pytorch/pytorch main via +# the Triton pin bump to 43422b04 ("Triton 3.8") on 2026-06-26, which ships +# in PyTorch 2.14.0.dev nightlies from that date onward. +# See https://github.com/triton-lang/triton/pull/9567. +from transformer_engine.pytorch import torch_version +_MIXED_FP8_MFMA_FIXED = torch_version() >= (2, 14) + +@pytest.mark.parametrize("M, K, N, in_dtype, out_dtype, col_a, col_b, use_bias, bias_dtype, grad", +[ (*shape, in_dtype, out_dtype, col_a, col_b, use_bias, bias_dtype, grad) + for shape in [(2304, 768, 4096), + (768, 768, 4096), + (768, 3072, 4096), + (229, 541, 541), + (71, 71, 3571), + (29, 29, 17389) + ] + for in_dtype, out_dtype in [('fp16', 'fp16'), + ('bf16', 'bf16'), + ('fp16', 'fp32'), + ('fp32', 'fp32'), + ('fp8e4', 'fp32'), + ('fp8e4', 'bf16'), + ('fp8e4', 'fp16'), + #('fp8e4', 'fp8e4'), + # TODO: d_amax compute seems to have some accuracy issues + ('fp8e5-fp8e4', 'fp32'), + ('fp8e4-fp8e5', 'fp32'), + ('fp8e5-fp8e4', 'bf16'), + ('fp8e4-fp8e5', 'bf16'), + ] + for col_a in [False, True] + for col_b in [False, True] + for use_bias in [True, False] + for bias_dtype in ['bf16'] + for grad in [True, False] + ] +) +def test_correctness(M, N, K, col_a, col_b, in_dtype, out_dtype, use_bias, bias_dtype, grad): + a_in_dtype, b_in_dtype = get_in_dtypes(in_dtype) + if is_fp8(in_dtype) and use_bias and grad: + pytest.skip('Skip tests for fp8 GEMM with BGRADB.') + + if is_mixed_fp8(in_dtype) and not _MIXED_FP8_MFMA_FIXED: + pytest.skip( + 'Mixed FP8 formats (e4m3 + e5m2) require Triton with ' + 'triton-lang/triton#9567 (only on Triton main; first ships in ' + 'PyTorch 2.14.0.dev nightlies from 2026-06-26+).' + ) + + if col_a and col_b: + pytest.skip('Skip tests for TT layout') + empty_tensor = torch.Tensor() + + a, a_fp16 = gen_input(K, M, a_in_dtype, col_a, 1, device='cuda') + b, b_fp16 = gen_input(N, K, b_in_dtype, col_b, 2, device='cuda') + a_fp32 = a.to(torch.float32) + b_fp32 = b.to(torch.float32) + # Allocates output. + torch_out_dtype = str_to_torch_dtype(out_dtype) + torch_bias_dtype = str_to_torch_dtype(bias_dtype) + c = torch.empty((N, M), device=a.device, dtype=torch_out_dtype) + if is_fp8(in_dtype): + A_scale_inverse = torch.randn((3,), dtype=torch.float32, device='cuda') + B_scale_inverse = torch.randn((3,), dtype=torch.float32, device='cuda') + else: + A_scale_inverse = empty_tensor + B_scale_inverse = empty_tensor + + D_scale = torch.empty((), dtype=torch.float32, device='cuda') + transa = col_a + transb = col_b + + A_type = torch_to_te_dtype(str_to_torch_dtype(a_in_dtype)) + B_type = torch_to_te_dtype(str_to_torch_dtype(b_in_dtype)) + D_type = torch_to_te_dtype(str_to_torch_dtype(out_dtype)) + if out_dtype in ('fp8e4', 'fp8e5'): + D_amax = torch.empty((), dtype=torch.float32, device='cuda') + else: + D_amax = empty_tensor + A_fp8_tensor = 0 + B_fp8_tensor = 0 + if use_bias: + if grad: + bias = torch.empty((N,), dtype=torch_bias_dtype, device='cuda') + else: + bias = torch.randn((M,), dtype=torch_bias_dtype, device='cuda') + else: + bias = empty_tensor + bias_type = torch_to_te_dtype(torch_bias_dtype) + pre_gelu_out = empty_tensor + workspace = empty_tensor + workspaceSize = 0 + accumulate = False + use_split_accumulator = False + + output_fp8 = is_fp8(out_dtype) + + ## b is (N, K) in row major and a is (K, M) in row major, + ## when doing GEMM in BLAS + ## a and b are swapped, so gemm_a is (M, K) in column major + ## and gemm_b is (K, N) is column major + torch_output = torch.matmul(b_fp32, a_fp32) + + if is_fp8(in_dtype): + # For f8 and inputs, multiplied by the scales + torch_output *= A_scale_inverse[A_fp8_tensor] * B_scale_inverse[B_fp8_tensor] + + if use_bias: + if grad: + torch_bias_gradient = b.sum(axis=1).to(torch_bias_dtype) + else: + torch_output += bias + + if output_fp8: + torch_output_amax = torch.max(torch.abs(torch_output)) + fp8_amax = 240.0 if out_dtype == 'fp8e4' else 57344.0 + D_scale = fp8_amax * 0.5 / torch_output_amax + torch_output *= D_scale + + torch_output = torch_output.to(torch_out_dtype) + + # Shape is different based on trans value + # for example, if transa == True, a is (K, M), + # we want the shape to be fed to te_gemm_triton + # to be (M, K) as te_gemm_triton will apply the + # the transpose. + a_col_major = a.T if transa else a + b_col_major = b.T if transb else b + te_gemm_triton(a_col_major, + A_scale_inverse, + A_fp8_tensor, + A_type, + transa, + b_col_major, + B_scale_inverse, + B_fp8_tensor, + B_type, + transb, + c, + D_scale, + D_type, + D_amax, + bias, + bias_type, + pre_gelu_out, + grad, + workspace, + workspaceSize, + accumulate, + use_split_accumulator) + + atol = 5e-3 + if torch_out_dtype == 'bf16' and use_bias: + atol = 8e-3 + rtol = 0 if torch.version.hip is None else 1e-2 + if output_fp8: + def check_if_adjacent_fp8(a, b, out_dtype): + m = 3 if out_dtype == 'fp8e4' else 2 + a = a.to(torch.float32) + b = b.to(torch.float32) + check = torch.abs(a-b) <= torch.exp(torch.floor(torch.log(torch.abs(a)))) + return torch.all(check) + assert check_if_adjacent_fp8(c, torch_output, out_dtype) + else: + torch.testing.assert_close(c.to(torch.float32), torch_output.to(torch.float32), atol=atol, rtol=rtol) + + if use_bias and grad: + torch.testing.assert_close(bias, torch_bias_gradient, atol=5e-3, rtol=rtol) + + if output_fp8: + torch.testing.assert_close(D_amax, torch_output_amax, atol=5e-3, rtol=rtol) diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index 12c3811c2..0a10bc653 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -512,7 +512,19 @@ def general_gemm( "beta": beta, } - out, bias_grad, gelu_input, extra_output = tex.generic_gemm(*args, **kwargs) + # ROCm-only backend: the Triton kernels use gfx942/gfx950-specific MFMA + # instructions and autotune configs, so refuse to enable on non-HIP builds. + use_gemm_triton = IS_HIP_EXTENSION and bool(int(os.environ.get("NVTE_USE_GEMM_TRITON", "0"))) + if use_gemm_triton: + # Lazy: only pull in Triton when the backend is opted into. Keeps + # `triton` off the module-import path when NVTE_USE_GEMM_TRITON is + # unset (the default), so stacks without pytorch-triton-rocm can + # still use the C++ hipBLASLt path. + from ..triton_kernels.gemm import te_generic_gemm_triton + out, bias_grad, gelu_input, extra_output = te_generic_gemm_triton(*args, **kwargs) + else: + out, bias_grad, gelu_input, extra_output = tex.generic_gemm(*args, **kwargs) + if IS_HIP_EXTENSION and use_bf16_tn_output_workaround: out = cast_if_needed(out, torch.float32) diff --git a/transformer_engine/pytorch/quantization.py b/transformer_engine/pytorch/quantization.py index a1436ff75..c6f520d6d 100644 --- a/transformer_engine/pytorch/quantization.py +++ b/transformer_engine/pytorch/quantization.py @@ -181,6 +181,22 @@ def check_recipe_support(recipe: Recipe) -> None: if not recipe_supported: raise RuntimeError(unsupported_reason) + # The Triton GEMM backend does not support mixed FP8 types due to a Triton + # compiler bug (triton-lang/triton#9567, not yet in pytorch-triton-rocm as of + # PyTorch 2.11). The HYBRID recipe uses e4m3 for forward and e5m2 for backward, + # producing mixed-type GEMMs during the backward pass. Only Format.E4M3 (which + # uses e4m3 for both forward and backward) is compatible with the Triton backend. + use_gemm_triton = bool(int(os.environ.get("NVTE_USE_GEMM_TRITON", "0"))) + if use_gemm_triton and recipe is not None and hasattr(recipe, "fp8_format"): + if recipe.fp8_format == Format.HYBRID: + raise ValueError( + "The Triton GEMM backend (NVTE_USE_GEMM_TRITON=1) does not support " + "Format.HYBRID because the backward pass produces mixed FP8 type GEMMs " + "(e5m2 x e4m3), which trigger a Triton compiler bug " + "(triton-lang/triton#9567). Use Format.E4M3 instead, or disable the " + "Triton backend (unset NVTE_USE_GEMM_TRITON)." + ) + def get_default_fp8_recipe() -> Recipe: """FP8 recipe with default args.""" diff --git a/transformer_engine/pytorch/triton_kernels/gemm/__init__.py b/transformer_engine/pytorch/triton_kernels/gemm/__init__.py new file mode 100644 index 000000000..07149c4a5 --- /dev/null +++ b/transformer_engine/pytorch/triton_kernels/gemm/__init__.py @@ -0,0 +1,29 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. + +"""Triton GEMM kernels (dense, non-grouped) for BF16/FP16/FP32/FP8/MXFP8.""" + +from .gemm_wrapper import te_gemm_triton, te_generic_gemm_triton, matmul, mxfp8_matmul +from .gemm_common import ( + is_fp8_dtype, + reinterpret_as_fp8_tensor, + getGemmOutputShape, + product, +) + +# Dtype conversions (torch <-> tex.DType, architecture-native FP8 dtypes) are +# NOT re-exported here -- import them directly from +# ``transformer_engine.pytorch.triton_kernels.common``, which is the +# authoritative source shared across all Triton kernel backends. + +__all__ = [ + "te_gemm_triton", + "te_generic_gemm_triton", + "matmul", + "mxfp8_matmul", + "is_fp8_dtype", + "reinterpret_as_fp8_tensor", + "getGemmOutputShape", + "product", +] diff --git a/transformer_engine/pytorch/triton_kernels/gemm/gemm_common.py b/transformer_engine/pytorch/triton_kernels/gemm/gemm_common.py new file mode 100644 index 000000000..72a7514c9 --- /dev/null +++ b/transformer_engine/pytorch/triton_kernels/gemm/gemm_common.py @@ -0,0 +1,172 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. + +"""Shared helpers for the Triton GEMM backend. + +Contains dtype conversion utilities, output-shape computation, and small +helpers used by the Python-side Triton GEMM path to reconstruct rowwise +data from a columnwise-only ``Float8TensorStorage`` and to pick the +correct pre-quantized copy from an ``MXFP8TensorStorage``. +""" + +import torch + +import transformer_engine_torch as tex + +# Reuse the shared architecture-native FP8 dtype helpers from the +# triton_kernels package to stay in sync with the norms / cast kernels. +from ..common import get_torch_e4m3_type, get_torch_e5m2_type + + +def is_fp8_dtype(dtype: tex.DType) -> bool: + """Whether a TE ``DType`` is one of the FP8 variants.""" + return dtype in (tex.DType.kFloat8E4M3, tex.DType.kFloat8E5M2) + + +def reinterpret_as_fp8_tensor(a: torch.Tensor, dtype: tex.DType) -> torch.Tensor: + """View a uint8 tensor as the architecture-native FP8 torch dtype. + + gfx942 (MI300/MI325) uses NANOO (``fnuz``) FP8 variants; gfx950 (MI350) + uses OCP-standard variants. Delegates dtype selection to + ``triton_kernels.common.get_torch_e4m3_type`` / ``_e5m2_type``. + """ + if dtype == tex.DType.kFloat8E4M3: + return a.view(dtype=get_torch_e4m3_type()) + if dtype == tex.DType.kFloat8E5M2: + return a.view(dtype=get_torch_e5m2_type()) + +def getGemmOutputShape(A, transa, B, transb): + """ + Compute output shape for GEMM following the C++ backend logic. + + Matches getGemmOutputShape in transformer_engine/pytorch/csrc/extensions/gemm.cpp + + Why Does This Preserve B's Batch Dimensions? + ============================================= + + This is a deliberate API design choice that makes the interface consistent and + predictable for neural network operations. + + Usage Patterns in Linear Layer: + + 1. Forward Pass (fprop) - Layout: TN (default) + output = general_gemm(weight, input) + - A = weight: [out_features, in_features] - no batch dims + - B = input: [batch, seq_len, in_features] - HAS batch dims + - Output: [batch, seq_len, out_features] - preserves B's batch + + 2. Input Gradient (dgrad) - Layout: NN + grad_input = general_gemm(weight, grad_output) + - A = weight: [out_features, in_features] - no batch dims + - B = grad_output: [batch, seq_len, out_features] - HAS batch dims + - Output: [batch, seq_len, in_features] - preserves B's batch + + 3. Weight Gradient (wgrad) - Layout: NT + grad_weight = general_gemm(input, grad_output) + - A = input: [batch, seq_len, in_features] - batch dims + - B = grad_output: [batch, seq_len, out_features] - batch dims + - Output: [out_features, in_features] - NO batch (transb=True) + + Key Insight: + The calling code consistently places the tensor with desired output batch + structure as the B operand. + - For fprop/dgrad: B has batch dimensions -> output keeps them + - For wgrad: Both have batch, use transb=True -> output flattens them (reduction over batch) + + Why This Convention? + - Consistency: Always put batched activations as B + - Predictability: Output shape always relates to B's structure + - Simplicity: Caller controls output shape by choosing B and transb + - Efficiency: Avoids extra reshapes in common cases + + Could It Be Different? + Yes! The API could preserve A's batch instead, but then all calling code would + need to swap operands. The math would work the same, just with reversed convention. + """ + # Handle both tensors and torch.Size objects + A_shape = A if isinstance(A, torch.Size) else A.shape + B_shape = B if isinstance(B, torch.Size) else B.shape + + # Calculate flattened dimensions (product of all leading dims) + A0 = product(A_shape[:-1]) # Product of all leading dims + A1 = A_shape[-1] + B0 = product(B_shape[:-1]) + B1 = B_shape[-1] + + # Construct output shape following C++ logic: + # if (transb) { ret = [B1] } + # else { ret = [B_shape[0], B_shape[1], ..., B_shape[-2]] } // Unflatten B0 + # if (transa) { ret.append(A0) } + # else { ret.append(A1) } + + ret = [] + + # First part: from B + if transb: + ret.append(B1) + else: + # Preserve B's batch structure (all dims except last) + for i in range(len(B_shape) - 1): + ret.append(B_shape[i]) + + # Second part: from A + if transa: + ret.append(A0) # Flattened A + else: + ret.append(A1) # A's last dim + + return torch.Size(ret) + +def product(shape): + ret = 1 + for i in shape: + ret *= i + return ret + + +def materialize_rowwise_from_columnwise(storage) -> torch.Tensor: + """Reconstruct the rowwise Float8 data buffer from a columnwise-only storage. + + Inverts ``tex.fp8_transpose``'s layout: fp8_transpose treats an n-D rowwise + tensor with shape (D0, ..., D_{n-2}, K) as 2-D (M, K) with + M = prod(D0..D_{n-2}), transposes to (K, M), and re-shapes the result to + [K, D0, ..., D_{n-2}]. To recover the original rowwise layout we rotate + the leading K dim back to the tail: [K, D0, ..., D_{n-2}] -> [D0, ..., D_{n-2}, K]. + + Args: + storage: A ``Float8TensorStorage`` whose ``_transpose`` (columnwise + data) is present and valid. + + Returns: + A contiguous rowwise buffer with the shape the tensor would have had + before it was columnwise-only. + """ + cw = storage._transpose + ndim = cw.dim() + if ndim == 2: + return cw.transpose(0, 1).contiguous() + perm = list(range(1, ndim)) + [0] + return cw.permute(*perm).contiguous() + + +def data_and_scale_for_transpose(storage, will_transpose: bool): + """Pick the pre-quantized MXFP8 copy that matches a BLAS transpose flag. + + MXFP8 cannot be safely re-transposed after quantization (the block-scale + layout is orientation-locked), so the storage keeps both rowwise and + columnwise pre-quantized copies. This picks the right one to hand to the + kernel. + + Args: + storage: An ``MXFP8TensorStorage`` with both rowwise and columnwise + data + scale populated. + will_transpose: If True, the kernel will transpose this operand -> + return the columnwise copy; else return the rowwise copy. + + Returns: + ``(data, scale_inv)`` -- both ``torch.Tensor``. + """ + if will_transpose: + return storage._columnwise_data, storage._columnwise_scale_inv + return storage._rowwise_data, storage._rowwise_scale_inv diff --git a/transformer_engine/pytorch/triton_kernels/gemm/gemm_kernels.py b/transformer_engine/pytorch/triton_kernels/gemm/gemm_kernels.py new file mode 100644 index 000000000..d6739d6b5 --- /dev/null +++ b/transformer_engine/pytorch/triton_kernels/gemm/gemm_kernels.py @@ -0,0 +1,367 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. + +"""Triton kernel definitions for dense GEMM. + +Contains the raw ``@triton.jit`` kernels used by ``gemm_wrapper.py``: + +- ``mxfp8_matmul_kernel``: block-scaled FP8 matmul using ``tl.dot_scaled`` +- ``matmul_kernel``: FP32/FP16/BF16/FP8 matmul with optional bias/BGRADB + epilogue and fused alpha/beta accumulation +""" + +import triton +import triton.language as tl + + +# MXFP8 (Microscaling FP8) Matmul Kernel and Wrapper +# Uses Triton's tl.dot_scaled() for native block-scaled FP8 matmul + +@triton.autotune( + configs=[ + # Simpler configs for MXFP8 - BLOCK_K must be multiple of 32 (VEC_SIZE) + triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 256, 'BLOCK_SIZE_K': 128, 'GROUP_SIZE_M': 4}), + triton.Config({'BLOCK_SIZE_M': 256, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 128, 'GROUP_SIZE_M': 4}), + triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 64, 'GROUP_SIZE_M': 4}), + triton.Config({'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 256, 'BLOCK_SIZE_K': 64, 'GROUP_SIZE_M': 4}), + ], + key=['M', 'N', 'K'], +) +@triton.heuristics({ + 'EVEN_K': lambda args: args['K'] % args['BLOCK_SIZE_K'] == 0, +}) +@triton.jit +def mxfp8_matmul_kernel( + # Data pointers + a_ptr, b_ptr, c_ptr, + # Scale pointers (E8M0 format, uint8) + a_scale_ptr, b_scale_ptr, + # Matrix dimensions + M, N, K, + # Data strides + stride_am, stride_ak, + stride_bk, stride_bn, + stride_cm, stride_cn, + # Scale strides + stride_a_scale_m, stride_a_scale_k, + stride_b_scale_n, stride_b_scale_k, + # Meta-parameters + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + GROUP_SIZE_M: tl.constexpr, + EVEN_K: tl.constexpr, + VEC_SIZE: tl.constexpr, # MXFP8_BLOCK_SCALING_SIZE (always 32) + FP8_FORMAT_A: tl.constexpr, # "e4m3" or "e5m2" + FP8_FORMAT_B: tl.constexpr, # "e4m3" or "e5m2" +): + """ + MXFP8 matmul kernel using tl.dot_scaled() for block-scaled FP8 computation. + + Scales are stored in E8M0 format (uint8 biased exponents) and converted to FP32. + """ + + # Program ID + pid = tl.program_id(axis=0) + num_pid_m = tl.cdiv(M, BLOCK_SIZE_M) + num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) + + # Swizzled block mapping for better L2 cache utilization + num_pid_in_group = GROUP_SIZE_M * num_pid_n + group_id = pid // num_pid_in_group + first_pid_m = group_id * GROUP_SIZE_M + group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) + pid_m = first_pid_m + (pid % group_size_m) + pid_n = (pid % num_pid_in_group) // group_size_m + + # Initialize accumulator + accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) + + # Compute block offsets + offs_am = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_bn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + offs_k = tl.arange(0, BLOCK_SIZE_K) + + # int64 promotion (see matmul_kernel A/B). + a_ptrs = a_ptr + (offs_am[:, None].to(tl.int64) * stride_am + offs_k[None, :] * stride_ak) + b_ptrs = b_ptr + (offs_k[:, None] * stride_bk + offs_bn[None, :].to(tl.int64) * stride_bn) + + # K-loop + num_k_blocks = tl.cdiv(K, BLOCK_SIZE_K) + for k in range(num_k_blocks): + # Load FP8 data + if EVEN_K: + a = tl.load(a_ptrs) + b = tl.load(b_ptrs) + else: + k_remaining = K - k * BLOCK_SIZE_K + mask_k = offs_k < k_remaining + a = tl.load(a_ptrs, mask=mask_k[None, :], other=0.0) + b = tl.load(b_ptrs, mask=mask_k[:, None], other=0.0) + + # Load E8M0 scales for this K-block + # We have: + # - A columnwise scales: [M//VEC_SIZE, K] - one scale per 32 rows for each column + # - B rowwise scales: [K, N//VEC_SIZE] - one scale per 32 columns for each row + # + # tl.dot_scaled expects: + # - A scales: [BLOCK_SIZE_M, BLOCK_SIZE_K // VEC_SIZE] - one scale per 32 elements along K + # - B scales: [BLOCK_SIZE_K // VEC_SIZE, BLOCK_SIZE_N] - one scale per 32 elements along K + # + # For columnwise A: each column has M//32 scales + # We need to gather the right scales for our block + + # A scales: we need scales for rows [pid_m*BLOCK_SIZE_M : (pid_m+1)*BLOCK_SIZE_M] + # and columns [k*BLOCK_SIZE_K : (k+1)*BLOCK_SIZE_K] + # From columnwise layout [M//VEC_SIZE, K], we need to: + # - Select row indices: [pid_m*BLOCK_SIZE_M//VEC_SIZE : (pid_m+1)*BLOCK_SIZE_M//VEC_SIZE] + # - Select column indices: [k*BLOCK_SIZE_K : (k+1)*BLOCK_SIZE_K] + # This gives us [BLOCK_SIZE_M//VEC_SIZE, BLOCK_SIZE_K] scales + # But we need [BLOCK_SIZE_M, BLOCK_SIZE_K//VEC_SIZE] for tl.dot_scaled + + # Actually, let's think about what the scales mean: + # Columnwise: scale[i,j] applies to data[i*32:(i+1)*32, j] + # So for block starting at row pid_m*BLOCK_SIZE_M, col k*BLOCK_SIZE_K: + # We need scales from row indices pid_m*BLOCK_SIZE_M//32 to (pid_m+1)*BLOCK_SIZE_M//32 + + # With reversed selection, A now has rowwise scales [M, K//32] + # For tl.dot_scaled we need [BLOCK_SIZE_M, BLOCK_SIZE_K//32] + k_block_start = k * (BLOCK_SIZE_K // VEC_SIZE) + offs_a_scale_k = k_block_start + tl.arange(0, BLOCK_SIZE_K // VEC_SIZE) + + a_scale_ptrs = a_scale_ptr + (offs_am[:, None].to(tl.int64) * stride_a_scale_m + + offs_a_scale_k[None, :] * stride_a_scale_k) + + # Check bounds for scale loading + # Use other=127 for out-of-bounds E8M0 scales (127 = scale of 1.0) + mask_a_scale_m = offs_am < M + mask_a_scale_k = offs_a_scale_k < (K // VEC_SIZE) + a_scale_mask = mask_a_scale_m[:, None] & mask_a_scale_k[None, :] + a_scale_e8m0 = tl.load(a_scale_ptrs, mask=a_scale_mask, other=127) + + # B scale layout: [N, K//32] (new dot_scaled API -- "Do NOT transpose rhs_scale") + # For tl.dot_scaled we need [BLOCK_SIZE_N, BLOCK_SIZE_K//32] + offs_b_scale_k = k_block_start + tl.arange(0, BLOCK_SIZE_K // VEC_SIZE) + b_scale_ptrs = b_scale_ptr + (offs_bn[:, None].to(tl.int64) * stride_b_scale_n + + offs_b_scale_k[None, :] * stride_b_scale_k) + + mask_b_scale_n = offs_bn < N + mask_b_scale_k = offs_b_scale_k < (K // VEC_SIZE) + b_scale_mask = mask_b_scale_n[:, None] & mask_b_scale_k[None, :] + b_scale_e8m0 = tl.load(b_scale_ptrs, mask=b_scale_mask, other=127) + + # Block-scaled matmul using tl.dot_scaled + # tl.dot_scaled expects E8M0 scales (uint8) and handles conversion internally + # a: [BLOCK_SIZE_M, BLOCK_SIZE_K] FP8 + # a_scale_e8m0: [BLOCK_SIZE_M, BLOCK_SIZE_K // VEC_SIZE] uint8 (E8M0) + # b: [BLOCK_SIZE_K, BLOCK_SIZE_N] FP8 (already in correct layout, no transpose needed) + # b_scale_e8m0: [BLOCK_SIZE_N, BLOCK_SIZE_K // VEC_SIZE] uint8 (E8M0) + # Note: rhs_scale uses [N, K//32] layout (NOT transposed) per new dot_scaled API + accumulator = tl.dot_scaled( + a, # [BLOCK_SIZE_M, BLOCK_SIZE_K] + a_scale_e8m0, # [BLOCK_SIZE_M, BLOCK_SIZE_K // VEC_SIZE] E8M0 + FP8_FORMAT_A, # "e4m3" or "e5m2" + b, # [BLOCK_SIZE_K, BLOCK_SIZE_N] - NO transpose + b_scale_e8m0, # [BLOCK_SIZE_K // VEC_SIZE, BLOCK_SIZE_N] E8M0 + FP8_FORMAT_B, # "e4m3" or "e5m2" + accumulator # [BLOCK_SIZE_M, BLOCK_SIZE_N] + ) + + # Advance data pointers + a_ptrs += BLOCK_SIZE_K * stride_ak + b_ptrs += BLOCK_SIZE_K * stride_bk + + # Store output (convert to target dtype) + offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + # int64 promotion (see matmul_kernel A/B). + c_ptrs = c_ptr + stride_cm * offs_cm[:, None].to(tl.int64) + stride_cn * offs_cn[None, :] + c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N) + + c = accumulator.to(c_ptr.type.element_ty) + tl.store(c_ptrs, c, mask=c_mask) + + +@triton.autotune( + configs=[ + triton.Config({'BLOCK_SIZE_M': 256, 'BLOCK_SIZE_N': 256, 'BLOCK_SIZE_K': 128, 'GROUP_SIZE_M': 4, 'waves_per_eu': 0}, num_warps=8), + triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 256, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 4, 'waves_per_eu': 0}, num_warps=8), + triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 256, 'BLOCK_SIZE_K': 16, 'GROUP_SIZE_M': 4, 'waves_per_eu': 2}, num_warps=4), + triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 1, 'waves_per_eu': 2}, num_warps=8), + triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 64, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 32, 'waves_per_eu': 2}, num_warps=4), + ], + # TODO: do we need to use different data types as key? + key=['M', 'N', 'K'], + # Ran into stream capture error when using cuda_graph, thus disabled. + #use_cuda_graph=True, + # With ACCUMULATE=True each benchmark iteration would add computed_c to + # the output again, multi-copying the result. Snapshot and restore c_ptr + # around every benchmark run. Harmless when ACCUMULATE=False (kernel + # overwrites regardless); cost is paid once per shape during warmup. + restore_value=['c_ptr'], +) +@triton.heuristics({ + 'EVEN_K': lambda args: args['K'] % args['BLOCK_SIZE_K'] == 0, +}) +@triton.jit +def matmul_kernel( + # Pointers to matrices + a_ptr, b_ptr, c_ptr, + # Pointers to scales + a_scale_ptr, b_scale_ptr, c_scale_ptr, + # Pointer to bias + bias_ptr, + # Pointer to amax + c_amax_ptr, + # GEMM output scale (α) and accumulate scale (β): D = α·(A·B) + bias + β·C + alpha, beta, + # Matrix dimensions + M, N, K, + # The stride variables represent how much to increase the ptr by when moving by 1 + # element in a particular dimension. E.g. `stride_am` is how much to increase `a_ptr` + # by to get the element one row down (A has M rows). + stride_am, stride_ak, + stride_bk, stride_bn, + stride_cm, stride_cn, + # Meta-parameters + BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr, BLOCK_SIZE_K: tl.constexpr, + EVEN_K: tl.constexpr, + GROUP_SIZE_M: tl.constexpr, + EPILOGUE: tl.constexpr, + # Whether multiplied by scale_a * scale_b + INPUT_FP8: tl.constexpr, + # Whether to output fp8 or not, if so, also calculate amax. + OUTPUT_FP8: tl.constexpr, + # β=1 accumulation: C := existing_C + α*A*B (used for fused wgrad accumulate) + ACCUMULATE: tl.constexpr, + # Fast-path toggle: skip `accumulator *= alpha` when α is known to be 1.0 + ALPHA_IS_ONE: tl.constexpr, +): + """Kernel for computing the matmul C = A x B. + A has shape (M, K), B has shape (K, N) and C has shape (M, N) + M = blas_n, K = blas_k, N = blas_m + """ + # ----------------------------------------------------------- + # Map program ids `pid` to the block of C it should compute. + # This is done in a grouped ordering to promote L2 data reuse. + # See above `L2 Cache Optimizations` section for details. + pid = tl.program_id(axis=0) + num_pid_m = tl.cdiv(M, BLOCK_SIZE_M) + num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) + if GROUP_SIZE_M == 1: + pid_m = pid // num_pid_n + pid_n = pid % num_pid_n + else: + num_pid_in_group = GROUP_SIZE_M * num_pid_n + group_id = pid // num_pid_in_group + first_pid_m = group_id * GROUP_SIZE_M + group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) + pid_m = first_pid_m + (pid % group_size_m) + pid_n = (pid % num_pid_in_group) // group_size_m + + # ---------------------------------------------------------- + # Create pointers for the first blocks of A and B. + # We will advance this pointer as we move in the K direction + # and accumulate + # `a_ptrs` is a block of [BLOCK_SIZE_M, BLOCK_SIZE_K] pointers + # `b_ptrs` is a block of [BLOCK_SIZE_K, BLOCK_SIZE_N] pointers + # See above `Pointer Arithmetics` section for details + offs_k = tl.arange(0, BLOCK_SIZE_K) + offs_am = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M + offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N + # Promote the M/N offsets to int64 before multiplying by a stride that may + # be as large as K or N: offs_am*stride_am overflows int32 whenever + # (M-1)*K > 2^31 (e.g. M=143616, K=15360 → 2.2G). The mask arithmetic + # (offs_* < M/N) stays int32 for speed; only the address computation + # widens. + a_ptrs = a_ptr + (offs_am[:, None].to(tl.int64) * stride_am + offs_k[None, :] * stride_ak) + b_ptrs = b_ptr + (offs_k[:, None] * stride_bk + offs_bn[None, :].to(tl.int64) * stride_bn) + + # ----------------------------------------------------------- + # Iterate to compute a block of the C matrix. + # We accumulate into a `[BLOCK_SIZE_M, BLOCK_SIZE_N]` block + # of fp32 values for higher accuracy. + # `accumulator` will be converted back to fp16 after the loop. + acc_dtype = tl.float32 if c_ptr.type.element_ty != tl.int8 else tl.int32 + accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=acc_dtype) + + if INPUT_FP8: + a_scale = tl.load(a_scale_ptr) + b_scale = tl.load(b_scale_ptr) + scale = a_scale * b_scale + + if OUTPUT_FP8: + c_scale = tl.load(c_scale_ptr) + + if EPILOGUE == 'BGRADB' and not INPUT_FP8: + bias_gradient = tl.zeros((BLOCK_SIZE_M,), dtype=tl.float32) + + for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)): + # Load the next block of A and B, generate a mask by checking the K dimension. + # If it is out of bounds, set it to 0. + if EVEN_K: + a = tl.load(a_ptrs) + b = tl.load(b_ptrs) + else: + a = tl.load(a_ptrs, mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, other=0.0) + b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_SIZE_K, other=0.0) + # We accumulate along the K dimension. + accumulator += tl.dot(a, b) + + if EPILOGUE == 'BGRADB' and not INPUT_FP8: + if pid_n == 0: + ## It is necessary to upcast to fp32 for reduction to ensure accuracy. + bias_gradient_partial = tl.sum(a.to(tl.float32), axis=1) + bias_gradient += bias_gradient_partial + + # Advance the ptrs to the next K block. + a_ptrs += BLOCK_SIZE_K * stride_ak + b_ptrs += BLOCK_SIZE_K * stride_bk + + + if EPILOGUE == 'BGRADB' and not INPUT_FP8: + if pid_n == 0: + offs_bias_gradient = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + bias_gradient_ptrs = bias_ptr + offs_bias_gradient + ## Though bias_gradient is fp32, type conversion will occur before store + tl.store(bias_gradient_ptrs, bias_gradient, mask=(offs_bias_gradient= 2.10 which ships a Triton version + # with the tl.dot_scaled() RHS scale layout bug fix. + # Earlier versions silently produce wrong results for non-uniform B scales. + from transformer_engine.pytorch import torch_version + if torch_version() < (2, 10): + raise RuntimeError( + f"Triton MXFP8 GEMM requires PyTorch >= 2.10 (found {torch_version()}). " + "Earlier versions contain a Triton compiler bug in tl.dot_scaled() that " + "produces incorrect results for the RHS scale operand. " + "Set NVTE_USE_GEMM_TRITON=0 to use the C++ GEMM backend instead." + ) + + # Validate both are MXFP8 + if a_kind != b_kind: + raise ValueError("Mixed MXFP8 and non-MXFP8 inputs not supported") + + # Sanity: both operands must have at least one pre-quantized copy. + if getattr(A, '_rowwise_data', None) is None and getattr(A, '_columnwise_data', None) is None: + raise RuntimeError("MXFP8Tensor has neither rowwise nor columnwise data") + if getattr(B, '_rowwise_data', None) is None and getattr(B, '_columnwise_data', None) is None: + raise RuntimeError("MXFP8Tensor has neither rowwise nor columnwise data") + + # Logical (rowwise-oriented) size for shape computation downstream. + # MXFP8 rowwise and columnwise share the same shape, so we pick whichever + # copy is populated. + A_size = (A._rowwise_data if A._rowwise_data is not None else A._columnwise_data).size() + B_size = (B._rowwise_data if B._rowwise_data is not None else B._columnwise_data).size() + + # IMPORTANT: Match the C++ CanonicalizeGemmInput logic for MXFP8 + # The C++ code selects data/scales based on the BLAS transpose flags: + # + # For A: + # - transa=True: Use rowwise data and scales + # - transa=False: Use columnwise data and scales + # For B: + # - transb=True: Use columnwise data and scales + # - transb=False: Use rowwise data and scales + + # Debug: print available data + import os + if os.getenv("DEBUG_MXFP8_SELECT"): + print(f"[DEBUG] MXFP8 data selection:") + print(f" A shape: {A_size}, transA={transa}") + if getattr(A, '_rowwise_data', None) is not None: + print(f" A rowwise: data {A._rowwise_data.shape}, scale {A._rowwise_scale_inv.shape}") + if getattr(A, '_columnwise_data', None) is not None: + print(f" A columnwise: data {A._columnwise_data.shape}, scale {A._columnwise_scale_inv.shape}") + print(f" B shape: {B_size}, transB={transb}") + if getattr(B, '_rowwise_data', None) is not None: + print(f" B rowwise: data {B._rowwise_data.shape}, scale {B._rowwise_scale_inv.shape}") + if getattr(B, '_columnwise_data', None) is not None: + print(f" B columnwise: data {B._columnwise_data.shape}, scale {B._columnwise_scale_inv.shape}") + + # MXFP8 Selection for BLAS API compatibility + # + # The API uses BLAS convention with column-major interpretation + # We need to select the right MXFP8 format based on BLAS transpose flags + # Following the C++ logic from CanonicalizeGemmInput: + # - When transA=True: use rowwise (will_transpose=False in the helper) + # - When transA=False: use columnwise (will_transpose=True) + # - When transB=True: use columnwise (will_transpose=True) + # - When transB=False: use rowwise (will_transpose=False) + A_data, a_scale_inv = data_and_scale_for_transpose(A, will_transpose=not transa) + B_data, b_scale_inv = data_and_scale_for_transpose(B, will_transpose=transb) + + # Debug output + if os.getenv("DEBUG_MXFP8_SELECT"): + print(f"[DEBUG] MXFP8 selection with logical transpose:") + print(f" transA={transa}, transB={transb}") + print(f" A selected: data {A_data.shape}, scale {a_scale_inv.shape}") + print(f" B selected: data {B_data.shape}, scale {b_scale_inv.shape}") + + a_fp8_dtype = A._fp8_dtype + b_fp8_dtype = B._fp8_dtype + a_nominal_dtype = getattr(A, 'dtype', torch.float32) + + input_mxfp8 = True + else: + # FP8 / regular path. _extract_fp8_operand pulls (data, scale, fp8 dtype, + # size) from a Float8 storage or a plain tensor; the columnwise-only + # Float8 case materializes the rowwise buffer once for reuse. + A_data, a_scale_inv, a_fp8_dtype, A_size = _extract_fp8_operand(A, a_kind) + B_data, b_scale_inv, b_fp8_dtype, B_size = _extract_fp8_operand(B, b_kind) + # A's nominal dtype is used downstream as an output-dtype fallback; + # B's isn't consumed. + a_nominal_dtype = getattr(A, 'dtype', None) if a_kind == "fp8" else A.dtype + input_mxfp8 = False + + # Mixed FP8 types (e.g. A=e4m3, B=e5m2) are not supported due to a Triton + # compiler bug: when the MFMA layout is transposed, operand B is packed using + # A's element type, and the instruction's format encoding doesn't account for + # the operand swap. This produces silently wrong results for all MFMA variants. + # Fixed upstream in triton-lang/triton PR #9567 (commit eaaa75cf5, 2026-02-27). + # Not yet included in any pytorch-triton-rocm release as of PyTorch 2.11. + # Expected in PyTorch 2.12+ once the Triton pin is bumped. + # TODO: Remove this guard once pytorch-triton-rocm includes the fix. + if (a_fp8_dtype is not None and b_fp8_dtype is not None + and a_fp8_dtype != b_fp8_dtype): + raise ValueError( + f"Mixed FP8 types (A={a_fp8_dtype}, B={b_fp8_dtype}) are not supported " + f"in the Triton GEMM backend due to a Triton compiler bug " + f"(triton-lang/triton#9567). Use the same FP8 format for both operands, " + f"or disable the Triton backend (unset NVTE_USE_GEMM_TRITON)." + ) + + # Reinterpret uint8 as native FP8 types for Triton + # The FP8 tensor data is stored as torch.uint8 but Triton needs torch.float8_e4m3fnuz + if a_fp8_dtype is not None: + A_data = reinterpret_as_fp8_tensor(A_data, a_fp8_dtype) + if b_fp8_dtype is not None: + B_data = reinterpret_as_fp8_tensor(B_data, b_fp8_dtype) + + # Compute dimensions from the logical (rowwise) size established above. + # A_size / B_size are picked in the FP8 / MXFP8 / regular branches -- + # they always exist even for a Float8TensorStorage without .shape. + # + # BLAS column-major interpretation: + # PyTorch tensors are row-major in memory, but BLAS interprets them as column-major. + # A PyTorch tensor with shape [X, Y] is seen by BLAS as column-major [Y, X]. + # + # For A with PyTorch shape [A0, A1]: + # - BLAS sees it as column-major with A0 columns and A1 rows + # - If transa=False: we use A as-is in column-major = [A1, A0] in BLAS = shape (A1 rows, A0 cols) + # So M = A1, K = A0 + # - If transa=True: we transpose in column-major = [A0, A1] in BLAS = shape (A0 rows, A1 cols) + # So M = A0, K = A1 + # + # For B with PyTorch shape [B0, B1]: + # - If transb=False: BLAS sees [B1, B0], so K = B1, N = B0... wait that's backwards + # + # Actually, let me use the test as reference. From test_gemm_triton.py line 126-127: + # a is (K, M) in PyTorch → BLAS column-major (M, K) + # b is (N, K) in PyTorch → BLAS column-major (K, N) + # So PyTorch (row_dim, col_dim) → BLAS column-major (col_dim, row_dim) + # + # For matrix A with PyTorch shape (A_shape[0], A_shape[1]): + # - BLAS column-major interpretation: (A_shape[1], A_shape[0]) + # - If transa=False: use as column-major (A_shape[1], A_shape[0]), so m=A_shape[1], k=A_shape[0] + # - If transa=True: transpose in column-major gives (A_shape[0], A_shape[1]), so m=A_shape[0], k=A_shape[1] + # + # Using product notation where A0 = product(A_shape[:-1]), A1 = A_shape[-1]: + A0 = product(A_size[:-1]) # First dim(s) + A1 = product(A_size[-1:]) # Last dim + B0 = product(B_size[:-1]) + B1 = product(B_size[-1:]) + + m = A0 if transa else A1 # Original code + k = A1 if transa else A0 # Original code + n = B1 if transb else B0 # Original code + + assert not (transa and transb), 'TT layout not allowed' + + ## general_gemm() follows BLAS convention: tensors are interpreted as column-major + ## PyTorch tensors are stored row-major in memory, but BLAS APIs treat them as column-major + ## Triton matmul kernel expects row-major layout + ## Convert using the standard trick: swap operands and transpose as needed + ## + ## For column-major interpretation: + ## TN: compute B @ A.T (transpose A, not B) + ## NN: compute B @ A (no transposes) + ## NT: compute B.T @ A (transpose B, not A) + + # For multi-dimensional tensors: flatten leading dims first, then transpose + # This implements "flattened multi-dimensional matmul" semantics + A_flat = A_data.reshape(-1, A_data.shape[-1]) # [prod(batch dims), last_dim] + B_flat = B_data.reshape(-1, B_data.shape[-1]) + + # For MXFP8, also flatten scale tensors to match data flattening + if input_mxfp8: + # MXFP8 scales must be flattened to match data + if a_scale_inv is not None and a_scale_inv.dim() > 2: + a_scale_inv = a_scale_inv.reshape(-1, a_scale_inv.shape[-1]) + if b_scale_inv is not None and b_scale_inv.dim() > 2: + b_scale_inv = b_scale_inv.reshape(-1, b_scale_inv.shape[-1]) + + # Skip padding handling for now - will be fixed in a later update + # from the ROCm Transformer Engine repo + + # Convert from BLAS column-major to Triton row-major by swapping operands + if input_mxfp8: + # MXFP8 follows the same pattern as regular FP8: + # 1. Swap operands (B becomes first, A becomes second) + # 2. Apply logical transpose based on the swapped flags + # 3. For MXFP8, transpose BOTH data and scales + + # First operand for Triton (originally B from BLAS) + # Apply transpose if transb is True + a_row_major = B_flat.T if transb else B_flat + a_scale_triton = b_scale_inv.T if (transb and b_scale_inv is not None) else b_scale_inv + + # Second operand for Triton (originally A from BLAS) + # Apply transpose if transa is True + b_row_major = A_flat.T if transa else A_flat + b_scale_triton = a_scale_inv.T if (transa and a_scale_inv is not None) else a_scale_inv + else: + # For regular FP8 and standard types, apply BLAS column-major to row-major conversion + # This swaps A and B + a_row_major = B_flat.T if transb else B_flat + b_row_major = A_flat.T if transa else A_flat + # Scales are swapped to match operand swap (B→a, A→b in row-major) + a_scale_triton = b_scale_inv + b_scale_triton = a_scale_inv + + # Bias / bias-gradient wiring for the regular/FP8 matmul kernel. + # MXFP8 takes a separate path below and does not yet support bias. + has_bias = bias is not None and bias.numel() > 0 + if has_bias and grad: + epilogue = 'BGRADB' + bias_grad = torch.empty_like(bias) + elif has_bias: + epilogue = 'BIAS' + bias_grad = None + else: + epilogue = 'DEFAULT' + bias_grad = None + + # Compute output shape (BLAS column-major convention; both branches identical) + D_shape = getGemmOutputShape(A_size, transa, B_size, transb) + + if D is None: + # Determine output dtype + if output_dtype is not None: + # Use explicitly provided output dtype (from TE_DType) + out_dtype = te_dtype_to_torch_dtype(output_dtype) + elif a_kind == "mxfp8": + # MXFP8 input: use nominal dtype + out_dtype = a_nominal_dtype + elif a_kind == "fp8": + # Regular FP8 input: use nominal dtype if available + if a_nominal_dtype is None: + raise RuntimeError( + "FP8 input detected (Float8TensorStorage without nominal dtype) but output_dtype " + "parameter is not provided. Please explicitly provide the output_dtype parameter " + "to general_gemm()." + ) + out_dtype = a_nominal_dtype + else: + # Regular input: use A's dtype + out_dtype = A_data.dtype + + D = torch.empty(D_shape, dtype=out_dtype, device=A_data.device) + + d_row_major = D.view(-1, D.shape[-1]) + + # Set FP8 flags + is_fp8_input = (a_kind == "fp8" and b_kind == "fp8") + is_mxfp8_input = (a_kind == "mxfp8" and b_kind == "mxfp8") + input_fp8 = is_fp8_input or is_mxfp8_input + output_fp8 = False # Not supporting FP8 output yet + + # Empty tensors for unused parameters (matching C++ empty tensor pattern) + D_scale = torch.Tensor() + D_amax = torch.Tensor() + # Bias tensor passed to the kernel: + # - BIAS: the actual bias, read and added to output + # - BGRADB: output buffer that receives the bias gradient + # - DEFAULT: empty (kernel ignores it) + if epilogue == 'BIAS': + bias_tensor = bias + elif epilogue == 'BGRADB': + bias_tensor = bias_grad + else: + bias_tensor = torch.Tensor() + + # Dispatch to appropriate kernel based on input type + if input_mxfp8: + # MXFP8 path: compute directly in row-major without BLAS column-major conversion + # + # The BLAS column-major conversion (swapping A and B) doesn't work well for MXFP8 + # because the scales are tied to specific data orientations. + # Instead, we compute the matmul directly based on what the user requested. + # + # User requested (in BLAS column-major): C = op(A) @ op(B) + # We need to figure out what that means in row-major and call the kernel. + # + # The kernel computes: C = A_kernel @ B_kernel in row-major + # + # Mapping: + # - NN layout (transa=False, transb=False): C = A @ B + # Row-major: C[M,N] = A[M,K] @ B[K,N] + # Use: A_data (no transpose), B_data (no transpose) + # + # - TN layout (transa=True, transb=False): C = A^T @ B + # Row-major: C[M,N] = A^T[M,K] @ B[K,N] where A is originally [K,M] + # Use: A_data^T or columnwise, B_data + # + # - NT layout (transa=False, transb=True): C = A @ B^T + # Row-major: C[M,N] = A[M,K] @ B^T[K,N] where B is originally [N,K] + # Use: A_data, B_data^T or columnwise + + # Use output dimensions to get correct M, N + # After operand swap and transpose handling: + # - a_row_major: [M, K] (first operand for Triton) + # - b_row_major: [K, N] (second operand for Triton) + # - d_row_major: [M, N] (output) + actual_m = d_row_major.shape[0] + actual_n = d_row_major.shape[1] + + # Verify operands are compatible for matmul + if a_row_major.shape[1] != b_row_major.shape[0]: + print(f"[ERROR] Dimension mismatch after swap/transpose:") + print(f" a_row_major: {a_row_major.shape}") + print(f" b_row_major: {b_row_major.shape}") + print(f" Cannot multiply: {a_row_major.shape} @ {b_row_major.shape}") + print(f" Original: A{A_size}, B{B_size}, trans={'T' if transa else 'N'}{'T' if transb else 'N'}") + assert False, f"Dimension mismatch: {a_row_major.shape} @ {b_row_major.shape}" + actual_k = a_row_major.shape[1] + + # Debug output + import os + if os.getenv("DEBUG_MXFP8_GEMM"): + print(f"\n[DEBUG] MXFP8 GEMM call:") + print(f" BLAS API: A{A_size}, B{B_size}, trans={'T' if transa else 'N'}{'T' if transb else 'N'}") + + # Identify the operation type based on shapes and transpose flags + op_type = "unknown" + if transa and not transb: + op_type = "fprop (TN)" + elif not transa and not transb: + op_type = "dgrad (NN)" + elif not transa and transb: + op_type = "wgrad (NT)" + + # For wgrad, provide more details about the tensor interpretation + if op_type == "wgrad (NT)": + print(f" Interpreting wgrad tensors:") + print(f" A (grad_output): original shape {A_size}") + print(f" After flatten: {A_flat.shape}") + print(f" B (input): original shape {B_size}") + print(f" After flatten: {B_flat.shape}") + print(f" Expected: dY^T @ X where dY=[batch*seq, out_feat], X=[batch*seq, in_feat]") + + print(f" Operation type: {op_type}") + print(f" Expected output shape: {D_shape}") + print(f" After operand swap and transpose for Triton:") + print(f" First operand: {a_row_major.shape} (from B, transb={transb})") + print(f" Second operand: {b_row_major.shape} (from A, transa={transa})") + print(f" Output: {d_row_major.shape}") + print(f" Actual dimensions: M={actual_m}, N={actual_n}, K={actual_k}") + if a_scale_triton is not None and b_scale_triton is not None: + print(f" Scale shapes: {a_scale_triton.shape}, {b_scale_triton.shape}") + # Check scale compatibility + expected_a_scale = (actual_m, actual_k // 32) + expected_b_scale = (actual_k // 32, actual_n) + if a_scale_triton.shape[:2] != expected_a_scale: + print(f" ⚠ First scale mismatch: {a_scale_triton.shape} vs expected {expected_a_scale}") + if b_scale_triton.shape[:2] != expected_b_scale: + print(f" ⚠ Second scale mismatch: {b_scale_triton.shape} vs expected {expected_b_scale}") + + # Call kernel with correct dimensions + # Use the actual output dimensions and inner dimension + # After swapping and transpose, we have: + # a_row_major: [M, K] (first operand) + # b_row_major: [K, N] (second operand) + # d_row_major: [M, N] (output) + + # Verify dimensions match + assert a_row_major.shape[1] == b_row_major.shape[0], \ + f"Inner dimensions don't match: {a_row_major.shape} @ {b_row_major.shape}" + + mxfp8_matmul( + a_row_major, a_scale_triton, # First operand (from B) + b_row_major, b_scale_triton, # Second operand (from A) + d_row_major, # Output + actual_m, actual_n, actual_k, # Use pre-computed dimensions + b_fp8_dtype, a_fp8_dtype # Swap FP8 formats to match swapped operands + ) + else: + # Call regular FP8 or standard matmul kernel + matmul(a_row_major, b_row_major, d_row_major, a_scale_triton, b_scale_triton, + D_scale, bias_tensor, D_amax, epilogue, input_fp8, output_fp8, + accumulate=accumulate, alpha=alpha, beta=beta) + + # Fused FP8 output quantization is not wired through this wrapper (see the + # `output_fp8 = False` above): the kernel produces D in `output_dtype` + # (typically fp32/bf16) rather than writing FP8 with scale + amax directly. + # When the caller passed a quantizer, apply it here so the returned tensor + # is the Float8Tensor / MXFP8Tensor the caller expects. This matches what + # the hipBLASLt fused path produces bit-for-bit: same fp32 accumulator -> + # same quantizer -> same FP8 payload. Skip DebugQuantizer because + # general_gemm strips it out before dispatching to us. + if quantizer is not None: + D = quantizer(D) + + return D, bias_grad, None, None