Skip to content

frost(sdpa): has_lse specialization for the FP8/MXFP8 SM100 flavors #574

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 1 commit into from
Aug 13, 2026
Merged

frost(sdpa): has_lse specialization for the FP8/MXFP8 SM100 flavors #574

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
7 changes: 3 additions & 4 deletions python/cudnn/frost/README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -99,10 +99,9 @@ Reading that sequence line by line:
still forwarded to the lowered C++ graph.
- **`get_workspace_size()` is honest.** For a python plan it returns
`CompiledPlan.get_workspace_size()`, the executor's real requirement (for
the graph above on an SM100 engine: the dummy-LSE scratch `b*h*s*4 = 16384`
bytes, because an inference graph has no Stats output and the SM100 kernels
always write an LSE; `lse_optional` adapters like SM120 compile the LSE
store out instead and report 0). `execute()` forwards the caller's
the dense graph above: 0 — every SM100/SM120 adapter is `lse_optional`, so
a stats-less inference graph compiles the LSE store out and binds no dummy
buffer at any level). `execute()` forwards the caller's
buffer through `ExecutionContext.workspace` and the executor carves its
scratch out of it in 128-byte-aligned chunks, never touching bytes at or
beyond the reported size -- no hidden per-execute allocation, stable
Expand Down
38 changes: 17 additions & 21 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 @@ -951,13 +951,16 @@ def compile(self) -> None:
self._compiled_kernel = "thd-deferred"
elif self._fp8:
# FP8/MXFP8 kernels are exact-match d128 (gated in check_support);
# their compile() has no envelope head-dim parameters.
# their compile() has no envelope head-dim parameters. has_lse=False
# (no Stats output) compiles the LSE store out — no dummy buffer at
# any level (the amax_s/amax_o atomicMax writes are independent).
self._compiled_kernel = self._k_mod.compile(
b=self.batch_size,
qh=self.h_q,
kh=self.h_kv,
sq=self.s_q_max,
skv=self.s_k_max,
has_lse=self.lse_desc is not None,
)
else:
# ENVELOPE: hand the f16/bf16 kernel the ACTUAL head dims so its
Expand All @@ -983,7 +986,7 @@ def scratch_workspace_bytes(self) -> int:
Fixed by the compiled geometry (call after ``check_support()``); 0 when
the path allocates nothing per execute. This is the api-level share of
a FROST executor's ``workspace_bytes`` (the engine lowering adds its
own chunks — dummy LSE, synthesized seq_len_kv — on top; see
own chunks — synthesized seq_len_kv — on top; see
``engines.lower_dsl_prefill``). When ``execute()`` is called WITHOUT a
workspace (standalone API use), these buffers are torch-allocated per
execute instead — the carve path is what the FROST dispatch uses.
Expand Down Expand Up @@ -1064,30 +1067,21 @@ def execute(
self.lse_desc is not None and lse_tensor is None,
"lse_tensor is required by this compiled specialization",
)
# Strict presence contract, both directions: the f16 kernels are
# compiled with has_lse keyed on sample_lse (no Stats output -> the
# LSE store is compiled out and there is no LSE slot to bind), and a
# THD lse_tensor is bound in its DECLARED packed layout (recorded at
# check_support) — so an lse_tensor without a sample_lse cannot be
# honored and is rejected rather than silently dropped. The FP8/MXFP8
# kernels (dense-only) still write an LSE unconditionally; their
# stats-less write lands in a cached write-only dummy (the FROST
# dispatch never reaches it: engines.lower_dsl_prefill carves the
# dummy from the caller's workspace instead).
# Strict presence contract, both directions: every SM100 kernel
# (f16/bf16, FP8, MXFP8) is compiled with has_lse keyed on sample_lse
# (no Stats output -> the LSE store is compiled out and there is no
# LSE slot to bind), and a THD lse_tensor is bound in its DECLARED
# packed layout (recorded at check_support) — so an lse_tensor without
# a sample_lse cannot be honored and is rejected rather than silently
# dropped.
self._value_error_if(
self.lse_desc is None and lse_tensor is not None and not self._fp8,
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",
)
if self.thd:
pass # bound in _execute_thd (declared packed layout)
elif lse_tensor is not None:
lse_tensor = self._checked_lse_view(lse_tensor)
elif self._fp8:
lse_tensor = self._dummy(
"lse",
q_tensor.device,
lambda: torch.empty((self.batch_size, self.h_q, self.s_q_max), dtype=torch.float32, device=q_tensor.device),
)

