[mxfp8 moe training] add CUDA kernel to quantize 3d tensor colwise

[danielvegamyhre]

Stacked PRs:

• [mxfp8 moe training] wrap 3d quantize tensor in custom ops and integrate it #3004
• ->[mxfp8 moe training] add CUDA kernel to quantize 3d tensor colwise #3002
• [moe training] add benchmarks for dsv3 236b, 671b shapes; reorganize benchmarks dir #2999
• [mxfp8 moe training] use dim1 cast cuda kernel for 3d weights by reshaping to 2d #2998

---

[mxfp8 moe training] add CUDA kernel to quantize 3d tensor colwise

Summary

• This PR adds a new CUDA kernel specifically for quantizing 3d expert weights shape (E,N,K) along the N dimension and writing directly to column major format.
  • Design: I create separate input/output TMA descriptors for each expert, and process each 2d expert in parallel using the same method that the 2d dim1 quantization kernel uses. The 2d kernel achieves 85% peak memory bandwidth utilization, so hopefully we can achieve similar perf for 3d.
• The existing methods for quantizing 3d expert weights both scale very poorly. I have verified this via benchmarking and traces (see previous PR), and hypothesize that it is due to required .contiguous() calls:
  • Using to_mx only scales along the last dim and requires contiguos inputs. So this requires transposing contiguous tensor (E,N,K) -> (E,K,N) then calling .contiguous() to scale along the N dim (needed for backwards)
  • Using the existing CUDA kernel for casting along dim1 is possible, by treating the 3d input tensor as a 2d tensor of shape (E*N, K). However, this produces a 2d output tensor in column major format, and there is no way to reshape and restride the tensor to be 3d again AND preserve the column major format, such that numerics are preserved. Thus, we have to transform the output to column major afterwards, requiring a .contiguous() call.

Test plan

• Added tests that verify numerical accuracy

Kernel microbenchmarks

Perf is decent for large E and abysmal for small E. Need to investigate this.

Update (9/15): NCU shows 3d kernel operating on (2,8192,5120) tensor is actually compute bound (??)

input_shape         to_mx_us    cuda_2d_us    cuda_3d_us    to_mx_gbps    cuda_2d_gbps    cuda_3d_gbps
----------------  ----------  ------------  ------------  ------------  --------------  --------------
(1, 8192, 5120)      117.92         69.776       242.112      1078.19          1822.11         525.128
(2, 8192, 5120)      431.264       105.536       249.728       589.615         2409.41        1018.23
(4, 8192, 5120)      848.992       195.584       297.376       599.015         2600.21        1710.16
(8, 8192, 5120)     1682.59        379.904       412.992       604.495         2677.3         2462.8
(16, 8192, 5120)    3350.53        775.984       615.536       607.139         2621.49        3304.82
(64, 8192, 5120)   13352          3150.66       1959.71        609.418         2582.62        4152.11

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

✅ No Failures

As of commit bb8b07f with merge base f75b251 ():
💚 Looks good so far! There are no failures yet. 💚

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

[danielvegamyhre]

@slayton58 @ngimel i would be curious to get your thoughts on ways to improve this kernel for quantizing 3d expert weights (E,N,K) along the N dim, where weights are contiguous. It uses nearly identical logic to the 2d dim1 cast kernel (which achieves ~85% mem bw utilization), yet the perf is much worse (~8% to 62% peak mem bw, depending on input size - see benchmarks in PR description).

I think the culprit might be how i'm allocating all the TMA descriptors and passing them in, the overhead might be too much for small E? NCU has not flagged anything particularly helpful so far. Strangely, for E=2 it shows the kernel is compute bound with 78% compute throughput % and 38% memory bandwidth %.

Additional context: torch.compile and handwritten triton kernels were both slow for mxfp8 quant for RHS operands where we scale colwise (32x1 granularity) e.g., (triton hit 56% peak mem bw). So I added a CUDA kernel here which I derived from a TE kernel which achieves ~85% peak mem bw (#2513). Basically we stripped out internal TE types, added support for different scale calculation modes (floor, rceil) to align with torchao numerics, then resolved some perf issues resulting from those changes to get reasonable perf.

Now, I'm finding quantizing 3d expert weights along dim1 is scaling extremely poorly as number of experts increases (see this PR's description for details, and see #2999 for benchmarks). So I added a similar CUDA kernel to our mxfp8_cuda extension specifically for quantizing 3d expert weights colwise and writing directly to col major format we need it in.

The first approach I tried was just updating the 2d kernel to handle 3d tensors by treating it as a 2d tensor of shape (E*N, K) but the coordinate mapping / pointer arithemetic became a complicated mess that wasn't working. So I made a new kernel, that is similar to the 2d kernel but passes in separate input/output TMA descriptors for each expert, then the kernel operates on each 2d expert with logical separation, in parallel.