DSA backward SM100: reject topk_length entries < 1 (a zero entry hangs the kernel)

python/cudnn/deepseek_sparse_attention/sparse_attention_backward/_interface_sm100.py

@@ -8,7 +8,7 @@
 import cutlass.cute as cute
 
 from cudnn.deepseek_sparse_attention.utils.compiler import compile_options
-from cudnn.deepseek_sparse_attention.utils.runtime import resolve_stream
+from cudnn.deepseek_sparse_attention.utils.runtime import resolve_stream, torch_stream_context
 from cudnn.deepseek_sparse_attention.utils.tensor_conversion import to_cute_tensor
 from .dsa_bwd_sm100 import FlashAttentionDSABackwardSm100
 
@@ -47,12 +47,29 @@ def flash_attn_bwd_sm100(
         attn_sink: (nheads,) float32
         topk_idxs: (total_S_q, topk_max) int32, global indices
         softmax_scale: float (default: 1/sqrt(headdim))
-        topk_length: (total_S_q,) int32, per-query valid count, optional
+        topk_length: (total_S_q,) int32, per-query valid count, optional.
+            Every entry must be >= 1 (see Raises).
         dq: pre-allocated (total_S_q, nheads, headdim), optional
         dkv: pre-allocated (total_S_kv, headdim), optional
 
     Returns:
         (dq, dkv, d_sink) -- flat layout gradients
+
+    Raises:
+        ValueError: if any entry of ``topk_length`` is < 1 (validated outside
+            CUDA graph capture only -- capture is detected on the resolved
+            launch stream, i.e. the explicit ``current_stream`` when one is
+            given, and the required device-to-host sync runs on that stream;
+            during capture the sync is not possible, and the caller must
+            uphold the precondition). A row with ``topk_length == 0`` gives the
+            kernel zero KV tiles to process, for which it has no defined
+            behavior: on the configuration we tested (B200, head_dim 512,
+            topk 512) the kernel never terminates -- the launch hangs at
+            100% SM until killed. Code inspection suggests a zero-tile
+            pipeline deadlock, and that a zero-tile token reaching the dq
+            epilogue would store accumulator TMEM no MMA has written. Drop
+            empty rows from the batch (or prevent them upstream of this
+            call).
     """
     total_S_q, num_head, head_dim = q.shape
     total_S_kv = kv.shape[0]
@@ -69,6 +86,31 @@ def flash_attn_bwd_sm100(
         tensors_to_check.append(topk_length)
     assert all(t.is_cuda for t in tensors_to_check)
 
+    current_stream = resolve_stream(current_stream)
+
+    if topk_length is not None:
+        # topk_length == 0 makes tile_count == 0 in the kernel, which has no
+        # defined behavior: measured on B200 (head_dim 512, topk 512) the
+        # launch hangs at 100% SM until killed. Code inspection suggests a
+        # zero-tile pipeline deadlock, and that a zero-tile token reaching
+        # the dq epilogue would store accumulator TMEM no MMA has written.
+        # Fail loudly here instead. The guard binds to the resolved launch
+        # stream: capture detection and the device-to-host sync both run on
+        # ``current_stream`` (not on whatever torch stream happens to be
+        # current), so an explicit-stream capture is skipped correctly and
+        # the sync is ordered after the caller's producer work on that
+        # stream. Skipped under CUDA graph capture, where the required
+        # device-to-host sync is illegal -- the precondition itself still
+        # applies (documented above).
+        with torch_stream_context(current_stream):
+            if not torch.cuda.is_current_stream_capturing() and bool((topk_length < 1).any()):
+                raise ValueError(
+                    "topk_length must be >= 1 for every query: a row with topk_length == 0 "
+                    "gives the SM100 backward kernel zero KV tiles, for which it has no "
+                    "defined behavior (observed to hang the launch). Drop empty rows from "
+                    "the batch or prevent them upstream."
+                )
+
     if softmax_scale is None:
         softmax_scale = 1.0 / math.sqrt(head_dim)
 
@@ -127,7 +169,6 @@ def flash_attn_bwd_sm100(
     problem_shape = (total_S_q, total_S_kv, head_dim, (num_head, batch_size))
 
     dtype = torch2cute_dtype_map[q.dtype]
-    current_stream = resolve_stream(current_stream)
 
     has_topk_length = topk_length is not None
     max_topk = topk_idxs.shape[1]

python/cudnn/deepseek_sparse_attention/sparse_attention_backward/api.py

@@ -152,6 +152,14 @@ def sparse_attention_backward_wrapper(
 
     Dispatches to SM90 or SM100 based on the active CUDA device. The returned
     ``d_sink`` is computed from ``attn_sink`` and ``dout``.
+
+    ``topk_length``, when provided, must be >= 1 for every query. The SM100
+    path validates this on the launch stream (the explicit ``stream`` when
+    one is given) outside CUDA graph capture and raises ``ValueError``
+    (under capture the required sync is not possible and the caller must
+    uphold the precondition): a row with ``topk_length == 0`` has no defined
+    kernel behavior -- on the configuration we tested the launch hangs until
+    killed (see ``_interface_sm100.flash_attn_bwd_sm100``).
     """
     key = (
         q.dtype,

test/python/fe_api/dsa/test_DSA_sparse_attention_backward.py

@@ -302,3 +302,139 @@ def rel_l2(a, b):
 
     assert rel_l2(dkv, dkv_ref) < 1e-4, "dkv parity vs stage-1 baseline"
     assert rel_l2(d_sink, d_sink_ref) < 1e-4, "d_sink parity vs stage-1 baseline"
+
+
+@pytest.mark.L0
+@torch_fork_set_rng(seed=0)
+def test_DSA_sparse_attention_backward_topk_length_zero_raises():
+    """Rows with topk_length == 0 must be rejected loudly on SM100.
+
+    A zero entry gives the kernel zero KV tiles for that token
+    (tile_count == 0 in dsa_bwd_sm100), which has no defined behavior:
+    measured at s_q=128, s_kv=4096, head_dim=512, topk=512 on B200, a call
+    whose topk_length contains zeros hangs until killed. The interface
+    validates the precondition (outside CUDA graph capture) and raises
+    before launch; this test exercises that guard path.
+
+    L0 is intentional: the guard raises before the interface consults its
+    compile cache, and ``SparseAttentionBackward.compile()`` defers kernel
+    compilation to ``execute()`` (see api.py), so this test pays no CuTe
+    compilation cost even on a cold cache.
+    """
+    try:
+        from cudnn import DSA
+    except ImportError:
+        pytest.skip("Environment not supported: cudnn[cutedsl] not installed")
+
+    major, minor = torch.cuda.get_device_capability()
+    if major * 10 + minor < 100:
+        pytest.skip("the topk_length >= 1 guard is implemented on the SM100 interface")
+
+    s_q, s_kv, h, d, d_v, topk = 8, 256, 64, 576, 512, 64
+    device = "cuda"
+    q = torch.randn(s_q, h, d, dtype=torch.bfloat16, device=device)
+    kv = torch.randn(s_kv, d, dtype=torch.bfloat16, device=device)
+    out = torch.randn(s_q, h, d_v, dtype=torch.bfloat16, device=device)
+    dout = torch.randn_like(out)
+    lse = torch.randn(s_q, h, dtype=torch.float32, device=device)
+    attn_sink = torch.randn(h, dtype=torch.float32, device=device)
+    topk_idxs = torch.stack([torch.randperm(s_kv, device=device)[:topk] for _ in range(s_q)]).to(torch.int32)
+
+    for bad_value in (0, -3):
+        topk_length = torch.full((s_q,), topk, dtype=torch.int32, device=device)
+        topk_length[3] = bad_value
+        with pytest.raises(ValueError, match="topk_length"):
+            DSA.sparse_attention_backward_wrapper(
+                q,
+                kv,
+                out,
+                dout,
+                lse,
+                attn_sink,
+                topk_idxs,
+                softmax_scale=1.0 / math.sqrt(d),
+                topk_length=topk_length,
+            )
+
+
+@pytest.mark.L1
+@torch_fork_set_rng(seed=0)
+def test_DSA_sparse_attention_backward_topk_length_guard_explicit_stream():
+    """The topk_length guard must bind to the resolved launch stream.
+
+    ``sparse_attention_backward_wrapper(..., stream=...)`` forwards an
+    explicit launch stream to the SM100 interface; the guard's capture
+    detection and its device-to-host sync must run on that stream, not on
+    whatever torch stream happens to be current at call time:
+
+    - outside capture, a zero row is rejected even when the launch targets
+      an explicit, non-current stream (the sync is ordered on that stream);
+    - while the explicit stream IS capturing and the torch-current stream
+      is a different, non-capturing stream, the guard must detect the
+      capture on the resolved stream and skip validation -- the documented
+      capture semantics (the caller upholds the precondition) -- instead of
+      syncing mid-capture. The captured graph is never replayed (a replay
+      would launch the kernel with a zero-length row).
+    """
+    try:
+        import cuda.bindings.driver as cuda_driver
+        from cudnn import DSA
+    except ImportError:
+        pytest.skip("Environment not supported: cudnn[cutedsl] not installed")
+
+    major, minor = torch.cuda.get_device_capability()
+    if major * 10 + minor < 100:
+        pytest.skip("the topk_length >= 1 guard is implemented on the SM100 interface")
+
+    s_q, s_kv, h, d, d_v, topk = 8, 256, 64, 576, 512, 64
+    device = "cuda"
+    q = torch.randn(s_q, h, d, dtype=torch.bfloat16, device=device)
+    kv = torch.randn(s_kv, d, dtype=torch.bfloat16, device=device)
+    out = torch.randn(s_q, h, d_v, dtype=torch.bfloat16, device=device)
+    dout = torch.randn_like(out)
+    lse = torch.randn(s_q, h, dtype=torch.float32, device=device)
+    attn_sink = torch.randn(h, dtype=torch.float32, device=device)
+    topk_idxs = torch.stack([torch.randperm(s_kv, device=device)[:topk] for _ in range(s_q)]).to(torch.int32)
+
+    side_stream = torch.cuda.Stream()
+
+    def call(topk_length):
+        return DSA.sparse_attention_backward_wrapper(
+            q,
+            kv,
+            out,
+            dout,
+            lse,
+            attn_sink,
+            topk_idxs,
+            softmax_scale=1.0 / math.sqrt(d),
+            topk_length=topk_length,
+            stream=cuda_driver.CUstream(side_stream.cuda_stream),
+        )
+
+    bad_topk_length = torch.full((s_q,), topk, dtype=torch.int32, device=device)
+    bad_topk_length[3] = 0
+
+    # Outside capture, an explicit non-current launch stream still rejects
+    # zero rows (raises before the interface's compile cache is consulted).
+    with pytest.raises(ValueError, match="topk_length"):
+        call(bad_topk_length)
+
+    # Warm the compile cache off the capture path: capture must not compile.
+    good_topk_length = torch.full((s_q,), topk, dtype=torch.int32, device=device)
+    call(good_topk_length)
+    torch.cuda.synchronize()
+
+    # Capture on the explicit stream while the torch-current stream is a
+    # different, non-capturing stream. The guard must see the capture on the
+    # resolved stream and skip its sync; the pre-fix guard consulted the
+    # torch-current stream instead, ran the validation mid-capture, and
+    # raised here. Relaxed capture mode keeps the interface's unrelated
+    # eager work (output/workspace allocation) legal during capture.
+    other_stream = torch.cuda.Stream()
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph, stream=side_stream, capture_error_mode="relaxed"):
+        with torch.cuda.stream(other_stream):
+            assert not torch.cuda.is_current_stream_capturing(), "the guard must not rely on the torch-current stream"
+            call(bad_topk_length)
+    del graph