if self._fp8 and self._pertensor:
# Per-tensor FP8 (sdpa_fp8): scalar descales fold into the softmax scale
Expand Down Expand Up @@ -1395,7 +1389,8 @@ def _execute_mxfp8(
sf_k_v = self._reshape_sf(sf_k, h_kv, n_kv_tiles, km.SF_SMEM_SIZE_K)
sf_v_v = self._reshape_sf(sf_v, h_kv, n_kv_tiles, km.SF_SMEM_SIZE_V)

lse = lse_tensor.reshape(b, h_q, sq)
# has_lse=False (no Stats output): the store is compiled out; bind None.
lse = lse_tensor.reshape(b, h_q, sq) if lse_tensor is not None else None
sinks_t = (
self._checked_sinks_1d(sinks) if sinks is not None else self._dummy("sinks", device, lambda: torch.zeros(h_q, dtype=torch.float32, device=device))
)
Expand Down Expand Up @@ -1488,7 +1483,8 @@ def _scalar(t, default=1.0):
O_view, o_needs_copy_back, O_scratch = self._to_bshd_writable(o_tensor)
O = O_scratch if o_needs_copy_back else O_view

lse = lse_tensor.reshape(b, h_q, sq)
# has_lse=False (no Stats output): the store is compiled out; bind None.
lse = lse_tensor.reshape(b, h_q, sq) if lse_tensor is not None else None
sinks_t = (
self._checked_sinks_1d(sinks) if sinks is not None else self._dummy("sinks", device, lambda: torch.zeros(h_q, dtype=torch.float32, device=device))
)
Expand Down
31 changes: 12 additions & 19 deletions python/cudnn/sdpa/fwd/engines.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -153,9 +153,10 @@ class Capabilities:
sink: bool = False
stats: bool = False
# The adapter accepts lse_tensor=None (its kernel None-specializes the LSE
# store), so a stats-less graph needs no dummy-LSE workspace chunk. Rows
# that keep False (the SM100 FP8/MXFP8 flavors) always write an LSE and get
# a carved dummy from lower_dsl_prefill when the graph has no Stats output.
# store), so a stats-less graph needs no dummy-LSE workspace chunk. Every
# lower_dsl_prefill row sets True (the SM100 FP8/MXFP8 flavors were the
# last holdouts); the SM80 row keeps the False default but lowers through
# its own adapter (lower_sm80_prefill), which never reads this flag.
lse_optional: bool = False
# THD lowerings assume FULLY-PACKED storage: the packed addressing is
# re-derived as prefix(lens) x token stride, and the graph's bound
Expand Down Expand Up @@ -451,6 +452,7 @@ def _sm100_mxfp8_spec(d: int) -> EngineSpec:
padded=True,
sink=True,
stats=True,
lse_optional=True,
sched_policies=frozenset({SCHED_NATURAL}),
tile_ms=frozenset({128}),
tile_ns=frozenset({128}),
Expand Down Expand Up @@ -488,6 +490,7 @@ def _sm100_fp8_spec(d: int) -> EngineSpec:
padded=True,
sink=True,
stats=True,
lse_optional=True,
# The fp8 kernel lacks the SEQ_Q_LENS_PRESENT epilogue trim, but its
# only reachable padded+stats population is KV-only padding with
# full-length seq_len_q (the fp8 suite; test_mhas_v2 fp8 padding
Expand Down Expand Up @@ -756,22 +759,15 @@ def lower_dsl_prefill(
# buffer is carved from the CALLER's workspace, so its size is fixed here at
# build time and recorded on the executor as ``workspace_bytes`` — that
# number is what the plan's CompiledPlan.get_workspace_size() reports.
# - dummy LSE (dense, stats absent, non-lse_optional adapters): the
# SM100 FP8/MXFP8 kernels always write an LSE; without a Stats output
# it lands in b*h_q*s_q fp32 scratch. lse_optional adapters (the f16
# flavors, SM120) compile the LSE store out instead and bind no
# buffer. (THD needs no engine-level LSE chunk either way.)
# - synthesized seq_len_kv (skv_tail_via_padding rows): b int32.
# - api-level scratch (api.scratch_workspace_bytes()): the dense padded
# [seq_kv|seq_q] combine and the THD metadata/LSE buffers.
dummy_lse_bytes = (
0
if (not spec.capabilities.stats or spec.capabilities.lse_optional or facts.stats_t is not None or facts.thd)
else ws_align(facts.b * facts.h_q * facts.s_q * 4)
)
# No dummy-LSE chunk: every lower_dsl_prefill row is lse_optional (the
# kernels None-specialize the LSE argument and compile the store out), so
# a stats-less graph binds no LSE buffer at any level.
synth_kv_bytes = ws_align(facts.b * 4) if synth_kv_padding else 0
api_scratch_bytes = api.scratch_workspace_bytes()
total_workspace_bytes = dummy_lse_bytes + synth_kv_bytes + api_scratch_bytes
total_workspace_bytes = synth_kv_bytes + api_scratch_bytes

binding = ga.SdpaBinding(
q=facts.q_t,
Expand Down Expand Up @@ -832,12 +828,9 @@ def _execute(variant_pack, workspace=None, stream=None):
import torch # execute path: a real tensor is about to be carved

carver = WorkspaceCarver(workspace, total_workspace_bytes, spec.name) if total_workspace_bytes else None
# Stats-less graphs bind lse_buf=None: every adapter here is
# lse_optional (the kernel compiles the LSE store out) — no dummy.
lse_buf = resolved.get(id(binding.stats)) if binding.stats is not None else None
if lse_buf is None and spec.capabilities.stats and not spec.capabilities.lse_optional and not facts.thd:
# Dummy LSE for stats-less dense graphs — carved, not allocated
# (uninitialized is fine: the kernel writes every row). lse_optional
# adapters take lse_tensor=None instead.
lse_buf = carver.take(facts.b * facts.h_q * facts.s_q, torch.float32)
# Shared presence-checked resolution (graph_analyzer.resolve_feature_operands);
# bias/block_mask/alibi stay None for these rows (mismatch gates them).
feature_ops = ga.resolve_feature_operands(facts, resolved)
Expand Down
48 changes: 30 additions & 18 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 @@ -30,7 +30,7 @@
import os
import sys
from functools import lru_cache
from typing import Callable, Tuple
from typing import Callable, Optional, Tuple


from cutlass.base_dsl.typing import Pointer # was the legacy DSL Pointer pre-DKG-bump
Expand Down Expand Up @@ -223,7 +223,7 @@ def _kernel(
tma_k_desc: cutlass.GridConstant[tmap.TensorMap],
tma_v_desc: cutlass.GridConstant[tmap.TensorMap],
tma_o_desc: cutlass.GridConstant[tmap.TensorMap],
lse_tensor: cute.Tensor,
lse_tensor: Optional[cute.Tensor],
sinks_tensor: cute.Tensor,
seq_kv_lens_tensor: cute.Tensor,
o_desc_words: cute.Tensor,
Expand Down Expand Up @@ -1522,7 +1522,7 @@ def _correction_warp_group(
tidx,
bars,
sched,
lse_tensor: cute.Tensor,
lse_tensor: Optional[cute.Tensor],
sinks_tensor: cute.Tensor,
seq_kv_lens_tensor,
n_q_supers,
Expand Down Expand Up @@ -1710,16 +1710,20 @@ def _correction_warp_group(
_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:
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
# has_lse=False: the Stats store is compiled out; the amax_s
# atomicMax is independent of it and always live.
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
nvvm.atomicrmw(nvvm.AtomicOp.MAX, _amax_s_ptr, _beta_bits)
else:
_row_valid = q_row_global < seqlen_q
if _row_valid:
lse_arr = cutlass.make_array_view(lse_tensor)
lse_row = lse_arr[batch_idx, head_idx, :]
lse_row[q_row_global] = lse_val
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
nvvm.atomicrmw(nvvm.AtomicOp.MAX, _amax_s_ptr, _beta_bits)

# amax_o = max over valid rows of |o_scaled| (the fp32 pre-cast output). Divided
Expand Down Expand Up @@ -1859,7 +1863,7 @@ def _host(
k_tensor: cute.Tensor,
v_tensor: cute.Tensor,
o_tensor: cute.Tensor,
lse_tensor: cute.Tensor,
lse_tensor: Optional[cute.Tensor],
sinks_tensor: cute.Tensor,
seq_kv_lens_tensor: cute.Tensor,
o_desc_words: cute.Tensor,
Expand Down Expand Up @@ -1965,10 +1969,13 @@ def _tma_swz(byte_w: int):


@lru_cache(maxsize=None)
def compile(b: int = 1, qh: int = 1, kh: int = 1, sq: int = 256, skv: int = 128) -> Callable: # noqa: A001
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."""
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_s/amax_o atomicMax writes are independent and unchanged."""
_fake_batch = 1 if CFG.THD_VARLEN else b
fake_q = cute.runtime.make_fake_compact_tensor(
STORAGE_DTYPE,
Expand All @@ -1994,12 +2001,17 @@ def compile(b: int = 1, qh: int = 1, kh: int = 1, sq: int = 256, skv: int = 128)
stride_order=(3, 2, 1, 0),
assumed_align=16,
)
fake_lse = cute.runtime.make_fake_compact_tensor(
cutlass.Float32,
(_fake_batch, qh, sq),
stride_order=(2, 1, 0),
assumed_align=16,
)
if not has_lse:
# No Stats output: the LSE argument is None-specialized and the store
# is compiled out entirely — no dummy buffer exists at any level.
fake_lse = None
else:
fake_lse = cute.runtime.make_fake_compact_tensor(
cutlass.Float32,
(_fake_batch, qh, sq),
stride_order=(2, 1, 0),
assumed_align=16,
)
# Always part of the ABI; unread when CFG.HAS_SINK == 0 (compile-time fold).
fake_sinks = cute.runtime.make_fake_compact_tensor(
cutlass.Float32,
Expand Down
60 changes: 40 additions & 20 deletions python/cudnn/sdpa/fwd/kernels/prefill_d128_mxfp8_sm100.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,7 @@
import os
import sys
from functools import lru_cache
from typing import Callable, Tuple
from typing import Callable, Optional, Tuple

from cutlass.experimental import primitives as nvvm
from cutlass.experimental.primitives import vote_sync, VoteSync
Expand Down Expand Up @@ -331,7 +331,7 @@ def _kernel(
tma_q_sf_desc: cutlass.GridConstant[tmap.TensorMap],
tma_k_sf_desc: cutlass.GridConstant[tmap.TensorMap],
tma_v_sf_desc: cutlass.GridConstant[tmap.TensorMap],
lse_tensor: cute.Tensor,
lse_tensor: Optional[cute.Tensor],
amax_o_tensor: cute.Tensor,
sinks_tensor: cute.Tensor,
seq_kv_lens_tensor: cute.Tensor,
Expand Down Expand Up @@ -1979,7 +1979,7 @@ def _correction_warp_group(
tidx,
bars,
sched,
lse_tensor: cute.Tensor,
lse_tensor: Optional[cute.Tensor],
amax_o_tensor: cute.Tensor,
sinks_tensor: cute.Tensor,
seq_kv_lens_tensor,
Expand Down Expand Up @@ -2141,16 +2141,20 @@ def _correction_warp_group(
_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:
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
# has_lse=False: the Stats store is compiled out; the amax_o
# atomicMax below is independent of it and always live.
if cutlass.const_expr(lse_tensor is not None):
if _row_valid:
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:
lse_arr = cutlass.make_array_view(lse_tensor)
lse_row = lse_arr[batch_idx, head_idx, :]
lse_row[q_row_global] = lse_val
if cutlass.const_expr(lse_tensor is not None):
if _row_valid:
lse_arr = cutlass.make_array_view(lse_tensor)
lse_row = lse_arr[batch_idx, head_idx, :]
lse_row[q_row_global] = lse_val

sO_sub_base = sO[qs].base

Expand Down Expand Up @@ -2234,7 +2238,7 @@ def _host(
sf_q_tensor: cute.Tensor,
sf_k_tensor: cute.Tensor,
sf_v_tensor: cute.Tensor,
lse_tensor: cute.Tensor,
lse_tensor: Optional[cute.Tensor],
amax_o_tensor: cute.Tensor,
sinks_tensor: cute.Tensor,
seq_kv_lens_tensor: cute.Tensor,
Expand Down Expand Up @@ -2387,11 +2391,22 @@ def _build_sf_desc(sf_tensor, num_tiles, sf_smem_size, num_rows_box, num_heads):


@lru_cache(maxsize=None)
def compile(
b: int = 1, qh: int = 1, kh: int = 1, sq: int = 256, skv: int = 128, total_q_sf_tiles: int = 0, total_kv_sf_tiles: int = 0 # noqa: A001
def compile( # noqa: A001
b: int = 1,
qh: int = 1,
kh: int = 1,
sq: int = 256,
skv: int = 128,
total_q_sf_tiles: int = 0,
total_kv_sf_tiles: int = 0,
has_lse: bool = True,
) -> Callable:
"""Compile a kernel with concrete dims; 3 SF tensors layout [B, H, num_seq_tiles, SF_SMEM_SIZE_*].

``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.

THD/varlen: q/k/v/o/lse + SF tensors are PACKED with batch dim 1; ``b`` is the
LOGICAL batch (sequence count). The SF tensors are per-sequence-TILE-padded so
their packed tile extent is ``total_q_sf_tiles`` / ``total_kv_sf_tiles`` (=
Expand Down Expand Up @@ -2455,12 +2470,17 @@ def compile(
assumed_align=16,
)

fake_lse = cute.runtime.make_fake_compact_tensor(
cutlass.Float32,
(_fake_batch, qh, sq),
stride_order=(2, 1, 0),
assumed_align=16,
)
if not has_lse:
# No Stats output: the LSE argument is None-specialized and the store
# is compiled out entirely — no dummy buffer exists at any level.
fake_lse = None
else:
fake_lse = cute.runtime.make_fake_compact_tensor(
cutlass.Float32,
(_fake_batch, qh, sq),
stride_order=(2, 1, 0),
assumed_align=16,
)
fake_amax_o = cute.runtime.make_fake_compact_tensor(
cutlass.Float32,
(1,),
Expand Down
Loading