[moe training] use custom ops instead of wrap_triton for fp8 rowwise kernels (#2734)

Author: danielvegamyhre

benchmarks/prototype/moe_training/benchmark_kernels.py

@@ -15,8 +15,8 @@
 from triton.testing import do_bench
 
 from torchao.prototype.moe_training.kernels.jagged_float8_scales import (
-    triton_fp8_col_major_jagged_colwise_scales,
-    triton_fp8_row_major_jagged_rowwise_scales,
+    triton_fp8_per_group_colwise_scales,
+    triton_fp8_per_group_rowwise_scales,
 )
 from torchao.prototype.moe_training.utils import (
     torch_to_float8_per_group_colwise,
@@ -114,13 +114,13 @@ def run_torch(
     def run_triton(
         input_row_major: torch.Tensor, input_col_major: torch.Tensor, offs: torch.Tensor
     ):
-        _ = triton_fp8_row_major_jagged_rowwise_scales(
+        _ = triton_fp8_per_group_rowwise_scales(
             input_row_major,
             offs,
             output_dtype=torch.float8_e4m3fn,
             round_scales_to_power_of_2=True,
         )
-        _ = triton_fp8_col_major_jagged_colwise_scales(
+        _ = triton_fp8_per_group_colwise_scales(
             input_col_major,
             offs,
             output_dtype=torch.float8_e4m3fn,

benchmarks/prototype/moe_training/benchmark_per_group_scaling_kernels.py

@@ -15,8 +15,8 @@
 from triton.testing import do_bench
 
 from torchao.prototype.moe_training.kernels.jagged_float8_scales import (
-    triton_fp8_col_major_jagged_colwise_scales,
-    triton_fp8_row_major_jagged_rowwise_scales,
+    triton_fp8_per_group_colwise_scales,
+    triton_fp8_per_group_rowwise_scales,
 )
 from torchao.prototype.moe_training.utils import (
     torch_to_float8_per_group_colwise,
@@ -114,13 +114,13 @@ def run_torch(
     def run_triton(
         input_row_major: torch.Tensor, input_col_major: torch.Tensor, offs: torch.Tensor
     ):
-        _ = triton_fp8_row_major_jagged_rowwise_scales(
+        _ = triton_fp8_per_group_rowwise_scales(
             input_row_major,
             offs,
             output_dtype=torch.float8_e4m3fn,
             round_scales_to_power_of_2=True,
         )
-        _ = triton_fp8_col_major_jagged_colwise_scales(
+        _ = triton_fp8_per_group_colwise_scales(
             input_col_major,
             offs,
             output_dtype=torch.float8_e4m3fn,

test/prototype/moe_training/test_kernels.py

@@ -17,8 +17,8 @@
     triton_fp8_rowwise_3d_transpose_rhs,
 )
 from torchao.prototype.moe_training.kernels.jagged_float8_scales import (
-    triton_fp8_col_major_jagged_colwise_scales,
-    triton_fp8_row_major_jagged_rowwise_scales,
+    triton_fp8_per_group_colwise_scales,
+    triton_fp8_per_group_rowwise_scales,
 )
 from torchao.prototype.moe_training.utils import (
     _is_column_major,
@@ -46,7 +46,7 @@ def test_row_major_with_jagged_rowwise_scales(round_scales_to_power_of_2: bool):
         target_dtype=torch.float8_e4m3fn,
         round_scales_to_power_of_2=round_scales_to_power_of_2,
     )
-    kernel_fp8_data, kernel_scales = triton_fp8_row_major_jagged_rowwise_scales(
+    kernel_fp8_data, kernel_scales = triton_fp8_per_group_rowwise_scales(
         x,
         colwise_offs,
         output_dtype=torch.float8_e4m3fn,
@@ -74,7 +74,7 @@ def test_column_major_with_jagged_colwise_scales(round_scales_to_power_of_2: boo
         target_dtype=torch.float8_e4m3fn,
         round_scales_to_power_of_2=round_scales_to_power_of_2,
     )
-    kernel_fp8_data, kernel_scales = triton_fp8_col_major_jagged_colwise_scales(
+    kernel_fp8_data, kernel_scales = triton_fp8_per_group_colwise_scales(
         x,
         rowwise_offs,
         output_dtype=torch.float8_e4m3fn,

test/prototype/moe_training/test_training.py

@@ -121,7 +121,7 @@ def moe_module_filter_fn(mod: nn.Module, cur_fqn: str) -> bool:
     )
 
     # validate param gradients
-    min_param_grad_sqnr = 25.0
+    min_param_grad_sqnr = 23.0
     for param1, param2 in zip(model.parameters(), ref_model.parameters()):
         param_grad_sqnr = compute_error(param1.grad, param2.grad)
         assert param_grad_sqnr.item() >= min_param_grad_sqnr, (

torchao/prototype/moe_training/kernels/__init__.py

@@ -2,8 +2,8 @@
     triton_fp8_rowwise_3d_transpose_rhs as triton_fp8_rowwise_3d_transpose_rhs,
 )
 from torchao.prototype.moe_training.kernels.jagged_float8_scales import (
-    triton_fp8_col_major_jagged_colwise_scales as triton_fp8_col_major_jagged_colwise_scales,
+    triton_fp8_per_group_colwise_scales as triton_fp8_per_group_colwise_scales,
 )
 from torchao.prototype.moe_training.kernels.jagged_float8_scales import (
-    triton_fp8_row_major_jagged_rowwise_scales as triton_fp8_row_major_jagged_rowwise_scales,
+    triton_fp8_per_group_rowwise_scales as triton_fp8_per_group_rowwise_scales,
 )

torchao/prototype/moe_training/kernels/float8_rowwise.py

@@ -42,10 +42,8 @@
     for stages in num_stages
 ]
 
-from torch.library import triton_op, wrap_triton
 
-
-@triton_op("torchao::triton_fp8_rowwise_transpose_rhs", mutates_args={})
+@torch.library.custom_op("torchao::triton_fp8_rowwise_transpose_rhs", mutates_args={})
 def triton_fp8_rowwise_3d_transpose_rhs(
     hp_tensor: torch.Tensor,  # (E, K, N)
     output_dtype: torch.dtype = torch.float8_e4m3fn,
@@ -80,7 +78,7 @@ def triton_fp8_rowwise_3d_transpose_rhs(
     )
 
     # compute scales
-    wrap_triton(_triton_fp8_rowwise_3d_transpose_scales_rhs_kernel)[grid](
+    _triton_fp8_rowwise_3d_transpose_scales_rhs_kernel[grid](
         hp_tensor,
         hp_tensor.stride(0),
         hp_tensor.stride(1),
@@ -100,7 +98,7 @@ def triton_fp8_rowwise_3d_transpose_rhs(
     )
 
     # perform casting
-    wrap_triton(_triton_fp8_rowwise_3d_transpose_cast_rhs_kernel)[grid](
+    _triton_fp8_rowwise_3d_transpose_cast_rhs_kernel[grid](
         hp_tensor,
         hp_tensor.stride(0),
         hp_tensor.stride(1),
@@ -124,6 +122,22 @@ def triton_fp8_rowwise_3d_transpose_rhs(
     return output_buffer, scales_buffer
 
 
+@triton_fp8_rowwise_3d_transpose_rhs.register_fake
+def _fake_triton_fp8_rowwise_3d_transpose_rhs(
+    hp_tensor: torch.Tensor,  # (E, K, N)
+    output_dtype: torch.dtype = torch.float8_e4m3fn,
+    round_scales_to_power_of_2: bool = False,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    assert hp_tensor.ndim == 3, "input tensor must be 3D"
+    e, k, n = hp_tensor.shape
+    output_buffer = torch.empty(
+        (e, n, k), dtype=output_dtype, device=hp_tensor.device
+    ).as_strided((e, n, k), (n * k, 1, n))
+
+    scales_buffer = torch.empty((e, k), dtype=torch.float32, device=hp_tensor.device)
+    return output_buffer, scales_buffer
+
+
 @triton.autotune(configs=kernel_configs_2D, key=["num_elements"])
 @triton.jit
 def _triton_fp8_rowwise_3d_transpose_scales_rhs_kernel(

torchao/prototype/moe_training/kernels/jagged_float8_scales.py

@@ -47,11 +47,11 @@
     for stages in num_stages
 ]
 
-from torch.library import triton_op, wrap_triton
 
-
-@triton_op("torchao::triton_fp8_row_major_jagged_rowwise_scales", mutates_args={})
-def triton_fp8_row_major_jagged_rowwise_scales(
+@torch.library.custom_op(
+    "torchao::triton_fp8_per_group_rowwise_scales", mutates_args={}
+)
+def triton_fp8_per_group_rowwise_scales(
     hp_tensor: torch.Tensor,
     offsets: torch.Tensor,
     output_dtype: torch.dtype = torch.float8_e4m3fn,
@@ -95,7 +95,7 @@ def triton_fp8_row_major_jagged_rowwise_scales(
         triton.cdiv(m, meta["BLOCK_SIZE"]),
         offsets.numel(),
     )
-    wrap_triton(_triton_fp8_row_major_jagged_rowwise_scales)[grid](
+    _triton_fp8_per_group_rowwise_scales_kernel[grid](
         hp_tensor,
         offsets,
         output_buffer,
@@ -117,6 +117,24 @@ def triton_fp8_row_major_jagged_rowwise_scales(
     return output_buffer, scales_buffer
 
 
+@triton_fp8_per_group_rowwise_scales.register_fake
+def _fake_triton_fp8_per_group_rowwise_scales_kernel(
+    hp_tensor: torch.Tensor,
+    offsets: torch.Tensor,
+    output_dtype: torch.dtype = torch.float8_e4m3fn,
+    round_scales_to_power_of_2: bool = False,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    assert hp_tensor.ndim == 2, "input tensor must be 2D"
+    m, k = hp_tensor.shape
+    n_groups = offsets.numel()
+    output = torch.empty_like(hp_tensor, dtype=output_dtype).as_strided(
+        (m, k),  # shape
+        (k, 1),  # stride
+    )
+    scales = torch.empty((m * n_groups), dtype=torch.float32, device=hp_tensor.device)
+    return output, scales
+
+
 # This kernel is used on grad_output.t() which has shape (K, M),
 # before the calculation `grad_B = grad_output_t @ input`.
 # However, in this code, we use the conventional dim names (M, K)
@@ -125,7 +143,7 @@ def triton_fp8_row_major_jagged_rowwise_scales(
 # to recompile on `token` dim (K, in this case) changes.
 @triton.autotune(configs=kernel_configs_2D, key=["M"])
 @triton.jit
-def _triton_fp8_row_major_jagged_rowwise_scales(
+def _triton_fp8_per_group_rowwise_scales_kernel(
     input_ptr,
     offsets_ptr,
     out_ptr,
@@ -215,8 +233,10 @@ def _triton_fp8_row_major_jagged_rowwise_scales(
         tl.store(out_ptr + out_offs, fp8_data, mask=block_mask)
 
 
-@triton_op("torchao::triton_fp8_col_major_jagged_colwise_scales", mutates_args={})
-def triton_fp8_col_major_jagged_colwise_scales(
+@torch.library.custom_op(
+    "torchao::triton_fp8_per_group_colwise_scales", mutates_args={}
+)
+def triton_fp8_per_group_colwise_scales(
     hp_tensor: torch.Tensor,
     offsets: torch.Tensor,
     output_dtype: torch.dtype = torch.float8_e4m3fn,
@@ -263,7 +283,7 @@ def triton_fp8_col_major_jagged_colwise_scales(
         triton.cdiv(n, meta["BLOCK_SIZE"]),
         offsets.numel(),
     )
-    wrap_triton(_triton_fp8_col_major_jagged_colwise_scales)[grid](
+    _triton_fp8_per_group_colwise_scales_kernel[grid](
         hp_tensor,
         offsets,
         output_buffer,
@@ -285,13 +305,33 @@ def triton_fp8_col_major_jagged_colwise_scales(
     return output_buffer, scales_buffer
 
 
+@triton_fp8_per_group_colwise_scales.register_fake
+def _fake_triton_fp8_per_group_colwise_scales(
+    hp_tensor: torch.Tensor,
+    offsets: torch.Tensor,
+    output_dtype: torch.dtype = torch.float8_e4m3fn,
+    round_scales_to_power_of_2: bool = False,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    assert hp_tensor.ndim == 2, "input tensor must be 2D"
+    k, n = hp_tensor.shape
+    n_groups = offsets.numel()
+    output_buffer = torch.empty_like(
+        hp_tensor, dtype=output_dtype, device=hp_tensor.device
+    ).as_strided(hp_tensor.size(), (1, k))
+
+    scales_buffer = torch.empty(
+        (n * n_groups), dtype=torch.float32, device=hp_tensor.device
+    )
+    return output_buffer, scales_buffer
+
+
 # This kernel is used on `input` which has shape (M, K),
 # before the calculation `grad_B = grad_output_t @ input`.
 # The tokens per expert will vary per iteration, so don't want
 # to recompile on `token` dim (M) changes.
 @triton.autotune(configs=kernel_configs_2D, key=["K"])
 @triton.jit
-def _triton_fp8_col_major_jagged_colwise_scales(
+def _triton_fp8_per_group_colwise_scales_kernel(
     input_ptr,
     offsets_ptr,
     out_ptr,

torchao/prototype/moe_training/scaled_grouped_mm.py

@@ -13,8 +13,8 @@
 from torchao.float8.float8_utils import tensor_to_scale, to_fp8_saturated
 from torchao.prototype.moe_training.conversion_utils import MoEScalingType
 from torchao.prototype.moe_training.kernels import (
-    triton_fp8_col_major_jagged_colwise_scales,
-    triton_fp8_row_major_jagged_rowwise_scales,
+    triton_fp8_per_group_colwise_scales,
+    triton_fp8_per_group_rowwise_scales,
     triton_fp8_rowwise_3d_transpose_rhs,
 )
 from torchao.prototype.moe_training.utils import (
@@ -230,15 +230,15 @@ def backward(ctx, grad_output: torch.Tensor):
         # Convert transpose of grad_output to float8, row-major for left operand of grouped GEMM
         # needed for grad_B: grad_output_t @ A
         grad_output_t_fp8_row_major, grad_output_t_scales = (
-            triton_fp8_row_major_jagged_rowwise_scales(
+            triton_fp8_per_group_rowwise_scales(
                 grad_output.transpose(-2, -1),
                 offs,
                 torch.float8_e4m3fn,
                 round_scales_to_power_of_2=True,
             )
         )
 
-        A_fp8_col_major, A_scales = triton_fp8_col_major_jagged_colwise_scales(
+        A_fp8_col_major, A_scales = triton_fp8_per_group_colwise_scales(
             A,
             offs,
             torch.float8_e4m3fn,