frost(sdpa): port #585 LPT schedulers to the SM107 fp8 sibling; defaults must obey declared knob domains

python/cudnn/frost/README.md

@@ -559,6 +559,17 @@ If a kernel cannot run the requested scheduler policy, the answer is "this
 engine cannot serve this plan", not "ran with a different policy". A knob
 object of the wrong operation's type is rejected outright.
 
+**Defaults obey the same domains as requests.** When no value is requested,
+the adapter's defaulting policy (`_causal_sched_policy` and friends) picks
+one -- and that pick must lie inside the served route's declared domain,
+falling back to the universal default otherwise. A default the kernel cannot
+decode is a plan-build failure, not a preference: issue #653 (the LPT
+scheduler landed in one FP8 kernel sibling but the auto causal policy chose
+it for both) is exactly the failure this rule closes. The kernel files
+declare their decode domain (`SUPPORTED_SCHED_POLICIES`), the adapter's
+routing table mirrors it, the engine rows must not declare beyond it, and
+`test_sdpa_fp8_sibling_parity` pins the three in lockstep.
+
 Generic discoverability survives without the enum: knob domains are ordinary
 dataclass fields on `Capabilities`, so "list every engine and the knobs it
 honors" is a `dataclasses.fields()` walk over the spec table.

python/cudnn/sdpa/fwd/api_dsl.py

@@ -128,6 +128,21 @@ def _torch_stream_context(current_stream: Optional[cuda.CUstream], device: torch
 _SCHED_L2_BUDGET_BYTES = 50 * 1024 * 1024
 
 
+def _sm100_sched_domain(rubin: bool) -> frozenset:
+    """Tile-scheduler policies the routed kernel file can DECODE.
+
+    One set per route (the SM107 sibling is hunk-symmetric with the SM100
+    kernel since the #585 port); kept as a function so a future divergence is
+    a one-line change HERE plus its kernel constant — the sibling-parity test
+    asserts this table equals each file's SUPPORTED_SCHED_POLICIES. Both the
+    explicit-request gate (check_support) and the defaulting clamp (compile)
+    read it: a default may never be a value the route could not honor as an
+    explicit request.
+    """
+    del rubin  # same domain on every SM100-family route today
+    return frozenset({SCHED_NATURAL, SCHED_LPT, SCHED_LPT_L2})
+
+
 def _causal_sched_policy(s_kv: int, d_qk: int, d_v: int, elem_bytes: int) -> int:
     """SCHED_LPT_L2 vs SCHED_LPT for a causal graph (see _SCHED_L2_BUDGET_BYTES)."""
     one_head_bytes = int(s_kv) * (int(d_qk) + int(d_v)) * int(elem_bytes)
@@ -801,8 +816,8 @@ def check_support(self) -> bool:
         )
         self.flavor = _pick_flavor(d_qk, d_v)
         self._value_error_if(
-            self.sched_policy != SCHED_NATURAL,
-            f"SM100 DSL SDPA only supports sched_policy={SCHED_NATURAL}",
+            self.sched_policy not in _sm100_sched_domain(self._device_cc == (10, 7)),
+            f"SM100 DSL SDPA serves sched_policy in {sorted(_sm100_sched_domain(self._device_cc == (10, 7)))}; got {self.sched_policy}",
         )
         for requested, supported, name in (
             (self.tile_m, 128, "tile_m"),
@@ -936,6 +951,12 @@ def compile(self) -> None:
                 d_v=d_v_sched,
                 elem_bytes=1 if self._fp8 else 2,
             )
