CuTe DSL: FP8 MMA segfaults on SM120 — MmaAtomSM80Type missing kind::f8f6f4 lowering

[blake-snc]

Summary

Using MmaAtomSM80Type.get() with FP8 types (Float8E4M3FN, shape (16, 8, 32)) causes a segfault during MLIR→PTX compilation when targeting SM120 (RTX 5090, GB10/DGX Spark, sm_121a). The MLIR backend appears to lack lowering support for SM120's mma.sync.aligned.kind::f8f6f4.m16n8k32 instruction variant.

Background

SM120 has an FP8 MMA instruction defined in C++ CuTe:

// cute/arch/mma_sm120.hpp lines 668-697
template <>
struct SM120_16x8x32_TN<float_e4m3_t, float_e4m3_t, float> {
    // PTX: mma.sync.aligned.kind::f8f6f4.m16n8k32.row.col.f32.e4m3.e4m3.f32
};

This is a different instruction from SM89's FP8 MMA (mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 — no kind:: prefix). SM120 has CUTE_ARCH_F8F6F4_MMA_ENABLED but not CUTE_ARCH_MMA_F32_SM89_ENABLED (per cute/arch/config.hpp).

The CuTe Python DSL currently only exposes:

• MmaF16BF16Op → MmaAtomSM80Type with F16/BF16 (works on SM120)
• MmaSM120BlockScaledOp → MmaAtomSM120BlockScaledType with FP4 block-scaled (works on SM120)

There is no exposed path for SM120's non-block-scaled FP8 MMA.

Reproducer

Bypassing MmaF16BF16Op's Python-level validation and calling MmaAtomSM80Type.get() directly with FP8 types:

import cutlass
from cutlass import cute
from cutlass.cute.nvgpu.warp.mma import WarpMmaOp, MmaF16BF16Trait, make_atom, _pack_shape
from cutlass.cute.typing import Float8E4M3FN, Float32
import cutlass._mlir.dialects.cute_nvgpu as _cute_nvgpu_ir
from dataclasses import dataclass
from typing import Type

@dataclass(frozen=True)
class MmaFP8Op(WarpMmaOp):
    ab_dtype: Type[cutlass.Numeric] = Float8E4M3FN
    acc_dtype: Type[cutlass.Numeric] = Float32
    shape_mnk: tuple = (16, 8, 32)

    def _make_trait(self, *, loc=None, ip=None, **kwargs):
        shape_mnk = _pack_shape(self.shape_mnk, loc=loc, ip=ip)
        ty = _cute_nvgpu_ir.MmaAtomSM80Type.get(
            shape_mnk.type.attribute,
            self.ab_dtype.mlir_type,
            self.ab_dtype.mlir_type,
            self.acc_dtype.mlir_type,
        )
        return MmaF16BF16Trait(make_atom(ty, loc=loc, ip=ip))

    def _verify_fragment_A(self, input, *, loc=None, ip=None): pass
    def _verify_fragment_B(self, input, *, loc=None, ip=None): pass

The MLIR type creation succeeds — a kernel that creates the atom without using it compiles fine. However, when the MMA is actually used in a compute path (cute.gemm), the MLIR→PTX lowering segfaults:

!!!!!!! Segfault encountered !!!!!!!
  File "./nptl/pthread_once.c", line 116, in __pthread_once_slow

Behavior

Step	Result
MmaAtomSM80Type.get(shape=(16,8,32), aType=f8E4M3FN, bType=f8E4M3FN, cType=f32)	Success (MLIR type created)
Create tiled_mma from the atom	Success
cute.gemm(tiled_mma, acc, fragA, fragB, acc) in kernel body	Segfault during cute.compile()

Expected Behavior

MmaAtomSM80Type.get() with FP8 types should either:

1. Lower to SM120's mma.sync.aligned.kind::f8f6f4.m16n8k32 PTX instruction when targeting sm_120a/sm_121a
2. Or raise a clear error (not segfault) indicating FP8 MMA is unsupported for the target arch

Environment

