Add CuTeDSL MXFP8 K-groups scale rearrange kernel

[alexsamardzic]

(Replaces #4604 due to incorrect ghstack base branch targeting.)

PR replaces the existing Triton MXFP8 K-groups scale rearrange kernel with a CuTeDSL implementation. It includes pytest coverage against a plain PyTorch reference and a validate/benchmark script for comparing correctness and performance against the current Triton path.

To test:

pytest -q test/prototype/moe_training/test_kernels.py::test_cutedsl_mx_block_rearrange_2d_K_groups

Benchmarking results, for M,K derived from some realistic models

MxK	Triton us	CuTeDSL us	CuTeDSL speedup
2048x4096	59.43	18.50	3.21
4096x4096	72.06	30.71	2.35
7168x4096	151.60	47.26	3.21
8192x4096	114.37	45.12	2.53
14336x4096	159.89	82.57	1.94
28672x4096	310.03	166.31	1.86
2048x8192	133.06	30.71	4.33
4096x8192	166.14	53.23	3.12
7168x8192	270.69	76.92	3.52
8192x8192	202.35	96.27	2.10
14336x8192	394.17	158.21	2.49
28672x8192	544.81	332.09	1.64

Benchmarking command to produce the table above

{
  for s in \
    2048,4096 4096,4096 7168,4096 8192,4096 14336,4096 28672,4096 \
    2048,8192 4096,8192 7168,8192 8192,8192 14336,8192 28672,8192; do
    IFS=, read -r r c <<< "$s"
    PYTHONPATH=$PWD python scripts/validate_k_groups_cutedsl.py \
      --rows $r --cols $c --groups 8 --multiple-of 1 \
      --cuda-graph-bench --graph-iters 1000
  done
} 2>/dev/null | perl -ne '
BEGIN {
  print "| MxK | Triton us | CuTeDSL us | CuTeDSL speedup |\n";
  print "|---:|---:|---:|---:|\n";
}
if (/shape=\((\d+),\s*(\d+)\)/) {
  $shape = "$1x$2";
  $triton_ok = undef;
  $cutedsl_ok = undef;
}
elsif (/triton_equal=(True|False)/) {
  $triton_ok = $1;
  die "triton_equal=False for $shape\n" if $triton_ok ne "True";
}
elsif (/cutedsl_equal=(True|False)/) {
  $cutedsl_ok = $1;
  die "cutedsl_equal=False for $shape\n" if $cutedsl_ok ne "True";
}
elsif (/triton_us=([\d.]+)/) { $triton = $1; }
elsif (/cutedsl_us=([\d.]+)/) {
  $cutedsl = $1;
  die "missing correctness lines for $shape\n"
    unless defined($triton_ok) && defined($cutedsl_ok);
  printf "| %s | %.2f | %.2f | %.2f |\n",
    $shape, $triton, $cutedsl, $triton / $cutedsl;
}
'

Benchmarking results, for M,K being powers of two

MxK	Triton us	CuTeDSL us	CuTeDSL speedup
128x32	4.10	6.15	0.67
128x64	4.61	6.15	0.75
128x128	4.10	6.15	0.67
128x256	8.19	6.15	1.33
128x512	12.29	6.45	1.91
128x1024	14.35	6.15	2.33
128x2048	28.69	8.19	3.50
128x4096	63.46	8.19	7.75
128x8192	153.53	8.19	18.75
128x16384	176.05	8.20	21.47
256x32	4.10	6.14	0.67
256x64	4.11	6.15	0.67
256x128	6.15	6.15	1.00
256x256	8.20	6.15	1.33
256x512	8.20	6.15	1.33
256x1024	16.40	7.62	2.15
256x2048	30.71	8.19	3.75
256x4096	40.98	8.19	5.00
256x8192	116.69	8.33	14.01
256x16384	216.44	10.25	21.12
512x32	4.10	6.15	0.67
512x64	4.11	6.15	0.67
512x128	6.15	6.21	0.99
512x256	8.19	6.15	1.33
512x512	14.33	6.93	2.07
512x1024	16.38	8.19	2.00
512x2048	22.55	8.19	2.75
512x4096	59.43	8.21	7.24
512x8192	106.55	10.25	10.40
512x16384	327.53	16.59	19.74
1024x32	4.10	6.14	0.67
1024x64	5.14	6.15	0.84
1024x128	6.15	6.22	0.99
1024x256	6.15	7.00	0.88
1024x512	10.24	8.19	1.25
1024x1024	16.27	8.20	1.98
1024x2048	32.79	10.24	3.20
1024x4096	49.13	12.29	4.00
1024x8192	88.03	16.30	5.40
1024x16384	219.04	30.71	7.13
2048x32	4.10	6.15	0.67
2048x64	4.28	6.15	0.70
2048x128	6.15	8.21	0.75
2048x256	8.19	8.20	1.00
2048x512	12.30	8.20	1.50
2048x1024	14.35	10.24	1.40
2048x2048	26.64	12.31	2.16
2048x4096	49.19	18.44	2.67
2048x8192	104.50	25.00	4.18
2048x16384	344.25	45.15	7.62
4096x32	4.10	6.15	0.67
4096x64	4.10	6.15	0.67
4096x128	6.15	8.20	0.75
4096x256	10.24	8.22	1.25
4096x512	16.19	10.24	1.58
4096x1024	24.59	12.32	2.00
4096x2048	30.71	19.17	1.60
4096x4096	83.94	28.69	2.93
4096x8192	174.58	45.30	3.85
4096x16384	281.60	107.16	2.63
8192x32	4.10	6.15	0.67
8192x64	6.15	8.20	0.75
8192x128	8.19	10.80	0.76
8192x256	12.30	12.29	1.00
8192x512	18.44	14.50	1.27
8192x1024	28.67	18.42	1.56
8192x2048	63.50	32.43	1.96
8192x4096	125.47	55.32	2.27
8192x8192	241.20	88.13	2.74
8192x16384	372.04	178.17	2.09
16384x32	6.15	8.19	0.75
16384x64	6.15	10.25	0.60
16384x128	10.24	15.12	0.68
16384x256	14.33	16.39	0.87
16384x512	24.59	20.55	1.20
16384x1024	52.62	32.73	1.61
16384x2048	118.66	43.01	2.76
16384x4096	244.33	106.63	2.29
16384x8192	421.36	197.93	2.13
16384x16384	627.75	354.11	1.77

[alexsamardzic]

Stack from ghstack (oldest at bottom):

• Wire MoE and MXFP8Linear training paths to CuTeDSL by default #4705
• -> Add CuTeDSL MXFP8 K-groups scale rearrange kernel #4704
• Add CuTeDSL MXFP8 pad/unpad kernels #4703
• Fix fake metadata for CuTeDSL 32x1 MXFP8 quantization #4702
• Add CuTeDSL MXFP8 M-groups scale rearrange kernel #4701

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/4704

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

• [ALERT ONLY] - System Maintenance

✅ You can merge normally! (1 Unrelated Failure)

As of commit 91b1436 with merge base 33ba9f7 ():

BROKEN TRUNK - The following job failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• Run PT2E GPU Tests / test-rocm (ROCM Nightly, linux.rocm.gpu.gfx942.1, --pre torch torchvision --index-url https://dow... / linux-job (gh) (trunk failure)

This comment was automatically generated by Dr. CI and updates every 15 minutes.