[gfx1250] Add MXFP4 GEMM support to FlyDSL backend by aris134 · Pull Request #786 · ROCm/TransformerEngine

aris134 · 2026-10-30T15:05:44Z

Description

Adds MXFP4 dense GEMM support for gfx1250 as part of the FlyDSL GEMM backend for training. Preliminary support does not have activation/bias, which is intended for a follow-up PR.

The GEMM backend requires plain scales and data, and performs scale shuffles and swizzles internally. The scales are shuffled into iteration-major order. The operand matrices are pre-swizzled at the scale-block level before the GEMM to work with TDM loads.

Fixes AMD-ROCm-Internal/frameworks-internal#394

Pipeline structure

All operands stream through LDS with no K-dependent allocation; total LDS use is a constant 144 KB.

Data tiles — 3-stage pipeline. A and B each get NUM_STAGES = 3 LDS pages (BLOCK_M×16 words per page). Each K-loop iteration issues one TDM load into page j % 3 while the MMA consumes an earlier page, giving ~2 iterations of load/compute overlap. In steady state a K-tile is loaded by 4 waves (TDM_GROUP_WAVES): 2 split A's 256 rows into halves (TDM_WAVES), 2 split B's.
Scales — 2-stage, grouped pipeline. SCALE_STAGES = 2 LDS pages each hold SCALE_GROUP = 8 K-tiles of scales (16 KB per operand). The prologue loads groups 0–1; thereafter, every 8th iteration the tile-load also streams the next scale group into page g % 2, sharing that iteration's wait slot. On those scale-refill iterations the data load switches from 2×half-tile to 1×full-tile to free two waves for the scale groups.
Registers — double-buffered. frag_A/B/sA/sB are 2-deep; tile k+1 is read into set k%2 while gemm_step(k) runs, interleaving next-tile LDS reads with current WMMAs.
Epilogue — 2-pass C store. The output tile is written in two C_HALF_ROWS = 128-row halves, staged register→LDS→global (TDM) through the LDS region reused from the K-loop, serialized so half 1 reuses half 0's buffer.

Performance

https://github.com/AMD-ROCm-Internal/frameworks-internal/issues/394#issuecomment-6525965976

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

…fx1250

alextmagro

Hi Aristotle, working on a review of the main loop now but a few integration comments to start.

alextmagro · 2026-11-03T18:23:06Z

        return (K, M), (N, K)
    if layout == "NT":  # A not transposed, B transposed
        return (K, M), (K, N)
+    if layout == "TT":  # both transposed


Do we need TT for anything? My understanding is that it is more efficient to swap the operands and do B@A and run a TT GEMM as an NN, then you get a row-major C instead.

Thanks for pointing this out. No, I don't actually think we need TT for anything here, so I dropped it in 126f1dd

alextmagro · 2026-11-03T18:23:23Z

+def create_mxfp4_tensors(M, K, N, layout):
+    """Create MXFP4 inputs and references dequantized from the buffers ``layout`` consumes.
+
+    MXFP4 rowwise and columnwise data are quantized separately (columnwise is the


comment trim

Done in 126f1dd

alextmagro · 2026-11-03T18:23:50Z

+    buffer the GEMM reads: rowwise for a transposed operand, columnwise otherwise.
+    """
+    # the MXFP4 quantizer takes BF16 input
+    A_bf16, B_bf16 = create_regular_tensors(M, K, N, layout, dtype=torch.bfloat16)


Are we limited to bf16 input? I think supporting fp32 and fp16 is important as well, especially since the test outputs all 3.

It seems the bf16 requirement is a hard-coded constraint.

TransformerEngine/transformer_engine/common/cast/mxfp4/quantize_mxfp4.cuh

Line 32 in 55919ff

"MXFP4 quantization only supports BF16 input (got ", to_string(input.dtype()), ")");

Extending support would probably fit better with a follow up PR perhaps?

alextmagro · 2026-11-03T18:26:24Z



+def _select_mxfp4_storage_for_layout(A, transa, B, transb):
+    """Select TE MXFP4 data and scales for a layout.


comment trim

Done in 126f1dd

alextmagro · 2026-11-03T18:28:14Z

+    B_words = as_words(B, N, k_words)
+
+    # physical chunks of this lane's logical chunks khalf and khalf + 2 (same for all its rows)
+    swz = (l16 // 4) % SWZ_CHUNKS


116 is ambiguous here, can we use a const var name?

Ah it is "L"16, that is a bit hard to read.

Yeah you're right, this isn't obvious, so I renamed it to lane_mod16 in 126f1dd

Add gfx1250 MXFP4 GEMM support to FlyDSL backend

bdac43a

aris134 requested review from Micky774, ipanfilo, matthiasdiener, wangye805 and wenchenvincent as code owners October 30, 2026 15:05

aris134 requested review from VeeraRajasekhar, alextmagro and sudhu2k October 30, 2026 15:06

clean up comments

512b0ff

aris134 self-assigned this Oct 30, 2026

aris134 added the ci-level 1 CI test level 1 label Oct 30, 2026

aris134 changed the title ~~Add gfx1250 MXFP4 GEMM support to FlyDSL backend~~ [gfx1250] Add MXFP4 GEMM support to FlyDSL backend Nov 2, 2026

aris134 added 2 commits November 2, 2026 18:56

restructure scale pipeline to remove K size limitation

ac84339

Merge remote-tracking branch 'origin/dev' into amartin/flydsl-mxfp4-g…

cfca1e6

…fx1250

alextmagro reviewed Nov 3, 2026

View reviewed changes

address review comments

126f1dd

aris134 requested a review from alextmagro November 3, 2026 20:48

Navigation Menu

Repository navigation

[gfx1250] Add MXFP4 GEMM support to FlyDSL backend - #786

[gfx1250] Add MXFP4 GEMM support to FlyDSL backend#786
aris134 wants to merge 5 commits into
devfrom
amartin/flydsl-mxfp4-gfx1250

aris134 commented Oct 30, 2026 •

edited

Loading

Uh oh!

alextmagro left a comment

Uh oh!

alextmagro Nov 3, 2026

Uh oh!

aris134 Nov 3, 2026

Uh oh!

alextmagro Nov 3, 2026

Uh oh!

aris134 Nov 3, 2026

Uh oh!

alextmagro Nov 3, 2026

Uh oh!

aris134 Nov 3, 2026

Uh oh!

alextmagro Nov 3, 2026

Uh oh!

aris134 Nov 3, 2026

Uh oh!

alextmagro Nov 3, 2026

Uh oh!

alextmagro Nov 3, 2026

Uh oh!

aris134 Nov 3, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants



		def _select_mxfp4_storage_for_layout(A, transa, B, transb):
		"""Select TE MXFP4 data and scales for a layout.

Conversation

aris134 commented Oct 30, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Pipeline structure

Performance

Type of change

Changes

Checklist:

Uh oh!

alextmagro left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

aris134 commented Oct 30, 2026 •

edited

Loading