Skip to content

[Common] Add MXFP8 2D quantization CuTeDSL kernels - #3843

Open
kainzhong wants to merge 38 commits into
NVIDIA:mainfrom
kainzhong:mxfp8_2d_quant
Open

kainzhong wants to merge 38 commits into
NVIDIA:mainfrom
kainzhong:mxfp8_2d_quant

Conversation

@kainzhong

Copy link
Copy Markdown
Collaborator

Description

Add MXFP8 2D quantization CuTeDSL kernels

Fixes # (issue)

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

Please list the changes introduced in this PR:

  • MXFP8 2D quantization CuTeDSL kernels

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

kainzhong and others added 30 commits October 7, 2026 07:22
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
kainzhong and others added 5 commits October 19, 2026 21:23
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
The more_cutedsl_mxfp8_specialized base was squash-merged as NVIDIA#3642, so
CuTeDSL conflicts were resolved by 3-way merging against bfdb240.
- quantize_mxfp8_common.py: keep the preferred_kernels restructure and
  carry over main's removal of "cast-only" from the comment.
- quantize_mxfp8_cutedsl.cuh: drop the 2D-unsupported guard that main
  had only edited, since 2D quantization is now implemented in CuTeDSL.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
@greptile-apps

greptile-apps Bot commented Nov 9, 2026 •

Copy link
Copy Markdown
Contributor

RetriggerView in GreptileConfidence Score: 5/5

[Medium impact] The PR appears safe to merge; both previous findings are addressed and no new actionable issues were found.

Summary

Adds CuTeDSL MXFP8 quantization with one shared scale per 32×32 block.

  • CuTeDSL now quantizes plain MXFP8 casts with one scale per 32-by-32 block.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart TD
  A[MXFP8 cast request] --> B{2D scaling requested?}
  B -->|Yes| C[2D tile math in general pipeline]
  B -->|No| D[Existing kernel selection]
  C --> E[FP8 values and repeated block scales]
  D --> F[FP8 values and directional scales]
Loading

Reviews (2) · Last reviewed commit: "[pre-commit.ci] auto fixes from pre-comm..." · Reviewed by Greptile

kainzhong and others added 2 commits November 10, 2026 00:04
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant