diff --git a/python/cudnn/sdpa/fwd/api_dsl.py b/python/cudnn/sdpa/fwd/api_dsl.py index 43c6d1f1b..3bf32c4fe 100644 --- a/python/cudnn/sdpa/fwd/api_dsl.py +++ b/python/cudnn/sdpa/fwd/api_dsl.py @@ -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, ) @@ -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 @@ -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}", @@ -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") @@ -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)", @@ -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. diff --git a/python/cudnn/sdpa/fwd/config_sm120.py b/python/cudnn/sdpa/fwd/config_sm120.py index 7b5d0ed41..d72dbbf51 100644 --- a/python/cudnn/sdpa/fwd/config_sm120.py +++ b/python/cudnn/sdpa/fwd/config_sm120.py @@ -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 @@ -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 @@ -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: diff --git a/python/cudnn/sdpa/fwd/engines.py b/python/cudnn/sdpa/fwd/engines.py index 926d83693..e0dc8dd05 100644 --- a/python/cudnn/sdpa/fwd/engines.py +++ b/python/cudnn/sdpa/fwd/engines.py @@ -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) @@ -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)" @@ -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}), diff --git a/python/cudnn/sdpa/fwd/kernels/prefill_f16_sm120.py b/python/cudnn/sdpa/fwd/kernels/prefill_f16_sm120.py index 610982532..cf3e1d3df 100644 --- a/python/cudnn/sdpa/fwd/kernels/prefill_f16_sm120.py +++ b/python/cudnn/sdpa/fwd/kernels/prefill_f16_sm120.py @@ -30,7 +30,9 @@ Constraints: * Supported input dtypes: Float16 and BFloat16, output dtype must match input dtype -* Head dimension must be a multiple of 16 between 16 and 256, inclusive +* Runtime head dimensions must be multiples of 8 up to 256 (TMA 16-byte + global-stride rule); the kernel compiles at head TILES rounded up to 16 and + TMA zero-fills the pad columns (the head-dim ENVELOPE) * Q heads must be divisible by the number of K/V heads * Q/K/V/O use compact BSHD storage * Supported CTA Q/KV tiles are 128 or 64 @@ -52,6 +54,7 @@ from cudnn.frost.tile_dsl.mma import ptx_mma_m16n8k16_f32 from cudnn.frost.tile_dsl.swizzle import swizzle_xor from cudnn.sdpa.fwd.config_sm120 import ( + HEAD_TILE_GRANULE, SEQ_KV_TILES as _SEQ_KV_TILES, SEQ_Q_TILES as _SEQ_Q_TILES, SUPPORTED_HEAD_TILES as _SUPPORTED_HEAD_TILES, @@ -132,6 +135,11 @@ def ceil_div(a: int, b: int) -> int: return (a + b - 1) // b +def round_up_head_tile(d: int) -> int: + """Round a head dim up to the head-tile granule — the ENVELOPE head tile.""" + return ceil_div(d, HEAD_TILE_GRANULE) * HEAD_TILE_GRANULE + + fmul2 = partial(prims.mul_packed_f32x2, ftz=False, rnd=prims.FPRoundingMode.RN) fma2 = partial(prims.fma_packed_f32x2, ftz=False, rnd=prims.FPRoundingMode.RN) @@ -173,6 +181,7 @@ def __init__( is_causal: bool = False, bottom_right: bool = False, window_size_left: int | None = None, + window_size_right: int | None = None, seq_q_lens_present: bool = False, seq_kv_lens_present: bool = False, has_sink: bool = False, @@ -192,6 +201,11 @@ def __init__( :param is_causal: Apply an upper causal bound to QK. :param bottom_right: Shift the causal diagonal by ``Skv - Sq``. :param window_size_left: Inclusive left-window offset, or ``None``. + :param window_size_right: Widen the causal diagonal to the right by + this many columns (inclusive; 0 = plain causal). Only meaningful + with ``is_causal`` — ``compile()`` maps the band model's + ``window_right`` to ``is_causal=(window_right is not None)`` plus + this offset. :param seq_q_lens_present: Read per-batch query lengths at runtime. :param seq_kv_lens_present: Read per-batch key/value lengths at runtime. :param has_sink: Fold the per-Q-head sink logit from the ``sinks`` @@ -207,9 +221,12 @@ def __init__( ``(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 + :param head_tile_qk: Q/K head TILE (the QK^T contraction width). Must + be a multiple of 16 between 16 and 256, inclusive. Runtime head + dims may be any multiple of 8 that rounds up to this tile — the + TMA descriptors keep the actual extents and zero-fill the pad + columns (the head-dim ENVELOPE). + :param head_tile_v: V/O head TILE (the P@V output width). Same constraint as ``head_tile_qk``. :param q_tile: Query sequence tile size. :param kv_tile: Key/value sequence tile size. @@ -224,6 +241,14 @@ def __init__( self.is_causal = is_causal self.bottom_right = bottom_right self.window_size_left = window_size_left + self.window_size_right = window_size_right + # Band model: a right bound implies the causal upper limit (compile() + # maps window_right -> is_causal), so the masking sites key off + # is_causal and add the (compile-time) widening. + self.right_slack = window_size_right if window_size_right is not None else 0 + # A translated diagonal (bottom-right anchoring or a right band) can + # straddle one more KV tile than the tile-aligned top-left one. + self.diag_shifted = bottom_right or self.right_slack > 0 self.seq_q_lens_present = seq_q_lens_present self.seq_kv_lens_present = seq_kv_lens_present self.has_sink = has_sink @@ -300,29 +325,62 @@ def load_one_kv_tile( batch_idx: cutlass.Int32, head_idx: cutlass.Int32, seq_coord: cutlass.Int32, + is_v: cutlass.Constexpr[bool], + envelope: cutlass.Constexpr[bool], ) -> None: - """Launch one TMA load for a complete K/V tile into swizzled SMEM. - - The tensor map exposes compact ``(B, S, H, D)`` storage through a - logical ``(B, H, I, S, C)`` view, where ``D = I * C``. Its TMA-order - dimensions are ``(C, S, I, H, B)``, so one rank-5 copy covers every - head chunk and uses coordinates ``(c, seq, i, head, batch)``. + """Launch the TMA load(s) for a complete K/V tile into swizzled SMEM. + + Exact head dims (``envelope=False``, the common case) issue ONE rank-5 + copy whose descriptor pre-splits the head dim into swizzle-span chunks + — the head boundary coincides with the tile so no zero-fill is needed. + + Envelope head dims (``envelope=True``, actual d < compile-time tile) + issue ``chunks`` copies over a rank-4 descriptor that keeps the ACTUAL + head extent as the innermost dimension, stepping the head coordinate + by ``chunk_elems``. Head columns at or past the actual extent are + outside that dimension, so the hardware zero-fills them (zero K columns + add exact zero terms to every Q@K^T; zero V columns produce O columns + the store guard clips). A single copy cannot serve this case: TMA + bounds-checks each coordinate against its OWN dimension, so a + chunk-dimension descriptor would fetch the next head's data instead of + zeros past d. :param s_dst: Swizzled SMEM destination tile. - :param tma_desc: K or V tensor map descriptor. + :param tma_desc: K or V tensor map descriptor (rank matches + ``envelope``). :param mbar: TMA completion mbarrier for this stream. :param batch_idx: Batch index. :param head_idx: Attention head index. :param seq_coord: Starting sequence row for the K/V tile. + :param is_v: Selects the V-side chunk geometry over the K-side one + (the two carry independent head tiles and swizzle spans). + :param envelope: Actual head dim < compile-time tile (zero-padded). """ + chunks = self.v_tma_swizzle_chunks if is_v else self.k_tma_swizzle_chunks + chunk_elems = self.v_swizzle_chunk_elems if is_v else self.k_swizzle_chunk_elems if prims.elect_sync(): - prims.mbarrier_arrive_expect_tx(mbar, tma_desc.global_tx_bytes()) - prims.cp_async_bulk_tensor_shared_cta_global( - s_dst, - tma_desc.get_ptr(), - (0, seq_coord, 0, head_idx, batch_idx), - mbar, - ) + if cutlass.const_expr(envelope): + # Every copy completes with its full box (OOB regions arrive as + # zeros but still count), so the expected transaction total is + # simply chunks x the per-copy box bytes. + prims.mbarrier_arrive_expect_tx(mbar, chunks * tma_desc.global_tx_bytes()) + for i in cutlass.range_constexpr(chunks): + prims.cp_async_bulk_tensor_shared_cta_global( + s_dst.subview(i * self.kv_tile * chunk_elems), + tma_desc.get_ptr(), + (i * chunk_elems, seq_coord, head_idx, batch_idx), + mbar, + ) + else: + # Rank-5 coordinates (c, seq, i, head, batch): one copy covers + # every head chunk of the tile. + prims.mbarrier_arrive_expect_tx(mbar, tma_desc.global_tx_bytes()) + prims.cp_async_bulk_tensor_shared_cta_global( + s_dst, + tma_desc.get_ptr(), + (0, seq_coord, 0, head_idx, batch_idx), + mbar, + ) @cute.jit def load_q_tile( @@ -495,9 +553,11 @@ def online_softmax( valid_cols = basic_params.seqlen_k if cutlass.const_expr(self.is_causal): + # Causal upper bound, widened right by the (compile-time) band + # slack — 0 for plain causal, R for diagonal_band_right_bound. valid_cols = cute.math.max( cutlass.Int32(0), - cute.math.min(diagonal_position + 1, basic_params.seqlen_k), + cute.math.min(diagonal_position + 1 + self.right_slack, basic_params.seqlen_k), ) first_valid_col = cutlass.Int32(0) @@ -750,8 +810,9 @@ def kernel( tidx, _, _ = cute.arch.thread_idx() q_tile_idx, batch_idx, head_idx = cute.arch.block_idx() if cutlass.const_expr(self.is_causal): - # Causal work grows with the Q tile. Launch long tiles first to - # avoid leaving a few expensive CTAs in the final scheduler waves. + # Diagonal-bounded work grows with the Q tile. Launch long tiles + # first to avoid leaving a few expensive CTAs in the final + # scheduler waves (right-band graphs share the causal shape). grid_q, _, _ = cute.arch.grid_dim() q_tile_idx = grid_q - q_tile_idx - 1 q_seq_idx = q_tile_idx * self.q_tile @@ -789,6 +850,10 @@ def kernel( num_heads_kv = k.shape[2] head_dim_qk = q.shape[3] head_dim_v = v.shape[3] + # ENVELOPE flags (static shapes): actual dim < compile-time tile means + # the TMA loads must zero-fill the pad columns via per-chunk copies. + k_envelope = head_dim_qk != self.head_tile_qk + v_envelope = head_dim_v != self.head_tile_v q_ptr = q.iterator.raw_ptr() o_ptr = o.iterator.raw_ptr() @@ -814,7 +879,7 @@ def kernel( if q_seq_idx >= seqlen_q: num_kv_tiles = cutlass.Int32(0) if cutlass.const_expr(self.is_causal): - causal_k_end = q_seq_idx + self.q_tile + causal_k_end = q_seq_idx + self.q_tile + self.right_slack if cutlass.const_expr(self.bottom_right): causal_k_end += seqlen_k - seqlen_q causal_k_end = cute.math.max(cutlass.Int32(0), cute.math.min(causal_k_end, seqlen_k)) @@ -884,6 +949,8 @@ def kernel( tma_batch_idx, kv_head_idx, kv_row_base + kv_seq_idx, + is_v=False, + envelope=k_envelope, ) self.load_one_kv_tile( sV, @@ -892,6 +959,8 @@ def kernel( tma_batch_idx, kv_head_idx, kv_row_base + kv_seq_idx, + is_v=True, + envelope=v_envelope, ) kv_seq_idx -= self.kv_tile @@ -907,6 +976,8 @@ def kernel( tma_batch_idx, kv_head_idx, kv_row_base + kv_seq_idx, + is_v=False, + envelope=k_envelope, ) prims.barrier_cta_sync( @@ -920,6 +991,8 @@ def kernel( tma_batch_idx, kv_head_idx, kv_row_base + kv_seq_idx, + is_v=True, + envelope=v_envelope, ) kv_seq_idx -= self.kv_tile # ///////////////////////////////////////////////////////////////////////////// @@ -990,8 +1063,11 @@ def kernel( mask_steps = 1 if cutlass.const_expr(self.is_causal): mask_steps = ceil_div(self.q_tile, self.kv_tile) - if cutlass.const_expr(self.bottom_right): - # The shifted diagonal can straddle one additional KV tile. + if cutlass.const_expr(self.diag_shifted): + # A translated diagonal (bottom-right anchoring or a right + # band) can straddle one additional KV tile; the frontier + # width itself is R-independent — the band only translates + # the diagonal. mask_steps = ceil_div(self.q_tile + self.kv_tile - 1, self.kv_tile) left_mask_steps = 1 if cutlass.const_expr(self.window_size_left is not None): @@ -1225,10 +1301,12 @@ def __call__( """ head_dim_qk = q.shape[3] head_dim_v = v.shape[3] - if cutlass.const_expr(head_dim_qk != k.shape[3] or head_dim_qk != self.head_tile_qk): - raise ValueError("runtime Q/K head dimensions must match the kernel head_tile_qk") - if cutlass.const_expr(head_dim_v != o.shape[3] or head_dim_v != self.head_tile_v): - raise ValueError("runtime V/O head dimensions must match the kernel head_tile_v") + if cutlass.const_expr(head_dim_qk != k.shape[3] or round_up_head_tile(head_dim_qk) != self.head_tile_qk): + raise ValueError("runtime Q/K head dimensions must round up (by the head-tile granule) to the kernel head_tile_qk") + if cutlass.const_expr(head_dim_v != o.shape[3] or round_up_head_tile(head_dim_v) != self.head_tile_v): + raise ValueError("runtime V/O head dimensions must round up (by the head-tile granule) to the kernel head_tile_v") + if cutlass.const_expr(head_dim_qk % 8 != 0 or head_dim_v % 8 != 0): + raise ValueError("head dimensions must be multiples of 8 (TMA 16-byte global-stride rule at 2 B/elem)") if cutlass.const_expr( q.shape[0] != k.shape[0] or k.shape[:3] != v.shape[:3] @@ -1269,66 +1347,42 @@ def __call__( raise ValueError("THD seq_kv_lens must be the (3*B+2,) metadata tensor") # Split D into I contiguous C-element chunks while preserving the - # compact (B, S, H, D) global-memory address calculation. TMA order - # (C, S, I, H, B) linearizes the SMEM destination as [I][kv_tile][C]. - k_tma_layout = cute.make_layout( - ( - k.shape[0], - k.shape[2], - self.k_tma_swizzle_chunks, - k.shape[1], - self.k_swizzle_chunk_elems, - ), - stride=( - k.shape[1] * k.shape[2] * head_dim_qk, - head_dim_qk, - self.k_swizzle_chunk_elems, - k.shape[2] * head_dim_qk, - 1, - ), - ) - k_tma_box = ( - 1, - 1, - self.k_tma_swizzle_chunks, - self.kv_tile, - self.k_swizzle_chunk_elems, - ) - tma_k_desc = cuda.create_tensor_map_tiled_from_view( - cute.make_tensor(k.iterator, k_tma_layout), - box_dims=k_tma_box, - stride_order=(4, 3, 2, 1, 0), - swizzle=self.k_tma_swizzle, - ) - v_tma_layout = cute.make_layout( - ( - v.shape[0], - v.shape[2], - self.v_tma_swizzle_chunks, - v.shape[1], - self.v_swizzle_chunk_elems, - ), - stride=( - v.shape[1] * v.shape[2] * head_dim_v, - head_dim_v, - self.v_swizzle_chunk_elems, - v.shape[2] * head_dim_v, - 1, - ), - ) - v_tma_box = ( - 1, - 1, - self.v_tma_swizzle_chunks, - self.kv_tile, - self.v_swizzle_chunk_elems, - ) - tma_v_desc = cuda.create_tensor_map_tiled_from_view( - cute.make_tensor(v.iterator, v_tma_layout), - box_dims=v_tma_box, - stride_order=(4, 3, 2, 1, 0), - swizzle=self.v_tma_swizzle, - ) + # per-tensor TMA descriptor over the compact (B, S, H, D) storage. + # + # Exact head dim (the common case): a rank-5 view pre-splits D into + # swizzle-span chunks as a descriptor dimension — dims (B, H, I, S, C) + # with D = I * C, TMA order (C, S, I, H, B) — so ONE copy covers the + # whole tile (the fast path; ~1-2% speedup upon exact-dim workloads). + # + # Envelope head dim (actual d < compile-time tile): the head boundary + # must be a single descriptor dimension for TMA's per-dimension bounds + # check to zero-fill past it, so a rank-4 view keeps the ACTUAL extent + # innermost — TMA order (D, S, H, B) — and load_one_kv_tile steps the + # head coordinate per swizzle-span chunk. + def kv_tma_desc(t, head_dim, head_tile, swizzle, swizzle_chunks, swizzle_chunk_elems): + if cutlass.const_expr(head_dim == head_tile): + layout = cute.make_layout( + (t.shape[0], t.shape[2], swizzle_chunks, t.shape[1], swizzle_chunk_elems), + stride=(t.shape[1] * t.shape[2] * head_dim, head_dim, swizzle_chunk_elems, t.shape[2] * head_dim, 1), + ) + box = (1, 1, swizzle_chunks, self.kv_tile, swizzle_chunk_elems) + stride_order = (4, 3, 2, 1, 0) + else: + layout = cute.make_layout( + (t.shape[0], t.shape[2], t.shape[1], head_dim), + stride=(t.shape[1] * t.shape[2] * head_dim, head_dim, t.shape[2] * head_dim, 1), + ) + box = (1, 1, self.kv_tile, swizzle_chunk_elems) + stride_order = (3, 2, 1, 0) + return cuda.create_tensor_map_tiled_from_view( + cute.make_tensor(t.iterator, layout), + box_dims=box, + stride_order=stride_order, + swizzle=swizzle, + ) + + tma_k_desc = kv_tma_desc(k, head_dim_qk, self.head_tile_qk, self.k_tma_swizzle, self.k_tma_swizzle_chunks, self.k_swizzle_chunk_elems) + tma_v_desc = kv_tma_desc(v, head_dim_v, self.head_tile_v, self.v_tma_swizzle, self.v_tma_swizzle_chunks, self.v_swizzle_chunk_elems) self.kernel( q, k, @@ -1394,6 +1448,7 @@ def compile( # noqa: A001 is_causal=PARAMS.window_right is not None, bottom_right=PARAMS.bottom_right, window_size_left=PARAMS.window_left, + window_size_right=PARAMS.window_right, seq_q_lens_present=PARAMS.seq_q_lens_present, seq_kv_lens_present=PARAMS.seq_kv_lens_present, has_sink=PARAMS.has_sink, @@ -1401,8 +1456,8 @@ def compile( # noqa: A001 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, + head_tile_qk=round_up_head_tile(d_qk), + head_tile_v=round_up_head_tile(d_v), q_tile=PARAMS.q_tile, kv_tile=PARAMS.kv_tile, ) diff --git a/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm120.py b/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm120.py index 8741e5488..023527490 100644 --- a/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm120.py +++ b/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm120.py @@ -98,6 +98,7 @@ def _ref_sdpa_full( is_causal: bool = False, causal_bottom_right: bool = False, window_size_left: int | None = None, + window_size_right: int | None = None, seq_q_lens: torch.Tensor | None = None, seq_kv_lens: torch.Tensor | None = None, sinks: torch.Tensor | None = None, @@ -106,7 +107,9 @@ def _ref_sdpa_full( """fp32 reference matching the SM120 DSL kernel's mask semantics. q/k/v are BHSD; GQA (h_q > h_kv) is handled by expanding K/V. ``sinks`` is one logit per Q head, joining the softmax as a virtual column with no - V row. + V row. ``window_size_right`` widens the diagonal to the right by R + columns (inclusive; keep ``j <= lim + R``) — cuDNN's + diagonal_band_right_bound, exclusive with ``is_causal``. """ b, h_q, s_q, _ = q.shape @@ -125,6 +128,8 @@ def _ref_sdpa_full( masked = (i >= q_lens.view(b, 1, 1, 1)) | (j >= kv_lens.view(b, 1, 1, 1)) if is_causal: masked = masked | (j > lim) + if window_size_right is not None: + masked = masked | (j > lim + window_size_right) if window_size_left is not None: masked = masked | (j < lim - window_size_left) scores = scores.masked_fill(masked, float("-inf")) @@ -157,6 +162,7 @@ def _run_case( is_causal: bool = False, causal_bottom_right: bool = False, window_size_left: int | None = None, + window_size_right: int | None = None, seq_q_lens: torch.Tensor | None = None, seq_kv_lens: torch.Tensor | None = None, scale: float | None = None, @@ -172,6 +178,7 @@ def _run_case( is_causal=is_causal, causal_bottom_right=causal_bottom_right, window_size_left=window_size_left, + window_size_right=window_size_right, seq_q_lens=seq_q_lens, seq_kv_lens=seq_kv_lens, sinks=sinks, @@ -196,6 +203,27 @@ def _run_case( torch.testing.assert_close(output.float(), expected, atol=0.1, rtol=5e-2) +def _apply_mask_kwargs(sdpa_kwargs, cudnn, *, is_causal, causal_bottom_right, window_size_left, window_size_right): + """Translate the reference mask vocabulary into graph sdpa kwargs. + + A right bound makes the mask a diagonal BAND (with the requested + alignment; a left bound rides along, cuDNN length L = offset W + 1); + otherwise the causal / sliding-window flags apply. + """ + if window_size_right is not None: + sdpa_kwargs["diagonal_band_right_bound"] = window_size_right + sdpa_kwargs["diagonal_alignment"] = cudnn.diagonal_alignment.BOTTOM_RIGHT if causal_bottom_right else cudnn.diagonal_alignment.TOP_LEFT + if window_size_left is not None: + sdpa_kwargs["diagonal_band_left_bound"] = window_size_left + 1 + else: + if causal_bottom_right: + sdpa_kwargs["use_causal_mask_bottom_right"] = True + elif is_causal: + sdpa_kwargs["use_causal_mask"] = True + if window_size_left is not None: + sdpa_kwargs["sliding_window_length"] = window_size_left + 1 + + def _run_dsl_graph( q_gpu: torch.Tensor, k_gpu: torch.Tensor, @@ -206,6 +234,7 @@ def _run_dsl_graph( is_causal: bool = False, causal_bottom_right: bool = False, window_size_left: int | None = None, + window_size_right: int | None = None, seq_q_lens: torch.Tensor | None = None, seq_kv_lens: torch.Tensor | None = None, sinks: torch.Tensor | None = None, @@ -250,12 +279,14 @@ def _run_dsl_graph( } variant_pack = {q: q_gpu, k: k_gpu, v: v_gpu} - if causal_bottom_right: - sdpa_kwargs["use_causal_mask_bottom_right"] = True - elif is_causal: - sdpa_kwargs["use_causal_mask"] = True - if window_size_left is not None: - sdpa_kwargs["sliding_window_length"] = window_size_left + 1 + _apply_mask_kwargs( + sdpa_kwargs, + cudnn, + is_causal=is_causal, + causal_bottom_right=causal_bottom_right, + window_size_left=window_size_left, + window_size_right=window_size_right, + ) if seq_q_lens is not None or seq_kv_lens is not None: assert seq_q_lens is not None and seq_kv_lens is not None seq_q = graph.tensor_like(seq_q_lens, name="seq_q") @@ -344,6 +375,7 @@ def _run_thd_case( is_causal: bool = False, causal_bottom_right: bool = False, window_size_left: int | None = None, + window_size_right: int | None = None, with_sink: bool = False, check_stats: bool = False, stats_layout: str = "token_major", @@ -406,12 +438,14 @@ def _run_thd_case( seq_len_q=sq, seq_len_kv=skv, ) - if causal_bottom_right: - sdpa_kwargs["use_causal_mask_bottom_right"] = True - elif is_causal: - sdpa_kwargs["use_causal_mask"] = True - if window_size_left is not None: - sdpa_kwargs["sliding_window_length"] = window_size_left + 1 + _apply_mask_kwargs( + sdpa_kwargs, + cudnn, + is_causal=is_causal, + causal_bottom_right=causal_bottom_right, + window_size_left=window_size_left, + window_size_right=window_size_right, + ) variant_pack = {tq: q_view, tk: k_view, tv: v_view, rq: q_ro, rk: k_ro, rv: v_ro, ro: o_ro, sq: sq_t, skv: skv_t} if sinks is not None: st = graph.tensor_like(sinks, name="sink") @@ -483,6 +517,7 @@ def _run_thd_case( is_causal=is_causal, causal_bottom_right=causal_bottom_right, window_size_left=window_size_left, + window_size_right=window_size_right, sinks=sinks, return_stats=check_stats, ) @@ -824,6 +859,14 @@ def test_dsl_sm120_execute_contract_mismatches(): with pytest.raises(ValueError, match="without an LSE output"): api.execute(q_tensor=q, k_tensor=k, v_tensor=v, o_tensor=o, seq_q_lens=seq_kv, seq_kv_lens=seq_kv, lse_tensor=lse_thd) + # Right-band contract (band model): window_size_right is the causal + # diagonal's right bound (0 = plain causal) and therefore requires + # is_causal; negative bounds are rejected. + with pytest.raises(ValueError, match="requires is_causal"): + SdpaFwdDslSm120(sample_q=q, sample_k=k, sample_v=v, sample_o=o, window_size_right=8).check_support() + with pytest.raises(ValueError, match="window_size_right must be >= 0"): + SdpaFwdDslSm120(sample_q=q, sample_k=k, sample_v=v, sample_o=o, is_causal=True, window_size_right=-1).check_support() + @pytest.mark.L0 @torch_fork_set_rng(seed=22) @@ -1123,6 +1166,126 @@ def test_dsl_sm120_causal_bottom_right(): ) +@pytest.mark.L0 +@pytest.mark.parametrize( + ("head_dim", "head_dim_v"), + [(72, 72), (104, 72), (8, 8), (200, 136)], + ids=["d72", "mixed104x72", "d8_min", "d200x136"], +) +@torch_fork_set_rng(seed=34) +def test_dsl_sm120_head_dim_envelope(head_dim: int, head_dim_v: int): + """Head dims that are multiples of 8 but not 16: the kernel compiles at + tiles rounded up to 16 and the per-chunk TMA copies zero-fill columns + past the actual extents — S, softmax, and P@V are bit-identical to the + unpadded problem, and O stores clip at the actual D_V.""" + _run_case(batch=2, h_q=4, h_kv=4, s_q=192, s_kv=256, head_dim=head_dim, head_dim_v=head_dim_v, is_causal=True) + + +@pytest.mark.L0 +@torch_fork_set_rng(seed=35) +def test_dsl_sm120_head_dim_envelope_features(): + """The envelope composed with the feature family: padded + stats (LSE + trim with pad columns), sink, and a THD ragged batch — plus a d=248 case + that forces the auto kv_tile=64 pick (per-chunk XOR phase at the smaller + tile).""" + seq_q_lens = torch.tensor([150, 96], dtype=torch.int32, device="cuda") + seq_kv_lens = torch.tensor([200, 128], dtype=torch.int32, device="cuda") + _run_case( + batch=2, + h_q=4, + h_kv=2, + s_q=192, + s_kv=256, + head_dim=104, + head_dim_v=72, + seq_q_lens=seq_q_lens, + seq_kv_lens=seq_kv_lens, + check_stats=True, + ) + _run_case(batch=2, h_q=4, h_kv=4, s_q=128, s_kv=128, head_dim=88, dtype=torch.bfloat16, with_sink=True, check_stats=True) + # Envelope pad columns and a ragged S_kv tail in the SAME rightmost + # tile: both zero-fill mechanisms at once, with the LSE checked. + _run_case(batch=2, h_q=4, h_kv=4, s_q=192, s_kv=300, head_dim=104, head_dim_v=72, check_stats=True) + _run_case(head_dim=248, head_dim_v=248, s_q=128, s_kv=128) # auto kv_tile=64 + _run_thd_case( + seq_q_lens=[130, 70], + seq_kv_lens=[130, 70], + h_q=4, + h_kv=2, + head_dim=104, + head_dim_v=72, + is_causal=True, + check_stats=True, + ) + + +@pytest.mark.L0 +@torch_fork_set_rng(seed=36) +def test_dsl_sm120_right_band(): + """diagonal_band_right_bound > 0: the causal machinery with the diagonal + widened right by a compile-time constant R (keep j <= diag + R, + inclusive). Cases: TOP_LEFT band; BOTTOM_RIGHT band (a band graph is + NON-causal in cuDNN's vocabulary — no causal flag involved); a full band + (left + right bounds); R + stats; R across a ragged S_kv tail.""" + + _run_case(batch=2, h_q=4, h_kv=4, s_q=256, s_kv=256, window_size_right=32) + _run_case(batch=2, h_q=4, h_kv=4, s_q=192, s_kv=320, causal_bottom_right=True, window_size_right=48) + _run_case(batch=2, h_q=4, h_kv=4, s_q=256, s_kv=256, window_size_left=64, window_size_right=32) + _run_case(batch=2, h_q=8, h_kv=2, s_q=256, s_kv=256, window_size_right=100, check_stats=True) + _run_case(batch=2, h_q=4, h_kv=4, s_q=256, s_kv=300, window_size_right=32) + # BR band + SWA + per-batch lengths: the runtime per-batch diagonal + # offset, the R-shifted right edge, and the UNSHIFTED left anchor in + # one launch. + seq_q_lens = torch.tensor([230, 120], dtype=torch.int32, device="cuda") + seq_kv_lens = torch.tensor([180, 240], dtype=torch.int32, device="cuda") + _run_case( + batch=2, + h_q=4, + h_kv=4, + s_q=256, + s_kv=256, + causal_bottom_right=True, + window_size_right=48, + window_size_left=96, + seq_q_lens=seq_q_lens, + seq_kv_lens=seq_kv_lens, + ) + # kv_tile=64 (auto-picked at d=248) with R a multiple of the tile: the + # 3-step masked frontier at the smaller tile. + _run_case(head_dim=248, head_dim_v=248, s_q=128, s_kv=128, window_size_right=64) + # Degenerate R >= S_kv: the widened bound clamps to full visibility. + _run_case(batch=2, h_q=4, h_kv=4, s_q=128, s_kv=128, window_size_right=200) + + +@pytest.mark.L1 +@torch_fork_set_rng(seed=37) +def test_dsl_sm120_thd_right_band(): + """THD + TOP_LEFT right band: per-sequence diagonals each widened by R, + with the ragged Stats checked.""" + + _run_thd_case(seq_q_lens=[130, 70], seq_kv_lens=[130, 70], window_size_right=24, check_stats=True) + # BOTTOM_RIGHT band under THD: each sequence's own diagonal, widened. + _run_thd_case(seq_q_lens=[100, 60], seq_kv_lens=[180, 120], causal_bottom_right=True, window_size_right=24, check_stats=True) + + +@pytest.mark.L0 +@torch_fork_set_rng(seed=33) +def test_dsl_sm120_ragged_skv_tail(): + """S_kv not a multiple of the KV tile, served natively (skv_tile=0): the + kernel's first masked step covers the partial rightmost tile in every + configuration — no padding mask, no synthesized lengths. + + Cases: dense unmasked; top-left causal with S_q > S_kv (the corner causal_covers_tail + excludes); causal + sliding window across a ragged tail; ragged tail with + the LSE checked.""" + + _run_case(batch=2, h_q=4, h_kv=4, s_q=256, s_kv=300, head_dim=128) + _run_case(batch=2, h_q=4, h_kv=4, s_q=384, s_kv=200, head_dim=128, is_causal=True) + _run_case(batch=2, h_q=4, h_kv=4, s_q=256, s_kv=300, head_dim=128, is_causal=True, window_size_left=96) + _run_case(batch=2, h_q=8, h_kv=2, s_q=192, s_kv=333, head_dim=64, check_stats=True) + _run_case(batch=2, h_q=4, h_kv=4, s_q=128, s_kv=40, head_dim=128) # num_kv_tiles == 1, tail-only tile + + @pytest.mark.L0 @torch_fork_set_rng(seed=6) def test_dsl_sm120_padded(): diff --git a/test/python/sdpa/frost/test_sdpa_graph_analyzer.py b/test/python/sdpa/frost/test_sdpa_graph_analyzer.py index f6b906ccc..66aff1bb7 100644 --- a/test/python/sdpa/frost/test_sdpa_graph_analyzer.py +++ b/test/python/sdpa/frost/test_sdpa_graph_analyzer.py @@ -618,10 +618,51 @@ def test_sm120_probe_rejects_on_sm100_family(): assert _SM120 not in _eligible(_mk_sm120_graph()) -def test_sm120_probe_rejects_head_dim_not_multiple_of_16(monkeypatch): +def test_sm120_probe_head_dim_envelope(monkeypatch): + # d_envelope: any multiple of 8 up to the 256 cap is served via TMA + # zero-padding; only sub-8 alignment (TMA 16-byte global-stride rule) + # stays ineligible. monkeypatch.setattr(ga, "_device_cc", lambda: (12, 0)) assert _SM120 in _eligible(_mk_sm120_graph(d=192)) - assert not _eligible(_mk_sm120_graph(d=136)) # multiple of 8, not of 16 + assert _SM120 in _eligible(_mk_sm120_graph(d=136)) # multiple of 8, not of 16 + assert not _eligible(_mk_sm120_graph(d=132)) # multiple of 4, not of 8 + + +def test_sm120_probe_accepts_right_band_widening(monkeypatch): + # diagonal_band_right_bound > 0 is served by the SM120 row (the causal + # machinery with a widened diagonal); the SM100 rows keep rejecting it. + monkeypatch.setattr(ga, "_device_cc", lambda: (12, 0)) + for align in (cudnn.diagonal_alignment.TOP_LEFT, cudnn.diagonal_alignment.BOTTOM_RIGHT): + g = _mk_graph() + q, k, v, dims, strides = _mk_qkv(g, d=128) + o, _ = g.sdpa( + name="s", + q=q, + k=k, + v=v, + attn_scale=0.1, + is_inference=True, + diagonal_band_right_bound=16, + diagonal_alignment=align, + ) + _finish_output(o, dims, strides) + assert _SM120 in _eligible(g), align + + +def test_sm120_probe_accepts_ragged_skv_without_padding_or_causal(monkeypatch): + # No KV-tail rule on the SM120 row (skv_tile=0): the kernel's first + # (masked) step always covers the rightmost — and therefore any partial — + # KV tile, so a dense unmasked graph with S_kv % 128 != 0 is served + # natively. The SM100 f16 row keeps rejecting this shape. + monkeypatch.setattr(ga, "_device_cc", lambda: (12, 0)) + g = _mk_graph() + s_kv, d = 300, 128 + q = g.tensor(dim=(B, H, S, d), stride=(S * H * d, d, H * d, 1), data_type=DTYPE, name="q") + k = g.tensor(dim=(B, H, s_kv, d), stride=(s_kv * H * d, d, H * d, 1), data_type=DTYPE, name="k") + v = g.tensor(dim=(B, H, s_kv, d), stride=(s_kv * H * d, d, H * d, 1), data_type=DTYPE, name="v") + o, _ = g.sdpa(name="s", q=q, k=k, v=v, attn_scale=0.1, is_inference=True) + _finish_output(o, (B, H, S, d), (S * H * d, d, H * d, 1)) + assert _SM120 in _eligible(g) def test_sm120_probe_accepts_mixed_head_dims(monkeypatch):