Skip to content

frost(spda): support ragged stats for SM120 frost sdpa forward engine #508

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
vedaanta merged 4 commits into from
Aug 7, 2026
Merged

frost(spda): support ragged stats for SM120 frost sdpa forward engine #508

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
eb4f9ea
Support token-major ragged LSE for SM120 frost sdpa forward engine
Aneureka Aug 7, 2026
b21391f
Support head-major (HT1) for ragged stats
Aneureka Aug 7, 2026
510a425
Update test/python/sdpa/frost/test_sdpa_fwd_dsl_sm120.py
Aneureka Aug 7, 2026
05f1e2b
Remove redundant reshape
Aneureka Aug 7, 2026
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
97 changes: 69 additions & 28 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 @@ -1470,6 +1470,8 @@ def _initialize_implementation(self) -> None:
self.compute_capability: Optional[tuple[int, int]] = None
self.head_dim_qk: Optional[int] = None
self.head_dim_v: Optional[int] = None
self.thd_stats_head_major = False
self.thd_stats_head_stride = 0
self._k_mod = None

def check_support(self) -> bool:
Expand Down Expand Up @@ -1523,13 +1525,21 @@ def check_support(self) -> bool:
self._check_tensor_shape(self.v_desc, (b, h_kv, s_kv, d_v), name="V")
self._check_tensor_shape(self.o_desc, (b, h_q, s_q, d_v), name="O")
if self.lse_desc is not None:
# THD stats are not plumbed (the kernel's packed (1, H, T) LSE does
# not match cuDNN's ragged Stats contract) — reject the request
# instead of silently never writing the user's LSE.
self._not_implemented_error_if(self.thd, "THD stats/LSE output is not plumbed yet; construct without sample_lse")
self._check_dtype(self.lse_desc, torch.float32, name="LSE")
self._check_tensor_shape(self.lse_desc, (b, h_q, s_q), name="LSE")
self._value_error_if(not self.lse_desc.is_contiguous(), "LSE must be contiguous on SM120 DSL")
if self.thd:
stride_h, stride_s = tuple(self.lse_desc.stride[1:])
token_major = (stride_h, stride_s) == (1, h_q)
head_major = not token_major and stride_s == 1 and stride_h >= 1
self._value_error_if(
not token_major and not head_major,
f"THD LSE must be packed token-major (stride_h == 1, stride_s == H) "
f"or head-major (stride_s == 1, stride_h == head_stride); got stride {self.lse_desc.stride}",
)
self.thd_stats_head_major = head_major
self.thd_stats_head_stride = int(stride_h) if head_major else 0
else:
self._value_error_if(not self.lse_desc.is_contiguous(), "LSE must be contiguous on SM120 DSL")

