diff --git a/ci/pytorch.sh b/ci/pytorch.sh index b5241f3d75..80eada7b6e 100755 --- a/ci/pytorch.sh +++ b/ci/pytorch.sh @@ -94,6 +94,7 @@ run_test_config(){ NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 run_default_fa_lbl "deterministic" 3 attention/test_attention.py -k "test_deterministic_bwd_ck" run_default_fa 1 attention/test_cp_utils.py run_default_fa 1 attention/test_kv_cache.py + run_default_fa 1 triton_kernels/test_blockwise_fp8.py run_default_fa 1 triton_kernels/test_cast.py run_default_fa 1 triton_kernels/test_cast_mxfp8.py run_default_fa 1 triton_kernels/test_cast_mxfp4.py diff --git a/tests/pytorch/test_grouped_linear.py b/tests/pytorch/test_grouped_linear.py index 51c16769a3..ead7966a57 100644 --- a/tests/pytorch/test_grouped_linear.py +++ b/tests/pytorch/test_grouped_linear.py @@ -58,6 +58,18 @@ seed = 1234 reset_rng_states() +if IS_HIP_EXTENSION: + from utils import EnvVarCleaner + + @pytest.fixture(autouse=True) + def reset_grouped_gemm_backend(): + # Snapshot/restore the process-global Triton grouped-GEMM env vars so a test + # that sets them (and may raise before its own cleanup) cannot leak the + # backend into unrelated tests. + env = EnvVarCleaner(["NVTE_USE_BLOCKWISE_GMM_TRITON", "NVTE_USE_GROUPED_GEMM_TRITON"]) + yield + + NVTE_TEST_NVINSPECT_ENABLED = int(os.environ.get("NVTE_TEST_NVINSPECT_ENABLED", "0")) if NVTE_TEST_NVINSPECT_ENABLED: @@ -300,8 +312,6 @@ def test_grouped_linear_accuracy( pytest.skip("Triton grouped gemm is only supported on HIP.") if IS_HIP_EXTENSION and dtype not in (torch.float32,) and fuse_wgrad_accumulation and not fp8: pytest.skip(f"ROCm does not support fused wgrad accumulation for {dtype}.") - if IS_HIP_EXTENSION and recipe is not None and recipe.float8_block_scaling(): - pytest.skip("ROCm grouped GEMM does not yet support FP8 block scaling.") if fp8 and fp8_model_params and NVTE_TEST_NVINSPECT_ENABLED: pytest.skip("FP8 parameters are not supported in debug mode.") if NVTE_TEST_NVINSPECT_ENABLED and delay_wgrad_compute: @@ -321,7 +331,10 @@ def test_grouped_linear_accuracy( ) if use_triton: - os.environ["NVTE_USE_GROUPED_GEMM_TRITON"] = "1" + if recipe is not None and recipe.float8_block_scaling(): + os.environ["NVTE_USE_BLOCKWISE_GMM_TRITON"] = "1" + else: + os.environ["NVTE_USE_GROUPED_GEMM_TRITON"] = "1" with quantized_model_init(enabled=fp8 and fp8_model_params, recipe=recipe): grouped_linear = GroupedLinear( @@ -385,9 +398,6 @@ def test_grouped_linear_accuracy( delay_wgrad_compute, ) - if use_triton: - os.environ.pop("NVTE_USE_GROUPED_GEMM_TRITON", None) - atol, rtol = 0, 0 if use_cutlass: atol, rtol = 1e-3, 1e-3 @@ -402,6 +412,40 @@ def test_grouped_linear_accuracy( torch.testing.assert_close(o, o_ref, rtol=rtol, atol=atol) +@pytest.mark.skipif(not IS_HIP_EXTENSION, reason="Blockwise FP8 grouped GEMM is ROCm-only.") +@pytest.mark.parametrize( + "in_features,out_features,expected", + [ + (256, 512, True), + (384, 512, True), # 384 % 128 == 0 (not a multiple of 256, still supported) + (192, 512, False), # in_features not a multiple of 128 + (256, 320, False), # out_features not a multiple of 128 + ], +) +def test_blockwise_fp8_gate_requires_128_aligned_features( + in_features, out_features, expected, monkeypatch +): + """The forward/dgrad kernel applies one scalar B scale per 128-wide N-tile, + so non-128-aligned in/out features must fall back to the default path.""" + from transformer_engine.pytorch.module.grouped_linear import _GroupedLinear + + monkeypatch.setenv("NVTE_USE_BLOCKWISE_GMM_TRITON", "1") + supported = _GroupedLinear._is_blockwise_fp8_triton_grouped_gemm_supported( + fp8=True, + recipe=recipe.Float8BlockScaling(), + use_bias=False, + backward_override=None, + cpu_offloading=False, + save_original_input=False, + debug=False, + unpad_output=False, + actual_m_splits=None, + in_features=in_features, + out_features=out_features, + ) + assert supported is expected + + @pytest.mark.skipif( torch.cuda.get_device_capability() != (9, 0), reason="Only enable CUTLASS grouped gemm on Hopper", @@ -472,6 +516,12 @@ def test_grouped_linear_accuracy_rocm_backends( ): pytest.skip("CK MXFP8 grouped GEMM only supported on gfx1250.") + if recipe is not None and recipe.float8_block_scaling(): + # The CUTLASS/HipKittens/CK grouped GEMM backends do not support FP8 block + # scaling (they abort at runtime). The blockwise path is covered by the + # default and Triton backends in test_grouped_linear_accuracy instead. + pytest.skip("CUTLASS/HipKittens/CK grouped GEMM backends do not support FP8 block scaling.") + monkeypatch.setenv("NVTE_USE_CUTLASS_GROUPED_GEMM", "1") monkeypatch.delenv("NVTE_USE_HIPKITTENS_GROUPED_GEMM", raising=False) monkeypatch.delenv("NVTE_USE_CK_GROUPED_GEMM", raising=False) @@ -754,8 +804,6 @@ def test_padding_grouped_linear_accuracy( ): if fp8_model_params and NVTE_TEST_NVINSPECT_ENABLED: pytest.skip("FP8 parameters are not supported in debug mode.") - if IS_HIP_EXTENSION and recipe is not None and recipe.float8_block_scaling(): - pytest.skip("ROCm grouped GEMM does not yet support FP8 block scaling.") skip_unsupported_backward_override( "grouped_linear", recipe, getattr(recipe, "backward_override", None) ) @@ -836,8 +884,6 @@ def test_padding_grouped_linear_accuracy_save_original_input( pytest.skip("FP8 parameters are not supported in debug mode.") if fp8 and recipe.delayed(): pytest.skip("DelayedScaling recipe is not supported with save_original_input") - if IS_HIP_EXTENSION and recipe is not None and recipe.float8_block_scaling(): - pytest.skip("ROCm grouped GEMM does not yet support FP8 block scaling.") skip_unsupported_backward_override( "grouped_linear", recipe, getattr(recipe, "backward_override", None) ) @@ -1953,6 +1999,69 @@ def test_grouped_linear_grouped_tensor_path_single_grouped_bias_delay_wgrad(monk grouped_linear.backward_dw() +@pytest.mark.skipif(not IS_HIP_EXTENSION, reason="Blockwise FP8 grouped GEMM is ROCm-only.") +@pytest.mark.skipif(not fp8_block_scaling_available, reason="FP8 block scaling is unsupported.") +def test_blockwise_fp8_weight_cache_reuse(monkeypatch): + """``is_first_microbatch`` caches the packed blockwise weight across microbatches. + + CI never passes ``is_first_microbatch``, so the blockwise cache-hit branch is + otherwise unexercised. Verify that (1) the first microbatch quantizes once and + stores the packed weight in the single workspace slot (``weight0``), (2) later + microbatches reuse it with no re-quantization, (3) the reuse is numerically + identical to quantizing fresh, and (4) an incompatible cached workspace fails loud + instead of running a wrong weight. + """ + monkeypatch.setenv("NVTE_USE_BLOCKWISE_GMM_TRITON", "1") + + from transformer_engine.pytorch.triton_kernels import blockwise_quantize as _bwq + + # Count quantization launches; the forward re-imports this symbol per call, so + # patching the source module is observed. + n_quant = {"count": 0} + _orig_quant = _bwq.quantize_fp8_blockwise_grouped_weight_qtensor + + def _counting_quant(*args, **kwargs): + n_quant["count"] += 1 + return _orig_quant(*args, **kwargs) + + monkeypatch.setattr( + _bwq, "quantize_fp8_blockwise_grouped_weight_qtensor", _counting_quant + ) + + num_gemms, in_features, out_features = 2, 256, 512 + fp8_recipe = recipe.Float8BlockScaling() + grouped_linear = GroupedLinear( + num_gemms, + in_features, + out_features, + bias=False, + params_dtype=torch.bfloat16, + device="cuda", + ).eval() + + m_splits = [32, 32] + x = torch.randn(sum(m_splits), in_features, dtype=torch.bfloat16, device="cuda") + + with torch.no_grad(), autocast(enabled=True, recipe=fp8_recipe): + out_first = grouped_linear(x, m_splits, is_first_microbatch=True) + assert n_quant["count"] == 1, "first microbatch should quantize the weight once" + assert "weight0" in grouped_linear._fp8_workspaces + cached = grouped_linear._fp8_workspaces["weight0"] + + with torch.no_grad(), autocast(enabled=True, recipe=fp8_recipe): + out_second = grouped_linear(x, m_splits, is_first_microbatch=False) + assert n_quant["count"] == 1, "later microbatch should reuse the cache, not re-quantize" + assert grouped_linear._fp8_workspaces["weight0"] is cached + torch.testing.assert_close(out_second, out_first, rtol=0, atol=0) + + # A live is_first_microbatch=False cache whose layout does not match the current + # weight must raise, not silently GEMM a wrong weight. + grouped_linear._fp8_workspaces["weight0"] = torch.empty(1, dtype=torch.uint8, device="cuda") + with torch.no_grad(), autocast(enabled=True, recipe=fp8_recipe): + with pytest.raises(RuntimeError, match="Cached blockwise weight workspace"): + grouped_linear(x, m_splits, is_first_microbatch=False) + + @pytest.mark.skipif(not _nvfp4_available, reason=_reason_for_no_nvfp4) def test_grouped_linear_grouped_tensor_path_skips_non_rht_nvfp4(monkeypatch): """Non-RHT NVFP4 falls back to the legacy path; check it stays numerically correct. diff --git a/tests/pytorch/triton_kernels/test_blockwise_fp8.py b/tests/pytorch/triton_kernels/test_blockwise_fp8.py new file mode 100644 index 0000000000..5a0d5d3176 --- /dev/null +++ b/tests/pytorch/triton_kernels/test_blockwise_fp8.py @@ -0,0 +1,286 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. + +"""Unit tests for ROCm Triton blockwise FP8 quantization and grouped GEMM.""" + +import pytest +import torch +from torch.utils.cpp_extension import IS_HIP_EXTENSION + +from transformer_engine.pytorch.triton_kernels.common import get_torch_e4m3_type +from transformer_engine.pytorch.triton_kernels.blockwise_quantize import ( + quantize_fp8_blockwise, + quantize_fp8_blockwise_weight, + quantize_fp8_blockwise_segment_m, +) +from transformer_engine.pytorch.triton_kernels.blockwise_fp8_grouped_gemm import ( + _grouped_gemm_fp8_blockwise_raw, + _grouped_gemm_fp8_blockwise_variable_k_raw, +) + +pytestmark = pytest.mark.skipif(not IS_HIP_EXTENSION, reason="ROCm Triton blockwise kernels only") + +BLOCK = 128 +FP8_DTYPE = get_torch_e4m3_type() +IN_DTYPES = [torch.bfloat16, torch.float16] + + +def _cdiv(n, d): + return (n + d - 1) // d + + +def _floor_to_pow2(scale: torch.Tensor) -> torch.Tensor: + # 0xFF800000 as signed int32: keep the exponent, zero the mantissa. + bits = scale.to(torch.float32).contiguous().view(torch.int32) + return (bits & -8388608).view(torch.float32) + + +def _group_offs(splits, device="cuda"): + offs = torch.zeros(len(splits) + 1, dtype=torch.int64, device=device) + offs[1:] = torch.cumsum(torch.tensor(splits, dtype=torch.int64, device=device), 0) + return offs + + +def _ref_rowwise_quantize( + x: torch.Tensor, dtype: torch.dtype, block: int = BLOCK, pow2: bool = False +): + m, n = x.shape + fp8_max = torch.finfo(dtype).max + nb = _cdiv(n, block) + x_pad = torch.nn.functional.pad(x.float(), (0, nb * block - n)) + tiles = x_pad.reshape(m, nb, block) + amax = tiles.abs().amax(dim=-1).clamp_min(1e-4) + scale = fp8_max / amax + if pow2: + scale = _floor_to_pow2(scale) + q = (tiles * scale.unsqueeze(-1)).clamp(-fp8_max, fp8_max).to(dtype) + q = q.reshape(m, nb * block)[:, :n].contiguous() + return q, (1.0 / scale).contiguous() + + +def _ref_colwise_quantize( + x: torch.Tensor, dtype: torch.dtype, block: int = BLOCK, pow2: bool = False +): + m, n = x.shape + fp8_max = torch.finfo(dtype).max + mb = _cdiv(m, block) + x_pad = torch.nn.functional.pad(x.float(), (0, 0, 0, mb * block - m)) + tiles = x_pad.reshape(mb, block, n) + amax = tiles.abs().amax(dim=1).clamp_min(1e-4) + scale = fp8_max / amax + if pow2: + scale = _floor_to_pow2(scale) + q = (tiles * scale.unsqueeze(1)).clamp(-fp8_max, fp8_max).to(dtype) + q = q.reshape(mb * block, n)[:m].contiguous() + return q, (1.0 / scale).contiguous() + + +def _ref_weight_quantize( + w: torch.Tensor, dtype: torch.dtype, block: int = BLOCK, pow2: bool = False +): + g, m, n = w.shape + fp8_max = torch.finfo(dtype).max + mb, nb = _cdiv(m, block), _cdiv(n, block) + w_pad = torch.nn.functional.pad(w.float(), (0, nb * block - n, 0, mb * block - m)) + tiles = w_pad.reshape(g, mb, block, nb, block) + amax = tiles.abs().amax(dim=(2, 4)).clamp_min(1e-4) + scale = fp8_max / amax + if pow2: + scale = _floor_to_pow2(scale) + q = (tiles * scale[:, :, None, :, None]).clamp(-fp8_max, fp8_max).to(dtype) + q = q.reshape(g, mb * block, nb * block)[:, :m, :n].contiguous() + return q, (1.0 / scale).contiguous() + + +def _dequant_rowwise(q, s, n, block=BLOCK): + return q.float() * s.repeat_interleave(block, dim=1)[:, :n] + + +def _dequant_colwise(q, s, m, block=BLOCK): + return q.float() * s.repeat_interleave(block, dim=0)[:m] + + +@pytest.mark.parametrize("shape", [(128, 256), (200, 128), (256, 384), (64, 192)]) +@pytest.mark.parametrize("dtype", IN_DTYPES, ids=str) +@pytest.mark.parametrize("axis", [0, 1]) +@pytest.mark.parametrize("pow2", [False, True]) +def test_quantize_fp8_blockwise(shape, dtype, axis, pow2): + x = torch.randn(*shape, dtype=dtype, device="cuda") + q, s = quantize_fp8_blockwise(x, FP8_DTYPE, axis=axis, block_size=BLOCK, pow2=pow2) + if axis == 1: + q_ref, s_ref = _ref_rowwise_quantize(x, FP8_DTYPE, pow2=pow2) + dq = _dequant_rowwise(q, s, x.shape[1]) + dq_ref = _dequant_rowwise(q_ref, s_ref, x.shape[1]) + else: + q_ref, s_ref = _ref_colwise_quantize(x, FP8_DTYPE, pow2=pow2) + dq = _dequant_colwise(q, s, x.shape[0]) + dq_ref = _dequant_colwise(q_ref, s_ref, x.shape[0]) + torch.testing.assert_close(s, s_ref, atol=1e-5, rtol=1e-4) + torch.testing.assert_close(dq, dq_ref, atol=0.10, rtol=0.10) + + +@pytest.mark.parametrize("shape", [(2, 256, 256), (3, 128, 256), (1, 200, 192), (256, 256)]) +@pytest.mark.parametrize("dtype", IN_DTYPES, ids=str) +@pytest.mark.parametrize("pow2", [False, True]) +def test_quantize_fp8_blockwise_weight(shape, dtype, pow2): + w = torch.randn(*shape, dtype=dtype, device="cuda") + q, s = quantize_fp8_blockwise_weight(w, FP8_DTYPE, block_size=BLOCK, pow2=pow2) + w3 = w if w.dim() == 3 else w.unsqueeze(0) + q_ref, s_ref = _ref_weight_quantize(w3, FP8_DTYPE, pow2=pow2) + if w.dim() == 2: + q_ref, s_ref = q_ref.squeeze(0), s_ref.squeeze(0) + torch.testing.assert_close(s, s_ref, atol=1e-5, rtol=1e-4) + n = w.shape[-1] + m = w.shape[-2] + dq = q.float() * ( + s.repeat_interleave(BLOCK, dim=-2).repeat_interleave(BLOCK, dim=-1)[..., :m, :n] + ) + dq_ref = q_ref.float() * ( + s_ref.repeat_interleave(BLOCK, dim=-2).repeat_interleave(BLOCK, dim=-1)[..., :m, :n] + ) + torch.testing.assert_close(dq, dq_ref, atol=0.10, rtol=0.10) + + +@pytest.mark.parametrize( + "m_splits", + [ + [80, 176], + [128, 128], + [0, 256], + [64, 0, 192], + [200, 56], + ], +) +@pytest.mark.parametrize("n", [128, 256, 192]) +@pytest.mark.parametrize("dtype", IN_DTYPES, ids=str) +def test_quantize_fp8_blockwise_segment_m(m_splits, n, dtype): + m = sum(m_splits) + x = torch.randn(m, n, dtype=dtype, device="cuda") + group_lens = torch.tensor(m_splits, dtype=torch.int64, device="cuda") + group_offs = _group_offs(m_splits) + x_fp8, scales, vk_lens, vk_offs = quantize_fp8_blockwise_segment_m( + x, FP8_DTYPE, BLOCK, group_lens, group_offs + ) + expected_lens = [((s + BLOCK - 1) // BLOCK) * BLOCK for s in m_splits] + expected_offs = [0] + for length in expected_lens: + expected_offs.append(expected_offs[-1] + length) + assert vk_lens.tolist() == expected_lens + assert vk_offs.tolist() == expected_offs + cursor = 0 + for valid, padded in zip(m_splits, expected_lens): + if padded > valid: + assert torch.all(x_fp8[cursor + valid : cursor + padded].float() == 0) + cursor += padded + assert x_fp8.shape[0] >= expected_offs[-1] + assert scales.dtype == torch.float32 + + +@pytest.mark.parametrize( + "splits,n,k", + [ + ([64, 96, 96], 256, 256), + ([128, 128], 256, 192), + ([80, 176], 192, 256), + ([0, 128, 128], 256, 256), + ([256], 128, 128), + ], +) +@pytest.mark.parametrize("trans_b", [True, False]) +@pytest.mark.parametrize("in_dtype", IN_DTYPES, ids=str) +@pytest.mark.parametrize("out_dtype", [torch.bfloat16, torch.float16], ids=str) +def test_grouped_gemm_fp8_blockwise_matches_dequant_ref(splits, n, k, trans_b, in_dtype, out_dtype): + g = len(splits) + m = sum(splits) + out_n = n if trans_b else k + if out_n % BLOCK != 0: + pytest.skip("blockwise grouped GEMM uses 128-wide N tiles") + a_k = k if trans_b else n + a = torch.randn(m, a_k, dtype=in_dtype, device="cuda") + b = torch.randn(g, n, k, dtype=in_dtype, device="cuda") + a_fp8, a_s = quantize_fp8_blockwise(a, FP8_DTYPE, axis=1, block_size=BLOCK) + b_fp8, b_s = quantize_fp8_blockwise_weight(b, FP8_DTYPE, block_size=BLOCK) + offs = _group_offs(splits) + + out = _grouped_gemm_fp8_blockwise_raw( + a_fp8, b_fp8, a_s, b_s, offs, trans_b=trans_b, out_dtype=out_dtype + ) + assert out.dtype == out_dtype + + a_dq = _dequant_rowwise(a_fp8, a_s, a_k) + b_dq = b_fp8.float() * ( + b_s.repeat_interleave(BLOCK, dim=1).repeat_interleave(BLOCK, dim=2)[:, :n, :k] + ) + ref = torch.zeros(m, out_n, dtype=torch.float32, device="cuda") + for i in range(g): + sl = slice(int(offs[i]), int(offs[i + 1])) + if sl.start == sl.stop: + continue + ref[sl] = a_dq[sl] @ (b_dq[i].T if trans_b else b_dq[i]) + torch.testing.assert_close(out.float(), ref, atol=0.15, rtol=0.05) + + +@pytest.mark.parametrize( + "splits,n,k", + [ + ([128, 128], 128, 128), + ([80, 176], 256, 128), + ([0, 256], 256, 128), + ([64, 96, 96], 128, 256), + # OUT_N == k == 192 is not a multiple of BLOCK_SIZE_N (128 or 256), so the + # wgrad kernel always hits an N tail tile regardless of the autotuned + # config. Covers tail-tile addressing / the un-wrapped store mask for both + # the fresh and ACCUMULATE paths. (Note: the wrapped-store race the mask + # removes converges to the correct value here, so this guards addressing, + # not the race itself.) + ([128, 128], 256, 192), + ([64, 96, 96], 128, 192), + ], +) +@pytest.mark.parametrize("accumulate", [False, True]) +@pytest.mark.parametrize("in_dtype", IN_DTYPES, ids=str) +@pytest.mark.parametrize("out_dtype", [torch.float32, torch.bfloat16], ids=str) +def test_variable_k_wgrad(splits, n, k, accumulate, in_dtype, out_dtype): + g = len(splits) + tokens = sum(splits) + dy = torch.randn(tokens, n, dtype=in_dtype, device="cuda") + x = torch.randn(tokens, k, dtype=in_dtype, device="cuda") + group_lens = torch.tensor(splits, dtype=torch.int64, device="cuda") + group_offs = _group_offs(splits) + go_col, go_s, _, vk_offs = quantize_fp8_blockwise_segment_m( + dy, FP8_DTYPE, BLOCK, group_lens, group_offs + ) + x_col, x_s, _, _ = quantize_fp8_blockwise_segment_m(x, FP8_DTYPE, BLOCK, group_lens, group_offs) + + fresh = _grouped_gemm_fp8_blockwise_variable_k_raw( + go_col, x_col, go_s, x_s, vk_offs, out_dtype=out_dtype, accumulate=False + ) + assert fresh.dtype == out_dtype + assert fresh.shape == (g, n, k) + + go_dq = _dequant_colwise(go_col, go_s, go_col.shape[0]) + x_dq = _dequant_colwise(x_col, x_s, x_col.shape[0]) + ref = torch.zeros(g, n, k, dtype=torch.float32, device="cuda") + for i in range(g): + sl = slice(int(vk_offs[i]), int(vk_offs[i + 1])) + if sl.start == sl.stop: + continue + ref[i] = go_dq[sl].T @ x_dq[sl] + torch.testing.assert_close(fresh.float(), ref, atol=0.15, rtol=0.05) + + if accumulate: + main_grad = torch.randn(g, n, k, dtype=out_dtype, device="cuda") + expected = main_grad.clone() + _grouped_gemm_fp8_blockwise_variable_k_raw( + go_col, + x_col, + go_s, + x_s, + vk_offs, + out=main_grad, + accumulate=True, + ) + torch.testing.assert_close( + main_grad.float(), (expected + fresh).float(), atol=1e-3, rtol=1e-4 + ) diff --git a/transformer_engine/pytorch/module/fp8_padding.py b/transformer_engine/pytorch/module/fp8_padding.py index 3a0073a492..41ff2607ba 100644 --- a/transformer_engine/pytorch/module/fp8_padding.py +++ b/transformer_engine/pytorch/module/fp8_padding.py @@ -6,7 +6,7 @@ """FP8 Padding API""" -from typing import List, Optional, Tuple +from typing import List, Optional, Tuple, Union import torch @@ -113,7 +113,8 @@ def forward( self, inp: torch.Tensor, m_splits: List[int], - ) -> Tuple[torch.Tensor, List[int]]: + m_splits_tensor: Optional[torch.Tensor] = None, + ) -> Tuple[torch.Tensor, Union[List[int], torch.Tensor]]: """ Apply the padding to the input. @@ -123,6 +124,10 @@ def forward( Input tensor. m_splits : List[int] List of integers representing the split of the input tensor. + m_splits_tensor : torch.Tensor, optional + Device copy of ``m_splits``. When provided, the padded split sizes are + returned as a tensor on the same device (rounded up on-device), instead + of a Python list, so the caller can keep the splits on the GPU. """ assert len(m_splits) == self.num_gemms, "Number of splits should match number of GEMMs." @@ -136,6 +141,8 @@ def forward( ] # no padding needed if m_splits == padded_m_splits: + if m_splits_tensor is not None: + return inp, m_splits_tensor return inp, m_splits is_grad_enabled = torch.is_grad_enabled() @@ -154,4 +161,7 @@ def forward( ) out = fn(*autograd_ctx, inp, non_tensor_args) + if m_splits_tensor is not None: + align = self.align_size + return out, ((m_splits_tensor + (align - 1)) // align) * align return out, padded_m_splits diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index f534da5c3b..56c1837660 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -399,6 +399,427 @@ def _forward_grouped_tensor( return out.view(-1, *inp.shape[1:-1], out.shape[-1]), new_workspaces + @staticmethod + def _packed_3d_view(tensors): + """[G, N, K] view of consecutive contiguous 2D buffers, or None.""" + g0 = tensors[0] + if g0 is None or g0.ndim != 2 or not g0.is_contiguous(): + return None + n, k = g0.shape + step = g0.numel() * g0.element_size() + nbytes_needed = g0.storage_offset() * g0.element_size() + len(tensors) * step + if g0.untyped_storage().size() < nbytes_needed: + return None + for i, g in enumerate(tensors): + if ( + g is None + or g.dtype != g0.dtype + or g.device != g0.device + or tuple(g.shape) != (n, k) + or not g.is_contiguous() + or g.untyped_storage().data_ptr() != g0.untyped_storage().data_ptr() + or g.data_ptr() != g0.data_ptr() + i * step + ): + return None + return g0.as_strided((len(tensors), n, k), (n * k, k, 1)) + + @staticmethod + def _expert_weights_as_3d(weights, dtype): + """[G, N, K] expert weights without stacking when storage is already packed. + + ``single_grouped_weight`` keeps experts as consecutive slices of one + ``GroupedTensor.rowwise_data`` buffer; those slices are a zero-copy view. + Discrete ``weight0..weightN`` still need a stack. + """ + packed = _GroupedLinear._packed_3d_view(weights) + if packed is not None: + if packed.dtype != dtype: + packed = packed.to(dtype) + return packed if packed.is_contiguous() else packed.contiguous() + return torch.stack([wt.to(dtype).contiguous() for wt in weights], 0).contiguous() + + @staticmethod + def _grouped_weight_blockwise_qtensor(weights, activation_dtype, dt, pow2_w): + """Packed 2D ``[G*N, K]`` grouped-weight ``Float8BlockwiseQTensor`` from a list + of per-expert blockwise ``Float8BlockwiseQTensor`` (``fp8_model_params``). + + When the per-expert ``_rowwise_data`` / ``_rowwise_scale_inv`` are consecutive + slices of one contiguous buffer (``single_grouped_weight``), the packed + ``[G*N, K]`` data and ``[G*(N/128), K/128]`` scales are reconstructed as a + zero-copy view. Discrete params fall back to a ``torch.stack`` copy. + """ + from ..triton_kernels.blockwise_fp8_grouped_gemm import _stack_weight_qtensors + from ..triton_kernels.blockwise_quantize import ( + wrap_fp8_blockwise_grouped_weight_qtensor, + ) + from ..triton_kernels.common import te_dtype_to_torch_dtype + + # Zero-copy [G, N, K] / [G, bn, bk] views when the per-expert payloads are + # consecutive slices of one buffer; the wrap builder flattens them to 2D. + b_fp8 = _GroupedLinear._packed_3d_view([w._rowwise_data for w in weights]) + b_scale = _GroupedLinear._packed_3d_view([w._rowwise_scale_inv for w in weights]) + if b_fp8 is None or b_scale is None: + b_fp8, b_scale = _stack_weight_qtensors(weights, te_dtype_to_torch_dtype(dt)) + return wrap_fp8_blockwise_grouped_weight_qtensor( + b_fp8, b_scale, activation_dtype, dt, pow2=pow2_w + ) + + @staticmethod + def _handle_fused_wgrad(weight, main_grad): + """Megatron DDP hook: mark wgrad as consumed and return a dummy (or None).""" + if hasattr(weight, "grad_added_to_main_grad"): + weight.grad_added_to_main_grad = True + shape = list(main_grad.shape) if main_grad is not None else list(weight.shape) + return get_dummy_wgrad( + shape, weight.dtype, zero=getattr(weight, "zero_out_wgrad", False) + ) + return None + + @staticmethod + def _is_blockwise_fp8_triton_grouped_gemm_supported( + *, + fp8, + recipe, + use_bias, + backward_override, + cpu_offloading, + save_original_input, + debug, + unpad_output, + actual_m_splits, + in_features, + out_features, + ) -> bool: + """ROCm Triton blockwise FP8 grouped GEMM: HIP, env, Float8BlockScaling 1x128/128x128.""" + return ( + IS_HIP_EXTENSION + and os.getenv("NVTE_USE_BLOCKWISE_GMM_TRITON", "0") == "1" + and fp8 + and recipe is not None + and recipe.float8_block_scaling() + and ( + recipe.x_block_scaling_dim, + recipe.w_block_scaling_dim, + recipe.grad_block_scaling_dim, + ) + == (1, 2, 1) + # The forward/dgrad kernel applies one scalar B scale per N-tile + # (BLOCK_SIZE_N == 128 == scale-block-N), so an N that is not a + # multiple of 128 would apply the wrong scale to the tail tile. + # Forward's N is out_features, dgrad's N is in_features. + and in_features % 128 == 0 + and out_features % 128 == 0 + and not ( + use_bias + or backward_override + or cpu_offloading + or save_original_input + or debug + or unpad_output + or actual_m_splits is not None + ) + ) + + @staticmethod + def _forward_blockwise_fp8_triton( + ctx, + *, + inp, + m_splits, + m_splits_tensor, + weights, + weight_quantizers, + activation_dtype, + is_grad_enabled, + fuse_wgrad_accumulation, + is_first_microbatch, + wgrad_store, + weight_workspaces=None, + cache_weight=False, + pow2_x=True, + pow2_w=True, + pow2_grad=True, + ): + """Blockwise FP8 grouped GEMM forward (ROCm Triton). + + Selected from :meth:`forward` under the ``Float8BlockScaling`` recipe when + :meth:`_is_blockwise_fp8_triton_grouped_gemm_supported` is true. Activations are + quantized rowwise (1x128 along K) and weights 128x128, both wrapped as + :class:`Float8BlockwiseQTensor`. When the weights are already blockwise-quantized + params (``fp8_model_params``) their stored rowwise data + scales are reused + directly instead of re-quantizing. The backward wgrad (see + :meth:`_backward_blockwise_fp8_triton`) uses a segment-padded columnwise operand + quantized from the original high-precision input. + """ + from ..triton_kernels.blockwise_quantize import ( + quantize_fp8_blockwise_act_operands, + quantize_fp8_blockwise_grouped_weight_qtensor, + ) + from ..triton_kernels.blockwise_fp8_grouped_gemm import ( + grouped_gemm_fp8_blockwise_triton_kernel, + ) + from ..tensor.float8_blockwise_tensor import Float8BlockwiseQTensor + + num_gemms = len(m_splits) + + # Resolve the FP8 dtype from the weight quantizer as a TE ``DType`` (the + # QTensor's native dtype). The Triton builders map it to the arch-correct + # torch FP8 dtype only at the kernel/``.view`` boundary. + if weight_quantizers and weight_quantizers[0] is not None: + dt = tex.DType(int(weight_quantizers[0].dtype)) + else: + dt = tex.DType.kFloat8E4M3 + in_features = weights[0].size(-1) + out_features = weights[0].size(0) + device = inp.device + + inp_shape = inp.shape + a = inp.reshape(-1, in_features).to(activation_dtype).contiguous() + + # Prefer the caller-provided GPU splits tensor (avoids a blocking H2D of + # pageable m_splits). Fall back to m_splits when it is not passed. + split_src = m_splits_tensor if m_splits_tensor is not None else m_splits + group_lens = split_src.to(device=device, dtype=torch.int64) + group_offs = torch.zeros(num_gemms + 1, dtype=torch.int64, device=device) + group_offs[1:] = torch.cumsum(group_lens, 0) + + # Quantize the activation rowwise (1x128 along K) as a Float8BlockwiseQTensor. + # When gradients are needed the same pass also emits the segment-padded + # columnwise operand for the variable-K wgrad into ``qa``'s columnwise slots. + qa = quantize_fp8_blockwise_act_operands( + a, + dt, + 128, + group_lens, + group_offs, + rowwise=True, + columnwise=is_grad_enabled, + pow2=pow2_x, + ) + + # Weight operand: a packed 2D [G*N, K] Float8BlockwiseQTensor. When the + # params are already blockwise-quantized (fp8_model_params) reuse their + # stored rowwise data + scales; otherwise quantize the high-precision + # weights once over the packed 3D view (zero-copy when + # ``single_grouped_weight``), no stack. + new_workspaces = [None] * num_gemms + if isinstance(weights[0], Float8BlockwiseQTensor): + # fp8_model_params: reuse the stored rowwise data + scales directly (no + # re-quantization). + qw = _GroupedLinear._grouped_weight_blockwise_qtensor( + weights, activation_dtype, dt, pow2_w + ) + else: + # High-precision weights: quantize the whole packed weight into a single + # Float8BlockwiseQTensor and cache it across microbatches. On a cache hit + # (``is_first_microbatch=False``) reuse the stored packed weight untouched so + # its identity is preserved; only a fresh quantization is staged for write-back. + update_ws = is_first_microbatch is None or is_first_microbatch + cached_qw = weight_workspaces[0] if weight_workspaces else None + if not update_ws: + expected_wshape = (num_gemms * out_features, in_features) + if not ( + isinstance(cached_qw, Float8BlockwiseQTensor) + and tuple(cached_qw.shape) == expected_wshape + ): + raise RuntimeError( + "Cached blockwise weight workspace is incompatible with the " + "current weights; expected a Float8BlockwiseQTensor of shape " + f"{expected_wshape}." + ) + qw = cached_qw + else: + w = _GroupedLinear._expert_weights_as_3d(weights, activation_dtype) + qw = quantize_fp8_blockwise_grouped_weight_qtensor(w, dt, pow2=pow2_w) + if cache_weight: + new_workspaces[0] = qw + + # Forward GEMM: out[seg] = A[seg] @ W[g]^T (trans_b=True). + out = grouped_gemm_fp8_blockwise_triton_kernel( + qa, qw, group_offs, trans_b=True, out_dtype=activation_dtype + ) + + if is_grad_enabled: + ctx.use_blockwise_fp8_triton = True + # The segment-padded columnwise activation for the variable-K wgrad rides + # in ``qa``'s columnwise slots (produced by the fused pass above). Save + # those raw fields plus the packed weight QTensor's 2D fields; both + # rewrapped into Float8BlockwiseQTensors in backward. + ctx.save_for_backward( + qa._columnwise_data, + qa._columnwise_scale_inv, + qw._rowwise_data, + qw._rowwise_scale_inv, + group_lens, + group_offs, + qa._vk_group_offs, + ) + ctx.dt = dt + ctx.pow2_x = pow2_x + ctx.pow2_w = pow2_w + ctx.pow2_grad = pow2_grad + ctx.activation_dtype = activation_dtype + ctx.num_gemms = num_gemms + ctx.inp_shape = inp_shape + ctx.out_features = out_features + ctx.requires_dgrad = inp.requires_grad + ctx.weight_requires_grad = weights[0].requires_grad + ctx.reduce_and_update_bwd_fp8_tensors = False + ctx.fuse_wgrad_accumulation = fuse_wgrad_accumulation + ctx.is_first_microbatch = is_first_microbatch + ctx.wgrad_store = wgrad_store + if fuse_wgrad_accumulation and ctx.weight_requires_grad: + ctx.origin_weight_refs = [weakref.ref(w) for w in weights] + ctx.origin_weights_overwrite_main_grad = getattr( + weights[0], "overwrite_main_grad", False + ) + if hasattr(weights[0], "__fsdp_param__"): + ctx.main_grad_funcs = [weights[i].get_main_grad for i in range(num_gemms)] + else: + ctx.main_grad_funcs = [ + lambda j=i: weights[j].main_grad for i in range(num_gemms) + ] + + return out.view(-1, *inp.shape[1:-1], out.shape[-1]), new_workspaces + + @staticmethod + def _backward_blockwise_fp8_triton(ctx, grad_output): + """Backward path paired with :meth:`_forward_blockwise_fp8_triton`.""" + from ..triton_kernels.blockwise_quantize import ( + quantize_fp8_blockwise_act_operands, + wrap_fp8_blockwise_grouped_weight_qtensor, + wrap_fp8_blockwise_segment_m_qtensor, + ) + from ..triton_kernels.blockwise_fp8_grouped_gemm import ( + grouped_gemm_fp8_blockwise_triton_kernel, + grouped_gemm_fp8_blockwise_variable_k_triton_kernel, + ) + + a_col, a_scol, b_fp8, b_scale, group_lens, group_offs, vk_offs = ctx.saved_tensors + dt = ctx.dt + g_out = grad_output.reshape(-1, ctx.out_features).contiguous() + + # Grad-output operands in one QTensor (``qdgrad``): rowwise for dgrad and/or + # segment-padded columnwise (variable-K) for wgrad. When both grads are + # needed they share a single fused pass over g_out; otherwise only the + # requested operand is produced. + qdgrad = quantize_fp8_blockwise_act_operands( + g_out, + dt, + 128, + group_lens, + group_offs, + rowwise=ctx.requires_dgrad, + columnwise=ctx.weight_requires_grad, + pow2=ctx.pow2_grad, + ) + + # dgrad: dX[seg] = dY[seg] @ W[g] (trans_b=False) -> [M_total, K]. + dgrad = None + if ctx.requires_dgrad: + # Rewrap the saved packed weight fields as a 2D [G*N, K] QTensor. + qw = wrap_fp8_blockwise_grouped_weight_qtensor( + b_fp8, b_scale, ctx.activation_dtype, dt, pow2=ctx.pow2_w + ) + dgrad = grouped_gemm_fp8_blockwise_triton_kernel( + qdgrad, + qw, + group_offs, + trans_b=False, + out_dtype=ctx.activation_dtype, + ) + + # wgrad: dW[g] = dY[g]^T @ X[g] (variable-K) -> [G, N, K]. + wgrad_list = [None] * ctx.num_gemms + if ctx.weight_requires_grad: + # qdgrad carries the segment-padded columnwise grad output (wgrad operand). + # Rewrap the saved activation columnwise fields as a columnwise-only QTensor. + a_operand = wrap_fp8_blockwise_segment_m_qtensor( + a_col, a_scol, vk_offs, ctx.activation_dtype, dt, pow2=ctx.pow2_x + ) + fuse = getattr(ctx, "fuse_wgrad_accumulation", False) + origin_weights = [None] * ctx.num_gemms + main_grads = [None] * ctx.num_gemms + accumulate = False + packed_out = None + if fuse: + origin_weight_refs = ctx.origin_weight_refs + ctx.origin_weight_refs = None + origin_weights = [ref() if ref is not None else None for ref in origin_weight_refs] + assert all( + w is not None for w in origin_weights + ), "weight was removed while fuse_wgrad_accumulation=True" + main_grads = [main_grad_func() for main_grad_func in ctx.main_grad_funcs] + for origin_weight, main_grad in zip(origin_weights, main_grads): + if main_grad is not None: + origin_weight.main_grad = main_grad + if ctx.is_first_microbatch is not None: + accumulate = not ctx.is_first_microbatch + else: + accumulate = True + if getattr(ctx, "origin_weights_overwrite_main_grad", False): + accumulate = False + packed_out = _GroupedLinear._packed_3d_view(main_grads) + wgrad_list = main_grads + out_dtype = ( + main_grads[0].dtype if main_grads[0] is not None else ctx.activation_dtype + ) + else: + out_dtype = ctx.activation_dtype + wgrad_3d = torch.empty( + (ctx.num_gemms, ctx.out_features, a_col.shape[1]), + dtype=out_dtype, + device=a_col.device, + ) + packed_out = wgrad_3d + wgrad_list = [wgrad_3d[g] for g in range(ctx.num_gemms)] + + def grouped_gemm_wgrad(*_unused): + dW = grouped_gemm_fp8_blockwise_variable_k_triton_kernel( + qdgrad, + a_operand, + out_dtype=out_dtype, + out=packed_out, + accumulate=accumulate and packed_out is not None and fuse, + ) + if fuse and packed_out is None: + for g, mg in enumerate(main_grads): + if mg is None: + continue + if accumulate: + mg.add_(dW[g]) + else: + mg.copy_(dW[g]) + # Signature matches WeightGradStore / GroupedLinear.backward_dw. + return None, [None] * ctx.num_gemms, None + + wgrad_store = getattr(ctx, "wgrad_store", None) + if wgrad_store is not None and wgrad_store.delay_wgrad_compute(): + # tensor_list[2] is the wgrad buffers backward_dw assigns to .grad. + wgrad_store.put( + [qdgrad._columnwise_data, a_operand._columnwise_data, wgrad_list], + grouped_gemm_wgrad, + ) + else: + grouped_gemm_wgrad() + + if fuse: + wgrad_list = [ + _GroupedLinear._handle_fused_wgrad(weight, main_grad) + for weight, main_grad in zip(origin_weights, main_grads) + ] + + grad_biases = [None] * ctx.num_gemms # bias not supported on this path + + # Grads match forward inputs: (inp, m_splits, non_tensor_args, *weights, *biases). + return ( + dgrad.view(ctx.inp_shape) if dgrad is not None else None, + None, # m_splits + None, # non_tensor_args + *wgrad_list, + *grad_biases, + ) + # pylint: disable=keyword-arg-before-vararg @staticmethod def forward( @@ -446,7 +867,12 @@ def forward( save_original_input = True # Check if Triton kernel should be used - use_grouped_gemm_triton = IS_HIP_EXTENSION and os.getenv("NVTE_USE_GROUPED_GEMM_TRITON", "0") == "1" and not fp8 and not fuse_wgrad_accumulation + use_grouped_gemm_triton = ( + IS_HIP_EXTENSION + and os.getenv("NVTE_USE_GROUPED_GEMM_TRITON", "0") == "1" + and not fp8 + and not fuse_wgrad_accumulation + ) num_gemms = len(m_splits) weights = weights_and_biases[:num_gemms] @@ -454,6 +880,43 @@ def forward( device = inp.device weight_requires_grad = weights[0].requires_grad + # Blockwise FP8 grouped GEMM (ROCm Triton) opt-in. Runs its own quantization + + # grouped GEMM and returns early, bypassing the default quantizer setup below. + blockwise_recipe = FP8GlobalStateManager.get_fp8_recipe() if fp8 else None + if _GroupedLinear._is_blockwise_fp8_triton_grouped_gemm_supported( + fp8=fp8, + recipe=blockwise_recipe, + use_bias=use_bias, + backward_override=backward_override, + cpu_offloading=cpu_offloading, + save_original_input=save_original_input, + debug=debug, + unpad_output=unpad_output, + actual_m_splits=actual_m_splits, + in_features=weights[0].size(-1), + out_features=weights[0].size(0), + ): + return _GroupedLinear._forward_blockwise_fp8_triton( + ctx, + inp=inp, + m_splits=m_splits, + m_splits_tensor=m_splits_tensor, + weights=weights, + weight_quantizers=weight_quantizers, + activation_dtype=activation_dtype, + is_grad_enabled=is_grad_enabled, + fuse_wgrad_accumulation=fuse_wgrad_accumulation, + is_first_microbatch=is_first_microbatch, + wgrad_store=wgrad_store, + weight_workspaces=weight_workspaces, + cache_weight=cache_weight, + # Match the recipe's scale rounding (pow2 by default) so the Triton + # path quantizes identically to the default Float8BlockScaling path. + pow2_x=blockwise_recipe.fp8_quant_fwd_inp.power_2_scale, + pow2_w=blockwise_recipe.fp8_quant_fwd_weight.power_2_scale, + pow2_grad=blockwise_recipe.fp8_quant_bwd_grad.power_2_scale, + ) + # Configure quantizers if save_original_input and isinstance(input_quantizers[0], Float8Quantizer): if FP8GlobalStateManager.get_fp8_recipe().custom(): @@ -553,12 +1016,19 @@ def forward( # tensors (like scales), but bulk allocation shares storage across all tensors, # so if scales can't be offloaded, nothing in the group can be offloaded. fused_padding_kwargs = {} - if actual_m_splits is not None and IS_HIP_EXTENSION \ - and inp_view.shape[0] == sum(actual_m_splits): + if ( + actual_m_splits is not None + and IS_HIP_EXTENSION + and inp_view.shape[0] == sum(actual_m_splits) + ): fused_padding_kwargs["valid_split_sections"] = actual_m_splits inputmats = tex.split_quantize( - inp_view, m_splits, input_quantizers, - disable_bulk_allocation=cpu_offloading, **fused_padding_kwargs) + inp_view, + m_splits, + input_quantizers, + disable_bulk_allocation=cpu_offloading, + **fused_padding_kwargs, + ) elif debug: inputmats = DebugQuantizer.multi_tensor_quantize( inp_view, input_quantizers, m_splits, activation_dtype @@ -632,12 +1102,17 @@ def forward( **kwargs, ) - output_unpadded = False - if unpad_output and actual_m_splits is not None and IS_HIP_EXTENSION and actual_m_splits != m_splits: + if ( + unpad_output + and actual_m_splits is not None + and IS_HIP_EXTENSION + and actual_m_splits != m_splits + ): out_unpadded = torch.empty( [sum(actual_m_splits), out.shape[-1]], - dtype=out.dtype, device=out.device, + dtype=out.dtype, + device=out.device, ) tex.fused_multi_row_unpadding(out, out_unpadded, m_splits, actual_m_splits) out = out_unpadded @@ -951,6 +1426,8 @@ def backward( ) -> Tuple[Union[torch.Tensor, None], ...]: # pylint: disable=missing-function-docstring with get_nvtx_range_context("_GroupedLinear_backward"): + if getattr(ctx, "use_blockwise_fp8_triton", False): + return _GroupedLinear._backward_blockwise_fp8_triton(ctx, grad_output) if ctx.use_grouped_tensor_path: return _GroupedLinear._backward_grouped_tensor(ctx, grad_output) @@ -958,7 +1435,7 @@ def backward( N = ctx.num_gemms num_inputs = ctx.num_input_tensors inputmats = saved_tensors[:num_inputs] - weights = saved_tensors[num_inputs: num_inputs + N] + weights = saved_tensors[num_inputs : num_inputs + N] saved_weights = saved_tensors[num_inputs + N : num_inputs + 2 * N] biases = saved_tensors[num_inputs + 2 * N : num_inputs + 3 * N] @@ -1010,13 +1487,19 @@ def backward( for i in range(ctx.num_gemms): grad_biases[i] = grad_output_mats[i].sum(dim=0) grad_output = tex.split_quantize( - grad_output_view, ctx.m_splits, - ctx.grad_output_quantizers, **bwd_fused_kwargs) + grad_output_view, + ctx.m_splits, + ctx.grad_output_quantizers, + **bwd_fused_kwargs, + ) else: # Multi-tensor quantize grad_output = tex.split_quantize( - grad_output_view, ctx.m_splits, - ctx.grad_output_quantizers, **bwd_fused_kwargs) + grad_output_view, + ctx.m_splits, + ctx.grad_output_quantizers, + **bwd_fused_kwargs, + ) elif ctx.debug: grad_output_mats = torch.split(grad_output_view, ctx.m_splits) for i in range(ctx.num_gemms): @@ -1032,9 +1515,9 @@ def backward( # wgrad GEMM. if not ctx.use_grouped_gemm_triton: grad_output = torch.split( - cast_if_needed(grad_output_view, ctx.activation_dtype), - ctx.m_splits, - ) + cast_if_needed(grad_output_view, ctx.activation_dtype), + ctx.m_splits, + ) else: grad_output = [cast_if_needed(grad_output_view, ctx.activation_dtype)] @@ -1106,14 +1589,21 @@ def backward( **kwargs, ) - if ctx.actual_m_splits is not None and ctx.actual_m_splits != ctx.m_splits \ - and not ctx.output_unpadded: + if ( + ctx.actual_m_splits is not None + and ctx.actual_m_splits != ctx.m_splits + and not ctx.output_unpadded + ): dgrad_unpadded = torch.empty( (sum(ctx.actual_m_splits), dgrad.shape[-1]), - dtype=dgrad.dtype, device=dgrad.device, + dtype=dgrad.dtype, + device=dgrad.device, ) tex.fused_multi_row_unpadding( - dgrad, dgrad_unpadded, ctx.m_splits, ctx.actual_m_splits, + dgrad, + dgrad_unpadded, + ctx.m_splits, + ctx.actual_m_splits, ) dgrad = dgrad_unpadded @@ -1140,7 +1630,7 @@ def backward( wgrad_list = torch.empty( (ctx.num_gemms, weights[0].size(0), weights[0].size(1)), dtype=ctx.activation_dtype, - device=ctx.device + device=ctx.device, ) if ctx.save_original_input: @@ -1159,12 +1649,15 @@ def backward( inputmats: list if ctx.fp8 and not ctx.debug: save_fused_kwargs = {} - if ctx.actual_m_splits is not None and IS_HIP_EXTENSION \ - and inp_view.shape[0] == sum(ctx.actual_m_splits): + if ( + ctx.actual_m_splits is not None + and IS_HIP_EXTENSION + and inp_view.shape[0] == sum(ctx.actual_m_splits) + ): save_fused_kwargs["valid_split_sections"] = ctx.actual_m_splits inputmats = tex.split_quantize( - inp_view, ctx.m_splits, ctx.input_quantizers, - **save_fused_kwargs) + inp_view, ctx.m_splits, ctx.input_quantizers, **save_fused_kwargs + ) elif ctx.debug: inputmats = DebugQuantizer.multi_tensor_quantize( inp_view, @@ -1998,6 +2491,34 @@ def _get_weight_tensors(self) -> List[Union[torch.Tensor, QuantizedTensorStorage if weight_tensors is None: # TODO(ksivaman): Remove this after GEMM integration. weight_tensors = grouped_weight.split_into_quantized_tensors() + # ROCm + CUDA fix: Split views don't carry the grad state that the autograd Function reads + # off ``weights[i]``, so mirror it from the grouped Parameter (needed on + # every backend; transitional until ``_get_weight_tensors`` returns the + # Parameter directly). + want_grad = grouped_weight.requires_grad + main_grad = getattr(grouped_weight, "main_grad", None) + per_expert_main_grad = None + if main_grad is not None: + # Must be an aliasing view (not a copy) so per-expert wgrad accumulates + # in place; requires a contiguous [num_gemms, out_features, in_features]. + try: + per_expert_main_grad = main_grad.view( + self.num_gemms, self.out_features, self.in_features + ) + except RuntimeError as e: + raise RuntimeError( + "single_grouped_weight with fuse_wgrad_accumulation requires main_grad to " + "be a contiguous [num_gemms * out_features, in_features] buffer so per-expert" + f" wgrad can accumulate in place (got shape {tuple(main_grad.shape)})." + ) from e + for i, w in enumerate(weight_tensors): + # ``requires_grad_`` only works on leaves. + if w.requires_grad != want_grad and w.is_leaf: + w.requires_grad_(want_grad) + if per_expert_main_grad is not None: + w.main_grad = per_expert_main_grad[i] + if hasattr(grouped_weight, "overwrite_main_grad"): + w.overwrite_main_grad = grouped_weight.overwrite_main_grad else: weight_tensors = [getattr(self, f"weight{i}") for i in range(self.num_gemms)] if not self.fp8 and any(isinstance(w, QuantizedTensorStorage) for w in weight_tensors): @@ -2018,6 +2539,13 @@ def _get_bias_tensors(self) -> List[torch.Tensor]: parts = grouped_bias.quantized_tensors if parts is None: parts = grouped_bias.split_into_quantized_tensors() + # Mirror ``requires_grad`` onto the per-expert views (bias never uses + # ``main_grad``: fuse-accumulation is weights-only). ``requires_grad_`` + # only works on leaves, so guard to avoid a RuntimeError on a non-leaf view. + want_grad = grouped_bias.requires_grad + for p in parts: + if p.requires_grad != want_grad and p.is_leaf: + p.requires_grad_(want_grad) return [p.reshape(-1) for p in parts] return [getattr(self, f"bias{i}") for i in range(self.num_gemms)] diff --git a/transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py index f7a3dae70b..3d4bab6d4a 100644 --- a/transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py @@ -1,3 +1,5 @@ +# This file was modified for portability to AMDGPU +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. # Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. @@ -34,6 +36,10 @@ class Float8BlockwiseQTensorStorage(QuantizedTensorStorage): _rowwise_scale_inv: Optional[torch.Tensor] _columnwise_scale_inv: Optional[torch.Tensor] _is_2D_scaled: bool + # ROCm-only: padded per-segment (MoE group) offsets for the variable-K grouped + # wgrad. Set only when the columnwise operand is segment-padded (columnwise data + # is [M_pad, N] rather than the standard transposed [N, M]); otherwise None. + _vk_group_offs: Optional[torch.Tensor] def __new__( cls, @@ -46,6 +52,7 @@ def __new__( is_2D_scaled: bool, *args, fake_dtype: Optional[torch.dtype] = None, + vk_group_offs: Optional[torch.Tensor] = None, **kwargs, ): if cls is Float8BlockwiseQTensorStorage: @@ -60,6 +67,7 @@ def __new__( instance._rowwise_scale_inv = rowwise_scale_inv instance._columnwise_scale_inv = columnwise_scale_inv instance._is_2D_scaled = is_2D_scaled + instance._vk_group_offs = vk_group_offs return instance diff --git a/transformer_engine/pytorch/triton_kernels/blockwise_fp8_grouped_gemm.py b/transformer_engine/pytorch/triton_kernels/blockwise_fp8_grouped_gemm.py new file mode 100644 index 0000000000..948e637e3b --- /dev/null +++ b/transformer_engine/pytorch/triton_kernels/blockwise_fp8_grouped_gemm.py @@ -0,0 +1,927 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +# +# Blockwise FP8 grouped GEMM Triton kernels (MoE), adapted from AMD Primus-Turbo +# (primus_turbo/triton/grouped_gemm/grouped_gemm_fp8_kernel.py). + +import contextlib +import os + +import torch +import triton +import triton.language as tl + +from .common import is_cdna3 +from .gmm.pid_preprocessing import remap_xcd_chunked + + +def get_num_cus() -> int: + return torch.cuda.get_device_properties(torch.cuda.current_device()).multi_processor_count + + +# ── AMD compiler knobs (scoped) ── +# The AMD codegen toggles live on ``triton.knobs.amd`` and are process-global: they +# are read when a kernel is compiled, so leaving them set would silently change the +# codegen of every unrelated Triton kernel compiled afterwards (cast, gmm, MXFP8, +# ...). We therefore apply them only around our own launches and restore the prior +# state on exit. Triton's setter keeps the backing env var in sync with the +# attribute, so we snapshot/restore both. +_AMD_KNOB_ENV = { + "use_async_copy": "TRITON_HIP_USE_ASYNC_COPY", + "scalarize_packed_fops": "AMDGCN_SCALARIZE_PACKED_FOPS", + "use_block_pingpong": "TRITON_HIP_USE_BLOCK_PINGPONG", +} + + +def _amd_knob_overrides(*, is_tn: bool) -> dict: + """Knob values to apply for the enclosed GEMM launch. + + gfx950 (CDNA4): force async_copy / scalarize / block_pingpong on. + gfx942 (CDNA3): async_copy / scalarize help NT/NN but regress TN/wgrad + ~5-8%, so gate them on the GEMM layout (block_pingpong left at its default). + """ + if is_cdna3(): + enable = not is_tn + return {"use_async_copy": enable, "scalarize_packed_fops": enable} + return { + "use_async_copy": True, + "scalarize_packed_fops": True, + "use_block_pingpong": True, + } + + +@contextlib.contextmanager +def _amd_compiler_knobs(*, is_tn: bool): + """Temporarily apply AMD Triton compiler knobs, restoring prior state on exit. + + Scoped rather than set once because the knobs are process-global and read at + kernel compile time; leaving them set would leak into other kernels' codegen. + """ + amd = getattr(getattr(triton, "knobs", None), "amd", None) + if amd is None: + yield + return + overrides = _amd_knob_overrides(is_tn=is_tn) + saved = { + name: (name in amd.__dict__, amd.__dict__.get(name), os.environ.get(env_key)) + for name, env_key in _AMD_KNOB_ENV.items() + if name in overrides + } + try: + for name, value in overrides.items(): + setattr(amd, name, value) + yield + finally: + for name, (had_override, override_val, env_val) in saved.items(): + if had_override: + amd.__dict__[name] = override_val + else: + amd.__dict__.pop(name, None) + env_key = _AMD_KNOB_ENV[name] + if env_val is None: + os.environ.pop(env_key, None) + else: + os.environ[env_key] = env_val + + +NUM_XCDS = 8 + +# First call per autotune key uses balanced group_offs so the cached config +# is not locked to one uneven MoE routing (Triton keys on G/N/K, not splits). +_grouped_blockwise_warmed = set() +_grouped_blockwise_vk_warmed = set() + + +def _get_grouped_blockwise_autotune_configs(): + """Curated fwd/dgrad tiles from Primus-Turbo (32-config sweep → these 8). + + BLOCK_SIZE_N is pinned to 128: the kernel loads one B scale per N-tile + (BN == SCALE_BLOCK_N). BN=256 applies the wrong scale to half the columns. + """ + return [ + triton.Config( + { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 8, + "CHUNK_SIZE": 32, + }, + num_warps=4, + num_stages=2, + ), + triton.Config( + { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "CHUNK_SIZE": 32, + }, + num_warps=4, + num_stages=2, + ), + triton.Config( + { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 8, + "CHUNK_SIZE": 64, + }, + num_warps=4, + num_stages=2, + ), + triton.Config( + { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "CHUNK_SIZE": 32, + }, + num_warps=8, + num_stages=2, + ), + triton.Config( + { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 8, + "CHUNK_SIZE": 32, + }, + num_warps=8, + num_stages=2, + ), + triton.Config( + { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "CHUNK_SIZE": 32, + }, + num_warps=8, + num_stages=1, + ), + triton.Config( + { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 8, + "CHUNK_SIZE": 64, + }, + num_warps=8, + num_stages=2, + ), + triton.Config( + { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "CHUNK_SIZE": 64, + }, + num_warps=8, + num_stages=1, + ), + ] + + +def _bwd_autotune_configs(): + """Curated variable-K wgrad tiles from Primus-Turbo. BLOCK_SIZE_K=128 always.""" + return [ + triton.Config( + { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "CHUNK_SIZE": 32, + }, + num_warps=8, + num_stages=2, + ), + triton.Config( + { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "CHUNK_SIZE": 32, + }, + num_warps=8, + num_stages=2, + ), + triton.Config( + { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 8, + "CHUNK_SIZE": 32, + }, + num_warps=8, + num_stages=2, + ), + triton.Config( + { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "CHUNK_SIZE": 32, + }, + num_warps=8, + num_stages=2, + ), + triton.Config( + { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 8, + "CHUNK_SIZE": 32, + }, + num_warps=8, + num_stages=2, + ), + triton.Config( + { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "CHUNK_SIZE": 32, + }, + num_warps=8, + num_stages=1, + ), + triton.Config( + { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "CHUNK_SIZE": 32, + }, + num_warps=8, + num_stages=1, + ), + triton.Config( + { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 8, + "CHUNK_SIZE": 32, + }, + num_warps=4, + num_stages=2, + ), + ] + + +# Blockwise grouped FP8 kernel and public entrypoint + + +@triton.autotune(configs=_get_grouped_blockwise_autotune_configs(), key=["G", "N", "K"]) +@triton.jit() +def _grouped_blockwise_fp8_persistent_gemm_kernel( + # Pointers + A, # [M_total, K] FP8 + B, # [G, ?, ?] FP8 + C, # [M_total, N] + A_scales_ptr, # [K//128, M_total] float32 (pre-transposed for coalesced access) + B_scales_ptr, # [G, ?, ?] float32 (block-wise, layout depends on trans_b) + group_offs_ptr, # [G+1] int64 + # Dimensions + G, # number of groups (runtime) + N, + K, + # Strides + stride_am, # A row stride + stride_bg, # B group stride: b.stride(0) + stride_bn, # B N-stride (within a group) + stride_cm, # C row stride + stride_cn, # C col stride + # A_scales strides (pre-transposed: [K//128, M_total]) + stride_as_k, # A_scales_t.stride(0) + stride_as_m, # A_scales_t.stride(1) + # B_scales strides + stride_bs_g, # B_scales.stride(0) — group stride + stride_bs_n, # stride along N-block dimension + stride_bs_k, # stride along K-block dimension + # Constexpr strides + stride_ak: tl.constexpr, + stride_bk: tl.constexpr, + # Tile config + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + GROUP_SIZE_M: tl.constexpr, + NUM_SMS: tl.constexpr, + NUM_XCDS: tl.constexpr, + CHUNK_SIZE: tl.constexpr, + EVEN_K: tl.constexpr, + CACHE_MODIFIER: tl.constexpr, +): + """Persistent grouped block-wise FP8 GEMM kernel (CPU-sync-free).""" + pid = tl.program_id(0) + if NUM_XCDS != 1: + pid = remap_xcd_chunked(pid, NUM_SMS, NUM_XCDS, CHUNK_SIZE) + + num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) + + # ── Compute total tiles across all groups ── + total_tiles: tl.int32 = 0 + for _g in range(G): + m_g = (tl.load(group_offs_ptr + _g + 1) - tl.load(group_offs_ptr + _g)).to(tl.int32) + total_tiles += tl.cdiv(m_g, BLOCK_SIZE_M) * num_pid_n + + tl.assume(stride_am > 0) + tl.assume(stride_ak > 0) + tl.assume(stride_bn > 0) + tl.assume(stride_bk > 0) + tl.assume(stride_cm > 0) + tl.assume(stride_cn > 0) + + acc_dtype = tl.float32 + + for global_tile_id in range(pid, total_tiles, NUM_SMS): + # ── Find group via linear scan (O(G)) ── + group_idx: tl.int32 = 0 + tile_start: tl.int32 = 0 + cumsum: tl.int32 = 0 + for _g in range(G): + m_g_i = (tl.load(group_offs_ptr + _g + 1) - tl.load(group_offs_ptr + _g)).to(tl.int32) + tiles_g = tl.cdiv(m_g_i, BLOCK_SIZE_M) * num_pid_n + new_cumsum = cumsum + tiles_g + if global_tile_id >= new_cumsum: + group_idx = _g + 1 + tile_start = new_cumsum + cumsum = new_cumsum + + # ── Group-local tile → (pid_m, pid_n) ── + local_tile = global_tile_id - tile_start + m_start_g = tl.load(group_offs_ptr + group_idx) # int64 + M_g = (tl.load(group_offs_ptr + group_idx + 1) - m_start_g).to(tl.int32) + tiles_m_g = tl.cdiv(M_g, BLOCK_SIZE_M) + + num_pid_in_group = GROUP_SIZE_M * num_pid_n + swizzle_group = local_tile // num_pid_in_group + first_pid_m = swizzle_group * GROUP_SIZE_M + group_size_m = min(tiles_m_g - first_pid_m, GROUP_SIZE_M) + pid_m = first_pid_m + ((local_tile % num_pid_in_group) % group_size_m) + pid_n = (local_tile % num_pid_in_group) // group_size_m + tl.assume(pid_m >= 0) + tl.assume(pid_n >= 0) + + # ── Address computation ── + rm = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M_g + rn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N + rk = tl.arange(0, BLOCK_SIZE_K) + rn = tl.max_contiguous(tl.multiple_of(rn, BLOCK_SIZE_N), BLOCK_SIZE_N) + + group_offset_b = group_idx.to(tl.int64) * stride_bg + A_BASE = A + m_start_g * stride_am + rm[:, None] * stride_am + rk[None, :] * stride_ak + B_BASE = B + group_offset_b + rk[:, None] * stride_bk + rn[None, :] * stride_bn + + # A_scales pointer: pre-transposed [K//128, M_total] + as_ptrs_base = A_scales_ptr + (m_start_g + rm.to(tl.int64)) * stride_as_m + + # B_scales pointer: B_scales[g, pn, ki] (2D block scaling) + bs_ptr_base = B_scales_ptr + group_idx.to(tl.int64) * stride_bs_g + pid_n * stride_bs_n + + # ── K-loop with block-wise scaling (EVEN_K pattern) ── + loop_k = tl.cdiv(K, BLOCK_SIZE_K) + if not EVEN_K: + # Full K-blocks only; the ragged tail is handled below. This can be 0 + # (K < BLOCK_SIZE_K), so only assume non-negativity. + loop_k -= 1 + tl.assume(loop_k >= 0) + else: + # K is a positive multiple of BLOCK_SIZE_K, so there is >= 1 full block. + tl.assume(loop_k >= 1) + + acc = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=acc_dtype) + + for ki in range(0, loop_k): + if stride_ak == 1: + a = tl.load(tl.multiple_of(A_BASE, (1, 16)), cache_modifier=CACHE_MODIFIER) + else: + a = tl.load(tl.multiple_of(A_BASE, (16, 1)), cache_modifier=CACHE_MODIFIER) + + if stride_bk == 1: + b = tl.load(tl.multiple_of(B_BASE, (16, 1)), cache_modifier=CACHE_MODIFIER) + else: + b = tl.load(tl.multiple_of(B_BASE, (1, 16)), cache_modifier=CACHE_MODIFIER) + + partial = tl.dot(a, b) + + # Block-wise scales: a_s is [BLOCK_M] vector, b_s is scalar + a_s = tl.load(as_ptrs_base + ki * stride_as_k) + b_s = tl.load(bs_ptr_base + ki * stride_bs_k) + acc += partial * (a_s * b_s)[:, None] + + A_BASE += BLOCK_SIZE_K * stride_ak + B_BASE += BLOCK_SIZE_K * stride_bk + + if not EVEN_K: + # ── Last partial K-block (masked) ── + rk_last = loop_k * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K) + A_LAST = ( + A + m_start_g * stride_am + rm[:, None] * stride_am + rk_last[None, :] * stride_ak + ) + B_LAST = B + group_offset_b + rk_last[:, None] * stride_bk + rn[None, :] * stride_bn + if stride_ak == 1: + A_LAST = tl.multiple_of(A_LAST, (1, 16)) + else: + A_LAST = tl.multiple_of(A_LAST, (16, 1)) + if stride_bk == 1: + B_LAST = tl.multiple_of(B_LAST, (16, 1)) + else: + B_LAST = tl.multiple_of(B_LAST, (1, 16)) + a = tl.load(A_LAST, mask=rk_last[None, :] < K, other=0.0, cache_modifier=CACHE_MODIFIER) + b = tl.load(B_LAST, mask=rk_last[:, None] < K, other=0.0, cache_modifier=CACHE_MODIFIER) + partial = tl.dot(a, b) + a_s = tl.load(as_ptrs_base + loop_k * stride_as_k) + b_s = tl.load(bs_ptr_base + loop_k * stride_bs_k) + acc += partial * (a_s * b_s)[:, None] + + # ── Store output ── + c = acc.to(C.type.element_ty) + rm_s = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M_g + rn_s = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N + rn_s = tl.max_contiguous(tl.multiple_of(rn_s, BLOCK_SIZE_N), BLOCK_SIZE_N) + c_mask = (rm_s[:, None] < M_g) & (rn_s[None, :] < N) + C_ = C + m_start_g * stride_cm + rm_s[:, None] * stride_cm + rn_s[None, :] * stride_cn + tl.store(C_, c, c_mask) + + +# ═══════════════════════════════════════════════════════════════════════════════ +# Blockwise FP8 Variable-K Backward Kernel (persistent, CPU-sync-free) +# +# Computes: C[g] = LHS[g]^T @ RHS[g] with 1D+1D block-wise scales +# ═══════════════════════════════════════════════════════════════════════════════ + + +def _grouped_gemm_fp8_blockwise_raw( + a: torch.Tensor, + b: torch.Tensor, + a_scales: torch.Tensor, + b_scales: torch.Tensor, + group_offs: torch.Tensor, + trans_b: bool = True, + out_dtype: torch.dtype = torch.bfloat16, + a_scales_pretransposed: bool = False, +) -> torch.Tensor: + """Persistent grouped block-wise FP8 GEMM (CPU-sync-free) using Triton. + + Computes: out[offs[g]:offs[g+1], :] = A[offs[g]:offs[g+1], :] @ B_view[g] + with block-wise scaling for each K-block. + + Args: + a: [M_total, K] FP8 input (trans_a=False always). + b: [G, N, K] (if trans_b=True) or [G, K, N] FP8 weights. + a_scales: [M_total, K//128] float32, block-wise scale for A. When + ``a_scales_pretransposed`` is True this is already ``[K//128, M_total]`` + (the GEMM-ready layout) and the internal transpose is skipped. + b_scales: [G, ceil(N/128), ceil(K/128)] or [G, ceil(K/128), ceil(N/128)] float32. + group_offs: [G+1] int64 prefix sum of group lengths. + trans_b: If True, b[g] is [N, K] (transposed). + out_dtype: Output dtype (default bfloat16). + + Returns: + [M_total, N] output in out_dtype. + """ + assert a.ndim == 2, f"a must be 2D, got {a.shape}" + assert b.ndim == 3, f"b must be 3D, got {b.shape}" + assert b_scales.ndim == 3, f"b_scales must be 3D, got {b_scales.shape}" + + M_total, K = a.shape + G = b.shape[0] + + if trans_b: + N = b.shape[1] + stride_bk = b.stride(2) + stride_bn = b.stride(1) + stride_bs_n = b_scales.stride(1) + stride_bs_k = b_scales.stride(2) + else: + N = b.shape[2] + stride_bk = b.stride(1) + stride_bn = b.stride(2) + stride_bs_n = b_scales.stride(2) + stride_bs_k = b_scales.stride(1) + + stride_bg = b.stride(0) + stride_ak = a.stride(1) + + out = torch.empty((M_total, N), device=a.device, dtype=out_dtype) + A_scales_t = a_scales if a_scales_pretransposed else a_scales.T.contiguous() + num_sms = get_num_cus() + even_k = K % 128 == 0 + + def _launch(c_out, offs): + _grouped_blockwise_fp8_persistent_gemm_kernel[(num_sms,)]( + a, + b, + c_out, + A_scales_t, + b_scales, + offs, + G, + N, + K, + a.stride(0), + stride_bg, + stride_bn, + c_out.stride(0), + c_out.stride(1), + A_scales_t.stride(0), + A_scales_t.stride(1), + b_scales.stride(0), + stride_bs_n, + stride_bs_k, + stride_ak=stride_ak, + stride_bk=stride_bk, + NUM_SMS=num_sms, + NUM_XCDS=NUM_XCDS, + EVEN_K=even_k, + CACHE_MODIFIER=".ca", + waves_per_eu=0, + matrix_instr_nonkdim=16, + kpack=1, + ) + + # Triton's autotune key is (G, N, K) only. Prime once with balanced offs + # so the cached config is not locked to the first uneven MoE routing. + warm_key = (G, N, K, even_k) + # trans_a is always False here: forward/dgrad are NT/NN (never TN). Scope the + # AMD knobs to these launches so they don't leak into other Triton kernels. + with _amd_compiler_knobs(is_tn=False): + if warm_key not in _grouped_blockwise_warmed: + _grouped_blockwise_warmed.add(warm_key) + per = M_total // max(G, 1) + bal_offs = torch.arange(G + 1, device=group_offs.device, dtype=group_offs.dtype) * per + bal_offs[-1] = M_total + _launch(torch.empty_like(out), bal_offs) + + _launch(out, group_offs) + return out + + +def _stack_weight_qtensors(weights, torch_fp8_dtype): + """Stack per-expert weight ``Float8BlockwiseQTensor`` into packed grouped + buffers: FP8 data ``[G, N, K]`` and scales ``[G, ceil(N/128), ceil(K/128)]``.""" + b_fp8 = torch.stack([w._rowwise_data.view(torch_fp8_dtype) for w in weights], 0) + b_scales = torch.stack([w._rowwise_scale_inv for w in weights], 0) + return b_fp8, b_scales + + +def grouped_gemm_fp8_blockwise_triton_kernel( + a, + b, + group_offs: torch.Tensor, + trans_b: bool = True, + out_dtype: torch.dtype = torch.bfloat16, +) -> torch.Tensor: + """Grouped block-wise FP8 GEMM over ``Float8BlockwiseQTensor`` operands. + + ``a`` is a 1x128 rowwise-scaled ``Float8BlockwiseQTensor`` (``[M_total, K]``); + its stored ``rowwise_scale_inv`` is already in the ``[K//128, M_total]`` GEMM + layout. ``b`` is the weight operand in one of three forms: + + * a packed weight ``Float8BlockwiseQTensor`` -- either 2D ``[G*N, K]`` (scale + ``[G*(N/128), K/128]``, reshaped to ``[G, N, K]`` here using ``G`` from + ``group_offs``) or already 3D ``[G, N, K]``, + * a list of per-expert 128x128 weight ``Float8BlockwiseQTensor`` (stacked here), or + * a pre-stacked ``(b_fp8, b_scales)`` tuple of raw tensors (used by the backward + pass to reuse the buffers built during the forward pass). + + Raw FP8 fields are extracted here and handed to + :func:`_grouped_gemm_fp8_blockwise_raw`. + """ + import transformer_engine_torch as tex + from ..tensor.float8_blockwise_tensor import Float8BlockwiseQTensor + from .common import te_dtype_to_torch_dtype + + a_dt = te_dtype_to_torch_dtype(tex.DType(int(a._fp8_dtype))) + a_fp8 = a._rowwise_data.view(a_dt) + a_scales = a._rowwise_scale_inv # already [K//128, M_total] + + if isinstance(b, Float8BlockwiseQTensor): + b_dt = te_dtype_to_torch_dtype(tex.DType(int(b._fp8_dtype))) + b_fp8 = b._rowwise_data.view(b_dt) + b_scales = b._rowwise_scale_inv + if b_fp8.dim() == 2: + # Packed 2D [G*N, K] grouped weight: recover [G, N, K] using G from offs. + groups = group_offs.shape[0] - 1 + n = b_fp8.shape[0] // groups + b_fp8 = b_fp8.view(groups, n, b_fp8.shape[1]) + bn = b_scales.shape[0] // groups + b_scales = b_scales.view(groups, bn, b_scales.shape[1]) + elif isinstance(b, (list, tuple)) and isinstance(b[0], Float8BlockwiseQTensor): + b_dt = te_dtype_to_torch_dtype(tex.DType(int(b[0]._fp8_dtype))) + b_fp8, b_scales = _stack_weight_qtensors(b, b_dt) + else: + # Pre-stacked raw (b_fp8, b_scales); b_fp8 already an FP8-typed tensor. + b_fp8, b_scales = b + + return _grouped_gemm_fp8_blockwise_raw( + a_fp8, + b_fp8, + a_scales, + b_scales, + group_offs, + trans_b=trans_b, + out_dtype=out_dtype, + a_scales_pretransposed=True, + ) + + +# ── Blockwise FP8 Variable-K Backward Public API ── + + +@triton.autotune( + configs=_bwd_autotune_configs(), + key=["G", "OUT_M", "OUT_N"], +) +@triton.jit() +def _grouped_blockwise_fp8_variable_k_gemm_kernel( + # C[g] = LHS_g^T @ RHS_g * block_scales + LHS, # [M_padded_total, OUT_M] FP8 + RHS, # [M_padded_total, OUT_N] FP8 + C, # [G, OUT_M, OUT_N] + LHS_scales_ptr, # [ceil(M_padded/128), OUT_M] float32 + RHS_scales_ptr, # [ceil(M_padded/128), OUT_N] float32 + group_offs_ptr, # [G+1] int64 (padded segment offsets, each aligned to 128) + G, # number of groups + OUT_M, + OUT_N, + # Strides + stride_lhs_m, + stride_rhs_m, + stride_cg, + stride_cm, + stride_cn, + # LHS_scales strides + stride_ls_0, + stride_ls_1, + # RHS_scales strides + stride_rs_0, + stride_rs_1, + # Constexpr strides + stride_lhs_n: tl.constexpr, + stride_rhs_n: tl.constexpr, + # Tile config + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + GROUP_SIZE_M: tl.constexpr, + NUM_SMS: tl.constexpr, + NUM_XCDS: tl.constexpr, + CHUNK_SIZE: tl.constexpr, + CACHE_MODIFIER: tl.constexpr, + ACCUMULATE: tl.constexpr, +): + """Persistent grouped block-wise FP8 variable-K GEMM kernel (backward, CPU-sync-free). + + All groups share the same output dims (OUT_M × OUT_N), only the inner product + dimension M_g varies per group. 1D+1D scale pattern for TN/CRR layout. + + NOTE: Data is segment-padded to BLOCK_SIZE_K (128) boundaries by + quant_fp8_blockwise_segment_m_impl, so M_g is always a multiple of + BLOCK_SIZE_K. No masking is needed in the K-loop. + """ + pid = tl.program_id(0) + if NUM_XCDS != 1: + pid = remap_xcd_chunked(pid, NUM_SMS, NUM_XCDS, CHUNK_SIZE) + + tiles_m = tl.cdiv(OUT_M, BLOCK_SIZE_M) + tiles_n = tl.cdiv(OUT_N, BLOCK_SIZE_N) + tiles_per_group = tiles_m * tiles_n + total_tiles = G * tiles_per_group + + tl.assume(stride_lhs_m > 0) + tl.assume(stride_lhs_n > 0) + tl.assume(stride_rhs_m > 0) + tl.assume(stride_rhs_n > 0) + tl.assume(stride_cm > 0) + tl.assume(stride_cn > 0) + + acc_dtype = tl.float32 + + for global_tile in range(pid, total_tiles, NUM_SMS): + # ── Map to (group, local_tile) ── + group_idx = global_tile // tiles_per_group + local_tile = global_tile - group_idx * tiles_per_group + + # ── Swizzle local tile → (pid_m, pid_n) ── + num_pid_in_group = GROUP_SIZE_M * tiles_n + swizzle_group = local_tile // num_pid_in_group + first_pid_m = swizzle_group * GROUP_SIZE_M + group_size_m = min(tiles_m - first_pid_m, GROUP_SIZE_M) + pid_m = first_pid_m + ((local_tile % num_pid_in_group) % group_size_m) + pid_n = (local_tile % num_pid_in_group) // group_size_m + tl.assume(pid_m >= 0) + tl.assume(pid_n >= 0) + + # ── Group boundaries ── + m_start = tl.load(group_offs_ptr + group_idx) # int64 + M_g = (tl.load(group_offs_ptr + group_idx + 1) - m_start).to(tl.int32) + + # ── Output indices ── + rm = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % OUT_M + rn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % OUT_N + rk = tl.arange(0, BLOCK_SIZE_K) + rn = tl.max_contiguous(tl.multiple_of(rn, BLOCK_SIZE_N), BLOCK_SIZE_N) + + # ── Base pointers ── + LHS_BASE = ( + LHS + m_start * stride_lhs_m + rm[:, None] * stride_lhs_n + rk[None, :] * stride_lhs_m + ) + RHS_BASE = ( + RHS + m_start * stride_rhs_m + rk[:, None] * stride_rhs_m + rn[None, :] * stride_rhs_n + ) + + scale_row_start = m_start // BLOCK_SIZE_K + + # ── K-loop over M_g with block-wise 1D+1D scaling ── + # M_g is always a multiple of BLOCK_SIZE_K (data padded), so no masking needed. + loop_k = M_g // BLOCK_SIZE_K + acc = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=acc_dtype) + + for k in range(loop_k): + if stride_lhs_n == 1: + a = tl.load( + tl.multiple_of(LHS_BASE, (16, 1)), + cache_modifier=CACHE_MODIFIER, + ) + else: + a = tl.load( + tl.multiple_of(LHS_BASE, (1, 16)), + cache_modifier=CACHE_MODIFIER, + ) + + if stride_rhs_n == 1: + b = tl.load( + tl.multiple_of(RHS_BASE, (1, 16)), + cache_modifier=CACHE_MODIFIER, + ) + else: + b = tl.load( + tl.multiple_of(RHS_BASE, (16, 1)), + cache_modifier=CACHE_MODIFIER, + ) + + partial = tl.dot(a, b) + + # 1D+1D block-wise scales + scale_row = scale_row_start + k + a_s = tl.load(LHS_scales_ptr + scale_row * stride_ls_0 + rm * stride_ls_1) + b_s = tl.load(RHS_scales_ptr + scale_row * stride_rs_0 + rn * stride_rs_1) + acc += partial * a_s[:, None] * b_s[None, :] + + LHS_BASE += BLOCK_SIZE_K * stride_lhs_m + RHS_BASE += BLOCK_SIZE_K * stride_rhs_m + + # ── Store output ── + rm_s = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + rn_s = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + rn_s = tl.max_contiguous(tl.multiple_of(rn_s, BLOCK_SIZE_N), BLOCK_SIZE_N) + c_mask = (rm_s[:, None] < OUT_M) & (rn_s[None, :] < OUT_N) + C_ = ( + C + + group_idx.to(tl.int64) * stride_cg + + rm_s[:, None] * stride_cm + + rn_s[None, :] * stride_cn + ) + c = acc.to(C.type.element_ty) + if ACCUMULATE: + c += tl.load(C_, mask=c_mask, other=0) + tl.store(C_, c, c_mask) + + +# ── Blockwise FP8 Forward Public API ── + + +def grouped_gemm_fp8_blockwise_variable_k_triton_kernel( + lhs, + rhs, + out_dtype: torch.dtype = torch.bfloat16, + out: torch.Tensor = None, + accumulate: bool = False, +) -> torch.Tensor: + """Variable-K grouped block-wise FP8 GEMM (backward, 1D+1D scaling) using Triton. + + ``lhs`` and ``rhs`` are ``Float8BlockwiseQTensor`` operands carrying the + segment-padded columnwise-quantized grad and activation in their columnwise + slots (``_columnwise_data`` ``[M_pad, N]``, ``_columnwise_scale_inv`` + ``[ceil(M_pad/128), N]``) plus the shared padded segment offsets in + ``_vk_group_offs``. Raw FP8 fields are extracted here (uint8 data viewed back to + the FP8 dtype) and handed to :func:`_grouped_gemm_fp8_blockwise_variable_k_raw`. + """ + import transformer_engine_torch as tex + from .common import te_dtype_to_torch_dtype + + lhs_dt = te_dtype_to_torch_dtype(tex.DType(int(lhs._fp8_dtype))) + rhs_dt = te_dtype_to_torch_dtype(tex.DType(int(rhs._fp8_dtype))) + return _grouped_gemm_fp8_blockwise_variable_k_raw( + lhs._columnwise_data.view(lhs_dt), + rhs._columnwise_data.view(rhs_dt), + lhs._columnwise_scale_inv, + rhs._columnwise_scale_inv, + lhs._vk_group_offs, + out_dtype=out_dtype, + out=out, + accumulate=accumulate, + ) + + +def _grouped_gemm_fp8_blockwise_variable_k_raw( + lhs: torch.Tensor, + rhs: torch.Tensor, + lhs_scales: torch.Tensor, + rhs_scales: torch.Tensor, + group_offs: torch.Tensor, + out_dtype: torch.dtype = torch.bfloat16, + out: torch.Tensor = None, + accumulate: bool = False, +) -> torch.Tensor: + """Variable-K grouped block-wise FP8 GEMM (backward, 1D+1D scaling) using Triton. + + Computes: C[g] = lhs[offs[g]:offs[g+1]]^T @ rhs[offs[g]:offs[g+1]] + with 1D+1D block-wise scaling applied in the K-loop. + + Output: [G, OUT_M, OUT_N]. When ``out`` is provided the kernel writes in-place + (and adds into it if ``accumulate``). ``out_dtype`` is ignored in that case. + """ + assert lhs.ndim == 2 and rhs.ndim == 2 + assert lhs.shape[0] == rhs.shape[0] + OUT_M = lhs.shape[1] + OUT_N = rhs.shape[1] + G = group_offs.shape[0] - 1 + + if out is None: + assert not accumulate, "accumulate=True requires an existing out tensor" + out = torch.empty((G, OUT_M, OUT_N), device=lhs.device, dtype=out_dtype) + else: + assert out.shape == ( + G, + OUT_M, + OUT_N, + ), f"out must be {(G, OUT_M, OUT_N)}, got {tuple(out.shape)}" + num_sms = get_num_cus() + + def _launch(c_out, offs, accumulate_flag): + _grouped_blockwise_fp8_variable_k_gemm_kernel[(num_sms,)]( + lhs, + rhs, + c_out, + lhs_scales, + rhs_scales, + offs, + G, + OUT_M, + OUT_N, + lhs.stride(0), + rhs.stride(0), + c_out.stride(0), + c_out.stride(1), + c_out.stride(2), + lhs_scales.stride(0), + lhs_scales.stride(1), + rhs_scales.stride(0), + rhs_scales.stride(1), + stride_lhs_n=lhs.stride(1), + stride_rhs_n=rhs.stride(1), + NUM_SMS=num_sms, + NUM_XCDS=NUM_XCDS, + CACHE_MODIFIER=".ca", + ACCUMULATE=accumulate_flag, + waves_per_eu=0, + matrix_instr_nonkdim=16, + kpack=1, + ) + + # Autotune key is (G, OUT_M, OUT_N). Prime on a scratch buffer with + # balanced padded offs so accumulate=True never benchmarks into ``out``. + warm_key = (G, OUT_M, OUT_N, out.dtype, accumulate) + # Variable-K wgrad is TN (lhs^T @ rhs). Scope the AMD knobs to these launches + # so they don't leak into other Triton kernels. + with _amd_compiler_knobs(is_tn=True): + if warm_key not in _grouped_blockwise_vk_warmed: + _grouped_blockwise_vk_warmed.add(warm_key) + m_padded = lhs.shape[0] + per = max((m_padded // max(G, 1)) // 128 * 128, 128) + bal_offs = torch.arange(G + 1, device=group_offs.device, dtype=group_offs.dtype) * per + bal_offs[-1] = m_padded + _launch(torch.zeros_like(out), bal_offs, False) + + _launch(out, group_offs, accumulate) + return out diff --git a/transformer_engine/pytorch/triton_kernels/blockwise_quantize.py b/transformer_engine/pytorch/triton_kernels/blockwise_quantize.py new file mode 100644 index 0000000000..9bc965c707 --- /dev/null +++ b/transformer_engine/pytorch/triton_kernels/blockwise_quantize.py @@ -0,0 +1,610 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +# +# Blockwise FP8 quantization Triton kernels (1x128 activation, 128x128 weight), +# adapted from AMD Primus-Turbo (primus_turbo/triton/quantization/quant_blockwise.py +# and primus_turbo/pytorch/kernels/quantization/quantization_impl.py). + +import torch +import triton +import triton.language as tl +import transformer_engine_torch as tex + +from .common import te_dtype_to_torch_dtype + +__all__ = [ + "quantize_fp8_blockwise", + "quantize_fp8_blockwise_weight", + "quantize_fp8_blockwise_segment_m", + "quantize_fp8_blockwise_act_qtensor", + "quantize_fp8_blockwise_grouped_weight_qtensor", + "wrap_fp8_blockwise_grouped_weight_qtensor", + "quantize_fp8_blockwise_segment_m_operand", + "wrap_fp8_blockwise_segment_m_qtensor", + "quantize_fp8_blockwise_act_and_segment_m", + "quantize_fp8_blockwise_act_operands", +] + + +@triton.jit +def _floor_to_pow2(scale): + scale_bits = scale.to(tl.uint32, bitcast=True) & 0xFF800000 + return scale_bits.to(tl.float32, bitcast=True) + + +@triton.jit +def compute_scale_and_quant(x_tile, x_tile_abs, axis, FP8_MAX, ROUND_POW2: tl.constexpr): + x_tile_max = tl.max(x_tile_abs, axis=axis, keep_dims=True) + x_tile_max = tl.maximum(x_tile_max, 1e-4) + x_scales_tile = FP8_MAX / x_tile_max + if ROUND_POW2: + x_scales_tile = _floor_to_pow2(x_scales_tile) + x_fp8_tile = x_tile * x_scales_tile + x_fp8_tile = tl.clamp(x_fp8_tile, min=-FP8_MAX, max=FP8_MAX) + return x_fp8_tile, x_scales_tile + + +@triton.jit +def quant_fp8_blockwise_kernel( + x_ptr, + x_fp8_ptr, + x_scales_ptr, + M, + N, + BLOCK_SIZE: tl.constexpr, + FP8_MAX: tl.constexpr, + AXIS: tl.constexpr, + ROUND_POW2: tl.constexpr, +): + pid_m = tl.program_id(axis=0) + pid_n = tl.program_id(axis=1) + offs_m = tl.cast(pid_m * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE), tl.int64) + offs_n = tl.cast(pid_n * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE), tl.int64) + mask = (offs_m[:, None] < M) & (offs_n[None, :] < N) + + x_ptrs = x_ptr + offs_m[:, None] * N + offs_n[None, :] + x_tile = tl.load(x_ptrs, mask=mask, other=0.0).to(tl.float32) + x_tile_abs = tl.abs(x_tile) + + x_fp8_tile, x_scales_tile = compute_scale_and_quant( + x_tile, x_tile_abs, AXIS, FP8_MAX, ROUND_POW2 + ) + + x_fp8_ptrs = x_fp8_ptr + offs_m[:, None] * N + offs_n[None, :] + tl.store(x_fp8_ptrs, x_fp8_tile.to(x_fp8_ptr.dtype.element_ty), mask=mask) + + if AXIS == 1: + scale_offs = offs_m * tl.cdiv(N, BLOCK_SIZE) + pid_n + scale_mask = offs_m < M + else: + scale_offs = pid_m * N + offs_n + scale_mask = offs_n < N + x_scales_tile_inv = tl.reshape(1.0 / x_scales_tile, BLOCK_SIZE) + tl.store(x_scales_ptr + scale_offs, x_scales_tile_inv, mask=scale_mask) + + +@triton.jit +def quant_fp8_blockwise_for_weight_kernel( + w_ptr, + w_fp8_ptr, + w_scales_ptr, + M, + N, + BLOCK_SIZE: tl.constexpr, + FP8_MAX: tl.constexpr, + ROUND_POW2: tl.constexpr, +): + bid = tl.program_id(axis=0) + pid_m = tl.program_id(axis=1) + pid_n = tl.program_id(axis=2) + + batch_offset_w = bid * M * N + batch_offset_scales = bid * tl.cdiv(M, BLOCK_SIZE) * tl.cdiv(N, BLOCK_SIZE) + + offs_m = pid_m * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + offs_n = pid_n * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = (offs_m[:, None] < M) & (offs_n[None, :] < N) + + w_ptrs = w_ptr + batch_offset_w + offs_m[:, None] * N + offs_n[None, :] + w_tile = tl.load(w_ptrs, mask=mask, other=0.0).to(tl.float32) + + w_tile_abs = tl.abs(w_tile) + w_tile_max = tl.max(w_tile_abs) + w_tile_max = tl.maximum(w_tile_max, 1e-4) + w_scales = FP8_MAX / w_tile_max + if ROUND_POW2: + w_scales = _floor_to_pow2(w_scales) + w_fp8_tile = w_tile * w_scales + w_fp8_tile = tl.clamp(w_fp8_tile, min=-FP8_MAX, max=FP8_MAX) + + w_fp8_ptrs = w_fp8_ptr + batch_offset_w + offs_m[:, None] * N + offs_n[None, :] + tl.store(w_fp8_ptrs, w_fp8_tile.to(w_fp8_ptr.dtype.element_ty), mask=mask) + scale_offs = batch_offset_scales + pid_m * tl.cdiv(N, BLOCK_SIZE) + pid_n + w_scales_inv = 1.0 / w_scales + tl.store(w_scales_ptr + scale_offs, w_scales_inv) + + +@triton.jit +def quant_fp8_blockwise_grouped_kernel( + x_ptr, # Input tensor [M_in, N] + x_fp8_row_ptr, # Rowwise output [M_in, N] (unused when ROWWISE is False) + x_scales_row_ptr, # Rowwise scales [M_in, ceil(N/BLOCK_SIZE)] (unused when ROWWISE is False) + x_fp8_col_ptr, # Colwise padded output [M_pad, N] (unused when COLUMNWISE is False) + x_scales_col_ptr, # Colwise scales [ceil(M_pad/BLOCK_SIZE), N] (unused when COLUMNWISE is False) + group_offs_ptr, # Original group offsets [B+1] + padded_group_offs_ptr, # Padded group offsets [B+1] + N, + num_groups, + BLOCK_SIZE: tl.constexpr, + FP8_MAX: tl.constexpr, + ROUND_POW2: tl.constexpr, + ROWWISE: tl.constexpr, + COLUMNWISE: tl.constexpr, +): + """Grouped blockwise quantize: rowwise and/or segment-padded columnwise. + + The grid iterates the *segment-padded* output tiles (each MoE group padded up + to a ``BLOCK_SIZE`` boundary along M). Each program reads one ``[BLOCK_SIZE, + BLOCK_SIZE]`` block of the original (unpadded) input once and, from that single + tile, emits the operand(s) selected by the ``constexpr`` flags: + + * ``COLUMNWISE`` (axis=0, BLOCKx1 along M) -- writes FP8 data + 1D scales to the + *segment-padded* positions; this is the per-segment M padding used by the + variable-K wgrad (each group starts on a BLOCK_SIZE boundary). + * ``ROWWISE`` (axis=1, 1xBLOCK along N) -- writes FP8 data + 1D scales to the + *original* (unpadded) positions (the forward/dgrad activation operand); + padding rows are skipped. + + With both flags set this is the fused activation + wgrad quantize (one HBM read + of the input); with only ``COLUMNWISE`` it is the plain segment-padded colwise + quantize. + """ + tl.static_assert(ROWWISE or COLUMNWISE, "at least one of ROWWISE/COLUMNWISE must be set") + pid_m = tl.program_id(axis=0) + pid_n = tl.program_id(axis=1) + + M_padded = tl.load(padded_group_offs_ptr + num_groups) + block_start = pid_m * BLOCK_SIZE + if block_start >= M_padded: + return + + group_id = 0 + for g in range(num_groups): + padded_start = tl.load(padded_group_offs_ptr + g) + padded_end = tl.load(padded_group_offs_ptr + g + 1) + if block_start >= padded_start and block_start < padded_end: + group_id = g + + orig_group_start = tl.load(group_offs_ptr + group_id) + orig_group_end = tl.load(group_offs_ptr + group_id + 1) + padded_group_start = tl.load(padded_group_offs_ptr + group_id) + + offs_m_out = tl.cast(pid_m * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE), tl.int64) + offs_n = tl.cast(pid_n * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE), tl.int64) + offs_m_in = orig_group_start + (offs_m_out - padded_group_start) + + row_valid = (offs_m_in >= orig_group_start) & (offs_m_in < orig_group_end) + mask = row_valid[:, None] & (offs_n[None, :] < N) + + x_ptrs = x_ptr + offs_m_in[:, None] * N + offs_n[None, :] + x_tile = tl.load(x_ptrs, mask=mask, other=0.0).to(tl.float32) + x_tile_abs = tl.abs(x_tile) + + if COLUMNWISE: + # axis=0 -> segment-padded output positions. + x_fp8_col_tile, x_scales_col_tile = compute_scale_and_quant( + x_tile, x_tile_abs, 0, FP8_MAX, ROUND_POW2 + ) + x_fp8_col_ptrs = x_fp8_col_ptr + offs_m_out[:, None] * N + offs_n[None, :] + col_out_mask = (offs_m_out[:, None] < M_padded) & (offs_n[None, :] < N) + tl.store( + x_fp8_col_ptrs, x_fp8_col_tile.to(x_fp8_col_ptr.dtype.element_ty), mask=col_out_mask + ) + col_scale_offs = pid_m * N + offs_n + col_scale_mask = (pid_m < tl.cdiv(M_padded, BLOCK_SIZE)) & (offs_n < N) + x_scales_col_tile_inv = tl.reshape(1.0 / x_scales_col_tile, BLOCK_SIZE) + tl.store(x_scales_col_ptr + col_scale_offs, x_scales_col_tile_inv, mask=col_scale_mask) + + if ROWWISE: + # axis=1 -> original (unpadded) positions; padding rows are skipped. + x_fp8_row_tile, x_scales_row_tile = compute_scale_and_quant( + x_tile, x_tile_abs, 1, FP8_MAX, ROUND_POW2 + ) + x_fp8_row_ptrs = x_fp8_row_ptr + offs_m_in[:, None] * N + offs_n[None, :] + tl.store(x_fp8_row_ptrs, x_fp8_row_tile.to(x_fp8_row_ptr.dtype.element_ty), mask=mask) + row_scale_offs = offs_m_in * tl.cdiv(N, BLOCK_SIZE) + pid_n + x_scales_row_tile_inv = tl.reshape(1.0 / x_scales_row_tile, BLOCK_SIZE) + tl.store(x_scales_row_ptr + row_scale_offs, x_scales_row_tile_inv, mask=row_valid) + + +# ----------------------------------------------------------------------------- +# Host launchers +# +# Plain functions, not ``torch.library.custom_op``. A custom_op wrapper may be +# needed if torch.compile hits the inductor ``identify_mutated_tensors`` bug +# seen on gfx942 + Triton 3.7. +# ----------------------------------------------------------------------------- + + +def quantize_fp8_blockwise( + x: torch.Tensor, dtype: torch.dtype, axis: int, block_size: int = 128, pow2: bool = False +): + """Single-direction blockwise quantize. axis=1 -> rowwise (1xB), axis=0 -> colwise (Bx1).""" + assert x.is_contiguous() and x.dim() == 2, "Input must be 2D and contiguous" + M, N = x.shape + fp8_max = torch.finfo(dtype).max + x_fp8 = torch.empty((M, N), dtype=dtype, device=x.device) + if axis == 1: + scales = torch.empty((M, triton.cdiv(N, block_size)), dtype=torch.float32, device=x.device) + else: + scales = torch.empty((triton.cdiv(M, block_size), N), dtype=torch.float32, device=x.device) + + grid = (triton.cdiv(M, block_size), triton.cdiv(N, block_size)) + quant_fp8_blockwise_kernel[grid]( + x, + x_fp8, + scales, + M, + N, + BLOCK_SIZE=block_size, + FP8_MAX=fp8_max, + AXIS=axis, + ROUND_POW2=pow2, + ) + return x_fp8, scales + + +def quantize_fp8_blockwise_weight( + w: torch.Tensor, dtype: torch.dtype, block_size: int = 128, pow2: bool = False +): + """128x128 weight blockwise quantize. w is [B, M, N] (or [M, N], promoted to B=1).""" + squeeze = False + if w.dim() == 2: + w = w.unsqueeze(0) + squeeze = True + assert w.is_contiguous() and w.dim() == 3, "Weight must be 3D [B,M,N] and contiguous" + B, M, N = w.shape + fp8_max = torch.finfo(dtype).max + + w_fp8 = torch.empty((B, M, N), dtype=dtype, device=w.device) + w_scales = torch.empty( + (B, triton.cdiv(M, block_size), triton.cdiv(N, block_size)), + dtype=torch.float32, + device=w.device, + ) + grid = (B, triton.cdiv(M, block_size), triton.cdiv(N, block_size)) + quant_fp8_blockwise_for_weight_kernel[grid]( + w, + w_fp8, + w_scales, + M, + N, + BLOCK_SIZE=block_size, + FP8_MAX=fp8_max, + ROUND_POW2=pow2, + ) + if squeeze: + return w_fp8.squeeze(0), w_scales.squeeze(0) + return w_fp8, w_scales + + +def quantize_fp8_blockwise_segment_m( + x, dtype, block_size, group_lens, group_offs, pow2: bool = False +): + """Colwise blockwise quantize with per-segment (MoE group) M padding. + + Returns (x_fp8 [M_pad, N], x_scales, var_k_group_lens [B], var_k_group_offs [B+1]). + Each group is padded up to a multiple of block_size. Allocates with an upper + bound (M + B*block_size) to avoid a device->host sync (graph-capture safe). + """ + assert x.is_contiguous() and x.dim() == 2, "Input must be 2D and contiguous" + M, N = x.shape + num_groups = group_lens.size(0) + fp8_max = torch.finfo(dtype).max + + var_k_group_lens = ((group_lens + block_size - 1) // block_size) * block_size + var_k_group_offs = torch.zeros(num_groups + 1, dtype=torch.int64, device=x.device) + var_k_group_offs[1:] = torch.cumsum(var_k_group_lens, dim=0) + + m_padded_max = M + num_groups * block_size + x_fp8 = torch.zeros((m_padded_max, N), dtype=dtype, device=x.device) + x_scales = torch.zeros( + (triton.cdiv(m_padded_max, block_size), N), dtype=torch.float32, device=x.device + ) + + grid = (triton.cdiv(m_padded_max, block_size), triton.cdiv(N, block_size)) + # Columnwise-only: the rowwise pointer args are unused (compiled out) and are + # given the colwise buffers as harmless placeholders. + quant_fp8_blockwise_grouped_kernel[grid]( + x, + x_fp8, + x_scales, + x_fp8, + x_scales, + group_offs, + var_k_group_offs, + N, + num_groups, + BLOCK_SIZE=block_size, + FP8_MAX=fp8_max, + ROUND_POW2=pow2, + ROWWISE=False, + COLUMNWISE=True, + ) + return x_fp8, x_scales, var_k_group_lens, var_k_group_offs + + +def _torch_fp8_dtype(te_fp8_dtype): + """Arch-correct torch FP8 dtype for a TE ``DType`` (kernel launch / ``.view``).""" + return te_dtype_to_torch_dtype(tex.DType(int(te_fp8_dtype))) + + +def _make_blockwise_qtensor( + x_fp8, + scale_inv, + *, + orig_shape, + orig_dtype, + te_fp8_dtype, + is_2D_scaled, + pow2, + columnwise_data=None, + columnwise_scale_inv=None, + vk_group_offs=None, +): + """Wrap raw Triton blockwise-quant outputs in a ``Float8BlockwiseQTensor``. + + ``scale_inv`` (rowwise) must already be in the TE (ROCm) GEMM layout: + ``[ceil(M/128), ceil(N/128)]`` for 2D (weight) scaling and ``[ceil(K/128), M]`` + for 1D (rowwise activation) scaling. + + The optional columnwise operand carries the **segment-padded** wgrad data + ``[M_pad, N]`` and its 1D block scales ``[ceil(M_pad/128), N]`` -- note this is + *not* the standard transposed ``[N, M]`` columnwise layout; the grouped + variable-K wgrad kernel consumes it together with ``vk_group_offs`` (the padded + per-segment offsets). Either operand may be omitted (``None``) when only the + other is needed. + """ + from ..tensor.float8_blockwise_tensor import Float8BlockwiseQTensor, Float8BlockQuantizer + + rowwise = x_fp8 is not None + columnwise = columnwise_data is not None + quantizer = Float8BlockQuantizer( + fp8_dtype=te_fp8_dtype, + rowwise=rowwise, + columnwise=columnwise, + block_scaling_dim=2 if is_2D_scaled else 1, + force_pow_2_scales=pow2, + ) + ref = x_fp8 if rowwise else columnwise_data + return Float8BlockwiseQTensor( + shape=tuple(orig_shape), + dtype=orig_dtype, + fp8_dtype=te_fp8_dtype, + rowwise_data=x_fp8.view(torch.uint8) if rowwise else None, + rowwise_scale_inv=scale_inv, + columnwise_data=columnwise_data.view(torch.uint8) if columnwise else None, + columnwise_scale_inv=columnwise_scale_inv, + quantizer=quantizer, + is_2D_scaled=is_2D_scaled, + device=ref.device, + vk_group_offs=vk_group_offs, + ) + + +def quantize_fp8_blockwise_act_qtensor(a, fp8_dtype, pow2: bool = False): + """1x128 rowwise activation quantize returning a ``Float8BlockwiseQTensor``. + + ``fp8_dtype`` is a TE ``DType`` (the QTensor's native dtype); it is mapped to + the arch-correct torch FP8 dtype only for the Triton launch. + """ + x_fp8, scales = quantize_fp8_blockwise( + a, _torch_fp8_dtype(fp8_dtype), axis=1, block_size=128, pow2=pow2 + ) + # Triton emits scales as [M, ceil(K/128)]; TE stores the 1D rowwise scale + # transposed as [ceil(K/128), M] (the GEMM-ready layout). + return _make_blockwise_qtensor( + x_fp8, + scales.t().contiguous(), + orig_shape=a.shape, + orig_dtype=a.dtype, + te_fp8_dtype=fp8_dtype, + is_2D_scaled=False, + pow2=pow2, + ) + + +def wrap_fp8_blockwise_segment_m_qtensor( + data, scale_inv, vk_group_offs, orig_dtype, fp8_dtype, pow2: bool = False +): + """Wrap segment-padded columnwise FP8 fields as a columnwise-only QTensor. + + The segment-padded wgrad operand lives in the ``columnwise_data`` / + ``columnwise_scale_inv`` slots (``[M_pad, N]`` data, ``[ceil(M_pad/128), N]`` + scales) with the padded per-segment offsets in ``_vk_group_offs``. The rowwise + operand is left ``None``; only the variable-K wgrad kernel reads this tensor. + """ + return _make_blockwise_qtensor( + None, + None, + orig_shape=tuple(data.shape), + orig_dtype=orig_dtype, + te_fp8_dtype=fp8_dtype, + is_2D_scaled=False, + pow2=pow2, + columnwise_data=data, + columnwise_scale_inv=scale_inv, + vk_group_offs=vk_group_offs, + ) + + +def quantize_fp8_blockwise_segment_m_operand( + x, fp8_dtype, block_size, group_lens, group_offs, pow2: bool = False +): + """Segment-padded columnwise quantize returning a columnwise-only QTensor. + + ``fp8_dtype`` is a TE ``DType``; the torch FP8 dtype is derived only for the + Triton launch. + """ + x_fp8, x_scales, _vk_lens, vk_offs = quantize_fp8_blockwise_segment_m( + x, _torch_fp8_dtype(fp8_dtype), block_size, group_lens, group_offs, pow2=pow2 + ) + return wrap_fp8_blockwise_segment_m_qtensor( + x_fp8, x_scales, vk_offs, x.dtype, fp8_dtype, pow2=pow2 + ) + + +def quantize_fp8_blockwise_act_and_segment_m( + a, fp8_dtype, block_size, group_lens, group_offs, pow2: bool = False +): + """Fused rowwise activation + segment-padded columnwise quantize in one pass. + + Reads ``a`` ``[M, K]`` from HBM once and returns a single + ``Float8BlockwiseQTensor`` carrying **both** operands the grouped + forward/backward consume: + + * rowwise (1x128 along K) in ``rowwise_data`` with the scale stored transposed + ``[ceil(K/128), M]`` (the forward/dgrad GEMM layout), identical to + :func:`quantize_fp8_blockwise_act_qtensor`. + * segment-padded columnwise (1x128 along M, each MoE group padded to + ``block_size``) in ``columnwise_data`` + ``_vk_group_offs`` (the variable-K + wgrad operand), identical to :func:`quantize_fp8_blockwise_segment_m_operand`. + + Both share the same ``pow2`` rounding and source ``a``; this replaces the two + separate passes (rowwise + segment-padded colwise) with a single launch. Only + worthwhile when the wgrad operand is needed (training). + """ + assert a.is_contiguous() and a.dim() == 2, "Input must be 2D and contiguous" + M, N = a.shape + num_groups = group_lens.size(0) + torch_fp8_dtype = _torch_fp8_dtype(fp8_dtype) + fp8_max = torch.finfo(torch_fp8_dtype).max + + var_k_group_lens = ((group_lens + block_size - 1) // block_size) * block_size + var_k_group_offs = torch.zeros(num_groups + 1, dtype=torch.int64, device=a.device) + var_k_group_offs[1:] = torch.cumsum(var_k_group_lens, dim=0) + + x_fp8_row = torch.empty((M, N), dtype=torch_fp8_dtype, device=a.device) + x_scales_row = torch.empty( + (M, triton.cdiv(N, block_size)), dtype=torch.float32, device=a.device + ) + + m_padded_max = M + num_groups * block_size + x_fp8_col = torch.zeros((m_padded_max, N), dtype=torch_fp8_dtype, device=a.device) + x_scales_col = torch.zeros( + (triton.cdiv(m_padded_max, block_size), N), dtype=torch.float32, device=a.device + ) + + grid = (triton.cdiv(m_padded_max, block_size), triton.cdiv(N, block_size)) + quant_fp8_blockwise_grouped_kernel[grid]( + a, + x_fp8_row, + x_scales_row, + x_fp8_col, + x_scales_col, + group_offs, + var_k_group_offs, + N, + num_groups, + BLOCK_SIZE=block_size, + FP8_MAX=fp8_max, + ROUND_POW2=pow2, + ROWWISE=True, + COLUMNWISE=True, + ) + + # Triton emits rowwise scales as [M, ceil(K/128)]; TE stores the 1D rowwise + # scale transposed as [ceil(K/128), M] (the GEMM-ready layout). The + # segment-padded columnwise operand rides along in the columnwise slots. + return _make_blockwise_qtensor( + x_fp8_row, + x_scales_row.t().contiguous(), + orig_shape=a.shape, + orig_dtype=a.dtype, + te_fp8_dtype=fp8_dtype, + is_2D_scaled=False, + pow2=pow2, + columnwise_data=x_fp8_col, + columnwise_scale_inv=x_scales_col, + vk_group_offs=var_k_group_offs, + ) + + +def quantize_fp8_blockwise_act_operands( + a, + fp8_dtype, + block_size, + group_lens, + group_offs, + *, + rowwise, + columnwise, + pow2: bool = False, +): + """Quantize a grouped activation into the operand(s) the caller needs. + + Dispatches over the (``rowwise``, ``columnwise``) request: + + * both -> :func:`quantize_fp8_blockwise_act_and_segment_m` (one fused HBM read). + * rowwise only -> :func:`quantize_fp8_blockwise_act_qtensor` (forward/dgrad). + * columnwise only -> :func:`quantize_fp8_blockwise_segment_m_operand` (wgrad). + * neither -> ``None``. + + ``block_size`` / ``group_lens`` / ``group_offs`` are only consumed by the + segment-padded columnwise path. + """ + if rowwise and columnwise: + return quantize_fp8_blockwise_act_and_segment_m( + a, fp8_dtype, block_size, group_lens, group_offs, pow2=pow2 + ) + if rowwise: + return quantize_fp8_blockwise_act_qtensor(a, fp8_dtype, pow2=pow2) + if columnwise: + return quantize_fp8_blockwise_segment_m_operand( + a, fp8_dtype, block_size, group_lens, group_offs, pow2=pow2 + ) + return None + + +def wrap_fp8_blockwise_grouped_weight_qtensor( + rowwise_data, rowwise_scale_inv, orig_dtype, fp8_dtype, pow2: bool = True +): + """Wrap packed grouped-weight FP8 data + scales as a 2D ``[G*N, K]`` QTensor. + + Because ``N`` (out_features) is a multiple of 128, expert boundaries align + with the 128x128 blocks, so a packed ``[G, N, K]`` weight is bit-identical to + a single 2D ``[G*N, K]`` blockwise tensor with the scale flattened to + ``[G*(N/128), K/128]`` -- and it dequantizes correctly. ``rowwise_data`` may + be ``[G, N, K]`` or already 2D ``[G*N, K]`` (scales correspondingly 3D or 2D). + """ + if rowwise_data.dim() == 3: + gg, nn, kk = rowwise_data.shape + rowwise_data = rowwise_data.reshape(gg * nn, kk) + rowwise_scale_inv = rowwise_scale_inv.reshape( + gg * rowwise_scale_inv.shape[1], rowwise_scale_inv.shape[2] + ) + return _make_blockwise_qtensor( + rowwise_data, + rowwise_scale_inv, + orig_shape=tuple(rowwise_data.shape), + orig_dtype=orig_dtype, + te_fp8_dtype=fp8_dtype, + is_2D_scaled=True, + pow2=pow2, + ) + + +def quantize_fp8_blockwise_grouped_weight_qtensor(w, fp8_dtype, pow2: bool = False): + """128x128 grouped-weight quantize of ``[G, N, K]`` -> 2D ``[G*N, K]`` QTensor. + + ``N`` must be a multiple of 128 so expert boundaries align with 128-blocks + (see :func:`wrap_fp8_blockwise_grouped_weight_qtensor`). One packed quantize + launch over ``[G, N, K]``; no per-expert launches, no stack. + """ + assert w.dim() == 3, "quantize_fp8_blockwise_grouped_weight_qtensor expects [G, N, K]" + assert w.shape[1] % 128 == 0, "grouped weight N (out_features) must be a multiple of 128" + x_fp8, scales = quantize_fp8_blockwise_weight( + w, _torch_fp8_dtype(fp8_dtype), block_size=128, pow2=pow2 + ) + return wrap_fp8_blockwise_grouped_weight_qtensor(x_fp8, scales, w.dtype, fp8_dtype, pow2=pow2)