Skip to content

feat(wgmma): add m64n64k32 FP8 E4M3 full-drain path - #1306

Draft
Vishalkulkarni45 wants to merge 6 commits into
NVlabs:mainfrom
Vishalkulkarni45:feat/fp8-wgmma
Draft

Vishalkulkarni45 wants to merge 6 commits into
NVlabs:mainfrom
Vishalkulkarni45:feat/fp8-wgmma

Conversation

@Vishalkulkarni45

@Vishalkulkarni45 Vishalkulkarni45 commented Sep 19, 2026

Copy link
Copy Markdown

Summary

Adds the shared-memory wgmma.mma_async.sync.aligned.m64n64k32.f32.e4m3.e4m3 path requested in #1281. This draft asks for review of the compiler implementation and H100 validation program before renting the GPU.

Refs #1281.

Status

  • E4M3 × E4M3, F32 accumulator, m64n64k32, fixed TN layout
  • All four layers: cuda-device, dialect-nvvm, mir-importer, and mir-lower
  • Full-drain lifetime: fence → MMA → commit_group → wait_group<0> with 32 tied F32 accumulators
  • Lowering and fail-closed rejection tests
  • Runnable FP8 example with valid SW32 shared-memory tiles, all-output F32 validation, and SUCCESS after validation
  • Comparable CUDA-event benchmark: two FP8 K=32 MMAs versus four BF16 K=16 MMAs for 64 × 64 × 64 work
  • H100 correctness run and SUCCESS output
  • H100 environment details and measured BF16/FP8 time and TFLOPS

Implementation

  • Adds the public intrinsic, NVVM pointer/value operations, MIR import, and canonical linear full-drain lowering.
  • Keeps all accumulator values inside one convergent inline-PTX region until wait_group<0>.
  • Rejects unsupported FP8 counted loops, partial waits, mixed variants, and noncanonical accumulators.
  • Fixes make_smem_desc for nonoverlapping K-major tiles: SW32, 32-byte K span, 256-byte group stride, and 256-byte base alignment.

H100 validation

cargo oxide run wgmma_mma_fp8 --arch sm_90a

The example validates all 4096 outputs against an F32 reference decoded from the represented E4M3 inputs. The benchmark uses 8192 CTAs, 10 alternating warmups, and 11 alternating samples of 100 CUDA-event-timed launches.

Path Average time TFLOPS Correctness
BF16 m64n64k16 × 4 Pending Pending Pending H100 run
FP8 m64n64k32 × 2 Pending Pending Pending H100 run

Local verification

  • 1,772 compiler/device tests passed.
  • 10 WGMMA PTX tests passed, including E4M3 chained MMA and zero-spill checks.
  • 4 FP8 example host tests passed.
  • sm_90a cubin materialization passed with CUDA 13.0.
  • Repository formatting, targeted Clippy, SPDX, dependency, book, and smoketest-contract checks passed.

The full just check suite and live GPU smoketests were not run.

Signed-off-by: vishal <vishalkoolkarni0045@gmail.com>
Signed-off-by: vishal <vishalkoolkarni0045@gmail.com>
Signed-off-by: vishal <vishalkoolkarni0045@gmail.com>
Signed-off-by: vishal <vishalkoolkarni0045@gmail.com>
Signed-off-by: vishal <vishalkoolkarni0045@gmail.com>
Signed-off-by: vishal <vishalkoolkarni0045@gmail.com>
@Vishalkulkarni45

Copy link
Copy Markdown
Author

@nihalpasham can u review this

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