Fix DSA offset alignment, stream handling, and CUDA Graph capture

docs/fe-oss-apis/dsa.md

@@ -228,10 +228,11 @@ indexer tower:
 `RuntimeError` rather than silently falling back.
 
 ```python
+grad_loss = torch.ones((), dtype=torch.float32, device=index_q.device)
 result = DSA.indexer_backward_wrapper(
     index_q, weights, index_k,
     attn_score, index_score, topk_indices,
-    sm_scale=1.0, loss_coeff=1.0, grad_loss=1.0, block_I=128,
+    grad_loss=grad_loss, sm_scale=1.0, loss_coeff=1.0, block_I=128,
 )
 d_index_q, d_weights, d_index_k = (
     result["d_index_q"], result["d_weights"], result["d_index_k"],
@@ -255,6 +256,7 @@ and denominators produced by Dense Indexer / Dense Attn Score Recompute.
 - **Constraints** — SM90 or SM100+, `H >= 64`, `ratio >= 1`
 
 ```python
+grad_loss = torch.ones((), dtype=torch.float32, device=index_q.device)
 dense_index = DSA.dense_indexer_score_recompute_wrapper(
     index_q, index_k.unsqueeze(2), weights,
     q_causal_offsets=q_causal_offsets,
@@ -268,7 +270,7 @@ result = DSA.dense_indexer_backward_wrapper(
     index_q, weights, index_k,
     dense_attn["out"], dense_attn["denom"],
     dense_index["out"], dense_index["denom"],
-    sm_scale=1.0, loss_coeff=1.0, grad_loss=1.0, block_I=128, ratio=1,
+    grad_loss=grad_loss, sm_scale=1.0, loss_coeff=1.0, block_I=128, ratio=1,
     q_causal_offsets=q_causal_offsets,
 )
 ```

python/cudnn/deepseek_sparse_attention/indexer_backward/api.py

@@ -12,7 +12,7 @@
 
 from __future__ import annotations
 
-from typing import Optional, Union
+from typing import Optional
 
 import torch
 import cuda.bindings.driver as cuda
@@ -29,10 +29,12 @@
 from .indexer_backward_sm90 import indexer_backward_sm90
 
 
-def _as_grad_loss_tensor(grad_loss: Union[float, torch.Tensor], device: torch.device) -> torch.Tensor:
-    if torch.is_tensor(grad_loss):
-        return grad_loss.detach().to(device=device, dtype=torch.float32, copy=False).reshape(1)
-    return torch.tensor([float(grad_loss)], dtype=torch.float32, device=device)
+def _validate_grad_loss_tensor(grad_loss: torch.Tensor, device: torch.device) -> torch.Tensor:
+    if not torch.is_tensor(grad_loss):
+        raise TypeError("grad_loss must be a torch.Tensor")
+    if grad_loss.numel() != 1 or grad_loss.dtype != torch.float32 or grad_loss.device != device:
+        raise ValueError(f"grad_loss must be a single-element float32 tensor on {device}")
+    return grad_loss.detach().view(1)
 
 
 def _contiguous_input(tensor: torch.Tensor) -> torch.Tensor:
@@ -223,16 +225,16 @@ def execute(
         attn_score: torch.Tensor,
         index_score: torch.Tensor,
         topk_indices: torch.Tensor,
+        grad_loss: torch.Tensor,
         loss_coeff: float = 1.0,
-        grad_loss: Union[float, torch.Tensor] = 1.0,
         current_stream: Optional[cuda.CUstream] = None,
     ) -> None:
         self._logger.debug("Entering execute")
         # ``grad_scale`` is forwarded as a runtime ``Float32`` arg; the
         # compiled kernel is reused when loss_coeff changes for the same
         # cached tensor shape.
         grad_scale = float(loss_coeff) / (int(self.batch) * int(self.seqlen))
-        grad_loss_tensor = _as_grad_loss_tensor(grad_loss, index_q.device)
+        grad_loss_tensor = _validate_grad_loss_tensor(grad_loss, index_q.device)
         self._compiled_kernel(
             index_q,
             weights,
@@ -371,17 +373,17 @@ def execute(
         attn_l1norm: torch.Tensor,
         index_score: torch.Tensor,
         index_lse: torch.Tensor,
+        grad_loss: torch.Tensor,
         loss_coeff: float = 1.0,
-        grad_loss: Union[float, torch.Tensor] = 1.0,
         cu_seqlens_q: Optional[torch.Tensor] = None,
         cu_seqlens_k: Optional[torch.Tensor] = None,
         q_causal_offsets: Optional[torch.Tensor] = None,
         current_stream: Optional[cuda.CUstream] = None,
     ) -> None:
         backend_stream = None if self._uses_current_stream_pipeline else current_stream
         with _torch_stream_context(backend_stream):
-            grad_loss_value = float(_as_grad_loss_tensor(grad_loss, index_q.device).item())
-            grad_scale = float(loss_coeff) * grad_loss_value / max(int(self.normalization_tokens), 1)
+            grad_loss_tensor = _validate_grad_loss_tensor(grad_loss, index_q.device)
+            grad_scale = float(loss_coeff) / max(int(self.normalization_tokens), 1)
 
             # Dense backward's dK path uses atomic/bulk reductions into fp32.
             d_index_k_target = d_index_k
@@ -402,6 +404,7 @@ def execute(
             attn_l1norm,
             index_score,
             index_lse,
+            grad_loss_tensor,
             grad_scale,
             cu_seqlens_q,
             cu_seqlens_k,
@@ -424,9 +427,9 @@ def indexer_backward_wrapper(
     attn_score: torch.Tensor,
     index_score: torch.Tensor,
     topk_indices: torch.Tensor,
+    grad_loss: torch.Tensor,
     sm_scale: float = 1.0,
     loss_coeff: float = 1.0,
-    grad_loss: Union[float, torch.Tensor] = 1.0,
     block_I: int = 128,
     topk_indices_global: bool = False,
     d_index_q: Optional[torch.Tensor] = None,
@@ -449,9 +452,9 @@ def indexer_backward_wrapper(
             weights-scaling trick.
         loss_coeff: coefficient scaling the KL-divergence loss in the
             forward (``indexer_loss = loss_coeff * kl.mean()``).
-        grad_loss: scalar gradient of the outer loss w.r.t. the KL term;
-            typically ``1.0`` when the KL loss is summed into the total
-            training loss with unit weight. Accepts a 0-D tensor or float.
+        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.
     """
     with _torch_stream_context(stream):
         if d_index_q is None:
@@ -532,9 +535,9 @@ def dense_indexer_backward_wrapper(
     attn_l1norm: torch.Tensor,
     index_score: torch.Tensor,
     index_lse: torch.Tensor,
+    grad_loss: torch.Tensor,
     sm_scale: float = 1.0,
     loss_coeff: float = 1.0,
-    grad_loss: Union[float, torch.Tensor] = 1.0,
     block_I: int = 128,
     ratio: int = 1,
     cu_seqlens_q: Optional[torch.Tensor] = None,
@@ -553,6 +556,10 @@ def dense_indexer_backward_wrapper(
     ``dense_attn_score_recompute_wrapper`` and
     ``dense_indexer_score_recompute_wrapper`` respectively. They are consumed
     in-place by the score-gradient precompute stage.
+
+    ``grad_loss`` must be a single-element float32 tensor on the same CUDA
+    device as ``index_q``. The kernel reads its value at runtime, including on
+    CUDA Graph replay.
     """
     major, _ = torch.cuda.get_device_capability()
     backend_stream = None if major == 9 else stream

python/cudnn/deepseek_sparse_attention/indexer_backward/dense_indexer_backward_sm100.py

@@ -147,6 +147,7 @@ def __call__(
         mCuSeqlensQ,
         mCuSeqlensK,
         mQCausalOffsets,
+        mGradLoss,
         grad_scale: Float32,
         max_seqlen_q: Int32,
         max_seqlen_k: Int32,
@@ -192,6 +193,7 @@ def __call__(
             mCuSeqlensQ,
             mCuSeqlensK,
             mQCausalOffsets,
+            mGradLoss,
             grad_scale,
             seqlen_k_pad,
             max_seqlen_q,
@@ -212,6 +214,7 @@ def kernel_score_grad(
         mCuSeqlensQ,
         mCuSeqlensK,
         mQCausalOffsets,
+        mGradLoss,
         grad_scale: Float32,
         seqlen_k_pad: Int32,
         seqlen_q_static: Int32,
@@ -233,6 +236,7 @@ def kernel_score_grad(
             seqlen_k_static,
         )
         q_causal_offset_b = Int32(0) if const_expr(mQCausalOffsets is None) else mQCausalOffsets[batch_idx]
+        grad_scale_f32 = Float32(grad_scale) * Float32(mGradLoss[0])
 
         # Out-of-range CTAs (seq_local >= seqlen_q_b in this batch): skip work.
         # CuTe DSL forbids early `return`, so wrap the body in an `if` block.
@@ -284,7 +288,7 @@ class SharedStorage:
                 target = tr / (denom_val + Float32(DENOM_EPS))
                 target_eff = target if target >= Float32(CLIP_PROB_MIN) else Float32(CLIP_PROB_MIN)
                 log_clip_mask = Float32(1.0) if score_minus_lse >= Float32(CLIP_LOG_MIN) else Float32(0.0)
-                g = -target_eff * log_clip_mask * grad_scale
+                g = -target_eff * log_clip_mask * grad_scale_f32
 
                 local_sum = local_sum + g
                 pos = pos + Int32(128)
@@ -324,7 +328,7 @@ class SharedStorage:
                     target = tr / (denom_val + Float32(DENOM_EPS))
                     target_eff = target if target >= Float32(CLIP_PROB_MIN) else Float32(CLIP_PROB_MIN)
                     log_clip_mask = Float32(1.0) if score_minus_lse >= Float32(CLIP_LOG_MIN) else Float32(0.0)
-                    g = -target_eff * log_clip_mask * grad_scale
+                    g = -target_eff * log_clip_mask * grad_scale_f32
 
                     grad_signal = g - predict * sum_grad
                     mPredict_b[seq_local, pos] = grad_signal
@@ -2007,6 +2011,7 @@ def _ensure_compiled_score_grad(
         CuSeqlensQ,
         CuSeqlensK,
         QCausalOffsets,
+        GradLoss,
         grad_scale,
         current_stream=None,
     ):
@@ -2037,6 +2042,7 @@ def _ensure_compiled_score_grad(
                 cuq_arg,
                 cuk_arg,
                 q_offsets_arg,
+                to_cute_tensor(GradLoss, assumed_align=4, leading_dim=0),
                 cutlass.Float32(float(grad_scale)),
                 cutlass.Int32(max_seqlen_q),
                 cutlass.Int32(max_seqlen_k),
@@ -2153,6 +2159,7 @@ def _run(
         AttnL1Norm,
         IdxScoreRaw,
         IdxLSE,
+        GradLoss,
         grad_scale,
         CuSeqlensQ=None,
         CuSeqlensK=None,
@@ -2165,6 +2172,9 @@ def _run(
         THD : pass packed tensors and CuSeqlens* (B+1,) int32.
         Kernel 3 (dK f32 → output dtype) is handled by the caller.
 
+        ``GradLoss`` is a single-element float32 tensor on the same CUDA device
+        as the inputs. Kernel 1 reads it at runtime, including during CUDA Graph
+        replay.
         ``grad_scale`` is a host scalar (Python float, ``loss_coeff /
         (b*sq)``); supplied at call time as a runtime ``Float32`` arg to
         kernel 1 so changing it does not require recompilation.
@@ -2177,6 +2187,8 @@ def _run(
             assert QCausalOffsets is not None, "offset-compiled kernel requires q_causal_offsets at runtime"
         else:
             assert QCausalOffsets is None, "non-offset compiled kernel must not receive q_causal_offsets"
+        if GradLoss.ndim == 0:
+            GradLoss = GradLoss.reshape(1)
         s = _resolve_stream(current_stream)
 
         # Kernel 1 (CuTe DSL): in-place overwrites IdxScoreRaw with grad_signal.
@@ -2189,6 +2201,7 @@ def _run(
             CuSeqlensQ,
             CuSeqlensK,
             QCausalOffsets,
+            GradLoss,
             grad_scale,
             current_stream=current_stream,
         )
@@ -2201,6 +2214,7 @@ def _run(
                 CuSeqlensQ,
                 CuSeqlensK,
                 QCausalOffsets,
+                GradLoss,
                 cutlass.Float32(float(grad_scale)),
                 cutlass.Int32(max_seqlen_q),
                 cutlass.Int32(max_seqlen_k),

python/cudnn/deepseek_sparse_attention/indexer_backward/dense_indexer_backward_sm90.py

@@ -67,6 +67,7 @@ def __call__(
         mCuSeqlensQ,
         mCuSeqlensK,
         mQCausalOffsets,
+        mGradLoss,
         grad_scale: Float32,
         max_seqlen_q: Int32,
         max_seqlen_k: Int32,
@@ -109,6 +110,7 @@ def __call__(
             mCuSeqlensQ,
             mCuSeqlensK,
             mQCausalOffsets,
+            mGradLoss,
             grad_scale,
             seqlen_k_pad,
             max_seqlen_q,
@@ -129,6 +131,7 @@ def kernel_score_grad(
         mCuSeqlensQ,
         mCuSeqlensK,
         mQCausalOffsets,
+        mGradLoss,
         grad_scale: Float32,
         seqlen_k_pad: Int32,
         seqlen_q_static: Int32,
@@ -148,6 +151,7 @@ def kernel_score_grad(
             seqlen_k_static,
         )
         q_causal_offset_b = Int32(0) if const_expr(mQCausalOffsets is None) else mQCausalOffsets[batch_idx]
+        grad_scale_f32 = Float32(grad_scale) * Float32(mGradLoss[0])
 
         if seq_local < seqlen_q_b:
             smem = cutlass.utils.SmemAllocator()
@@ -195,7 +199,7 @@ class SharedStorage:
                 target = attn_raw / (attn_l1_val + Float32(EPS))
                 target_eff = target if target >= Float32(CLIP_PROB_MIN) else Float32(CLIP_PROB_MIN)
                 log_clip_mask = Float32(1.0) if score_minus_lse >= Float32(CLIP_LOG_MIN) else Float32(0.0)
-                local_sum = local_sum + (-target_eff * log_clip_mask * grad_scale)
+                local_sum = local_sum + (-target_eff * log_clip_mask * grad_scale_f32)
                 pos = pos + Int32(128)
 
             warp_sum = cute.arch.warp_reduction_sum(local_sum)
@@ -221,7 +225,7 @@ class SharedStorage:
                     target = attn_raw / (attn_l1_val + Float32(EPS))
                     target_eff = target if target >= Float32(CLIP_PROB_MIN) else Float32(CLIP_PROB_MIN)
                     log_clip_mask = Float32(1.0) if score_minus_lse >= Float32(CLIP_LOG_MIN) else Float32(0.0)
-                    g = -target_eff * log_clip_mask * grad_scale
+                    g = -target_eff * log_clip_mask * grad_scale_f32
                     mIdxScore_b[seq_local, pos] = g - predict * sum_grad
                 else:
                     mIdxScore_b[seq_local, pos] = Float32(0.0)
@@ -352,6 +356,7 @@ def _ensure_compiled_score_grad(
         CuSeqlensQ,
         CuSeqlensK,
         QCausalOffsets,
+        GradLoss,
         grad_scale,
         current_stream=None,
     ):
@@ -369,6 +374,7 @@ def _ensure_compiled_score_grad(
                 cuq_arg,
                 cuk_arg,
                 q_offsets_arg,
+                to_cute_tensor(GradLoss, assumed_align=4, leading_dim=0),
                 cutlass.Float32(float(grad_scale)),
                 cutlass.Int32(seqlen),
                 cutlass.Int32(seqlen_k),
@@ -381,12 +387,15 @@ def _run_score_grad_only(
         IdxLSE,
         AttnScoreRaw,
         AttnL1Norm,
+        GradLoss,
         grad_scale,
         CuSeqlensQ=None,
         CuSeqlensK=None,
         QCausalOffsets=None,
         current_stream=None,
     ):
+        if GradLoss.ndim == 0:
+            GradLoss = GradLoss.reshape(1)
         if is_varlen:
             assert CuSeqlensQ is not None and CuSeqlensK is not None, "THD-compiled score-grad kernel requires cu_seqlens_q/k at runtime"
         else:
@@ -404,6 +413,7 @@ def _run_score_grad_only(
             CuSeqlensQ,
             CuSeqlensK,
             QCausalOffsets,
+            GradLoss,
             grad_scale,
             current_stream=current_stream,
         )
@@ -415,6 +425,7 @@ def _run_score_grad_only(
             CuSeqlensQ,
             CuSeqlensK,
             QCausalOffsets,
+            GradLoss,
             cutlass.Float32(float(grad_scale)),
             cutlass.Int32(seqlen),
             cutlass.Int32(seqlen_k),
@@ -488,6 +499,7 @@ def _run(
         attn_l1norm,
         idx_scores_raw,
         idx_lse,
+        GradLoss,
         grad_scale,
         CuSeqlensQ=None,
         CuSeqlensK=None,
@@ -503,6 +515,7 @@ def _run(
             idx_lse,
             attn_scores_raw,
             attn_l1norm,
+            GradLoss,
             grad_scale,
             CuSeqlensQ,
             CuSeqlensK,

python/cudnn/deepseek_sparse_attention/score_recompute/_interface_sm90.py

@@ -383,7 +383,7 @@ def _dense_score_recompute(
         out_cute = to_cute_tensor(out)
         weights_cute = to_cute_tensor(weights_or_lse)
         denom_cute = to_cute_tensor(denom_out)
-        q_offsets_cute = to_cute_tensor(q_causal_offsets, leading_dim=0) if q_causal_offsets is not None else None
+        q_offsets_cute = to_cute_tensor(q_causal_offsets, assumed_align=4, leading_dim=0) if q_causal_offsets is not None else None
 
         kernel_obj = DenseScoreRecomputeSm90(
             dtype,

python/cudnn/deepseek_sparse_attention/utils/runtime.py

@@ -56,5 +56,5 @@ def torch_stream_context(current_stream: Optional[cuda.CUstream] = None) -> Iter
     if current_stream is None:
         yield
         return
-    with torch.cuda.stream(torch.cuda.ExternalStream(int(current_stream))):
+    with torch.cuda.stream(torch.cuda.get_stream_from_external(int(current_stream))):
         yield