diff --git a/benchmark/dsa/README.md b/benchmark/dsa/README.md index 34fab1d48..440abcd54 100644 --- a/benchmark/dsa/README.md +++ b/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 diff --git a/docs/fe-oss-apis/dsa.md b/docs/fe-oss-apis/dsa.md index 68a86a727..c97cb623a 100644 --- a/docs/fe-oss-apis/dsa.md +++ b/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( 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..5282b190e 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 @@ -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 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 2c03b0877..25f9748c7 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 @@ -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)