• Device: NVIDIA GB10 (DGX Spark), SM121a
• CUDA: 13.0
• nvidia-cutlass-dsl: 4.3.5 (pip)
• Target arch: CUTE_DSL_ARCH=sm_121a
• Python: 3.13.3

Impact

SM120's non-block-scaled FP8 MMA (m16n8k32 kind::f8f6f4) offers 2x theoretical throughput over BF16 (m16n8k16) for attention kernels. This is the highest-impact optimization available for SM120 flash attention, and it's currently blocked by this lowering gap.

The C++ CuTe headers already define the instruction (SM120_16x8x32_TN with 77 specializations for different FP8/FP6/FP4 type combinations). The Python DSL just needs the MLIR lowering to be connected.

---

Contributed by Second Nature Computing — tested on DGX Spark hardware

[blake-snc]

Update: Inline PTX Workaround Validated

While investigating the segfault, we found that llvm.inline_asm() in the CuTe DSL can be used to emit the kind::f8f6f4 instruction directly, bypassing the broken MmaAtomSM80Type lowering:

from cutlass.cutlass_dsl import T, dsl_user_op
from cutlass._mlir import ir
from cutlass._mlir.dialects import llvm
from cutlass.cute.typing import Int32, Float32

@dsl_user_op
def mma_f8f6f4_m16n8k32_e4m3(a0, a1, a2, a3, b0, b1, c0, c1, c2, c3, *, loc=None, ip=None):
    return_struct = llvm.inline_asm(
        ir.Type.parse("!llvm.struct<(f32,f32,f32,f32)>"),
        [Int32(a0).ir_value(loc=loc, ip=ip), Int32(a1).ir_value(loc=loc, ip=ip),
         Int32(a2).ir_value(loc=loc, ip=ip), Int32(a3).ir_value(loc=loc, ip=ip),
         Int32(b0).ir_value(loc=loc, ip=ip), Int32(b1).ir_value(loc=loc, ip=ip),
         Float32(c0).ir_value(loc=loc, ip=ip), Float32(c1).ir_value(loc=loc, ip=ip),
         Float32(c2).ir_value(loc=loc, ip=ip), Float32(c3).ir_value(loc=loc, ip=ip)],
        "mma.sync.aligned.kind::f8f6f4.m16n8k32.row.col.f32.e4m3.e4m3.f32 "
        "{$0, $1, $2, $3}, {$4, $5, $6, $7}, {$8, $9}, {$10, $11, $12, $13};",
        "=f,=f,=f,=f,r,r,r,r,r,r,f,f,f,f",
        has_side_effects=True, asm_dialect=0, loc=loc, ip=ip,
    )
    return tuple(Float32(llvm.extractvalue(T.f32(), return_struct, [i], loc=loc, ip=ip)) for i in range(4))

Validated on GB10 (SM121a): Compilation succeeds, execution produces correct results (32.0 for K=32 all-ones inputs).

This confirms the instruction itself works fine on SM120 — the issue is specifically in the MmaAtomSM80Type MLIR lowering path, which doesn't know how to emit kind::f8f6f4 when targeting SM120.

We'll use this inline PTX approach to build an FP8 flash attention kernel for SM120 in PR #3030 while the proper MLIR lowering fix is pending.

---

Contributed by Second Nature Computing

[shubaoyu2]

did you try MmaAtomSM89Type instead of MmaAtomSM80Type? F8 is not invented in sm80

[blake-snc]

Good question — I tried this. Unfortunately MmaAtomSM89Type doesn't exist in the CuTe DSL. The only available MLIR MMA types are:

• MmaAtomSM80Type (warp-level F16/BF16, m16n8k8/m16n8k16)
• MmaAtomSM90Type (warpgroup WGMMA)
• MmaAtomSM100UMMAType / MmaAtomSM100UMMABlockScaledType (tcgen05)
• MmaAtomSM120BlockScaledType (block-scaled FP4/FP8)

>>> from cutlass._mlir.dialects import cute_nvgpu as ir
>>> [a for a in dir(ir) if 'MmaAtom' in a]
['MmaAtomSM100UMMABlockScaledType', 'MmaAtomSM100UMMAType',
 'MmaAtomSM120BlockScaledType', 'MmaAtomSM80Type', 'MmaAtomSM90Type',
 'MmaAtomTypeInterface']

