From f418fd8251f6aef9fd7ff8b6ed1330084708cc0a Mon Sep 17 00:00:00 2001 From: Vedaanta Agarwalla Date: Tue, 18 Aug 2026 16:59:59 -0700 Subject: [PATCH] frost(sdpa): port the #585 LPT schedulers to the SM107 fp8 sibling; defaults must obey declared knob domains (fixes #653) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two changes that are one rule: a scheduler the served kernel cannot decode must be impossible to reach — whether requested or defaulted. 1. #585 port (fixes #653): the LPT/LPT_L2 tile schedulers landed in prefill_d128_fp8_sm100.py but the hunk-symmetric port to the SM107 sibling was never done, so the adapter's auto causal sched policy picked SCHED_LPT_L2 and every causal/masked per-tensor FP8 graph on cc10.7 failed plan build (29 of 33 suite cases). Ported: qh_per_kh/seqlen_kv threaded through the five warp groups and appended at all ten scheduler decode call sites, mirroring the SM100 file position for position. 2. Domain honesty end to end: - Both fp8 kernel files declare SUPPORTED_SCHED_POLICIES (decode truth at the source). - The adapter gains _sm100_sched_domain(); check_support now admits explicit NATURAL/LPT/LPT_L2 requests (was NATURAL-only — under- declared), and the causal defaulting heuristic CLAMPS its pick into the served route's domain, falling back to NATURAL. A default may never be a value the route could not honor as an explicit request. - The three SM100-family engine rows declare the true request domain {NATURAL, LPT, LPT_L2} (previously {NATURAL} while the lowering auto-picked LPT — dishonest in both directions). SM120 rows untouched. - frost/README knob-channel section documents the rule. 3. Sibling-lockstep fences (test_sdpa_fp8_sibling_parity.py): source-scan tests (no GPU/DSL needed) assert every decode call site carries the LPT args and the sched-domain constants match; module tests assert compile() signatures in lockstep (#574's fence), kernel domains == adapter table, and CFG divergence is exactly the intentional Rubin set {TILE_K_HW_BMM1, TILE_K_HW_BMM2, STAGES_KV}. Negative-tested: removing one call site's args turns the fence red. SM100 box: 40 passed (5 parity + 2 routing + 33 fp8 e2e). Rubin board validation of the causal suite to follow on the PR. Co-Authored-By: Claude Fable 5 --- python/cudnn/frost/README.md | 11 ++ python/cudnn/sdpa/fwd/api_dsl.py | 25 +++- python/cudnn/sdpa/fwd/engines.py | 8 +- .../fwd/kernels/prefill_d128_fp8_sm100.py | 7 +- .../fwd/kernels/prefill_d128_fp8_sm107.py | 39 ++++- .../frost/test_sdpa_fp8_sibling_parity.py | 133 ++++++++++++++++++ 6 files changed, 215 insertions(+), 8 deletions(-) create mode 100644 test/python/sdpa/frost/test_sdpa_fp8_sibling_parity.py diff --git a/python/cudnn/frost/README.md b/python/cudnn/frost/README.md index 9ba0dbc7f..d6619f05f 100644 --- a/python/cudnn/frost/README.md +++ b/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. diff --git a/python/cudnn/sdpa/fwd/api_dsl.py b/python/cudnn/sdpa/fwd/api_dsl.py index cc3b844c0..2c00c628f 100644 --- a/python/cudnn/sdpa/fwd/api_dsl.py +++ b/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], diff --git a/python/cudnn/sdpa/fwd/engines.py b/python/cudnn/sdpa/fwd/engines.py index 81be0e9f0..f4d6ad1ab 100644 --- a/python/cudnn/sdpa/fwd/engines.py +++ b/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}), diff --git a/python/cudnn/sdpa/fwd/kernels/prefill_d128_fp8_sm100.py b/python/cudnn/sdpa/fwd/kernels/prefill_d128_fp8_sm100.py index 84d2e5c12..4e0efbe54 100644 --- a/python/cudnn/sdpa/fwd/kernels/prefill_d128_fp8_sm100.py +++ b/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 diff --git a/python/cudnn/sdpa/fwd/kernels/prefill_d128_fp8_sm107.py b/python/cudnn/sdpa/fwd/kernels/prefill_d128_fp8_sm107.py index b6717d7c9..4454cf41e 100644 --- a/python/cudnn/sdpa/fwd/kernels/prefill_d128_fp8_sm107.py +++ b/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) diff --git a/test/python/sdpa/frost/test_sdpa_fp8_sibling_parity.py b/test/python/sdpa/frost/test_sdpa_fp8_sibling_parity.py new file mode 100644 index 000000000..adc9dc448 --- /dev/null +++ b/test/python/sdpa/frost/test_sdpa_fp8_sibling_parity.py @@ -0,0 +1,133 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Sibling-lockstep fences for the per-tensor FP8 kernel pair. + +``prefill_d128_fp8_sm107.py`` is a hunk-symmetric sibling of the SM100 file +with a small set of INTENTIONAL divergences (Rubin K=64 MMA geometry, deeper +KV ring, baked LDTM/rowsum). Everything else must move in lockstep — and has +silently not, three times: #574 (compile-signature skew), #602 (amax ABI +drift), #585/#653 (the LPT scheduler landed in one file only, breaking every +causal FP8 graph on the sibling's arch). These tests turn that class of +drift into a red test at the PR that causes it. + +The source-scan tests need no GPU and no DSL — they run everywhere. +""" + +import inspect +import re +from pathlib import Path + +import pytest + +from frost_test_utils import requires_dsl + +from cudnn.sdpa.fwd.config_sm100 import TemplateParams + +pytestmark = pytest.mark.L0 + +_E4M3, _BF16_OUT = 0, 2 +_SIBLINGS = ("prefill_d128_fp8_sm100.py", "prefill_d128_fp8_sm107.py") +# CFG fields the sibling intentionally re-derives (dataclasses.replace at +# import). Anything else differing is drift. +_CFG_DIVERGENCE_ALLOWLIST = {"TILE_K_HW_BMM1", "TILE_K_HW_BMM2", "STAGES_KV"} + + +def _kernel_dir() -> Path: + from cudnn.sdpa.fwd import api_dsl + + return Path(api_dsl.__file__).parent / "kernels" + + +def _decode_calls(src: str): + """Every ``_dispatch_decode_*`` call block in a kernel source.""" + out = [] + for m in re.finditer(r"_dispatch_decode_(initial|payload)\(", src): + depth, i = 0, m.end() - 1 + while True: + c = src[i] + if c == "(": + depth += 1 + elif c == ")": + depth -= 1 + if depth == 0: + break + i += 1 + out.append(src[m.start() : i]) + return out + + +def test_sched_decode_call_sites_carry_lpt_args(): + """#653's fence: every scheduler decode call in BOTH files must pass + qh_per_kh and seqlen_kv (SCHED_LPT_L2 needs them at every call site; a + file whose calls lack them silently loses LPT support).""" + counts = {} + for f in _SIBLINGS: + src = (_kernel_dir() / f).read_text() + calls = _decode_calls(src) + assert calls, f"{f}: no scheduler decode call sites found (parser drift?)" + for blk in calls: + assert "qh_per_kh" in blk and "seqlen_kv" in blk, f"{f}: decode call missing LPT args:\n{blk}" + counts[f] = len(calls) + assert counts[_SIBLINGS[0]] == counts[_SIBLINGS[1]], f"decode call-site count drift: {counts}" + + +def test_sched_domain_declared_in_both_sources(): + """The decode-domain constant must exist in both files and match — the + adapter's defaulting clamp and the engine rows are tied to it.""" + doms = {} + for f in _SIBLINGS: + src = (_kernel_dir() / f).read_text() + m = re.search(r"SUPPORTED_SCHED_POLICIES = frozenset\(\{([^}]*)\}\)", src) + assert m, f"{f}: SUPPORTED_SCHED_POLICIES declaration missing" + doms[f] = {x.strip() for x in m.group(1).split(",") if x.strip()} + assert doms[_SIBLINGS[0]] == doms[_SIBLINGS[1]], f"sched domain drift: {doms}" + + +def _load(rubin): + from cudnn.sdpa.fwd.api_dsl import _load_sm100_kernel_module + + return _load_sm100_kernel_module( + (128, 128), + TemplateParams(dtype_qkv=_E4M3, dtype_o=_BF16_OUT), + fp8=True, + pertensor=True, + rubin=rubin, + ) + + +@requires_dsl +def test_compile_signatures_in_lockstep(): + """#574's fence: the shared adapter calls both modules' ``compile`` with + the same keywords — the signatures must be identical.""" + sig100 = inspect.signature(_load(rubin=False).compile) + sig107 = inspect.signature(_load(rubin=True).compile) + assert list(sig100.parameters) == list(sig107.parameters), f"compile param drift: {list(sig100.parameters)} vs {list(sig107.parameters)}" + + +@requires_dsl +def test_sched_domains_match_adapter_table(): + """#1's invariant: file constants == the adapter's routing-level domain + table, for both routes. A default may never be a value the served kernel + cannot decode.""" + from cudnn.sdpa.fwd.api_dsl import _sm100_sched_domain + + for rubin in (False, True): + mod = _load(rubin=rubin) + assert mod.SUPPORTED_SCHED_POLICIES == _sm100_sched_domain(rubin), ( + f"rubin={rubin}: kernel file declares {sorted(mod.SUPPORTED_SCHED_POLICIES)}, " f"adapter table says {sorted(_sm100_sched_domain(rubin))}" + ) + + +@requires_dsl +def test_cfg_divergence_is_exactly_the_allowlist(): + """Every CFG field must match between siblings except the declared Rubin + re-derivations — a new silent divergence is drift until allowlisted.""" + import dataclasses + + cfg100 = _load(rubin=False).CFG + cfg107 = _load(rubin=True).CFG + diverged = {f.name for f in dataclasses.fields(cfg100) if getattr(cfg100, f.name) != getattr(cfg107, f.name)} + assert diverged == _CFG_DIVERGENCE_ALLOWLIST, ( + f"CFG divergence drift: unexpected={sorted(diverged - _CFG_DIVERGENCE_ALLOWLIST)}, " f"missing={sorted(_CFG_DIVERGENCE_ALLOWLIST - diverged)}" + )