[QST] SM120 FP8 MMA with FP32 accumulate, plain form against block-scaled form with unit scales

Open
  • Is something wrong?
#3,732 6 comments 0 reactions 0 assignees View on GitHub

Maintainers usually reply within 4 days

Nobody has claimed this yet.

Assessment

Difficulty
4/5
Estimated time
3-5 days
Newbie friendliness
43/100
Issue type
Feature
Clarity
Mostly clear
Activity status
Active
Tech stack
cpp
Domain
backend, performance

Research direction

Read the SM120 selectors in sm120_mma_builder.inl, sm120_array_mma_builder.inl, and sm120_blockwise_mma_builder.inl, along with CuTe DSL MmaFP8Op and the reported CUTLASS GEMM measurements. Determine whether the E4M3 × E4M3 FP32-accumulate case should select the block-scaled form with unit scales, accounting for the reported swizzle caveat. Done means the default decision is supported by correctness and performance evidence for the scoped SM120 case.

Written by the indexing model from the issue text.

Description

What is your question?

On an RTX 5080, the plain FP8 MMA with FP32 accumulate runs at half the rate of the block-scaled FP8 MMA with FP32 accumulate. When all scale bytes are 0x7F (UE8M0 1.0), the two forms give equal bits.

My FP8 attention kernel in CuTe DSL used MmaFP8Op. I changed it to MmaMXF8Op with constant unit scales and no other change. The kernel time went from 123.6 ms to 83.6 ms and the output is byte-equal (sgl-project/sglang#43390). So a user can do this today. My question is about the default.

  1. Is there a reason to keep the plain form for E4M3 x E4M3 with FP32 accumulate on SM120?
  2. If there is no reason, do you want MmaFP8Op and rr_op_selector_sm120 to emit the block-scaled form with unit scales for this case? I can send a PR.
Related report

The same behavior is reported in flashinfer-ai/flashinfer#5963 for the FlashInfer SM120 FP8 attention kernel, with RTX 5090 numbers and a bit-identity check. That issue also points to the MXFP8 attention work in flashinfer-ai/flashinfer#3628 and flashinfer-ai/flashinfer#4149. The numbers below add an RTX 5080 and a CUTLASS GEMM measurement to that report.

Scope

E4M3 x E4M3, FP32 accumulate, one RTX 5080 (driver 595.99.02, CUDA 13.3, CUTLASS ffa119a12). I did not measure other types or other cards.

Raw MMA rate

Independent MMA chains, no memory access, 84 SMs. The rate is the same from 4 to 32 warps per SM.

instruction TFLOP/s
kind::f8f6f4.m16n8k32 ... f32.e4m3.e4m3.f32 246
kind::mxf8f6f4.block_scale.scale_vec::1X.m16n8k32 ... f32.e4m3.e4m3.f32.ue8m0 489
Bit equality
  • Random finite inputs. 0 of 704,643,072 accumulator registers differ.
  • Negative control with scale 2.0. 704,643,071 registers differ.
  • Edge cases (subnormal, ±448, NaN bytes, accumulator start at zero, FLT_MAX, subnormal, infinity, NaN). 30 classes, 176,160,768 registers each, 0 bits differ.
CUTLASS GEMM

I used the SM120 CollectiveBuilder with OpClassTensorOp, pingpong schedule, tile 128x128x128, FP32 output. Then I changed only the instruction in SM120_16x8x32_TN<float_e4m3_t, float_e4m3_t, float> to the block-scaled form. The scale 0x7F7F7F7F stays in a register. There is no scale factor tensor and no additional load. The patch is 34 lines.

I measured 30 shapes from real models. All 30 give 0 different output words. Five examples:

shape source M x N x K plain patched speedup
Wan 2.2 14B, to_q 32760 x 5120 x 5120 236.1 398.6 1.69x
Llama 3.3 70B, MLP up 4096 x 28672 x 8192 234.7 398.1 1.70x
FLUX.1-dev, FFN up 4608 x 12288 x 3072 231.8 393.3 1.70x
MiniMax-H3, FFN up 30272 x 14336 x 5376 234.6 363.5 1.55x
Gemma 3 27B, MLP up 128 x 21504 x 5376 195.8 203.4 1.04x

TFLOP/s. The speedup is 1.55x to 1.72x for the 22 shapes with M of 4096 or more. It is 1.04x to 1.82x for the 8 shapes with M = 128, where the low end is DRAM-bound.

Nsight Compute on the stock plain kernel shows the tensor pipe 96% busy. The MMA instruction rate is 47.9% of the peak.

One caution. These numbers use max_swizzle_size = 8. With the default of 1, three large shapes become DRAM-bound and the patched kernel is slower than the plain kernel (0.88x). I can open a different issue for that with the data.

Where the plain form is used
  • CuTe DSL MmaFP8Op
  • rr_op_selector_sm120 in sm120_mma_builder.inl, sm120_array_mma_builder.inl and sm120_blockwise_mma_builder.inl

cuBLASLt on the same card uses QMMA.SF.16832.F32.E4M3.E4M3.E8 in its FP8 GEMM kernel and no plain QMMA.

I can share the benchmark source, the patch and the raw results if they are useful.

Dominant language
C++
Stars
10.7k
Forks
2.2k
Avg merge
6d 10h
Merged PRs (30d)
6

Getting set up

This project ships no dev container, Dockerfile or contributing guide, so setting up is up to you: start from its README, and see our first-contribution guide for the general steps.

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

More from NVIDIA/cutlass

All issues in NVIDIA/cutlass

Similar issues

More C++ issues

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.