Expand Triton autotune configs for MoE FP8 kernels to improve AMD GPU performance (#3952)

* Expand Triton autotune configs for MoE FP8 rowwise and jagged scales kernels

The existing autotune configs for the MoE training FP8 kernels use a
single configuration each (e.g., num_warps=4, num_stages=4, one block
size), which prevents Triton's autotuner from finding better configs
for different hardware targets.

Expand the search space to cover:
- Multiple num_warps values (4, 8) to better saturate both NVIDIA
  (warp size 32) and AMD (wavefront size 64) GPU compute units
- Multiple num_stages values for software pipelining flexibility
  across different cache hierarchies
- Multiple block sizes to adapt to varying matrix dimensions

This is complementary to PR #3945 (relaxed atomics on AMDGPU) and
targets the same kernels.

* Gate expanded autotune configs to AMD only, preserve original NVIDIA configs

H100 benchmarks showed ~18% regression on the atomic kernel with the
expanded search space. The autotuner appears to pick suboptimal configs
from the larger candidate set on NVIDIA. Gate the expanded configs
behind torch.version.hip so AMD gets the broader search (4-7% faster
on MI250X) while NVIDIA keeps the original tuned configs.

* Widen autotune search space and add N_GROUPS to scales kernel autotuning key

Two improvements based on MI300X (gfx942) benchmarking:

1. float8_rowwise.py: Widen block size search space for AMD GPUs.
   - Atomic configs: add BLOCK_SIZE_N=256 and BLOCK_SIZE_K=64
   - Reduction configs: add BLOCK_SIZE_N=128, BLOCK_SIZE_K=64, and num_stages=2,4
   - Yields 1.5-2.2x speedup on MI300X for the atomic kernel and
     1.05-1.25x for the reduction kernel across Llama4 MoE shapes.

2. jagged_float8_scales.py: Add N_GROUPS to autotuning key for both
   rowwise and colwise scales kernels. The previous key (M or K only)
   caused the autotuner to cache a single config across all n_groups
   values, but optimal tile sizes differ significantly by n_groups.
   This eliminates cross-n_groups interference and allows each n_groups
   value to independently find its best config.

Author: brucechanglongxu

torchao/prototype/moe_training/kernels/float8_rowwise.py

@@ -33,21 +33,26 @@
         torch.float64: tl.float64,
     }
 
