This repository implements K-Sparse Attention (KSA): a simple sparse attention for Transformer language models. On small-scale experiments, KSA yields an order of magnitude reduction in the number of keys and values that need to be read per query with no hit to perplexity. This significantly speeds up long-context inference without loss in quality.

KSA strikes a good balance between hybrid models like SSMs or sliding window attention and dense quadratic attention. Like dense attention, the memory use is unbounded. This is important for in-context learning. However, like sliding window attention, the computation and memory movement is (almost) linear (a small quadratic term plus large linear term).

KSA uses a top-k operation to compute indices and scores. These scores are then used to select the relevant keys and values to attend to for each query. Hence each query attends to a fixed number of keys and values which dramatically speeds up long-context inference.

The algorithm is motivated by DeepSeek Sparse Attention (DSA) but is (in my biased opinion) better in a few ways:

• Simplicity: The algorithm is quite simple. It is a straight-forward modification of the familiar multi-headed attention (or MQA, GQA, etc). It doesn't require you to use multi-headed latent attention (MLA) which is already considerably more complicated than regular attention.

• Training: Unlike DSA, you can pretrain with KSA from scratch. You don't need any additional loss functions to constrain the sparsity. You don't need to warmup the model with dense attention. It just works.

• Compatibility: For inference, KSA works with existing fused attention operations. For training it can work with fused attention operations provided they support gradients with respect to an additive floating point mask.

Some simple Python code is the best way to describe the algorithm. The KSA operation for generation (assuming a query sequence length of one) is below:

def topk_attention(q: array, k: array, v: array, scale: float, K: int):
 """
 Inputs:
 - q: array queries with shape (B, N_Q + 1, 1, D)
 - k: array keys with shape (B, N_KV + 1, L, D)
 - v: array values with shape (B, N_KV, L, D)
 - scale: floating point scale
 - K: sparsity parameter for topk

 The dimensions are:
 - B: batch size
 - N_Q: number of query heads
 - N_KV: number of key-value heads
 - L: context length
 - D: head dimension
 """
 # q and k have one extra head used for computing top-k
 # scores and indices
 q, gate_q = split(q, [-1], axis=1) # (B, N_Q, 1, D), (B, 1, 1, D)
 k, gate_k = split(k, [-1], axis=1) # (B, N_KV, L, D), (B, 1, L, D)
 # Compute the top-k indices and scores
 scores = (gate_q * scale) @ gate_k.swapaxes(-1, -2) # (B, 1, 1, L)
 topk_scores, topk_indices = topk(scores, K) # (B, 1, 1, k)
 topk_indices = topk_indices.squeeze(2)[..., None] # (B, 1, k, 1)
 # Select top-k keys and values
 k = take_along_axis(k, topk_indices, axis=2) # (B, N_KV, K, D)
 v = take_along_axis(v, topk_indices, axis=2) # (B, N_KV, K, D)
 # Pass topk keys, values and the topk scores as an additive
 # mask to the fused implementation
 return scaled_dot_product_attention(
 q, k, v, scale=scale, mask=topk_scores,
 )

Aside from the sparse attention, the model architecture is identical to the dense Qwen 3.0 models. The pretraining corpus is DCLM 1.0. The remaining details on the pretraining recipe are easy to find in the single file pretrain.py.

The experiments below are for a 600M parameter model. The pretraining context length is 4096. Each experiments was run on an H100 (via Nebius) with the MLX CUDA back-end. Only the value of k for the number of top keys and values is varied. The baseline is using standard dense attention. I also compared against a standard sliding window attention with sliding window size of 64.

Benchmarks were run on a MacBook with the base M5 chip (10 GPU cores, 10 CPU cores). The model is quantized using 8-bits per weight. The figure below shows that as you increase context length, generation tokens-per-second degrades more slowly for KSA compared to the dense baseline.

The figure below shows the speedup of KSA relative to the dense baseline. As you increase the context length, the speedup increases.

pip install -r requirements.txt

To run on a single GPU:

python pretrain.py --config configs/base_600m.yaml

To run on multiple GPUs:

mlx.launch -n 8 pretrain.py --config configs/base_600m.yaml