Skip to content

frost(sdpa): remove the legacy (pre-envelope) THD leg from the FP8/MXFP8 SM100/SM107 kernels #622

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 2 commits into from
Aug 17, 2026
Merged

frost(sdpa): remove the legacy (pre-envelope) THD leg from the FP8/MXFP8 SM100/SM107 kernels #622

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
95d9661
frost(sdpa): remove the legacy (pre-envelope) THD leg from the FP8/MX…
vedaanta Aug 17, 2026
0d6d69a
frost(sdpa): drop the freed THD ABI slots from the FP8/MXFP8 hosts
vedaanta Aug 17, 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
27 changes: 14 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 @@ -846,6 +846,14 @@ def check_support(self) -> bool:
(self.cu_seq_q_lens or self.cu_seq_kv_lens) and not self.thd,
"cu_seq_len_* is THD-only (the dense kernels have no CU read mode yet)",
)
# The engine specs already declare thd=False for the FP8/MXFP8 cells
# (graph-API routing); this gate covers direct construction.
self._not_implemented_error_if(
self.thd and self._fp8,
"THD/varlen is not supported by the SM100 FP8/MXFP8 kernels (the "
"legacy THD leg was removed; the port will follow the write_thd_meta "
"envelope design — issue #552)",
)
# Dense padded-Q trim backstops (engines.lower_dsl_prefill never sets
# these combinations; a direct caller could).
self._value_error_if(
Expand Down Expand Up @@ -937,9 +945,10 @@ def compile(self) -> None:
# compile as dynamic extents — issue #552), so compile HERE like
# every dense specialization; execute()'s lru-cached call re-binds
# this artifact. (The all-KV-zero clamp swaps the K/V strides and
# mints its own entry on first hit.) FP8 THD is not wired
# (_execute_thd is f16-only); keep the deferred sentinel there.
self._compiled_kernel = "thd-deferred" if self._fp8 else self._k_mod.compile(**self._thd_compile_kwargs())
# mints its own entry on first hit.) FP8/MXFP8 THD is rejected in
# check_support (the legacy THD leg was removed), so this branch is
# f16-only.
self._compiled_kernel = self._k_mod.compile(**self._thd_compile_kwargs())
elif self._fp8:
# FP8/MXFP8 kernels are exact-match d128 (gated in check_support);
# their compile() has no envelope head-dim parameters. has_lse=False
Expand Down Expand Up @@ -1462,7 +1471,7 @@ def _execute_mxfp8(
import cutlass

if self.thd:
raise NotImplementedError("Frost MXFP8: THD/varlen execute is not wired yet (dense d128 only for v1)")
raise NotImplementedError("Frost MXFP8: the legacy THD leg was removed (dense d128 only); see issue #552")
if sf_q is None or sf_k is None or sf_v is None:
raise ValueError("Frost MXFP8 execute requires sf_q/sf_k/sf_v (block-scale descale tensors)")

Expand Down Expand Up @@ -1493,7 +1502,6 @@ def _execute_mxfp8(
if seq_kv_lens is not None
else self._dummy("seq_kv", device, lambda: torch.zeros(b, dtype=torch.int32, device=device))
)
o_desc_dummy = self._dummy("o_desc", device, lambda: torch.zeros(1, dtype=torch.int64, device=device))

amax_o_buf = amax_o.reshape(-1)[:1] if amax_o is not None else self._dummy("amax_o", device, lambda: torch.zeros(1, dtype=torch.float32, device=device))
# Must be enqueued on the SAME stream as the kernel launch below, else the
Expand All @@ -1514,12 +1522,8 @@ def _execute_mxfp8(
amax_o_buf,
sinks_t,
seq_kv_t,
o_desc_dummy,
(b, h_q, h_kv, sq, skv, 0),
cutlass.Float32(scale_softmax_log2),
cutlass.Int32(0), # n_thd_units (dense)
cutlass.Int32(0), # total_q_sf_tiles (dense — kernel folds it out)
cutlass.Int32(0), # total_kv_sf_tiles
stream=current_stream,
)
if o_needs_copy_back:
Expand Down Expand Up @@ -1556,7 +1560,7 @@ def _execute_fp8(
import cutlass

if self.thd:
raise NotImplementedError("Frost per-tensor FP8: THD/varlen execute is not wired yet (dense d128 only)")
raise NotImplementedError("Frost per-tensor FP8: the legacy THD leg was removed (dense d128 only); see issue #552")

def _scalar(t, default=1.0):
return float(t.reshape(-1)[0].item()) if t is not None else default
Expand Down Expand Up @@ -1586,7 +1590,6 @@ def _scalar(t, default=1.0):
if seq_kv_lens is not None
else self._dummy("seq_kv", device, lambda: torch.zeros(b, dtype=torch.int32, device=device))
)
o_desc_dummy = self._dummy("o_desc", device, lambda: torch.zeros(1, dtype=torch.int64, device=device))

# amax_o: the kernel atomicMax'es into this buffer, so it MUST start
# at 0. It accumulates max|o_scaled| (pre-cast, exact even for FP8 O);
Expand All @@ -1605,11 +1608,9 @@ def _scalar(t, default=1.0):
lse,
sinks_t,
seq_kv_t,
o_desc_dummy,
(b, h_q, h_kv, sq, skv, 0),
cutlass.Float32(scale_softmax_log2),
cutlass.Float32(o_scale_fused),
cutlass.Int32(0), # n_thd_units (dense)
amax_o_buf,
stream=current_stream,
)
Expand Down
129 changes: 33 additions & 96 deletions python/cudnn/sdpa/fwd/kernels/prefill_d128_fp8_sm100.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,13 +18,11 @@
``CFG.TILE_K_HW_BMM2`` (→ 4 k-steps at TILE_K_HW=32). Confirmed by the
cuDNN f8 reference (UTCMMA_TILE_K=32, BMM_XMMAS_K=4, kind::f8f6f4).

THD / varlen (``CFG.THD_VARLEN=1``) is supported (E4M3/E5M2) — FP8 is
element-addressed (no block-scale SF), so it rides the same THD path as f16:
packed ``[1,T,H,D]`` + ``cu_seqlens`` coord offset (both Q slabs under
TILES_Q=2), per-batch O TMA-descriptor array (shared
``kernels/dsl/common/sdpa/thd.py``), packed ``[1,QH,T]`` LSE. Dense path
byte-identical. Public-API quantize→attention(layout="thd") glue is a
follow-up (needs a THD-aware quantizer); drive via the kernel / the DSL backend.
THD / varlen is NOT supported here: the legacy (pre-#606) THD leg was removed
(issue #552) — it was never wired (engine spec thd=False; the adapter rejects
FP8 THD in check_support). A future port must follow the device-built-metadata
plus plan-time-envelope design (``write_thd_meta``, issue #552 / PRs #606, #608)
used by the f16 kernels; ``CFG.THD_VARLEN=1`` now fails loudly at trace time.
"""

import os
Expand Down Expand Up @@ -93,7 +91,7 @@
from cudnn.frost.tile_dsl.regtile import RegTile, vec_concat
from cudnn.frost.tile_dsl.mma import mma_ss, mma_ts_step
from cudnn.frost.tile_dsl.tma import tma_load_tile, tma_store_tile, tma_store_commit, tma_store_wait
from cudnn.frost.tile_dsl.handles import MmaDesc, SmemTile, GmemTileTma, tma_slice_runtime_desc
from cudnn.frost.tile_dsl.handles import MmaDesc, SmemTile, GmemTileTma
from cudnn.frost.tile_dsl.tmem import tmem_alloc, tmem_dealloc
from cudnn.frost.tile_dsl.mask import (
apply_mask_chunk,
Expand Down Expand Up @@ -172,12 +170,10 @@
_resolve_seqlen_kv = _sdpa_h.resolve_seqlen_kv
_resolve_seqlen_q = _sdpa_h.resolve_seqlen_q

# THD / varlen — shared helpers (FP8 element-addressed like f16, per-tensor
# dequant scalars, no block-scale SF). Gated by CFG.THD_VARLEN (folds out).
# TILES_Q=2: q_seq_off applies to BOTH Q slabs + both O-store slabs.
from cudnn.sdpa.fwd.kernels.thd_sm100 import build_o_descs_kernel as _build_o_descs_kernel, TENSOR_MAP_QWORDS

_TENSOR_MAP_QWORDS = TENSOR_MAP_QWORDS
# Flat-grid decode dispatch + seq-offset helper from the shared factory. The
# legacy THD leg was removed (issue #552); _thd_tma_offsets folds to
# (0, 0, batch_idx) at THD_VARLEN=0, so the dense TMA coords below are
# byte-identical.
_dispatch_decode_initial = _sdpa_h.dispatch_decode_initial
_dispatch_decode_payload = _sdpa_h.dispatch_decode_payload
_thd_tma_offsets = _sdpa_h.thd_tma_offsets
Expand Down Expand Up @@ -226,7 +222,6 @@ def _kernel(
lse_tensor: Optional[cute.Tensor],
sinks_tensor: cute.Tensor,
seq_kv_lens_tensor: cute.Tensor,
o_desc_words: cute.Tensor,
seqlen_q: cutlass.Int32,
seqlen_kv: cutlass.Int32,
n_q_supers: cutlass.Int32,
Expand Down Expand Up @@ -503,7 +498,6 @@ def _kernel(
n_batch=n_batch,
cta_in_pair=cta_in_pair,
seq_kv_lens_tensor=seq_kv_lens_tensor,
o_desc_words=o_desc_words,
)

else: # warp_idx == CFG.SCHED_WARP_ID
Expand Down Expand Up @@ -611,11 +605,8 @@ def _tmaldg_warp_group(
bars.mb_k_full[kv_state.idx].arrive(n_bytes=kTmaTransactionBytes, pred=nvvm.elect_sync())
tma_load_tile(
sK[kv_state.idx],
# THD: prologue K load MUST apply the per-sequence kv offset
# (kv_seq_off) + packed batch coord (tma_batch), like the mainloop
# K load + the V loads. The old `+ K_ROW_OFFSET_PEER, batch_idx`
# is byte-identical for dense but reads the wrong packed location
# for THD batch>=1 (see the f16 kernel's prologue K-load fix).
# kv_seq_off / tma_batch fold to (0, batch_idx) — dense-identity
# with the THD leg removed.
tma_k(cutlass.Int32(0), kv_head_idx, kv_row_base + K_ROW_OFFSET_PEER + kv_seq_off, tma_batch),
bars.mb_k_full[kv_state.idx].smem_ptr,
cta_group=CFG.CTA_MMA,
Expand Down Expand Up @@ -735,7 +726,6 @@ def _tmastg_warp_group(
n_batch,
cta_in_pair,
seq_kv_lens_tensor,
o_desc_words,
):
"""Persistent O-store warp; tiles claimed via scheduler's try_cancel.async."""
o_full_phase = cutlass.Int32(0)
Expand Down Expand Up @@ -764,18 +754,10 @@ def _tmastg_warp_group(
bars.mb_o_full[qs].wait(o_full_phase)

# O TMA params follow O's swizzle, not V's (V and O swizzles may differ).
if cutlass.const_expr(CFG.THD_VARLEN):
# THD: store each Q slab through this batch's pre-built descriptor
# (seq extent = S_q_b → box past S_q_b OOB-clipped). q_row coord
# sequence-local; batch → 0. Both slabs share one descriptor.
o_desc_ptr = (o_desc_words.iterator.raw_ptr() + batch_idx * cutlass.Int32(_TENSOR_MAP_QWORDS)).tospace(cutlass.AddressSpace.generic)
o_slice = tma_slice_runtime_desc(o_desc_ptr, cutlass.Int32(0), head_idx, q_row_base + cutlass.Int32(qs * CFG.TILE_M), cutlass.Int32(0))
tma_store_tile(sO[qs], o_slice)
else:
tma_store_tile(
sO[qs],
tma_o(cutlass.Int32(0), head_idx, q_row_base + cutlass.Int32(qs * CFG.TILE_M), batch_idx),
)
tma_store_tile(
sO[qs],
tma_o(cutlass.Int32(0), head_idx, q_row_base + cutlass.Int32(qs * CFG.TILE_M), batch_idx),
)

tma_store_commit()
tma_store_wait(0)
Expand Down Expand Up @@ -1692,24 +1674,12 @@ def _correction_warp_group(

# cga2 OOB-row guard: cluster Q rows can exceed seqlen_q.
q_row_global = q_super_idx * cutlass.Int32(CFG.TILES_Q * CFG.TILE_M) + cutlass.Int32(qs * CFG.TILE_M) + tid_in_wg
if cutlass.const_expr(CFG.THD_VARLEN):
# THD: sequence-local row; LSE packed [1,QH,T] → [0, head, cu_q[b]+local].
_cu = cutlass.make_array_view(seq_kv_lens_tensor)
_cu_q_b = cutlass.Int32(_cu[n_batch + batch_idx])
_s_q_b = cutlass.Int32(_cu[n_batch + batch_idx + cutlass.Int32(1)]) - _cu_q_b
_row_valid = q_row_global < _s_q_b
if _row_valid:
if cutlass.const_expr(lse_tensor is not None):
lse_arr = cutlass.make_array_view(lse_tensor)
lse_row = lse_arr[cutlass.Int32(0), head_idx, :]
lse_row[_cu_q_b + q_row_global] = lse_val
else:
_row_valid = q_row_global < seqlen_q
if _row_valid:
if cutlass.const_expr(lse_tensor is not None):
lse_arr = cutlass.make_array_view(lse_tensor)
lse_row = lse_arr[batch_idx, head_idx, :]
lse_row[q_row_global] = lse_val
_row_valid = q_row_global < seqlen_q
if _row_valid:
if cutlass.const_expr(lse_tensor is not None):
lse_arr = cutlass.make_array_view(lse_tensor)
lse_row = lse_arr[batch_idx, head_idx, :]
lse_row[q_row_global] = lse_val

# amax_o = max over valid rows of |o_scaled| (the fp32 pre-cast output). Divided
# by scale_o in api to give the pre-quant output amax (cuDNN FP8 ref, in-kernel).
Expand Down Expand Up @@ -1851,14 +1821,15 @@ def _host(
lse_tensor: Optional[cute.Tensor],
sinks_tensor: cute.Tensor,
seq_kv_lens_tensor: cute.Tensor,
o_desc_words: cute.Tensor,
problem_size: Tuple[int, int, int, int, int, int],
scale_softmax_log2: cutlass.Float32,
o_scale_fused: cutlass.Float32,
n_thd_units: cutlass.Int32,
amax_o_tensor: cute.Tensor,
stream: _cuda_driver.CUstream = None,
) -> None:
# Legacy THD leg removed — fail loudly at trace time (issue #552).
if cutlass.const_expr(CFG.THD_VARLEN):
raise NotImplementedError("prefill_d128_fp8_sm100: the legacy THD leg was removed; port the write_thd_meta envelope design (issue #552) instead")
B, QH, KH, SQ, SKV, _ = problem_size

# K box rows are per-CTA (TILE_N/CTA_MMA); O box inner must match O's swizzle, not V's.
Expand Down Expand Up @@ -1909,26 +1880,7 @@ def _tma_swz(byte_w: int):
q_clusters = (SQ + rows_per_cluster - 1) // rows_per_cluster
grid_q_supers = q_clusters * CFG.CTA_MMA
q_supers = grid_q_supers
if cutlass.const_expr(CFG.THD_VARLEN):
# THD: build the per-batch O descriptor array, then launch the exact
# flat batch-outermost grid (n_thd_units host-computed); grid_x = units*CGA_M.
# Works at cga1 (CGA_M=1) and cga2.
# Per-token element stride of packed O (QH * d_v) — NOT CFG.TILE_O,
# which is only coincidentally right at QH == 1 and otherwise lands
# every batch >= 1's descriptor base inside earlier batches' rows
# (the f16 kernel carries the same fix).
_build_o_descs_kernel(
o_tensor,
tma_o_desc,
o_desc_words,
seq_kv_lens_tensor,
cutlass.Int32(QH),
cutlass.Int32(B),
cutlass.Int32(o_tensor.stride[1]),
).launch(grid=(1, 1, 1), block=(32, 1, 1), stream=stream)
grid_shape = (n_thd_units * cutlass.Int32(CFG.CGA_M), cutlass.Int32(1), cutlass.Int32(1))
else:
grid_shape = (grid_q_supers, QH, B) if cutlass.const_expr(CFG.SCHEDULER_POLICY == SCHED_NATURAL) else (grid_q_supers * QH * B, 1, 1)
grid_shape = (grid_q_supers, QH, B) if cutlass.const_expr(CFG.SCHEDULER_POLICY == SCHED_NATURAL) else (grid_q_supers * QH * B, 1, 1)
_kernel(
tma_q_desc,
tma_k_desc,
Expand All @@ -1937,7 +1889,6 @@ def _tma_swz(byte_w: int):
lse_tensor,
sinks_tensor,
seq_kv_lens_tensor,
o_desc_words,
cutlass.Int32(SQ),
cutlass.Int32(SKV),
cutlass.Int32(q_supers),
Expand All @@ -1959,32 +1910,30 @@ def _tma_swz(byte_w: int):
def compile(b: int = 1, qh: int = 1, kh: int = 1, sq: int = 256, skv: int = 128, has_lse: bool = True) -> Callable: # noqa: A001
"""Compile with ALL dims concrete — pins TMA strides at compile time.

THD/varlen: q/k/v/o/lse PACKED with batch dim 1; ``b`` = logical batch.
``has_lse=False`` compiles the LSE store out (the kernel specializes on a
``None`` LSE argument) — callers without a Stats output pass no LSE buffer
at all; the amax_o atomicMax write is independent and unchanged."""
_fake_batch = 1 if CFG.THD_VARLEN else b
fake_q = cute.runtime.make_fake_compact_tensor(
STORAGE_DTYPE,
(_fake_batch, sq, qh, CFG.TILE_K),
(b, sq, qh, CFG.TILE_K),
stride_order=(3, 2, 1, 0),
assumed_align=16,
)
fake_k = cute.runtime.make_fake_compact_tensor(
STORAGE_DTYPE,
(_fake_batch, skv, kh, CFG.TILE_K),
(b, skv, kh, CFG.TILE_K),
stride_order=(3, 2, 1, 0),
assumed_align=16,
)
fake_v = cute.runtime.make_fake_compact_tensor(
STORAGE_DTYPE,
(_fake_batch, skv, kh, CFG.TILE_O),
(b, skv, kh, CFG.TILE_O),
stride_order=(3, 2, 1, 0),
assumed_align=16,
)
fake_o = cute.runtime.make_fake_compact_tensor(
OUT_STORAGE_DTYPE,
(_fake_batch, sq, qh, CFG.TILE_O),
(b, sq, qh, CFG.TILE_O),
stride_order=(3, 2, 1, 0),
assumed_align=16,
)
Expand All @@ -1995,7 +1944,7 @@ def compile(b: int = 1, qh: int = 1, kh: int = 1, sq: int = 256, skv: int = 128,
else:
fake_lse = cute.runtime.make_fake_compact_tensor(
cutlass.Float32,
(_fake_batch, qh, sq),
(b, qh, sq),
stride_order=(2, 1, 0),
assumed_align=16,
)
Expand All @@ -2006,20 +1955,10 @@ def compile(b: int = 1, qh: int = 1, kh: int = 1, sq: int = 256, skv: int = 128,
stride_order=(0,),
assumed_align=16,
)
# Always part of the ABI; unread when CFG.SEQ_KV_LENS_PRESENT == 0. THD
# overloads it as [seq_kv_lens(B)|cu_q(B+1)|cu_k(B+1)] (len 3B+2).
_skv_len = (3 * b + 2) if CFG.THD_VARLEN else b
# Always part of the ABI; unread when CFG.SEQ_KV_LENS_PRESENT == 0.
fake_seq_kv_lens = cute.runtime.make_fake_compact_tensor(
cutlass.Int32,
(_skv_len,),
stride_order=(0,),
assumed_align=16,
)
# Per-batch O TMA-descriptor array (16 int64 = 128 B each) + 1 pad slot.
_odesc_len = (b * _TENSOR_MAP_QWORDS + _TENSOR_MAP_QWORDS) if CFG.THD_VARLEN else 1
fake_o_desc = cute.runtime.make_fake_compact_tensor(
cutlass.Int64,
(_odesc_len,),
(b,),
stride_order=(0,),
assumed_align=16,
)
Expand All @@ -2038,11 +1977,9 @@ def compile(b: int = 1, qh: int = 1, kh: int = 1, sq: int = 256, skv: int = 128,
fake_lse,
fake_sinks,
fake_seq_kv_lens,
fake_o_desc,
(b, qh, kh, sq, skv, 0),
cutlass.Float32(0.0),
cutlass.Float32(0.0),
cutlass.Int32(0),
fake_amax_o,
stream=cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=False),
options="--enable-tvm-ffi",
Expand Down
Loading