dsa(indexer_backward): optimize SM100 gradient kernels

docs/fe-oss-apis/dsa.md

@@ -305,8 +305,9 @@ and normalization semantics differ from the indexer path.
 Three-stage sparse top-K pipeline that produces the training gradients for the
 indexer tower:
 
-1. `ScoreGradSm90` / `ScoreGradSm100` (kernel 1) — in-place score-grad precompute from
-   `attn_score` (target) and `index_score` (predict).
+1. `ScoreGradSm90` / `ScoreGradSm100` (kernel 1) — in-place score-grad precompute
+   that overwrites `attn_score` (target) and reads `index_score` (predict)
+   without modifying it.
 2. `IndexerBackwardSm90` / `IndexerBackwardSm100` (kernel 2) — three
    warp-specialised GEMMs produce `d_index_q`, `d_weights`, and a
    `dIndexK_f32` accumulator.
@@ -331,11 +332,11 @@ d_index_q, d_weights, d_index_k = (
 
 When compressed forward returns its fused `softmax`, backward can skip both
 indexer Q@K score recompute and the separate logits softmax. Pass `softmax`
-directly as `index_score`; backward consumes and overwrites this buffer, so
-pass `softmax.clone()` if it must be preserved. `attn_score` must use the same
-valid-slot mask. Because compressed forward returns global indices by default,
-also pass `topk_indices_global=True` unless forward used
-`topk_indices_global=False`. The public sparse `indexer_backward_wrapper` has a
+directly as `index_score`; backward treats this buffer as read-only.
+`attn_score` must use the same valid-slot mask. Because compressed forward
+returns global indices by default, also pass `topk_indices_global=True` unless
+forward used `topk_indices_global=False`. The public sparse
+`indexer_backward_wrapper` has a
 BSHD-shaped interface; BF16 THD tensors can use zero-copy `B=1` views (squeeze
 the singleton K head and add a batch dimension) together with global Top-K
 indices. FP8 and MXFP8 indexer backward are not currently supported because
@@ -388,10 +389,10 @@ the default backend.
   full dQ/dK contribution (`g' * w`, not scaled by `S`) is what differs. The
   default backend's multiply preserves subnormals, so reaching it takes an
   exact `sm_scale * S` below `2^-150` (~7.0e-46).
-- **Scratch / workspace behavior** — `attn_score`/`index_score` are consumed
-  in place exactly like the default backend (`attn_score` is left holding
-  kernel 1's `grad_signal`; `sm_scale` folds inside kernel 2 without touching
-  the buffer). The backend owns one piece of per-plan workspace — the
+- **Scratch / workspace behavior** — `attn_score` is consumed in place and left
+  holding kernel 1's `grad_signal`, while `index_score` is read-only and
+  preserved. `sm_scale` folds inside kernel 2 without touching either buffer.
+  The backend owns one piece of per-plan workspace — the
   dynamic-ticket counter — allocated on first execute and reused by every
   later one. A BF16 `d_index_k` additionally needs a `B * S_k * D` fp32
   accumulator (2 MiB = 2,097,152 bytes at B=1, S_k=4096, D=128, growing with
@@ -403,7 +404,7 @@ the default backend.
 - **Concurrency** — executions sharing one plan must not overlap on the
   device (the ticket counter is per-plan workspace). One plan serves one
   device: the workspace is device-resident, and execution rejects tensors on
-  any other device before touching the score buffers. The wrapper keys its
+  any other device before overwriting `attn_score`. The wrapper keys its
   plan cache on the CUDA device and on the **resolved** stream (`stream` when
   given, otherwise `torch.cuda.current_stream()` at call time), so calls that
   differ in device or stream get a private plan and private workspace; calls

python/cudnn/deepseek_sparse_attention/indexer_backward/api.py

@@ -169,11 +169,11 @@ class IndexerBackward(APIBase):
 
     1. Runtime signature re-validation of the tensors against the descriptors
        captured at plan-build time (dtype, shape, and stride/layout),
-       *before* any kernel launch — kernel 1 mutates the score buffers in
+       *before* any kernel launch — kernel 1 overwrites ``attn_score`` in
        place, so a mismatched reuse of an exported plan must raise here to
        avoid a fail-dirty.
-    2. Kernel 1 — in-place score-grad precompute (overwrites ``attn_score`` and
-       ``index_score``).
+    2. Kernel 1 — in-place score-grad precompute (overwrites ``attn_score``
+       with ``grad_signal`` and reads ``index_score`` without modifying it).
     3. (conditional) fp32 ``d_index_k`` zero-init — when the ``d_index_k``
        output buffer is fp32, ``execute`` zeroes it on the selected stream
        before the GEMM, because the dK epilogue atomic-adds into it. This is a
@@ -191,8 +191,9 @@ class IndexerBackward(APIBase):
     of the 24 fp32 significand bits — measured ~675x lower gradient-matrix
     representation error than the default single-bf16 rounding);
     ``d_weights`` is additionally reduced deterministically (bitwise
-    run-to-run, as is ``d_index_q``). Same wrapper contract and in-place
-    score consumption. Output dtype selects output precision: ``d_weights``
+    run-to-run, as is ``d_index_q``). Same wrapper contract: ``attn_score`` is
+    the in-place score scratch and ``index_score`` remains read-only. Output
+    dtype selects output precision: ``d_weights``
     and ``d_index_k`` accept fp32 buffers, which receive the fp32
     accumulators directly — **the headline accuracy gains require fp32
     output buffers**; bf16 output buffers round the result and keep only the
