[FeathBow]

Why

sm_89 adds an FP8 warp MMA, mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32, which the sm_89 target from #259 could not select: T.cuda.sm80.Mma takes BF16 only, so an E4M3 product in registers had no instruction to lower to. It was suggested in #256 as the change to make after the sm_89 documents.

What

T.cuda.sm89.Mma declares the instruction: 16 x 8 x 32, E4M3 A and B in registers, and an F32 accumulator in registers laid out as the SM80 atom's. Its A and B fragments follow the PTX fragment table and match CuTe's MMA_Traits<SM89_16x8x32_F32E4M3E4M3F32_TN> thread for thread. It requires mma.sync, which sm_89, sm_90 and sm_100 all list and all run in the E4M3 form. TiledMma accepts it after the SM80 atom, and the register fragment pattern both atoms build moves into mma_atom.py.

The runtime keeps one ops::mma entry. The element type of A selects the CuTe atom; B must hold the same type, C must hold f32, and each fragment must have the selected atom's length, or a static_assert fails. Codegen emits the same call for both atoms.

schedule facts names each capability once per instruction, so T.tiled_mma still lists wgmma.mma_async, mma.sync with three atoms behind it.

Tests:

• tests/ir/cuda/atoms/test_mma.py compares both atoms' fragments with CuTe's traits
• tests/ops/tir/cuda/test_mma.py runs an E4M3 product on the GPU against the F32 product of the same values, once for H200 and once for sm_89 through the sm89 marker
• tests/ops/ir/test_matmul.py types an E4M3 register by register matmul on sm_89
• tests/runtime/test_cuda_static_asserts.py checks that ops::mma refuses an E4M3 A with a BF16 B, a C that is not F32, and a B with the BF16 atom's length
• two schedule goldens and the refused atom set gain the new atom

Contract

docs/spec/tir.md lists T.cuda.sm89.Mma among the public declarations and the declarations the CUDA emitter accepts. docs/spec/runtime.md §2.6.5 states how ops::mma selects its atom and each atom's fragment sizes.

Risk

mma.sync names the instruction family, so every architecture document that lists it admits the E4M3 atom. A document for an architecture below 8.9, such as a user supplied sm_80 document, would admit an instruction that hardware lacks. No installed document is affected.

[FeathBow]

Test plan

These ran on a single RTX 6000 Ada Generation (compute capability 8.9) with Python 3.12, PyTorch 2.11 for CUDA 12.9 and nvcc 12.8:

• The four kernels marked sm89 run on this card, including the new E4M3 one, and all four pass.
• pytest tests ran on main and on this change; the table is below.
• pre-commit run --files with every changed file listed passes.

Before this change the new tests fail. The two files that use T.cuda.sm89 do not collect, the sm_89 typing case is refused, the static assert message differs, and the schedule goldens and candidate reports lack the new atom. Each of the three new refusal cases compiles once the check it pins (B's element type, C's element type, B's length) is removed from atom_shaped_v.

A/B on a single sm_89 GPU

pytest tests	main	this change
collected	2176	2184
passed	2142	2149
failed	21	22
skipped	13	13

The new nodes are the sm80 and sm89 CuTe comparisons (replacing one unparametrized test), the sm_89 typing case, the facts describe T.cuda.sm89.Mma case, the three refusal cases, and the E4M3 GPU test for each target. All of them pass except the H200 node of the E4M3 GPU test, which fails on this GPU like every kernel built for sm_90, with no kernel image is available for execution on the device.

The other 21 failures are the same tests on both sides: 19 launch kernels built for sm_90, and 2 need a second GPU.

[zhen8838]

LGTM