+            # Defaults must stay inside the served route's declared domain —
+            # a heuristic choice the kernel cannot decode is a plan-build
+            # failure, not a preference (this clamp is what turned #653's
+            # failure mode into a clean NATURAL fallback).
+            if sched_policy not in _sm100_sched_domain(self._device_cc == (10, 7)):
+                sched_policy = SCHED_NATURAL
         params = Sm100TemplateParams(
             dtype_qkv=_SM100_DTYPE_QKV_CODE[self.dtype],
             dtype_o=_SM100_DTYPE_QKV_CODE[self.dtype_o],

python/cudnn/sdpa/fwd/engines.py

@@ -32,7 +32,7 @@
 
 import cudnn
 
-from cudnn.frost.tile_dsl.constants import SCHED_NATURAL
+from cudnn.frost.tile_dsl.constants import SCHED_LPT, SCHED_LPT_L2, SCHED_NATURAL
 from cudnn.frost.buffers import CUTEDSL_MIN_VERSION, cutedsl_state, cutedsl_too_old
 from cudnn.sdpa import graph_analyzer as ga
 
@@ -437,7 +437,7 @@ def _sm100_spec(d: int, d_v: Optional[int] = None) -> EngineSpec:
             # FP8/MXFP8 rows stay on the strict BSHD gate until their padded /
             # scale-factor paths are validated against relaxed layouts.
             layouts=frozenset({"bshd", "dense_flex"}),
-            sched_policies=frozenset({SCHED_NATURAL}),
+            sched_policies=frozenset({SCHED_NATURAL, SCHED_LPT, SCHED_LPT_L2}),
             tile_ms=frozenset({128}),
             tile_ns=frozenset({128}),
             cgas=frozenset({2}),
@@ -472,7 +472,7 @@ def _sm100_mxfp8_spec(d: int) -> EngineSpec:
             sink=True,
             stats=True,
             lse_optional=True,
-            sched_policies=frozenset({SCHED_NATURAL}),
+            sched_policies=frozenset({SCHED_NATURAL, SCHED_LPT, SCHED_LPT_L2}),
             tile_ms=frozenset({128}),
             tile_ns=frozenset({128}),
             cgas=frozenset({2}),
@@ -522,7 +522,7 @@ def _sm100_fp8_spec(d: int) -> EngineSpec:
             # race was fixed with the mb_stats_read barrier (verified on the
             # gated 132/192/200-cluster repros, 3x each).
             skv_tail_via_padding=True,
-            sched_policies=frozenset({SCHED_NATURAL}),
+            sched_policies=frozenset({SCHED_NATURAL, SCHED_LPT, SCHED_LPT_L2}),
             tile_ms=frozenset({128}),
             tile_ns=frozenset({128}),
             cgas=frozenset({2}),

python/cudnn/sdpa/fwd/kernels/prefill_d128_fp8_sm100.py

@@ -47,12 +47,17 @@
 
 from dataclasses import dataclass
 
-from cudnn.sdpa.fwd.config_sm100 import TemplateParams, make_cfg_d128
+from cudnn.sdpa.fwd.config_sm100 import SCHED_LPT, SCHED_LPT_L2, SCHED_NATURAL, TemplateParams, make_cfg_d128
 
 # The template loader (api_dsl._load_kernel_module) injects FROST_TEMPLATE_PARAMS
 # as a module global before this body runs; the default keeps direct import usable.
 PARAMS: TemplateParams = globals().get("FROST_TEMPLATE_PARAMS", TemplateParams())
 CFG, _TMA = make_cfg_d128(PARAMS)
+
+# Tile-scheduler policies this kernel file DECODES. The adapter's defaulting
+# heuristic must choose within this set and the engine row must not declare
+# beyond it — test_sdpa_fp8_sibling_parity pins all three in lockstep.
+SUPPORTED_SCHED_POLICIES = frozenset({SCHED_NATURAL, SCHED_LPT, SCHED_LPT_L2})
 Cfg = type(CFG)
 TMA_QK_ITERS = _TMA.QK_ITERS
 TMA_VO_ITERS = _TMA.VO_ITERS

python/cudnn/sdpa/fwd/kernels/prefill_d128_fp8_sm107.py

@@ -60,12 +60,17 @@
 
 from dataclasses import dataclass
 
-from cudnn.sdpa.fwd.config_sm100 import TemplateParams, make_cfg_d128
+from cudnn.sdpa.fwd.config_sm100 import SCHED_LPT, SCHED_LPT_L2, SCHED_NATURAL, TemplateParams, make_cfg_d128
 
 # The template loader (api_dsl._load_kernel_module) injects FROST_TEMPLATE_PARAMS
 # as a module global before this body runs; the default keeps direct import usable.
 PARAMS: TemplateParams = globals().get("FROST_TEMPLATE_PARAMS", TemplateParams())
 CFG, _TMA = make_cfg_d128(PARAMS)
+
+# Tile-scheduler policies this kernel file DECODES. The adapter's defaulting
+# heuristic must choose within this set and the engine row must not declare
+# beyond it — test_sdpa_fp8_sibling_parity pins all three in lockstep.
+SUPPORTED_SCHED_POLICIES = frozenset({SCHED_NATURAL, SCHED_LPT, SCHED_LPT_L2})
 # Rubin geometry, baked post-validation (this module is only ever loaded for
 # cc10.7 by the adapter): dense-FP8 K=64 steps and the 9-stage KV ring. The
 # TMA iteration constants depend only on TILE_K/TILE_O/BPE/swizzle, so _TMA
@@ -488,6 +493,7 @@ def _kernel(
             n_batch=n_batch,
             leader_cta_id=leader_cta_id,
             cta_in_pair=cta_in_pair,
+            qh_per_kh=qh_per_kh,
         )
 
     elif warp_idx >= CFG.SOFTMAX_WG1_BASE and warp_idx < CFG.SOFTMAX_WG1_BASE + CFG.SOFTMAX_WG_WARPS:
@@ -507,6 +513,7 @@ def _kernel(
             n_batch=n_batch,
             leader_cta_id=leader_cta_id,
             cta_in_pair=cta_in_pair,
+            qh_per_kh=qh_per_kh,
         )
 
     elif warp_idx >= CFG.CORR_WARP_BASE and warp_idx < CFG.CORR_WARP_BASE + CFG.CORRECTION_WARPS:
@@ -530,6 +537,7 @@ def _kernel(
             cta_id_x=cta_id_x,
             o_scale_fused=o_scale_fused,
             amax_o_tensor=amax_o_tensor,
+            qh_per_kh=qh_per_kh,
         )
 
     # cga2 non-leader runs quiet body (alloc+dealloc only); cga1 folds to full path.
@@ -553,6 +561,7 @@ def _kernel(
                     n_batch=n_batch,
                     mcast_mask=mcast_mask,
                     cta_in_pair=cta_in_pair,
+                    qh_per_kh=qh_per_kh,
                 )
             else:
                 _mma_warp_quiet(tmem_ptr_i32, bars)
@@ -573,6 +582,7 @@ def _kernel(
                 n_batch=n_batch,
                 mcast_mask=mcast_mask,
                 cta_in_pair=cta_in_pair,
+                qh_per_kh=qh_per_kh,
             )
 
     elif warp_idx == CFG.TMALDG_WARP_ID:
@@ -614,6 +624,8 @@ def _kernel(
             n_batch=n_batch,
             cta_in_pair=cta_in_pair,
             seq_kv_lens_tensor=seq_kv_lens_tensor,
+            seqlen_kv=seqlen_kv,
+            qh_per_kh=qh_per_kh,
         )
 
     else:  # warp_idx == CFG.SCHED_WARP_ID
@@ -669,6 +681,8 @@ def _tmaldg_warp_group(
         n_qh,
         n_batch,
         seq_kv_lens_tensor,
+        qh_per_kh,
+        seqlen_kv,
     )
     # GQA: K/V are indexed by kv-head.
     kv_head_idx = cute.arch.make_warp_uniform(head_idx // qh_per_kh)
@@ -802,6 +816,8 @@ def _tmaldg_warp_group(
             n_qh,
             n_batch,
             seq_kv_lens_tensor,
+            qh_per_kh,
+            seqlen_kv,
         )
         kv_head_idx = cute.arch.make_warp_uniform(head_idx // qh_per_kh)
         # q_row_base after decode drives ptxas R2UR (keeps nxt_q live before back-edge).
@@ -842,6 +858,8 @@ def _tmastg_warp_group(
     n_batch,
     cta_in_pair,
     seq_kv_lens_tensor,
+    seqlen_kv,
+    qh_per_kh,
 ):
     """Persistent O-store warp; tiles claimed via scheduler's try_cancel.async."""
     o_full_phase = cutlass.Int32(0)
@@ -857,6 +875,8 @@ def _tmastg_warp_group(
         n_qh,
         n_batch,
         seq_kv_lens_tensor,
+        qh_per_kh,
+        seqlen_kv,
     )
     is_valid_tile = cutlass.Int32(1)
     sched_state = PipelineState.start()
@@ -894,6 +914,8 @@ def _tmastg_warp_group(
             n_qh,
             n_batch,
             seq_kv_lens_tensor,
+            qh_per_kh,
+            seqlen_kv,
         )
         is_valid_tile = nxt_v & cutlass.Int32(1)
         sched_state = advance(sched_state, CFG.SCHEDULER_STAGES)
@@ -960,6 +982,7 @@ def _mma_warp_group(
     n_batch,
     mcast_mask,
     cta_in_pair,
+    qh_per_kh,
 ):
     """Unified MMA warp (cga1 / cga2-leader; MASK_NONE/PADDED/CAUSAL/SWA).
 
@@ -1069,6 +1092,8 @@ def _mma_warp_group(
             n_qh,
             n_batch,
             seq_kv_lens_tensor,
+            qh_per_kh,
+            seqlen_kv,
         )
         eff_seqlen_kv = _resolve_seqlen_kv(seq_kv_lens_tensor, batch_idx, seqlen_kv)
         eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch)
@@ -1268,6 +1293,8 @@ def _mma_warp_group(
                 n_qh,
                 n_batch,
                 seq_kv_lens_tensor,
+                qh_per_kh,
+                seqlen_kv,
             )
             is_valid_tile = nxt_v & cutlass.Int32(1)
             eff_seqlen_kv = _resolve_seqlen_kv(seq_kv_lens_tensor, batch_idx, seqlen_kv)
@@ -1462,6 +1489,7 @@ def _softmax_warp_group(
     n_batch,
     leader_cta_id,
     cta_in_pair,
+    qh_per_kh,
 ):
     """Softmax warp group: online softmax per kv iter, one lane per S_acc row.
 
@@ -1503,6 +1531,8 @@ def _softmax_warp_group(
         n_qh,
         n_batch,
         seq_kv_lens_tensor,
+        qh_per_kh,
+        seqlen_kv,
     )
     is_valid_tile = cutlass.Int32(1)
     sched_state = PipelineState.start()
@@ -1635,6 +1665,8 @@ def _softmax_warp_group(
             n_qh,
             n_batch,
             seq_kv_lens_tensor,
+            qh_per_kh,
+            seqlen_kv,
         )
         is_valid_tile = nxt_v & cutlass.Int32(1)
         sched_state = advance(sched_state, CFG.SCHEDULER_STAGES)
@@ -1663,6 +1695,7 @@ def _correction_warp_group(
     cta_id_x,
     o_scale_fused,
     amax_o_tensor,
+    qh_per_kh,
 ):
     """Correction warp group: 4 warps × 32 lanes = 128, one lane per O row.
 
@@ -1706,6 +1739,8 @@ def _correction_warp_group(
         n_qh,
         n_batch,
         seq_kv_lens_tensor,
+        qh_per_kh,
+        seqlen_kv,
     )
     is_valid_tile = cutlass.Int32(1)
     sched_state = PipelineState.start()
@@ -1968,6 +2003,8 @@ def _correction_warp_group(
             n_qh,
             n_batch,
             seq_kv_lens_tensor,
+            qh_per_kh,
+            seqlen_kv,
         )
         is_valid_tile = nxt_v & cutlass.Int32(1)
         sched_state = advance(sched_state, CFG.SCHEDULER_STAGES)