Fix DSA backward boundaries and relax SM100 BSA FP8 shape limits

docs/fe-oss-apis/bsa.md

@@ -34,9 +34,9 @@ pip install nvidia-cudnn-frontend
 ```
 
 The package-wide `nvidia-cutlass-dsl[cu13]>=4.6.2` dependency floor applies to
-BSA. The Sage FP8 API below performs an additional runtime check and requires
-`nvidia-cutlass-dsl>=4.6.1`, which the package floor already satisfies; the
-check stays in place for environments that force an older DSL.
+BSA. Sage FP8 relies on functionality introduced in CuTe DSL 4.6.1, which the
+package floor already satisfies. A defensive runtime check remains for
+environments that force an older DSL.
 
 ## Forward
 
@@ -167,19 +167,19 @@ backward implementation.
 
 The architecture-specific FP8 contracts are:
 
-- SM100/SM103 requires `B=1`, `H` equal to 4 or 8, and both sequence lengths
-  to be multiples of 64. It uses fixed `block_sparse_num` with full 64-token
-  KV blocks; `q2k_block_nums` and `block_sizes` are not supported. Split-KV is
-  selected internally, and the public FP8 API does not expose `kv_splits` or
-  `use_clc`.
+- SM100/SM103 accepts any positive batch and head counts and requires both
+  sequence lengths to be multiples of 64. It uses fixed `block_sparse_num`
+  with full 64-token KV blocks; `q2k_block_nums` and `block_sizes` are not
+  supported. Split-KV is selected internally, and the public FP8 API does not
+  expose `kv_splits` or `use_clc`.
 - SM120 accepts any positive batch and head counts, non-aligned Q/KV sequence
   tails, fixed or per-query-block counts, and `block_sizes` shaped `(N_kv,)`,
   `(B, N_kv)`, or `(B, H, N_kv)`. It does not use split-KV.
 
-`block_sparse_attention_fp8_forward` requires CuTe DSL 4.6.1 or newer at call
-time. Its internal Q/K/V quantizer is also implemented in CuTe DSL and is
-loaded lazily, so importing `cudnn` still works with the package-wide CuTe DSL
-4.5 dependency floor.
+`block_sparse_attention_fp8_forward` relies on functionality introduced in
+CuTe DSL 4.6.1; package-supported installations provide CuTe DSL 4.6.2 or
+newer. Its internal Q/K/V quantizer is also implemented in CuTe DSL and is
+loaded lazily, so importing `cudnn` does not eagerly import it.
 
 ## Backward
 
@@ -223,7 +223,7 @@ therefore requires full physical KV blocks and `block_sizes=None`.
 | SM90 | 64 | FP16, BF16 | each of 64, 96, 128 | MHA, GQA, MQA |
 | SM100/SM103 | 128 | FP16, BF16 | QK=V=64, 96, or 128 | MHA, GQA, MQA |
 | SM100/SM103 | 64 (explicit) | BF16 | QK=128, V=128 | MHA |
-| SM100/SM103 | 64 | BF16 / FP8 E4M3 | QK=128, V=128 | MHA (B=1, H=4 or 8) |
+| SM100/SM103 | 64 | BF16 / FP8 E4M3 | QK=128, V=128 | MHA |
 | SM120 | 64 | FP16, BF16 | QK=128, V=128 | MHA, GQA, MQA |
 | SM120 | 64 | BF16 / FP8 E4M3 | QK=128, V=128 | MHA |
 

python/cudnn/block_sparse_attention/_fp8_quant.py

@@ -407,8 +407,6 @@ def _quantize_sage_bhsd(
     if seqlen_k < 1:
         raise ValueError("Sage FP8 quantization requires positive batch, head, and sequence counts")
     is_sm120 = torch.cuda.get_device_capability(q_bhsd.device)[0] == 12
-    if not is_sm120 and (batch_size != 1 or num_heads not in (4, 8)):
-        raise ValueError("Sage FP8 v1 requires B=1, H in {4, 8}, and D=128")
     if not is_sm120 and (seqlen_q % 64 or seqlen_k % 64):
         raise ValueError("Q and K/V sequence lengths must be multiples of 64")
 

python/cudnn/block_sparse_attention/_interface.py

@@ -1354,8 +1354,6 @@ def bsa_attn_fwd_blk64_cutedsl(
     assert k_bhsd.shape == (batch_size, num_head_kv, seqlen_k, head_dim)
     assert v_bhsd.shape == (batch_size, num_head_kv, seqlen_k, head_dim_v)
     if is_sage_fp8:
-        assert batch_size == 1, "SM100 FP8 requires batch size 1"
-        assert num_head in (4, 8), "SM100 FP8 supports H=4 or H=8"
         assert seqlen_q % 64 == 0 and seqlen_k % 64 == 0, "SM100 FP8 requires Sq and Sk to be multiples of 64"
         q_scale = maybe_contiguous(q_scale)
         k_scale = maybe_contiguous(k_scale)

python/cudnn/block_sparse_attention/api.py

@@ -323,8 +323,6 @@ def block_sparse_attention_fp8_forward(
     if arch_family not in {10, 11, 12}:
         raise RuntimeError(f"Sage FP8 block sparse attention requires SM100-SM120, found SM{arch}")
     if arch_family in {10, 11}:
-        if batch != 1 or heads not in {4, 8}:
-            raise NotImplementedError("SM100/SM110 Sage FP8 requires B=1 and H in {4, 8}")
         if seqlen_q % 64 or seqlen_k % 64:
             raise NotImplementedError("SM100/SM110 Sage FP8 requires Sq and Sk to be multiples of 64")
         if q2k_block_nums is not None or block_sizes is not None:

python/cudnn/deepseek_sparse_attention/README.md

@@ -10,4 +10,4 @@
 
 ## Acknowledgements
 
-The DSA/CSA kernels were a collaborative effort, jointly developed by: Hongxiao Bai, Jiayu Sun and Jie Fang
+The DSA/CSA kernels were a collaborative effort, jointly developed by: Hongxiao Bai, Jiayu Sun, Jerry Chen, Jie Fang, Qi Zhang, Bin Chai and Yue Weng

python/cudnn/deepseek_sparse_attention/sparse_attention_backward/_interface_sm90.py

@@ -64,11 +64,10 @@ def flash_attn_bwd_sm90(
         dkv: pre-allocated (total_S_kv, headdim), optional
         d_sink: pre-allocated (nheads,), optional
         topk_idxs: (total_S_q, topk_max) int32, global KV indices.
-            Compact (`topk_length` given): prefix `[0, topk_length)` is valid.
-            Non-compact (`topk_length` is None): a negative entry is padding.
-            Positive indices must be in `[0, S_kv)`; oversized `topk_length`
-            is out of contract.
-        topk_length: (total_S_q,) int32, per-query valid prefix length, optional
+            Entries outside `[0, S_kv)` are ignored in both compact and
+            non-compact modes.
+        topk_length: (total_S_q,) int32, optional per-query valid prefix
+            length, clamped to `[0, topk_max]` by the kernel.
         need_d_sink: return and compute d_sink when True
 
     Returns: