Fix stream-ordering and input-validation holes in the SM100 DSA backward interface

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
 
@@ -39,10 +39,10 @@ def flash_attn_bwd_sm100(
     Internally wraps as batch=1 for the CuTe DSL kernel.
 
     Args:
-        q: (total_S_q, nheads, headdim) bfloat16
-        kv: (total_S_kv, headdim) bfloat16  (K=V, MQA h_kv=1)
-        out: (total_S_q, nheads, headdim_v) bfloat16
-        dout: (total_S_q, nheads, headdim_v) bfloat16
+        q: (total_S_q, nheads, headdim) float16 or bfloat16
+        kv: (total_S_kv, headdim) float16 or bfloat16  (K=V, MQA h_kv=1)
+        out: (total_S_q, nheads, headdim_v) float16 or bfloat16
+        dout: (total_S_q, nheads, headdim_v) float16 or bfloat16
         lse: (total_S_q, nheads) float32, FlashMLA KV-only LSE excluding sink
         attn_sink: (nheads,) float32
         topk_idxs: (total_S_q, topk_max) int32, global indices
@@ -56,6 +56,10 @@ def flash_attn_bwd_sm100(
     """
     total_S_q, num_head, head_dim = q.shape
     total_S_kv = kv.shape[0]
+    # Mirror the check_support gate: the SM100 kernel is tiled only for
+    # head_dim in {512, 576}; any other value indexes shared memory out of
+    # bounds and crashes inside the kernel.
+    assert head_dim in (512, 576), f"head_dim must be 512 or 576, got {head_dim}"
     head_dim_v = 512 if head_dim == 576 else head_dim
     device = q.device
 
@@ -67,7 +71,21 @@ def flash_attn_bwd_sm100(
     tensors_to_check = [q, kv, out, dout, lse, attn_sink, topk_idxs]
     if topk_length is not None:
         tensors_to_check.append(topk_length)
-    assert all(t.is_cuda for t in tensors_to_check)
+    assert all(t.is_cuda and t.device == device for t in tensors_to_check), f"all inputs must be CUDA tensors on {device}"
+
+    # Cross-tensor shape validation: every tensor below is indexed with
+    # coordinates derived from q, so a mismatched shape silently reads or
+    # writes out of place instead of failing.
+    assert kv.ndim == 2 and kv.shape[1] == head_dim, f"kv shape mismatch: expected (total_S_kv, {head_dim}), got {tuple(kv.shape)}"
+    expected_o_shape = (total_S_q, num_head, head_dim_v)
+    assert out.shape == expected_o_shape, f"out shape mismatch: expected {expected_o_shape}, got {tuple(out.shape)}"
+    assert dout.shape == expected_o_shape, f"dout shape mismatch: expected {expected_o_shape}, got {tuple(dout.shape)}"
+    assert lse.shape == (total_S_q, num_head), f"lse shape mismatch: expected {(total_S_q, num_head)}, got {tuple(lse.shape)}"
+    assert attn_sink.shape == (num_head,), f"attn_sink shape mismatch: expected {(num_head,)}, got {tuple(attn_sink.shape)}"
+    assert topk_idxs.ndim == 2 and topk_idxs.shape[0] == total_S_q, f"topk_idxs shape mismatch: expected ({total_S_q}, topk_max), got {tuple(topk_idxs.shape)}"
+    if topk_length is not None:
+        assert topk_length.dtype == torch.int32, f"topk_length dtype mismatch: expected torch.int32, got {topk_length.dtype}"
+        assert topk_length.shape == (total_S_q,), f"topk_length shape mismatch: expected {(total_S_q,)}, got {tuple(topk_length.shape)}"
 
     if softmax_scale is None:
         softmax_scale = 1.0 / math.sqrt(head_dim)
@@ -76,58 +94,75 @@ def flash_attn_bwd_sm100(
     num_head_blocks = (num_head + block_tile - 1) // block_tile
     batch_size = 1
 
-    # Ensure contiguous
-    q, kv, out, dout = [t.contiguous() for t in (q, kv, out, dout)]
-    lse = lse.contiguous()
-
-    # Allocate output tensors
-    if dq is None:
-        dq = torch.empty_like(q)
-    else:
-        assert dq.shape == q.shape, f"dq shape mismatch: expected {q.shape}, got {dq.shape}"
-        assert dq.dtype == q.dtype, f"dq dtype mismatch: expected {q.dtype}, got {dq.dtype}"
-        assert dq.device == device, f"dq device mismatch: expected {device}, got {dq.device}"
-    if dkv is None:
-        dkv = torch.zeros(total_S_kv, head_dim, dtype=kv.dtype, device=device)
-    else:
-        expected_dkv_shape = (total_S_kv, head_dim)
-        assert dkv.shape == expected_dkv_shape, f"dkv shape mismatch: expected {expected_dkv_shape}, got {dkv.shape}"
-        assert dkv.dtype == kv.dtype, f"dkv dtype mismatch: expected {kv.dtype}, got {dkv.dtype}"
-        assert dkv.device == device, f"dkv device mismatch: expected {device}, got {dkv.device}"
-        dkv.fill_(0)
-    d_sink = torch.zeros_like(attn_sink)
-
-    # Allocate workspace tensors
-    acc_dtype = cutlass.Float32
-    ws_lse_odo_shape = FlashAttentionDSABackwardSm100._get_workspace_size_LSE_OdO(
-        total_S_q,
-        head_dim,
-        num_head,
-        batch_size,
-        acc_dtype,
-    )
-    workspace_LSE_OdO = torch.zeros(
-        *ws_lse_odo_shape,
-        dtype=torch.uint8,
-        device=device,
-    )
-
-    ws_dkv_shape = FlashAttentionDSABackwardSm100._get_workspace_size_dKV(
-        total_S_kv,
-        head_dim,
-        batch_size,
-        acc_dtype,
-    )
-    workspace_dKV = torch.zeros(
-        *ws_dkv_shape,
-        dtype=torch.uint8,
-        device=device,
-    )
+    current_stream = resolve_stream(current_stream)
+
+    # Normalize inputs and allocate outputs/workspaces on the execution stream:
+    # the kernel below launches on `current_stream`, so the semantically
+    # required zero-initialization of dkv/d_sink and both workspaces (and any
+    # contiguity copies) must be stream-ordered with it, not with the ambient
+    # torch stream the caller happens to be on.
+    with torch_stream_context(current_stream):
+        # Ensure contiguous
+        q, kv, out, dout = [t.contiguous() for t in (q, kv, out, dout)]
+        lse = lse.contiguous()
+        attn_sink = attn_sink.contiguous()
+        topk_idxs = topk_idxs.contiguous()
+        if topk_length is not None:
+            topk_length = topk_length.contiguous()
+
+        # Allocate output tensors
+        if dq is None:
+            dq = torch.empty_like(q)
+        else:
+            assert dq.shape == q.shape, f"dq shape mismatch: expected {q.shape}, got {dq.shape}"
+            assert dq.dtype == q.dtype, f"dq dtype mismatch: expected {q.dtype}, got {dq.dtype}"
+            assert dq.device == device, f"dq device mismatch: expected {device}, got {dq.device}"
+            # The compile cache is keyed without output strides, so a caller
+            # provided output must match the contiguous layout the kernel was
+            # compiled for (it is not copied: that would break out-parameter
+            # identity).
+            assert dq.is_contiguous(), "dq must be contiguous"
+        if dkv is None:
+            dkv = torch.zeros(total_S_kv, head_dim, dtype=kv.dtype, device=device)
+        else:
+            expected_dkv_shape = (total_S_kv, head_dim)
+            assert dkv.shape == expected_dkv_shape, f"dkv shape mismatch: expected {expected_dkv_shape}, got {dkv.shape}"
+            assert dkv.dtype == kv.dtype, f"dkv dtype mismatch: expected {kv.dtype}, got {dkv.dtype}"
+            assert dkv.device == device, f"dkv device mismatch: expected {device}, got {dkv.device}"
+            assert dkv.is_contiguous(), "dkv must be contiguous"
+            dkv.fill_(0)
+        d_sink = torch.zeros_like(attn_sink)
+
+        # Allocate workspace tensors
+        acc_dtype = cutlass.Float32
+        ws_lse_odo_shape = FlashAttentionDSABackwardSm100._get_workspace_size_LSE_OdO(
+            total_S_q,
+            head_dim,
+            num_head,
+            batch_size,
+            acc_dtype,
+        )
+        workspace_LSE_OdO = torch.zeros(
+            *ws_lse_odo_shape,
+            dtype=torch.uint8,
+            device=device,
+        )
+
+        ws_dkv_shape = FlashAttentionDSABackwardSm100._get_workspace_size_dKV(
+            total_S_kv,
+            head_dim,
+            batch_size,
+            acc_dtype,
+        )
+        workspace_dKV = torch.zeros(
+            *ws_dkv_shape,
+            dtype=torch.uint8,
+            device=device,
+        )
 
     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]
@@ -149,6 +184,7 @@ def flash_attn_bwd_sm100(
         workspace_dKV_tensor = to_cute_tensor(workspace_dKV)
 
         kernel_obj = FlashAttentionDSABackwardSm100(
+            element_dtype=dtype,
             head_dim=head_dim,
             head_dim_v=head_dim_v,
             block_tile=block_tile,

python/cudnn/deepseek_sparse_attention/sparse_attention_backward/api.py

@@ -20,8 +20,8 @@
 class SparseAttentionBackward(APIBase):
     def __init__(
         self,
-        sample_q: torch.Tensor,  # (total_S_q, H, D) BF16
-        sample_kv: torch.Tensor,  # (total_S_kv, D) BF16 (K=V)
+        sample_q: torch.Tensor,  # (total_S_q, H, D) FP16/BF16
+        sample_kv: torch.Tensor,  # (total_S_kv, D) FP16/BF16 (K=V)
         sample_out: torch.Tensor,  # (total_S_q, H, D_v)
         sample_dout: torch.Tensor,  # (total_S_q, H, D_v)
         sample_lse: torch.Tensor,  # (total_S_q, H) FP32, KV-only LSE
@@ -69,6 +69,79 @@ def check_support(self) -> bool:
         self._check_dtype(self.lse_desc, torch.float32, name="LSE")
         self._check_dtype(self.attn_sink_desc, torch.float32, name="attn_sink")
         self._check_dtype(self.topk_idxs_desc, torch.int32, name="topk_idxs")
+        self._check_dtype(self.out_desc, self.q_desc.dtype, name="out", extra_error_msg="out must have same dtype as Q")
+        self._check_dtype(self.dout_desc, self.q_desc.dtype, name="dout", extra_error_msg="dout must have same dtype as Q")
+        if self.topk_length_desc is not None:
+            self._check_dtype(self.topk_length_desc, torch.int32, name="topk_length")
+
+        # Device placement + cross-tensor device consistency. The SM90/SM100
+        # kernels are CUDA-only and reject CPU or cross-device inputs at
+        # execution time (see the is_cuda / same-device assert in
+        # ``_interface_sm100.flash_attn_bwd_sm100``), so a placement mismatch
+        # must fail the support gate here rather than compile/launch and crash.
+        ref_device = self.q_desc.device
+        descriptors = [
+            self.q_desc,
+            self.kv_desc,
+            self.out_desc,
+            self.dout_desc,
+            self.lse_desc,
+            self.attn_sink_desc,
+            self.topk_idxs_desc,
+        ]
+        if self.topk_length_desc is not None:
+            descriptors.append(self.topk_length_desc)
+        self._value_error_if(
+            ref_device.type != "cuda",
+            f"Q must live on CUDA, got {ref_device}",
+        )
+        self._value_error_if(
+            any(desc.device != ref_device for desc in descriptors),
+            f"All inputs must share Q's device {ref_device}, got {[desc.device for desc in descriptors]}",
+        )
+
+        # Cross-tensor shape contract: every companion tensor is indexed with
+        # coordinates derived from Q, so a mismatched shape silently reads or
+        # writes out of place at execution time instead of failing.
+        total_s_q, num_heads, head_dim = self.q_desc.shape
+        # The SM100 kernel is tiled only for head_dim in {512, 576} (the 576
+        # MLA case splits QK=576 / V=512); any other head_dim compiles to a
+        # layout that indexes shared memory out of bounds and crashes.
+        self._value_error_if(
+            head_dim not in (512, 576),
+            f"head_dim must be 512 or 576, got {head_dim}",
+        )
+        head_dim_v = 512 if head_dim == 576 else head_dim
+        expected_o_shape = (total_s_q, num_heads, head_dim_v)
+        self._value_error_if(
+            self.kv_desc.shape[1] != head_dim,
+            f"KV must have shape (total_S_kv, {head_dim}), got {self.kv_desc.shape}",
+        )
+        self._value_error_if(
+            self.out_desc.shape != expected_o_shape,
+            f"out must have shape {expected_o_shape}, got {self.out_desc.shape}",
+        )
+        self._value_error_if(
+            self.dout_desc.shape != expected_o_shape,
+            f"dout must have shape {expected_o_shape}, got {self.dout_desc.shape}",
+        )
+        self._value_error_if(
+            self.lse_desc.shape != (total_s_q, num_heads),
+            f"LSE must have shape {(total_s_q, num_heads)}, got {self.lse_desc.shape}",
+        )
+        self._value_error_if(
+            self.attn_sink_desc.shape != (num_heads,),
+            f"attn_sink must have shape {(num_heads,)}, got {self.attn_sink_desc.shape}",
+        )
+        self._value_error_if(
+            self.topk_idxs_desc.ndim != 2 or self.topk_idxs_desc.shape[0] != total_s_q,
+            f"topk_idxs must have shape ({total_s_q}, topk_max), got {self.topk_idxs_desc.shape}",
+        )
+        if self.topk_length_desc is not None:
+            self._value_error_if(
+                self.topk_length_desc.shape != (total_s_q,),
+                f"topk_length must have shape {(total_s_q,)}, got {self.topk_length_desc.shape}",
+            )
 
         self._is_supported = True
         return True

python/cudnn/deepseek_sparse_attention/sparse_attention_backward/dsa_bwd_sm100.py

@@ -4,7 +4,7 @@
 
 import cutlass
 import cutlass.cute as cute
-from cutlass.cute.typing import Float32, Int32, BFloat16, Int64
+from cutlass.cute.typing import Float32, Int32, Int64
 import cutlass.pipeline as pipeline
 from cutlass.cute.nvgpu import OperandMajorMode, cpasync, tcgen05
 import cutlass.utils.blackwell_helpers as sm100_utils
@@ -17,6 +17,7 @@ class FlashAttentionDSABackwardSm100:
 
     def __init__(
         self,
+        element_dtype: Type[cutlass.Numeric],
         head_dim: int,
         head_dim_v: int,
         block_tile: int,
@@ -40,9 +41,11 @@ def __init__(
         self.QdS_cta_tiler = (head_dim_main, block_tile, block_tile)
         self.cluster_shape_mn = (1, 1)
 
-        self.element_dtype = BFloat16
-        # User constraint: dKV accumulation must stay FP32. BFloat16 is only
-        # used for element/output storage.
+        if element_dtype not in [cutlass.Float16, cutlass.BFloat16]:
+            raise ValueError(f"Unsupported element dtype: {element_dtype}")
+        self.element_dtype = element_dtype
+        # dKV accumulation stays FP32; element_dtype controls element/output
+        # storage.
         self.acc_dtype = Float32
 
         # =============== Sum OdO ================

test/python/fe_api/dsa/dsa_utils.py

@@ -16,6 +16,11 @@ def _require_sm90():
         pytest.skip("SM90 GPU required")
 
 
+def _require_sm100():
+    if not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 10:
+        pytest.skip("SM100 GPU required")
+
+
 # Parameterization marks shared by every DSA test
 DSA_PARAM_MARKS = [
     pytest.mark.parametrize("dtype", [torch.bfloat16]),