diff --git a/tests/pytorch/mxfp4/test_mxfp4_gemm_exact.py b/tests/pytorch/mxfp4/test_mxfp4_gemm_exact.py index 116a566bf..f08427031 100644 --- a/tests/pytorch/mxfp4/test_mxfp4_gemm_exact.py +++ b/tests/pytorch/mxfp4/test_mxfp4_gemm_exact.py @@ -18,9 +18,14 @@ import transformer_engine_torch as tex from transformer_engine.pytorch.tensor.mxfp4_tensor import MXFP4Quantizer from transformer_engine.pytorch.custom_recipes.quantization_mxfp4 import MXFP4QuantizerRef +from transformer_engine.pytorch.utils import get_device_compute_capability recipe_available, reason_for_no_recipe = te.is_mxfp4_available(return_reason=True) +requires_gfx950 = pytest.mark.skipif( + not recipe_available or get_device_compute_capability() != (9, 5), + reason="AITER and hipBLASLt MXFP4 GEMMs require gfx950", +) try: import aiter # noqa: F401 @@ -122,6 +127,7 @@ def check_mxfp4_gemm_versus_reference( # AITER a4w4 backend: shuffled FP4 data + swizzled scales, TN only. Requires the aiter package. @pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +@requires_gfx950 @pytest.mark.skipif(not _aiter_available, reason="aiter package not available") @pytest.mark.parametrize("M, K, N", _SHAPES) @pytest.mark.parametrize("out_dtype", [torch.bfloat16], ids=str) @@ -146,6 +152,7 @@ def test_mxfp4_gemm_aiter_versus_reference( # hipBLASLt F4F4 backend: plain FP4 data, plain or pre-swizzled UE8M0 scales @pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +@requires_gfx950 @pytest.mark.parametrize("M, K, N", _SHAPES) @pytest.mark.parametrize("out_dtype", [torch.bfloat16], ids=str) @pytest.mark.parametrize("layout", ["TN", "NN", "NT", "TT"]) diff --git a/tests/pytorch/mxfp4/test_mxfp4_module_exact.py b/tests/pytorch/mxfp4/test_mxfp4_module_exact.py index d3d19784d..1adca5645 100644 --- a/tests/pytorch/mxfp4/test_mxfp4_module_exact.py +++ b/tests/pytorch/mxfp4/test_mxfp4_module_exact.py @@ -19,6 +19,7 @@ ) from transformer_engine.pytorch.custom_recipes.quantization_mxfp4 import MXFP4QuantizerRef from transformer_engine.pytorch.quantization import FP8GlobalStateManager +from transformer_engine.pytorch.utils import get_device_compute_capability recipe_available, reason_for_no_recipe = te.is_mxfp4_available(return_reason=True) @@ -30,6 +31,16 @@ except ImportError: _aiter_available = False +_is_gfx1250 = recipe_available and get_device_compute_capability() == (12, 5) +requires_mxfp4_gemm_backend = pytest.mark.skipif( + not (_is_gfx1250 or _aiter_available), reason="aiter package not available" +) + + +def skip_if_unsupported_shape(in_features: int, out_features: int): + if _is_gfx1250 and (in_features % 256 or out_features % 256): + pytest.skip("FlyDSL MXFP4 GEMM needs in_features and out_features divisible by 256") + class GetRecipes: @staticmethod @@ -244,7 +255,7 @@ def check_mxfp4_module_versus_reference( @pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) -@pytest.mark.skipif(not _aiter_available, reason="aiter package not available") +@requires_mxfp4_gemm_backend @pytest.mark.parametrize( "in_features, out_features", [ @@ -270,6 +281,7 @@ def test_mxfp4_linear_versus_reference( """Test MXFP4 Linear module against reference implementation.""" if x_dtype != torch.bfloat16: pytest.skip("MXFP4 quantization is only supported for bfloat16 input") + skip_if_unsupported_shape(in_features, out_features) check_mxfp4_module_versus_reference( module_class=te.Linear, @@ -412,12 +424,13 @@ def check_mxfp4_layernorm_linear_versus_reference( @pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) -@pytest.mark.skipif(not _aiter_available, reason="aiter package not available") +@requires_mxfp4_gemm_backend @pytest.mark.parametrize( "in_features, out_features", [ (128, 256), (256, 128), + (512, 512), ], ) @pytest.mark.parametrize("bias", [False], ids=["no_bias"]) @@ -434,6 +447,7 @@ def test_mxfp4_layernorm_linear_versus_reference( num_steps: int, use_hadamard: bool, ): + skip_if_unsupported_shape(in_features, out_features) check_mxfp4_layernorm_linear_versus_reference( in_features=in_features, out_features=out_features, diff --git a/tests/pytorch/test_gemm_backends.py b/tests/pytorch/test_gemm_backends.py index fedfdea54..eaa50c1a2 100644 --- a/tests/pytorch/test_gemm_backends.py +++ b/tests/pytorch/test_gemm_backends.py @@ -14,7 +14,8 @@ - The ``test_triton_*`` family runs when ``pytorch-triton-rocm`` is importable (the Triton GEMM path has no arch gate, so it runs on gfx942 and gfx950). - The ``test_flydsl_*`` family runs when the version-gated ``flydsl`` package is - importable and the device is gfx950 (the only arch the FlyDSL path supports). + importable and the device is gfx950; the ``test_flydsl_*_mxfp4`` tests run on + gfx1250 instead (the only arch the FlyDSL MXFP4 path currently supports). Triton coverage: fp32 / fp16 / bf16 / same-format FP8 / mixed FP8 (skipped for a compiler bug) / MXFP8 across TN / NN / NT, plus bias / bias-grad epilogues and a @@ -32,7 +33,8 @@ config: fp32 / fp16 / bf16, same- and mixed-format tensor-wise FP8, same- and mixed-format MXFP8, TN / NN / NT, batched multidim FP8, the fused forward BIAS epilogue, and the fused GELU_AUX / GELU_AUX_BIAS epilogues. BGRADB and DGELU are -not implemented on any FlyDSL path yet. FlyDSL kernels require tile-aligned +not implemented on any FlyDSL path yet. MXFP4 (gfx1250) covers TN / NN / NT +with FP32 / BF16 / FP16 output against a torch.matmul reference only. FlyDSL kernels require tile-aligned launch dimensions, so its shapes are aligned to the 256x256x128 kernel contract rather than reusing the odd-sized Triton edge-mask cases. """ @@ -47,6 +49,7 @@ from transformer_engine.pytorch import Float8Tensor, is_mxfp8_available from transformer_engine.pytorch.cpp_extensions.gemm import general_gemm from transformer_engine.pytorch.tensor.float8_tensor import Float8Quantizer +from transformer_engine.pytorch.tensor.mxfp4_tensor import MXFP4Quantizer from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Quantizer, MXFP8Tensor from transformer_engine.pytorch import torch_version from transformer_engine.pytorch.utils import get_gemm_backend @@ -93,7 +96,14 @@ def _device_capability(): reason="FlyDSL GEMM tests require the flydsl package on gfx950", ) -if _flydsl_available: +# FlyDSL MXFP4 is gfx1250-only (matching the C++ dispatch). +_flydsl_mxfp4_available = _CAP == (12, 5) and importlib.util.find_spec("flydsl") is not None +requires_flydsl_mxfp4_backend = pytest.mark.skipif( + not _flydsl_mxfp4_available, + reason="FlyDSL MXFP4 GEMM tests require the flydsl package on gfx1250", +) + +if _flydsl_available or _flydsl_mxfp4_available: pytest.importorskip("flydsl", reason="FlyDSL package is not installed") # gfx942's Triton fp32 matmul has a stable numerical divergence from @@ -162,6 +172,15 @@ def _device_capability(): (512, 512, 1024), # M != N ] +FLYDSL_MXFP4_SHAPES = [ + (512, 512, 512), + (512, 1024, 512), + (512, 512, 1024), # M != N + (512, 640, 512), # odd K tile count + (512, 4096, 512), # K > 2048: scales streamed in the K loop + (512, 3200, 512), # streamed scales with a partial last scale group +] + LAYOUTS = ["TN", "NN", "NT"] # Same-format FP8 combos, shared by both backends. @@ -334,6 +353,36 @@ def create_mxfp8_tensors( return A_mxfp8, B_mxfp8, A_mxfp8.dequantize(), B_mxfp8.dequantize() +_E2M1_MAGNITUDES = [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0] + + +def _dequantize_mxfp4(data, scale): + """Dequantize packed FP4 ``data`` [rows, K/2] with E8M0 ``scale`` (padded) to FP32.""" + data = data.view(torch.uint8) + rows, k = data.shape[0], data.shape[1] * 2 + codes = torch.stack([data & 0xF, data >> 4], dim=-1).reshape(rows, k).long() + values = torch.tensor(_E2M1_MAGNITUDES, device=data.device)[codes & 7] + values = torch.where(codes & 8 != 0, -values, values) + scales = torch.exp2(scale[:rows, : k // 32].float() - 127.0) + return values * scales.repeat_interleave(32, dim=1) + + +def create_mxfp4_tensors(M, K, N, layout): + """Create MXFP4 inputs and references dequantized from the buffers ``layout`` consumes.""" + # the MXFP4 quantizer takes BF16 input + A_bf16, B_bf16 = create_regular_tensors(M, K, N, layout, dtype=torch.bfloat16) + quantizer = MXFP4Quantizer(fp4_dtype=tex.DType.kFloat4E2M1, rowwise=True, columnwise=True) + A_mxfp4, B_mxfp4 = quantizer(A_bf16), quantizer(B_bf16) + + def reference(t, rowwise): + if rowwise: + return _dequantize_mxfp4(t._rowwise_data, t._rowwise_scale_inv) + return _dequantize_mxfp4(t._columnwise_data, t._columnwise_scale_inv).T + + transa, transb = layout[0] == "T", layout[1] == "T" + return A_mxfp4, B_mxfp4, reference(A_mxfp4, transa), reference(B_mxfp4, not transb) + + # ============================================================================== # Triton backend (NVTE_GEMM_BACKEND=TRITON) # ============================================================================== @@ -1094,6 +1143,46 @@ def test_flydsl_mxfp8_unsupported_shape_falls_back(): assert_gemm_close(output, expected, atol=5e-3, rtol=1e-2) +@requires_flydsl_mxfp4_backend +@pytest.mark.parametrize("M, K, N", FLYDSL_MXFP4_SHAPES) +@pytest.mark.parametrize("layout", LAYOUTS) +@pytest.mark.parametrize( + "out_dtype", [torch.float32, torch.bfloat16, torch.float16], ids=["fp32", "bf16", "fp16"] +) +def test_flydsl_vs_pytorch_mxfp4(M, K, N, layout, out_dtype): + """Test MXFP4 FlyDSL GEMMs (gfx1250) against an FP32 PyTorch reference.""" + torch.manual_seed(42) + + A_mxfp4, B_mxfp4, A_deq, B_deq = create_mxfp4_tensors(M, K, N, layout) + + output = flydsl_call_gemm(A_mxfp4, B_mxfp4, layout, out_dtype=out_dtype, use_flydsl=True) + expected = compute_pytorch_reference(A_deq, B_deq, layout) + + assert output.dtype == out_dtype + if out_dtype == torch.float32: + assert_gemm_close(output, expected, atol=1e-3, rtol=1e-4) + else: + assert_gemm_close(output, expected, atol=5e-3, rtol=1e-2) + + +@requires_flydsl_mxfp4_backend +def test_flydsl_mxfp4_unsupported_shape_raises(): + """MXFP4 configs FlyDSL cannot run raise FlyDSLUnsupportedError. + + ``general_gemm`` raises it, as gfx1250 has no other MXFP4 GEMM. + """ + from transformer_engine.pytorch.flydsl_kernels.gemm import ( + FlyDSLUnsupportedError, + te_generic_gemm_flydsl, + ) + + torch.manual_seed(42) + # M=128 is not a multiple of the 256-wide kernel tile + A_mxfp4, B_mxfp4, _, _ = create_mxfp4_tensors(128, 512, 256, "TN") + with pytest.raises(FlyDSLUnsupportedError): + te_generic_gemm_flydsl(A_mxfp4, True, B_mxfp4, False, None, None, tex.DType.kBFloat16) + + # --- Approach 2: FlyDSL vs native C++ ``generic_gemm`` reference -------------- diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index ddb953138..92dac2268 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -437,6 +437,56 @@ def _validate_native_gemm_output_quantizer(quantization_params): _flydsl_import_warned = False +def _try_flydsl_gemm(*args, **kwargs): + """Run te_generic_gemm_flydsl; return None if the caller should fall back.""" + try: + # Lazy import keeps FlyDSL off the normal Transformer Engine + # import path. It is done inside the try so a wheel built without + # flydsl (NVTE_GEMM_BACKEND!=FLYDSL at build time) degrades to the + # default backend instead of raising a bare ImportError. + from ..flydsl_kernels.gemm import ( + FlyDSLUnsupportedError, + te_generic_gemm_flydsl, + ) + + return te_generic_gemm_flydsl( + *args, + **kwargs, + ) + except ImportError as exc: + # NVTE_GEMM_BACKEND=FLYDSL was requested but the flydsl package is + # missing or too old (see flydsl_kernels.gemm._MIN_FLYDSL). This + # is a misconfiguration, not an unsupported GEMM config, so always + # warn (once) regardless of the opt-in fallback flag before + # degrading to the default backend. + global _flydsl_import_warned + if not _flydsl_import_warned: + _flydsl_import_warned = True + warnings.warn( + "[FLYDSL WARNING]: NVTE_GEMM_BACKEND=FLYDSL but the flydsl " + "package is unavailable; falling back to the default backend. " + f"Install a supported version (e.g. `pip install flydsl`) to " + f"enable it. Reason: {exc}", + UserWarning, + stacklevel=3, + ) + return None + except FlyDSLUnsupportedError as exc: + warn_fallback = os.environ.get( + "NVTE_FLYDSL_GEMM_WARN_FALLBACK", + "0", + ).lower() not in ("", "0", "false", "no", "off") + + if warn_fallback: + warnings.warn( + "[FLYDSL WARNING]: FlyDSL GEMM does not support this configuration; " + f"falling back to the default backend. Reason: {exc}", + UserWarning, + stacklevel=3, + ) + return None + + def general_gemm( A: torch.Tensor, B: torch.Tensor, @@ -540,9 +590,39 @@ def general_gemm( # Use bfloat16 as default bias_dtype bias_dtype = TE_DType[torch.bfloat16 if bias is None else bias.dtype] - # MXFP4 GEMM: route to AITER a4w4 ASM kernels, unless the hipBLASLt backend is - # opted in via NVTE_ROCM_USE_HIPBLASLT_MXFP4 + # MXFP4 GEMM: FlyDSL on gfx1250 (no fallback), else route to AITER a4w4 ASM kernels, + # unless the hipBLASLt backend is opted in via NVTE_ROCM_USE_HIPBLASLT_MXFP4 if isinstance(A, MXFP4TensorStorage) or isinstance(B, MXFP4TensorStorage): + if IS_HIP_EXTENSION and get_device_compute_capability() == (12, 5): + from ..flydsl_kernels.gemm import te_generic_gemm_flydsl + + out, bias_grad, gelu_input, extra_output = te_generic_gemm_flydsl( + A, + transa, + B, + transb, + out, + quantization_params, + TE_DType[out_dtype] if out_dtype is not None else None, + bias, + bias_dtype, + gelu, + gelu_in, + grad, + workspace, + workspace.shape[0], + accumulate, + use_split_accumulator, + comm_overlap=ub, + comm_type=ub_type, + extra_output=extra_output, + bulk_overlap=bulk_overlap, + alpha=alpha, + beta=beta, + ) + if debug_quantizer is not None: + out = debug_quantizer.process_gemm_output(out) + return out, bias_grad, gelu_input, extra_output if not bool(int(os.environ.get("NVTE_ROCM_USE_HIPBLASLT_MXFP4", "0"))): result = mxfp4_gemm( A, @@ -625,55 +705,9 @@ def general_gemm( and get_device_compute_capability() == (9, 5) and gemm_backend == "flydsl" ) - if use_gemm_flydsl: - try: - # Lazy import keeps FlyDSL off the normal Transformer Engine - # import path. It is done inside the try so a wheel built without - # flydsl (NVTE_GEMM_BACKEND!=FLYDSL at build time) degrades to the - # default backend instead of raising a bare ImportError. - from ..flydsl_kernels.gemm import ( - FlyDSLUnsupportedError, - te_generic_gemm_flydsl, - ) - - out, bias_grad, gelu_input, extra_output = te_generic_gemm_flydsl( - *args, - **kwargs, - ) - except ImportError as exc: - # NVTE_GEMM_BACKEND=FLYDSL was requested but the flydsl package is - # missing or too old (see flydsl_kernels.gemm._MIN_FLYDSL). This - # is a misconfiguration, not an unsupported GEMM config, so always - # warn (once) regardless of the opt-in fallback flag before - # degrading to the default backend. - global _flydsl_import_warned - if not _flydsl_import_warned: - _flydsl_import_warned = True - warnings.warn( - "[FLYDSL WARNING]: NVTE_GEMM_BACKEND=FLYDSL but the flydsl " - "package is unavailable; falling back to the default backend. " - f"Install a supported version (e.g. `pip install flydsl`) to " - f"enable it. Reason: {exc}", - UserWarning, - stacklevel=2, - ) - - out, bias_grad, gelu_input, extra_output = tex.generic_gemm(*args, **kwargs) - except FlyDSLUnsupportedError as exc: - warn_fallback = os.environ.get( - "NVTE_FLYDSL_GEMM_WARN_FALLBACK", - "0", - ).lower() not in ("", "0", "false", "no", "off") - - if warn_fallback: - warnings.warn( - "[FLYDSL WARNING]: FlyDSL GEMM does not support this configuration; " - f"falling back to the default backend. Reason: {exc}", - UserWarning, - stacklevel=2, - ) - - out, bias_grad, gelu_input, extra_output = tex.generic_gemm(*args, **kwargs) + result = _try_flydsl_gemm(*args, **kwargs) if use_gemm_flydsl else None + if result is not None: + out, bias_grad, gelu_input, extra_output = result else: out, bias_grad, gelu_input, extra_output = tex.generic_gemm(*args, **kwargs) else: diff --git a/transformer_engine/pytorch/flydsl_kernels/__init__.py b/transformer_engine/pytorch/flydsl_kernels/__init__.py index 0057be1e2..d464e3cad 100644 --- a/transformer_engine/pytorch/flydsl_kernels/__init__.py +++ b/transformer_engine/pytorch/flydsl_kernels/__init__.py @@ -1,4 +1,4 @@ # Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. # # See LICENSE for license information. -"""FlyDSL kernels for ROCm gfx950 GEMM replacement.""" +"""FlyDSL kernels for ROCm gfx950 and gfx1250 (MXFP4) GEMM replacement.""" diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/__init__.py b/transformer_engine/pytorch/flydsl_kernels/gemm/__init__.py index aa20712c7..c47a43d07 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/__init__.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/__init__.py @@ -2,7 +2,7 @@ # # See LICENSE for license information. -"""FlyDSL GEMM kernels (dense, non-grouped) for BF16/FP16/FP32/FP8/MXFP8.""" +"""FlyDSL dense GEMM kernels: BF16/FP16/FP32/FP8/MXFP8 (gfx950), MXFP4 (gfx1250).""" # Supported flydsl release series. Single source of truth for the version # requirement, since flydsl is user-installed (not a declared dependency). A diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py index 4a1dbb09a..9be2c8512 100644 --- a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py @@ -15,6 +15,7 @@ from .fp32_gemm import fp32_matmul from .fp8_gemm import fp8_matmul from .mxfp8_gemm import mxfp8_matmul +from .mxfp4_gemm import mxfp4_matmul # TODO: Some backend-independent GEMM wrapper utilities overlap with the # Triton GEMM backend (PR#667), including operand classification, logical @@ -172,6 +173,16 @@ def _classify_input(t): except ImportError: pass + try: + from transformer_engine.pytorch.tensor.storage.mxfp4_tensor_storage import ( + MXFP4TensorStorage, + ) + + if isinstance(t, MXFP4TensorStorage): + return "mxfp4", t + except ImportError: + pass + try: from transformer_engine.pytorch.quantized_tensor import ( QuantizedTensorStorage, @@ -181,8 +192,8 @@ def _classify_input(t): raise ValueError( "The FlyDSL GEMM backend does not support " f"{type(t).__name__}. Only Float8Tensor / " - "Float8TensorStorage and MXFP8Tensor / " - "MXFP8TensorStorage are implemented." + "Float8TensorStorage, MXFP8Tensor / " + "MXFP8TensorStorage and MXFP4TensorStorage are implemented." ) except ImportError: pass @@ -1008,6 +1019,88 @@ def _run_mxfp8( return D, aux +def _select_mxfp4_storage_for_layout(A, transa, B, transb): + """Select TE MXFP4 data and scales for a layout.""" + + def select(t, rowwise, name): + orientation = "rowwise" if rowwise else "columnwise" + data = t._rowwise_data if rowwise else t._columnwise_data + scale = t._rowwise_scale_inv if rowwise else t._columnwise_scale_inv + if data is None or scale is None: + raise FlyDSLUnsupportedError( + f"FlyDSL MXFP4 requires existing {name} {orientation} data and scales" + ) + if t._shuffle_rowwise_data if rowwise else t._shuffle_columnwise_data: + raise FlyDSLUnsupportedError( + f"FlyDSL MXFP4 requires plain FP4 data; {name} {orientation} data is shuffled" + ) + if t._with_gemm_swizzled_scales: + raise FlyDSLUnsupportedError( + f"FlyDSL MXFP4 requires plain E8M0 scales; {name} scales are GEMM-swizzled" + ) + if t._fp4_dtype != tex.DType.kFloat4E2M1: + raise FlyDSLUnsupportedError(f"FlyDSL MXFP4 requires E2M1 {name}, got {t._fp4_dtype}") + flatten = _flatten_rowwise if rowwise else _flatten_columnwise + return flatten(data.view(torch.uint8), f"{name}.{orientation}"), scale, data + + A_data, A_scale, _ = select(A, bool(transa), "A") + B_data, B_scale, B_payload = select(B, not bool(transb), "B") + return A_data, A_scale, B_data, B_scale, B_payload + + +def _run_mxfp4( + A, + transa, + B, + transb, + D, + *, + output_dtype: torch.dtype, +): + """Dispatch MXFP4 (gfx1250) through the single TN kernel for TN/NN/NT.""" + layout = f"{'T' if transa else 'N'}{'T' if transb else 'N'}" + A_data, A_scale, B_data, B_scale, B_payload = _select_mxfp4_storage_for_layout( + A, transa, B, transb + ) + + # Kernel operand ownership is always swapped relative to TE + a_flydsl, a_scale = B_data, B_scale + b_flydsl, b_scale = A_data, A_scale + m, n = a_flydsl.shape[0], b_flydsl.shape[0] + if a_flydsl.shape[1] != b_flydsl.shape[1]: + raise FlyDSLUnsupportedError( + f"FlyDSL MXFP4 {layout} selected incompatible payloads: " + f"a={tuple(a_flydsl.shape)} and b={tuple(b_flydsl.shape)}" + ) + if a_flydsl.device != b_flydsl.device: + raise ValueError( + f"FlyDSL MXFP4 {layout} operands must be on the same device, got " + f"{a_flydsl.device} and {b_flydsl.device}" + ) + + if D is not None: + output_shape = torch.Size(D.shape) + elif transb: + output_shape = torch.Size((m, n)) + else: + output_shape = torch.Size((*B_payload.shape[:-1], n)) + if _product(output_shape) != m * n: + raise FlyDSLUnsupportedError( + f"FlyDSL MXFP4 {layout} logical output shape " + f"{tuple(output_shape)} does not match kernel output {(m, n)}" + ) + + D = _validate_or_allocate_output( + D, + shape=output_shape, + dtype=output_dtype, + device=a_flydsl.device, + backend_name=f"MXFP4 {layout}", + ) + mxfp4_matmul(a_flydsl, a_scale, b_flydsl, b_scale, D.view(m, n)) + return D + + def _select_fp8_storage_for_layout(A, transa, B, transb): """Select existing TE FP8 backings and normalize to one core contract. @@ -1245,6 +1338,7 @@ def te_generic_gemm_flydsl( TT is intentionally rejected. Supported dtypes: + - MXFP4 E2M1 A/B (gfx1250) with FP16, BF16, or FP32 output - MXFP8 E4M3/E5M2 A/B combinations with FP16, BF16, or FP32 output - tensor-wise E4M3/E5M2 FP8 A/B combinations with FP16, BF16, or FP32 output - BF16 input with FP16, BF16, or FP32 output @@ -1273,12 +1367,12 @@ def te_generic_gemm_flydsl( ): raise FlyDSLUnsupportedError("FlyDSL GEMM does not support comm+GEMM overlap") - if transa and transb: - raise FlyDSLUnsupportedError("FlyDSL GEMM does not support transa=True, transb=True (TT)") - a_kind, _ = _classify_input(A) b_kind, _ = _classify_input(B) + if transa and transb: + raise FlyDSLUnsupportedError("FlyDSL GEMM does not support transa=True, transb=True (TT)") + _validate_common_epilogue( quantizer=quantizer, bias=bias, @@ -1289,6 +1383,35 @@ def te_generic_gemm_flydsl( beta=beta, ) + if a_kind == "mxfp4" or b_kind == "mxfp4": + if a_kind != b_kind: + raise ValueError("Mixed MXFP4 and non-MXFP4 FlyDSL GEMM inputs are not supported") + if bias is not None: + raise FlyDSLUnsupportedError("FlyDSL MXFP4 GEMM fused bias is not implemented") + if gelu: + raise FlyDSLUnsupportedError("FlyDSL MXFP4 GEMM fused GELU is not implemented") + + mxfp4_output_dtypes = { + None: torch.bfloat16, + tex.DType.kFloat16: torch.float16, + tex.DType.kBFloat16: torch.bfloat16, + tex.DType.kFloat32: torch.float32, + } + if output_dtype not in mxfp4_output_dtypes: + raise FlyDSLUnsupportedError( + f"FlyDSL MXFP4 supports FP16, BF16, or FP32 output, got {output_dtype}" + ) + + D = _run_mxfp4( + A, + transa, + B, + transb, + D, + output_dtype=mxfp4_output_dtypes[output_dtype], + ) + return D, None, None, None + if a_kind == "mxfp8" or b_kind == "mxfp8": # Validate both are MXFP8 if a_kind != b_kind: diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp4_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp4_gemm.py new file mode 100644 index 000000000..f6e72cd74 --- /dev/null +++ b/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp4_gemm.py @@ -0,0 +1,691 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# License for AMD contributions = MIT. See LICENSE for more information + +"""FlyDSL MXFP4 TN GEMM for gfx1250. + +D[M, N] = a[M, K] @ b[N, K]^T with K-contiguous packed FP4 operands and E8M0 scales along K. TE's +MXFP4 columnwise storage is a materialized transpose, so this one kernel serves TN/NN/NT. +Tiles are staged with TDM into a 3-stage LDS pipeline and C is written back through LDS by TDM. +""" + +import torch + +import flydsl.compiler as flyc +import flydsl.expr as fx +from flydsl.expr import const_expr, range_constexpr + +from .exceptions import FlyDSLUnsupportedError + +MXFP4_BLOCK_SIZE = 32 +BLOCK_M = 256 +BLOCK_N = 256 +BLOCK_K = 128 +NUM_WAVES = 8 +WAVE_32 = 32 +NUM_THREADS = NUM_WAVES * WAVE_32 +NUM_STAGES = 3 # LDS pages per operand +# K-loop tile j is loaded by wave group j % 2, two waves per operand +TDM_GROUP_WAVES = NUM_WAVES // 2 +TDM_WAVES = 2 +PRO_SCALE_WAVES = 2 +SCALE_GROUP = 6 +SCALE_STAGES = 2 +LDS_B_START = 64 * 1024 +C_HALF_ROWS = BLOCK_M // 2 # C is stored in two halves +C_WAVE_ROWS = C_HALF_ROWS // NUM_WAVES +# A/B LDS rows: 16 words padded to 20 by the TDM, so ds_load_b128 reads avoid bank conflicts +ROW_WORDS = 16 +LDS_ROW_WORDS = 20 +SCHED_TAIL_MMAS = 8 +K_LOOP_BODY = SCALE_GROUP +OUT_KINDS = {torch.float32: 0, torch.bfloat16: 1, torch.float16: 2} + + +@flyc.kernel +def gemm_kernel( + A: fx.Tensor, + B: fx.Tensor, + A_s: fx.Tensor, + B_s: fx.Tensor, + C: fx.Tensor, + M: fx.Int32, + N: fx.Int32, + K: fx.Int32, + As_sw: fx.Int32, + Bs_sw: fx.Int32, + K_SHORT: fx.Constexpr[int], + K_REM: fx.Constexpr[int], + OUT_KIND: fx.Constexpr[int], +): + + def as_words(t, rows, words_per_row): + src = fx.get_iter(t) + ptr = fx.recast_iter(fx.PointerType.get(fx.Int32.ir_type, src.memspace, 4), src) + return fx.make_view(ptr, fx.make_layout((rows, words_per_row), (words_per_row, 1))) + + out_fx = (fx.Float32, fx.BFloat16, fx.Float16)[OUT_KIND] + out_bytes = 4 if OUT_KIND == 0 else 2 + c_stage_bytes = C_HALF_ROWS * BLOCK_N * out_bytes + + tid = fx.thread_idx.x + bid = fx.block_idx.x + blocks_n = N // BLOCK_N + bid_m = bid // blocks_n + bid_n = bid % blocks_n + + # wave-uniform, so TDM descriptors derived from it stay in SGPRs + lane = tid % 32 + wave = fx.Int32(fx.rocdl.readfirstlane(fx.Int32.ir_type, tid // 32)) + lane_mod16, khalf = lane % 16, lane // 16 + wm, wn = wave % 2, wave // 2 + + k_words = K // 8 + K_iters = K // BLOCK_K + # global offsets in 64 bits: rows * k_words and rows * N exceed 2^31 for large shapes + k_words64, N64 = fx.Int64(k_words), fx.Int64(N) + A_words = as_words(A, M, k_words) + B_words = as_words(B, N, k_words) + + # this lane's 16 B chunk khalf (khalf + 2 is 8 words on) + a_lds_lo = (wm * 16 + lane_mod16) * LDS_ROW_WORDS + 4 * khalf + b_lds_lo = (wn * 16 + lane_mod16) * LDS_ROW_WORDS + 4 * khalf + as_lds_base = (wm * 16 + lane_mod16) * SCALE_GROUP + bs_lds_base = (wn * 16 + lane_mod16) * SCALE_GROUP + + # tile k lives in LDS page k % NUM_STAGES + smem = fx.SharedAllocator(static=False) + PAGE_A, PAGE_B = BLOCK_M * LDS_ROW_WORDS, BLOCK_N * LDS_ROW_WORDS + lds_A = smem.allocate(fx.Array[fx.Int32, NUM_STAGES * PAGE_A]).peek() + # B starts at 64 KB so all its pages are reachable with immediate offsets from one base register + smem.allocate(fx.Array[fx.Int32, LDS_B_START // 4 - NUM_STAGES * PAGE_A]) + lds_B = smem.allocate(fx.Array[fx.Int32, NUM_STAGES * PAGE_B]).peek() + box_A = fx.make_layout((BLOCK_M, ROW_WORDS), (LDS_ROW_WORDS, 1)) + sA = fx.make_view( + lds_A.ptr, + fx.make_layout(((BLOCK_M, ROW_WORDS), NUM_STAGES), ((LDS_ROW_WORDS, 1), PAGE_A)), + ) + box_B = fx.make_layout((BLOCK_N, ROW_WORDS), (LDS_ROW_WORDS, 1)) + sB = fx.make_view( + lds_B.ptr, + fx.make_layout(((BLOCK_N, ROW_WORDS), NUM_STAGES), ((LDS_ROW_WORDS, 1), PAGE_B)), + ) + pad = dict(pad_interval=ROW_WORDS, pad_amount=LDS_ROW_WORDS - ROW_WORDS) + + PAGE_S = BLOCK_M * SCALE_GROUP + lds_As = smem.allocate(fx.Array[fx.Int32, SCALE_STAGES * PAGE_S]).peek() + lds_Bs = smem.allocate(fx.Array[fx.Int32, SCALE_STAGES * PAGE_S]).peek() + + # C is staged from LDS offset 0 after the K loop + if smem.allocated_bytes < c_stage_bytes: + smem.allocate(c_stage_bytes - smem.allocated_bytes) + lds_C_ptr = fx.recast_iter(out_fx, smem.base_ptr) + + sA_base_ptr = fx.get_iter(sA) + sB_base_ptr = fx.get_iter(sB) + sAs_base_ptr = lds_As.ptr + sBs_base_ptr = lds_Bs.ptr + + # K-loop TDM atoms over this wave's half tile; later K tiles are reached with imm_offset + tdm_wave = wave % TDM_WAVES + HALF_ROWS = BLOCK_M // TDM_WAVES + half_box = fx.make_layout((HALF_ROWS, 16), (16, 1)) + gA_half = fx.make_view( + fx.add_offset( + fx.get_iter(A_words), fx.Int64(bid_m * BLOCK_M + tdm_wave * HALF_ROWS) * k_words64 + ), + half_box, + ) + gB_half = fx.make_view( + fx.add_offset( + fx.get_iter(B_words), fx.Int64(bid_n * BLOCK_N + tdm_wave * HALF_ROWS) * k_words64 + ), + half_box, + ) + tdm_A = fx.rocdl.make_tdm_atom( + gA_half, [None, None], strides=[fx.Int64(k_words), None], num_warps=1, **pad + ) + tdm_B = fx.rocdl.make_tdm_atom( + gB_half, [None, None], strides=[fx.Int64(k_words), None], num_warps=1, **pad + ) + + def lds_half(base_ptr, page): + return fx.make_view( + fx.add_offset(base_ptr, page * PAGE_A + tdm_wave * HALF_ROWS * LDS_ROW_WORDS), + half_box, + ) + + # prologue atoms: whole tiles 0 and 1 + tdm_A1w, cA1w = fx.rocdl.cdna5.make_tiled_tdm_atom( + fx.rocdl.TensorLoad(), A_words, box_A, (BLOCK_M, 16), num_warps=1 + ) + blk_cA1w = fx.zipped_divide(cA1w, (BLOCK_M, 16))[None, (bid_m, None)] + tAsA1w, tAgA1w = fx.rocdl.cdna5.tdm_partition(tdm_A1w, 0, fx.make_layout(1, 1), sA, blk_cA1w) + tdm_B1w, cB1w = fx.rocdl.cdna5.make_tiled_tdm_atom( + fx.rocdl.TensorLoad(), B_words, box_B, (BLOCK_N, 16), num_warps=1 + ) + blk_cB1w = fx.zipped_divide(cB1w, (BLOCK_N, 16))[None, (bid_n, None)] + tBsB1w, tBgB1w = fx.rocdl.cdna5.tdm_partition(tdm_B1w, 0, fx.make_layout(1, 1), sB, blk_cB1w) + + # scale group g of a row block is a [rows][SCALE_GROUP] box at column g * SCALE_GROUP + def scale_ptr(t): + src = fx.get_iter(t) + return fx.recast_iter(fx.PointerType.get(fx.Int32.ir_type, src.memspace, 4), src) + + # scale offsets stay 32-bit (rows * K / 128 words; the host checks they fit) + sA_blk = fx.add_offset(scale_ptr(A_s), bid_m * BLOCK_M * As_sw) + sB_blk = fx.add_offset(scale_ptr(B_s), bid_n * BLOCK_N * Bs_sw) + + grp_box = fx.make_layout((BLOCK_M, SCALE_GROUP), (SCALE_GROUP, 1)) + + gAs_grp = fx.make_view(sA_blk, grp_box) + gBs_grp = fx.make_view(sB_blk, grp_box) + tdm_As_grp = fx.rocdl.make_tdm_atom( + gAs_grp, [None, None], strides=[fx.Int64(As_sw), None], num_warps=1 + ) + tdm_Bs_grp = fx.rocdl.make_tdm_atom( + gBs_grp, [None, None], strides=[fx.Int64(Bs_sw), None], num_warps=1 + ) + + def copy_scale_group(blk, sw, lds_base, g, full=False, page=None): + page = g % SCALE_STAGES if page is None else page + lds = fx.make_view(fx.add_offset(lds_base, page * PAGE_S), grp_box) + if const_expr(full): + atom, view = (tdm_As_grp, gAs_grp) if blk is sA_blk else (tdm_Bs_grp, gBs_grp) + fx.copy(atom, view, lds, imm_offset=fx.Int64(g) * (SCALE_GROUP * 4)) + else: + view = fx.make_view(fx.add_offset(blk, g * SCALE_GROUP), grp_box) + atom = fx.rocdl.make_tdm_atom( + view, [None, K_iters - g * SCALE_GROUP], strides=[fx.Int64(sw), None], num_warps=1 + ) + fx.copy(atom, view, lds) + + # whole tiles, for K-loop iterations that also load a scale group + gA_full = fx.make_view( + fx.add_offset(fx.get_iter(A_words), fx.Int64(bid_m * BLOCK_M) * k_words64), box_A + ) + gB_full = fx.make_view( + fx.add_offset(fx.get_iter(B_words), fx.Int64(bid_n * BLOCK_N) * k_words64), box_B + ) + tdm_A_full = fx.rocdl.make_tdm_atom( + gA_full, [None, None], strides=[fx.Int64(k_words), None], num_warps=1, **pad + ) + tdm_B_full = fx.rocdl.make_tdm_atom( + gB_full, [None, None], strides=[fx.Int64(k_words), None], num_warps=1, **pad + ) + + bC = fx.slice(fx.zipped_divide(C, (BLOCK_M, BLOCK_N)), (None, (bid_m, bid_n))) + + mma_atom = fx.make_mma_atom(fx.rocdl.WMMAScale(16, 16, 128, fx.Float4E2M1FN)) + tiled_mma = fx.make_tiled_mma(mma_atom, fx.make_layout((2, 4, 1), (1, 2, 0))) + thr_mma = tiled_mma.thr_slice(tid) + + copy_atom_i32 = fx.make_copy_atom(fx.UniversalCopy128b(), fx.Int32) + copy_atom_s32 = fx.make_copy_atom(fx.UniversalCopy32b(), fx.Int32) + copy_atom_c = fx.make_copy_atom( + fx.UniversalCopy32b() if OUT_KIND == 0 else fx.UniversalCopy16b(), out_fx + ) + + def new_frags(): + return ( + fx.make_rmem_tensor(fx.make_layout((8, 8, 1), (1, 8, 0)), fx.Int32), + fx.make_rmem_tensor(fx.make_layout((8, 4, 1), (1, 8, 0)), fx.Int32), + fx.make_rmem_tensor(fx.make_layout((1, 8, 1), (0, 1, 0)), fx.Int32), + fx.make_rmem_tensor(fx.make_layout((1, 4, 1), (0, 1, 0)), fx.Int32), + ) + + A_REPS, B_REPS = 8, 4 + acc_tile = fx.make_layout((8, 1, 1), (1, 0, 0)) + scale_tile = fx.make_layout((1, 1, 1), (0, 0, 0)) + frag_acc = [ + [fx.make_rmem_tensor(acc_tile, fx.Float32) for _ in range(B_REPS)] for _ in range(A_REPS) + ] + + def frag_view(t, off, layout): + return fx.make_view(fx.add_offset(fx.get_iter(t), off), layout) + + def frag_view_ptr(ptr, off): + return fx.make_view(fx.add_offset(ptr, off), fx.make_layout(1, 1)) + + # one TDM op per wave: tiles 0 and 1 and the scales + def issue_prologue_tdm(): + if wave < 4: + if wave == 0: + fx.copy(tdm_A1w, tAgA1w[None, 0], tAsA1w[None, 0]) + else: + if wave == 1: + fx.copy(tdm_B1w, tBgB1w[None, 0], tBsB1w[None, 0]) + else: + if wave == 2: + fx.copy(tdm_A1w, tAgA1w[None, 1], tAsA1w[None, 1]) + else: + fx.copy(tdm_B1w, tBgB1w[None, 1], tBsB1w[None, 1]) + else: + if wave < 6: + if wave == 4: + copy_scale_group(sA_blk, As_sw, sAs_base_ptr, 0) + else: + if const_expr(K_SHORT == 0 or K_SHORT > SCALE_GROUP): + copy_scale_group(sA_blk, As_sw, sAs_base_ptr, 1) + else: + if wave == 6: + copy_scale_group(sB_blk, Bs_sw, sBs_base_ptr, 0) + else: + if const_expr(K_SHORT == 0 or K_SHORT > SCALE_GROUP): + copy_scale_group(sB_blk, Bs_sw, sBs_base_ptr, 1) + + n_dsrd = 3 * (A_REPS + B_REPS) + n_mma = A_REPS * B_REPS + + # interleave tile k+1's LDS reads with gemm(k)'s WMMAs + def hot_loop_scheduler(): + n_spread = n_mma - SCHED_TAIL_MMAS + issued = 0 + for i in range_constexpr(n_spread): + want = (i + 1) * n_dsrd // n_spread + if want > issued: + fx.rocdl.sched_dsrd(want - issued) + issued = want + fx.rocdl.sched_mfma(1) + fx.rocdl.sched_mfma(SCHED_TAIL_MMAS) + fx.rocdl.sched_barrier(0) + + def gemm_step(src): + frag_A, frag_B, frag_sA, frag_sB = src + for i in range_constexpr(A_REPS): + for j in range_constexpr(B_REPS): + fx.gemm( + mma_atom, + frag_acc[i][j], + [frag_view(frag_A, 8 * i, acc_tile), frag_view(frag_sA, i, scale_tile)], + [frag_view(frag_B, 8 * j, acc_tile), frag_view(frag_sB, j, scale_tile)], + frag_acc[i][j], + ) + + def rt_ptrs0(): + return [fx.get_iter(v) for v in (gA_half, gB_half, gA_full, gB_full)] + [sA_blk, sB_blk] + + def rt_advance(ptrs, tiles): + tw = tiles * (BLOCK_K // 8) + return [fx.add_offset(p, tw) for p in ptrs[:4]] + [ + fx.add_offset(p, tiles) for p in ptrs[4:] + ] + + def rt_views(ptrs): + layouts = (half_box, half_box, box_A, box_B, grp_box, grp_box) + return tuple(fx.make_view(p, l) for p, l in zip(ptrs, layouts)) + + def issue_tdm_rt(kb, vw, ks, scales, exact): + vA_half, vB_half, vA_full, vB_full, vAs, vBs = vw + js = ks + NUM_STAGES + page, group = js % NUM_STAGES, js % 2 + k_off = fx.Int64(js * (BLOCK_K // 2)) + if wave // TDM_GROUP_WAVES == group: + gw = wave % TDM_GROUP_WAVES + if const_expr(scales): + gi = ks // SCALE_GROUP + 1 + gpage = (kb // SCALE_GROUP + gi) % SCALE_STAGES + s_off = fx.Int64(gi * SCALE_GROUP * 4) + if gw < 2: + if gw == 0: + fx.copy( + tdm_A_full, + vA_full, + fx.make_view(fx.add_offset(sA_base_ptr, page * PAGE_A), box_A), + imm_offset=k_off, + ) + else: + if const_expr(exact): + g = kb // SCALE_GROUP + gi + copy_scale_group(sA_blk, As_sw, sAs_base_ptr, g, page=gpage) + else: + lds = fx.make_view(fx.add_offset(sAs_base_ptr, gpage * PAGE_S), grp_box) + fx.copy(tdm_As_grp, vAs, lds, imm_offset=s_off) + else: + if gw == 2: + fx.copy( + tdm_B_full, + vB_full, + fx.make_view(fx.add_offset(sB_base_ptr, page * PAGE_B), box_B), + imm_offset=k_off, + ) + else: + if const_expr(exact): + g = kb // SCALE_GROUP + gi + copy_scale_group(sB_blk, Bs_sw, sBs_base_ptr, g, page=gpage) + else: + lds = fx.make_view(fx.add_offset(sBs_base_ptr, gpage * PAGE_S), grp_box) + fx.copy(tdm_Bs_grp, vBs, lds, imm_offset=s_off) + else: + if gw < TDM_WAVES: + fx.copy(tdm_A, vA_half, lds_half(sA_base_ptr, page), imm_offset=k_off) + else: + fx.copy(tdm_B, vB_half, lds_half(sB_base_ptr, page), imm_offset=k_off) + + def lds_to_regs_rt(r, dst, bases): + a_lo, b_lo, s_a, s_b, s_d = bases + page = r % NUM_STAGES + frag_A, frag_B, frag_sA, frag_sB = dst + for h in range_constexpr(2): + fx.copy( + copy_atom_i32, + fx.make_view( + fx.add_offset(sA_base_ptr, a_lo + (page * PAGE_A + 8 * h)), + fx.make_layout((4, 8), (1, 32 * LDS_ROW_WORDS)), + ), + frag_view(frag_A, 4 * h, fx.make_layout((4, 8), (1, 8))), + ) + fx.copy( + copy_atom_i32, + fx.make_view( + fx.add_offset(sB_base_ptr, b_lo + (page * PAGE_B + 8 * h)), + fx.make_layout((4, 4), (1, 64 * LDS_ROW_WORDS)), + ), + frag_view(frag_B, 4 * h, fx.make_layout((4, 4), (1, 8))), + ) + if const_expr(r // SCALE_GROUP % SCALE_STAGES == 1): + s_a, s_b = s_a + s_d, s_b + s_d + col = r % SCALE_GROUP + for i in range_constexpr(A_REPS): + fx.copy( + copy_atom_s32, + frag_view_ptr(sAs_base_ptr, s_a + (col + i * 32 * SCALE_GROUP)), + frag_view(frag_sA, i, fx.make_layout(1, 1)), + ) + fx.rocdl.sched_barrier(0x7FF) + for j in range_constexpr(B_REPS): + fx.copy( + copy_atom_s32, + frag_view_ptr(sBs_base_ptr, s_b + (col + j * 64 * SCALE_GROUP)), + frag_view(frag_sB, j, fx.make_layout(1, 1)), + ) + fx.rocdl.sched_barrier(0x7FF) + + def wait_tile_rt(ks, pending): + if const_expr(pending == 0): + fx.rocdl.s_wait_tensorcnt(0) + else: + if wave // TDM_GROUP_WAVES == (ks + 2) % 2: + fx.rocdl.s_wait_tensorcnt(1) + else: + fx.rocdl.s_wait_tensorcnt(0) + + def k_step_rt(ks, cur, bases, issue=None, scales=False, exact=False, pending=1): + fx.rocdl.sched_barrier(0) + wait_tile_rt(ks, pending) + fx.rocdl.s_wait_dscnt(0) + fx.barrier() + fx.rocdl.sched_barrier(0) + if const_expr(issue is not None): + issue_tdm_rt(issue[0], issue[1], ks, scales, exact) + nxt = new_frags() + lds_to_regs_rt(ks + 1, nxt, bases) + gemm_step(cur) + hot_loop_scheduler() + return nxt + + def carry_values(fr): + frag_A, frag_B, frag_sA, frag_sB = fr + return ( + [frag_view(frag_A, 8 * i, acc_tile).load() for i in range_constexpr(A_REPS)] + + [frag_view(frag_B, 8 * j, acc_tile).load() for j in range_constexpr(B_REPS)] + + [frag_sA.load(), frag_sB.load()] + + [acc.load() for acc_row in frag_acc for acc in acc_row] + ) + + def from_carry(vals): + fr = new_frags() + frag_A, frag_B, frag_sA, frag_sB = fr + for i in range_constexpr(A_REPS): + frag_view(frag_A, 8 * i, acc_tile).store(vals[i]) + for j in range_constexpr(B_REPS): + frag_view(frag_B, 8 * j, acc_tile).store(vals[A_REPS + j]) + frag_sA.store(vals[A_REPS + B_REPS]) + frag_sB.store(vals[A_REPS + B_REPS + 1]) + accs = vals[A_REPS + B_REPS + 2 :] + for i in range_constexpr(A_REPS): + for j in range_constexpr(B_REPS): + frag_acc[i][j].store(accs[i * B_REPS + j]) + return fr + + for acc_row in frag_acc: + for acc in acc_row: + acc.fill(0) + + issue_prologue_tdm() + if wave // TDM_WAVES == 1: + fx.rocdl.s_wait_tensorcnt(1) + else: + fx.rocdl.s_wait_tensorcnt(0) + fx.barrier() + bases = (a_lds_lo, b_lds_lo, as_lds_base, bs_lds_base, fx.Int32(PAGE_S)) + cur = new_frags() + lds_to_regs_rt(0, cur, bases) + + fx.rocdl.sched_barrier(0) + fx.rocdl.s_wait_tensorcnt(0) + fx.rocdl.s_wait_dscnt(0) + fx.barrier() + fx.rocdl.sched_barrier(0) + views0 = rt_views(rt_ptrs0()) + for ks in range_constexpr(-1, 1): + if const_expr(K_SHORT == 0 or ks + NUM_STAGES < K_SHORT): + issue_tdm_rt(0, views0, ks, False, False) + nxt = new_frags() + lds_to_regs_rt(1, nxt, bases) + gemm_step(cur) + hot_loop_scheduler() + cur = nxt + + if const_expr(K_SHORT == 0): + # kb & 2^30 is zero, but keeps the carried LDS bases from being hoisted + n_body = K_iters // K_LOOP_BODY - 2 + init = carry_values(cur) + list(bases) + rt_ptrs0() + for i, state in range(0, n_body, 1, init=init): + nc = len(state) - 11 + cur = from_carry(state[:nc]) + a_lo, b_lo, s_a, s_b, s_d = (fx.Int32(v) for v in state[nc : nc + 5]) + bases = (a_lo, b_lo, s_a, s_b, s_d) + ptrs = list(state[nc + 5 :]) + kb = fx.Int32(i) * K_LOOP_BODY + issue = (kb, rt_views(ptrs)) + for u in range_constexpr(K_LOOP_BODY): + at_group = (1 + u) % SCALE_GROUP == 0 + cur = k_step_rt(1 + u, cur, bases, issue, scales=at_group) + zero = kb & (1 << 30) + bases_next = [a_lo + zero, b_lo + zero, s_a + s_d, s_b + s_d, -s_d] + results = yield carry_values(cur) + bases_next + rt_advance(ptrs, K_LOOP_BODY) + nc = len(results) - 11 + cur = from_carry(results[:nc]) + bases = tuple(fx.Int32(v) for v in results[nc : nc + 5]) + kb = n_body * K_LOOP_BODY + issue = (kb, rt_views(list(results[nc + 5 :]))) + R = 2 * SCALE_GROUP + K_REM - 4 + else: + issue = (0, views0) + R = max(K_SHORT - 4, 0) + for u in range_constexpr(R): + at_group = (1 + u) % SCALE_GROUP == 0 and u < R - 3 + cur = k_step_rt(1 + u, cur, bases, issue, scales=at_group, exact=u >= R - 8) + ks_last = R + 3 if K_SHORT == 0 else K_SHORT - 1 + for ks in range_constexpr(max(ks_last - 2, 1), ks_last): + cur = k_step_rt(ks, cur, bases, pending=int(ks < ks_last - 1)) + gemm_step(cur) + x = bases[0] - wm * 16 * LDS_ROW_WORDS + lane_mod16, khalf = x // LDS_ROW_WORDS, x % LDS_ROW_WORDS // 4 + + # C store: each wave writes C_WAVE_ROWS rows of each half + c_wave_box = fx.make_layout((C_WAVE_ROWS, BLOCK_N), (BLOCK_N, 1)) + c_wave_row = bid_m * BLOCK_M + wave * C_WAVE_ROWS + gC_wave = [ + fx.make_view( + fx.add_offset( + fx.get_iter(C), + fx.Int64(c_wave_row + h * C_HALF_ROWS) * N64 + fx.Int64(bid_n * BLOCK_N), + ), + c_wave_box, + ) + for h in range(2) + ] + tdm_C = fx.rocdl.make_tdm_atom(gC_wave[0], [None, None], strides=[N64, None], num_warps=1) + lds_C_wave = fx.make_view(fx.add_offset(lds_C_ptr, wave * C_WAVE_ROWS * BLOCK_N), c_wave_box) + + if const_expr(OUT_KIND == 0): + frag_c = frag_acc + else: + frag_c = [[fx.make_fragment_like(acc, out_fx) for acc in acc_row] for acc_row in frag_acc] + for i in range_constexpr(A_REPS): + for j in range_constexpr(B_REPS): + frag_c[i][j].store(frag_acc[i][j].load().to(out_fx)) + + # accumulator (v, j, i) is tile row wm*16 + 8*khalf + v + 32*i, column wn*16 + lane_mod16 + 64*j + fx.rocdl.s_wait_dscnt(0) + fx.barrier() + lane_c = (wm * 16 + 8 * khalf) * BLOCK_N + wn * 16 + lane_mod16 + for h in range_constexpr(2): + if h > 0: + fx.rocdl.s_wait_tensorcnt(0) + fx.barrier() + # sched_barrier(0x7FF) keeps the stores from being merged into ds_store_2addr + for ii in range_constexpr(4): + for j in range_constexpr(4): + for v in range_constexpr(8): + src_c = frag_view(frag_c[4 * h + ii][j], v, fx.make_layout(1, 1)) + dst_c = fx.make_view( + fx.add_offset(lds_C_ptr, lane_c + (32 * ii + v) * BLOCK_N + 64 * j), + fx.make_layout(1, 1), + ) + fx.copy(copy_atom_c, src_c, dst_c) + fx.rocdl.sched_barrier(0x7FF) + fx.rocdl.s_wait_dscnt(0) + fx.barrier() + fx.copy(tdm_C, lds_C_wave, gC_wave[h]) + fx.rocdl.s_wait_tensorcnt(0) + + +@flyc.jit +def launch_gemm( + A: fx.Tensor, + B: fx.Tensor, + A_s: fx.Tensor, + B_s: fx.Tensor, + C: fx.Tensor, + M: fx.Int32, + N: fx.Int32, + K: fx.Int32, + As_sw: fx.Int32, + Bs_sw: fx.Int32, + K_SHORT: fx.Constexpr[int], + K_REM: fx.Constexpr[int], + OUT_KIND: fx.Constexpr[int], + stream: fx.Stream = fx.Stream(None), +): + grid_tot = (M // BLOCK_M) * (N // BLOCK_N) + gemm_kernel(A, B, A_s, B_s, C, M, N, K, As_sw, Bs_sw, K_SHORT, K_REM, OUT_KIND).launch( + grid=(grid_tot, 1, 1), block=(NUM_THREADS, 1, 1), stream=stream + ) + + +def _scale_row_words(scale, rows, K): + """Row stride in words of TE's row-major E8M0 scales (possibly padded), read in place.""" + if scale.dtype != torch.uint8 or scale.ndim != 2: + raise FlyDSLUnsupportedError( + f"FlyDSL MXFP4 expects 2D uint8 E8M0 scales, got {scale.dtype} {tuple(scale.shape)}" + ) + if scale.shape[0] < rows or scale.shape[1] < K // MXFP4_BLOCK_SIZE: + raise FlyDSLUnsupportedError( + f"FlyDSL MXFP4 scales {tuple(scale.shape)} do not cover " + f"({rows}, {K // MXFP4_BLOCK_SIZE})" + ) + if scale.stride(1) != 1 or scale.stride(0) % 4 != 0 or scale.data_ptr() % 4 != 0: + raise FlyDSLUnsupportedError( + f"FlyDSL MXFP4 needs word-aligned row-major scales, got strides {scale.stride()}" + ) + if rows * (scale.stride(0) // 4) > 2**31 - 1: + raise FlyDSLUnsupportedError(f"FlyDSL MXFP4 scale offsets exceed 32 bits ({rows} rows)") + return scale.stride(0) // 4 + + +def k_loop_constexprs(K): + """(K_SHORT, K_REM): the K tile count if the kernel needs no K loop, else K tiles mod + K_LOOP_BODY, so that kernels are compiled per remainder, not per K.""" + k_iters = K // BLOCK_K + if k_iters < 2 * K_LOOP_BODY: + return k_iters, 0 + return 0, k_iters % K_LOOP_BODY + + +def check_mxfp4_shape(m, n, k): + """Raise FlyDSLUnsupportedError unless the kernel supports C[m, n] with reduction depth k.""" + if m % BLOCK_M or n % BLOCK_N: + raise FlyDSLUnsupportedError( + f"FlyDSL MXFP4 needs M, N multiples of {BLOCK_M}, got M={m}, N={n}" + ) + if k % BLOCK_K or k < 2 * BLOCK_K: + raise FlyDSLUnsupportedError( + f"FlyDSL MXFP4 needs K a multiple of {BLOCK_K}, at least {2 * BLOCK_K}, got K={k}" + ) + + +def mxfp4_matmul( + a: torch.Tensor, + a_scale: torch.Tensor, + b: torch.Tensor, + b_scale: torch.Tensor, + D: torch.Tensor, + stream=None, +): + """D[M, N] = dequant(a, a_scale) @ dequant(b, b_scale)^T. + + a [M, K/2] and b [N, K/2] packed FP4; scales E8M0 covering [M, K/32] and [N, K/32]; + D float32, bfloat16 or float16. + """ + if a.ndim != 2 or b.ndim != 2 or D.ndim != 2: + raise ValueError( + "FlyDSL MXFP4 expects rank-2 a, b, D; got " + f"{tuple(a.shape)}, {tuple(b.shape)}, {tuple(D.shape)}" + ) + m, k2 = a.shape + n, kb2 = b.shape + k = 2 * k2 + if kb2 != k2: + raise ValueError(f"Incompatible MXFP4 operands: a{tuple(a.shape)} and b{tuple(b.shape)}") + if tuple(D.shape) != (m, n): + raise ValueError(f"D shape {tuple(D.shape)} != expected {(m, n)}") + if D.dtype not in OUT_KINDS: + raise FlyDSLUnsupportedError( + f"FlyDSL MXFP4 output must be float32, bfloat16 or float16, got {D.dtype}" + ) + if not (a.is_contiguous() and b.is_contiguous() and D.is_contiguous()): + raise FlyDSLUnsupportedError("FlyDSL MXFP4 requires contiguous a, b and D") + if len({a.device, b.device, a_scale.device, b_scale.device, D.device}) != 1: + raise ValueError("a, b, scales and D must be on the same device") + check_mxfp4_shape(m, n, k) + + if stream is None: + stream = torch.cuda.current_stream(a.device) + launch_gemm( + a.view(torch.uint8), + b.view(torch.uint8), + a_scale, + b_scale, + D, + m, + n, + k, + _scale_row_words(a_scale, m, k), + _scale_row_words(b_scale, n, k), + *k_loop_constexprs(k), + OUT_KINDS[D.dtype], + stream=stream, + ) + return D + + +__all__ = [ + "BLOCK_M", + "BLOCK_N", + "BLOCK_K", + "MXFP4_BLOCK_SIZE", + "check_mxfp4_shape", + "mxfp4_matmul", +] diff --git a/transformer_engine/pytorch/quantization.py b/transformer_engine/pytorch/quantization.py index 176ee9e07..8e42ae47e 100644 --- a/transformer_engine/pytorch/quantization.py +++ b/transformer_engine/pytorch/quantization.py @@ -218,9 +218,15 @@ def _compute_mxfp4_support() -> Tuple[bool, str]: """Return if mxfp4 support is available""" if IS_HIP_EXTENSION: gpu_arch = get_device_compute_capability() - if gpu_arch in ((9, 5),): # TODO: enable for gfx1250 when GEMM is available + if gpu_arch == (9, 5): return True, "" - return False, "Device arch gfx95x is required for MXFP4 execution." + if gpu_arch == (12, 5): + try: + from .flydsl_kernels.gemm import te_generic_gemm_flydsl # noqa: F401 + except ImportError as e: + return False, f"MXFP4 on gfx1250 requires the FlyDSL GEMM backend: {e}" + return True, "" + return False, "Device arch gfx950 or gfx1250 is required for MXFP4 execution." return False, "Only ROCm supports MXFP4" @@ -1815,19 +1821,26 @@ def make_quantizers(self) -> list: # scales are plain (VEC32_UE8M0) by default, but can optionally be pre-swizzled into # the 32x8 tile order that hipBLASLt's BLK32_UE8M0_32_8_EXT (mode 1001) reads in place # -- opt in via the recipe's use_swizzled_scales flag. The AITER a4w4 backend always - # needs the 16x16 weight shuffle and swizzled scales. - use_hipblaslt = bool(int(os.environ.get("NVTE_ROCM_USE_HIPBLASLT_MXFP4", "0"))) + # needs the 16x16 weight shuffle and swizzled scales. FlyDSL (gfx1250) takes plain data + # and scales. + use_flydsl = IS_HIP_EXTENSION and get_device_compute_capability() == (12, 5) + use_hipblaslt = not use_flydsl and bool( + int(os.environ.get("NVTE_ROCM_USE_HIPBLASLT_MXFP4", "0")) + ) use_swizzled = use_hipblaslt and self.recipe.use_swizzled_scales # AITER path swizzles scales; hipBLASLt path swizzles only when the recipe opts in. # FP4 data shuffle stays off on the hipBLASLt path regardless - swizzled_scales = use_swizzled if use_hipblaslt else True + if use_flydsl: + swizzled_scales = False + else: + swizzled_scales = use_swizzled if use_hipblaslt else True if self.mode == "forward": def _make_quantizer(idx: int): is_activation = idx % 3 == 0 is_weight = idx % 3 == 1 - if use_hipblaslt: + if use_hipblaslt or use_flydsl: shuffle_rowwise_data = False shuffle_columnwise_data = False elif is_activation: