diff --git a/docs/fe-oss-apis/dsa.md b/docs/fe-oss-apis/dsa.md index 68a86a727..3efa871cc 100644 --- a/docs/fe-oss-apis/dsa.md +++ b/docs/fe-oss-apis/dsa.md @@ -78,6 +78,9 @@ DSA.sparse_attention_backward_wrapper DSA.IndexerForward DSA.indexer_forward_wrapper +DSA.IndexerForwardLean +DSA.indexer_forward_lean_wrapper + DSA.IndexerTopK DSA.indexer_top_k_wrapper @@ -162,6 +165,70 @@ result = DSA.indexer_forward_wrapper( scores = result["scores"] ``` +#### SM100 lean fast path (`IndexerForwardLean`) + +On SM100, `indexer_forward_wrapper` transparently dispatches eligible +configurations to a specialized persistent kernel +(`DSA.IndexerForwardLean` / `DSA.indexer_forward_lean_wrapper`) that is +substantially faster than the generic SM100 kernel for the large-`S_q` +H=64 regime. The lean path is **additive**: any configuration outside its +support gate keeps using the existing kernel unchanged, and setting the +environment variable `CUDNNFE_DSA_INDEXER_FWD_DISABLE_LEAN` (to any +non-empty value) forces the legacy kernel for every transparently +dispatched configuration. Calls made explicitly through +`indexer_forward_lean_wrapper` do not consult the variable; in particular +its THD mode has no legacy fallback (THD lean scores use a different +output layout — see below), so opting a THD call site out of the lean +path means calling `indexer_forward_wrapper` instead. + +The lean gate (`IndexerForwardLean.check_support()`) requires all of: + +- `head_dim == 128`, `qhead_per_kv_head == 64`, `H_kv == 1` +- uniform-length batched BSHD (any `B`) on the transparent wrapper. + THD/varlen (`cu_seqlens_*`) is served by the same lean kernel, but only + through the explicit `indexer_forward_lean_wrapper`: per-row absolute + compressed-KV windows carry the segment isolation and the ratio-causal + mask, and the scores come back in a global-compressed-KV-column + `(T_q, m_total)` layout (each segment's finite scores sit in its own + absolute column block), which intentionally differs from the legacy + `(total_q, max_seqlen_k)` local-column layout — + `indexer_forward_wrapper` therefore keeps routing THD calls to the + legacy kernel (`q_causal_offsets` with THD also stays legacy) +- BF16 `q`/`k`; BF16 **or** FP32 `w` (both ingested directly — BF16 + weights are up-converted in-kernel, exactly); FP32 scores +- contiguous inputs with 16-byte-aligned base pointers (a TMA + requirement — storage-offset views that break the alignment fall back + to the legacy kernel, they are never copied); `S_q` a multiple of 4; + any `S_k >= 1` +- a saturated persistent grid: `S_q / 4 >= LEAN_MIN_WAVES * sm_count` + (see `api_lean.py`; below that q-tile count the static schedule loses + its load-balance advantage and the legacy kernel is used) +- default tuning parameters on the wrapper call (`m_block_size=128`, + `n_block_size=128`, `q_stage=2`, `kv_stage=4`) + +`sm_scale` and `q_causal_offsets` are fully supported (offsets are folded +into the per-row visibility windows on the host, in int64; offset windows +are rebuilt per call, while the offset-free windows are cached per shape). +One documented dispatch bound: the legacy kernel evaluates +`(q_causal_offset + i + 1)` in int32 on device, so offsets with +`offset + S_q + 1 > INT32_MAX` stay on the legacy path (keeping the +dispatched behavior identical to legacy for extreme-but-valid offsets); +enforcing the bound reads the small `(B,)` offsets tensor once per +offsets call. +The lean path compiles per `(S_q, S_k, W dtype, sm_scale variant, +device)` (by design — the schedule is static; the batch size and `ratio` +never affect the generated code, so one compiled kernel serves every `B` +of a shape), and both the compiled-kernel and dispatch-verdict caches are +bounded thread-safe LRUs. It is intended for the training-style regime +where shapes repeat. A configuration the support gate accepts is expected +to compile: a JIT failure raises `RuntimeError` (with the escape hatch +named) instead of silently falling back, so real lean-path bugs cannot +hide behind the legacy kernel. Scores returned by the lean path are +always contiguous, and match +the legacy kernel's values to FP32 accumulation-order rounding (both +paths use BF16 tensor-core products with FP32 accumulation; the lean +epilogue uses a fixed reduction order and is deterministic run-to-run). + ### 3. Indexer Top-K Radix top-K kernel for selecting candidate KV indices from indexer scores, diff --git a/docs/fe-oss-apis/overview.md b/docs/fe-oss-apis/overview.md index 2bc3ccfb2..b0fb5167e 100644 --- a/docs/fe-oss-apis/overview.md +++ b/docs/fe-oss-apis/overview.md @@ -21,6 +21,7 @@ This folder documents the Python FE APIs implemented under `python/cudnn`. For d - [Grouped GEMM + Quant (Unified)](gemm_fusions/grouped_gemm_quant_unified.md) - [Grouped GEMM + Wgrad](gemm_fusions/grouped_gemm_wgrad.md) - [Block Sparse Attention (BSA)](bsa.md) +- [DeepSeek Sparse Attention (DSA)](dsa.md) - [Native Sparse Attention (NSA)](nsa.md) - [RMSNorm + RHT + Amax](rmsnorm_rht_amax.md) - [SDPA Forward FE OSS API (SM100, D=256)](https://docs.nvidia.com/deeplearning/cudnn/frontend/latest/operations/Attention.html#sdpa-forward-fe-oss-sm100-d256) diff --git a/python/cudnn/deepseek_sparse_attention/__init__.py b/python/cudnn/deepseek_sparse_attention/__init__.py index 7f900037f..ba58f3ad3 100644 --- a/python/cudnn/deepseek_sparse_attention/__init__.py +++ b/python/cudnn/deepseek_sparse_attention/__init__.py @@ -5,6 +5,8 @@ "sparse_attention_backward_wrapper": (".sparse_attention_backward", "sparse_attention_backward_wrapper"), "IndexerForward": (".indexer_forward", "IndexerForward"), "indexer_forward_wrapper": (".indexer_forward", "indexer_forward_wrapper"), + "IndexerForwardLean": (".indexer_forward", "IndexerForwardLean"), + "indexer_forward_lean_wrapper": (".indexer_forward", "indexer_forward_lean_wrapper"), "IndexerTopK": (".indexer_top_k", "IndexerTopK"), "indexer_top_k_wrapper": (".indexer_top_k", "indexer_top_k_wrapper"), "local_to_global_wrapper": (".indexer_top_k", "local_to_global_wrapper"), diff --git a/python/cudnn/deepseek_sparse_attention/indexer_forward/__init__.py b/python/cudnn/deepseek_sparse_attention/indexer_forward/__init__.py index 4b8df4c14..a97c0736a 100644 --- a/python/cudnn/deepseek_sparse_attention/indexer_forward/__init__.py +++ b/python/cudnn/deepseek_sparse_attention/indexer_forward/__init__.py @@ -1,3 +1,9 @@ from .api import IndexerForward, indexer_forward_wrapper +from .api_lean import IndexerForwardLean, indexer_forward_lean_wrapper -__all__ = ["IndexerForward", "indexer_forward_wrapper"] +__all__ = [ + "IndexerForward", + "indexer_forward_wrapper", + "IndexerForwardLean", + "indexer_forward_lean_wrapper", +] diff --git a/python/cudnn/deepseek_sparse_attention/indexer_forward/api.py b/python/cudnn/deepseek_sparse_attention/indexer_forward/api.py index d35e0da8b..5eb326f2f 100644 --- a/python/cudnn/deepseek_sparse_attention/indexer_forward/api.py +++ b/python/cudnn/deepseek_sparse_attention/indexer_forward/api.py @@ -9,6 +9,7 @@ from __future__ import annotations import logging +import os from typing import Optional import torch @@ -26,9 +27,14 @@ from .indexer_fwd_sm100 import IndexerForwardSm100 from ._interface import indexer_fwd as indexer_fwd_sm100 from ._interface_sm90 import indexer_fwd as indexer_fwd_sm90 +from . import api_lean as _api_lean TMA_ALIGN_ELEMS = 4 # FP32 output => seqlen_k padded to multiples of 4 (16 B) +# escape hatch: set to any non-empty value to keep every configuration on +# the legacy kernel (the lean fast path is dispatched transparently below) +_DISABLE_LEAN_ENV = "CUDNNFE_DSA_INDEXER_FWD_DISABLE_LEAN" + class IndexerForward(APIBase): """SM100+ APIBase implementation used by ``indexer_forward_wrapper``. @@ -239,6 +245,13 @@ def indexer_forward_wrapper( positions outside the valid KV range with -inf. ``q_causal_offsets`` may specify the global uncompressed token index for each batch/THD segment's local q[0]. + + On SM100, uniform-length BSHD configurations inside the lean + specialization (``head_dim == 128``, ``qhead_per_kv_head == 64``, + ``h_kv == 1``, default tuning parameters, saturated grid) are served by + the lean fast-path kernel (:mod:`.api_lean`); set the + ``CUDNNFE_DSA_INDEXER_FWD_DISABLE_LEAN`` environment variable to force + the legacy kernel for every configuration. """ if device_major() == 9: unsupported = [] @@ -273,6 +286,68 @@ def indexer_forward_wrapper( ) return TupleDict(scores=scores) + # SM100 additive lean fast path: uniform-length BSHD configurations + # inside the lean specialization (H=64/D=128 MQA, BF16 Q/K, default + # tuning parameters, saturated persistent grid — see + # api_lean.IndexerForwardLean.check_support) are routed to the lean + # kernel transparently. Every other configuration — THD/varlen, H=32, + # non-default tuning knobs, SM90/SM110+, small grids, misaligned or + # non-contiguous tensors, ... — falls through to the legacy path + # below, which is unchanged. + if ( + cu_seqlens_q is None + and cu_seqlens_k is None + and m_block_size == 128 + and n_block_size == 128 + and q_stage == 2 + and kv_stage == 4 + and not os.getenv(_DISABLE_LEAN_ENV) + ): + # cheap metadata screen for q_causal_offsets BEFORE any lean work + # (in particular before the lean JIT): malformed offsets fall + # through so the legacy path raises its own errors (it validates + # with the same shared helper). + offsets_ok = q_causal_offsets is None or ( + q.ndim == 4 + and q_causal_offsets.dtype == torch.int32 + and q_causal_offsets.ndim == 1 + and q_causal_offsets.shape[0] == q.shape[0] + and q_causal_offsets.is_cuda + and q_causal_offsets.device == q.device + ) + lean_api = None + if offsets_ok: + try: + # _maybe_lean_api runs cheap non-throwing eligibility + # checks, then the cached check_support verdict + JIT. + # Only validation ValueErrors are caught (malformed + # cross-tensor combos: legacy reproduces its own error + # surface); a compile failure for an accepted config + # raises RuntimeError and deliberately escapes — a silent + # fallback would hide real lean-path bugs + # (CUDNNFE_DSA_INDEXER_FWD_DISABLE_LEAN is the escape + # hatch). + lean_api = _api_lean._maybe_lean_api(q, k, w, ratio, qhead_per_kv_head, sm_scale=sm_scale) + except ValueError: + lean_api = None + if lean_api is not None and not _api_lean._q_causal_offsets_within_int32(q_causal_offsets, q.shape[1]): + # legacy evaluates (offset + q_token + 1) in int32 on device; + # offsets large enough for that to overflow stay on the legacy + # path so dispatched behavior is unchanged (documented bound; + # costs one small D2H read, only on the offsets path). + lean_api = None + if lean_api is not None: + return _api_lean.indexer_forward_lean_wrapper( + q, + k, + w, + ratio=ratio, + qhead_per_kv_head=qhead_per_kv_head, + sm_scale=sm_scale, + q_causal_offsets=q_causal_offsets, + stream=stream, + ) + # BSHD and THD both go through indexer_fwd (it branches on cu_seqlens # internally): it owns output allocation + TMA padding and uses a single # shape-agnostic compile cache. cu_seqlens_q/k / max_seqlen_q/k are None for diff --git a/python/cudnn/deepseek_sparse_attention/indexer_forward/api_lean.py b/python/cudnn/deepseek_sparse_attention/indexer_forward/api_lean.py new file mode 100644 index 000000000..7015aca5f --- /dev/null +++ b/python/cudnn/deepseek_sparse_attention/indexer_forward/api_lean.py @@ -0,0 +1,1195 @@ +"""APIBase wrapper for the lean H=64/D=128 indexer-forward score kernel. + +``IndexerForwardLean`` is an additive fast path next to +:class:`cudnn.deepseek_sparse_attention.indexer_forward.IndexerForward`: +``check_support()`` returns ``True`` only for the configuration the lean +schedule is specialized for (``head_dim == 128``, ``qhead_per_kv_head == +64``, ``h_kv == 1``, BF16 Q/K with BF16/FP32 W and FP32 scores, and a +q-tile count that keeps the persistent grid saturated). Two input layouts +are served, both on the *same* compiled kernel (the schedule is a pure +per-row-window sweep — it never sees a batch or segment axis): + +* **Uniform BSHD** ``(B, S_q, H, D)``: one persistent launch per batch + entry, uniform ``ks == 0`` ratio-causal windows. This is the path + ``indexer_forward_wrapper`` dispatches to transparently. + +* **THD / varlen (ragged packed)** ``(T_q, H, D)`` + ``cu_seqlens_q/k``: + one persistent launch over the whole packed problem, with per-row + *absolute* compressed-KV windows ``ks = cu_seqlens_k[seg]``, + ``ke = ks + visible`` (see ``_thd_ratio_causal_windows``, the exact + integer mirror of gtp ``csa_host_index_math``). The windows carry all + segment isolation *and* ratio-causal masking, so no kernel change is + needed. THD scores use **global compressed-KV columns** + ``(T_q, m_total)`` (a query in segment ``b`` has finite scores only in + its own segment's column block ``[cu_seqlens_k[b], cu_seqlens_k[b+1])``); + this differs from the legacy kernel's local ``(total_q, max_seqlen_k)`` + layout, so THD is exposed only through the explicit + ``indexer_forward_lean_wrapper`` — ``indexer_forward_wrapper`` still + routes every THD call to the (local-column) legacy path, which is never + modified. +""" + +from __future__ import annotations + +import math +import threading +from collections import OrderedDict +from typing import Optional + +import torch +import cuda.bindings.driver as cuda + +import cutlass +import cutlass.cute as cute +from cutlass.cute.runtime import make_fake_stream + +from cudnn.api_base import APIBase, TupleDict + +from cudnn.deepseek_sparse_attention.utils.compiler import compile_options +from cudnn.deepseek_sparse_attention.utils.runtime import ( + resolve_stream, + torch_stream_context, + validate_q_causal_offsets, +) + +from .indexer_fwd_sm100_lean import IndexerForwardSm100Lean + +# Lean-path dispatch constants (documented, fixed for now): +# +# LEAN_MIN_WAVES: the lean schedule uses a static single-wave persistent +# grid of min(sm_count, num_m_tiles) CTAs with reversed-LPT tile order. +# Dispatch to lean requires +# num_m_tiles >= LEAN_MIN_WAVES * sm_count +# where num_m_tiles = s_q // LEAN_TILE_TOKENS is the number of persistent +# q-tiles (metadata-only; sm_count is queried from the target device). +# +# Calibration (B200, 148 SMs, nsys pure-kernel medians, lean vs legacy +# IndexerForwardSm100, BF16 H=64 D=128 MQA, 50 warmup / 100 reps): +# +# ratio 2, S_k = S_q/2: ratio 4, S_k = S_q/4: +# S_q waves speedup S_q waves speedup +# 256 0.43 1.27x 1024 1.73 1.40x +# 512 0.86 1.38x 4096 6.92 1.48x +# 768 1.30 1.45x 8192 13.84 1.47x +# 1024 1.73 1.43x +# 1536 2.59 1.37x +# 2048 3.46 1.33x +# 3072 5.19 1.43x +# 4096 6.92 1.47x +# 5120 8.65 1.46x +# 6144 10.38 1.44x +# 8192 13.84 1.44x +# +# The lean schedule won at every measured tile count, including sub-wave +# grids, so the gate is not a measured perf boundary. It is kept at one +# full wave because (a) below one wave the persistent grid no longer +# saturates the GPU, which is the regime the schedule was designed and +# validated for, and (b) the lean path compiles per (S_q, S_k, W dtype, +# sm_scale variant, device) — the schedule is static — so auto-dispatching +# arbitrarily small dynamic shapes would trade a ~1-2 us per-call win for +# a JIT compile per novel shape. Callers who want lean below the gate can +# invoke indexer_forward_lean_wrapper explicitly after lowering +# LEAN_MIN_WAVES. +LEAN_MIN_WAVES = 1 +# query tokens per q-tile: the kernel packs TQ=4 tokens x 64 heads into +# its N=256 MMA tile. +LEAN_TILE_TOKENS = 4 +# KV rows per MMA tile in the lean kernel: multiples compile with a fully +# static K extent; ragged extents use a dynamic K dim + TMA zero-fill. +_LEAN_KV_TILE_ROWS = 128 +# TMA descriptors are built with assumed_align=16. Contiguity does NOT +# imply base-pointer alignment (storage-offset views), so the runtime +# data_ptr() of every kernel operand is checked against this. +_TMA_MIN_ALIGN_BYTES = 16 +# Legacy evaluates (q_causal_offset + q_token + 1) in int32 on device; +# the lean host path evaluates the same window in int64 and clamps. The +# two agree only while the legacy intermediate cannot overflow, so the +# dispatcher keeps offsets beyond this bound on the legacy path. +_INT32_MAX = 2**31 - 1 +# Bound + eviction policy follow python/cudnn/graph.py's graph_cache +# precedent (maxsize=256). +_LEAN_CACHE_MAXSIZE = 256 + + +class _LruDict: + """Bounded thread-safe LRU mapping. + + Mirrors the bound of ``graph.py``'s ``graph_cache`` (maxsize 256) with + true LRU eviction and a lock, since dispatch may be called from + multiple threads. + """ + + def __init__(self, maxsize: int = _LEAN_CACHE_MAXSIZE): + self._data: OrderedDict = OrderedDict() + self._lock = threading.Lock() + self._maxsize = maxsize + + def get(self, key, default=None): + with self._lock: + if key not in self._data: + return default + self._data.move_to_end(key) + return self._data[key] + + def put(self, key, value) -> None: + with self._lock: + self._data[key] = value + self._data.move_to_end(key) + while len(self._data) > self._maxsize: + self._data.popitem(last=False) + + +class IndexerForwardLean(APIBase): + """SM100 lean APIBase implementation used by ``indexer_forward_lean_wrapper``. + + Computes ``scores[b, i, j] = sm_scale * sum_h relu(q[b,i,h,:] . + k[b,j,0,:]) * w[b,i,h]`` for ``j`` inside the per-row visibility + window ``[ks[b, i], ke[b, i])`` (positions inside swept KV tiles but + outside the window are ``-inf``; the wrapper pre-fills the output so + tiles the kernel never sweeps also read ``-inf``). The wrapper builds + ratio-causal windows ``ke = clamp((q_causal_offsets[b] + i + 1) // + ratio, 0, S_k)`` matching the legacy mask semantics. + + The compiled kernel is batch-size independent (it runs on flattened + per-batch views); ``execute()`` reads the batch size from the runtime + tensors, so one compiled instance serves every ``B`` of the same + ``(S_q, S_k)``. + """ + + def __init__( + self, + sample_q: torch.Tensor, # (B, S_q, H_q, D) BF16 + sample_k: torch.Tensor, # (B, S_k, H_kv, D) BF16 + sample_w: torch.Tensor, # (B, S_q, H_q) BF16 or FP32 + sample_out: torch.Tensor, # (B, S_q, S_k) FP32 + ratio: int = 4, + qhead_per_kv_head: Optional[int] = None, + sm_scale: float = 1.0, + cu_seqlens_q: Optional[torch.Tensor] = None, + cu_seqlens_k: Optional[torch.Tensor] = None, + max_seqlen_q: Optional[int] = None, + max_seqlen_k: Optional[int] = None, + ): + super().__init__() + self._kernel = IndexerForwardSm100Lean + + # THD/varlen (ragged packed) mode is a purely additive branch: when + # ``cu_seqlens_q`` is provided the sample tensors are packed + # ``(T_q, H, D)`` / ``(m_total, H_kv, D)`` / ``(T_q, H)`` / + # ``(T_q, m_total)`` and the per-row absolute compressed-KV windows + # carry all segment isolation + ratio-causal masking. When it is + # None every code path below is byte-identical to the BSHD schedule + # this PR shipped. + self._thd = cu_seqlens_q is not None + self.cu_seqlens_q = cu_seqlens_q + self.cu_seqlens_k = cu_seqlens_k + self.max_seqlen_q = max_seqlen_q + self.max_seqlen_k = max_seqlen_k + + self.q_desc = self._make_tensor_desc(sample_q, name="sample_q") + self.k_desc = self._make_tensor_desc(sample_k, name="sample_k") + self.w_desc = self._make_tensor_desc(sample_w, name="sample_w") + self.o_desc = self._make_tensor_desc(sample_out, name="sample_out") + + self.ratio = int(ratio) + # sm_scale selects the compiled VARIANT only (whether the epilogue + # multiply exists at all); the value itself is a runtime kernel + # argument that callers pass to execute() per call. The cached + # instance is never mutated after compile. + self.sm_scale = float(sm_scale) + self._kernel_applies_sm_scale = False + self.qhead_per_kv_head = qhead_per_kv_head + + self.batch_size = None # sample batch size (execute() is B-agnostic) + self.s_q = None + self.s_k = None + self.h_q = None + self.h_kv = None + self.head_dim = None + self.sm_count = None + self.target_device: Optional[torch.device] = None + + def _unsupported(self, msg: str) -> bool: + """Soft-fail a lean-path gate: log and report unsupported.""" + self._logger.debug("IndexerForwardLean unsupported: %s", msg) + self._is_supported = False + return False + + def check_support(self) -> bool: + """Lean-path dispatch gate. + + Malformed inputs (rank/shape mismatches between Q/K/W/Out) raise + ``ValueError`` exactly like ``IndexerForward.check_support``; + configurations that are merely outside the lean specialization + return ``False`` so the caller can fall back to the legacy path. + + Descriptors carry no storage pointers, so base-pointer alignment + (TMA needs 16 B) is a runtime concern checked in ``execute()`` and + by the dispatcher, not here. Meta-device sample tensors are + accepted for metadata-only support checks; all real (CUDA) + descriptors must agree on one device, and compute capability plus + sm_count are queried from that resolved device (falling back to + the current device when every descriptor is meta). + """ + self._logger.debug("Entering check_support") + if self._thd: + return self._check_support_thd() + self._value_error_if( + self.q_desc.ndim != 4, + f"Q must be 4-D (B, S_q, H_q, D), got {self.q_desc.shape}", + ) + self._value_error_if( + self.k_desc.ndim != 4, + f"K must be 4-D (B, S_k, H_kv, D), got {self.k_desc.shape}", + ) + self._value_error_if( + self.w_desc.ndim != 3, + f"W must be 3-D (B, S_q, H_q), got {self.w_desc.shape}", + ) + self._value_error_if( + self.o_desc.ndim != 3, + f"Out must be 3-D (B, S_q, S_k), got {self.o_desc.shape}", + ) + + b, s_q, h_q, d = self.q_desc.shape + b_k, s_k, h_kv, d_k = self.k_desc.shape + b_o, s_q_out, s_k_out = self.o_desc.shape + self._value_error_if(b != b_k, f"Batch size mismatch Q={b} vs K={b_k}") + self._value_error_if(b != b_o, f"Batch size mismatch Q={b} vs Out={b_o}") + self._value_error_if(s_q != s_q_out, f"S_q mismatch Q={s_q} vs Out={s_q_out}") + self._value_error_if(d != d_k, f"Head dim mismatch Q={d} vs K={d_k}") + self._value_error_if( + self.w_desc.shape != (b, s_q, h_q), + f"W must have shape (B, S_q, H_q) = ({b}, {s_q}, {h_q}), got {self.w_desc.shape}", + ) + + qhpkv = self.qhead_per_kv_head if self.qhead_per_kv_head is not None else (h_q // h_kv) + self._value_error_if( + qhpkv * h_kv != h_q, + f"qhead_per_kv_head * h_kv != h_q ({qhpkv} * {h_kv} != {h_q})", + ) + self.qhead_per_kv_head = qhpkv + + # ---- lean-specialization gates (soft: False -> legacy fallback) ---- + if d != 128: + return self._unsupported(f"head_dim must be 128, got {d}") + if qhpkv != 64: + return self._unsupported(f"qhead_per_kv_head must be 64, got {qhpkv}") + if h_kv != 1: + return self._unsupported(f"h_kv must be 1, got {h_kv}") + if self.q_desc.dtype != torch.bfloat16 or self.k_desc.dtype != torch.bfloat16: + return self._unsupported(f"Q/K must be bfloat16, got {self.q_desc.dtype}/{self.k_desc.dtype}") + if self.w_desc.dtype not in (torch.bfloat16, torch.float32): + return self._unsupported(f"W must be bfloat16 or float32, got {self.w_desc.dtype}") + if self.o_desc.dtype != torch.float32: + return self._unsupported(f"Out must be float32, got {self.o_desc.dtype}") + if not math.isfinite(self.sm_scale): + return self._unsupported(f"sm_scale must be finite, got {self.sm_scale}") + if self.ratio < 1: + return self._unsupported(f"ratio must be >= 1, got {self.ratio}") + if s_q % LEAN_TILE_TOKENS != 0: + return self._unsupported(f"S_q must be a multiple of {LEAN_TILE_TOKENS}, got {s_q}") + if s_k < 1 or s_k_out != s_k: + return self._unsupported(f"Out column dim must equal S_k >= 1, got S_k={s_k} Out={s_k_out}") + for desc, name in ( + (self.q_desc, "Q"), + (self.k_desc, "K"), + (self.w_desc, "W"), + (self.o_desc, "Out"), + ): + if not desc.is_contiguous(): + return self._unsupported(f"{name} must be contiguous") + + # device resolution: all runtime (CUDA) descriptors on one device; + # meta descriptors are metadata-only stand-ins and pin nothing. + devices = {desc.device for desc in (self.q_desc, self.k_desc, self.w_desc, self.o_desc)} + if any(dev.type not in ("cuda", "meta") for dev in devices): + return self._unsupported(f"Q/K/W/Out must be CUDA tensors, got devices {sorted(str(dev) for dev in devices)}") + cuda_devices = {dev for dev in devices if dev.type == "cuda"} + if len(cuda_devices) > 1: + return self._unsupported(f"Q/K/W/Out must share one CUDA device, got {sorted(str(dev) for dev in cuda_devices)}") + if not torch.cuda.is_available(): + return self._unsupported("CUDA is not available") + if cuda_devices: + target = next(iter(cuda_devices)) + if target.index is None: + target = torch.device("cuda", torch.cuda.current_device()) + else: + target = torch.device("cuda", torch.cuda.current_device()) + + # full (major, minor) from the TARGET device (not the current one, + # and not the process-wide lru-cached device_major()). Every + # SM10.x minor exposes the tcgen05 UMMA/TMEM/LDTM features and the + # 227 KB CTA smem carveout this schedule needs; SM90 and SM110+ + # stay on the legacy path. + major, minor = torch.cuda.get_device_capability(target) + if major != 10: + return self._unsupported(f"lean schedule requires SM100-class compute capability, found SM{major}.{minor} on {target}") + + # occupancy gate: keep the static persistent grid saturated + # (metadata-only; sm_count comes from the target device runtime, + # not GPU memory) + num_m_tiles = s_q // LEAN_TILE_TOKENS + sm_count = torch.cuda.get_device_properties(target).multi_processor_count + if num_m_tiles < LEAN_MIN_WAVES * sm_count: + return self._unsupported( + f"num_m_tiles={num_m_tiles} < LEAN_MIN_WAVES*sm_count=" + f"{LEAN_MIN_WAVES}*{sm_count} — static reversed-LPT grid " + f"needs >= {LEAN_MIN_WAVES} q-tiles per SM" + ) + + self.batch_size = b + self.s_q = s_q + self.s_k = s_k + self.h_q = h_q + self.h_kv = h_kv + self.head_dim = d + self.sm_count = sm_count + self.target_device = target + self._is_supported = True + return True + + def _check_support_thd(self) -> bool: + """THD / varlen (ragged packed) dispatch gate. + + Additive to :meth:`check_support`; only reachable when + ``cu_seqlens_q`` was provided. Packed layout: + + Q (T_q, H_q, D) W (T_q, H_q) + K (m_total, H_kv, D) Out (T_q, m_total) + + Malformed inputs (rank/shape mismatches, missing/ill-typed + cu_seqlens, non-monotonic offsets, or a ``cu_seqlens[-1]`` that + disagrees with the packed extents) raise ``ValueError`` — the same + hard-failure surface ``IndexerForward``/``indexer_fwd`` produce for + malformed THD input. Configurations that are merely outside the lean + specialization return ``False`` (the explicit wrapper then raises a + fall-back-to-legacy error; there is no transparent THD dispatch). + """ + self._logger.debug("Entering _check_support_thd") + self._value_error_if( + self.q_desc.ndim != 3, + f"THD Q must be 3-D (T_q, H_q, D), got {self.q_desc.shape}", + ) + self._value_error_if( + self.k_desc.ndim != 3, + f"THD K must be 3-D (m_total, H_kv, D), got {self.k_desc.shape}", + ) + self._value_error_if( + self.w_desc.ndim != 2, + f"THD W must be 2-D (T_q, H_q), got {self.w_desc.shape}", + ) + self._value_error_if( + self.o_desc.ndim != 2, + f"THD Out must be 2-D (T_q, m_total), got {self.o_desc.shape}", + ) + + t_q, h_q, d = self.q_desc.shape + m_total, h_kv, d_k = self.k_desc.shape + t_q_o, m_total_o = self.o_desc.shape + self._value_error_if(d != d_k, f"Head dim mismatch Q={d} vs K={d_k}") + self._value_error_if( + self.w_desc.shape != (t_q, h_q), + f"THD W must have shape (T_q, H_q) = ({t_q}, {h_q}), got {self.w_desc.shape}", + ) + self._value_error_if( + t_q_o != t_q, + f"THD Out row dim must equal T_q ({t_q}), got {t_q_o}", + ) + self._value_error_if( + m_total_o != m_total, + f"THD Out column dim must equal m_total = K rows ({m_total}), got {m_total_o}", + ) + + # ---- cu_seqlens structural validation (hard: ValueError) ---- + cu_q, cu_k = self.cu_seqlens_q, self.cu_seqlens_k + self._value_error_if( + cu_k is None, + "THD input requires both cu_seqlens_q and cu_seqlens_k", + ) + for cu, name in ((cu_q, "cu_seqlens_q"), (cu_k, "cu_seqlens_k")): + self._value_error_if(cu.dtype != torch.int32, f"{name} must be int32, got {cu.dtype}") + self._value_error_if(cu.ndim != 1, f"{name} must be 1-D, got {cu.ndim}-D") + self._value_error_if(cu.stride(0) != 1, f"{name} must be contiguous") + # device checks mirror the bwd THD gate: reject CPU/foreign-device + # cu_seqlens HERE with a clear message instead of letting + # _execute_thd fail later on the internally-built ks/ke tensors. + self._value_error_if( + cu.device.type not in ("cuda", "meta"), + f"{name} must be a CUDA tensor, got device {cu.device}", + ) + self._value_error_if( + cu.device.type == "cuda" and self.q_desc.device.type == "cuda" and cu.device != self.q_desc.device, + f"{name} must be on the same CUDA device as Q ({self.q_desc.device}), got {cu.device}", + ) + self._value_error_if( + cu_q.shape[0] != cu_k.shape[0], + f"cu_seqlens_q ({cu_q.shape[0]}) and cu_seqlens_k ({cu_k.shape[0]}) must have the same length", + ) + self._value_error_if( + cu_q.shape[0] < 2, + f"cu_seqlens must have length batch+1 >= 2, got {cu_q.shape[0]}", + ) + n_seg = cu_q.shape[0] - 1 + + # value validation needs the offsets on host (small int32, one D2H + # per dispatch verdict — THD verdicts are not shape-cached, so a + # caller that mutates cu_seqlens contents is always re-validated). + cu_q_host = cu_q.tolist() if cu_q.device.type != "meta" else None + cu_k_host = cu_k.tolist() if cu_k.device.type != "meta" else None + if cu_q_host is not None: + self._value_error_if(cu_q_host[0] != 0, f"cu_seqlens_q must start at 0, got {cu_q_host[0]}") + self._value_error_if( + any(cu_q_host[i + 1] < cu_q_host[i] for i in range(n_seg)), + f"cu_seqlens_q must be non-decreasing, got {cu_q_host}", + ) + self._value_error_if( + cu_q_host[-1] != t_q, + f"cu_seqlens_q[-1] ({cu_q_host[-1]}) must equal packed T_q ({t_q})", + ) + if cu_k_host is not None: + self._value_error_if(cu_k_host[0] != 0, f"cu_seqlens_k must start at 0, got {cu_k_host[0]}") + self._value_error_if( + any(cu_k_host[i + 1] < cu_k_host[i] for i in range(n_seg)), + f"cu_seqlens_k must be non-decreasing, got {cu_k_host}", + ) + self._value_error_if( + cu_k_host[-1] != m_total, + f"cu_seqlens_k[-1] ({cu_k_host[-1]}) must equal packed K rows / m_total ({m_total})", + ) + + qhpkv = self.qhead_per_kv_head if self.qhead_per_kv_head is not None else (h_q // h_kv) + self._value_error_if( + qhpkv * h_kv != h_q, + f"qhead_per_kv_head * h_kv != h_q ({qhpkv} * {h_kv} != {h_q})", + ) + self.qhead_per_kv_head = qhpkv + + # ---- lean-specialization gates (soft: False -> caller falls back) ---- + if d != 128: + return self._unsupported(f"head_dim must be 128, got {d}") + if qhpkv != 64: + return self._unsupported(f"qhead_per_kv_head must be 64, got {qhpkv}") + if h_kv != 1: + return self._unsupported(f"h_kv must be 1, got {h_kv}") + if self.q_desc.dtype != torch.bfloat16 or self.k_desc.dtype != torch.bfloat16: + return self._unsupported(f"Q/K must be bfloat16, got {self.q_desc.dtype}/{self.k_desc.dtype}") + if self.w_desc.dtype not in (torch.bfloat16, torch.float32): + return self._unsupported(f"W must be bfloat16 or float32, got {self.w_desc.dtype}") + if self.o_desc.dtype != torch.float32: + return self._unsupported(f"Out must be float32, got {self.o_desc.dtype}") + if not math.isfinite(self.sm_scale): + return self._unsupported(f"sm_scale must be finite, got {self.sm_scale}") + if self.ratio < 1: + return self._unsupported(f"ratio must be >= 1, got {self.ratio}") + # The static reversed-LPT grid tiles the packed q axis in TQ-token + # tiles exactly as BSHD does; T_q (not any per-segment length) is the + # only q-axis divisibility the schedule needs. Segment boundaries + # that fall mid-tile are still correct (the per-row window masks each + # of the tile's 4 tokens to its own segment's KV block); they only + # widen that one tile's union KV sweep, which is bounded by the + # larger adjacent segment and negligible for realistic packings. + if t_q % LEAN_TILE_TOKENS != 0: + return self._unsupported(f"packed T_q must be a multiple of {LEAN_TILE_TOKENS}, got {t_q}") + if m_total < 1: + return self._unsupported(f"packed K rows / m_total must be >= 1, got {m_total}") + for desc, name in ( + (self.q_desc, "Q"), + (self.k_desc, "K"), + (self.w_desc, "W"), + (self.o_desc, "Out"), + ): + if not desc.is_contiguous(): + return self._unsupported(f"{name} must be contiguous") + + # device resolution mirrors the BSHD gate (cu_seqlens do not pin a + # device; the packed operands do). + devices = {desc.device for desc in (self.q_desc, self.k_desc, self.w_desc, self.o_desc)} + if any(dev.type not in ("cuda", "meta") for dev in devices): + return self._unsupported(f"Q/K/W/Out must be CUDA tensors, got devices {sorted(str(dev) for dev in devices)}") + cuda_devices = {dev for dev in devices if dev.type == "cuda"} + if len(cuda_devices) > 1: + return self._unsupported(f"Q/K/W/Out must share one CUDA device, got {sorted(str(dev) for dev in cuda_devices)}") + if not torch.cuda.is_available(): + return self._unsupported("CUDA is not available") + if cuda_devices: + target = next(iter(cuda_devices)) + if target.index is None: + target = torch.device("cuda", torch.cuda.current_device()) + else: + target = torch.device("cuda", torch.cuda.current_device()) + + major, minor = torch.cuda.get_device_capability(target) + if major != 10: + return self._unsupported(f"lean schedule requires SM100-class compute capability, found SM{major}.{minor} on {target}") + + num_m_tiles = t_q // LEAN_TILE_TOKENS + sm_count = torch.cuda.get_device_properties(target).multi_processor_count + if num_m_tiles < LEAN_MIN_WAVES * sm_count: + return self._unsupported( + f"num_m_tiles={num_m_tiles} < LEAN_MIN_WAVES*sm_count=" + f"{LEAN_MIN_WAVES}*{sm_count} — static reversed-LPT grid " + f"needs >= {LEAN_MIN_WAVES} q-tiles per SM" + ) + + # THD folds the whole packed problem into ONE launch (S_q = T_q, + # S_k = m_total); the compiled kernel is the shared BSHD schedule. + self.batch_size = n_seg + self.s_q = t_q + self.s_k = m_total + self.h_q = h_q + self.h_kv = h_kv + self.head_dim = d + self.sm_count = sm_count + self.target_device = target + self._is_supported = True + return True + + def compile(self) -> None: + self._logger.debug("Entering compile") + self._ensure_support_checked() + if self._compiled_kernel is not None: + return + + # the sm_scale == 1.0 variant compiles the epilogue multiply out + # entirely by passing sm_scale=None (an absent optional folds out of + # the kernel parameter layout at trace time, keeping the production + # instruction stream identical to the scale-free schedule; + # _maybe_lean_api keys the cache on this variant choice) + self._kernel_applies_sm_scale = self.sm_scale != 1.0 + + s_q, s_k, h, d = self.s_q, self.s_k, self.h_q, self.head_dim + # generated code depends only on this key: shapes enter through the + # fake tensors below (flattened per-batch views — B never appears), + # the sm_scale variant through the optional scalar, and the device + # through sm_count/toolchain target. ratio and B are runtime-side. + compile_key = (self.target_device.index, s_q, s_k, self.w_desc.dtype, self._kernel_applies_sm_scale) + cached = _lean_compile_cache.get(compile_key) + if cached is not None: + self._compiled_kernel = cached + self._logger.debug("Kernel fetched from compile cache") + return + + kernel_obj = self._kernel( + num_heads=self.qhead_per_kv_head, + head_dim=self.head_dim, + sm_count=self.sm_count, + ) + + # flattened single-batch views the kernel operates on (B > 1 loops + # per-batch launches over the same compiled kernel) + fake_q = self._make_fake_cute_tensor(torch.bfloat16, (s_q * h, d), (d, 1), assumed_align=16) + if s_k % _LEAN_KV_TILE_ROWS == 0: + # tile-aligned K: fully static extent (best address codegen) + fake_k = self._make_fake_cute_tensor(torch.bfloat16, (s_k, d), (d, 1), assumed_align=16) + else: + # ragged K: the row count is a dynamic dim, so the TMA + # descriptor carries the true runtime extent and the trailing + # partial 128-row KV tile is zero-filled by the TMA hardware + # instead of tripping the static tile-divisibility check + fake_k = self._make_fake_cute_compact_tensor( + torch.bfloat16, + (s_k, d), + stride_order=(1, 0), + assumed_align=16, + dynamic_mode=0, + divisibility=1, + ) + fake_w = self._make_fake_cute_tensor(self.w_desc.dtype, (s_q, h), (h, 1), assumed_align=16) + fake_ks = self._make_fake_cute_tensor(torch.int32, (s_q,), (1,), assumed_align=16) + fake_ke = self._make_fake_cute_tensor(torch.int32, (s_q,), (1,), assumed_align=16) + fake_out = self._make_fake_cute_tensor(torch.float32, (s_q, s_k), (s_k, 1), assumed_align=16) + + fake_stream = make_fake_stream(use_tvm_ffi_env_stream=False) + + # Compile-failure policy: a config check_support() accepted MUST + # compile; if it does not, that is a lean-path bug and it raises + # loudly (a silent legacy fallback here would hide real compiler + # or schedule bugs). CUDNNFE_DSA_INDEXER_FWD_DISABLE_LEAN is the + # escape hatch while such a bug is being fixed. + try: + _compiled_kernel = cute.compile( + kernel_obj, + fake_q, + fake_k, + fake_w, + fake_ks, + fake_ke, + fake_out, + cutlass.Float32(self.sm_scale) if self._kernel_applies_sm_scale else None, + fake_stream, + options=compile_options(), + ) + except Exception as exc: + raise RuntimeError( + f"IndexerForwardLean failed to compile a configuration check_support() accepted " + f"(S_q={s_q}, S_k={s_k}, W dtype={self.w_desc.dtype}, scaled={self._kernel_applies_sm_scale}, " + f"device={self.target_device}). This is a lean fast-path bug. For transparently " + f"dispatched BSHD calls, set CUDNNFE_DSA_INDEXER_FWD_DISABLE_LEAN=1 to force the legacy " + f"kernel while it is investigated; explicit indexer_forward_lean_wrapper THD calls have " + f"no legacy fallback (the global-column THD score layout differs from the legacy " + f"local-column output) — switch such call sites to indexer_forward_wrapper." + ) from exc + + def tensor_api(q_flat, k, w, ks, ke, out, sm_scale, stream): + # The kernel sweeps whole 128-column KV tiles of each q-tile's + # union visibility window: in-window positions get scores, + # out-of-window positions inside swept tiles get -inf, and + # never-swept tiles are left untouched. Callers that depend on + # -inf there must pre-fill the output (the wrapper does). + return _compiled_kernel(q_flat, k, w, ks, ke, out, sm_scale, stream) + + _lean_compile_cache.put(compile_key, tensor_api) + self._compiled_kernel = tensor_api + self._logger.debug("Kernel compiled successfully") + + def execute( + self, + q: torch.Tensor, # (B, S_q, H_q, D) BF16, contiguous + k: torch.Tensor, # (B, S_k, 1, D) BF16, contiguous + w: torch.Tensor, # (B, S_q, H_q) BF16/FP32, contiguous + ks: torch.Tensor, # (B, S_q) INT32 per-row window start (rows contiguous) + ke: torch.Tensor, # (B, S_q) INT32 per-row window end (exclusive) + out: torch.Tensor, # (B, S_q, S_k) FP32, pre-filled with -inf + sm_scale: Optional[float] = None, + current_stream: Optional[cuda.CUstream] = None, + ) -> None: + """Run the compiled kernel, one persistent launch per batch entry. + + ``sm_scale`` is a per-call runtime argument (``None`` uses the + construction-time default); the cached instance itself is never + mutated here, so concurrent same-shape callers with different + scales do not race. ``B`` is read from the runtime tensors — the + compiled kernel operates on flattened per-batch views and is + batch-size independent. + """ + self._logger.debug("Entering execute") + current_stream = resolve_stream(current_stream) + if self._compiled_kernel is None: + raise ValueError("IndexerForwardLean kernel not compiled") + scale_value = float(self.sm_scale if sm_scale is None else sm_scale) + if not math.isfinite(scale_value): + raise ValueError(f"sm_scale must be finite, got {scale_value}") + if scale_value != 1.0 and not self._kernel_applies_sm_scale: + raise ValueError( + "IndexerForwardLean was compiled as the sm_scale == 1.0 variant " + f"(the epilogue multiply is compiled out); got sm_scale={scale_value}. " + "Build a separate instance for scaled execution." + ) + if self._thd: + self._execute_thd(q, k, w, ks, ke, out, scale_value, current_stream) + return + b = int(q.shape[0]) if q.ndim == 4 else 0 + s_q, s_k, h_q, d = self.s_q, self.s_k, self.h_q, self.head_dim + for t, name, shape, dtype in ( + (q, "q", (b, s_q, h_q, d), torch.bfloat16), + (k, "k", (b, s_k, self.h_kv, d), torch.bfloat16), + (w, "w", (b, s_q, h_q), self.w_desc.dtype), + (ks, "ks", (b, s_q), torch.int32), + (ke, "ke", (b, s_q), torch.int32), + (out, "out", (b, s_q, s_k), torch.float32), + ): + if tuple(t.shape) != shape: + raise ValueError(f"{name} must have shape {shape}, got {tuple(t.shape)}") + if t.dtype != dtype: + raise ValueError(f"{name} must have dtype {dtype}, got {t.dtype}") + if not t.is_cuda or t.device != q.device: + raise ValueError(f"{name} must be a CUDA tensor on {q.device}, got {t.device}") + # ks/ke may be batch-expanded views (stride 0 on B); only their + # per-batch rows must be dense. Everything else is viewed flat. + if name in ("ks", "ke"): + if t.stride(-1) != 1: + raise ValueError(f"{name} rows must be contiguous, got stride {tuple(t.stride())}") + elif not t.is_contiguous(): + raise ValueError(f"{name} must be contiguous, got stride {tuple(t.stride())}") + if t.data_ptr() % _TMA_MIN_ALIGN_BYTES: + raise ValueError(f"{name} base pointer must be {_TMA_MIN_ALIGN_BYTES}-byte aligned for TMA, got 0x{t.data_ptr():x}") + sm_scale_arg = cutlass.Float32(scale_value) if self._kernel_applies_sm_scale else None + # one single-wave persistent launch per batch; K/windows/output are + # per-batch, so the batch dim folds into sequential launches on the + # same stream (uniform-length BSHD only — the gate rejects ragged) + for bi in range(b): + self._compiled_kernel( + q[bi].view(s_q * h_q, d), + k[bi].view(s_k, d), + w[bi].view(s_q, h_q), + ks[bi], + ke[bi], + out[bi].view(s_q, s_k), + sm_scale_arg, + current_stream, + ) + + def _execute_thd( + self, + q: torch.Tensor, # (T_q, H_q, D) BF16, contiguous + k: torch.Tensor, # (m_total, 1, D) BF16, contiguous + w: torch.Tensor, # (T_q, H_q) BF16/FP32, contiguous + ks: torch.Tensor, # (T_q,) INT32 absolute compressed-KV window start + ke: torch.Tensor, # (T_q,) INT32 absolute compressed-KV window end (excl.) + out: torch.Tensor, # (T_q, m_total) FP32, pre-filled with -inf + scale_value: float, + current_stream: Optional[cuda.CUstream], + ) -> None: + """Single persistent launch over the whole packed THD problem. + + ``S_q == T_q`` and ``S_k == m_total`` fold the batch away entirely: + the kernel is a per-row-window sweep, so the packed problem is just + one big uniform launch on the shared compiled schedule. Segment + structure lives only in the ``ks``/``ke`` absolute windows built by + ``_thd_ratio_causal_windows``. + """ + s_q, s_k, h_q, d = self.s_q, self.s_k, self.h_q, self.head_dim + for t, name, shape, dtype in ( + (q, "q", (s_q, h_q, d), torch.bfloat16), + (k, "k", (s_k, self.h_kv, d), torch.bfloat16), + (w, "w", (s_q, h_q), self.w_desc.dtype), + (ks, "ks", (s_q,), torch.int32), + (ke, "ke", (s_q,), torch.int32), + (out, "out", (s_q, s_k), torch.float32), + ): + if tuple(t.shape) != shape: + raise ValueError(f"{name} must have shape {shape}, got {tuple(t.shape)}") + if t.dtype != dtype: + raise ValueError(f"{name} must have dtype {dtype}, got {t.dtype}") + if not t.is_cuda or t.device != q.device: + raise ValueError(f"{name} must be a CUDA tensor on {q.device}, got {t.device}") + if not t.is_contiguous(): + raise ValueError(f"{name} must be contiguous, got stride {tuple(t.stride())}") + if t.data_ptr() % _TMA_MIN_ALIGN_BYTES: + raise ValueError(f"{name} base pointer must be {_TMA_MIN_ALIGN_BYTES}-byte aligned for TMA, got 0x{t.data_ptr():x}") + sm_scale_arg = cutlass.Float32(scale_value) if self._kernel_applies_sm_scale else None + self._compiled_kernel( + q.view(s_q * h_q, d), + k.view(s_k, d), + w.view(s_q, h_q), + ks, + ke, + out.view(s_q, s_k), + sm_scale_arg, + current_stream, + ) + + +# module-level bounded LRU caches (thread-safe, maxsize per graph.py's +# graph_cache precedent): +# _dispatch_cache: verdict key -> compiled IndexerForwardLean | None +# _lean_compile_cache: codegen key -> compiled kernel callable (shared +# across instances; B and ratio never affect codegen) +# _window_cache: (S_q, S_k, ratio, device) -> single-row ks/ke + event +_dispatch_cache = _LruDict() +_lean_compile_cache = _LruDict() +_window_cache = _LruDict() +_MISS = object() +# serializes verdict construction + JIT so concurrent same-config callers +# cannot compile the same kernel twice +_lean_build_lock = threading.Lock() + + +def _sample_out_like(b: int, s_q: int, s_k: int) -> torch.Tensor: + """Metadata-only stand-in for the (B, S_q, S_k) FP32 output. + + ``check_support`` only consumes shape/stride/dtype, so describe the + output on the meta device instead of allocating the real buffer (which + can be GBs) just to decide dispatch. + """ + return torch.empty(b, s_q, s_k, dtype=torch.float32, device="meta") + + +def _q_causal_offsets_within_int32(q_causal_offsets: Optional[torch.Tensor], s_q: int) -> bool: + """True when the legacy kernel's int32 window arithmetic cannot overflow. + + Legacy computes ``(q_causal_offset + q_token + 1) // ratio`` in Int32 on + device (``indexer_fwd_sm100.py``), which wraps for offsets near + INT32_MAX; the lean host path computes the same windows in int64 and + clamps, which does not. Additive discipline: extreme-but-valid offsets + where ``offset + S_q + 1 > INT32_MAX`` stay on the legacy path so + dispatched behavior is unchanged. Negative offsets cannot underflow + (``offset + i + 1 >= INT32_MIN + 1``) and clamp identically on both + paths. + + NOTE: this reads the small ``(B,)`` offsets tensor (one device sync), + and only runs on the offsets path for otherwise lean-eligible configs. + """ + if q_causal_offsets is None or q_causal_offsets.numel() == 0: + return True + off_max = int(q_causal_offsets.max().item()) + return off_max + s_q + 1 <= _INT32_MAX + + +def _maybe_lean_api( + q: torch.Tensor, + k: torch.Tensor, + w: torch.Tensor, + ratio: int, + qhead_per_kv_head: Optional[int], + sm_scale: float = 1.0, +) -> Optional[IndexerForwardLean]: + """Return a compiled ``IndexerForwardLean`` when the config is inside the + lean specialization, else ``None`` (dispatching callers fall back to + legacy). + + Structure (dispatch exception boundary): + + - Cheap NON-THROWING per-call guards run first and are never cached: + rank, same-CUDA-device, dtype, batch/W-shape consistency, contiguity, + base-pointer TMA alignment (a per-pointer property — storage-offset + views defeat metadata caching), finite ``sm_scale`` (a non-finite + scale must neither poison nor hit the cached finite-variant + verdicts), and ``ratio >= 1``. Anything they reject goes to legacy, + which produces its own error surface for malformed inputs. + - The cached verdict is keyed ONLY by metadata that determines the + ``check_support`` verdict and the generated code. ``B``, ``ratio`` + and the exact ``sm_scale`` value are excluded: the kernel compiles + on flattened per-batch views (B-independent — ``execute()`` reads B + from the runtime tensors), ratio only shapes the host-built windows, + and only the ``sm_scale == 1.0`` variant choice reaches codegen. + - ``check_support`` may still raise ``ValueError`` for cross-tensor + mismatches (mirroring ``IndexerForward.check_support``); dispatching + callers catch exactly that and let the legacy path raise its own + errors. ``compile()`` failures for accepted configs raise + ``RuntimeError`` and are deliberately NOT caught (see compile()). + """ + # ---- cheap, non-throwing per-call eligibility (False -> legacy) ---- + if q.ndim != 4 or k.ndim != 4 or w.ndim != 3: + return None + if not (q.is_cuda and k.is_cuda and w.is_cuda) or not (q.device == k.device == w.device): + return None + if q.dtype != torch.bfloat16 or k.dtype != torch.bfloat16 or w.dtype not in (torch.bfloat16, torch.float32): + return None + b, s_q, h_q, d = q.shape + if k.shape[0] != b or tuple(w.shape) != (b, s_q, h_q): + return None + if not (q.is_contiguous() and k.is_contiguous() and w.is_contiguous()): + return None + if (q.data_ptr() % _TMA_MIN_ALIGN_BYTES) or (k.data_ptr() % _TMA_MIN_ALIGN_BYTES) or (w.data_ptr() % _TMA_MIN_ALIGN_BYTES): + return None + if not math.isfinite(sm_scale): + return None + if int(ratio) < 1: + return None + + s_k, h_kv = k.shape[1], k.shape[2] + key = ( + s_q, + s_k, + h_q, + h_kv, + d, + w.dtype, + qhead_per_kv_head, + float(sm_scale) != 1.0, + q.device.index, + ) + hit = _dispatch_cache.get(key, _MISS) + if hit is not _MISS: + return hit + + with _lean_build_lock: + hit = _dispatch_cache.get(key, _MISS) + if hit is not _MISS: + return hit + api = IndexerForwardLean( + sample_q=q, + sample_k=k, + sample_w=w, + sample_out=_sample_out_like(b, s_q, s_k), + ratio=ratio, + qhead_per_kv_head=qhead_per_kv_head, + sm_scale=sm_scale, + ) + if api.check_support(): + api.compile() + else: + api = None + _dispatch_cache.put(key, api) + return api + + +def _ratio_causal_windows( + batch: int, + s_q: int, + s_k: int, + ratio: int, + device, + q_causal_offsets: Optional[torch.Tensor] = None, + stream: Optional[cuda.CUstream] = None, +) -> tuple: + """(B, S_q) int32 window tensors for the legacy ratio-causal mask: + ``ks = 0``, ``ke[b, i] = clamp((q_causal_offsets[b] + i + 1) // ratio, + 0, S_k)`` (offsets are 0 when omitted; int64 host math, clamped before + the int32 cast). + + The offset-free windows are cached per ``(S_q, S_k, ratio, device)`` + as single-row tensors and batch-expanded per call (metadata-only). + They are produced asynchronously on the first caller's stream, so a + CUDA event recorded right after production is waited on by every + consuming stream before reuse — the cheapest correct option (one + ``wait_event`` per hit) versus per-stream rebuilds or a producer-side + sync. Offset windows are rebuilt each call because the offsets tensor + contents may change between calls with identical metadata. + """ + if q_causal_offsets is None: + key = (s_q, s_k, int(ratio), device.index) + hit = _window_cache.get(key) + if hit is None: + with torch_stream_context(stream): + ks_row = torch.zeros(1, s_q, dtype=torch.int32, device=device) + ke_row = (((torch.arange(s_q, device=device) + 1) // int(ratio)).clamp_(0, s_k)).to(torch.int32).view(1, s_q) + ready = torch.cuda.Event() + ready.record(torch.cuda.current_stream(device)) + hit = (ks_row, ke_row, ready) + _window_cache.put(key, hit) + else: + ks_row, ke_row, ready = hit + with torch_stream_context(stream): + torch.cuda.current_stream(device).wait_event(ready) + return ks_row.expand(batch, s_q), ke_row.expand(batch, s_q) + with torch_stream_context(stream): + ks = torch.zeros(1, s_q, dtype=torch.int32, device=device).expand(batch, s_q) + pos = torch.arange(1, s_q + 1, device=device, dtype=torch.int64).view(1, s_q) + ke = ((q_causal_offsets.view(batch, 1).to(torch.int64) + pos) // int(ratio)).clamp_(0, s_k).to(torch.int32) + return ks, ke + + +def _thd_ratio_causal_windows( + cu_seqlens_q: torch.Tensor, + cu_seqlens_k: torch.Tensor, + total_q: int, + ratio: int, + stream: Optional[cuda.CUstream] = None, +) -> tuple: + """(T_q,) int32 ABSOLUTE compressed-KV windows for a ragged packed batch. + + Exact integer mirror of gtp ``csa_host_index_math`` (tcp_rank = 0). For + the packed query row ``i`` in segment ``sid``: + + sid = searchsorted(cu_seqlens_q[1:], i, right=True) clamped [0, S-1] + off = i - cu_seqlens_q[sid] # local position + m_seg = cu_seqlens_k[sid+1] - cu_seqlens_k[sid] # compressed KV len + vis = min((off + 1) // ratio, m_seg) # ratio-causal visible + ks[i] = cu_seqlens_k[sid] # absolute segment base + ke[i] = ks[i] + vis # absolute end (excl.) + + ``ks``/``ke`` are absolute columns into the packed compressed-KV buffer, + so they encode segment isolation AND the ratio causal mask in one shot: + a row in segment ``b`` can only produce finite scores in columns + ``[cu_seqlens_k[b], cu_seqlens_k[b+1])`` and only up to its visible + prefix. ``total_q`` is passed in (== packed ``q.shape[0]``) so the arange + extent needs no device sync. + """ + device = cu_seqlens_q.device + with torch_stream_context(stream): + cu_q = cu_seqlens_q.to(torch.int64) + cu_k = cu_seqlens_k.to(torch.int64) + n_seg = cu_q.shape[0] - 1 + i = torch.arange(total_q, device=device, dtype=torch.int64) + sid = torch.searchsorted(cu_q[1:], i, right=True).clamp_(max=n_seg - 1) + off = i - cu_q[sid] + m_seg = cu_k[sid + 1] - cu_k[sid] + vis = torch.minimum((off + 1) // int(ratio), m_seg) + ks = cu_k[sid].to(torch.int32).contiguous() + ke = (cu_k[sid] + vis).to(torch.int32).contiguous() + return ks, ke + + +def _maybe_lean_api_thd( + q: torch.Tensor, + k: torch.Tensor, + w: torch.Tensor, + cu_seqlens_q: torch.Tensor, + cu_seqlens_k: torch.Tensor, + ratio: int, + qhead_per_kv_head: Optional[int], + sm_scale: float = 1.0, + max_seqlen_q: Optional[int] = None, + max_seqlen_k: Optional[int] = None, +) -> Optional[IndexerForwardLean]: + """Compiled ``IndexerForwardLean`` for a packed THD problem, or ``None``. + + Cheap non-throwing eligibility (packed rank, device, dtype, contiguity, + TMA alignment, finite scale, ratio) runs first; anything it rejects + returns ``None`` (the wrapper then raises its fall-back-to-legacy error). + THD verdicts are deliberately NOT metadata-cached: the accept/reject + decision depends on the cu_seqlens *contents* (monotonicity, extent + consistency), so ``check_support`` re-validates them every call. The + expensive JIT is still shared through ``_lean_compile_cache`` (the codegen + key is shape-only: T_q, m_total, W dtype, scale variant, device). + """ + if q.ndim != 3 or k.ndim != 3 or w.ndim != 2: + return None + if not (q.is_cuda and k.is_cuda and w.is_cuda) or not (q.device == k.device == w.device): + return None + if q.dtype != torch.bfloat16 or k.dtype != torch.bfloat16 or w.dtype not in (torch.bfloat16, torch.float32): + return None + t_q, h_q, d = q.shape + if tuple(w.shape) != (t_q, h_q): + return None + if not (q.is_contiguous() and k.is_contiguous() and w.is_contiguous()): + return None + if (q.data_ptr() % _TMA_MIN_ALIGN_BYTES) or (k.data_ptr() % _TMA_MIN_ALIGN_BYTES) or (w.data_ptr() % _TMA_MIN_ALIGN_BYTES): + return None + if not math.isfinite(sm_scale): + return None + if int(ratio) < 1: + return None + + m_total = int(k.shape[0]) + with _lean_build_lock: + api = IndexerForwardLean( + sample_q=q, + sample_k=k, + sample_w=w, + sample_out=torch.empty(t_q, m_total, dtype=torch.float32, device="meta"), + ratio=ratio, + qhead_per_kv_head=qhead_per_kv_head, + sm_scale=sm_scale, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + max_seqlen_q=max_seqlen_q, + max_seqlen_k=max_seqlen_k, + ) + if api.check_support(): + api.compile() + return api + return None + + +def _indexer_forward_lean_thd( + q: torch.Tensor, + k: torch.Tensor, + w: torch.Tensor, + cu_seqlens_q: Optional[torch.Tensor], + cu_seqlens_k: Optional[torch.Tensor], + ratio: int, + qhead_per_kv_head: Optional[int], + sm_scale: float, + max_seqlen_q: Optional[int], + max_seqlen_k: Optional[int], + q_causal_offsets: Optional[torch.Tensor], + stream: Optional[cuda.CUstream], +) -> TupleDict: + """Explicit THD (ragged packed) lean path — GLOBAL compressed-KV columns. + + ``q`` ``(T_q, H, D)`` / ``k`` ``(m_total, H_kv, D)`` / ``w`` ``(T_q, H)``; + returns ``{'scores': (T_q, m_total) FP32}`` where row ``i`` (in segment + ``b``) is ``-inf`` everywhere except its own segment's compressed-KV + column block, ratio-causally masked. ``q_causal_offsets`` is not folded + on this path yet (documented gate); pass such shapes to + ``indexer_forward_wrapper`` (legacy) instead. + """ + if q_causal_offsets is not None: + raise ValueError( + "indexer_forward_lean_wrapper: q_causal_offsets is not supported on the THD " "lean path; use indexer_forward_wrapper (legacy) for offset varlen" + ) + if cu_seqlens_q is None or cu_seqlens_k is None: + raise ValueError("THD input requires both cu_seqlens_q and cu_seqlens_k") + api = _maybe_lean_api_thd( + q, + k, + w, + cu_seqlens_q, + cu_seqlens_k, + ratio, + qhead_per_kv_head, + sm_scale=sm_scale, + max_seqlen_q=max_seqlen_q, + max_seqlen_k=max_seqlen_k, + ) + if api is None: + raise ValueError( + "indexer_forward_lean_wrapper: THD configuration is outside the lean specialization " + "(rank/dtype/contiguity/alignment/H=64-D=128-MQA/saturated-grid gate); " + "use indexer_forward_wrapper instead" + ) + total_q = int(q.shape[0]) + m_total = int(k.shape[0]) + current_stream = resolve_stream(stream) + ks, ke = _thd_ratio_causal_windows(cu_seqlens_q, cu_seqlens_k, total_q, ratio, stream) + with torch_stream_context(stream): + out = torch.full((total_q, m_total), float("-inf"), dtype=torch.float32, device=q.device) + with torch.cuda.nvtx.range("indexer_fwd_lean_thd_kernel"): + api.execute(q, k, w, ks, ke, out, sm_scale=sm_scale, current_stream=current_stream) + return TupleDict(scores=out) + + +def indexer_forward_lean_wrapper( + q: torch.Tensor, + k: torch.Tensor, + w: torch.Tensor, + ratio: int = 4, + qhead_per_kv_head: Optional[int] = None, + sm_scale: float = 1.0, + q_causal_offsets: Optional[torch.Tensor] = None, + stream: Optional[cuda.CUstream] = None, + cu_seqlens_q: Optional[torch.Tensor] = None, + cu_seqlens_k: Optional[torch.Tensor] = None, + max_seqlen_q: Optional[int] = None, + max_seqlen_k: Optional[int] = None, +) -> TupleDict: + """High-level wrapper for the lean H=64/D=128 path (BSHD or THD). + + **BSHD** ``q (B, S_q, H, D)`` / ``k (B, S_k, H_kv, D)`` / ``w (B, S_q, H)`` + (``cu_seqlens_*`` omitted): allocates and ``-inf``-pre-fills the + ``(B, S_q, S_k)`` FP32 score buffer (always contiguous — the lean kernel + has no TMA-store padding constraint), builds the ratio causal windows + (folding ``q_causal_offsets`` when given, in int64 host math), and runs + the lean kernel once per batch. Returns ``{'scores': (B, S_q, S_k)}``; + positions outside the mask ``j < (q_causal_offsets[b] + i + 1) // ratio`` + are ``-inf``. + + **THD / varlen** ``q (T_q, H, D)`` / ``k (m_total, H_kv, D)`` / + ``w (T_q, H)`` with ``cu_seqlens_q`` and ``cu_seqlens_k`` (both int32, + ``batch+1``): one packed launch with per-row absolute compressed-KV + windows. Returns ``{'scores': (T_q, m_total)}`` in GLOBAL compressed-KV + columns — row ``i`` (segment ``b``) is ``-inf`` outside its own block + ``[cu_seqlens_k[b], cu_seqlens_k[b+1])`` and beyond its ratio-causal + visible prefix. This global-column THD layout differs from the legacy + ``(total_q, max_seqlen_k)`` local-column output, so ``q_causal_offsets`` + is unsupported here and ``indexer_forward_wrapper`` never auto-routes THD + to this path (it keeps every THD call on the legacy kernel). + + Raises ``ValueError`` if the config is outside the lean specialization + (including non-contiguous or non-16-byte-aligned inputs). All tensor + arguments must live on one CUDA device, and ``stream`` (when given) must + belong to that device — the same contract as the legacy wrapper + (CUstream carries no queryable device identity). + """ + if not math.isfinite(sm_scale): + raise ValueError(f"sm_scale must be finite, got {sm_scale}") + if cu_seqlens_q is not None or cu_seqlens_k is not None: + return _indexer_forward_lean_thd( + q, + k, + w, + cu_seqlens_q, + cu_seqlens_k, + ratio=ratio, + qhead_per_kv_head=qhead_per_kv_head, + sm_scale=sm_scale, + max_seqlen_q=max_seqlen_q, + max_seqlen_k=max_seqlen_k, + q_causal_offsets=q_causal_offsets, + stream=stream, + ) + api = _maybe_lean_api(q, k, w, ratio, qhead_per_kv_head, sm_scale=sm_scale) + if api is None: + raise ValueError( + "indexer_forward_lean_wrapper: configuration is outside the lean specialization " + "(shape/dtype/contiguity/alignment/device gate); use indexer_forward_wrapper instead" + ) + b, s_q = q.shape[0], q.shape[1] + s_k = k.shape[1] + + current_stream = resolve_stream(stream) + q_causal_offsets = validate_q_causal_offsets(q_causal_offsets, int(b), q.device, stream=current_stream) + ks, ke = _ratio_causal_windows(b, s_q, s_k, ratio, q.device, q_causal_offsets, stream) + with torch_stream_context(stream): + # -inf pre-fill: tiles the kernel never sweeps keep -inf (same + # contract as the legacy wrapper's pre-fill) + out = torch.full((b, s_q, s_k), float("-inf"), dtype=torch.float32, device=q.device) + with torch.cuda.nvtx.range("indexer_fwd_lean_kernel"): + # sm_scale is passed through as a per-call runtime argument; the + # cached api instance is immutable after compile + api.execute(q, k, w, ks, ke, out, sm_scale=sm_scale, current_stream=current_stream) + return TupleDict(scores=out) diff --git a/python/cudnn/deepseek_sparse_attention/indexer_forward/indexer_fwd_sm100_lean.py b/python/cudnn/deepseek_sparse_attention/indexer_forward/indexer_fwd_sm100_lean.py new file mode 100644 index 000000000..3f148f57f --- /dev/null +++ b/python/cudnn/deepseek_sparse_attention/indexer_forward/indexer_fwd_sm100_lean.py @@ -0,0 +1,627 @@ +"""Lean SM100 indexer-forward score kernel — H=64 / D=128 specialization. + +Persistent tcgen05 CuTe DSL kernel producing dense lightning-indexer scores + + out[i, j] = sm_scale * sum_h relu( dot(q[i, h, :], kv[j, :]) ) * w[i, h] + +for ``j`` in the per-row visibility window ``[ks[i], ke[i])`` (both clamped +to ``S_k``). ``sm_scale`` is a runtime scalar applied to the fp32 +head-reduced score (same placement as the legacy kernel: post head-reduce, +pre causal mask); the ``sm_scale == 1.0`` variant compiles the multiply +out entirely (``sm_scale=None`` — an absent optional folds out of the +parameter layout at trace time), keeping the production instruction +stream identical to the scale-free schedule. ``w`` may be BF16 +or FP32; BF16 weights are up-converted to FP32 in the staging copy (exact), +so the register math is identical for both ingest dtypes. Positions inside +a swept 128-column KV tile but outside the row's window are written +``-inf``; columns in tiles the kernel never sweeps are left untouched — +callers that rely on ``-inf`` there must pre-fill the output +(``indexer_forward_lean_wrapper`` does). ``S_k`` may be any positive size: +the KV TMA descriptor carries the true extent, so partial trailing tiles +are zero-filled by the TMA hardware and the fp32 stores are bounds-guarded +(``col < S_k``); when ``S_k`` is a multiple of the 128-row KV tile the +wrapper compiles with a fully static K extent instead. + +Schedule (the lean fast path for the ``qhead_per_kv_head == 64`` case): + + * swapAB UMMA: A = one dense 128-row KV tile (M = 128, TMA), B = the + tile's TQ=4 tokens x H=64 packed (token, head) query rows (N = 256, + TMA, loaded once per tile), K = head_dim 128; fp32 accumulation in + TMEM, 2 x 256-column slot ring (512 TMEM columns). + * Static reversed-LPT persistent grid: ``min(sm_count, num_tiles)`` CTAs; + CTA ``b`` handles linear tile ids ``b, b+G, b+2G, ...`` mapped in + reverse so block 0 takes the largest causal KV window (LPT balance on + the triangular work distribution). No dynamic tile scheduler. + * 12 warps / 384 threads: warp 0 = TMA load (Q double-buffered with + next-block prefetch, KV ``kv_stage``-deep), warp 1 = UMMA, warps 2-3 + idle; two epilogue warpgroups (warps 4-11) each own one fixed TMEM + slot and drain alternating KV tiles. + * Raw-mbarrier rolling split-LDTM drain: one full/empty mbarrier pair + + 1-bit phase per epilogue warpgroup; each 64-column token chunk is one + Ld32x32b x64 LDTM, and every LDTM after the first is issued behind the + previous chunk's FMA reduction; the TMEM slot is released after the + last fence, before the last reduction, so the UMMA warp refills it + while the epilogue finishes math out of registers. + * fp32 relu-weight head-sum in registers (packed f32x2 FMA, fixed + reduction order — deterministic run-to-run, no atomics). + +Numerics: bf16 tensor-core products with fp32 accumulation and an fp32 +relu/weight/head-sum epilogue. Register budget: 128 x 24 (load/MMA +warpgroup) + 256 x 216 (epilogue warpgroups) = 58368 <= 64K. Shared +memory: 2 x 64 KB sQ + 3 x 32 KB sK + 2 KB sW = 226 KB <= the 227 KB CTA +limit. + +Shape support is validated by ``IndexerForwardLean.check_support()`` in +``api_lean.py``; the asserts here are compile-time backstops only. +""" + +import cuda.bindings.driver as cuda + +from typing import Optional + +import cutlass +import cutlass.cute as cute +import cutlass.pipeline as pipeline +import cutlass.utils as utils +import cutlass.utils.blackwell_helpers as sm100_utils +from cutlass import Int32, const_expr +from cutlass.cute.nvgpu import cpasync, tcgen05 +from cutlass.cute.typing import BFloat16, Float32 + + +class IndexerForwardSm100Lean: + """Persistent swapAB dense-TMA score kernel, one 4-token tile per trip.""" + + def __init__(self, num_heads: int, head_dim: int, sm_count: int): + # check_support() gates dispatch; these are compile-time backstops. + assert num_heads == 64, "lean schedule is specialized for H=64" + assert head_dim == 128, "lean schedule is specialized for D=128" + assert sm_count > 0 + self.num_heads = num_heads + self.head_dim = head_dim + self.sm_count = sm_count + self.io_dtype = BFloat16 + + self.tq = 4 # query tokens per tile + self.n_block = 128 # KV rows per MMA tile (M) + self.n_cols = self.tq * num_heads # MMA N = 256 packed query rows + self.k_block = head_dim # single K chunk + self.kv_stage = 3 # KV TMA pipeline depth + self.q_stage = 2 # Q double-buffer w/ next-block prefetch + # (M, N, K) = (kv block, tq*heads, head_dim) + self.mma_tiler = (self.n_block, self.n_cols, self.k_block) + + self.load_warp_id = 0 + self.mma_warp_id = 1 + # 12 warps: WG0 = load + MMA + 2 idle; WG1/WG2 = epilogue. + self.epi_warp_ids = (4, 5, 6, 7, 8, 9, 10, 11) + self.num_warps = 12 + self.threads_per_cta = 32 * self.num_warps + + SM100_TMEM_COLS = 512 + self.num_tmem_slots = SM100_TMEM_COLS // self.n_cols # 2 + self.tmem_alloc_cols = SM100_TMEM_COLS + + # one Ld32x32b x64 LDTM per 64-column token chunk + self.epi_rep = 64 + + # register redistribution across 384 threads: + # 128 x 24 + 256 x 216 = 58368 <= 64K regs. + self.num_regs_wg0 = 24 + self.num_regs_epi = 216 + + # participants: mma warp + 8 epilogue warps + self.tmem_alloc_barrier = pipeline.NamedBarrier(barrier_id=2, num_threads=32 * (1 + len(self.epi_warp_ids))) + self.epi_sync_barrier = pipeline.NamedBarrier(barrier_id=3, num_threads=32 * len(self.epi_warp_ids)) + + # ----------------------------------------------------------------- + # host side + # ----------------------------------------------------------------- + @cute.jit + def __call__( + self, + mQ: cute.Tensor, # (S*H, D) bf16 (q.view(S*H, D), row-major) + mKV: cute.Tensor, # (SKV, D) bf16 + mW: cute.Tensor, # (S, H) fp32 or bf16 + mKS: cute.Tensor, # (S,) int32 + mKE: cute.Tensor, # (S,) int32 + mOut: cute.Tensor, # (S, SKV) fp32 + sm_scale: Optional[Float32], # None -> the multiply is compiled out + stream: cuda.CUstream, + ): + num_tiles = cute.size(mQ.shape[0]) // self.n_cols # S // TQ + mQ_v = cute.make_tensor( + mQ.iterator, + cute.make_layout( + (self.n_cols, self.head_dim, num_tiles), + stride=(self.head_dim, 1, self.n_cols * self.head_dim), + ), + ) + + cta_group = tcgen05.CtaGroup.ONE + tiled_mma = sm100_utils.make_trivial_tiled_mma( + self.io_dtype, + tcgen05.OperandMajorMode.K, + tcgen05.OperandMajorMode.K, + Float32, + cta_group, + self.mma_tiler[:2], + ) + # epilogue-only tiled mma (no MMA issued with it): the drain runs in + # per-token 64-column chunks ((128, 64) fragment/copy layouts). + tiled_mma_epi = sm100_utils.make_trivial_tiled_mma( + BFloat16, + tcgen05.OperandMajorMode.K, + tcgen05.OperandMajorMode.K, + Float32, + cta_group, + (self.n_block, self.num_heads), + ) + + sK_layout = sm100_utils.make_smem_layout_a(tiled_mma, self.mma_tiler, self.io_dtype, self.kv_stage) + sQ_layout = sm100_utils.make_smem_layout_b(tiled_mma, self.mma_tiler, self.io_dtype, self.q_stage) + + cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1, 1, 1)), (tiled_mma.thr_id.shape,)) + tma_load_op = cpasync.CopyBulkTensorTileG2SOp(cta_group) + tma_atom_K, tma_tensor_K = cute.nvgpu.make_tiled_tma_atom_A( + tma_load_op, + mKV, + cute.select(sK_layout, mode=[0, 1, 2]), + self.mma_tiler, + tiled_mma, + cluster_layout_vmnk.shape, + ) + tma_atom_Q, tma_tensor_Q = cute.nvgpu.make_tiled_tma_atom_B( + tma_load_op, + mQ_v, + cute.select(sQ_layout, mode=[0, 1, 2]), + self.mma_tiler, + tiled_mma, + cluster_layout_vmnk.shape, + ) + self.tma_copy_q_bytes = self.n_cols * self.head_dim * (self.io_dtype.width // 8) + self.tma_copy_k_bytes = self.n_block * self.k_block * (self.io_dtype.width // 8) + + num_ctas = cutlass.min(Int32(self.sm_count), num_tiles) + self.kernel( + tiled_mma, + tiled_mma_epi, + tma_atom_K, + tma_tensor_K, + tma_atom_Q, + tma_tensor_Q, + mW, + mKS, + mKE, + mOut, + sm_scale, + sQ_layout, + sK_layout, + ).launch( + grid=(num_ctas, 1, 1), + block=[self.threads_per_cta, 1, 1], + cluster=(1, 1, 1), + stream=stream, + ) + + # ----------------------------------------------------------------- + # device side + # ----------------------------------------------------------------- + @cute.kernel + def kernel( + self, + tiled_mma: cute.TiledMma, + tiled_mma_epi: cute.TiledMma, + tma_atom_K: cute.CopyAtom, + tma_tensor_K: cute.Tensor, + tma_atom_Q: cute.CopyAtom, + tma_tensor_Q: cute.Tensor, + mW: cute.Tensor, # (S, H) fp32 or bf16 + mKS: cute.Tensor, # (S,) int32 + mKE: cute.Tensor, # (S,) int32 + mOut: cute.Tensor, # (S, SKV) fp32 + sm_scale: Optional[Float32], + sQ_layout: cute.ComposedLayout, + sK_layout: cute.ComposedLayout, + ): + bidx, _, _ = cute.arch.block_idx() + gdim, _, _ = cute.arch.grid_dim() + tidx, _, _ = cute.arch.thread_idx() + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + + skv = Int32(cute.size(mOut.shape[1])) + # persistent schedule: CTA b handles linear ids b, b+G, b+2G, ... + # mapped in *reverse* (largest KV window first, LPT balance). + num_tiles = Int32(cute.size(mW.shape[0])) // self.tq + trips = cute.ceil_div(num_tiles - Int32(bidx), gdim) + + if warp_idx == self.load_warp_id: + cpasync.prefetch_descriptor(tma_atom_K) + cpasync.prefetch_descriptor(tma_atom_Q) + + @cute.struct + class SharedStorage: + Q_mbar_ptr: cute.struct.MemRange[cutlass.Int64, 2 * self.q_stage] + K_mbar_ptr: cute.struct.MemRange[cutlass.Int64, 2 * self.kv_stage] + S_mbar_ptr: cute.struct.MemRange[cutlass.Int64, 2 * self.num_tmem_slots] + tmem_holding_buf: cutlass.Int32 + sW: cute.struct.Align[cute.struct.MemRange[Float32, 2 * self.tq * self.num_heads], 16] + sQ: cute.struct.Align[cute.struct.MemRange[self.io_dtype, cute.cosize(sQ_layout)], 1024] + sK: cute.struct.Align[cute.struct.MemRange[self.io_dtype, cute.cosize(sK_layout)], 1024] + + smem = utils.SmemAllocator() + storage = smem.allocate(SharedStorage) + + # ---- pipelines ---- + pipe_Q = pipeline.PipelineTmaUmma.create( + barrier_storage=storage.Q_mbar_ptr.data_ptr(), + num_stages=self.q_stage, + producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), + consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), + tx_count=self.tma_copy_q_bytes, + defer_sync=True, + ) + pipe_K = pipeline.PipelineTmaUmma.create( + barrier_storage=storage.K_mbar_ptr.data_ptr(), + num_stages=self.kv_stage, + producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), + consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), + tx_count=self.tma_copy_k_bytes, + defer_sync=True, + ) + pipe_S = pipeline.PipelineUmmaAsync.create( + barrier_storage=storage.S_mbar_ptr.data_ptr(), + num_stages=self.num_tmem_slots, + producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), + consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 128), + defer_sync=True, + ) + + tmem = utils.TmemAllocator( + storage.tmem_holding_buf, + barrier_for_retrieve=self.tmem_alloc_barrier, + allocator_warp_id=self.epi_warp_ids[0], + ) + # raw handle to the pipe_S mbarrier ring for the epilogue's minimal + # consumer state machine (hoisted here: the SharedStorage python + # object cannot cross the dynamic warpgroup branch) + s_mbar_base = storage.S_mbar_ptr.data_ptr().align(min_align=8) + + pipeline.pipeline_init_arrive(is_relaxed=True) + + sQ = storage.sQ.get_tensor(sQ_layout.outer, swizzle=sQ_layout.inner) + sK = storage.sK.get_tensor(sK_layout.outer, swizzle=sK_layout.inner) + sW = storage.sW.get_tensor(cute.make_layout((2 * self.tq * self.num_heads,))) + + pipeline.pipeline_init_wait() + + # accumulator reference layouts + thr_mma = tiled_mma.get_slice(0) + acc_shape = tiled_mma.partition_shape_C(cute.select(self.mma_tiler, mode=[0, 1])) + acc_fake = tiled_mma.make_fragment_C(cute.append(acc_shape, self.num_tmem_slots)) + # each slot drains in TQ 64-column token chunks + epi_chunks = self.num_tmem_slots * self.tq + acc_shape_epi = tiled_mma_epi.partition_shape_C((self.n_block, self.num_heads)) + acc_fake_epi = tiled_mma_epi.make_fragment_C(cute.append(acc_shape_epi, epi_chunks)) + + wg_idx = tidx // 128 + + # ============================================================= + # Warpgroup 0: load warp + MMA warp (+ 2 idle warps) + # ============================================================= + if wg_idx == 0: + cute.arch.setmaxregister_decrease(self.num_regs_wg0) + + if warp_idx == self.load_warp_id: + # --- partitions (loop-invariant) --- + gQ = cute.local_tile( + tma_tensor_Q, + cute.select(self.mma_tiler, mode=[1, 2]), + (None, None, None), + ) + tSgQ = thr_mma.partition_B(gQ) + tQsQ, tQgQ = cpasync.tma_partition( + tma_atom_Q, + 0, + cute.make_layout(1), + cute.group_modes(sQ, 0, 3), + cute.group_modes(tSgQ, 0, 3), + ) + gK = cute.local_tile( + tma_tensor_K, + cute.select(self.mma_tiler, mode=[0, 2]), + (None, None), + ) + tSgK = thr_mma.partition_A(gK) + tKsK, tKgK = cpasync.tma_partition( + tma_atom_K, + 0, + cute.make_layout(1), + cute.group_modes(sK, 0, 3), + cute.group_modes(tSgK, 0, 3), + ) + q_prod = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.q_stage) + k_prod = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.kv_stage) + # Q of the first block + tile0 = num_tiles - Int32(1) - Int32(bidx) + pipe_Q.producer_acquire(q_prod) + cute.copy( + tma_atom_Q, + tQgQ[(None, 0, 0, tile0)], + tQsQ[(None, q_prod.index)], + tma_bar_ptr=pipe_Q.producer_get_barrier(q_prod), + ) + q_prod.advance() + for w in cutlass.range(trips): + linear = Int32(bidx) + w * gdim + # prefetch the NEXT block's Q ahead of this block's + # KV tiles (so MMA never stalls on Q at a block boundary) + nxt = linear + gdim + if nxt < num_tiles: + ntile = num_tiles - Int32(1) - nxt + pipe_Q.producer_acquire(q_prod) + cute.copy( + tma_atom_Q, + tQgQ[(None, 0, 0, ntile)], + tQsQ[(None, q_prod.index)], + tma_bar_ptr=pipe_Q.producer_get_barrier(q_prod), + ) + q_prod.advance() + + tile_idx = num_tiles - Int32(1) - linear + q0 = tile_idx * self.tq + ks_min = skv + ke_max = Int32(0) + for r in cutlass.range_constexpr(self.tq): + ks_min = cutlass.min(ks_min, cutlass.min(Int32(mKS[q0 + r]), skv)) + ke_max = cutlass.max(ke_max, cutlass.min(Int32(mKE[q0 + r]), skv)) + k_tile0 = ks_min // self.n_block + span = cutlass.max(ke_max - k_tile0 * self.n_block, Int32(0)) + n_iters = cute.arch.make_warp_uniform(cute.ceil_div(span, self.n_block)) + k_tile0 = cute.arch.make_warp_uniform(k_tile0) + + for i in cutlass.range(n_iters): + pipe_K.producer_acquire(k_prod) + cute.copy( + tma_atom_K, + tKgK[(None, k_tile0 + i, 0)], + tKsK[(None, k_prod.index)], + tma_bar_ptr=pipe_K.producer_get_barrier(k_prod), + ) + k_prod.advance() + + elif warp_idx == self.mma_warp_id: + tmem.wait_for_alloc() + tmem_base = tmem.retrieve_ptr(Float32) + tAcc = cute.make_tensor(tmem_base, acc_fake.layout) + + tSrK = tiled_mma.make_fragment_A(sK) + tSrQ = tiled_mma.make_fragment_B(sQ) + + q_cons = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.q_stage) + k_cons = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.kv_stage) + s_prod = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.num_tmem_slots) + + for w in cutlass.range(trips): + linear = Int32(bidx) + w * gdim + tile_idx = num_tiles - Int32(1) - linear + q0 = tile_idx * self.tq + ks_min = skv + ke_max = Int32(0) + for r in cutlass.range_constexpr(self.tq): + ks_min = cutlass.min(ks_min, cutlass.min(Int32(mKS[q0 + r]), skv)) + ke_max = cutlass.max(ke_max, cutlass.min(Int32(mKE[q0 + r]), skv)) + k_tile0 = ks_min // self.n_block + span = cutlass.max(ke_max - k_tile0 * self.n_block, Int32(0)) + n_iters = cute.arch.make_warp_uniform(cute.ceil_div(span, self.n_block)) + + pipe_Q.consumer_wait(q_cons) + for i in cutlass.range(n_iters): + pipe_S.producer_acquire(s_prod) + acc = tAcc[(None, None, None, s_prod.index)] + tiled_mma.set(tcgen05.Field.ACCUMULATE, False) + pipe_K.consumer_wait(k_cons) + for kb in cutlass.range(0, cute.size(tSrQ, mode=[2]), unroll_full=True): + cute.gemm( + tiled_mma, + acc, + tSrK[(None, None, kb, k_cons.index)], + tSrQ[(None, None, kb, q_cons.index)], + acc, + ) + tiled_mma.set(tcgen05.Field.ACCUMULATE, True) + pipe_K.consumer_release(k_cons) + k_cons.advance() + pipe_S.producer_commit(s_prod) + s_prod.advance() + pipe_Q.consumer_release(q_cons) + q_cons.advance() + + # ============================================================= + # Warpgroups 1-2: epilogue — each owns one TMEM slot, drains + # alternating KV tiles (per-warpgroup UMMA ping-pong) + # ============================================================= + else: + cute.arch.setmaxregister_increase(self.num_regs_epi) + if warp_idx == self.epi_warp_ids[0]: + tmem.allocate(self.tmem_alloc_cols) + tmem.wait_for_alloc() + tmem_base = tmem.retrieve_ptr(Float32) + tAccE = cute.make_tensor(tmem_base, acc_fake_epi.layout) + + tidx_wg = tidx % 128 + epi_wg = wg_idx - 1 # 0 or 1 == owned TMEM slot + + thr_mma_epi = tiled_mma_epi.get_slice(tidx_wg) + cS = cute.make_identity_tensor((self.n_block, self.num_heads)) + tAcc0 = tAccE[(None, None, None, 0)] + # Compile-time layout probe: a probe tiled-copy built once out + # here derives this thread's KV row within the tile, the + # per-chunk register fragment shape, and the per-chunk TMEM + # source partitions (tTR_srcs below) — all of which feed the + # persistent loop. The tiled-copy actually issuing the LDTMs + # is re-created INSIDE the tile loop so no !cute.tiled_copy + # value is live across a dynamic back-edge (DSL 4.5.x + # make_tmem_copy loop-carry limitation); the probe itself is + # trace-time only and costs nothing at runtime. + epi_rep = self.epi_rep + probe_atom = cute.make_copy_atom( + tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(epi_rep)), + Float32, + ) + thr_probe = tcgen05.make_tmem_copy(probe_atom, tAcc0).get_slice(tidx_wg) + tScS = thr_probe.partition_D(thr_mma_epi.partition_C(cS)) + # this thread's KV row within the tile (== TMEM lane) + row = cute.get(tScS[0], mode=[0]) + + tSrS_shape = thr_probe.partition_D(cute.make_identity_tensor(tAcc0.shape)).shape + # rolling split-LDTM lookahead: two live fragments so every LDTM + # after the first is in flight behind the previous chunk's FMA + # reduction. Chunk computations are independent -> the reduction + # order (and hence the result) matches a serial per-chunk drain + # bit for bit. + tSrS_a = cute.make_rmem_tensor(tSrS_shape, Float32) + tSrS_b = cute.make_rmem_tensor(tSrS_shape, Float32) + rW = cute.make_rmem_tensor((self.num_heads,), Float32) + + mW_flat = cute.make_tensor( + mW.iterator, + cute.make_layout((cute.size(mW.shape[0]) * self.num_heads,)), + ) + # WG-private weight staging buffer (1KB apart -> align holds) + sW_wg = (sW.iterator + epi_wg * (self.tq * self.num_heads)).align(16) + + # Raw-mbarrier drain state: each epilogue WG owns exactly ONE + # of the two 256-column TMEM slots, so the generic pipeline + # consumer state machine (phase XORs, index SELs, advances) is + # replaced by one fixed full/empty mbarrier pair and a 1-bit + # phase. Semantically identical to consumer_wait(full, + # phase0-start) / arrive(empty): same mbarriers, same counts + # (128 arrives per stage = this WG). Hoisted per-chunk TMEM + # partitions (tq of them, fixed per WG) remove the per-tile + # dynamic TMEM layout arithmetic. Plain tensors cross the + # dynamic back-edge fine (the 4.5.x ICE is tiled-copy-only). + tTR_srcs = [] + for cc in cutlass.range_constexpr(self.tq): + tTR_srcs.append( + thr_probe.partition_S( + tAccE[ + ( + None, + None, + None, + Int32(epi_wg) * self.tq + cc, + ) + ] + ) + ) + full_bar = s_mbar_base + Int32(epi_wg) + empty_bar = s_mbar_base + (Int32(epi_wg) + self.num_tmem_slots) + phase = Int32(0) + + ks_q = cute.make_rmem_tensor((self.tq,), Int32) + ke_q = cute.make_rmem_tensor((self.tq,), Int32) + + g_par = Int32(0) # parity of the global KV-tile counter + + for w in cutlass.range(trips): + linear = Int32(bidx) + w * gdim + tile_idx = num_tiles - Int32(1) - linear + q0 = tile_idx * self.tq + ks_min = skv + ke_max = Int32(0) + for r in cutlass.range_constexpr(self.tq): + ks_min = cutlass.min(ks_min, cutlass.min(Int32(mKS[q0 + r]), skv)) + ke_max = cutlass.max(ke_max, cutlass.min(Int32(mKE[q0 + r]), skv)) + k_tile0 = ks_min // self.n_block + span = cutlass.max(ke_max - k_tile0 * self.n_block, Int32(0)) + n_iters = cute.arch.make_warp_uniform(cute.ceil_div(span, self.n_block)) + k_tile0 = cute.arch.make_warp_uniform(k_tile0) + + # stage this block's weights into the WG-private buffer. + # pre-barrier = WAR guard (a lagging warp of THIS WG may + # still read the previous block's weights); post = RAW. + # BF16 weights are up-converted here (exact), so the fp32 + # register math below is dtype-independent. + cute.arch.barrier(barrier_id=4 + epi_wg, number_of_threads=128) + wi = tidx_wg + while wi < self.tq * self.num_heads: + sW[epi_wg * (self.tq * self.num_heads) + wi] = Float32(mW_flat[q0 * self.num_heads + wi]) + wi += 128 + cute.arch.barrier(barrier_id=4 + epi_wg, number_of_threads=128) + + # per-row visibility windows for the -inf mask on + # out-of-window positions inside swept tiles + for r in cutlass.range_constexpr(self.tq): + ks_q[r] = Int32(mKS[q0 + r]) + ke_q[r] = Int32(mKE[q0 + r]) + + # this WG's KV tiles: local i with (g_par+i) % 2 == epi_wg + i0 = (g_par + Int32(epi_wg)) % Int32(2) + n_my = cutlass.max((n_iters - i0 + Int32(1)) // Int32(2), Int32(0)) + for t in cutlass.range(n_my): + i = i0 + 2 * t + col = (k_tile0 + i) * self.n_block + row + # wait this WG's fixed slot (1-bit phase) + cute.arch.mbarrier_wait(full_bar, phase) + _atomV = cute.make_copy_atom( + tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(epi_rep)), + Float32, + ) + _t2rV = tcgen05.make_tmem_copy(_atomV, tAcc0) + # rolling split-LDTM drain: ld c0, then per chunk + # {fence; issue the NEXT chunk's LDTM (or release the + # slot after the last fence); reduce the current chunk} + # — every LDTM after the first flies behind the previous + # chunk's FMA chain, and the UMMA warp refills the slot + # during the last reduction. + cute.copy(_t2rV, tTR_srcs[0], tSrS_a) + for qg in cutlass.range_constexpr(self.tq): + tSrSv = tSrS_a if qg % 2 == 0 else tSrS_b + nSrSv = tSrS_b if qg % 2 == 0 else tSrS_a + cute.arch.fence_view_async_tmem_load() + if const_expr(qg == self.tq - 1): + cute.arch.mbarrier_arrive(empty_bar) + else: + cute.copy( + _t2rV, + tTR_srcs[qg + 1], + nSrSv, + ) + sW_g = cute.make_tensor( + sW_wg + qg * self.num_heads, + cute.make_layout((self.num_heads,)), + ) + cute.autovec_copy(sW_g, rW) + acc2 = (Float32(0.0), Float32(0.0)) + for j in cutlass.range_constexpr(0, self.num_heads, 2): + v0 = cute.arch.fmax(tSrSv[j], Float32(0.0)) + v1 = cute.arch.fmax(tSrSv[j + 1], Float32(0.0)) + acc2 = cute.arch.fma_packed_f32x2( + (v0, v1), + (rW[j], rW[j + 1]), + acc2, + rnd="rn", + ) + score = acc2[0] + acc2[1] + if const_expr(sm_scale is not None): + # sm_scale on the fp32 head-reduced score + # (legacy placement: post-reduce, pre-mask) + score = score * sm_scale + # -inf on out-of-window positions inside the swept + # tile (in-window values are untouched by the mask) + ksq = Int32(ks_q[qg]) + keq = Int32(ke_q[qg]) + masked = Float32(float("-inf")) + if (col >= ksq) and (col < keq): + masked = score + score = masked + if col < skv: + mOut[q0 + qg, col] = score + phase = phase ^ Int32(1) + g_par = (g_par + n_iters) % Int32(2) + + # all epilogue TMEM reads done before dealloc + self.epi_sync_barrier.arrive_and_wait() + if warp_idx == self.epi_warp_ids[0]: + cute.arch.dealloc_tmem(tmem_base, self.tmem_alloc_cols) diff --git a/test/python/fe_api/dsa/test_DSA_indexer_forward_lean.py b/test/python/fe_api/dsa/test_DSA_indexer_forward_lean.py new file mode 100644 index 000000000..da8b814c8 --- /dev/null +++ b/test/python/fe_api/dsa/test_DSA_indexer_forward_lean.py @@ -0,0 +1,747 @@ +import pytest +import torch + +from test_utils import torch_fork_set_rng + +from fe_api.dsa.dsa_reference import check_ref_indexer_forward, ref_indexer_forward + +INT32_MAX = 2**31 - 1 + + +def _import_lean(): + try: + from cudnn.deepseek_sparse_attention.indexer_forward import api_lean + + return api_lean + except ImportError: + pytest.skip("Environment not supported: cudnn[cutedsl] not installed") + + +def _require_sm100(): + if not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 10: + pytest.skip("SM100 GPU required") + + +def _lean_min_s_q(api_lean) -> int: + """Smallest S_q that saturates the lean static persistent grid.""" + sm_count = torch.cuda.get_device_properties(torch.cuda.current_device()).multi_processor_count + return api_lean.LEAN_MIN_WAVES * sm_count * api_lean.LEAN_TILE_TOKENS + + +def _meta(shape, dtype, stride=None): + """Metadata-only sample tensor (meta device) for check_support tests.""" + if stride is None: + stride = [] + acc = 1 + for s in reversed(shape): + stride.append(acc) + acc *= s + stride = tuple(reversed(stride)) + return torch.empty_strided(shape, stride, dtype=dtype, device="meta") + + +def _meta_samples( + api_lean, b=1, s_q=None, s_k=None, h_q=64, h_kv=1, d=128, q_dtype=torch.bfloat16, w_dtype=torch.bfloat16, o_dtype=torch.float32, s_k_out=None +): + if s_q is None: + s_q = _lean_min_s_q(api_lean) + if s_k is None: + s_k = max(s_q // 2, 1) + if s_k_out is None: + s_k_out = s_k + q = _meta((b, s_q, h_q, d), q_dtype) + k = _meta((b, s_k, h_kv, d), q_dtype) + w = _meta((b, s_q, h_q), w_dtype) + o = _meta((b, s_q, s_k_out), o_dtype) + return q, k, w, o + + +def _alloc_inputs(b, s_q, s_k, h_q=64, h_kv=1, d=128, w_dtype=torch.bfloat16): + q = torch.randn(b, s_q, h_q, d, dtype=torch.bfloat16, device="cuda") + k = torch.randn(b, s_k, h_kv, d, dtype=torch.bfloat16, device="cuda") + w = torch.randn(b, s_q, h_q, dtype=torch.bfloat16, device="cuda").to(w_dtype) + return q, k, w + + +def _spy_lean_execute(monkeypatch, api_lean): + """Count IndexerForwardLean.execute invocations (per-batch launches count once).""" + calls = [] + original = api_lean.IndexerForwardLean.execute + + def spy(self, *args, **kwargs): + calls.append(1) + return original(self, *args, **kwargs) + + monkeypatch.setattr(api_lean.IndexerForwardLean, "execute", spy) + return calls + + +# --------------------------------------------------------------------------- +# check_support: pass / soft-fail / hard-fail +# --------------------------------------------------------------------------- + + +@pytest.mark.L0 +def test_DSA_indexer_forward_lean_check_support_pass(): + api_lean = _import_lean() + _require_sm100() + for kwargs in ( + {}, + {"b": 3}, + {"w_dtype": torch.float32}, + {"s_k": 999}, # S_k need not be a multiple of the 128-column KV tile + ): + q, k, w, o = _meta_samples(api_lean, **kwargs) + api = api_lean.IndexerForwardLean(q, k, w, o, ratio=2) + assert api.check_support() is True, f"expected supported for {kwargs}" + + +@pytest.mark.L0 +@pytest.mark.parametrize( + "case", + [ + "h32", # qhead_per_kv_head == 32 stays on the legacy kernel + "hkv2", # h_kv != 1 + "d64", # head_dim != 128 + "q_fp16", + "w_fp64", + "o_bf16", + "sq_nonmult4", # S_q not a multiple of the 4-token tile + "grid_small", # grid too small for the static single-wave schedule + "skout_mismatch", # Out column dim must equal S_k + ], +) +def test_DSA_indexer_forward_lean_check_support_false(case): + api_lean = _import_lean() + _require_sm100() + min_s_q = _lean_min_s_q(api_lean) + kwargs = { + "h32": {"h_q": 32}, + "hkv2": {"h_q": 128, "h_kv": 2}, + "d64": {"d": 64}, + "q_fp16": {"q_dtype": torch.float16}, + "w_fp64": {"w_dtype": torch.float64}, + "o_bf16": {"o_dtype": torch.bfloat16}, + # derived from the device's SM count (not hardcoded): a saturating + # tile count made non-multiple-of-4 + "sq_nonmult4": {"s_q": min_s_q * 8 + 2}, + "grid_small": {"s_q": 8}, + "skout_mismatch": {"s_k": 128, "s_k_out": 132}, + }[case] + q, k, w, o = _meta_samples(api_lean, **kwargs) + api = api_lean.IndexerForwardLean(q, k, w, o, ratio=2) + assert api.check_support() is False + + +@pytest.mark.L0 +def test_DSA_indexer_forward_lean_check_support_noncontiguous_false(): + api_lean = _import_lean() + _require_sm100() + q, k, w, o = _meta_samples(api_lean) + s = q.shape + q_pad = _meta(s, torch.bfloat16, stride=(s[1] * s[2] * (s[3] + 8), s[2] * (s[3] + 8), s[3] + 8, 1)) + api = api_lean.IndexerForwardLean(q_pad, k, w, o, ratio=2) + assert api.check_support() is False + + +@pytest.mark.L0 +def test_DSA_indexer_forward_lean_check_support_raises(): + api_lean = _import_lean() + _require_sm100() + q, k, w, o = _meta_samples(api_lean) + with pytest.raises(ValueError): + api_lean.IndexerForwardLean(q[0], k, w, o, ratio=2).check_support() # Q not 4-D + k2 = _meta((2, k.shape[1], 1, 128), torch.bfloat16) + with pytest.raises(ValueError): + api_lean.IndexerForwardLean(q, k2, w, o, ratio=2).check_support() # batch mismatch + w2 = _meta((1, q.shape[1], 32), torch.bfloat16) + with pytest.raises(ValueError): + api_lean.IndexerForwardLean(q, k, w2, o, ratio=2).check_support() # W shape mismatch + + +# --------------------------------------------------------------------------- +# numerics: lean wrapper vs the pure-torch reference +# --------------------------------------------------------------------------- + + +@pytest.mark.L0 +@torch_fork_set_rng(seed=0) +@pytest.mark.parametrize("batch,with_offsets", [(1, False), (2, True)]) +def test_DSA_indexer_forward_lean_wrapper_numerics(batch, with_offsets): + api_lean = _import_lean() + _require_sm100() + ratio = 4 + s_q = _lean_min_s_q(api_lean) + s_k = s_q // 2 + q, k, w = _alloc_inputs(batch, s_q, s_k) + q_causal_offsets = None + if with_offsets: + q_causal_offsets = torch.arange(3, 3 + 253 * batch, 253, dtype=torch.int32, device="cuda") + + result = api_lean.indexer_forward_lean_wrapper(q, k, w, ratio=ratio, q_causal_offsets=q_causal_offsets) + scores = result["scores"] + assert scores.shape == (batch, s_q, s_k) + assert scores.is_contiguous() + check_ref_indexer_forward(q, k, w, scores, ratio, q_causal_offsets=q_causal_offsets) + + +@pytest.mark.L0 +@torch_fork_set_rng(seed=1) +def test_DSA_indexer_forward_lean_wrapper_sm_scale_and_fp32_w(): + api_lean = _import_lean() + _require_sm100() + ratio, sm_scale = 4, 0.125 + s_q = _lean_min_s_q(api_lean) + s_k = s_q // 2 + q, k, w = _alloc_inputs(1, s_q, s_k, w_dtype=torch.float32) + + scores = api_lean.indexer_forward_lean_wrapper(q, k, w, ratio=ratio, sm_scale=sm_scale)["scores"] + ref = ref_indexer_forward(q, k, w, ratio) + finite = torch.isfinite(ref) + assert torch.equal(torch.isneginf(scores), torch.isneginf(ref)) + torch.testing.assert_close(scores[finite], ref[finite] * sm_scale, atol=1e-4, rtol=1e-4) + + +@pytest.mark.L0 +@torch_fork_set_rng(seed=2) +@pytest.mark.parametrize( + "s_k,ratio,offset_mode", + [ + (1, 1, None), # single-column K: one 1-row partial KV tile + (127, 1, None), # partial tile just below the 128-row boundary + (129, 1, None), # one full tile + a 1-row partial tail + (999, 1, None), # heavy clamping + partial tail, visibility sweeps it + (999, 4, "full"), # offsets push every row's window to the full S_k + ], +) +def test_DSA_indexer_forward_lean_wrapper_ragged_s_k_tail(s_k, ratio, offset_mode): + """Ragged S_k where the visibility window actually REACHES the physical + tail columns of the trailing partial KV tile (ratio=1 rows sweep ke + across the 128-column tile boundary and up to S_k; the offsets variant + drives every row's window to S_k), exercising the TMA zero-fill and the + bounds-guarded fp32 tail stores.""" + api_lean = _import_lean() + _require_sm100() + min_s_q = _lean_min_s_q(api_lean) + # ratio=1 windows reach column S_k-1 once i+1 >= S_k, so make S_q >= S_k + s_q = max(min_s_q, (s_k + 4) // 4 * 4) + q, k, w = _alloc_inputs(1, s_q, s_k) + q_causal_offsets = None + if offset_mode == "full": + q_causal_offsets = torch.full((1,), ratio * s_k, dtype=torch.int32, device="cuda") + + scores = api_lean.indexer_forward_lean_wrapper(q, k, w, ratio=ratio, q_causal_offsets=q_causal_offsets)["scores"] + assert scores.shape == (1, s_q, s_k) + # prove the tail was reached: the last row's window covers column S_k-1 + assert torch.isfinite(scores[0, -1, s_k - 1]), "tail column never became visible — test would be vacuous" + if offset_mode == "full": + assert torch.isfinite(scores).all(), "full-visibility offsets must produce finite scores everywhere" + check_ref_indexer_forward(q, k, w, scores, ratio, q_causal_offsets=q_causal_offsets) + + +@pytest.mark.L0 +@torch_fork_set_rng(seed=7) +def test_DSA_indexer_forward_lean_wrapper_all_masked_windows(): + """Offsets so negative that every per-row window is empty: the kernel + sweeps nothing and the -inf pre-fill must survive untouched.""" + api_lean = _import_lean() + _require_sm100() + ratio = 4 + s_q = _lean_min_s_q(api_lean) + s_k = s_q // 2 + q, k, w = _alloc_inputs(1, s_q, s_k) + q_causal_offsets = torch.full((1,), -(s_q + ratio), dtype=torch.int32, device="cuda") + + scores = api_lean.indexer_forward_lean_wrapper(q, k, w, ratio=ratio, q_causal_offsets=q_causal_offsets)["scores"] + assert torch.isneginf(scores).all() + ref = ref_indexer_forward(q, k, w, ratio, q_causal_offsets=q_causal_offsets) + assert torch.equal(torch.isneginf(scores), torch.isneginf(ref)) + + +@pytest.mark.L0 +@torch_fork_set_rng(seed=8) +def test_DSA_indexer_forward_lean_wrapper_w_dtype_bitwise_and_determinism(): + """BF16 W is up-converted in-kernel exactly, so BF16-W and FP32-W runs + must be bitwise identical; the fixed-order epilogue reduction must be + deterministic run-to-run.""" + api_lean = _import_lean() + _require_sm100() + ratio = 2 + s_q = _lean_min_s_q(api_lean) + s_k = s_q // 2 + q, k, w = _alloc_inputs(1, s_q, s_k) + + scores_bf16w = api_lean.indexer_forward_lean_wrapper(q, k, w, ratio=ratio)["scores"] + scores_fp32w = api_lean.indexer_forward_lean_wrapper(q, k, w.float().contiguous(), ratio=ratio)["scores"] + assert torch.equal(scores_bf16w, scores_fp32w), "BF16-W ingest must be bitwise identical to FP32-W" + + scores_again = api_lean.indexer_forward_lean_wrapper(q, k, w, ratio=ratio)["scores"] + assert torch.equal(scores_bf16w, scores_again), "lean kernel must be deterministic run-to-run" + + +# --------------------------------------------------------------------------- +# transparent dispatch through indexer_forward_wrapper +# --------------------------------------------------------------------------- + + +@pytest.mark.L0 +@torch_fork_set_rng(seed=3) +def test_DSA_indexer_forward_dispatch_routes_lean(monkeypatch): + api_lean = _import_lean() + from cudnn import DSA + + _require_sm100() + monkeypatch.delenv("CUDNNFE_DSA_INDEXER_FWD_DISABLE_LEAN", raising=False) + ratio, sm_scale = 4, 0.25 + s_q = _lean_min_s_q(api_lean) + s_k = s_q // 2 + q, k, w = _alloc_inputs(1, s_q, s_k) + q_causal_offsets = torch.tensor([7], dtype=torch.int32, device="cuda") + + calls = _spy_lean_execute(monkeypatch, api_lean) + dispatched = DSA.indexer_forward_wrapper(q, k, w, ratio=ratio, sm_scale=sm_scale, q_causal_offsets=q_causal_offsets)["scores"] + assert len(calls) == 1, "lean fast path was not dispatched" + + # the family wrapper must return exactly what the lean wrapper returns + # (ratio / sm_scale / q_causal_offsets plumbed through unchanged) + direct = api_lean.indexer_forward_lean_wrapper(q, k, w, ratio=ratio, sm_scale=sm_scale, q_causal_offsets=q_causal_offsets)["scores"] + assert torch.equal(dispatched, direct) + + ref = ref_indexer_forward(q, k, w, ratio, q_causal_offsets=q_causal_offsets) + finite = torch.isfinite(ref) + assert torch.equal(torch.isneginf(dispatched), torch.isneginf(ref)) + torch.testing.assert_close(dispatched[finite], ref[finite] * sm_scale, atol=1e-4, rtol=1e-4) + + +@pytest.mark.L0 +@torch_fork_set_rng(seed=9) +def test_DSA_indexer_forward_dispatch_shares_compile_across_batch(monkeypatch): + """The lean kernel runs on flattened per-batch views, so one compiled + instance must serve every B of the same (S_q, S_k): the B=3 call after a + B=1 call must trigger zero additional cute.compile invocations AND be + numerically correct (empirical batch-independence of the codegen key).""" + api_lean = _import_lean() + from cudnn import DSA + + _require_sm100() + monkeypatch.delenv("CUDNNFE_DSA_INDEXER_FWD_DISABLE_LEAN", raising=False) + ratio = 4 + s_q = _lean_min_s_q(api_lean) + s_k = s_q // 2 + 16 # unique S_k so this test owns its compile-cache entry + + compile_calls = [] + original_compile = api_lean.cute.compile + + def compile_spy(*args, **kwargs): + compile_calls.append(1) + return original_compile(*args, **kwargs) + + monkeypatch.setattr(api_lean.cute, "compile", compile_spy) + calls = _spy_lean_execute(monkeypatch, api_lean) + + q1, k1, w1 = _alloc_inputs(1, s_q, s_k) + scores1 = DSA.indexer_forward_wrapper(q1, k1, w1, ratio=ratio)["scores"] + assert len(calls) == 1, "lean fast path was not dispatched for B=1" + compiles_after_b1 = len(compile_calls) + assert compiles_after_b1 >= 1, "expected the fresh (S_q, S_k) to JIT once" + check_ref_indexer_forward(q1, k1, w1, scores1, ratio) + + q3, k3, w3 = _alloc_inputs(3, s_q, s_k) + scores3 = DSA.indexer_forward_wrapper(q3, k3, w3, ratio=ratio)["scores"] + assert len(calls) == 2, "lean fast path was not dispatched for B=3" + assert len(compile_calls) == compiles_after_b1, "B must not appear in the codegen key — no recompile for a new batch size" + check_ref_indexer_forward(q3, k3, w3, scores3, ratio) + + +@pytest.mark.L0 +@torch_fork_set_rng(seed=10) +def test_DSA_indexer_forward_dispatch_int32_offset_bound(monkeypatch): + """Legacy evaluates (offset + i + 1) in int32 on device; the dispatcher + must keep offsets that would overflow that arithmetic on the legacy path + (offset + S_q + 1 > INT32_MAX) and may dispatch anything below the bound.""" + api_lean = _import_lean() + from cudnn import DSA + + _require_sm100() + monkeypatch.delenv("CUDNNFE_DSA_INDEXER_FWD_DISABLE_LEAN", raising=False) + ratio = 4 + s_q = _lean_min_s_q(api_lean) + s_k = s_q // 2 + q, k, w = _alloc_inputs(1, s_q, s_k) + + calls = _spy_lean_execute(monkeypatch, api_lean) + + # largest accepted offset: offset + S_q + 1 == INT32_MAX -> lean + off_ok = torch.full((1,), INT32_MAX - s_q - 1, dtype=torch.int32, device="cuda") + scores_ok = DSA.indexer_forward_wrapper(q, k, w, ratio=ratio, q_causal_offsets=off_ok)["scores"] + assert len(calls) == 1, "boundary-accepted offsets must still dispatch lean" + check_ref_indexer_forward(q, k, w, scores_ok, ratio, q_causal_offsets=off_ok) + + # one past the bound: legacy int32 window math could overflow -> legacy + # (legacy itself still stays exactly at INT32_MAX here, so its output + # remains reference-correct; the conservative bound is on the dispatch) + off_over = torch.full((1,), INT32_MAX - s_q, dtype=torch.int32, device="cuda") + scores_over = DSA.indexer_forward_wrapper(q, k, w, ratio=ratio, q_causal_offsets=off_over)["scores"] + assert len(calls) == 1, "offsets past the int32 bound must stay on the legacy path" + check_ref_indexer_forward(q, k, w, scores_over, ratio, q_causal_offsets=off_over) + + +@pytest.mark.L0 +@torch_fork_set_rng(seed=4) +def test_DSA_indexer_forward_dispatch_env_disable(monkeypatch): + api_lean = _import_lean() + from cudnn import DSA + + _require_sm100() + ratio = 4 + s_q = _lean_min_s_q(api_lean) + s_k = s_q // 2 + q, k, w = _alloc_inputs(1, s_q, s_k) + q_causal_offsets = torch.tensor([4], dtype=torch.int32, device="cuda") + + monkeypatch.setenv("CUDNNFE_DSA_INDEXER_FWD_DISABLE_LEAN", "1") + calls = _spy_lean_execute(monkeypatch, api_lean) + scores = DSA.indexer_forward_wrapper(q, k, w, ratio=ratio, q_causal_offsets=q_causal_offsets)["scores"] + assert len(calls) == 0, "lean fast path must stay off when disabled by env" + check_ref_indexer_forward(q, k, w, scores, ratio, q_causal_offsets=q_causal_offsets) + + +@pytest.mark.L0 +@torch_fork_set_rng(seed=5) +def test_DSA_indexer_forward_dispatch_keeps_legacy_for_h32(monkeypatch): + api_lean = _import_lean() + from cudnn import DSA + + _require_sm100() + ratio = 4 + # S_q saturates the lean grid, so qhead_per_kv_head=32 is the ONLY + # reason this config is rejected (isolates the H gate) + b, s_q = 1, _lean_min_s_q(api_lean) + s_k = s_q // 2 + q, k, w = _alloc_inputs(b, s_q, s_k, h_q=32) + q_causal_offsets = torch.full((b,), 4, dtype=torch.int32, device="cuda") + + calls = _spy_lean_execute(monkeypatch, api_lean) + scores = DSA.indexer_forward_wrapper(q, k, w, ratio=ratio, qhead_per_kv_head=32, q_causal_offsets=q_causal_offsets)["scores"] + assert len(calls) == 0, "qhead_per_kv_head=32 must keep the legacy path" + check_ref_indexer_forward(q, k, w, scores, ratio, q_causal_offsets=q_causal_offsets) + + +@pytest.mark.L0 +@torch_fork_set_rng(seed=11) +def test_DSA_indexer_forward_dispatch_keeps_legacy_for_non_default_knobs(monkeypatch): + api_lean = _import_lean() + from cudnn import DSA + + _require_sm100() + ratio = 4 + # lean-eligible shape, but a non-default tuning knob must force legacy + s_q = _lean_min_s_q(api_lean) + s_k = s_q // 2 + q, k, w = _alloc_inputs(1, s_q, s_k) + + calls = _spy_lean_execute(monkeypatch, api_lean) + scores = DSA.indexer_forward_wrapper(q, k, w, ratio=ratio, kv_stage=2)["scores"] + assert len(calls) == 0, "non-default tuning knobs must keep the legacy path" + check_ref_indexer_forward(q, k, w, scores, ratio) + + +@pytest.mark.L0 +@torch_fork_set_rng(seed=6) +def test_DSA_indexer_forward_dispatch_keeps_legacy_for_thd(monkeypatch): + api_lean = _import_lean() + from cudnn import DSA + + _require_sm100() + device = torch.device("cuda") + shapes = [(8, 64), (12, 72)] + ratio, h_q, h_kv, d = 4, 64, 1, 128 + q_lengths = [s_q for s_q, _ in shapes] + k_lengths = [s_k for _, s_k in shapes] + cu_seqlens_q = torch.tensor([0, *torch.tensor(q_lengths).cumsum(0).tolist()], dtype=torch.int32, device=device) + cu_seqlens_k = torch.tensor([0, *torch.tensor(k_lengths).cumsum(0).tolist()], dtype=torch.int32, device=device) + total_q, total_k = int(cu_seqlens_q[-1]), int(cu_seqlens_k[-1]) + q = torch.randn(total_q, h_q, d, dtype=torch.bfloat16, device=device) + k = torch.randn(total_k, h_kv, d, dtype=torch.bfloat16, device=device) + w = torch.randn(total_q, h_q, dtype=torch.bfloat16, device=device) + + calls = _spy_lean_execute(monkeypatch, api_lean) + scores = DSA.indexer_forward_wrapper( + q, + k, + w, + ratio=ratio, + qhead_per_kv_head=h_q, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + max_seqlen_q=max(q_lengths), + max_seqlen_k=max(k_lengths), + )["scores"] + torch.cuda.synchronize() + assert len(calls) == 0, "THD/varlen must keep the legacy path" + + cu_q_host = cu_seqlens_q.tolist() + cu_k_host = cu_seqlens_k.tolist() + for batch, (_, s_k) in enumerate(shapes): + q0, q1 = cu_q_host[batch : batch + 2] + k0, k1 = cu_k_host[batch : batch + 2] + check_ref_indexer_forward( + q[q0:q1].unsqueeze(0), + k[k0:k1].unsqueeze(0), + w[q0:q1].unsqueeze(0), + scores[q0:q1, :s_k].unsqueeze(0), + ratio, + ) + + +# --------------------------------------------------------------------------- +# THD / varlen (ragged packed) lean fast path — GLOBAL compressed-KV columns +# --------------------------------------------------------------------------- + + +def _cu(seg_lens, device="cuda"): + """(len+1,) int32 cu_seqlens from per-segment lengths.""" + cs = torch.tensor(seg_lens, dtype=torch.int64).cumsum(0).tolist() + return torch.tensor([0, *cs], dtype=torch.int32, device=device) + + +def _alloc_thd(seg_q, seg_k, h_q=64, h_kv=1, d=128, w_dtype=torch.bfloat16): + """Packed THD inputs for unequal segments plus their cu_seqlens.""" + t_q = int(sum(seg_q)) + m_total = int(sum(seg_k)) + q = torch.randn(t_q, h_q, d, dtype=torch.bfloat16, device="cuda") + k = torch.randn(m_total, h_kv, d, dtype=torch.bfloat16, device="cuda") + w = torch.randn(t_q, h_q, dtype=torch.bfloat16, device="cuda").to(w_dtype) + return q, k, w, _cu(seg_q), _cu(seg_k) + + +def _meta_samples_thd(seg_q, seg_k, h_q=64, h_kv=1, d=128, q_dtype=torch.bfloat16, w_dtype=torch.bfloat16, o_dtype=torch.float32): + t_q = int(sum(seg_q)) + m_total = int(sum(seg_k)) + q = _meta((t_q, h_q, d), q_dtype) + k = _meta((m_total, h_kv, d), q_dtype) + w = _meta((t_q, h_q), w_dtype) + o = _meta((t_q, m_total), o_dtype) + return q, k, w, o, _cu(seg_q), _cu(seg_k) + + +def _thd_segments_saturating(api_lean, n_seg=3): + """Unequal q-segments whose T_q saturates the lean grid (each %4).""" + base = _lean_min_s_q(api_lean) # already a multiple of LEAN_TILE_TOKENS + seg_q = [base + 4 * i for i in range(n_seg)] # unequal, each %4, sums > base + seg_k = [max(s // 2, 4) for s in seg_q] # ratio-2-like geometry, unequal + return seg_q, seg_k + + +@pytest.mark.L0 +def test_DSA_indexer_forward_lean_thd_check_support_pass(): + api_lean = _import_lean() + _require_sm100() + seg_q, seg_k = _thd_segments_saturating(api_lean) + for w_dtype in (torch.bfloat16, torch.float32): + q, k, w, o, cu_q, cu_k = _meta_samples_thd(seg_q, seg_k, w_dtype=w_dtype) + api = api_lean.IndexerForwardLean(q, k, w, o, ratio=2, cu_seqlens_q=cu_q, cu_seqlens_k=cu_k) + assert api.check_support() is True, f"expected THD supported for w_dtype={w_dtype}" + assert api._thd is True and api.s_q == int(sum(seg_q)) and api.s_k == int(sum(seg_k)) + + +@pytest.mark.L0 +@pytest.mark.parametrize("case", ["h32", "d64", "q_fp16", "grid_small", "tq_nonmult4"]) +def test_DSA_indexer_forward_lean_thd_check_support_false(case): + api_lean = _import_lean() + _require_sm100() + seg_q, seg_k = _thd_segments_saturating(api_lean) + kwargs = { + "h32": dict(h_q=32), + "d64": dict(d=64), + "q_fp16": dict(q_dtype=torch.float16), + "grid_small": None, # handled below + "tq_nonmult4": None, + }[case] + if case == "grid_small": + seg_q, seg_k = [8], [4] + kwargs = {} + elif case == "tq_nonmult4": + seg_q = [_lean_min_s_q(api_lean) + 2] # T_q not a multiple of 4 + seg_k = [64] + kwargs = {} + q, k, w, o, cu_q, cu_k = _meta_samples_thd(seg_q, seg_k, **kwargs) + api = api_lean.IndexerForwardLean(q, k, w, o, ratio=2, cu_seqlens_q=cu_q, cu_seqlens_k=cu_k) + assert api.check_support() is False + + +@pytest.mark.L0 +def test_DSA_indexer_forward_lean_thd_check_support_raises(): + api_lean = _import_lean() + _require_sm100() + seg_q, seg_k = _thd_segments_saturating(api_lean, n_seg=2) + t_q, m_total = int(sum(seg_q)), int(sum(seg_k)) + + # missing cu_seqlens_k + q, k, w, o, cu_q, cu_k = _meta_samples_thd(seg_q, seg_k) + with pytest.raises(ValueError): + api_lean.IndexerForwardLean(q, k, w, o, ratio=2, cu_seqlens_q=cu_q).check_support() + + # cu_seqlens_k[-1] disagrees with packed K rows (m_total) + q, k, w, o, cu_q, _ = _meta_samples_thd(seg_q, seg_k) + bad_k = torch.tensor([0, seg_k[0], m_total + 8], dtype=torch.int32, device="cuda") + with pytest.raises(ValueError): + api_lean.IndexerForwardLean(q, k, w, o, ratio=2, cu_seqlens_q=cu_q, cu_seqlens_k=bad_k).check_support() + + # non-monotonic / non-zero-start cu_seqlens_q + q, k, w, o, _, cu_k = _meta_samples_thd(seg_q, seg_k) + nonmono = torch.tensor([seg_q[0], 0, t_q], dtype=torch.int32, device="cuda") + with pytest.raises(ValueError): + api_lean.IndexerForwardLean(q, k, w, o, ratio=2, cu_seqlens_q=nonmono, cu_seqlens_k=cu_k).check_support() + + +def _check_thd_scores(api_lean, q, k, w, cu_q, cu_k, scores, ratio, sm_scale=1.0): + """Per-segment: global-column block matches the BSHD oracle; every + column outside a row's own segment block is -inf (segment isolation).""" + t_q, m_total = int(cu_q[-1]), int(cu_k[-1]) + assert tuple(scores.shape) == (t_q, m_total) + assert scores.is_contiguous() + cu_q_host, cu_k_host = cu_q.tolist(), cu_k.tolist() + n_seg = len(cu_q_host) - 1 + for b in range(n_seg): + q0, q1 = cu_q_host[b], cu_q_host[b + 1] + k0, k1 = cu_k_host[b], cu_k_host[b + 1] + if q1 <= q0: + continue + block = scores[q0:q1, k0:k1].unsqueeze(0) + ref = ref_indexer_forward(q[q0:q1].unsqueeze(0), k[k0:k1].unsqueeze(0), w[q0:q1].unsqueeze(0), ratio) + finite = torch.isfinite(ref) + assert torch.equal(torch.isneginf(block), torch.isneginf(ref)), f"seg {b} mask mismatch" + torch.testing.assert_close(block[finite], ref[finite] * sm_scale, atol=1e-4, rtol=1e-4) + # segment isolation: everything OUTSIDE this segment's KV block, for + # these rows, must be -inf (a query in segment b sees only seg b's KV) + outside = scores[q0:q1].clone() + outside[:, k0:k1] = float("-inf") + assert torch.isneginf(outside).all(), f"seg {b} leaked finite scores into another segment's columns" + + +@pytest.mark.L0 +@torch_fork_set_rng(seed=20) +@pytest.mark.parametrize("ratio", [2, 4]) +def test_DSA_indexer_forward_lean_thd_numerics_ragged(ratio): + """Ragged multi-segment THD vs the per-segment fp32 oracle + strict + segment isolation, at the same 1e-4 tolerance the BSHD/legacy suite uses.""" + api_lean = _import_lean() + _require_sm100() + seg_q, seg_k = _thd_segments_saturating(api_lean, n_seg=3) + q, k, w, cu_q, cu_k = _alloc_thd(seg_q, seg_k) + scores = api_lean.indexer_forward_lean_wrapper( + q, + k, + w, + ratio=ratio, + cu_seqlens_q=cu_q, + cu_seqlens_k=cu_k, + max_seqlen_q=max(seg_q), + max_seqlen_k=max(seg_k), + )["scores"] + _check_thd_scores(api_lean, q, k, w, cu_q, cu_k, scores, ratio) + + +@pytest.mark.L0 +@torch_fork_set_rng(seed=21) +def test_DSA_indexer_forward_lean_thd_sm_scale_and_fp32_w(): + api_lean = _import_lean() + _require_sm100() + ratio, sm_scale = 4, 0.125 + seg_q, seg_k = _thd_segments_saturating(api_lean, n_seg=2) + q, k, w, cu_q, cu_k = _alloc_thd(seg_q, seg_k, w_dtype=torch.float32) + scores = api_lean.indexer_forward_lean_wrapper( + q, + k, + w, + ratio=ratio, + sm_scale=sm_scale, + cu_seqlens_q=cu_q, + cu_seqlens_k=cu_k, + max_seqlen_q=max(seg_q), + max_seqlen_k=max(seg_k), + )["scores"] + _check_thd_scores(api_lean, q, k, w, cu_q, cu_k, scores, ratio, sm_scale=sm_scale) + + +@pytest.mark.L0 +@torch_fork_set_rng(seed=22) +def test_DSA_indexer_forward_lean_thd_single_segment_equals_bshd(): + """A single-segment THD problem (cu_seqlens = [0, T]) has ks == 0 windows + identical to the B=1 BSHD path, so the shared kernel must produce a + BITWISE-identical score matrix — proving THD reuses the exact schedule.""" + api_lean = _import_lean() + _require_sm100() + ratio = 4 + s_q = _lean_min_s_q(api_lean) + s_k = s_q // 2 + q, k, w = _alloc_inputs(1, s_q, s_k) # BSHD (1, s_q, H, D) + + bshd = api_lean.indexer_forward_lean_wrapper(q, k, w, ratio=ratio)["scores"][0] + + q_thd = q.view(s_q, 64, 128).contiguous() + k_thd = k.view(s_k, 1, 128).contiguous() + w_thd = w.view(s_q, 64).contiguous() + cu_q = torch.tensor([0, s_q], dtype=torch.int32, device="cuda") + cu_k = torch.tensor([0, s_k], dtype=torch.int32, device="cuda") + thd = api_lean.indexer_forward_lean_wrapper( + q_thd, + k_thd, + w_thd, + ratio=ratio, + cu_seqlens_q=cu_q, + cu_seqlens_k=cu_k, + max_seqlen_q=s_q, + max_seqlen_k=s_k, + )["scores"] + assert thd.shape == (s_q, s_k) + assert torch.equal(thd, bshd), "single-segment THD must be bitwise identical to B=1 BSHD" + + +@pytest.mark.L0 +@torch_fork_set_rng(seed=23) +def test_DSA_indexer_forward_lean_thd_q_causal_offsets_rejected(): + """q_causal_offsets is intentionally unsupported on the THD lean path.""" + api_lean = _import_lean() + _require_sm100() + seg_q, seg_k = _thd_segments_saturating(api_lean, n_seg=2) + q, k, w, cu_q, cu_k = _alloc_thd(seg_q, seg_k) + off = torch.zeros(2, dtype=torch.int32, device="cuda") + with pytest.raises(ValueError): + api_lean.indexer_forward_lean_wrapper( + q, + k, + w, + ratio=4, + cu_seqlens_q=cu_q, + cu_seqlens_k=cu_k, + max_seqlen_q=max(seg_q), + max_seqlen_k=max(seg_k), + q_causal_offsets=off, + ) + + +@pytest.mark.L0 +@torch_fork_set_rng(seed=24) +def test_DSA_indexer_forward_lean_thd_empty_and_tiny_segments(): + """Empty query segment + a single-token KV segment: windows collapse to + empty ([-inf] rows) or 1-column, and the ragged K tail is exercised.""" + api_lean = _import_lean() + _require_sm100() + ratio = 2 + base = _lean_min_s_q(api_lean) + seg_q = [base, 0, 8] # middle segment has zero queries + seg_k = [base // 2, 16, 1] + q, k, w, cu_q, cu_k = _alloc_thd(seg_q, seg_k) + scores = api_lean.indexer_forward_lean_wrapper( + q, + k, + w, + ratio=ratio, + cu_seqlens_q=cu_q, + cu_seqlens_k=cu_k, + max_seqlen_q=max(seg_q), + max_seqlen_k=max(seg_k), + )["scores"] + _check_thd_scores(api_lean, q, k, w, cu_q, cu_k, scores, ratio)