From 7bdd9a6529152be39e422172ad5953a51c4d75f3 Mon Sep 17 00:00:00 2001 From: Masaki Kozuki Date: Tue, 18 Nov 2025 11:22:13 +0900 Subject: [PATCH 1/8] add `F.scaled_grouped_mm` Signed-off-by: Masaki Kozuki --- thunder/executors/torchex.py | 31 +++++++++++++ thunder/tests/test_ops.py | 75 +++++++++++++++++++++++++++++++ thunder/torch/__init__.py | 87 ++++++++++++++++++++++++++++++++++++ 3 files changed, 193 insertions(+) diff --git a/thunder/executors/torchex.py b/thunder/executors/torchex.py index d6fb0b31c1..e76c11d3f3 100644 --- a/thunder/executors/torchex.py +++ b/thunder/executors/torchex.py @@ -1531,6 +1531,8 @@ def _copy_with_setitem_impl(a, key, value): baddbmm = _register_torch_operation("baddbmm") if LooseVersion(torch.__version__) >= "2.8": _grouped_mm = _register_torch_operation("_grouped_mm") +if _scaled_grouped_mm_available := hasattr(torch.nn.functional, "scaled_grouped_mm"): + scaled_grouped_mm = _register_torch_operation("scaled_grouped_mm", module=torch.nn.functional) convolution = _register_torch_operation("convolution") conv1d = _register_torch_operation("conv1d", module=torch.nn.functional) conv2d = _register_torch_operation("conv2d", module=torch.nn.functional) @@ -1827,11 +1829,40 @@ def _grouped_mm_checker(a: TensorProxy, b: TensorProxy, offsets: TensorProxy) -> return a.dtype == dtypes.bfloat16 and b.dtype == dtypes.bfloat16 and offsets.dtype == dtypes.int32 +def _scaled_grouped_mm_checker( + mat_a: TensorProxy, + mat_b: TensorProxy, + scale_a, + scale_recipe_a, + scale_b, + scale_recipe_b, + swizzle_a=None, + swizzle_b=None, + bias: None | TensorProxy = None, + offs: None | TensorProxy = None, + output_dtype: None | dtypeLike = None, + contraction_dim: Sequence[int] | tuple[int, ...] = (), + use_fast_accum: bool = False, +) -> bool: + if offs is None: + return False + if isinstance(offs, TensorProxy): + return utils.is_integer_dtype(offs.dtype) + return True + + _register_implementation(ltorch.baddbmm, baddbmm, checker=_always_executable) _register_implementation(ltorch.bmm, bmm, checker=_always_executable) if LooseVersion(torch.__version__) >= "2.8": _register_implementation(prims._grouped_mm, _grouped_mm, checker=_grouped_mm_checker) _register_implementation(ltorch._grouped_mm, _grouped_mm, checker=_grouped_mm_checker) + +if _scaled_grouped_mm_available: + _register_implementation( + ltorch.scaled_grouped_mm, + scaled_grouped_mm, + checker=_scaled_grouped_mm_checker, + ) _register_implementation(ltorch.convolution, checker=_always_executable, execution_transform=_convolution_transform) _register_implementation(ltorch.conv1d, conv1d, checker=_always_executable) _register_implementation(ltorch.conv2d, conv2d, checker=_always_executable) diff --git a/thunder/tests/test_ops.py b/thunder/tests/test_ops.py index 62479e1223..dbcaa6e34c 100644 --- a/thunder/tests/test_ops.py +++ b/thunder/tests/test_ops.py @@ -4,6 +4,7 @@ import numpy as np import pytest import torch +from torch.nn.functional import ScalingType from torch.testing import assert_close import thunder @@ -419,6 +420,80 @@ def fn(a): assert_close(b, b_ref) +# NOTE: The following tests exercise torch.nn.functional.scaled_grouped_mm via Thunder. +# They validate that the Thunder tracing mirrors eager execution for representative 2D/2D +# and 2D/3D grouped matmul shapes that correspond to the accepted combinations in +# thunder.core.prims._grouped_mm_meta. These scenarios mirror the small tensor smoke tests +# in PyTorch's scaled matmul CUDA suite ([pytorch/test_scaled_matmul_cuda.py](https://github.com/pytorch/pytorch/blob/main/test/test_scaled_matmul_cuda.py)). +@requiresCUDA +def test_scaled_grouped_mm_2d2d_tensorwise(): + device = "cuda" + group_sizes = [8, 12] + m = 6 + n = 10 + k = sum(group_sizes) + + mat_a = torch.randn(m, k, device=device, dtype=torch.float16) + mat_b = torch.randn(k, n, device=device, dtype=torch.float16) + offs = torch.tensor(group_sizes, device=device, dtype=torch.int32).cumsum(0) + scale_a = torch.ones(1, device=device, dtype=torch.float32) + scale_b = torch.ones(1, device=device, dtype=torch.float32) + + def fn(a, b, scale_a, scale_b, offs): + return torch.nn.functional.scaled_grouped_mm( + a, + b, + scale_a, + ScalingType.TensorWise, + scale_b, + ScalingType.TensorWise, + offs=offs, + output_dtype=torch.bfloat16, + ) + + eager = fn(mat_a, mat_b, scale_a, scale_b, offs) + compiled = thunder.jit(fn) + result = compiled(mat_a, mat_b, scale_a, scale_b, offs) + + assert_close(result, eager) + assert_consistency_of_compiletime_and_runtime(compiled, result) + + +@requiresCUDA +def test_scaled_grouped_mm_2d3d_tensorwise(): + device = "cuda" + group_sizes = [5, 7] + groups = len(group_sizes) + total_rows = sum(group_sizes) + k = 16 + n = 9 + + mat_a = torch.randn(total_rows, k, device=device, dtype=torch.float16) + mat_b = torch.randn(groups, k, n, device=device, dtype=torch.float16) + offs = torch.tensor(group_sizes, device=device, dtype=torch.int32).cumsum(0) + scale_a = torch.ones(1, device=device, dtype=torch.float32) + scale_b = torch.ones(1, device=device, dtype=torch.float32) + + def fn(a, b, scale_a, scale_b, offs): + return torch.nn.functional.scaled_grouped_mm( + a, + b, + scale_a, + ScalingType.TensorWise, + scale_b, + ScalingType.TensorWise, + offs=offs, + output_dtype=torch.bfloat16, + ) + + eager = fn(mat_a, mat_b, scale_a, scale_b, offs) + compiled = thunder.jit(fn) + result = compiled(mat_a, mat_b, scale_a, scale_b, offs) + + assert_close(result, eager) + assert_consistency_of_compiletime_and_runtime(compiled, result) + + # https://github.com/Lightning-AI/lightning-thunder/issues/1857 def test_max_with_int(): def f(x, ids): diff --git a/thunder/torch/__init__.py b/thunder/torch/__init__.py index 090f46643e..b0db69ba44 100644 --- a/thunder/torch/__init__.py +++ b/thunder/torch/__init__.py @@ -5806,6 +5806,93 @@ def _grouped_mm( return prims._grouped_mm(a, b, offsets) +if hasattr(torch.nn.functional, "scaled_grouped_mm"): + + @torchsymbol( + torch.nn.functional.scaled_grouped_mm, + id="torch.nn.functional.scaled_grouped_mm", + is_method=False, + is_prim=True, + ) + def scaled_grouped_mm( + mat_a: TensorProxy, + mat_b: TensorProxy, + scale_a, + scale_recipe_a, + scale_b, + scale_recipe_b, + swizzle_a=None, + swizzle_b=None, + bias: None | TensorProxy = None, + offs: None | TensorProxy = None, + output_dtype: dtypeLike = torch.bfloat16, + contraction_dim: Sequence[int] | tuple[int, ...] = (), + use_fast_accum: bool = False, + ) -> TensorProxy: + utils.check( + offs is not None, + lambda: "scaled_grouped_mm currently requires `offs`.", + ) + utils.check_type(offs, TensorProxy) + utils.check( + bias is None, + lambda: "scaled_grouped_mm currently doesn't support `bias`.", + ) + utils.check( + len(contraction_dim) == 0, + lambda: f"scaled_grouped_mm currently expects an empty `contraction_dim`, but got {contraction_dim}.", + ) + + utils.check_type(mat_a, TensorProxy) + utils.check_type(mat_b, TensorProxy) + utils.check(mat_a.ndim in (2, 3), lambda: f"Expected mat_a to have 2 or 3 dimensions, got {mat_a.ndim}") + utils.check(mat_b.ndim in (2, 3), lambda: f"Expected mat_b to have 2 or 3 dimensions, got {mat_b.ndim}") + + utils.check(offs.ndim == 1, lambda: f"`offs` must be a vector, got shape {offs.shape}") + if mat_a.ndim == 2 and mat_b.ndim == 2: + utils.check( + mat_a.shape[1] == mat_b.shape[0], + lambda: f"Inner dimension mismatch: {mat_a.shape} vs {mat_b.shape}", + ) + out_shape = (offs.shape[0], mat_a.shape[0], mat_b.shape[1]) + elif mat_a.ndim == 3 and mat_b.ndim == 2: + utils.check( + mat_a.shape[2] == mat_b.shape[1], + lambda: f"Inner dimension mismatch: {mat_a.shape} vs {mat_b.shape}", + ) + utils.check( + mat_a.shape[0] == offs.shape[0], + lambda: f"Group count mismatch: {mat_a.shape} vs {offs.shape}", + ) + out_shape = (mat_a.shape[1], mat_b.shape[1]) + elif mat_a.ndim == 2 and mat_b.ndim == 3: + utils.check( + mat_a.shape[1] == mat_b.shape[1], + lambda: f"Inner dimension mismatch: {mat_a.shape} vs {mat_b.shape}", + ) + utils.check( + mat_b.shape[0] == offs.shape[0], + lambda: f"Group count mismatch: {mat_b.shape} vs {offs.shape}", + ) + out_shape = (mat_a.shape[0], mat_b.shape[2]) + else: + utils.check(False, lambda: f"Unexpected shape combination: {mat_a.shape} and {mat_b.shape}") + + utils.check_same_dtype(mat_a, mat_b) + utils.check( + mat_a.dtype in dtypes.float_math_dtypes, + lambda: f"`mat_a` must be 16-bit float or higher, got {mat_a.dtype}", + ) + utils.check( + utils.is_integer_dtype(offs.dtype), + lambda: f"`offs` must be integers, got {offs.dtype}", + ) + utils.check_same_device(mat_a, mat_b) + + target_dtype = to_dtype(output_dtype) if output_dtype is not None else mat_a.dtype + return TensorProxy(like=mat_a, shape=out_shape, dtype=target_dtype) + + @torchsymbol(torch.logsumexp, is_method=True) def logsumexp(a: TensorLike, /, dim: int | Sequence[int], keepdim: bool = False) -> TensorLike: input_max = amax(a, dim, keepdim=True) From c5d4212798968a969824a4b062b66f116e7f63ee Mon Sep 17 00:00:00 2001 From: Masaki Kozuki Date: Mon, 17 Nov 2025 19:37:44 -0800 Subject: [PATCH 2/8] add mxfp8 scaled grouped mm test Signed-off-by: Masaki Kozuki --- thunder/tests/test_ops.py | 83 ++++++++++++++++++++++++++++++++++++++- thunder/torch/__init__.py | 5 ++- 2 files changed, 85 insertions(+), 3 deletions(-) diff --git a/thunder/tests/test_ops.py b/thunder/tests/test_ops.py index dbcaa6e34c..e7701bb78a 100644 --- a/thunder/tests/test_ops.py +++ b/thunder/tests/test_ops.py @@ -4,8 +4,13 @@ import numpy as np import pytest import torch -from torch.nn.functional import ScalingType +from torch.nn.functional import ScalingType, SwizzleType from torch.testing import assert_close +from torch.testing._internal.common_cuda import ( + PLATFORM_SUPPORTS_FP8_GROUPED_GEMM, + PLATFORM_SUPPORTS_MXFP8_GROUPED_GEMM, +) +from torch.testing._internal.common_quantized import to_blocked, to_mxfp import thunder import thunder.core.devices as devices @@ -425,8 +430,14 @@ def fn(a): # and 2D/3D grouped matmul shapes that correspond to the accepted combinations in # thunder.core.prims._grouped_mm_meta. These scenarios mirror the small tensor smoke tests # in PyTorch's scaled matmul CUDA suite ([pytorch/test_scaled_matmul_cuda.py](https://github.com/pytorch/pytorch/blob/main/test/test_scaled_matmul_cuda.py)). +F8_GROUPED_MSG = "FP8 grouped is only supported on SM90 and MI300+ devices" +MXFP8_GROUPED_MSG = "MXFP8 grouped GEMM is only supported when PyTorch is built with USE_FBGEMM_GENAI=1 on SM100+" + + @requiresCUDA def test_scaled_grouped_mm_2d2d_tensorwise(): + if not bool(PLATFORM_SUPPORTS_FP8_GROUPED_GEMM): + pytest.skip(F8_GROUPED_MSG) device = "cuda" group_sizes = [8, 12] m = 6 @@ -461,6 +472,8 @@ def fn(a, b, scale_a, scale_b, offs): @requiresCUDA def test_scaled_grouped_mm_2d3d_tensorwise(): + if not bool(PLATFORM_SUPPORTS_FP8_GROUPED_GEMM): + pytest.skip(F8_GROUPED_MSG) device = "cuda" group_sizes = [5, 7] groups = len(group_sizes) @@ -494,6 +507,74 @@ def fn(a, b, scale_a, scale_b, offs): assert_consistency_of_compiletime_and_runtime(compiled, result) +@requiresCUDA +def test_scaled_grouped_mm_2d2d_mxfp8_blockwise(): + if not bool(PLATFORM_SUPPORTS_MXFP8_GROUPED_GEMM): + pytest.skip(MXFP8_GROUPED_MSG) + + # if dtypes.float8_e4m3fn not in dtypes.float_math_dtypes: + # pytest.skip("Thunder float8_e4m3fn inputs are not yet supported") + + device = "cuda" + torch.manual_seed(0) + + group_sizes = [64, 32] + m = 128 + n = 96 + total_k = sum(group_sizes) + + raw_offs = torch.tensor(group_sizes, device=device, dtype=torch.int32) + offs = torch.cumsum(raw_offs, dim=0, dtype=torch.int32) + + def round_up(x: int, y: int) -> int: + return ((x + y - 1) // y) * y + + def quantize_to_mxfp8(mat: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + segments: list[torch.Tensor] = [] + scale_segments: list[torch.Tensor] = [] + start = 0 + for end in offs.tolist(): + segment = mat[:, start:end].contiguous() + scale, lowp = to_mxfp(segment, format="mxfp8") + scale_segments.append(to_blocked(scale)) + segments.append(lowp) + start = end + lowp_full = torch.cat(segments, dim=1) + rows_rounded = round_up(mat.shape[0], 128) + scale_full = torch.cat(scale_segments, dim=0).reshape(rows_rounded, -1) + return lowp_full, scale_full + + mat_a_hp = torch.randn(m, total_k, device=device, dtype=torch.bfloat16) * 0.1 + mat_b_hp = torch.randn(n, total_k, device=device, dtype=torch.bfloat16) * 0.01 + + mat_a_lp, scale_a = quantize_to_mxfp8(mat_a_hp) + mat_b_lp_rows, scale_b = quantize_to_mxfp8(mat_b_hp) + mat_b_lp = mat_b_lp_rows.transpose(0, 1) + + swizzle = SwizzleType.SWIZZLE_32_4_4 if torch.version.cuda else SwizzleType.NO_SWIZZLE + + def fn(mat_a, mat_b, scale_a, scale_b, offs): + return torch.nn.functional.scaled_grouped_mm( + mat_a, + mat_b, + scale_a, + ScalingType.BlockWise1x32, + scale_b, + ScalingType.BlockWise1x32, + swizzle_a=swizzle, + swizzle_b=swizzle, + offs=offs, + output_dtype=torch.bfloat16, + ) + + eager = fn(mat_a_lp, mat_b_lp, scale_a, scale_b, offs) + jitted = thunder.jit(fn) + result = jitted(mat_a_lp, mat_b_lp, scale_a, scale_b, offs) + + assert_close(result, eager) + assert_consistency_of_compiletime_and_runtime(jitted, result) + + # https://github.com/Lightning-AI/lightning-thunder/issues/1857 def test_max_with_int(): def f(x, ids): diff --git a/thunder/torch/__init__.py b/thunder/torch/__init__.py index b0db69ba44..c84ccee8b6 100644 --- a/thunder/torch/__init__.py +++ b/thunder/torch/__init__.py @@ -5879,9 +5879,10 @@ def scaled_grouped_mm( utils.check(False, lambda: f"Unexpected shape combination: {mat_a.shape} and {mat_b.shape}") utils.check_same_dtype(mat_a, mat_b) + allowed_input_dtypes = dtypes.float_math_dtypes | dtypes.float_8bit_dtypes utils.check( - mat_a.dtype in dtypes.float_math_dtypes, - lambda: f"`mat_a` must be 16-bit float or higher, got {mat_a.dtype}", + mat_a.dtype in allowed_input_dtypes, + lambda: f"`mat_a` must be a floating dtype, got {mat_a.dtype}", ) utils.check( utils.is_integer_dtype(offs.dtype), From d1788bfc3eb70e19916b7161f478727346fc9aec Mon Sep 17 00:00:00 2001 From: Masaki Kozuki Date: Mon, 17 Nov 2025 19:49:41 -0800 Subject: [PATCH 3/8] guard `scaled_grouped_mm` Signed-off-by: Masaki Kozuki --- thunder/tests/test_ops.py | 300 +++++++++++++++++++------------------- 1 file changed, 149 insertions(+), 151 deletions(-) diff --git a/thunder/tests/test_ops.py b/thunder/tests/test_ops.py index e7701bb78a..e7bbc870ec 100644 --- a/thunder/tests/test_ops.py +++ b/thunder/tests/test_ops.py @@ -4,13 +4,16 @@ import numpy as np import pytest import torch -from torch.nn.functional import ScalingType, SwizzleType -from torch.testing import assert_close -from torch.testing._internal.common_cuda import ( - PLATFORM_SUPPORTS_FP8_GROUPED_GEMM, - PLATFORM_SUPPORTS_MXFP8_GROUPED_GEMM, -) -from torch.testing._internal.common_quantized import to_blocked, to_mxfp + +HAS_SCALED_GROUPED_MM = hasattr(torch.nn.functional, "scaled_grouped_mm") + +if HAS_SCALED_GROUPED_MM: + from torch.nn.functional import ScalingType, SwizzleType + from torch.testing._internal.common_cuda import ( + PLATFORM_SUPPORTS_FP8_GROUPED_GEMM, + PLATFORM_SUPPORTS_MXFP8_GROUPED_GEMM, + ) + from torch.testing._internal.common_quantized import to_blocked, to_mxfp import thunder import thunder.core.devices as devices @@ -425,154 +428,149 @@ def fn(a): assert_close(b, b_ref) -# NOTE: The following tests exercise torch.nn.functional.scaled_grouped_mm via Thunder. -# They validate that the Thunder tracing mirrors eager execution for representative 2D/2D -# and 2D/3D grouped matmul shapes that correspond to the accepted combinations in -# thunder.core.prims._grouped_mm_meta. These scenarios mirror the small tensor smoke tests -# in PyTorch's scaled matmul CUDA suite ([pytorch/test_scaled_matmul_cuda.py](https://github.com/pytorch/pytorch/blob/main/test/test_scaled_matmul_cuda.py)). -F8_GROUPED_MSG = "FP8 grouped is only supported on SM90 and MI300+ devices" -MXFP8_GROUPED_MSG = "MXFP8 grouped GEMM is only supported when PyTorch is built with USE_FBGEMM_GENAI=1 on SM100+" - - -@requiresCUDA -def test_scaled_grouped_mm_2d2d_tensorwise(): - if not bool(PLATFORM_SUPPORTS_FP8_GROUPED_GEMM): - pytest.skip(F8_GROUPED_MSG) - device = "cuda" - group_sizes = [8, 12] - m = 6 - n = 10 - k = sum(group_sizes) - - mat_a = torch.randn(m, k, device=device, dtype=torch.float16) - mat_b = torch.randn(k, n, device=device, dtype=torch.float16) - offs = torch.tensor(group_sizes, device=device, dtype=torch.int32).cumsum(0) - scale_a = torch.ones(1, device=device, dtype=torch.float32) - scale_b = torch.ones(1, device=device, dtype=torch.float32) - - def fn(a, b, scale_a, scale_b, offs): - return torch.nn.functional.scaled_grouped_mm( - a, - b, - scale_a, - ScalingType.TensorWise, - scale_b, - ScalingType.TensorWise, - offs=offs, - output_dtype=torch.bfloat16, - ) - - eager = fn(mat_a, mat_b, scale_a, scale_b, offs) - compiled = thunder.jit(fn) - result = compiled(mat_a, mat_b, scale_a, scale_b, offs) - - assert_close(result, eager) - assert_consistency_of_compiletime_and_runtime(compiled, result) - - -@requiresCUDA -def test_scaled_grouped_mm_2d3d_tensorwise(): - if not bool(PLATFORM_SUPPORTS_FP8_GROUPED_GEMM): - pytest.skip(F8_GROUPED_MSG) - device = "cuda" - group_sizes = [5, 7] - groups = len(group_sizes) - total_rows = sum(group_sizes) - k = 16 - n = 9 - - mat_a = torch.randn(total_rows, k, device=device, dtype=torch.float16) - mat_b = torch.randn(groups, k, n, device=device, dtype=torch.float16) - offs = torch.tensor(group_sizes, device=device, dtype=torch.int32).cumsum(0) - scale_a = torch.ones(1, device=device, dtype=torch.float32) - scale_b = torch.ones(1, device=device, dtype=torch.float32) - - def fn(a, b, scale_a, scale_b, offs): - return torch.nn.functional.scaled_grouped_mm( - a, - b, - scale_a, - ScalingType.TensorWise, - scale_b, - ScalingType.TensorWise, - offs=offs, - output_dtype=torch.bfloat16, - ) - - eager = fn(mat_a, mat_b, scale_a, scale_b, offs) - compiled = thunder.jit(fn) - result = compiled(mat_a, mat_b, scale_a, scale_b, offs) - - assert_close(result, eager) - assert_consistency_of_compiletime_and_runtime(compiled, result) +if HAS_SCALED_GROUPED_MM: + # NOTE: The following tests exercise torch.nn.functional.scaled_grouped_mm via Thunder. + # They validate that the Thunder tracing mirrors eager execution for representative 2D/2D + # and 2D/3D grouped matmul shapes that correspond to the accepted combinations in + # thunder.core.prims._grouped_mm_meta. These scenarios mirror the small tensor smoke tests + # in PyTorch's scaled matmul CUDA suite ([pytorch/test_scaled_matmul_cuda.py](https://github.com/pytorch/pytorch/blob/main/test/test_scaled_matmul_cuda.py)). + F8_GROUPED_MSG = "FP8 grouped is only supported on SM90 and MI300+ devices" + MXFP8_GROUPED_MSG = "MXFP8 grouped GEMM is only supported when PyTorch is built with USE_FBGEMM_GENAI=1 on SM100+" + + @requiresCUDA + def test_scaled_grouped_mm_2d2d_tensorwise(): + if not bool(PLATFORM_SUPPORTS_FP8_GROUPED_GEMM): + pytest.skip(F8_GROUPED_MSG) + device = "cuda" + group_sizes = [8, 12] + m = 6 + n = 10 + k = sum(group_sizes) + + mat_a = torch.randn(m, k, device=device, dtype=torch.float16) + mat_b = torch.randn(k, n, device=device, dtype=torch.float16) + offs = torch.tensor(group_sizes, device=device, dtype=torch.int32).cumsum(0) + scale_a = torch.ones(1, device=device, dtype=torch.float32) + scale_b = torch.ones(1, device=device, dtype=torch.float32) + + def fn(a, b, scale_a, scale_b, offs): + return torch.nn.functional.scaled_grouped_mm( + a, + b, + scale_a, + ScalingType.TensorWise, + scale_b, + ScalingType.TensorWise, + offs=offs, + output_dtype=torch.bfloat16, + ) + eager = fn(mat_a, mat_b, scale_a, scale_b, offs) + jitted = thunder.jit(fn) + result = jitted(mat_a, mat_b, scale_a, scale_b, offs) + + torch.testing.assert_close(result, eager) + assert_consistency_of_compiletime_and_runtime(jitted, result) + + @requiresCUDA + def test_scaled_grouped_mm_2d3d_tensorwise(): + if not bool(PLATFORM_SUPPORTS_FP8_GROUPED_GEMM): + pytest.skip(F8_GROUPED_MSG) + device = "cuda" + group_sizes = [5, 7] + groups = len(group_sizes) + total_rows = sum(group_sizes) + k = 16 + n = 9 + + mat_a = torch.randn(total_rows, k, device=device, dtype=torch.float16) + mat_b = torch.randn(groups, k, n, device=device, dtype=torch.float16) + offs = torch.tensor(group_sizes, device=device, dtype=torch.int32).cumsum(0) + scale_a = torch.ones(1, device=device, dtype=torch.float32) + scale_b = torch.ones(1, device=device, dtype=torch.float32) + + def fn(a, b, scale_a, scale_b, offs): + return torch.nn.functional.scaled_grouped_mm( + a, + b, + scale_a, + ScalingType.TensorWise, + scale_b, + ScalingType.TensorWise, + offs=offs, + output_dtype=torch.bfloat16, + ) -@requiresCUDA -def test_scaled_grouped_mm_2d2d_mxfp8_blockwise(): - if not bool(PLATFORM_SUPPORTS_MXFP8_GROUPED_GEMM): - pytest.skip(MXFP8_GROUPED_MSG) - - # if dtypes.float8_e4m3fn not in dtypes.float_math_dtypes: - # pytest.skip("Thunder float8_e4m3fn inputs are not yet supported") - - device = "cuda" - torch.manual_seed(0) - - group_sizes = [64, 32] - m = 128 - n = 96 - total_k = sum(group_sizes) - - raw_offs = torch.tensor(group_sizes, device=device, dtype=torch.int32) - offs = torch.cumsum(raw_offs, dim=0, dtype=torch.int32) - - def round_up(x: int, y: int) -> int: - return ((x + y - 1) // y) * y - - def quantize_to_mxfp8(mat: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: - segments: list[torch.Tensor] = [] - scale_segments: list[torch.Tensor] = [] - start = 0 - for end in offs.tolist(): - segment = mat[:, start:end].contiguous() - scale, lowp = to_mxfp(segment, format="mxfp8") - scale_segments.append(to_blocked(scale)) - segments.append(lowp) - start = end - lowp_full = torch.cat(segments, dim=1) - rows_rounded = round_up(mat.shape[0], 128) - scale_full = torch.cat(scale_segments, dim=0).reshape(rows_rounded, -1) - return lowp_full, scale_full - - mat_a_hp = torch.randn(m, total_k, device=device, dtype=torch.bfloat16) * 0.1 - mat_b_hp = torch.randn(n, total_k, device=device, dtype=torch.bfloat16) * 0.01 - - mat_a_lp, scale_a = quantize_to_mxfp8(mat_a_hp) - mat_b_lp_rows, scale_b = quantize_to_mxfp8(mat_b_hp) - mat_b_lp = mat_b_lp_rows.transpose(0, 1) - - swizzle = SwizzleType.SWIZZLE_32_4_4 if torch.version.cuda else SwizzleType.NO_SWIZZLE - - def fn(mat_a, mat_b, scale_a, scale_b, offs): - return torch.nn.functional.scaled_grouped_mm( - mat_a, - mat_b, - scale_a, - ScalingType.BlockWise1x32, - scale_b, - ScalingType.BlockWise1x32, - swizzle_a=swizzle, - swizzle_b=swizzle, - offs=offs, - output_dtype=torch.bfloat16, - ) + eager = fn(mat_a, mat_b, scale_a, scale_b, offs) + jitted = thunder.jit(fn) + result = jitted(mat_a, mat_b, scale_a, scale_b, offs) + + torch.testing.assert_close(result, eager) + assert_consistency_of_compiletime_and_runtime(jitted, result) + + @requiresCUDA + def test_scaled_grouped_mm_2d2d_mxfp8_blockwise(): + if not bool(PLATFORM_SUPPORTS_MXFP8_GROUPED_GEMM): + pytest.skip(MXFP8_GROUPED_MSG) + + device = "cuda" + torch.manual_seed(0) + + group_sizes = [64, 32] + m = 128 + n = 96 + total_k = sum(group_sizes) + + raw_offs = torch.tensor(group_sizes, device=device, dtype=torch.int32) + offs = torch.cumsum(raw_offs, dim=0, dtype=torch.int32) + + def round_up(x: int, y: int) -> int: + return ((x + y - 1) // y) * y + + def quantize_to_mxfp8(mat: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + segments: list[torch.Tensor] = [] + scale_segments: list[torch.Tensor] = [] + start = 0 + for end in offs.tolist(): + segment = mat[:, start:end].contiguous() + scale, lowp = to_mxfp(segment, format="mxfp8") + scale_segments.append(to_blocked(scale)) + segments.append(lowp) + start = end + lowp_full = torch.cat(segments, dim=1) + rows_rounded = round_up(mat.shape[0], 128) + scale_full = torch.cat(scale_segments, dim=0).reshape(rows_rounded, -1) + return lowp_full, scale_full + + mat_a_hp = torch.randn(m, total_k, device=device, dtype=torch.bfloat16) * 0.1 + mat_b_hp = torch.randn(n, total_k, device=device, dtype=torch.bfloat16) * 0.01 + + mat_a_lp, scale_a = quantize_to_mxfp8(mat_a_hp) + mat_b_lp_rows, scale_b = quantize_to_mxfp8(mat_b_hp) + mat_b_lp = mat_b_lp_rows.transpose(0, 1) + + swizzle = SwizzleType.SWIZZLE_32_4_4 if torch.version.cuda else SwizzleType.NO_SWIZZLE + + def fn(mat_a, mat_b, scale_a, scale_b, offs): + return torch.nn.functional.scaled_grouped_mm( + mat_a, + mat_b, + scale_a, + ScalingType.BlockWise1x32, + scale_b, + ScalingType.BlockWise1x32, + swizzle_a=swizzle, + swizzle_b=swizzle, + offs=offs, + output_dtype=torch.bfloat16, + ) - eager = fn(mat_a_lp, mat_b_lp, scale_a, scale_b, offs) - jitted = thunder.jit(fn) - result = jitted(mat_a_lp, mat_b_lp, scale_a, scale_b, offs) + eager = fn(mat_a_lp, mat_b_lp, scale_a, scale_b, offs) + jitted = thunder.jit(fn) + result = jitted(mat_a_lp, mat_b_lp, scale_a, scale_b, offs) - assert_close(result, eager) - assert_consistency_of_compiletime_and_runtime(jitted, result) + torch.testing.assert_close(result, eager) + assert_consistency_of_compiletime_and_runtime(jitted, result) # https://github.com/Lightning-AI/lightning-thunder/issues/1857 From 04baea1ec132e28eb50775ad085c81b9ce8c63d6 Mon Sep 17 00:00:00 2001 From: Masaki Date: Tue, 18 Nov 2025 13:07:11 +0900 Subject: [PATCH 4/8] Import assert_close from torch.testing --- thunder/tests/test_ops.py | 1 + 1 file changed, 1 insertion(+) diff --git a/thunder/tests/test_ops.py b/thunder/tests/test_ops.py index e7bbc870ec..03bc14a130 100644 --- a/thunder/tests/test_ops.py +++ b/thunder/tests/test_ops.py @@ -4,6 +4,7 @@ import numpy as np import pytest import torch +from torch.testing import assert_close HAS_SCALED_GROUPED_MM = hasattr(torch.nn.functional, "scaled_grouped_mm") From 9239321ed873dee6cd1fb58108e874c89898a458 Mon Sep 17 00:00:00 2001 From: Masaki Kozuki Date: Mon, 17 Nov 2025 23:01:36 -0800 Subject: [PATCH 5/8] fix test Signed-off-by: Masaki Kozuki --- thunder/tests/test_ops.py | 49 ++++++++++++++++++++------------------- 1 file changed, 25 insertions(+), 24 deletions(-) diff --git a/thunder/tests/test_ops.py b/thunder/tests/test_ops.py index 03bc14a130..9cc0ff319f 100644 --- a/thunder/tests/test_ops.py +++ b/thunder/tests/test_ops.py @@ -439,29 +439,30 @@ def fn(a): MXFP8_GROUPED_MSG = "MXFP8 grouped GEMM is only supported when PyTorch is built with USE_FBGEMM_GENAI=1 on SM100+" @requiresCUDA - def test_scaled_grouped_mm_2d2d_tensorwise(): + def test_scaled_grouped_mm_3d2d_rowwise(): if not bool(PLATFORM_SUPPORTS_FP8_GROUPED_GEMM): pytest.skip(F8_GROUPED_MSG) device = "cuda" - group_sizes = [8, 12] - m = 6 - n = 10 - k = sum(group_sizes) + group_sizes = [16, 16] + groups = len(group_sizes) + total_rows = sum(group_sizes) + k = 16 + n = 16 - mat_a = torch.randn(m, k, device=device, dtype=torch.float16) - mat_b = torch.randn(k, n, device=device, dtype=torch.float16) - offs = torch.tensor(group_sizes, device=device, dtype=torch.int32).cumsum(0) - scale_a = torch.ones(1, device=device, dtype=torch.float32) - scale_b = torch.ones(1, device=device, dtype=torch.float32) + mat_a = torch.randn(total_rows, k, device=device, dtype=torch.bfloat16).to(torch.float8_e4m3fn) + mat_b = torch.randn(groups, n, k, device=device, dtype=torch.bfloat16).to(torch.float8_e4m3fn) + offs = torch.tensor(group_sizes, device=device, dtype=torch.int32).cumsum(0, dtype=torch.int32) + scale_a = torch.ones(total_rows, device=device, dtype=torch.float32) + scale_b = torch.ones(groups, n, device=device, dtype=torch.float32) def fn(a, b, scale_a, scale_b, offs): return torch.nn.functional.scaled_grouped_mm( a, - b, + b.transpose(-2, -1), scale_a, - ScalingType.TensorWise, + ScalingType.RowWise, scale_b, - ScalingType.TensorWise, + ScalingType.RowWise, offs=offs, output_dtype=torch.bfloat16, ) @@ -474,30 +475,30 @@ def fn(a, b, scale_a, scale_b, offs): assert_consistency_of_compiletime_and_runtime(jitted, result) @requiresCUDA - def test_scaled_grouped_mm_2d3d_tensorwise(): + def test_scaled_grouped_mm_2d3d_rowwise(): if not bool(PLATFORM_SUPPORTS_FP8_GROUPED_GEMM): pytest.skip(F8_GROUPED_MSG) device = "cuda" - group_sizes = [5, 7] + group_sizes = [8, 8] groups = len(group_sizes) total_rows = sum(group_sizes) k = 16 - n = 9 + n = 16 - mat_a = torch.randn(total_rows, k, device=device, dtype=torch.float16) - mat_b = torch.randn(groups, k, n, device=device, dtype=torch.float16) - offs = torch.tensor(group_sizes, device=device, dtype=torch.int32).cumsum(0) - scale_a = torch.ones(1, device=device, dtype=torch.float32) - scale_b = torch.ones(1, device=device, dtype=torch.float32) + mat_a = torch.randn(total_rows, k, device=device, dtype=torch.bfloat16).to(torch.float8_e4m3fn) + mat_b = torch.randn(groups, n, k, device=device, dtype=torch.bfloat16).to(torch.float8_e4m3fn) + offs = torch.tensor(group_sizes, device=device, dtype=torch.int32).cumsum(0, dtype=torch.int32) + scale_a = torch.ones(total_rows, device=device, dtype=torch.float32) + scale_b = torch.ones(groups, n, device=device, dtype=torch.float32) def fn(a, b, scale_a, scale_b, offs): return torch.nn.functional.scaled_grouped_mm( a, - b, + b.transpose(-2, -1), scale_a, - ScalingType.TensorWise, + ScalingType.RowWise, scale_b, - ScalingType.TensorWise, + ScalingType.RowWise, offs=offs, output_dtype=torch.bfloat16, ) From 77603f1574d9e013aef8949380f77032d12ef512 Mon Sep 17 00:00:00 2001 From: Masaki Kozuki Date: Mon, 17 Nov 2025 23:26:31 -0800 Subject: [PATCH 6/8] fix shape check cond and add tests to cover it Signed-off-by: Masaki Kozuki --- thunder/tests/test_ops.py | 39 ++++++++++++++++++++++++++------------- thunder/torch/__init__.py | 2 +- 2 files changed, 27 insertions(+), 14 deletions(-) diff --git a/thunder/tests/test_ops.py b/thunder/tests/test_ops.py index 9cc0ff319f..e697249d03 100644 --- a/thunder/tests/test_ops.py +++ b/thunder/tests/test_ops.py @@ -439,15 +439,20 @@ def fn(a): MXFP8_GROUPED_MSG = "MXFP8 grouped GEMM is only supported when PyTorch is built with USE_FBGEMM_GENAI=1 on SM100+" @requiresCUDA - def test_scaled_grouped_mm_3d2d_rowwise(): + @pytest.mark.parametrize( + "group_sizes,k,n", + [ + ([8, 8], 16, 16), + ([16, 16], 16, 16), + ], + ) + def test_scaled_grouped_mm_2d3d_rowwise(group_sizes, k, n): + """Test 2D x 3D grouped matmul with various dimensions.""" if not bool(PLATFORM_SUPPORTS_FP8_GROUPED_GEMM): pytest.skip(F8_GROUPED_MSG) device = "cuda" - group_sizes = [16, 16] groups = len(group_sizes) total_rows = sum(group_sizes) - k = 16 - n = 16 mat_a = torch.randn(total_rows, k, device=device, dtype=torch.bfloat16).to(torch.float8_e4m3fn) mat_b = torch.randn(groups, n, k, device=device, dtype=torch.bfloat16).to(torch.float8_e4m3fn) @@ -475,21 +480,29 @@ def fn(a, b, scale_a, scale_b, offs): assert_consistency_of_compiletime_and_runtime(jitted, result) @requiresCUDA - def test_scaled_grouped_mm_2d3d_rowwise(): + @pytest.mark.parametrize( + "group_sizes,m,k,n", + [ + ([8, 8], 16, 32, 16), # k != n to catch the dimension check bug + ([8, 8], 16, 16, 16), # k == n edge case + ], + ) + def test_scaled_grouped_mm_3d2d_rowwise(group_sizes, m, k, n): + """Test 3D x 2D grouped matmul with various dimensions. + + Note: k != n in first test case specifically catches the bug where + mat_a.shape[2] was incorrectly compared with mat_b.shape[1]. + """ if not bool(PLATFORM_SUPPORTS_FP8_GROUPED_GEMM): pytest.skip(F8_GROUPED_MSG) device = "cuda" - group_sizes = [8, 8] groups = len(group_sizes) - total_rows = sum(group_sizes) - k = 16 - n = 16 - mat_a = torch.randn(total_rows, k, device=device, dtype=torch.bfloat16).to(torch.float8_e4m3fn) - mat_b = torch.randn(groups, n, k, device=device, dtype=torch.bfloat16).to(torch.float8_e4m3fn) + mat_a = torch.randn(groups, m, k, device=device, dtype=torch.bfloat16).to(torch.float8_e4m3fn) + mat_b = torch.randn(n, k, device=device, dtype=torch.bfloat16).to(torch.float8_e4m3fn) offs = torch.tensor(group_sizes, device=device, dtype=torch.int32).cumsum(0, dtype=torch.int32) - scale_a = torch.ones(total_rows, device=device, dtype=torch.float32) - scale_b = torch.ones(groups, n, device=device, dtype=torch.float32) + scale_a = torch.ones(groups, m, device=device, dtype=torch.float32) + scale_b = torch.ones(n, device=device, dtype=torch.float32) def fn(a, b, scale_a, scale_b, offs): return torch.nn.functional.scaled_grouped_mm( diff --git a/thunder/torch/__init__.py b/thunder/torch/__init__.py index c84ccee8b6..dacd58c116 100644 --- a/thunder/torch/__init__.py +++ b/thunder/torch/__init__.py @@ -5857,7 +5857,7 @@ def scaled_grouped_mm( out_shape = (offs.shape[0], mat_a.shape[0], mat_b.shape[1]) elif mat_a.ndim == 3 and mat_b.ndim == 2: utils.check( - mat_a.shape[2] == mat_b.shape[1], + mat_a.shape[2] == mat_b.shape[0], lambda: f"Inner dimension mismatch: {mat_a.shape} vs {mat_b.shape}", ) utils.check( From 269d05620d355572640b98079c61b60ea9fa97ef Mon Sep 17 00:00:00 2001 From: Masaki Date: Tue, 18 Nov 2025 20:08:45 +0900 Subject: [PATCH 7/8] User permalink --- thunder/tests/test_ops.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/thunder/tests/test_ops.py b/thunder/tests/test_ops.py index e697249d03..659ac99eaf 100644 --- a/thunder/tests/test_ops.py +++ b/thunder/tests/test_ops.py @@ -434,7 +434,7 @@ def fn(a): # They validate that the Thunder tracing mirrors eager execution for representative 2D/2D # and 2D/3D grouped matmul shapes that correspond to the accepted combinations in # thunder.core.prims._grouped_mm_meta. These scenarios mirror the small tensor smoke tests - # in PyTorch's scaled matmul CUDA suite ([pytorch/test_scaled_matmul_cuda.py](https://github.com/pytorch/pytorch/blob/main/test/test_scaled_matmul_cuda.py)). + # in PyTorch's scaled matmul CUDA suite ([pytorch/test_scaled_matmul_cuda.py](https://github.com/pytorch/pytorch/blob/2f023bf7/test/test_scaled_matmul_cuda.py)). F8_GROUPED_MSG = "FP8 grouped is only supported on SM90 and MI300+ devices" MXFP8_GROUPED_MSG = "MXFP8 grouped GEMM is only supported when PyTorch is built with USE_FBGEMM_GENAI=1 on SM100+" From 27f3de3f0705e9fda3c2af5372a58938c118be51 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 15 Dec 2025 03:55:29 +0000 Subject: [PATCH 8/8] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- thunder/tests/test_ops.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/thunder/tests/test_ops.py b/thunder/tests/test_ops.py index 1721f489bf..6eca2522fd 100644 --- a/thunder/tests/test_ops.py +++ b/thunder/tests/test_ops.py @@ -589,6 +589,8 @@ def fn(mat_a, mat_b, scale_a, scale_b, offs): torch.testing.assert_close(result, eager) assert_consistency_of_compiletime_and_runtime(jitted, result) + + def _cuda_version_tuple() -> tuple[int, int] | None: if torch.version.cuda is None: return None