mx: triton kernel to cast to mx and write in col-major

[vkuzo]

Summary:

Implements a triton kernel for a cast to mxfp8 from a row-major input across dim1, which is 3.5x to 4.5x faster than what compile can generate today. Note that this is a prototype kernel, and I expect to (a) improve it in future PRs and (b) delete it in ~weeks when we have compile support for this.

An integration into MXLinear will follow in a separate PR.

Example of tiling (simplified for small example size):

        Example tiling for n_rows==8, n_cols=8, ROW_TILE_SIZE=4, COL_TILE_SIZE=4, INNER_BLOCK_SIZE=2,
        pid_row=0, pid_col=0:

        Input (row-major)

        cols      0  1  2  3  4  5  6  7
        --------------------------------
        rows 0 |  0  1  2  3
             1 |  8  9 10 11
             2 | 16 17 18 19
             3 | 24 25 26 27
             4 |
             5 |
             6 |
             7 |

        Output (row-major of transpose), ids are from input

        cols      0  1  2  3  4  5  6  7
        --------------------------------
        rows 0 |  0  8 16 24
             1 |  1  9 17 25
             2 |  2 10 18 26
             3 |  3 11 19 27
             4 |
             5 |
             6 |
             7 |

        Output (scales), s(0, 8) means the scale used to cast elements 0 and 8

        rows           0          1  ...      4  ...       31
        ------------------------------------------------------
                  s(0, 8)  s(16, 24) ... s(1, 9) ... s(19, 27)

Test Plan:

// tests pass
pytest test/prototype/mx_formats/test_custom_cast.py -s -x -k triton_mxfp8_dim1

// performance compile vs triton: https://www.internalfb.com/phabricator/paste/view/P1762691809
// * 4k by 4k tensor: about a 3.6x speedup
// * 16k by 16k tensor: about a 4.7x speedup

Reviewers:

Subscribers:

Tasks:

Tags:

[vkuzo]

Stack from ghstack (oldest at bottom):

• mx roofline: adjust mxfp8 formulas #1953
• print MX config when printing MXLinear and MXInferenceLinear #1947
• integrate mx dim1 triton kernel into MXLinear #1943
• -> mx: triton kernel to cast to mx and write in col-major #1932

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 1 New Failure

As of commit 7ecd79f with merge base 3fb1665 ():

NEW FAILURE - The following job has failed:

• Run TorchAO Experimental Tests / test-mps-ops (macos-m1-stable) (gh)
  ModuleNotFoundError: No module named 'importlib_metadata'

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

[eellison]

i think we should just be doing integer division instead of floor + /

[eellison]

looks good ! next step: compile to generate this

[eellison]

in the launcher, should we assert divisibility of block sizes, so we hard error for this case ?

[vkuzo]

yes, for now I hackly assert that on L1319:L1322

[eellison]

since only the data_ptr of output_col_major is used when you pass it into triton, you could initialize it with the correct strides

[eellison]

same thing here..