Add a triton kernel for swizziling

[drisspg]

Stacked PRs:

• Add subclass based method for inference w/ MXFP8 #2132
• ->Add a triton kernel for swizziling #2168

---

LLama 70b feed forward w/ this kernel:
https://fburl.com/125yv8hh
~632 µs
without:
https://fburl.com/a21gwjmc
~ 801 µs

BF16 for reference:
https://fburl.com/2lgn9xkx

VLLM

Once I figure out why this isnt' always working, I am getting

============ Serving Benchmark Result ============
Successful requests:                     1024      
Benchmark duration (s):                  11.37     
Total input tokens:                      225502    
Total generated tokens:                  189380    
Request throughput (req/s):              90.09     
Output token throughput (tok/s):         16661.09  
Total Token throughput (tok/s):          36500.08  
---------------Time to First Token----------------
Mean TTFT (ms):                          965.17    
Median TTFT (ms):                        880.49    
P99 TTFT (ms):                           1492.64   
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          28.41     
Median TPOT (ms):                        29.23     
P99 TPOT (ms):                           44.61     
---------------Inter-token Latency----------------
Mean ITL (ms):                           20.41     
Median ITL (ms):                         14.86     
P99 ITL (ms):                            49.08     
==================================================

Future

We are swizziling the scales for all the weights as well we can shave some micros seconds if we keep around a cached variant

[pytorch-bot]

🔗 Helpful Links

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

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

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

[drisspg]

Well thats fun

No correctness issues w/ TRITON_INTERPRET=1 and sanitize is clean

[drisspg]

All good now:

=========================================================================================== warnings summary ===========================================================================================
test/prototype/mx_formats/test_mx_linear.py::test_linear_compile[False-False-mxfp8_emulated-hp_dtype0]
  /home/drisspg/meta/pytorch/torch/_inductor/compile_fx.py:246: UserWarning: TensorFloat32 tensor cores for float32 matrix multiplication available but not enabled. Consider setting `torch.set_float32_matmul_precision('high')` for better performance.
    warnings.warn(

-- Docs: https://docs.pytest.org/en/stable/how-to/capture-warnings.html
======================================================================== 257 passed, 63 skipped, 1 warning in 116.63s (0:01:56) ========================================================================
❯ 

TLDR: dont use reorder=True w/ tl.reshape

[vkuzo]

lgtm!

[drisspg]