Skip to content

sdpa fp8 sm107: port the has_lse specialization; add a static SMEM guard #579

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

sdpa fp8 sm107: port the has_lse specialization; add a static SMEM guard #579

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
67 changes: 49 additions & 18 deletions python/cudnn/sdpa/fwd/kernels/prefill_d128_fp8_sm107.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -46,7 +46,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 @@ -76,6 +76,25 @@

CFG = _dc.replace(CFG, TILE_K_HW_BMM1=64, TILE_K_HW_BMM2=64, STAGES_KV=9)
Cfg = type(CFG)

# Static SMEM accounting for the post-override geometry. The 9-stage ring
# with BF16/FP16 O (~241 KiB) exceeds the STANDARD sm_10x 227 KiB per-CTA
# opt-in, and is legal on GR100 only through the sm107 oversized-SMEM
# launch mode (function attribute 16), which the required internal
# cutlass-dsl toolchain enables for its sm_107a kernels — board-validated
# e4m3->bf16 across the full suite. The guard below is against the GR100
# HARDWARE budget so a future geometry bump fails at import with a clear
# message instead of at launch.
_GR100_SMEM_BUDGET = 320 * 1024 # usable per-CTA carveout (327 KiB capacity minus reserves)
_SMEM_BYTES = (
CFG.TILES_Q * CFG.TILE_M * CFG.TILE_K * CFG.BPE # Q slabs
+ CFG.STAGES_KV * (CFG.TILE_N * CFG.TILE_K // CFG.CTA_MMA) * CFG.BPE # K ring
+ CFG.STAGES_KV * (CFG.TILE_O * CFG.TILE_N // CFG.CTA_MMA) * CFG.BPE # V ring
+ CFG.TILES_Q * CFG.TILE_M * CFG.TILE_O * CFG.BPE_O # O slabs
+ 2048 # barriers + scheduler + tmem-ptr slack (upper bound)
)
if _SMEM_BYTES > _GR100_SMEM_BUDGET:
raise ValueError(f"prefill_d128_fp8_sm107: SMEM {_SMEM_BYTES} B exceeds the GR100 budget ({_GR100_SMEM_BUDGET} B) — shrink STAGES_KV")
TMA_QK_ITERS = _TMA.QK_ITERS
TMA_VO_ITERS = _TMA.VO_ITERS
TMA_QK_GRANU_ELEMS = _TMA.QK_GRANU_ELEMS
Expand Down Expand Up @@ -246,7 +265,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 @@ -1548,7 +1567,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 @@ -1736,16 +1755,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 @@ -1885,7 +1908,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 @@ -1994,10 +2017,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 @@ -2023,12 +2049,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