From fcb380c942d8a75d993eeba01e733cd831b2a8bd Mon Sep 17 00:00:00 2001 From: kunlunl Date: Tue, 7 Jul 2026 02:32:15 -0700 Subject: [PATCH 1/5] Fix SM90 query offset alignment Signed-off-by: kunlunl --- .../score_recompute/_interface_sm90.py | 2 +- .../dsa/test_DSA_dense_score_recompute.py | 46 +++++++++++++++++++ 2 files changed, 47 insertions(+), 1 deletion(-) diff --git a/python/cudnn/deepseek_sparse_attention/score_recompute/_interface_sm90.py b/python/cudnn/deepseek_sparse_attention/score_recompute/_interface_sm90.py index bd86382af1..29fd6d94cc 100644 --- a/python/cudnn/deepseek_sparse_attention/score_recompute/_interface_sm90.py +++ b/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, diff --git a/test/python/fe_api/dsa/test_DSA_dense_score_recompute.py b/test/python/fe_api/dsa/test_DSA_dense_score_recompute.py index 55aace4647..90d1537934 100644 --- a/test/python/fe_api/dsa/test_DSA_dense_score_recompute.py +++ b/test/python/fe_api/dsa/test_DSA_dense_score_recompute.py @@ -85,6 +85,52 @@ def _allocate(cfg, score_type: str): return q, k, lse +@pytest.mark.L0 +@torch_fork_set_rng(seed=0) +def test_DSA_dense_score_recompute_thd_q_causal_offsets_alignment(): + try: + from cudnn import DSA + from cuda.bindings import driver as cuda + except ImportError: + pytest.skip("Environment not supported: cudnn[cutedsl] not installed") + + if torch.cuda.get_device_capability()[0] != 9: + pytest.skip("This regression is specific to the SM90 THD adapter") + + batch, seqlen_q, seqlen_k, heads, head_dim = 2, 256, 1024, 32, 128 + # Use a non-default stream so this test exercises only the offset contract. + torch_stream = torch.cuda.Stream() + with torch.cuda.stream(torch_stream): + q = torch.randn(batch, seqlen_q, heads, head_dim, dtype=torch.bfloat16, device="cuda") + k = torch.randn(batch, seqlen_k, 1, head_dim, dtype=torch.bfloat16, device="cuda") + weights = torch.randn(batch, seqlen_q, heads, dtype=torch.bfloat16, device="cuda") + cu_seqlens_q = torch.arange(0, (batch + 1) * seqlen_q, seqlen_q, dtype=torch.int32, device="cuda") + cu_seqlens_k = torch.arange(0, (batch + 1) * seqlen_k, seqlen_k, dtype=torch.int32, device="cuda") + q_causal_offsets = cu_seqlens_q[:-1] + + assert q_causal_offsets[1:].data_ptr() % 4 == 0 + assert q_causal_offsets[1:].data_ptr() % 16 != 0 + + result = DSA.dense_indexer_score_recompute_wrapper( + q.flatten(0, 1), + k.flatten(0, 1), + weights.flatten(0, 1), + qhead_per_kv_head=heads, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + max_seqlen_q=seqlen_q, + max_seqlen_k=seqlen_k, + q_causal_offsets=q_causal_offsets, + stream=cuda.CUstream(torch_stream.cuda_stream), + ) + + torch_stream.synchronize() + assert result["out"].shape == (batch * seqlen_q, seqlen_k) + assert result["denom"].shape == (batch * seqlen_q,) + assert torch.isfinite(result["out"]).any() + assert torch.isfinite(result["denom"]).all() + + @pytest.mark.L0 @torch_fork_set_rng(seed=0) @with_dsa_score_recompute_params From 908b427cfa031781df23c631ea50d255821ead28 Mon Sep 17 00:00:00 2001 From: kunlunl Date: Tue, 7 Jul 2026 02:32:15 -0700 Subject: [PATCH 2/5] Preserve the CUDA default stream Signed-off-by: kunlunl --- .../deepseek_sparse_attention/utils/runtime.py | 2 +- test/python/fe_api/dsa/test_DSA_runtime.py | 17 +++++++++++++++++ 2 files changed, 18 insertions(+), 1 deletion(-) create mode 100644 test/python/fe_api/dsa/test_DSA_runtime.py diff --git a/python/cudnn/deepseek_sparse_attention/utils/runtime.py b/python/cudnn/deepseek_sparse_attention/utils/runtime.py index d9984a258f..145fba12cc 100644 --- a/python/cudnn/deepseek_sparse_attention/utils/runtime.py +++ b/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 diff --git a/test/python/fe_api/dsa/test_DSA_runtime.py b/test/python/fe_api/dsa/test_DSA_runtime.py new file mode 100644 index 0000000000..1b4611f219 --- /dev/null +++ b/test/python/fe_api/dsa/test_DSA_runtime.py @@ -0,0 +1,17 @@ +import pytest +import torch + + +@pytest.mark.L0 +def test_torch_stream_context_preserves_legacy_default_stream(): + try: + from cuda.bindings import driver as cuda + from cudnn.deepseek_sparse_attention.utils.runtime import torch_stream_context + except ImportError: + pytest.skip("Environment not supported: cudnn[cutedsl] not installed") + + outer_stream = torch.cuda.Stream() + with torch.cuda.stream(outer_stream): + with torch_stream_context(cuda.CUstream(0)): + assert torch.cuda.current_stream().cuda_stream == 0 + assert torch.cuda.current_stream().cuda_stream == outer_stream.cuda_stream From 7cb157055d6011ed26a259db30b2a5c572ee1892 Mon Sep 17 00:00:00 2001 From: kunlunl Date: Tue, 7 Jul 2026 03:08:16 -0700 Subject: [PATCH 3/5] Make dense indexer backward graph safe Signed-off-by: kunlunl --- .../indexer_backward/api.py | 5 +- .../dense_indexer_backward_sm100.py | 15 +++- .../dense_indexer_backward_sm90.py | 17 +++- .../dsa/test_DSA_dense_indexer_backward.py | 80 +++++++++++++++++++ 4 files changed, 111 insertions(+), 6 deletions(-) diff --git a/python/cudnn/deepseek_sparse_attention/indexer_backward/api.py b/python/cudnn/deepseek_sparse_attention/indexer_backward/api.py index 023c0e20ae..d000fea02a 100644 --- a/python/cudnn/deepseek_sparse_attention/indexer_backward/api.py +++ b/python/cudnn/deepseek_sparse_attention/indexer_backward/api.py @@ -380,8 +380,8 @@ def execute( ) -> 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 = _as_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 +402,7 @@ def execute( attn_l1norm, index_score, index_lse, + grad_loss_tensor, grad_scale, cu_seqlens_q, cu_seqlens_k, diff --git a/python/cudnn/deepseek_sparse_attention/indexer_backward/dense_indexer_backward_sm100.py b/python/cudnn/deepseek_sparse_attention/indexer_backward/dense_indexer_backward_sm100.py index 7e81122115..b6c79b215c 100644 --- a/python/cudnn/deepseek_sparse_attention/indexer_backward/dense_indexer_backward_sm100.py +++ b/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, @@ -2177,6 +2184,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 +2198,7 @@ def _run( CuSeqlensQ, CuSeqlensK, QCausalOffsets, + GradLoss, grad_scale, current_stream=current_stream, ) @@ -2201,6 +2211,7 @@ def _run( CuSeqlensQ, CuSeqlensK, QCausalOffsets, + GradLoss, cutlass.Float32(float(grad_scale)), cutlass.Int32(max_seqlen_q), cutlass.Int32(max_seqlen_k), diff --git a/python/cudnn/deepseek_sparse_attention/indexer_backward/dense_indexer_backward_sm90.py b/python/cudnn/deepseek_sparse_attention/indexer_backward/dense_indexer_backward_sm90.py index 5e7a186962..83c80e71d7 100644 --- a/python/cudnn/deepseek_sparse_attention/indexer_backward/dense_indexer_backward_sm90.py +++ b/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, diff --git a/test/python/fe_api/dsa/test_DSA_dense_indexer_backward.py b/test/python/fe_api/dsa/test_DSA_dense_indexer_backward.py index f8325a73fb..08010a746a 100644 --- a/test/python/fe_api/dsa/test_DSA_dense_indexer_backward.py +++ b/test/python/fe_api/dsa/test_DSA_dense_indexer_backward.py @@ -154,3 +154,83 @@ def test_DSA_dense_indexer_backward_wrapper( grad_scale=grad_scale_expected, q_causal_offsets=q_causal_offsets, ) + + +@pytest.mark.L0 +@torch_fork_set_rng(seed=0) +def test_DSA_dense_indexer_backward_cuda_graph(): + if torch.cuda.get_device_capability()[0] < 9: + pytest.skip("Dense indexer backward requires SM90+") + + try: + from cudnn import DSA + except ImportError: + pytest.skip("Environment not supported: cudnn[cutedsl] not installed") + + cfg = { + "b": 1, + "s_q": 128, + "s_kv": 128, + "head_dim": 128, + "qhead_per_kv_head": 64, + } + ( + index_q, + weights, + index_k, + attn_score, + attn_l1norm, + index_score, + index_lse, + ) = _allocate(cfg, sm_scale=1.0, ratio=1, q_causal_offsets=None) + attn_score_source = attn_score.clone() + index_score_source = index_score.clone() + d_index_q = torch.empty_like(index_q) + d_weights = torch.empty_like(weights) + d_index_k = torch.empty_like(index_k, dtype=torch.float32) + grad_loss = torch.ones(1, dtype=torch.float32, device="cuda") + + def run(): + attn_score.copy_(attn_score_source) + index_score.copy_(index_score_source) + DSA.dense_indexer_backward_wrapper( + index_q, + weights, + index_k, + attn_score, + attn_l1norm, + index_score, + index_lse, + loss_coeff=float(cfg["b"] * cfg["s_q"]), + grad_loss=grad_loss, + block_I=128, + ratio=1, + d_index_q=d_index_q, + d_weights=d_weights, + d_index_k=d_index_k, + ) + + warmup_stream = torch.cuda.Stream() + warmup_stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(warmup_stream): + run() + torch.cuda.current_stream().wait_stream(warmup_stream) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run() + + for scale in (0.5, 1.5): + grad_loss.fill_(scale) + graph.replay() + check_ref_dense_indexer_backward( + index_q, + weights, + index_k, + attn_score_source, + attn_l1norm, + d_index_q, + d_weights, + d_index_k, + grad_scale=scale, + ) From 1d75788f3df47e4ed3c727e921d670b485b26f15 Mon Sep 17 00:00:00 2001 From: kunlunl Date: Wed, 8 Jul 2026 00:55:59 -0700 Subject: [PATCH 4/5] Require a tensor grad_loss for indexer backward Signed-off-by: kunlunl --- docs/fe-oss-apis/dsa.md | 6 ++-- .../indexer_backward/api.py | 34 +++++++++++-------- .../dsa/test_DSA_dense_indexer_backward.py | 4 +-- .../fe_api/dsa/test_DSA_indexer_backward.py | 4 +-- 4 files changed, 28 insertions(+), 20 deletions(-) diff --git a/docs/fe-oss-apis/dsa.md b/docs/fe-oss-apis/dsa.md index 9ec5f3400b..6791f0537e 100644 --- a/docs/fe-oss-apis/dsa.md +++ b/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, ) ``` diff --git a/python/cudnn/deepseek_sparse_attention/indexer_backward/api.py b/python/cudnn/deepseek_sparse_attention/indexer_backward/api.py index d000fea02a..3fa6c2e3cd 100644 --- a/python/cudnn/deepseek_sparse_attention/indexer_backward/api.py +++ b/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,8 +225,8 @@ 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") @@ -232,7 +234,7 @@ def execute( # 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,8 +373,8 @@ 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, @@ -380,7 +382,7 @@ def execute( ) -> None: backend_stream = None if self._uses_current_stream_pipeline else current_stream with _torch_stream_context(backend_stream): - grad_loss_tensor = _as_grad_loss_tensor(grad_loss, index_q.device) + 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. @@ -425,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, @@ -450,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: @@ -533,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, @@ -554,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 diff --git a/test/python/fe_api/dsa/test_DSA_dense_indexer_backward.py b/test/python/fe_api/dsa/test_DSA_dense_indexer_backward.py index 08010a746a..2976552707 100644 --- a/test/python/fe_api/dsa/test_DSA_dense_indexer_backward.py +++ b/test/python/fe_api/dsa/test_DSA_dense_indexer_backward.py @@ -89,8 +89,8 @@ def test_DSA_dense_indexer_backward_wrapper( s_q_cfg = cfg["s_q"] q_causal_offsets = torch.full((b_cfg,), 8, dtype=torch.int32, device="cuda") loss_coeff = float(b_cfg * s_q_cfg) - grad_loss = 1.0 - grad_scale_expected = (loss_coeff / (b_cfg * s_q_cfg)) * grad_loss + grad_loss = torch.ones((), dtype=torch.float32, device="cuda") + grad_scale_expected = loss_coeff / (b_cfg * s_q_cfg) ( index_q, diff --git a/test/python/fe_api/dsa/test_DSA_indexer_backward.py b/test/python/fe_api/dsa/test_DSA_indexer_backward.py index 722ecd16f3..0f41283306 100644 --- a/test/python/fe_api/dsa/test_DSA_indexer_backward.py +++ b/test/python/fe_api/dsa/test_DSA_indexer_backward.py @@ -91,8 +91,8 @@ def test_DSA_indexer_backward_wrapper( b_cfg = cfg["b"] s_q_cfg = cfg["s_q"] loss_coeff = float(b_cfg * s_q_cfg) - grad_loss = 1.0 - grad_scale_expected = (loss_coeff / (b_cfg * s_q_cfg)) * grad_loss # = 1.0 + grad_loss = torch.ones((), dtype=torch.float32, device="cuda") + grad_scale_expected = loss_coeff / (b_cfg * s_q_cfg) # = 1.0 ( index_q, From facd8b0faf9d0709bd9dfed10705355f50c1257d Mon Sep 17 00:00:00 2001 From: kunlunl Date: Wed, 8 Jul 2026 01:15:37 -0700 Subject: [PATCH 5/5] Address remaining DSA review comments Signed-off-by: kunlunl --- .../indexer_backward/dense_indexer_backward_sm100.py | 3 +++ test/python/fe_api/dsa/test_DSA_dense_score_recompute.py | 1 + 2 files changed, 4 insertions(+) diff --git a/python/cudnn/deepseek_sparse_attention/indexer_backward/dense_indexer_backward_sm100.py b/python/cudnn/deepseek_sparse_attention/indexer_backward/dense_indexer_backward_sm100.py index b6c79b215c..9b12433fd4 100644 --- a/python/cudnn/deepseek_sparse_attention/indexer_backward/dense_indexer_backward_sm100.py +++ b/python/cudnn/deepseek_sparse_attention/indexer_backward/dense_indexer_backward_sm100.py @@ -2172,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. diff --git a/test/python/fe_api/dsa/test_DSA_dense_score_recompute.py b/test/python/fe_api/dsa/test_DSA_dense_score_recompute.py index 90d1537934..14d7512b4d 100644 --- a/test/python/fe_api/dsa/test_DSA_dense_score_recompute.py +++ b/test/python/fe_api/dsa/test_DSA_dense_score_recompute.py @@ -128,6 +128,7 @@ def test_DSA_dense_score_recompute_thd_q_causal_offsets_alignment(): assert result["out"].shape == (batch * seqlen_q, seqlen_k) assert result["denom"].shape == (batch * seqlen_q,) assert torch.isfinite(result["out"]).any() + assert (torch.isfinite(result["out"]) | torch.isneginf(result["out"])).all() assert torch.isfinite(result["denom"]).all()