From d7b76ff6370e3d19a82aa0d4981275b1a48d1090 Mon Sep 17 00:00:00 2001 From: zky <51477259+zkyue@users.noreply.github.com> Date: Thu, 23 Jul 2026 11:40:07 +0000 Subject: [PATCH 1/3] DSA backward SM100: reject topk_length entries < 1 at the interface A topk_length entry of 0 gives the SM100 kernel zero KV tiles for that token (tile_count = ceil_div(0, block_tile) = 0 in dsa_bwd_sm100), a case the kernel has no defined behavior for. Measured on B200 (head_dim 512, head_dim_v 512, h 64, topk 512, s_q 128, s_kv 4096): a call whose topk_length contains zero entries (3 of 128 rows) never terminates -- the launch hangs at 100% SM until killed -- while the identical call with all lengths >= 1 completes normally (~40 s including first-call JIT compilation). Code inspection suggests a zero-tile pipeline deadlock, and that a zero-tile token reaching the dq epilogue would store accumulator TMEM that no MMA instruction has written. The mathematically correct dq for a query attending to nothing is exactly 0. Make the failure loud at the interface layer instead of hanging the device: flash_attn_bwd_sm100 raises ValueError when any topk_length entry is < 1 (zero entries are the measured hazard; negatives are rejected by the same check). The check needs a device-to-host sync, so it runs outside CUDA graph capture only -- under capture the caller must uphold the precondition, which is documented on the interface and on sparse_attention_backward_wrapper. Add a test covering 0 and a negative entry. No kernel change: whether topk_length == 0 should instead be supported input (skipping the pipeline handshake and zero-filling dq for zero-tile tokens) is left as an upstream semantic decision; this is minimal interface validation and is trivially removable in that world. Signed-off-by: zky <51477259+zkyue@users.noreply.github.com> --- .../_interface_sm100.py | 34 ++++++++++++- .../sparse_attention_backward/api.py | 7 +++ .../dsa/test_DSA_sparse_attention_backward.py | 48 +++++++++++++++++++ 3 files changed, 88 insertions(+), 1 deletion(-) diff --git a/python/cudnn/deepseek_sparse_attention/sparse_attention_backward/_interface_sm100.py b/python/cudnn/deepseek_sparse_attention/sparse_attention_backward/_interface_sm100.py index 3030f8198..604abdeb4 100644 --- a/python/cudnn/deepseek_sparse_attention/sparse_attention_backward/_interface_sm100.py +++ b/python/cudnn/deepseek_sparse_attention/sparse_attention_backward/_interface_sm100.py @@ -47,12 +47,27 @@ 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 -- during capture the required + device-to-host 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 +84,23 @@ def flash_attn_bwd_sm100( tensors_to_check.append(topk_length) assert all(t.is_cuda for t in tensors_to_check) + if topk_length is not None and not torch.cuda.is_current_stream_capturing(): + # 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. Skipped under CUDA graph capture, where + # the required device-to-host sync is illegal -- the precondition + # itself still applies (documented above). + if 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) diff --git a/python/cudnn/deepseek_sparse_attention/sparse_attention_backward/api.py b/python/cudnn/deepseek_sparse_attention/sparse_attention_backward/api.py index 7ce340394..375dc945c 100644 --- a/python/cudnn/deepseek_sparse_attention/sparse_attention_backward/api.py +++ b/python/cudnn/deepseek_sparse_attention/sparse_attention_backward/api.py @@ -152,6 +152,13 @@ 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 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, 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..13fb638e2 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 @@ -302,3 +302,51 @@ 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. + """ + try: + from cudnn import DSA + except ImportError: + pytest.skip("Environment not supported: cudnn[cutedsl] not installed") + + major, _ = torch.cuda.get_device_capability() + if major != 10: + 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, + ) From 19fb64f743172c9111403fb963d0aee57dd93675 Mon Sep 17 00:00:00 2001 From: zky <51477259+zkyue@users.noreply.github.com> Date: Fri, 24 Jul 2026 03:00:00 +0000 Subject: [PATCH 2/3] Address review feedback: bind the topk_length guard to the resolved stream, test gating _interface_sm100: resolve the launch stream before the topk_length guard and evaluate both the capture check and the device-to-host sync inside torch_stream_context(current_stream). Previously the guard consulted torch.cuda.is_current_stream_capturing() on whatever torch stream was current: with an explicit current_stream, a capture on that stream was not detected (validation ran mid-capture, syncing a stream it must not), and outside capture the sync was not ordered against the caller producer work on the launch stream. Add a test that exercises both sides on a real explicit side stream: outside capture a zero row is still rejected; during capture on the explicit stream -- while the torch-current stream is a different, non-capturing stream -- the guard detects the capture on the resolved stream and skips validation per the documented capture semantics (caller upholds the precondition), the launch is recorded into the graph, and the graph is never replayed. Verified red on the previous guard placement (ValueError raised mid-capture) and green with this one. Test gating: replace the bare major == 10 check with the suite's DSA capability gate (major * 10 + minor < 100, matching test_DSA_sparse_attention_backward_staged_store and dsa_init's min_compute_capability=100). The cudnn[cutedsl] import skip remains the backend gate: this CuTe-DSL-only path does not depend on cudnn.backend_version(). L0 classification of the zero-raises test: kept, with the reasoning documented in the test docstring -- SparseAttentionBackward.compile() defers compilation to execute(), and the guard raises before the interface consults its compile cache, so the validation-only test pays no CuTe compilation cost even on a cold cache. test/python fe_api/dsa: 4 failed, 28 passed, 4 skipped; the failure set is identical to the pre-change baseline. Signed-off-by: zky <51477259+zkyue@users.noreply.github.com> --- .../_interface_sm100.py | 41 +++++---- .../sparse_attention_backward/api.py | 3 +- .../dsa/test_DSA_sparse_attention_backward.py | 92 ++++++++++++++++++- 3 files changed, 117 insertions(+), 19 deletions(-) diff --git a/python/cudnn/deepseek_sparse_attention/sparse_attention_backward/_interface_sm100.py b/python/cudnn/deepseek_sparse_attention/sparse_attention_backward/_interface_sm100.py index 604abdeb4..28b5e8741 100644 --- a/python/cudnn/deepseek_sparse_attention/sparse_attention_backward/_interface_sm100.py +++ b/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 @@ -57,9 +57,11 @@ def flash_attn_bwd_sm100( Raises: ValueError: if any entry of ``topk_length`` is < 1 (validated outside - CUDA graph capture only -- during capture the required - device-to-host sync is not possible, and the caller must uphold - the precondition). A row with ``topk_length == 0`` gives the + 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 @@ -84,22 +86,30 @@ def flash_attn_bwd_sm100( tensors_to_check.append(topk_length) assert all(t.is_cuda for t in tensors_to_check) - if topk_length is not None and not torch.cuda.is_current_stream_capturing(): + 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. Skipped under CUDA graph capture, where - # the required device-to-host sync is illegal -- the precondition - # itself still applies (documented above). - if 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." - ) + # 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) @@ -159,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] diff --git a/python/cudnn/deepseek_sparse_attention/sparse_attention_backward/api.py b/python/cudnn/deepseek_sparse_attention/sparse_attention_backward/api.py index 375dc945c..0040f458d 100644 --- a/python/cudnn/deepseek_sparse_attention/sparse_attention_backward/api.py +++ b/python/cudnn/deepseek_sparse_attention/sparse_attention_backward/api.py @@ -154,7 +154,8 @@ def sparse_attention_backward_wrapper( ``d_sink`` is computed from ``attn_sink`` and ``dout``. ``topk_length``, when provided, must be >= 1 for every query. The SM100 - path validates this outside CUDA graph capture and raises ``ValueError`` + 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 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 13fb638e2..58a0e22cd 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 @@ -315,14 +315,19 @@ def test_DSA_sparse_attention_backward_topk_length_zero_raises(): 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, _ = torch.cuda.get_device_capability() - if major != 10: + 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 @@ -350,3 +355,86 @@ def test_DSA_sparse_attention_backward_topk_length_zero_raises(): softmax_scale=1.0 / math.sqrt(d), topk_length=topk_length, ) + + +@pytest.mark.L0 +@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 From 46d7187f1310502872edffc4db391b1e1725d002 Mon Sep 17 00:00:00 2001 From: zky <51477259+zkyue@users.noreply.github.com> Date: Fri, 24 Jul 2026 03:39:29 +0000 Subject: [PATCH 3/3] Mark the explicit-stream guard test L1 (it warms the compile cache before capturing) Signed-off-by: zky <51477259+zkyue@users.noreply.github.com> --- test/python/fe_api/dsa/test_DSA_sparse_attention_backward.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 58a0e22cd..cb11218fa 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 @@ -357,7 +357,7 @@ def test_DSA_sparse_attention_backward_topk_length_zero_raises(): ) -@pytest.mark.L0 +@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.