for label, val in (
("B", b),
Expand Down Expand Up @@ -1704,10 +1714,17 @@ def execute(
"this specialization was compiled without sink support; construct the API with has_sink=True",
)
self._check_seq_lens_contract(seq_q_lens, seq_kv_lens)
self._value_error_if(
self.lse_desc is not None and lse_tensor is None,
"lse_tensor is required by this compiled specialization",
)
self._value_error_if(
self.lse_desc is None and lse_tensor is not None,
"this specialization was compiled without an LSE output; construct the API with sample_lse",
)
scale_val = self.scale_softmax if scale_softmax is None or scale_softmax == 0.0 else float(scale_softmax)
scale_softmax_log2 = scale_val * math.log2(math.e)
if self.thd:
self._not_implemented_error_if(lse_tensor is not None, "THD stats/LSE output is not plumbed yet")
self._execute_thd(
q_tensor,
k_tensor,
Expand All @@ -1717,18 +1734,11 @@ def execute(
sinks,
seq_kv_lens,
seq_q_lens,
lse_tensor=lse_tensor,
workspace=workspace,
current_stream=current_stream,
)
return
self._value_error_if(
self.lse_desc is not None and lse_tensor is None,
"lse_tensor is required by this compiled specialization",
)
self._value_error_if(
self.lse_desc is None and lse_tensor is not None,
"this specialization was compiled without an LSE output; construct the API with sample_lse",
)
lse = self._checked_lse_view(lse_tensor) if lse_tensor is not None else None
sinks_t = self._checked_sinks_1d(sinks) if sinks is not None else None
seq_q_lens = (
Expand Down Expand Up @@ -1777,15 +1787,20 @@ def execute(
if o_needs_copy_back:
o_view.copy_(o_scratch)

def _execute_thd(self, q_buf, k_buf, v_buf, o_buf, scale_softmax_log2, sinks, seq_kv_lens, seq_q_lens, workspace=None, current_stream=None):
def _execute_thd(
self, q_buf, k_buf, v_buf, o_buf, scale_softmax_log2, sinks, seq_kv_lens, seq_q_lens, lse_tensor=None, workspace=None, current_stream=None
):
"""THD (ragged) execute: packed ``(1, T, H, D)`` views + cu_seqlens.

Serves the same fully-packed contract as the SM100 THD path
(``ragged_offset == cumsum(seq_len) * H * D`` from 0, multiplier 1);
the offsets are re-derived from ``seq_len_q``/``seq_len_kv``. The two
``.tolist()`` D2H syncs are inherent — the packed totals and the
longest sequence's Q length are runtime values that size the
per-execute compile and grid.
per-execute compile and grid. ``lse_tensor``, when given, is the
caller's ragged Stats buffer, in its declared layout: token-major
packed ``(T, H)`` in the first ``T*H`` elements, or head-major
``(H, head_stride)`` with tokens contiguous within each head row.
"""

b, qh, kh = self.batch_size, self.h_q, self.h_kv
Expand Down Expand Up @@ -1822,15 +1837,36 @@ def _execute_thd(self, q_buf, k_buf, v_buf, o_buf, scale_softmax_log2, sinks, se

if t_q == 0:
return
lse = None
lse_valid = None # the valid-region view (first t_q tokens) the t_kv == 0 fill writes
if lse_tensor is not None:
if self.thd_stats_head_major:
head_stride = self.thd_stats_head_stride
self._value_error_if(
head_stride < t_q,
f"head-major THD LSE head_stride ({head_stride}) must cover the packed Q token total ({t_q})",
)
lse = lse_tensor.as_strided((qh, head_stride), (head_stride, 1), lse_tensor.storage_offset())
lse_valid = lse_tensor.as_strided((qh, t_q), (head_stride, 1), lse_tensor.storage_offset())
else:
lse = lse_tensor.as_strided((t_q, qh), (qh, 1), lse_tensor.storage_offset())
lse_valid = lse
if t_kv == 0:
# Every row is dead: O := 0 (THD stats are not plumbed). A
# zero-token K/V view cannot back a TMA descriptor, so short-cut.
# Every row is dead: O := 0, LSE := -inf (or the sink alone —
# its column keeps the denominator alive). A zero-token K/V view
# cannot back a TMA descriptor, so short-cut both.
o_buf.as_strided((t_q * qh * d_v,), (1,), o_buf.storage_offset()).zero_()
if lse_valid is not None:
if sinks is not None:
sinks_v = self._checked_sinks_1d(sinks)
sinks_v = sinks_v.reshape(qh, 1).expand(qh, t_q) if self.thd_stats_head_major else sinks_v.reshape(1, qh).expand(t_q, qh)
lse_valid.copy_(sinks_v)
else:
lse_valid.fill_(float("-inf"))
return

# THD stats are not plumbed: the kernel compiles with has_lse=False, so
# there is no packed-LSE buffer (dummy or otherwise) to bind. Sinks are
# None-specialized the same way when the graph has no sink token.
# Sinks are None-specialized like the LSE when the graph has no sink
# token.
sinks_t = self._checked_sinks_1d(sinks) if sinks is not None else None
seq_q_dummy = self._dummy(
"seq_q_lens",
Expand Down Expand Up @@ -1858,14 +1894,16 @@ def _packed(buf, tokens):
d_qk=d_qk,
d_v=d_v,
max_sq=max_sq,
has_lse=False,
has_lse=self.lse_desc is not None,
lse_head_major=self.thd_stats_head_major,
lse_head_stride=self.thd_stats_head_stride,
)
fn(
_packed(q_buf, t_q),
_packed(k_buf, t_kv),
_packed(v_buf, t_kv),
_packed(o_buf, t_q),
None,
lse,
sinks_t,
seq_q_dummy,
meta,
Expand All @@ -1876,11 +1914,14 @@ def _packed(buf, tokens):
def scratch_workspace_bytes(self) -> int:
if self.thd:
# [slq32 | slk32 | meta(seq_kv, cu_q, cu_k)].
# No packed-LSE chunk: THD stats are not plumbed and the kernel is
# compiled with has_lse=False, so no LSE buffer exists at all. No
# sinks-dummy chunk either: the kernel None-specializes on sinks.
# No O-descriptor chunk: SM120 stores O with plain guarded GMEM
# stores, so THD needs no per-sequence tensor maps.
# No packed-LSE chunk: with a Stats output the kernel writes the
# caller's ragged Stats buffer directly (token-major (T, H) or
# head-major (H, head_stride)); without one it compiles with
# has_lse=False and no LSE buffer exists at all. No sinks-dummy
# chunk either: the kernel
# None-specializes on sinks. No O-descriptor chunk: SM120 stores O
# with plain guarded GMEM stores, so THD needs no per-sequence
# tensor maps.
b = self.batch_size
return 2 * ws_align(b * 4) + ws_align((3 * b + 2) * 4)
return 0
Expand Down
4 changes: 2 additions & 2 deletions python/cudnn/sdpa/fwd/engines.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -484,6 +484,7 @@ def _sm120_spec() -> EngineSpec:
padded_stats=True,
thd=True,
thd_bottom_right=True,
thd_stats=True,
layouts=frozenset({"bshd", "dense_flex"}),
sched_policies=frozenset({SCHED_NATURAL}),
tile_ms=frozenset({64, 128}),
Expand Down Expand Up @@ -691,8 +692,7 @@ def _execute(variant_pack, workspace=None, stream=None):
k_tensor=k_buf,
v_tensor=v_buf,
o_tensor=o_buf,
# THD ignores lse_tensor (the packed LSE is api-level scratch).
lse_tensor=lse_buf.reshape(facts.b, facts.h_q, facts.s_q) if lse_buf is not None else None,
lse_tensor=lse_buf,
scale_softmax=facts.scale,
sinks=sinks_buf,
seq_kv_lens=seq_kv_buf,
Expand Down
67 changes: 53 additions & 14 deletions python/cudnn/sdpa/fwd/kernels/prefill_f16_sm120.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -179,6 +179,7 @@ def __init__(
thd_varlen: bool = False,
thd_batch: int = 1,
thd_max_sq: int = 0,
thd_lse_head_major: bool = False,
head_tile_qk: int = 128,
head_tile_v: int = 128,
kv_tile: int = SEQ_KV_TILES[0],
Expand All @@ -202,6 +203,10 @@ def __init__(
grid covers ``ceil(thd_max_sq / q_tile)`` tiles per sequence.
:param thd_batch: THD only: the real sequence count B.
:param thd_max_sq: THD only: the longest sequence's Q length.
:param thd_lse_head_major: THD only: the packed LSE is head-major
``(H, head_stride)`` (FlashAttention's ``softmax_lse`` layout; tokens
contiguous within a head, ``head_stride >= T``) instead of the default
token-major ``(T, H)``.
:param head_tile_qk: Q/K head dimension (the QK^T contraction width).
Must be a multiple of 16 between 16 and 256, inclusive.
:param head_tile_v: V/O head dimension (the P@V output width). Same
Expand All @@ -225,6 +230,7 @@ def __init__(
self.thd_varlen = thd_varlen
self.thd_batch = thd_batch
self.thd_max_sq = thd_max_sq
self.thd_lse_head_major = thd_lse_head_major

self.head_tile_qk = head_tile_qk
self.head_tile_v = head_tile_v
Expand Down Expand Up @@ -729,7 +735,9 @@ def kernel(
:param k: Key tensor.
:param v: Value tensor.
:param o: Output tensor.
:param lse: ``(B, H, Sq)`` fp32 log-sum-exp output, or ``None`` to
:param lse: fp32 log-sum-exp output — ``(B, H, Sq)`` dense; packed
token-major ``(T, H)`` or head-major ``(H, head_stride)`` (per
``thd_lse_head_major``) under ``thd_varlen``; or ``None`` to
compile the LSE store out (the DSL specializes on ``None``).
:param sinks: ``(H,)`` fp32 per-Q-head sink logits; ``None`` iff the
kernel is configured without ``has_sink``.
Expand Down Expand Up @@ -1093,12 +1101,19 @@ def kernel(
lse_q_idx = q_seq_idx + q_warp_row0 + (lane // 4) + row_half * 8
lse_out = cutlass.Float32(row_lse[row_half])
if cutlass.const_expr(self.thd_varlen):
# Packed (1, H, T) LSE: rows past this sequence's Q
# length belong to the NEXT sequence — never written,
# and there is no padded region to trim.
# Packed ragged-Stats LSE, written directly in the
# caller's declared layout: token-major (T, H) or
# head-major (H, head_stride).
# Rows past this sequence's Q length belong to the
# NEXT sequence — never written, and there is no
# padded region to trim.
if lse_q_idx < seqlen_q:
lse_row = lse_arr[0, head_idx, :]
lse_row[q_row_base + lse_q_idx] = lse_out
if cutlass.const_expr(self.thd_lse_head_major):
lse_row = lse_arr[head_idx, :]
lse_row[q_row_base + lse_q_idx] = lse_out
else:
lse_row = lse_arr[q_row_base + lse_q_idx, :]
lse_row[head_idx] = lse_out
else:
# Rows at/past this batch's Q length trim to -inf.
if lse_q_idx >= seqlen_q:
Expand Down Expand Up @@ -1197,7 +1212,9 @@ def __call__(
:param k: Key tensor with shape ``(B, Sk, H, D)``.
:param v: Value tensor with shape ``(B, Sk, H, D)``.
:param o: Output tensor with shape ``(B, Sq, H, D)``.
:param lse: ``(B, H, Sq)`` fp32 log-sum-exp output, or ``None`` to
:param lse: fp32 log-sum-exp output — ``(B, H, Sq)`` dense; packed
token-major ``(T, H)`` or head-major ``(H, head_stride)`` (per
``thd_lse_head_major``) under ``thd_varlen``; or ``None`` to
compile the LSE store out entirely (no dummy buffer needed).
:param sinks: ``(H,)`` fp32 per-Q-head sink logits; must be ``None``
exactly when the kernel is configured without ``has_sink``.
Expand Down Expand Up @@ -1225,10 +1242,22 @@ def __call__(
if cutlass.const_expr(not self.is_layout_supported(tensor.shape, tensor.stride)):
raise ValueError(f"{name} must use compact BSHD storage")
if cutlass.const_expr(lse is not None):
if cutlass.const_expr(lse.shape != (q.shape[0], q.shape[2], q.shape[1])):
raise ValueError("LSE must have shape (B, H, Sq)")
if cutlass.const_expr(lse.stride != (q.shape[2] * q.shape[1], q.shape[1], 1)):
raise ValueError("LSE must be compact row-major")
if cutlass.const_expr(self.thd_varlen):
if cutlass.const_expr(self.thd_lse_head_major):
if cutlass.const_expr(lse.shape[0] != q.shape[2] or lse.shape[1] < q.shape[1]):
raise ValueError("head-major THD LSE must have shape (H, head_stride) with head_stride >= T")
if cutlass.const_expr(lse.stride != (lse.shape[1], 1)):
raise ValueError("head-major THD LSE must be compact row-major")
else:
if cutlass.const_expr(lse.shape != (q.shape[1], q.shape[2])):
raise ValueError("THD LSE must have shape (T, H)")
if cutlass.const_expr(lse.stride != (q.shape[2], 1)):
raise ValueError("THD LSE must be compact token-major")
else:
if cutlass.const_expr(lse.shape != (q.shape[0], q.shape[2], q.shape[1])):
raise ValueError("LSE must have shape (B, H, Sq)")
if cutlass.const_expr(lse.stride != (q.shape[2] * q.shape[1], q.shape[1], 1)):
raise ValueError("LSE must be compact row-major")
if cutlass.const_expr(self.has_sink != (sinks is not None)):
raise ValueError("sinks must be provided exactly when the kernel is configured with has_sink")
if cutlass.const_expr(sinks is not None and sinks.shape != (q.shape[2],)):
Expand Down Expand Up @@ -1339,6 +1368,8 @@ def compile( # noqa: A001
d_v: int = 128,
max_sq: int = 0,
has_lse: bool = True,
lse_head_major: bool = False,
lse_head_stride: int = 0,
) -> Callable:
"""Compile and cache one architecture-specific compact BSHD shape.

Expand All @@ -1351,7 +1382,10 @@ def compile( # noqa: A001

``has_lse=False`` compiles the LSE store out (the kernel specializes on a
``None`` LSE argument) — callers that don't want stats pass no LSE buffer
at all instead of a dummy.
at all instead of a dummy. THD LSE is token-major ``(T, H)`` by default;
``lse_head_major=True`` switches to head-major ``(H, lse_head_stride)``
(FlashAttention's ``softmax_lse`` layout), where ``lse_head_stride`` is the
caller-declared head-row stride (``>= T``, a shape — part of the cache key).
"""

kernel = SM120FusedMultiHeadAttentionForward(
Expand All @@ -1366,6 +1400,7 @@ def compile( # noqa: A001
thd_varlen=PARAMS.thd_varlen,
thd_batch=b,
thd_max_sq=max_sq,
thd_lse_head_major=lse_head_major,
head_tile_qk=d_qk,
head_tile_v=d_v,
q_tile=PARAMS.q_tile,
Expand Down Expand Up @@ -1396,11 +1431,15 @@ def compile( # noqa: A001
stride_order=(3, 2, 1, 0),
assumed_align=16,
)
if PARAMS.thd_varlen:
fake_lse_shape = (qh, lse_head_stride) if lse_head_major else (sq, qh)
else:
fake_lse_shape = (fake_batch, qh, sq)
fake_lse = (
cute.runtime.make_fake_compact_tensor(
cutlass.Float32,
(fake_batch, qh, sq),
stride_order=(2, 1, 0),
fake_lse_shape,
stride_order=(1, 0) if PARAMS.thd_varlen else (2, 1, 0),
assumed_align=4,
)
if has_lse
Expand Down
Loading