DSA indexer forward: add a lean SM100 fast path for the head_dim=128 / qhead_per_kv_head=64 regime (up to 1.46x)

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,

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)

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"),

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",
+]

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