[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 actionable issue was found.

Summary

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

• Rowwise MXFP8 casts can use register-based kernels for more input layouts.
• Bidimensional MXFP8 casts can quantize from a register-loaded tile.
• MXFP8 calls now select specialized kernels when their inputs fit.

Previous-thread context:

• The added BF16 cases address the missing launch-shape coverage.
• kainzhong accepted the lack of exact kernel-selection assertions because these tests check correctness, not performance. Greptile withdrew that concern.
• kainzhong explicitly accepts positive-zero versus negative-zero differences in the PR description. Both backend tests now allow that difference.
• The coverage and signed-zero threads were manually resolved without explanatory replies.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart TD
  A[MXFP8 quantization call] --> B[Check configuration and shape]
  B --> C[Rowwise two-lane G2R]
  B --> D[Bidimensional G2R]
  B --> E[Rowwise one-lane G2R]
  B --> F[Bidimensional G2S]
  B --> G[General G2S fallback]

Loading

_{Reviews (18) · Last reviewed commit: "Merge branch 'main' into more_cutedsl_mx..." · Reviewed by Greptile}

[tdophung]

for the case of WITH_ACT/WITH_DACT, the scale computed by this CuteDSL function is different from how thee CUDA kernel fused case compute it. CUDA version ignores NaN and the scale = e8m0(fmaxf(amax * max_norm_rcp, 0)), so it scale will be 0 when there is NaN. However this function in cuteDSL will return 0xFF for the scale here.
I know that the cast-only version in CUDA also uses the same scale calculation, which returns 0xFF. And numerically this should not matter, but since you have the byte-for-byte test between CuteDSL and CUDA, thought I might call it out for the fused activation case.

[Oleg-Goncharov]

I’d keep the new behavior and define all-NaN blocks to produce a NaN scale (0xFF), consistently with the cast-only path. The fused CUDA reduction currently loses the all-NaN condition by starting fmaxf from zero.
As we’re planning to transition to CuTeDSL kernels (and eventually retire the legacy CUDA implementations), we should probably prefer to standardize this behavior. And the remaining inconsistencies between backends can be fixed in a separate fix PR

[kainzhong]

OK so... I'm just not gonna change anything?

[tdophung]

Yeah you do not need to change anything then, per Oleg. I'll fix the inconsistencies in the CUDA sidee in the next PR. Keeping this thread open and not resolved so I remember it though

[Oleg-Goncharov]

LGTM

[tdophung]

Approve pending CI (and #3743 + #3737 )