🔗 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 b64758e 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.

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

Discussed offline, likely culprit is cudaMallocManaged calls on the hot path https://github.com/pytorch/ao/pull/3002/files#diff-7ddc6623d9efea4ee4f4bdb3cdd7ef16ec3d3bc8bc974be85311125e464efb3dR1326, we should be able to create just a single descriptor and do the loads using the correct offsets.

Discussed offline, likely culprit is cudaMallocManaged calls on the hot path https://github.com/pytorch/ao/pull/3002/files#diff-7ddc6623d9efea4ee4f4bdb3cdd7ef16ec3d3bc8bc974be85311125e464efb3dR1326, we should be able to create just a single descriptor and do the loads using the correct offsets.

I've been trying this today, progress so far is:

• Scales for all experts are correct
• Quantized data is correct for the expert=0 subtensor, but is all 0s for all other experts

This indicates the new input TMA descriptor and async loads are working properly, but the output TMA descriptor and/or async stores are not.

I think this is because the input tensor is in simple row major format, which can easily be represented in a TMA descriptor with shape (E*N, K) with stride K. However, the output data needs to be in "column major PER expert" format, so strides (N*K, 1, N). The 2d output TMA descriptor + ptx::cp_async_bulk_tensor_2d_global_to_shared do not seem capable of representing this layout so far (or could be a skill issue on my part, haha)

I will push a WIP PR on top of this stack to show the difference between this multiple-tma-descriptor approach, which is at least functionally correct, versus the single tma descriptor approach.

@danielvegamyhre I think there's a couple of options, can we do something like:

# out : [E, N, K], stride: [N*K, 1, N]#
out.transpose(2,1) # Now [E, K, N], stride: [N*K, N, 1], representable by row-major TMA descriptor (2d or otherwise)
# compute
out.transpose(1, 2) # Back to original form

or, we can try using 3d TMAs directly - ptx::cp_async_bulk_tensor_3d_global_to_shared is the relevant read invocation, and a 3d descriptor would have to be created for this.
or (finally) we can ignore TMA writes, and use regular global stores (so STG) - we shouldn't need the async, pipeline-able part of the TMA instruction.

@slayton58 row major -> transpose to per expert col major is a great idea! Trying it now.

or, we can try using 3d TMAs directly

Yeah I've considered this, it is probably the "proper" way to do it but would require a larger refactor, hopefully the transpose method works.

can we just merge the top of stack w/ this pR