Views
No views yet
pip install -e .1queries = torch.randn(batch, query_heads, seq_len, head_dim, device=device, dtype=dtype).contiguous()
2keys = torch.randn(batch, kv_heads, seq_len, head_dim, device=device, dtype=dtype).contiguous()
3values = torch.randn(batch, kv_heads, seq_len, head_dim, device=device, dtype=dtype).contiguous()
4head_cls = torch.randn(kv_heads, head_dim, device=device, dtype=dtype).contiguous()
5
6model = dash_attn(
7 chunk_size=chunk_size,
8 enable_gqa=True,
9 estimate_diagonal=True,
10 return_active_blocks=True,
11)
12
13out, active_blocks = model(queries, keys, values, head_cls)python ./example/run_niah.pydash_attn.prefill.summarize_chunk and dash_attn.decoding.summarize_chunk build one learned key summary per KV chunk.score_blocks computes sparse chunk supports and routing priors from query-to-summary scores.full_attn applies token-level attention only over routed chunks, using the Stage 1 prior to preserve differentiability through the hierarchy.dash_attn.dash_attn_interface.dash_attn. It supports both prefill and decoding: prefill summarizes the current sequence and stores complete chunk summaries, while decoding reuses the chunk-summary cache and appends newly completed chunks.1from dash_attn import dash_attn
2
3attn = dash_attn(
4 chunk_size=64,
5 enable_gqa=True,
6 estimate_diagonal=True,
7 scaling_factor=1.0,
8 return_active_blocks=False,
9)| Argument | Description |
|---|---|
chunk_size | Number of tokens per routed KV chunk. |
enable_gqa | Enables grouped-query attention support when query heads outnumber KV heads. |
estimate_diagonal | Includes special handling for the current or near-diagonal chunk. |
scaling_factor | Scales routing logits before sparse block selection; this is the main knob for sparsity. |
return_active_blocks | Returns the number of active routed blocks per token for sparsity analysis. |
max_chunks | Preallocated chunk-summary cache capacity used during decoding. |
sigma | Controls the strength of the Stage 1 routing prior used by Stage 2. |
[batch, heads, seq_len, head_dim] layout for queries, keys, and values; head_cls has shape [kv_heads, head_dim].dash_attn.models.llama. LlamaConfig defaults to attn_implementation="dash_attn" and adds DashAttention-specific fields such as chunk_size, estimate_diagonal, sigma, and scaling_factor.1from dash_attn.models.llama import LlamaForCausalLM
2
3model = LlamaForCausalLM.from_pretrained(
4 "fasa-org/MiniCPM-4-8B-DashAttention",
5 attn_implementation="dash_attn",
6 torch_dtype="auto",
7)return_active_blocks=True, then read model.get_active_blocks().example/run_niah.py runs a needle-in-a-haystack style generation example and reports measured sparsity.test/test_smoke.py checks the standalone DashAttention kernel wrapper.test/test_llama_dash_attn.py checks the Llama integration and active-block reporting.pytest| Model | Link |
|---|---|
| 8B-FullAttn | Hugging Face |
| 8B-InfLLMv2 | Hugging Face |
| 8B-NSA | Hugging Face |
| 8B-DashAttention | Hugging Face |
1@article{dash-attention,
2 title={DashAttention: Differentiable and Adaptive Sparse Hierarchical Attention},
3 author={Huang, Yuxiang and Gon{\c{c}}alves, Nuno M. T. and Alvetreti, Federico and Li, Lei and Han, Xu and Ponti, Edoardo M. and Martins, Andr{\'e} F. T. and Treviso, Marcos V.},
4 journal={arXiv preprint arXiv:2605.18753},
5 year={2026}
6}