@@ -270,13 +271,13 @@ def __init__(
     def _validate_plan_shapes_and_layout(self) -> None:
         """Semantic shape + layout validation of the plan's tensor descriptors.
 
-        Runs before kernel 1 mutates the score buffers in place.
+        Runs before kernel 1 overwrites ``attn_score`` in place.
         ``check_support`` otherwise only checks dtypes and ``execute``'s
         signature check only compares each runtime tensor against the recorded
         descriptor, so without this a directly-built / first-call plan whose
         ``d_weights`` / ``d_index_q`` / ``d_index_k`` / score / top-k shape is
         already inconsistent with ``index_q`` would pass validation, kernel 1
-        would corrupt the scores, and only the GEMM would fault or corrupt
+        would overwrite ``attn_score``, and only the GEMM would fault or corrupt
         memory. The relationships enforced (per the API contract):
 
           index_q      (B, S_q, H, D)      weights   (B, S_q, H)
@@ -366,7 +367,7 @@ def check_support(self) -> bool:
             # full metadata matrix: cross-tensor shapes, output dtypes,
             # devices, contiguity — the backend binds true views and uses
             # 16B-aligned row pointers, so reject bad metadata before
-            # kernel 1 mutates the score buffers.
+            # kernel 1 overwrites ``attn_score``.
             self._check_tensor_shape(self.w_desc, (b, s_q, h), name="weights")
             self._check_tensor_shape(self.ik_desc, (b, s_k, d), name="index_k")
             self._check_tensor_shape(self.diq_desc, (b, s_q, h, d), name="d_index_q")
@@ -409,9 +410,9 @@ def check_support(self) -> bool:
         self._check_dtype(self.idx_score_desc, torch.float32, name="index_score")
         self._check_dtype(self.topk_desc, torch.int32, name="topk_indices")
         # Output-dtype contract, validated here — before compile()/execute(),
-        # i.e. before kernel 1 mutates attn_score/index_score in place (no
-        # fail-dirty on a bad output dtype). ``d_index_q`` is bf16-only (TMA
-        # store of the input dtype). ``d_index_k`` accepts bf16 or fp32 on both
+        # i.e. before kernel 1 overwrites ``attn_score`` (``index_score`` is
+        # read-only; no fail-dirty on a bad output dtype). ``d_index_q`` is
+        # bf16-only (TMA store of the input dtype). ``d_index_k`` accepts bf16 or fp32 on both
         # arches (the GEMM accumulates dK in fp32; a bf16 buffer gets a
         # trailing cast). ``d_weights`` is bf16-only on the default backend:
         # its kernel's dW store rounds the fp32 accumulator to bf16, so an
@@ -435,8 +436,8 @@ def check_support(self) -> bool:
                 extra_error_msg="fp32 buffers receive the fp32 accumulator directly",
             )
         # Semantic shape relationships + compact-layout requirement, validated
-        # here (before compile()/execute(), i.e. before kernel 1 mutates the
-        # score buffers). Guards a directly-built / first-call plan whose output
+        # here (before compile()/execute(), i.e. before kernel 1 overwrites
+        # ``attn_score``). Guards a directly-built / first-call plan whose output
         # or score shapes are inconsistent, or whose layout is non-compact.
         self._validate_plan_shapes_and_layout()
         self._is_supported = True
@@ -496,7 +497,7 @@ def _check_execute_signature(self, *entries) -> None:
         Each entry is ``(tensor, descriptor, name)``. Guards against reusing a
         directly-built/exported plan with a tensor whose dtype, shape, or
         stride/layout differs from what it was compiled for — kernel 1
-        mutates the score buffers in place, so a mismatch must raise *before*
+        overwrites ``attn_score`` in place, so a mismatch must raise *before*
         the pipeline starts (no fail-dirty).
         """
         for tensor, desc, name in entries:
@@ -537,10 +538,10 @@ def execute(
         self._logger.debug("Entering execute")
         # Stage 1: runtime signature re-validation against the descriptors
         # captured at plan-build time, BEFORE any kernel launch. The plan is
-        # specialized to the sample dtypes and shapes. Kernel 1
-        # mutates ``attn_score`` / ``index_score`` in place before the GEMM
-        # runs, so reusing an exported plan with a mismatched signature would
-        # otherwise fail (or corrupt memory) only after the scores were
+        # specialized to the sample dtypes and shapes. Kernel 1 overwrites
+        # ``attn_score`` in place and treats ``index_score`` as read-only before
+        # the GEMM runs, so reusing an exported plan with a mismatched signature
+        # would otherwise fail (or corrupt memory) only after ``attn_score`` was
         # already mutated. Raise a clean ValueError here — no fail-dirty. The
         # signature-keyed wrapper cache never hits this, but a directly-
         # built/exported plan can.
@@ -563,8 +564,8 @@ def execute(
         )
         # One plan serves one device: the v2 backend's per-plan
         # workspace (the ticket counter) lives on the sample tensors'
-        # device, so reject cross-device execution before kernel 1 mutates
-        # the score buffers. The backend re-checks every
+        # device, so reject cross-device execution before kernel 1 overwrites
+        # ``attn_score``. The backend re-checks every
         # tensor against index_q.device, so validating index_q covers all.
         if self.use_v2 and index_q.device != self.iq_desc.device:
             raise ValueError(
@@ -785,9 +786,9 @@ def indexer_backward_wrapper(
 
     ``index_score`` contains the predict softmax aligned slot-for-slot with
     ``topk_indices``; the fused softmax returned by compressed indexer forward
-    can be passed directly. ``attn_score`` and ``index_score`` are consumed
-    in-place: the kernel overwrites ``attn_score`` with ``grad_signal`` and
-    ``index_score`` with ``sum_grad`` during the score-grad precompute stage.
+    can be passed directly. The kernel overwrites ``attn_score`` with
+    ``grad_signal`` during the score-grad precompute stage and treats
+    ``index_score`` as read-only.
 
     Args:
         topk_indices_global: whether ``topk_indices`` already contains global
@@ -801,7 +802,7 @@ def indexer_backward_wrapper(
         index_score: FP32 predict probabilities over the selected indexer
             logits. The caller must use the same valid-slot mask when
             constructing ``attn_score`` so target and predict describe
-            identical slots. This buffer is overwritten in-place.
+            identical slots. This buffer is preserved.
         grad_loss: single-element float32 tensor on the same CUDA device as
             ``index_q``. The kernel reads its value at runtime, including on
             CUDA Graph replay.
@@ -817,9 +818,9 @@ def indexer_backward_wrapper(
             outside its envelope — SM100, H == 64, D == 128, block_I == 128,
             topk % 128 == 0 with 128 <= topk <= 2048, sm_scale > 0, contiguous
             same-device tensors — and never silently falls back to the default
-            backend. It keeps the wrapper contract (same inputs, outputs, and
-            in-place score consumption: ``attn_score`` is left holding
-            exactly kernel 1's ``grad_signal``; ``sm_scale`` folds in-kernel)
+            backend. It keeps the wrapper contract (same inputs and outputs:
+            ``attn_score`` is left holding exactly kernel 1's ``grad_signal``,
+            ``index_score`` remains read-only, and ``sm_scale`` folds in-kernel)
             and computes the GEMM stage with fp32 weights and a two-term
             bf16 expansion of the gradient matrix, plus deterministic
             (bitwise run-to-run) ``d_weights`` / ``d_index_q``.