Skip to content

add MXFP8 pre-swizzling for gfx1250 GEMM#568

Open
matthiasdiener wants to merge 28 commits into
devfrom
mdiener/mxfp8-swizzle
Open

add MXFP8 pre-swizzling for gfx1250 GEMM#568
matthiasdiener wants to merge 28 commits into
devfrom
mdiener/mxfp8-swizzle

Conversation

@matthiasdiener
Copy link
Copy Markdown
Contributor

@matthiasdiener matthiasdiener commented Apr 29, 2026

Description

Fixes https://github.com/ROCm/frameworks-internal/issues/16428

This was manually tested on gfx1250.

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:

  • Change A
  • Change B

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

@matthiasdiener matthiasdiener self-assigned this Apr 29, 2026
@matthiasdiener matthiasdiener added the ci-level 1 CI test level 1 label Apr 29, 2026
@matthiasdiener matthiasdiener force-pushed the mdiener/mxfp8-swizzle branch from ddf19da to 313a6b7 Compare May 3, 2026 22:06
@matthiasdiener matthiasdiener requested a review from alextmagro May 4, 2026 16:33
Copy link
Copy Markdown
Contributor

@alextmagro alextmagro left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hi Matthias, a few comments. I also assume you are still planning on adding in the hooks to scale swizzle when we're on gfx1250? I believe there were hooks in all of common, pytorch and jax. These PRs removed them, so would be a partial revert.

#420
#424
#442

Comment thread transformer_engine/common/hadamard_transform/hadamard_transform.cu Outdated
Comment thread transformer_engine/common/swizzle/swizzle.cu
Comment thread transformer_engine/common/swizzle/swizzle.cu Outdated
Comment thread tests/cpp/operator/test_swizzle.cu
Comment thread transformer_engine/common/gemm/rocm_gemm.cu Outdated
Comment thread transformer_engine/pytorch/csrc/quantizer.cpp Outdated
@matthiasdiener matthiasdiener changed the title [proof-of-concept] add MXFP8 pre-swizzling for gfx1250 add MXFP8 pre-swizzling for gfx1250 GEMM May 13, 2026
@matthiasdiener
Copy link
Copy Markdown
Contributor Author

I also assume you are still planning on adding in the hooks to scale swizzle when we're on gfx1250? I believe there were hooks in all of common, pytorch and jax. These PRs removed them, so would be a partial revert.

#420 #424 #442

The hooks should be re-added in 384d590.

@matthiasdiener matthiasdiener requested a review from alextmagro May 14, 2026 20:20
@matthiasdiener matthiasdiener marked this pull request as ready for review May 14, 2026 20:21
Comment thread tests/cpp/operator/test_cublaslt_gemm.cu Outdated
Comment thread tests/cpp/operator/test_cublaslt_gemm.cu Outdated
Comment thread transformer_engine/pytorch/csrc/extensions/swizzle.cpp
@matthiasdiener matthiasdiener requested a review from alextmagro May 19, 2026 19:54
@alextmagro
Copy link
Copy Markdown
Contributor

LGTM

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

Labels

ci-level 1 CI test level 1

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants