Skip to content

frost(sdpa): add more features like padded head dim, skv_tile=0, diagonal_band_right_bound to sm120 frost sdpa prefill kernel #531

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Aneureka merged 1 commit into from
Aug 10, 2026
Merged

frost(sdpa): add more features like padded head dim, skv_tile=0, diagonal_band_right_bound to sm120 frost sdpa prefill kernel #531

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
38 changes: 25 additions & 13 deletions python/cudnn/sdpa/fwd/api_dsl.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,9 +28,10 @@
)
from cudnn.sdpa.fwd.config_sm100 import TemplateParams as Sm100TemplateParams
from cudnn.sdpa.fwd.config_sm120 import (
HEAD_TILE_GRANULE as _SM120_HEAD_TILE_GRANULE,
SEQ_KV_TILES as _SM120_KV_TILES,
SEQ_Q_TILES as _SM120_Q_TILES,
SUPPORTED_HEAD_TILES as _SM120_SUPPORTED_HEAD_TILES,
SUPPORTED_HEAD_TILE_MAX as _SM120_HEAD_TILE_MAX,
TemplateParams as Sm120TemplateParams,
smem_bytes as _sm120_smem_bytes,
)
Expand Down Expand Up @@ -1523,8 +1524,9 @@ class SdpaFwdDslSm120(SdpaFwdDsl):
``execute()`` normalizes to the kernel's compact-BSHD storage via
``_to_bshd`` / ``_to_bshd_writable`` — zero-copy when the tensor already
is BSHD-compact, one gather / scatter copy otherwise. The kernel supports
FP16/BF16 MHA, GQA, and MQA; head dimensions from 16 through 256 in
increments of 16; top-left or bottom-right causal masks; left sliding
FP16/BF16 MHA, GQA, and MQA; head dimensions in multiples of 8 through
256 (ENVELOPE: the kernel compiles at tiles rounded up to 16 and TMA
zero-fills the pad columns); top-left or bottom-right causal masks; left sliding
windows; optional per-batch query and key/value lengths; optional
per-Q-head attention-sink logits; and THD (ragged / fully packed
variable-length) batches, whose per-shape compile is deferred to
Expand All @@ -1550,10 +1552,6 @@ def check_support(self) -> bool:
if self.thd:
self._value_error_if(self.seq_q_lens_present, "seq_q_lens_present is dense-only (THD carries per-sequence Q lengths via cu_seqlens)")
self.seq_kv_lens_present = True
self._value_error_if(
self.window_size_right is not None,
"SM120 DSL SDPA: window_size_right (causal right-band widening) is not plumbed for this kernel",
)
self._value_error_if(
self.sched_policy is not None and self.sched_policy != SCHED_NATURAL,
f"SM120 DSL SDPA only supports sched_policy={SCHED_NATURAL}",
Expand Down Expand Up @@ -1630,12 +1628,12 @@ def check_support(self) -> bool:
f"H_q ({h_q}) must be divisible by H_kv ({h_kv}) for GQA / MQA",
)
self._value_error_if(
d_q not in _SM120_SUPPORTED_HEAD_TILES,
f"D_QK ({d_q}) must be one of {_SM120_SUPPORTED_HEAD_TILES}",
d_q % 8 != 0 or not 0 < d_q <= _SM120_HEAD_TILE_MAX,
f"D_QK ({d_q}) must be a multiple of 8 (TMA 16-byte global-stride rule at 2 B/elem) and <= {_SM120_HEAD_TILE_MAX}",
)
self._value_error_if(
d_v not in _SM120_SUPPORTED_HEAD_TILES,
f"D_V ({d_v}) must be one of {_SM120_SUPPORTED_HEAD_TILES}",
d_v % 8 != 0 or not 0 < d_v <= _SM120_HEAD_TILE_MAX,
f"D_V ({d_v}) must be a multiple of 8 (TMA 16-byte global-stride rule at 2 B/elem) and <= {_SM120_HEAD_TILE_MAX}",
)

self.dtype = self._check_dtype(self.q_desc, [torch.float16, torch.bfloat16], name="Q")
Expand Down Expand Up @@ -1665,12 +1663,20 @@ def check_support(self) -> bool:
)
self._value_error_if(
self.causal_bottom_right and not self.is_causal,
"causal_bottom_right requires is_causal=True",
"causal_bottom_right requires is_causal=True (a band graph arrives as is_causal with its right bound)",
)
self._value_error_if(
self.window_size_left is not None and self.window_size_left < 0,
f"window_size_left must be non-negative, got {self.window_size_left}",
)
self._value_error_if(
self.window_size_right is not None and self.window_size_right < 0,
f"window_size_right must be >= 0; got {self.window_size_right}",
)
self._value_error_if(
self.window_size_right is not None and not self.is_causal,
"window_size_right widens the causal diagonal and requires is_causal=True",
)
self._value_error_if(
self.seq_q_lens_present and not self.seq_kv_lens_present,
"seq_q_lens_present requires seq_kv_lens_present (padding mask)",
Expand All @@ -1688,8 +1694,14 @@ def check_support(self) -> bool:
arch = f"sm_{self.compute_capability[0]}{self.compute_capability[1]}"
smem_capacity_bytes = cutlass.utils.get_smem_capacity_in_bytes(arch)

# SMEM tiles are sized at the ENVELOPE-padded head tiles (rounded up
# to the head-tile granule), not the actual dims — the kernel stages
# full tiles and the TMA zero-fills the pad columns.
d_qp = -(-d_q // _SM120_HEAD_TILE_GRANULE) * _SM120_HEAD_TILE_GRANULE
d_vp = -(-d_v // _SM120_HEAD_TILE_GRANULE) * _SM120_HEAD_TILE_GRANULE

def _smem_bytes(kv_tile: int) -> int:
return _sm120_smem_bytes(d_q, d_v, self.q_tile, kv_tile, self.dtype.itemsize)
return _sm120_smem_bytes(d_qp, d_vp, self.q_tile, kv_tile, self.dtype.itemsize)

if self.tile_n is None:
# Pick the largest KV tile that fits this device.
Expand Down
15 changes: 9 additions & 6 deletions python/cudnn/sdpa/fwd/config_sm120.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,10 @@

SEQ_Q_TILES = (128, 64)
SEQ_KV_TILES = (128, 64)
SUPPORTED_HEAD_TILES = tuple(range(16, 257, 16))
HEAD_TILE_GRANULE = 16
SUPPORTED_HEAD_TILE_MIN = 16
SUPPORTED_HEAD_TILE_MAX = 256
SUPPORTED_HEAD_TILES = tuple(range(SUPPORTED_HEAD_TILE_MIN, SUPPORTED_HEAD_TILE_MAX + 1, HEAD_TILE_GRANULE))

# SMEM the SM120 parts expose to a kernel. The adapter asks cutlass for the
# authoritative number at build time; this constant lets the ranking answer
Expand Down Expand Up @@ -44,9 +47,9 @@ class TemplateParams:
dtype_qkv: int = DTYPE_FP16
# The mask is ONE diagonal band (same model as config_sm100 / the analyzer
# facts): per-side offsets from the diagonal, None = unbounded on that
# side. This kernel serves window_right in {None, 0} only (plain causal;
# right-band widening is not plumbed here). bottom_right anchors the
# band's diagonal at the bottom-right corner.
# side. window_right = 0 is plain causal; window_right > 0 widens the
# diagonal right by R columns (cuDNN's diagonal_band_right_bound).
# bottom_right anchors the band's diagonal at the bottom-right corner.
window_left: int | None = None
window_right: int | None = None
bottom_right: bool = False
Expand All @@ -68,8 +71,8 @@ def validate_params(params: TemplateParams) -> None:

if params.dtype_qkv not in (DTYPE_BF16, DTYPE_FP16):
raise ValueError(f"SM120 SDPA: dtype_qkv must be DTYPE_BF16 ({DTYPE_BF16}) or DTYPE_FP16 ({DTYPE_FP16}); got {params.dtype_qkv}")
if params.window_right not in (None, 0):
raise ValueError(f"SM120 SDPA: window_right must be None (unbounded) or 0 (causal) — right-band widening is not plumbed; got {params.window_right}")
if params.window_right is not None and params.window_right < 0:
raise ValueError(f"SM120 SDPA: window_right must be None (unbounded) or >= 0 (0 = plain causal); got {params.window_right}")
if params.bottom_right and params.window_right is None:
raise ValueError("SM120 SDPA: bottom_right anchors the band's diagonal and requires a right bound (window_right)")
if params.window_left is not None and params.window_left < 0:
Expand Down
20 changes: 16 additions & 4 deletions python/cudnn/sdpa/fwd/engines.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -109,8 +109,8 @@ class Capabilities:
# global-stride rule at 2 bytes/elem) via TMA zero-padding — the kernel's
# descriptors carry the ACTUAL extents, so padded contraction columns load
# as exact zeros (S/softmax unchanged) and O stores past d_v are
# OOB-clipped. False = only the native dims above are eligible (FP8/MXFP8;
# SM120, whose lowering has no zero-padding path wired yet).
# OOB-clipped. False = only the native dims above are eligible (FP8/MXFP8,
# whose SF plumbing is not audited for zero-padding).
d_envelope: bool = False
dtypes: frozenset = frozenset({cudnn.data_type.HALF, cudnn.data_type.BFLOAT16}) # cudnn.data_type, see graph_analyzer
is_mxfp8: bool = False # block-scale MXFP8 engine (FP8 in + per-32-block E8M0 SF)
Expand Down Expand Up @@ -308,6 +308,8 @@ def mismatch(capabilities: Capabilities, facts: "ga.SdpaGraphFacts", knobs: Opti
if fact and not cap:
return f"graph uses {label}, which this engine does not support"

if facts.right_band_widening and facts.right_bound is not None and facts.right_bound < 0:
return f"negative diagonal_band_right_bound ({facts.right_bound}) is not supported"
if facts.bottom_right:
if not (facts.causal or facts.right_band_widening):
return "bottom-right alignment requires a causal upper bound (plain or right-widened)"
Expand Down Expand Up @@ -470,26 +472,36 @@ def _sm100_fp8_spec(d: int) -> EngineSpec:


def _sm120_spec() -> EngineSpec:
from cudnn.sdpa.fwd.config_sm120 import SUPPORTED_HEAD_TILES

return EngineSpec(
name="sdpa_fwd_prefill_sm120",
capabilities=Capabilities(
sm_lo=_BLACKWELL_GEFORCE[0],
sm_hi=_BLACKWELL_GEFORCE[1],
phase="prefill",
d_qk=frozenset(range(16, 257, 16)),
d_v=frozenset(range(16, 257, 16)),
d_qk=frozenset(SUPPORTED_HEAD_TILES),
d_v=frozenset(SUPPORTED_HEAD_TILES),
d_envelope=True,
dtypes=frozenset({cudnn.data_type.HALF, cudnn.data_type.BFLOAT16}),
causal=True,
bottom_right=True,
bottom_right_with_swa=True,
bottom_right_padded_seq_q=True,
swa=True,
right_band_widening=True,
padded=True,
sink=True,
stats=True,
lse_optional=True,
padded_stats=True,
thd=True,
# No KV-tail rule: the kernel walks KV tiles right-to-left and its
# first (masked) step always covers the rightmost — and therefore
# any partial — tile, comparing columns against seqlen_k regardless
# of mask flags. Ragged S_kv is served natively with no synthesized
# padding and no padded-path cost.
skv_tile=0,
layouts=frozenset({"bshd", "dense_flex"}),
sched_policies=frozenset({SCHED_NATURAL}),
tile_ms=frozenset({64, 128}),
Expand Down
Loading