DSA: fix 576-wide d_sink, test/document top-k index semantics, B300 benchmarks

benchmark/dsa/README.md

@@ -88,6 +88,72 @@ repository.
 |     8192 |      8192 | 1024 |  4.538 |     605.76 |
 |     8192 |      8192 | 2048 |  8.562 |     642.06 |
 
+### B300
+
+Generated on an NVIDIA B300 (`nheads=64`, bf16, attention sink and
+`topk_length` enabled, `warmup=10`, `repeat=50`), using `torch 2.13.0+cu130`,
+`nvidia-cutlass-dsl 4.5.2`, CUDA 13.0, and `nvidia-cudnn-frontend` built from
+this repository.
+
+`d_qk = d_v = 512`:
+
+| seqlen_q | seqlen_kv | topk | BWD ms | BWD TFLOPS |
+|---------:|----------:|-----:|-------:|-----------:|
+|     1024 |      1024 |  512 |  0.344 |     500.08 |
+|     2048 |      2048 |  512 |  0.648 |     530.63 |
+|     4096 |      4096 |  512 |  1.238 |     555.26 |
+|     8192 |      8192 |  512 |  2.406 |     571.31 |
+|    16384 |     16384 |  512 |  4.789 |     574.04 |
+|    32768 |     32768 |  512 | 10.415 |     527.85 |
+|     2048 |      2048 | 2048 |  1.993 |     689.57 |
+|     4096 |      4096 | 2048 |  3.949 |     696.08 |
+|     8192 |      8192 | 2048 |  7.914 |     694.63 |
+|    16384 |     16384 | 2048 | 15.906 |     691.24 |
+|    32768 |     32768 | 2048 | 33.351 |     659.37 |
+
+`d_qk = 576` (512 value dims + 64 RoPE dims, `d_v = 512`):
+
+| seqlen_q | seqlen_kv | topk | BWD ms | BWD TFLOPS |
+|---------:|----------:|-----:|-------:|-----------:|
+|     2048 |      2048 | 2048 |  2.638 |     560.12 |
+|     4096 |      4096 | 2048 |  5.214 |     566.76 |
+|     8192 |      8192 | 2048 | 10.312 |     573.13 |
+|    16384 |     16384 | 2048 | 20.620 |     573.21 |
+|    32768 |     32768 | 2048 | 47.195 |     500.89 |
+
+### Tensor-pipe (MMA) utilization
+
+Nsight Compute speed-of-light numbers for the main warp-specialized backward
+kernel (`kernel_cutlass_bwd_*`, 97–98% of the pass at `topk=2048`), collected
+on the same B300 with
+`ncu --profile-from-start off -k regex:kernel_cutlass_bwd --metrics
+sm__pipe_tensor_cycles_active.avg.pct_of_peak_sustained_elapsed`. ncu locks
+clocks to base during collection, so percentages are comparable across
+configs while the wall-clock numbers above reflect boost clocks. All tensor
+pipe activity is HMMA.
+
+| d_qk | topk | seqlen 1k | 2k | 4k | 8k | 16k | 32k |
+|-----:|-----:|----------:|-----:|-----:|-----:|-----:|-----:|
+|  512 |  512 |      36.7 | 37.5 | 38.3 | 38.7 | 38.7 | 38.4 |
+|  512 | 2048 |         — | 43.7 | 44.1 | 44.2 | 44.5 | 44.3 |
+|  576 | 2048 |         — | 36.8 | 36.9 | 37.1 | 37.2 | 34.2 |
+
+Observations:
+
+- MMA utilization is essentially **sequence-length invariant** at a fixed
+  top-k: the grid parallelizes over query tokens, so utilization is set by
+  the per-CTA inner loop (`topk / 64` MMA iterations amortizing a fixed
+  prologue/epilogue), not by problem size. `topk` is the utilization knob.
+- The kernel is memory-pipe-heavy rather than tensor-bound: L1TEX/memory
+  speed-of-light (62–68%) exceeds SM (49–57%) and MMA (34–45%) everywhere;
+  the top-k gather of KV rows plus fp32 dKV accumulation is the busiest pipe.
+- At `seqlen 32k` the KV working set (32k x 512 x 2B = 32 MB) exceeds L2 and
+  DRAM utilization jumps (1.8% -> 12.4% at `d_qk=512/topk=2048`,
+  6.5% -> 22.7% at `topk=512`, 1.6% -> 15.1% at `d_qk=576`), which is the
+  wall-clock dip in the tables above.
+- The 576-wide latent runs ~7 SOL points below 512 at the same top-k: the
+  non-128-aligned 64-dim RoPE tail takes separate tail-tile MMAs.
+
 ## Profiling
 
 `profile` mode runs a single warmed-up backward call (using the first value

docs/fe-oss-apis/dsa.md

@@ -119,6 +119,15 @@ Backward pass for DeepSeek Sparse Attention. Expects the forward outputs
   - `topk_length` (optional): `(total_S_q,)` INT32 — per-query valid count
 - **Outputs** — tuple `(dq, dkv, d_sink)`
 - **Constraints** — SM90 or SM100; SM90 supports the FlashMLA DSA shape with `head_dim ∈ {512, 576}`
+- **Top-k index semantics** — without `topk_length`, invalid `topk_idxs`
+  entries must be `-1` sentinels (the FlashMLA convention). With
+  `topk_length`, the first `topk_length[i]` entries of each row must all be
+  valid indices — this compact-prefix fast path skips per-row sentinel
+  checks, so `-1` entries inside the prefix are not allowed. `topk_length`
+  values must be `>= 1`; to make a query row inert (e.g. padding rows for
+  fixed-shape capture), use `topk_length = 1` with any valid index and a
+  zeroed `dout` row, which yields an exactly-zero `dq` row and no `dkv` /
+  `d_sink` contribution.
 
 ```python
 result = DSA.sparse_attention_backward_wrapper(

python/cudnn/deepseek_sparse_attention/sparse_attention_backward/_interface_sm100.py

@@ -197,6 +197,18 @@ def flash_attn_bwd_sm100(
             current_stream,
         )
 
+    if head_dim != head_dim_v:
+        # The kernel's sum_dSink postprocess only runs when D_qk == D_v; for
+        # the 576-wide MLA latent it leaves d_sink zeroed. The sink gradient
+        # has a closed form from tensors already at hand:
+        #   d_sink_h = -sum_t sigmoid(sink_h - lse[t,h]) * (out . dout)[t,h]
+        # which is the derivative of the sink-normalized softmax through the
+        # KV-only LSE. Precision matches the kernel's own D_qk == D_v
+        # sum_dSink path (both consume 16-bit out/dout).
+        odo = (out.float() * dout.float()).sum(-1)
+        p_sink = torch.sigmoid(attn_sink.unsqueeze(0) - lse)
+        d_sink.copy_(-(p_sink * odo).sum(0))
+
     return dq, dkv, d_sink
 
 

test/python/fe_api/dsa/test_DSA_sparse_attention_backward.py

@@ -216,3 +216,47 @@ def fail_on_pad(*args, **kwargs):
         atol=5e-2,
         rtol=5e-2,
     )
+
+
+@pytest.mark.L0
+@torch_fork_set_rng(seed=0)
+def test_DSA_sparse_attention_backward_sentinel_matches_topk_length():
+    """-1-sentinel indices (FlashMLA convention) match the compact-prefix path.
+
+    Without ``topk_length``, invalid ``topk_idxs`` entries are ``-1``
+    sentinels and the kernel zero-fills those rows; with ``topk_length``,
+    valid indices must be a compact prefix and per-row sentinel checks are
+    skipped. The same valid index set must produce the same gradients either
+    way. The two modes are different compiled variants, so this also guards
+    the variant selection.
+    """
+    if not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] < 10:
+        pytest.skip("SM100+ GPU required")
+
+    s_q, s_kv, h, d, topk = 512, 2048, 64, 512, 256
+    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)
+    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)
+    topk_length = torch.randint(1, topk + 1, (s_q,), dtype=torch.int32, device=device)
+
+    # Sentinel variant: same valid prefix, -1 beyond each row's length.
+    slot = torch.arange(topk, device=device).unsqueeze(0).expand(s_q, topk)
+    idx_sentinel = torch.where(slot < topk_length.unsqueeze(1), topk_idxs, torch.full_like(topk_idxs, -1))
+
+    out, lse = ref_sparse_attention_forward(q, kv, attn_sink, topk_idxs, topk_length=topk_length)
+    dout = torch.randn_like(out)
+
+    try:
+        from cudnn import DSA
+
+        r_len = DSA.sparse_attention_backward_wrapper(q, kv, out, dout, lse, attn_sink, topk_idxs, topk_length=topk_length)
+        r_sen = DSA.sparse_attention_backward_wrapper(q, kv, out, dout, lse, attn_sink, idx_sentinel, topk_length=None)
+    except (ImportError, ValueError, NotImplementedError, RuntimeError) as e:
+        pytest.skip(f"Unsupported testcase: {e}")
+    torch.cuda.synchronize()
+
+    torch.testing.assert_close(r_len["dq"], r_sen["dq"], atol=1e-2, rtol=1e-2)
+    torch.testing.assert_close(r_len["dkv"], r_sen["dkv"], atol=1e-2, rtol=1e-2)
+    torch.testing.assert_close(r_len["d_sink"], r_sen["d_sink"], atol=1e-2, rtol=1e-2)