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

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.

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:

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.