diff --git a/python/cudnn/deepseek_sparse_attention/sparse_attention_backward/dsa_bwd_sm100.py b/python/cudnn/deepseek_sparse_attention/sparse_attention_backward/dsa_bwd_sm100.py index 5d78f31f0..c9de516ec 100644 --- a/python/cudnn/deepseek_sparse_attention/sparse_attention_backward/dsa_bwd_sm100.py +++ b/python/cudnn/deepseek_sparse_attention/sparse_attention_backward/dsa_bwd_sm100.py @@ -788,6 +788,25 @@ def bwd( max_seqlen_q, max_seqlen_kv, head_dim, (num_heads, batch_size) = problem_shape + if cutlass.const_expr(mTopkLength is not None): + topk = mTopkLength[token_idx] + else: + topk = mTopkIdxs.shape[0] + + # topk is CTA-uniform (one value per query token). Handle an empty + # sparse row before initializing any async pipeline or allocating + # TMEM: the row contributes nothing to dKV and its dQ tile is zero. + # Treat malformed negative lengths as empty as well so they cannot + # enter the same zero-tile pipeline path. + if topk <= 0: + for linear_idx in cutlass.range(tidx, self.head_dim * self.block_tile, self.threads_per_cta): + head_offset = linear_idx // self.head_dim + dim_idx = linear_idx % self.head_dim + head_idx = head_block_idx * self.block_tile + head_offset + if head_idx < num_heads: + mdQ[dim_idx, head_idx, (token_idx, batch_idx)] = mdQ.element_type(0.0) + cute.arch.nvvm.exit() + if warp_idx == self.load_warp_id: cpasync.prefetch_descriptor(tma_atom_Q) cpasync.prefetch_descriptor(tma_atom_dO) @@ -881,11 +900,6 @@ def bwd( pipeline.pipeline_init_wait() - if cutlass.const_expr(mTopkLength is not None): - topk = mTopkLength[token_idx] - else: - topk = mTopkIdxs.shape[0] - tile_count = cute.ceil_div(topk, self.block_tile) if warp_idx == self.load_warp_id: diff --git a/test/python/fe_api/dsa/test_DSA_sparse_attention_backward.py b/test/python/fe_api/dsa/test_DSA_sparse_attention_backward.py index d0e6548e4..7f029c696 100644 --- a/test/python/fe_api/dsa/test_DSA_sparse_attention_backward.py +++ b/test/python/fe_api/dsa/test_DSA_sparse_attention_backward.py @@ -136,6 +136,118 @@ def test_DSA_sparse_attention_backward_wrapper( ) +@pytest.mark.L0 +@torch_fork_set_rng(seed=433) +@pytest.mark.parametrize( + "head_dim,num_heads,topk_length_values", + [ + pytest.param(512, 64, (-3, 0, 1, 63, 64, 65, 128), id="d512-mixed"), + pytest.param(576, 32, None, id="d576-all-empty"), + ], +) +def test_DSA_sparse_attention_backward_sm100_zero_topk_length(head_dim, num_heads, topk_length_values): + """SM100 must treat a nonpositive top-k length as an empty attention row.""" + if not torch.cuda.is_available(): + pytest.skip("SM100 GPU required") + major, minor = torch.cuda.get_device_capability() + if major * 10 + minor < 100: + pytest.skip("zero top-k length regression test targets the SM100 kernel") + + try: + from cudnn import DSA + except ImportError: + pytest.skip("Environment not supported: cudnn[cutedsl] not installed") + + device = torch.device("cuda") + s_q = len(topk_length_values) if topk_length_values is not None else 2 + s_kv, topk = 256, 128 + softmax_scale = 1.0 / math.sqrt(head_dim) + + q = torch.randn(s_q, num_heads, head_dim, dtype=torch.bfloat16, device=device) / 10 + kv = torch.randn(s_kv, head_dim, dtype=torch.bfloat16, device=device) / 10 + attn_sink = torch.randn(num_heads, dtype=torch.float32, device=device) + base_topk_idxs = torch.stack([torch.randperm(s_kv, device=device)[:topk] for _ in range(s_q)]).to(torch.int32) + + # The D512 mixed case covers defensive negative handling and both sides of + # the 64-row tile boundary. The all-empty cases make every expected + # gradient exactly zero; D576/H32 additionally covers the feature tail and + # a partial 64-head CTA. + topk_length_cases = [torch.zeros(s_q, dtype=torch.int32, device=device)] + if topk_length_values is not None: + topk_length_cases.insert(0, torch.tensor(topk_length_values, dtype=torch.int32, device=device)) + positions = torch.arange(topk, device=device).unsqueeze(0) + + for topk_length in topk_length_cases: + topk_idxs = base_topk_idxs.clone() + topk_idxs[positions >= topk_length.unsqueeze(1)] = -1 + empty_rows = topk_length <= 0 + # Invalid values in ignored slots exercise that the empty fast path + # does not consult topk_idxs before exiting. + topk_idxs[empty_rows] = torch.iinfo(torch.int32).max + + out, lse = ref_sparse_attention_forward( + q, + kv, + attn_sink, + topk_idxs, + topk_length=topk_length, + softmax_scale=softmax_scale, + ) + assert torch.equal(out[empty_rows], torch.zeros_like(out[empty_rows])) + assert torch.isneginf(lse[empty_rows]).all() + dout = torch.randn_like(out) + + # In particular, the empty-row fast path must overwrite caller-owned + # storage; torch.empty_like() could otherwise hide a missing dQ store. + dq_buffer = torch.full_like(q, float("nan")) + dkv_buffer = torch.full_like(kv, float("nan")) + result = DSA.sparse_attention_backward_wrapper( + q, + kv, + out, + dout, + lse, + attn_sink, + topk_idxs, + softmax_scale=softmax_scale, + topk_length=topk_length, + dq=dq_buffer, + dkv=dkv_buffer, + ) + # Reaching the synchronization is also the regression check for the + # empty-CTA barrier deadlock. + torch.cuda.synchronize() + + dq, dkv, d_sink = result["dq"], result["dkv"], result["d_sink"] + assert dq.data_ptr() == dq_buffer.data_ptr() + assert dkv.data_ptr() == dkv_buffer.data_ptr() + assert torch.isfinite(dq).all() + assert torch.isfinite(dkv).all() + assert torch.isfinite(d_sink).all() + assert torch.equal(dq[empty_rows], torch.zeros_like(dq[empty_rows])) + + check_ref_dsa_sparse_attention_backward( + q, + kv, + attn_sink, + topk_idxs, + out, + dout, + lse, + dq, + dkv, + d_sink, + softmax_scale=softmax_scale, + topk_length=topk_length, + atol=5e-2, + rtol=5e-2, + ) + + if torch.all(empty_rows): + assert torch.equal(dkv, torch.zeros_like(dkv)) + assert torch.equal(d_sink, torch.zeros_like(d_sink)) + + @pytest.mark.L0 @torch_fork_set_rng(seed=385) def test_DSA_sparse_attention_backward_qh32_uses_per_query_topk_without_padding(monkeypatch):