triton kernel to cast to mx across dim0 and dim1

[vkuzo]

Summary:

This is a start of the short term workaround for pytorch/pytorch#148682 . tl;dr; is that today torch.compile is not yet generating good fusions for the cast to mx across dim0 and dim1, so we can add custom kernels to do this until the torch.compile support is better. We take a simple strategy to minimize overall reads and writes: read the input in 32x32 tiles, and write out the row-major and col-major casted values and scales directly from each tile.

The kernel added in this PR is ok-ish: it's ~2x faster than torch.compile across a range of shapes with MKN from 1024 to 16384, but it only hits around (redacted) 8 TBPS peak mem bandwidth. Impact on e2e benchmarks of fwd+bwd of MXLinear:

# K = 1024: speedup over bf16 0.19 -> 0.28
# K = 16384: speedup over bf16 0.91 -> 1.07

Some future improvements we can make:

1. add an outer_tile and set it to 128, keeping inner tile at 32 (see below, will likely get us to redacted of peak bandwidth). 128 because outer_tile 256+ seems to run out of shared memory on B200.
2. write out swizzled scales (not currently in this benchmark, but will help with e2e)
3. see if we can make the col-major output and row-major and col-major writes coalesced (didn't do any research on how yet)

For now I want to check this in as just a kernel + benchmark without an integration to MXLinear. If we make this kernel significantly faster, we can integrate.

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:

[vkuzo]

Stack from ghstack (oldest at bottom):

• -> triton kernel to cast to mx across dim0 and dim1 #1869

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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 3f33752 with merge base ddb7f83 ():
💚 Looks good so far! There are no failures yet. 💚

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

[eellison]

lol i wasnt sure what to call this either..

[eellison]

looks good ! did not review indexing math too closely but hoepfully will be replaced with torch.compile soon.

[eellison]

Because the triton is just accessing data_ptr you could initialize these as the correct layout if you wanted.

[eellison]

any reason to not just use integer division here instead of / and floor ?

[eellison]

Should we add some runtime asserts about block divisibility ?

[eellison]

I would suspect theres a better way of getting the scale avoiding doing a float log. i guess this is what the TODO is for.

[vkuzo]

closing this in favor of #1932, I split the dim1 logic so we can have this PR as an easy historical reference