Add SM120 FROST SDPA backward engine (sdpa_bwd_sm120)

python/cudnn/engines/engine_ids.py

@@ -33,7 +33,7 @@
 LINEAR_ATTENTION_ID_BASE = PYTHON_ENGINE_ID_BASE + 100  # 20_100..20_199
 FROST_GEMM_ID_BASE = PYTHON_ENGINE_ID_BASE + 200  # 20_200..20_299
 FROST_SDPA_FWD_ID_BASE = PYTHON_ENGINE_ID_BASE + 300  # 20_300..20_399
-FROST_SDPA_BWD_ID_BASE = PYTHON_ENGINE_ID_BASE + 400  # reserved
+FROST_SDPA_BWD_ID_BASE = PYTHON_ENGINE_ID_BASE + 400  # 20_400..20_499
 OUT_OF_TREE_ID_BASE = PYTHON_ENGINE_ID_BASE + 10_000  # 30_000+, register_backend()
 
 # The delegating entry: the backend picks among candidates it holds but does not

python/cudnn/engines/manifest.py

@@ -43,7 +43,7 @@
 from dataclasses import dataclass
 from typing import Any, Dict, Optional, Tuple
 
-from .engine_ids import FROST_GEMM_ID_BASE, FROST_SDPA_FWD_ID_BASE, LINEAR_ATTENTION_ID_BASE
+from .engine_ids import FROST_GEMM_ID_BASE, FROST_SDPA_BWD_ID_BASE, FROST_SDPA_FWD_ID_BASE, LINEAR_ATTENTION_ID_BASE
 
 _LOG = logging.getLogger("cudnn.engines.manifest")
 
@@ -101,6 +101,7 @@ def matches(self, node_types: frozenset, sm: Optional[int]) -> bool:
 _GEMM_ANCHOR = frozenset({"MATMUL", "MATMUL_FP8", "MOE_GROUPED_MATMUL"})
 _GEMM_CLOSURE = _GEMM_ANCHOR | frozenset({"POINTWISE", "REDUCTION", "RESHAPE", "BLOCK_SCALE_QUANTIZE", "BLOCK_SCALE_DEQUANTIZE"})
 _SDPA_FWD = frozenset({"SDPA", "SDPA_FP8", "SDPA_MXFP8"})
+_SDPA_BWD = frozenset({"SDPA_BWD"})
 _GDN = frozenset({"GDN", "GDN_BWD"})
 _GDN2 = frozenset({"GDN2", "GDN2_BWD"})
 _KDA = frozenset({"KDA", "KDA_BWD"})
@@ -176,6 +177,18 @@ def matches(self, node_types: frozenset, sm: Optional[int]) -> bool:
         sm_lo=100,
         opt_in=True,
     ),
+    EngineRow(
+        FROST_SDPA_BWD_ID_BASE + 0,
+        "frost_sdpa_bwd",
+        "cudnn.sdpa.bwd.engine",
+        "FrostSdpaBwdEngines",
+        _SDPA_BWD,
+        id_hi=FROST_SDPA_BWD_ID_BASE + 100,
+        # TODO: widen when an SM100/SM80 spec lands
+        sm_lo=120,
+        sm_hi=121,
+        opt_in=True,
+    ),
 )
 
 

python/cudnn/frost/tile_dsl/mma.py

@@ -12,6 +12,36 @@
 from .swizzle import swizzle_xor_128b, swizzle_lin_128b
 
 
