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..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 @@ -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] 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..0040f458d 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,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, 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..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 @@ -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