Add Native Sparse Attention (NSA) Triton kernels and module

[abtonmoy]

Summary

Adds Native Sparse Attention (NSA, arXiv:2502.11089) as a
Triton-kernel-backed module. Closes #656.

NSA is a natively-trainable sparse attention mechanism built for grouped-query attention. Every
query attends through three gated branches:

• compression — coarse attention over learned block-level KV summaries,
• selection — fine attention over the top-n KV blocks chosen per query (the sparse core),
• sliding window — recent-token attention,

blended by learned per-head gates. All query heads in a KV group share one block selection, which
is what keeps the sparse kernel hardware-aligned.

This PR provides:

• LigerNativeSparseAttention (nn.Module) and a functional entry point liger_native_sparse_attention,
• three FlashAttention-2 Triton kernels (forward + backward) for the compression, selection, and
  sliding-window branches, in ops/nsa_{compressed,selected,sliding}_attention.py,
• a pure-PyTorch reference of the full mechanism used as the correctness oracle and as the CPU / fp64
  fallback,
• unit tests and a sequence-length benchmark against dense causal attention.

Details

What is kernelized, and what is not. The three attention branches carry the O(S²) cost, so they
run as fused kernels. Compression φ (a small MLP over strided KV blocks), the selection scoring
(Eq. 9/10) and top-n routing (Eq. 11/12), and the gate blend stay in PyTorch — they are cheap and
the routing is non-differentiable by design (discrete block selection carries no gradient; the block
scorer is still trained through the differentiable compression branch). The module dispatches the
attention branches to the kernels on an accelerator and to the reference on CPU; both paths are
numerically equivalent and fully differentiable. use_kernel=True/False/None overrides or auto-selects.

Backward. Each kernel uses the standard FlashAttention-2 recompute-from-LSE backward, split into
a query-parallel dQ kernel and a KV-parallel dK/dV kernel. Because block selection is a fixed boolean
mask (no gradient), dK/dV can be accumulated with a single deterministic write per tile — there are no
atomic_add calls anywhere, so the backward is bit-for-bit reproducible.

Portability. The kernels follow the conventions already used across ops/: fp32 dot accumulators,
head dim padded to a power of two for the MMA, exp2/log2 via the libdevice import ladder, int64
offset arithmetic, no bf16/fp16 atomics, and warp counts capped for HIP. fp32 dots use
input_precision="ieee" so the fp32 path matches the reference tightly rather than silently dropping
to TF32 (this does not affect the fp16/bf16 paths, which use native tensor cores).

Scope. This is the parallel training / prefill path over dense
[batch, heads, seq_len, head_dim] tensors, with Dk == Dv. Autoregressive decode and variable-length
(cu_seqlens) inputs are out of scope here. The reference implementations
XunhaoLai/native-sparse-attention-triton
(Apache-2.0) and fla-org/native-sparse-attention
(MIT) were consulted for the kernel structure; no code was copied.

Candidate follow-ups (kept out of this PR to keep it reviewable): group-centric KV reuse in the
selected kernel, a fused online top-k to avoid materializing the compression score matrix, a Triton
compression φ, and Dk != Dv support.

Testing Done

67 unit tests in test/transformers/test_native_sparse_attention.py:

• The vectorized reference is checked against an independent, per-token naive oracle (forward and
  backward) across fp64/fp32/fp16/bf16, GQA groups 1/2/4, several sequence lengths, and edge shapes;
  plus degenerate-case milestones (a branch reduces to ordinary causal attention), exact hand-computed
  block selection, no-future-leakage (causality), and empty-compression short sequences.
• Each Triton branch kernel is checked against its reference branch (forward and backward), and the
  full module's kernel path is checked against its torch path on the output and on every parameter
  gradient. fp32 matches to ~1e-3 on gradients; fp16/bf16 within the usual low-precision tolerances.
  The backward is asserted bit-identical across repeated runs.

Benchmark (benchmark/scripts/benchmark_native_sparse_attention.py), bf16, hidden 1024, GQA-4,
l=32, d=16, l'=64, n=16, w=512, full forward+backward, vs dense causal attention:

seq_len	dense (ms)	NSA (ms)	speedup	dense (MB)	NSA (MB)	mem ratio
512	1.81	9.81	0.18×	107	55	1.97×
1024	9.47	7.95	1.19×	333	81	4.14×
2048	35.65	7.57	4.71×	1218	151	8.07×
4096	140.72	24.60	5.72×	4722	440	10.72×

Speed crosses over around 1k tokens and grows to ~5.7× at 4k; memory is lower at every length (up to
~10.7×), and dense runs out of memory past 4k on this device. The short-sequence slowdown is expected:
at small S the fixed routing overhead dominates and w >= S means the window branch is still full
attent (Plots attached.)

ion.

• Hardware Type: RTX 4060 Laptop GPU (8 GB), CUDA 12.6, Triton 3.x
• run make test to ensure correctness
• run make checkstyle to ensure code style
• run make test-convergence — not applicable; this is a standalone attention module, not a
  drop-in replacement for a modeled layer's forward, so there is no convergence target to compare.

[abtonmoy]

This implements #656: NSA as a Triton-kernel-backed module (compression, top-n block selection, and sliding-window branches), with a pure-PyTorch reference used as the correctness oracle and CPU/fp64 fallback.

Validated on an RTX 4060: 67 unit tests pass (reference vs an independent naive oracle, and each kernel + the full module vs the reference, across fp32/fp16/bf16, GQA groups, and edge shapes; backward asserted bit-identical). ruff check/format are clean. The sequence-length benchmark vs dense causal attention shows the branches crossing over around 1k tokens and reaching ~5.7x faster / ~10.7x less memory at 4k (numbers and plots in the description).

@Tcc0403 would appreciate a review when you have a chance. Happy to adjust scope or split anything out.