[Feature Request] DeepSeek Native Sparse Attention

[xysmlx]

https://arxiv.org/abs/2502.11089

P0

[michaelwithu]

Operator Description

The Native Sparse Attention is a highly optimized neural network operator designed for efficient sequence generation in autoregressive transformer-based models, such as large language models (LLMs). It computes attention over a sparse, pre-defined subset of key-value pairs—instead of the full sequence—significantly reducing both memory bandwidth and compute requirements. By integrating directly with the KV Cache and supporting variable-length sequences, it enables fast, memory-efficient decoding while maintaining model quality through structured sparsity patterns (e.g., sliding window or block-sparse attention)..

Implementation Plan

Split Native Sparse Attention into the following components:

• Op1: Mean Pooling
• Op2: NSA Compression
• Op3: NSA Top-K
• Op4: NSA Selection Attention
• Op5: Sliding Window Attention

Then, implement multi-level interface wrappers for each kernel.

1. Kernel Implementation (L1)

• Kernel: Implement TileLang kernel in top/kernels/deepseek_nsa/<kernel_name>.py
• Verification: Pass functional correctness tests

2. Op Definition (L2)

• Interface: Define torch.ops interface in top/ops/deepseek_nsa.py
• Unit Tests: Implement tests/ops/test_<op_name>.py (Compare vs PyTorch Ref/FLA-Triton Ref)
  • FP16 (close: 1e-3)
  • BF16 (close: 1.6e-2)
• Benchmarks: Implement benchmarks/deepseek_nsa/deepseek_nsa.py
  • Latency
  • TFLOPS
  • DRAM Bandwidth

3. Functional API (L3)

• Autograd: Implement torch.autograd.Function in top/functions/deepseek_nsa.py
• GradCheck: Pass torch.autograd.gradcheck

4. Layer Wrapper (L4)

• Module: Implement nn.Module in top/layers/deepseek_nsa.py
• Type Hints: Complete type hints (strict mode checked)

5. Benchmark Results

Shape	FP16 Latency (ms)	FP16 TFLOPS	FP16 DRAM Bandwidth (GB/s)	BF16 Latency (ms)	BF16 TFLOPS	BF16 DRAM Bandwidth (GB/s)
(B=1, S=1024, ...)	...	...	...	...	...	...

Reference

https://arxiv.org/pdf/2502.11089