Add sparse MLA backward (DSA top-k), sliding-window attention + attention sink, and block-sparse MLA forward/backward for SM100

[kmswin1]

Summary

This PR adds three training-oriented kernel extensions, targeting SM100 (B200/B300) with the SM90 dense-decode path also covered. Commits are split per feature — happy to split into separate PRs if that is easier to review.

1. Sparse MLA backward (per-token top-k, DSA) — a CUTLASS-3 warp-specialized SM100 backward for absorbed MLA (d_qk=576, d_v=512), completing the training loop for the existing flash_mla_sparse_fwd. FA3-style: D_QK split across 3 cluster CTAs (192-col slices), D_V in 4 chunks per K-tile; CTA = one Q-token iterating its top-k K-tiles; dK/dV scattered with FP32 atomicAdd, dQ accumulated via SM90_TMA_REDUCE_ADD into a caller-provided fp32 workspace. Optionally emits a fused kl_target for indexer losses. Exposed as flash_mla_sparse_bwd (base-2 LSE contract).

2. Sliding-window attention (SWA) on dense MLA — runtime window_size (FlashAttention-style (left, right), causal left window, right must be 0, (-1,-1) disables) threaded through SM100 dense prefill fwd+bwd (MLA 192/128 and generic 128/128) and the SM90 dense decode kernel. The mask predicate is applied identically in fwd and bwd so gradients stay consistent; the forward additionally skips K-tiles entirely below the window (get_trip_start shifts the TMA loads and softmax cursor together), so a windowed prefill does proportionally less work instead of just masking (~20x fewer K-tiles at 64K with a 128-wide window). Disabled is bit-identical to plain causal.

   Includes an optional gpt-oss-style per-head attention sink (sink_bias on flash_attn_varlen_func, [h_q]): a value-less virtual key folded entirely into the existing kernels — forward folds it into the softmax denominator at finalization (O and LSE come back sink-aware), backward accumulates d(sink) inside the existing sum_OdO pass. No extra launches, ~0% overhead. Decode reuses it via attn_sink on flash_mla_with_kvcache.

3. Block-sparse MLA forward + KV-outer backward (non-absorbed 192/128 training shape) — a matched fwd/bwd pair forked from the dense MLA kernels where each query-block attends only its selected key-blocks (q2k), giving O(s·selected) compute. The backward keeps the dense kernel's KV-outer structure via a k2q reverse CSR (K-block → attending Q-blocks): a K-block CTA iterates only the Q-blocks that selected it, accumulates dK/dV in TMEM and stores once — no atomic scatter, bitwise-deterministic dK/dV. With k2q_row_ptr == nullptr it falls back to the dense path, so one kernel serves dense / dense+SWA / block-sparse / block-sparse+SWA. expand_block_selection reconciles the fwd (q_block=256) and bwd (q_block=64) granularities exactly for an end-to-end training step.

API additions

• flash_mla_sparse_bwd(q, kv, o, do, indices, lse, dq_acc_ws, dkv_acc_ws, sm_scale, d_v, fuse_reducesum)
• window_size=(left, right) on flash_attn_varlen_func / flash_attn_varlen_{qkv,kv}packed_func and flash_mla_with_kvcache
• sink_bias on flash_attn_varlen_func (autograd populates its grad); attn_sink on dense flash_mla_with_kvcache
• block_sparse_prefill_fwd / block_sparse_prefill_bwd / expand_block_selection

Benchmarks (B200, 192/128 MLA)

Selection = global sink + 4-block window (~5 selected K-blocks ≈ 640 keys, constant in seqlen); dense is O(s²), block-sparse O(s·selected), so the speedup grows with context. Running the block-sparse kernel with a full-causal selection reproduces the stock dense time to ~1%, confirming the gain is the sparsity, not a kernel artifact.

Forward (tests/bench_block_sparse_fwd_vs_dense.py, h=8):

seqlen	dense (stock, ms)	block-sparse (ms)	speedup
16,384	0.56	0.13	~4x
65,536	9.0	0.48	~19x
131,072	35.9	0.92	~39x
262,144	144	~1.7	~80-90x

Backward (tests/bench_block_sparse_bwd_scale.py, h=16):

seqlen	dense (stock, ms)	block-sparse (ms)	speedup
16,384	4.09	1.26	3.2x
65,536	69.0	4.82	14.3x
131,072	273.6	9.42	29.1x
262,144	1093.8	18.7	58.4x

The benchmark uses a fixed sink+window pattern to show scaling behavior, not a claim that every workload sees ~80x; the value is contingent on a trained block-indexer picking the important K-blocks.

Numerics

• Sparse bwd: dKV mean_abs 0.010 vs reference (bf16 noise floor in dense regions 0.002), dQ max_abs 0.058 at T=4096 / topk=2048. The dKV atomicAdd scatter is non-deterministic (run-to-run dKV cos ~0.983).
• SWA: window=0 regression matches plain causal; W ∈ {128, 384} validated fwd+bwd vs the SDPA reference in tests/test_fmha_sm100.py (which already models windowed attention), fixed + varlen, MLA (192/128) and generic (128/128) shapes.
• Attention sink: O/LSE/dQ/dK/dV/d_sink cos ≥ 0.99999 vs a sink-column reference; overhead ~0%.
• Block-sparse: fwd O/LSE cos 0.999998, bwd dQ/dK/dV cos 0.999997 vs the autograd oracle (full-causal and scattered selections); dK/dV bitwise-deterministic across runs.

Tests

tests/test_swa_correctness.py, tests/test_attention_sink.py, tests/test_block_sparse_{e2e,integrated,existing,fwd_192_integrated,fwd_inference,sink}.py, tests/ref_block_sparse_mla.py (oracle), plus scaling benches (tests/bench_swa_scaling.py, tests/bench_attention_sink_ablation.py, tests/bench_block_sparse_*.py, tests/bench_sparse_bwd_576.py).

Notes for reviewers

• One commit per feature (sparse bwd / SWA+sink / block-sparse); happy to split into separate PRs.
• The per-token sparse backward currently supports h_q=64, d_qk=576 (DSA absorbed-MLA shape) on SM100 only.
• The block-sparse kernels are built into the same extension (flash_mla.cuda) via csrc/api/api.cpp; the standalone PYBIND11_MODULE blocks in the pybind TUs are compiled out (#ifdef *_STANDALONE).