Add SM120 per-tensor FP8 (e4m3) SDPA-forward engine

include/cudnn_frontend/node/sdpa_support_surface.h

@@ -162,10 +162,6 @@ SDPA_attributes::validate_sdpa_support_surface(const detail::Context& context,
     if (mma_core_mode == DataType_t::FP8_E4M3 || mma_core_mode == DataType_t::FP8_E5M2) {
         // FP8 specific validation
 
-        RETURN_CUDNN_FRONTEND_ERROR_IF((prop_major == 12) && is_ragged,
-                                       error_code_t::GRAPH_NOT_SUPPORTED,
-                                       "sdpa fp8 with THD not supported for sm120 yet.");
-
         // version specific validation
         RETURN_CUDNN_FRONTEND_ERROR_IF(detail::get_backend_version() < 90100,
                                        error_code_t::GRAPH_NOT_SUPPORTED,

python/cudnn/engines/manifest.py

@@ -179,6 +179,7 @@ def offered_ids(self) -> Dict[str, int]:
             "sdpa_fwd_prefill_sm100_d128_fp8": EngineSlot(4, opt_in=True),
             "sdpa_fwd_prefill_sm120": EngineSlot(5, opt_in=True),
             "sdpa_fwd_prefill_sm100_d192_d128": EngineSlot(6, opt_in=True),
+            "sdpa_fwd_prefill_sm120_fp8": EngineSlot(7, opt_in=True),
         },
         analyzer=("cudnn.sdpa.graph_analyzer", "analyze"),
         heuristics=("cudnn.sdpa.fwd.heuristics", "recommend"),

python/cudnn/frost/tile_dsl/mma.py

@@ -42,6 +42,80 @@ def ptx_mma_m16n8k16_f32(
     )
 
 
+@cute.jit
+def ptx_mma_m16n8k32_e4m3_f32(
+    a0: cutlass.Int32,
+    a1: cutlass.Int32,
+    a2: cutlass.Int32,
+    a3: cutlass.Int32,
+    b0: cutlass.Int32,
+    b1: cutlass.Int32,
+    c0: cutlass.Float32,
+    c1: cutlass.Float32,
+    c2: cutlass.Float32,
+    c3: cutlass.Float32,
+) -> tuple[cutlass.Float32, cutlass.Float32, cutlass.Float32, cutlass.Float32]:
+    """``mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32`` (SM89+/SM120).
+
+    The C operands travel as ``Int32`` bit patterns and are ``mov.b32``'d into
+    ``.f32`` temps inside the asm block: cutlass-dsl 4.7.0's ``inline_ptx``
+    fails in libNVVM when a compile-time-constant ``Float32`` reaches
+    ``read_only_args``, and an accumulator's zero-init can fold to a constant.
+    """
+    return cute.arch.inline_ptx(
+        "{ .reg .f32 fc<4>; "
+        "mov.b32 fc0, {$r6}; mov.b32 fc1, {$r7}; mov.b32 fc2, {$r8}; mov.b32 fc3, {$r9}; "
+        "mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 "
+        "{{$w0},{$w1},{$w2},{$w3}}, {{$r0},{$r1},{$r2},{$r3}}, {{$r4},{$r5}}, {fc0,fc1,fc2,fc3}; }",
+        write_only_types=[
+            cutlass.Float32,
+            cutlass.Float32,
+            cutlass.Float32,
+            cutlass.Float32,
+        ],
+        read_only_args=[
+            a0,
+            a1,
+            a2,
+            a3,
+            b0,
+            b1,
+            c0.bitcast(cutlass.Int32),
+            c1.bitcast(cutlass.Int32),
+            c2.bitcast(cutlass.Int32),
+            c3.bitcast(cutlass.Int32),
+        ],
+    )
+
+
+@cute.jit
+def ptx_cvt_e4m3x2(hi: cutlass.Float32, lo: cutlass.Float32) -> cutlass.Uint16:
+    """Pack two fp32 into e4m3 bytes: low byte = e4m3(lo), byte 1 = e4m3(hi).
+
+    ``cvt.rn.satfinite.e4m3x2.f32`` matches torch's ``.to(float8_e4m3fn)``
+    bit-exactly. Operands ride as Int32 bit patterns for the same
+    constant-operand ``inline_ptx`` reason as :func:`ptx_mma_m16n8k32_e4m3_f32`.
+
+    Stays 16-bit so two results pair into one MMA operand register with
+    :func:`pack_f8x2_pairs`.
+    """
+    return cute.arch.inline_ptx(
+        "{ .reg .f32 fa, fb; " "mov.b32 fa, {$r0}; mov.b32 fb, {$r1}; " "cvt.rn.satfinite.e4m3x2.f32 {$w0}, fa, fb; }",
+        write_only_types=[cutlass.Uint16],
+        read_only_args=[hi.bitcast(cutlass.Int32), lo.bitcast(cutlass.Int32)],
+    )
+
+
+@cute.jit
+def pack_f8x2_pairs(pair0: cutlass.Uint16, pair1: cutlass.Uint16) -> cutlass.Int32:
+    """Two e4m3x2 halves into one 32-bit MMA A/B operand (pair0 = low half)."""
+    return cute.arch.inline_ptx(
+        "mov.b32 $0, {$1, $2};",
+        write_only_types=[cutlass.Int32],
+        read_only_args=[pair0, pair1],
+    )
+
+
 @cute.jit
 def mma_ss(desc, desc_a_base, desc_b_base, tmem_c, tmem_sf_a=None, tmem_sf_b=None, accumulate: bool = False, k_start: int = 0, k_count=None):
     if cutlass.const_expr(desc.cta_group == 1):