Skip to content

[Dialect] Add FP8 support to GFX120X WMMA atom - #1074

Open
big-yellow-duck wants to merge 1 commit into
ROCm:mainfrom
big-yellow-duck:codex/gfx120x-fp8-wmma-atom
Open

[Dialect] Add FP8 support to GFX120X WMMA atom#1074
big-yellow-duck wants to merge 1 commit into
ROCm:mainfrom
big-yellow-duck:codex/gfx120x-fp8-wmma-atom

Conversation

@big-yellow-duck

Copy link
Copy Markdown

Motivation

The GFX120X WMMA atom added in #943 currently exposes the RDNA4 wave32 FP16 and BF16 paths, but not the hardware FP8 path. As a result, FlyDSL kernels targeting gfx1200/gfx1201 cannot use the regular MMA atom API for E4M3FN WMMA and must bypass it with a raw ROCDL operation.

This completes the floating-point surface of the existing GFX120X atom by adding E4M3FN inputs with FP32 accumulation.

Technical Details

  • Accept f8E4M3FN x f8E4M3FN -> f32 for the GFX120X 16x16x16 WMMA atom.
  • Lower both memref and SSA atom calls to rocdl.wmma.f32.16x16x16.fp8_fp8.
  • Use the RDNA4 v8 operand ABI: eight FP8 values per lane packed as vector<2xi32>.
  • Retain validation that GFX120X supports only 16x16x16, matched operand types, FP32 accumulation, and no floating-point sign/clamp controls.
  • Document the supported FP8 path and extend positive and negative lowering coverage.

This is compiler/atom support only. Optimized block-scaled GEMM kernels and shape routing are intentionally left for a separate follow-up.

Test Plan

  • Build FlyDSL with cmake --build build-fly -j2.
  • Run the positive GFX120X WMMA lowering test through fly-opt and FileCheck.
  • Run the split-input negative verifier test through fly-opt and FileCheck.
  • Run pytest -q tests/kernels/test_rdna4_wmma_atom.py on gfx1201.
  • Run Python/C++ formatting checks and git diff --check.

Test Result

  • Build passed.
  • Positive and negative MLIR tests passed.
  • GFX120X device atom tests passed: 3 tests.
  • Targeted RDNA4 FP8 suite passed: 34 tests.
  • Formatting and whitespace checks passed.

Submission Checklist

  • Looked over the contributing guidelines.
  • Added tests for the new behavior.
  • Updated relevant documentation.

@coderfeli
coderfeli requested a review from sjfeng1999 August 27, 2026 11:16
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