[Common] More CuTeDSL MXFP8 specialized kernels

[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: (BF16 only, since the new G2R kernels only support BF16, and activation is gelu)

https://claude.ai/artifact/SzJ4HBGpo5KRS2VzSpbim5?sk=g7SsHtfq1MQxEzpL4z-Rpg

Non-Swizzled scale

\	Rowwise	Columnwise	Both
Plain
ACT
DACT

Swizzled scale

\	Rowwise	Columnwise	Both
Plain
ACT
DACT

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 changes appear safe to merge; the previous plain-cast mismatch is fixed and no new actionable issue was found.

Summary

Adds specialized CuTeDSL MXFP8 kernels and splits the existing implementation into shared dispatch code and separate kernel files.

• Rowwise MXFP8 casts get dedicated register-memory kernels.
• Bidimensional MXFP8 casts can use a register-memory kernel.
• MXFP8 dispatch routes eligible casts to specialized kernels.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart TD
  A[MXFP8 quantization request] --> B[Shared configuration and dispatch]
  B --> C{Specialized kernel supports this call?}
  C -->|Yes| D[Rowwise or bidimensional kernel]
  C -->|No| E[General shared-memory kernel]
  D --> F[FP8 values and block scales]
  E --> F

Loading

_{Reviews (11) · Last reviewed commit: "nit" · Reviewed by Greptile}