mx_block_rearrange_2d_M_groups_cuda does not support scales_tensor with M-group size = 44

[eee4017]

Description

We are hitting a runtime failure when using the newly added CUDA kernel mx_block_rearrange_2d_M_groups_cuda introduced in PR #3546.

Our workload requires a specific per-group size (44), which appears to be unsupported by the kernel. The failure manifests as a CUDA driver error at runtime.

Input configuration:

scales_tensor: torch.Size([98048, 44])
dtype = torch.float8_e8m0fnu

input_group_end_offsets: torch.Size([8])
dtype = torch.int32

chunk_width = 64
chunks_per_tb = 8

Error:

CUDA workload error: CUDA Driver Error at /workspace/ao/torchao/csrc/cuda/mx_kernels/mx_block_rearrange_2d_M_groups.cu:431
invalid argument

cc @danielvegamyhre @drisspg

[danielvegamyhre]

at first glance this looks like a validation bug but will investigate.

in practice you'll need to pad group sizes to be multiple of 32, the scaling block size for mxfp8.
this is because in the backward pass, for wgrad = grad_out^T @ input_act the groups are along the contracting dim, where the 1x32 (LHS) and 32x1 (RHS) scaling blocks are applied, and we can't have a scaling block cross group boundaries.

for example, in TorchTitan, we do this here, we fuse this padding to the kernel which does "permutation."

A better name would be gather_and_fill or similar, really. what it does is generate indexes to rearrange/gather tokens from layout:

[rank0 tokens for expert 0, rank0 tokens for expert 1, rank1 tokens for expert0, rank1 tokens for expert1, ...]

to layout:

[tokens for expert 0 from all ranks, tokens for expert 1 from all ranks, ...] - INCLUDING adding padding tokens as necessary to ensure group sizes are divisible by 32.

If you want, we have our own version of this operation in torchao you can use:

ao/torchao/prototype/moe_training/ep/permute.py

Line 167 in c6d21f9

def _permute_bf16(

i will figure out why this kernel is not handling this properly, though.

[danielvegamyhre]

@eee4017 as a temporary mitigation, you can set this use_cuda_kernel_for_blocked_layout to False:

ao/torchao/prototype/moe_training/scaled_grouped_mm.py

Line 308 in c6d21f9

use_cuda_kernel_for_blocked_layout: bool = True,

This will use a triton kernel instead, which is slower but doesn't have this issue.

[danielvegamyhre]

@eee4017 follow up questions for you:

1. are you using _to_mxfp8_then_scaled_grouped_mm directly in the modeling code, or using quantize_(...) to do model conversion?
2. can you share how you're using this feature? this is helpful for us to do planning, prioritization, etc