+@cute.jit
+def ptx_mma_m16n8k16_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,
+    ab_dtype: cutlass.Constexpr[Type[cutlass.Numeric]],
+) -> tuple[cutlass.Float32, cutlass.Float32, cutlass.Float32, cutlass.Float32]:
+    """``mma.sync.aligned.m16n8k16.row.col.f32.{f16|bf16}.{f16|bf16}.f32``."""
+    if cutlass.const_expr(ab_dtype != cutlass.Float16 and ab_dtype != cutlass.BFloat16):
+        raise TypeError(f"Invalid A/B dtype: {ab_dtype}")
+    ab_tag = "f16" if cutlass.const_expr(ab_dtype == cutlass.Float16) else "bf16"
+    return cute.arch.inline_ptx(
+        f"mma.sync.aligned.m16n8k16.row.col.f32.{ab_tag}.{ab_tag}.f32 {{$0,$1,$2,$3}}, {{$4,$5,$6,$7}}, {{$8,$9}}, {{$10,$11,$12,$13}};",
+        write_only_types=[
+            cutlass.Float32,
+            cutlass.Float32,
+            cutlass.Float32,
+            cutlass.Float32,
+        ],
+        read_only_args=[a0, a1, a2, a3, b0, b1, c0, c1, c2, c3],
+    )
+
+
 @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):

python/cudnn/frost/tile_dsl/swizzle.py

@@ -15,6 +15,49 @@ def swizzle_xor_128b(row, col_elem, *, elem_bytes: cutlass.Constexpr[int] = 2):
     return swz_chunk * chunk_elems + in_chunk
 
 
+@cute.jit
+def swizzle_xor_64b(row, col_elem, *, elem_bytes: cutlass.Constexpr[int] = 2):
+    chunk_elems = 16 // elem_bytes
+    chunk_idx = col_elem // chunk_elems
+    in_chunk = col_elem % chunk_elems
+    swz_chunk = chunk_idx ^ ((row >> 1) & 3)
+    return swz_chunk * chunk_elems + in_chunk
+
+
+@cute.jit
+def swizzle_xor_32b(row, col_elem, *, elem_bytes: cutlass.Constexpr[int] = 2):
+    chunk_elems = 16 // elem_bytes
+    chunk_idx = col_elem // chunk_elems
+    in_chunk = col_elem % chunk_elems
+    swz_chunk = chunk_idx ^ ((row >> 2) & 1)
+    return swz_chunk * chunk_elems + in_chunk
+
+
+@cute.jit
+def swizzle_xor(
+    row: cutlass.Int32,
+    col: cutlass.Int32,
+    row_stride: cutlass.Constexpr[int],
+    elem_bytes: cutlass.Constexpr[int],
+) -> cutlass.Int32:
+    """Return the physical SMEM column for an XOR-swizzled row-major tile.
+
+    The XOR is applied at the 16-byte boundary for all element widths.
+    ``elem_bytes`` selects the element-domain shift and swizzle chunk size.
+    """
+    row_stride_bytes = cutlass.const_expr(row_stride * elem_bytes)
+    if cutlass.const_expr(row_stride_bytes % 128 == 0):
+        chunk_elems = 128 // elem_bytes
+        swizzled = swizzle_xor_128b(row, col % chunk_elems, elem_bytes=elem_bytes)
+    elif cutlass.const_expr(row_stride_bytes % 64 == 0):
+        chunk_elems = 64 // elem_bytes
+        swizzled = swizzle_xor_64b(row, col % chunk_elems, elem_bytes=elem_bytes)
+    else:
+        chunk_elems = 32 // elem_bytes
+        swizzled = swizzle_xor_32b(row, col % chunk_elems, elem_bytes=elem_bytes)
+    return (col // chunk_elems) * chunk_elems + swizzled
+
+
 @cute.jit
 def swizzle_lin_128b(lin, *, row_stride_log2: cutlass.Constexpr[int], elem_bytes: cutlass.Constexpr[int] = 2):
     chunk_log2 = cutlass.const_expr((16 // elem_bytes).bit_length() - 1)

python/cudnn/sdpa/bwd/__init__.py

@@ -2,8 +2,12 @@
 # SPDX-License-Identifier: Apache-2.0
 
 from .api import SdpabwdSm100D256, sdpa_bwd_wrapper_sm100_d256
+from .api_dsl import SdpaBwdDsl, SdpaBwdDslSm120, sdpa_bwd_wrapper_dsl_sm120
 
 __all__ = [
     "SdpabwdSm100D256",
     "sdpa_bwd_wrapper_sm100_d256",
+    "SdpaBwdDsl",
+    "SdpaBwdDslSm120",
+    "sdpa_bwd_wrapper_dsl_sm120",
 ]