[kainzhong]

Description

Add specialized MXFP8 CuTeDSL quantization kernels ported from #3459
And also some refactoring of the existing kernel -- so quantize_mxfp8.py becomes

• quantize_mxfp8_common.py: kernel config, signature, dispatching and registeration
• quantize_mxfp8_g2s_general.py: the same general MXFP8 kernel that handles all cases
• quantize_mxfp8_g2r_rowwise_1lane.py: the same old rowwise specialized kernel
• quantize_mxfp8_g2s_bidim.py: the same old bidim specialized kernel

g2s and g2r mean the kernel uses SMEM / RMEM to quantize

Note: CuTeDSL cannot exactly mirror CUDA's implementation due to a bug in cutlass: NVIDIA/cutlass#3726

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Changes

• Split the original MXFP8 CuTeDSL file into 3 different files (generic, bidim specialized, rowwise specialized)
• Added new kernels

Difference with the CUDA path:

• G2R bidim kernel also works with swizzled scale
• G2R bidim kernel also handles colwise only quantization
• G2R kernels also support ACT / DACT fusion
• Due to CuTeDSL's G2R activation fusion uses different paths than CUDA, sometimes they disagree with +0 and -0. I think both should be legal so I relaxed bit-exact tests to tolerate +0 and -0

Benchmark result (full sweep on the tip of this branch including things that aren't changed by this PR):

https://claude.ai/artifact/GsCNMy9S4YTxHnR7W3b1Vo?sk=DtZFlcQxJpBxCOlmq-W8_A

Or download it as a static HTML:
Quantization Benchmark .html

To examine which kernel is dispatched, expand "Data table" and you can see kernel names

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

[greptile-apps]

Confidence Score: 5/5

[Medium impact] The PR appears safe to merge; no new blocking issue was found.

Summary

Adds BF16 register-based MXFP8 kernels and splits the existing kernels into separate files.

• Rowwise MXFP8 casts now have register-based fast paths.
• Bidirectional MXFP8 casts now have a register-based fast path.
• MXFP8 calls now fall back to the general kernel when no specialized path fits.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart TD
  A[MXFP8 call] --> B{Eligible specialized kernel?}
  B -->|BF16 rowwise| C[Two-lane register kernel]
  B -->|BF16 columnwise or both| D[Bidimensional register kernel]
  B -->|Other supported plain casts| E[Existing specialized kernels]
  B -->|No| F[General shared-memory kernel]
  C --> G[FP8 data and block scales]
  D --> G
  E --> G
  F --> G

Loading

_{Reviews (15) · Last reviewed commit: "can always use static tile num" · Reviewed by Greptile}