MiniCPM-4-8B-NSA

<div align="center"> <h1>DashAttention</h1> <p><strong>Differentiable and Adaptive Sparse Hierarchical Attention</strong></p> </div> <div align="center" style="line-height: 1;"> <a href="https://github.com/fasa-org/dash-attention" style="margin: 2px;"> <img alt="Code" src="https://img.shields.io/badge/GitHub-100000?style=for-the-badge&logo=github&logoColor=white" style="display: inline-block; vertical-align: middle;"/> </a> <a href="https://huggingface.co/collections/fasa-org/dashattention" style="margin: 2px;"> <img alt="Hugging Face" src="https://img.shields.io/badge/DashAttention-fcd022?style=for-the-badge&logo=huggingface&logoColor=000&labelColor" style="display: inline-block; vertical-align: middle;"/> </a> <a href="https://arxiv.org/abs/2605.18753" style="margin: 2px;"> <img alt="Paper" src="https://img.shields.io/badge/Paper-2605.18753-b31b1b.svg" style="display: inline-block; vertical-align: middle;"/> </a> </div>

Installation

For the usage of DashAttention kernels and running the example, please run the following script:

pip install -e .

For benchmark environment setup, please refer to each corresponding folder.

Usage

The dash attention's interface can be used as follows:

queries = torch.randn(batch, query_heads, seq_len, head_dim, device=device, dtype=dtype).contiguous()
keys = torch.randn(batch, kv_heads, seq_len, head_dim, device=device, dtype=dtype).contiguous()
values = torch.randn(batch, kv_heads, seq_len, head_dim, device=device, dtype=dtype).contiguous()
head_cls = torch.randn(kv_heads, head_dim, device=device, dtype=dtype).contiguous()

model = dash_attn(
    chunk_size=chunk_size,
    enable_gqa=True,
    estimate_diagonal=True,
    return_active_blocks=True,
)

out, active_blocks = model(queries, keys, values, head_cls)

We also provide an example on how to use DashAttention in Llama-architecture models in here.

python ./example/run_niah.py

Documentation

DashAttention implements the attention mechanism introduced in DashAttention: Differentiable and Adaptive Sparse Hierarchical Attention. The method replaces fixed-budget top-k block routing with an adaptive, differentiable sparse router, then refines the selected regions with token-level softmax attention.

How it works

The implementation follows the three-stage hierarchy described in the paper:

1. Local chunk summarization: dash_attn.prefill.summarize_chunk and dash_attn.decoding.summarize_chunk build one learned key summary per KV chunk.
2. Entmax block routing: score_blocks computes sparse chunk supports and routing priors from query-to-summary scores.
3. Prior-induced sparse softmax: full_attn applies token-level attention only over routed chunks, using the Stage 1 prior to preserve differentiability through the hierarchy.

The public kernel wrapper is 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.

Core API

from dash_attn import dash_attn

attn = dash_attn(
    chunk_size=64,
    enable_gqa=True,
    estimate_diagonal=True,
    scaling_factor=1.0,
    return_active_blocks=False,
)

Important arguments:

Inputs are expected in [batch, heads, seq_len, head_dim] layout for queries, keys, and values; head_cls has shape [kv_heads, head_dim].

Llama integration

DashAttention includes a Llama-compatible modeling implementation in 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.

from dash_attn.models.llama import LlamaForCausalLM

model = LlamaForCausalLM.from_pretrained(
    "fasa-org/MiniCPM-4-8B-DashAttention",
    attn_implementation="dash_attn",
    torch_dtype="auto",
)

To inspect routing behavior, call the model with return_active_blocks=True, then read model.get_active_blocks().

Examples and tests

• 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.

Run the test suite with:

pytest

The current kernels require CUDA-capable hardware.

Models

We release our 8B models for reproducibility.

The base models we use are MiniCPM4-1B-Base, MiniCPM4-3B-Base, and MiniCPM4-8B-Base.

Benchmarks

• Performance: Please refer to README.

License

This project is released under the BSD-3-Clause License.

Acknowledgement

This repository is developed with the aid of RULER, OLMES, InfLLMv2, and NSA-triton.

Citation

@article{dash-attention,
  title={DashAttention: Differentiable and Adaptive Sparse Hierarchical Attention},
  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.},
  journal={arXiv preprint arXiv:2605.18753},
  year={2026}
}