[inductor] tune_scaled_grouped_mm fails with memory layout assertion, despite memory layout assertions prior to op call passing

[danielvegamyhre]

🐛 Describe the bug

In the torchao float8 MoE training prototype, we have an autograd func that performs a differentiable scaled grouped mm, with dynamic fp8 rowwise quantization on inputs before every GEMM in forward + backward.

Recently, torch.compile started throwing an error when compiling this, when it wasn't ~1-2 months ago. It seems to be related to the addition of autotuner support for scaled grouped mm.

Specifically, the error is asserting the "B" tensor is one of the scaled grouped GEMMs is not in column memory major layout. However, my code has assertions before every torch._scaled_grouped_mm call validating that the "A" tensor is row major and the "B" tensor is column major - and these assertions do not trigger. The code works as expected in eager mode.

Error:

torch._inductor.exc.InductorError: LoweringException: RuntimeError: mat_b must be col_major, got stride [67108864, 8192, 1]
  target: aten._scaled_grouped_mm.default

Full stack trace: https://www.internalfb.com/phabricator/paste/view/P1844809161

Repro

• Checkout PR in torchao (skip if PR is landed by the time you see this): [float8 moe training] make using triton kernels for per-group scaling configurable ao#2405
• Run python benchmark_scaled_grouped_mm.py --compile from ~/ao/torchao/prototype/moe_training/benchmarks

Versions

torch: 2.8.0.dev20250611+cu126
torchao: 0.11.0+gitcdced21f

cc @chauhang @penguinwu @voznesenskym @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @ipiszy @chenyang78 @kadeng @muchulee8 @amjames @aakhundov

[danielvegamyhre]

cc @alexsamardzic

[alexsamardzic]

After a quick assessment: I tried to put a print statement before torch._scaled_grouped_gemm call in _Float8GroupedMM.forward, in order to check strides of the operands - the problem disappears, and B strides printed show that its layout is indeed column-major. However, print introduces a graph break, thus my best guess here is that B is not turning out column-major because of fusion. Are we sure that assert statements actually stay in the compiled code (I'd assume they would introduce a graph break too if so)?

The backward pass for this variant of _grouped_gemm is tested here, and it works fine. The derivative for _scaled_grouped_gemm is not implemented in PyTorch core, so corresponding test could not have been written there, but single Triton kernel is handling both grouped and scaled grouped variants, so there should be no difference whatsoever.
Thus, I don't see an immediate problem with compiled implementation of torch._scaled_grouped_gemm per se - so I'd check what happens during the fusion, and if you still find that the Triton kernel, or maybe related Inductor code, should be changed, please let me know.

[danielvegamyhre]

This was fixed by #159134