The C++ SM89 atoms exist in cute/arch/mma_sm89.hpp (SM89_16x8x32_F32E4M3E4M3F32_TN etc.) but they're not exposed in the Python DSL layer.

Also worth noting: the SM89 atoms emit mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32, which is a different instruction from SM120's mma.sync.aligned.kind::f8f6f4.m16n8k32 — so even if they were exposed, they wouldn't test the kind::f8f6f4 lowering path. The inline PTX workaround from the earlier comment remains the only way to access this instruction from the DSL today.

---

Contributed by Second Nature Computing

[shubaoyu2]

Got it, thanks. We need to make some change to our close sourced code of cutedsl, and maybe bring it out with 4.5 version. And you're right, mixed precision f8 dense mma is a SM120 thing not SM89

[shubaoyu2]

Hi @blake-snc , I want to make sure that you do need kind::f8f6f4, right? SM89 can support things like e4m3xe5m2, the kind::f8f6f4 on SM120 is for things like f4xf8 / f6xf8 / f8xf8 mma.

[blake-snc]

did you try MmaAtomSM89Type instead of MmaAtomSM80Type? F8 is not invented in sm80

You're right — I should correct my earlier claim. For e4m3×e4m3 specifically, SM120 does support the SM89-style instruction via backward compatibility. I validated both instructions side-by-side on GB10 (SM121a):

$ nvcc -gencode arch=compute_121a,code=sm_121a -o test test_sm89_vs_f8f6f4.cu && ./test

Device: NVIDIA GB10 (SM 121)

=== Test 1: All-ones (K=32, thread 0) ===
  SM89 instruction:     [32.0, 32.0, 32.0, 32.0]
  kind::f8f6f4:         [32.0, 32.0, 32.0, 32.0]
  Match: YES

=== Test 2: Varied data (thread 0) ===
  SM89 instruction:     [47.0000, 48.0000, 49.0000, 50.0000]
  kind::f8f6f4:         [47.0000, 48.0000, 49.0000, 50.0000]
  Match: YES

=== Test 3: All-ones, all 32 warp threads ===
  All 128 outputs match: YES (max_diff = 0.000000)

=== SUMMARY ===
SM89-style  (mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32):             WORKS
SM120-style (mma.sync.aligned.kind::f8f6f4.m16n8k32.row.col.f32.e4m3.e4m3.f32): WORKS
Results identical: YES

Both produce identical results across all warp threads with both uniform and varied data.

For our use cases (FP8 flash attention, MLA decode), e4m3×e4m3 is what we need — so an MmaAtomSM89Type in the DSL would fully unblock us on SM120. The mixed-precision combinations unique to kind::f8f6f4 (f8×f6, f4×f8, etc.) would still be locked out of the DSL ecosystem, but SM120's block-scaled path (MmaSM120BlockScaledOp) already covers FP4/FP8 GEMM workloads, so the gap is narrow in practice.

Thanks for pushing on this — happy to test anything else on our GB10 hardware.

---

Contributed by Second Nature Computing

[Alimadcorp]

❤️

[blake-snc]

Closing out the status on this issue: SM89-style FP8 MMA instructions (SM89_16x8x32_F32E4M3E4M3F32_TN etc.) work correctly on SM120/SM121 via backward compatibility — validated side-by-side against kind::f8f6f4 on GB10 (SM121a), both produce identical results.

The original segfault was from calling MmaAtomSM80Type with FP8 types, which is the wrong API — SM80 atoms only cover F16/BF16. The kind::f8f6f4 instruction itself works via inline PTX (llvm.inline_asm()), also validated.

For CuTe DSL usage: MmaAtomSM89Type doesn't exist in the Python DSL layer (no entry in cute_nvgpu dialect), so either inline PTX or the C++ SM89 atoms (via the cute.nvgpu C++ path) are the current options for FP8 MMA on SM120.

[github-actions]