-    block_sizes_n = [128]  # large dim (output_features)
-    block_sizes_k = [128]  # small dim (input_features)
-    num_warps = [4]
-    num_stages = [4]
-    atomic_kernel_configs_2D = [
-        triton.Config(
-            {"BLOCK_SIZE_N": block_size_n, "BLOCK_SIZE_K": block_size_k},
-            num_warps=warps,
-            num_stages=stages,
-        )
-        for block_size_n in block_sizes_n
-        for block_size_k in block_sizes_k
-        for warps in num_warps
-        for stages in num_stages
-    ]
+    if torch.version.hip is not None:
+        atomic_kernel_configs_2D = [
+            triton.Config(
+                {"BLOCK_SIZE_N": block_size_n, "BLOCK_SIZE_K": block_size_k},
+                num_warps=warps,
+                num_stages=stages,
+            )
+            for block_size_n in [64, 128, 256]
+            for block_size_k in [64, 128]
+            for warps in [4, 8]
+            for stages in [2, 4]
+        ]
+    else:
+        atomic_kernel_configs_2D = [
+            triton.Config(
+                {"BLOCK_SIZE_N": 128, "BLOCK_SIZE_K": 128},
+                num_warps=4,
+                num_stages=4,
+            )
+        ]
 
     @torch.library.custom_op(
         "torchao::triton_fp8_rowwise_transpose_rhs", mutates_args={}
@@ -283,23 +288,26 @@ def _triton_fp8_rowwise_3d_transpose_cast_rhs_kernel(
         output_mask = (n_offs[:, None] < N) & (k_offs[None, :] < K)
         tl.store(output_ptr + output_offs, output_data, mask=output_mask)
 
-    block_sizes_n = [
-        64,
-    ]  # large dim (output_features)
-    block_sizes_k = [128]  # small dim (input_features)
-    num_warps = [8]
-    num_stages = [6]
-    reduction_kernel_configs_2D = [
-        triton.Config(
-            {"BLOCK_SIZE_N": block_size_n, "BLOCK_SIZE_K": block_size_k},
-            num_warps=warps,
-            num_stages=stages,
-        )
-        for block_size_n in block_sizes_n
-        for block_size_k in block_sizes_k
-        for warps in num_warps
-        for stages in num_stages
-    ]
+    if torch.version.hip is not None:
+        reduction_kernel_configs_2D = [
+            triton.Config(
+                {"BLOCK_SIZE_N": block_size_n, "BLOCK_SIZE_K": block_size_k},
+                num_warps=warps,
+                num_stages=stages,
+            )
+            for block_size_n in [32, 64, 128]
+            for block_size_k in [64, 128]
+            for warps in [4, 8]
+            for stages in [2, 4, 6]
+        ]
+    else:
+        reduction_kernel_configs_2D = [
+            triton.Config(
+                {"BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 128},
+                num_warps=8,
+                num_stages=6,
+            )
+        ]
 
     @triton.autotune(configs=reduction_kernel_configs_2D, key=["K", "N"])
     @triton.jit

torchao/prototype/moe_training/kernels/jagged_float8_scales.py

@@ -38,21 +38,26 @@
         torch.float64: tl.float64,
     }
 
-    block_sizes = [32]  # [16, 32, 64]
-    block_sizes_iter = [128]  # [64, 128, 256]
-    num_warps = [4]
-    num_stages = [3]
-    kernel_configs_2D = [
-        triton.Config(
-            {"BLOCK_SIZE": block_size, "BLOCK_SIZE_ITER": block_size_iter},
-            num_warps=warps,
-            num_stages=stages,
-        )
-        for block_size in block_sizes
-        for block_size_iter in block_sizes_iter
-        for warps in num_warps
-        for stages in num_stages
-    ]
+    if torch.version.hip is not None:
+        kernel_configs_2D = [
+            triton.Config(
+                {"BLOCK_SIZE": block_size, "BLOCK_SIZE_ITER": block_size_iter},
+                num_warps=warps,
+                num_stages=stages,
+            )
+            for block_size in [32, 64]
+            for block_size_iter in [64, 128]
+            for warps in [4, 8]
+            for stages in [2, 3]
+        ]
+    else:
+        kernel_configs_2D = [
+            triton.Config(
+                {"BLOCK_SIZE": 32, "BLOCK_SIZE_ITER": 128},
+                num_warps=4,
+                num_stages=3,
+            )
+        ]
 
     @torch.library.custom_op(
         "torchao::triton_fp8_per_group_rowwise_scales", mutates_args={}
@@ -108,6 +113,7 @@ def triton_fp8_per_group_rowwise_scales(
             scales_buffer,
             m,
             k,
+            n_groups,
             hp_tensor.stride(0),
             hp_tensor.stride(1),
             output_buffer.stride(0),
@@ -147,7 +153,7 @@ def _fake_triton_fp8_per_group_rowwise_scales_kernel(
     # so the kernel is easily interpretable in a standalone fasion.
     # The tokens per expert will vary per iteration, so don't want
     # to recompile on `token` dim (K, in this case) changes.
-    @triton.autotune(configs=kernel_configs_2D, key=["M"])
+    @triton.autotune(configs=kernel_configs_2D, key=["M", "N_GROUPS"])
     @triton.jit
     def _triton_fp8_per_group_rowwise_scales_kernel(
         input_ptr,
@@ -156,6 +162,7 @@ def _triton_fp8_per_group_rowwise_scales_kernel(
         scales_ptr,
         M: int,
         K: int,
+        N_GROUPS: int,
         stride_input_row: int,
         stride_input_col: int,
         stride_output_row: int,
@@ -299,6 +306,7 @@ def triton_fp8_per_group_colwise_scales(
             scales_buffer,
             k,
             n,
+            n_groups,
             hp_tensor.stride(0),
             hp_tensor.stride(1),
             output_buffer.stride(0),
@@ -336,7 +344,7 @@ def _fake_triton_fp8_per_group_colwise_scales(
     # 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.autotune(configs=kernel_configs_2D, key=["K", "N_GROUPS"])
     @triton.jit
     def _triton_fp8_per_group_colwise_scales_kernel(
         input_ptr,
@@ -345,6 +353,7 @@ def _triton_fp8_per_group_colwise_scales_kernel(
         scales_ptr,
         K: int,
         N: int,
+        N_GROUPS: int,
         stride_input_row: int,
         stride_input_col: int,
         stride_output_row: int,