…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 pytorch#3945 (relaxed atomics on AMDGPU) and
targets the same kernels.

… 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.