This issue has been labeled inactive-30d due to no recent activity in the past 30 days. Please close this issue if no further response or action is needed. Otherwise, please respond with a comment indicating any updates or changes to the original issue and/or confirm this issue still needs to be addressed. This issue will be labeled inactive-90d if there is no activity in the next 60 days.

[blake-snc]

Reproduced on cutlass-dsl 4.4.2 (current latest)

Re-validated this issue end-to-end on cutlass-dsl 4.4.2 / CUTLASS main f74fea9c on NVIDIA GB10 (DGX Spark, sm_121a) with three distinct reproducers. All three confirm the segfault is still present.

Reproducer 1 — Realistic FP8 attention path

Took the upstream examples/python/CuTeDSL/ampere/flash_attention_v2.py example, substituted MmaFP8Op (per the issue body) for MmaF16BF16Op, relaxed the dtype gates, and widened the SMEM swizzle from M=3 to M=4 to satisfy 128-bit cp.async alignment for FP8 elements. Same kernel structure as the stock F16 example, dtype is the only meaningful change.

>>> cute.compile() with dtype=Float16  → compile OK
>>> cute.compile() with dtype=Float8E4M3FN
!!!!!!! Segfault encountered !!!!!!!
  File "./nptl/pthread_once.c", line 116, in __pthread_once_slow
  File "./nptl/pthread_create.c", line 447, in start_thread
  File "../sysdeps/unix/sysv/linux/aarch64/clone3.S", line 76, in thread_start

Reproducer 2 — Minimal isolated MMA repro (clearest evidence)

Same per-thread register-only kernel structure for F16 and FP8, identical fragment shape construction, only the MMA atom differs (MmaF16BF16Op vs MmaFP8Op). F16 reaches the cute.gemm rank verification and reports a clean MLIR error; FP8 segfaults before any verification message is printed.

>>> [f16 control]  DSLRuntimeError: ICE IR Verification Failed
                   Caused exception: 'cute.gemm' op expects rank 2 or 3 for A, but got: 4
>>> [fp8 minimal repro]
!!!!!!! Segfault encountered !!!!!!!
  File "./nptl/pthread_once.c", line 116, in __pthread_once_slow

This is the strongest signal: identical code path, only the dtype + atom differ, F16 reaches frontend verification while FP8 crashes earlier in the backend.

Reproducer 3 — Diff of cutlass-dsl 4.3.5 → 4.4.2

python/CuTeDSL/cutlass/cute/nvgpu/warp/mma.py shows:

• No new frontend gate rejecting FP8 in MmaAtomSM80Type. The bypass technique from this issue's reproducer still constructs successfully.
• Added MmaSM120BlockScaledOp for block-scaled FP4/FP8 MMA (mma.sync.aligned.block_scale instructions). This is a separate code path from non-block-scaled FP8 (mma.sync.aligned.kind::f8f6f4.m16n8k32).
• Backend binary _mlir/_mlir_libs/_cutlass_ir.cpython-313-aarch64-linux-gnu.so differs but does not appear to fix the FP8 lowering segfault for MmaAtomSM80Type based on the reproducer behavior above.

Adjacent context

CUTLASS PR #3030 demonstrates that the underlying PTX instruction (mma.sync.aligned.kind::f8f6f4.m16n8k32) works correctly on SM120 hardware end-to-end, accessed via inline PTX in C++ CuTe rather than the Python DSL. That confirms the bug is specifically in the Python DSL's MLIR lowering, not in the hardware or the instruction itself, and bounds the scope of the fix.

Net

The lowering gap reported in this issue persists in 4.4.2. Adding MmaSM120BlockScaledOp covers the block-scaled FP8/FP4 path but the non-block-scaled FP8 path through MmaAtomSM80Type still segfaults. SM120's mma.sync.aligned.kind::f8f6f4.m16n8k32 instruction (the highest-impact optimization for SM120 attention at 2x BF16 throughput) remains inaccessible from the Python DSL.

The two reproducers are minimal Python files (~60 and ~100 lines respectively), happy to share if helpful for triage.

---

Re-tested by Second Nature Computing on DGX Spark hardware, 2026-04-25.