Prefrontal-AI/latent-sparse-attention
Machine LearningResearch code for Latent Sparse Attention, an SGLang plugin that speeds up long-context prefill, decode and training.
What it is
Latent Sparse Attention (LSA) is a research codebase implementing the method described in Latent Sparse Attention: Towards Subquadratic Ultra-Long Context Attention. Instead of scoring every key, queries score centroids of hierarchically clustered keys in a learned latent space, then read only the selected long-range keys plus the sink and local window. The attention kernel itself is unchanged; LSA decides which keys the kernel reads, so prefill, decode and training can get cheaper at long context while accuracy stays near the original model.
Two variants are implemented. D-LSA (Distilled) adapts a pre-trained model by distilling its indexer and can run without training (a PCA map is also supported). N-LSA (Native) is trained from scratch on next-token loss and matches full attention accuracy in the reported experiments. The repository also contains the baseline implementations (full attention, SWA, NSA, PISA, block sparse), the evaluation harness, and the training and distillation pipelines behind the paper.
This is for researchers and infrastructure engineers who work with long-context serving or attention research and already have GPU and cluster resources. It is not a drop-in library for typical application developers.
Quick start
Install with Python 3.12 and uv:
git clone https://github.com/Prefrontal-AI/latent-sparse-attention.git
cd latent-sparse-attention
uv sync --locked --extra evaluation --extra training
uv run lsa --help
The training extra builds causal-conv1d and needs the CUDA toolkit; without nvcc, add --no-install-package causal-conv1d. For evaluation only, use uv sync --locked --extra evaluation (source:17).
Serving runs inside the project's SGLang fork image, which provides the context-parallel and long-context support needed for 10M-token runs. LSA itself is installed into that image rather than baked into it. Pull the container and mount this checkout, a directory of distilled maps, and your Hugging Face cache (source:18):
docker pull ghcr.io/prefrontal-ai/sglang:lsa-v0.1.0
docker run -it --rm --gpus '"device=0"' --shm-size 32g --network host -e HF_HOME=/hf \
-v $PWD:/workspace/latent-sparse-attention:ro -v <maps directory>:/maps:ro -v $HF_HOME:/hf \
ghcr.io/prefrontal-ai/sglang:lsa-v0.1.0 bash
Docker mount sources must be real directories, not symlinks. Inside the container, install the checkout and launch the server with an lsa block (source:21):
pip install --no-deps /workspace/latent-sparse-attention
python -m sglang.launch_server --model-path Qwen/Qwen3.8-27b --tp 1 \
--attention-backend fa3 --page-size 1 --disable-prefill-cuda-graph \
--json-model-override-args '{"lsa": {"projections": "/maps/qwen38_27b.pt", "query_heads": "all"}}'
--disable-prefill-cuda-graph is required because key selection during prefill happens on the host and cannot be captured by a CUDA graph; without it the server crashes at startup. Wait for the server log line armed and for /health, then test with a request:
curl -s localhost:30000/v1/completions -H 'Content-Type: application/json' \
-d '{"model": "Qwen/Qwen3.8-27b", "prompt": "The capital of France is", "max_tokens": 32}'
Without the armed line, the server is running the checkpoint's own attention. The projections file is a distilled map; use the d-lsa-weights-v1 release or train one via the distillation pipeline. Distributed runs use the same repository on Slurm with lsa --site <name>; the site's root, datasets and Hugging Face cache are set in a YAML file (source:49).
Constraints and requirements
- Serving requires the project's SGLang fork and matching container image; the
lsa-v0.1.0fork tag is used for all served runs (source:7). B300 (Blackwell) nodes need the separately tagged imagelsa-v0.1.0-b300. - The paper's results were measured on NVIDIA H200 GPUs. D-LSA with Qwen3.8-27B reported one RULER evaluation averaging 0.974 at 4k and 0.968 at 32k over 13 tasks of 100 rows, taking 1 hour 52 minutes on one H200 (source:28). These are paper measurements, not guarantees for other hardware or models.
- Distillation and native 128k decay need prepared data: LongDocs assembled with
lsa data prepare, agent traces from thetraces-v1release, and the RULER sources (source:37). The native checkpoints release ships weights trained with the paper's own trainer, which is not part of this repository; tensor names have been remapped to this codebase and a reload checks the mapping (source:48). - The license is Apache-2.0, with third-party code and datasets keeping their own licenses (source:58).
HotRepo's take
The repository appeared in HotRepo's new-project list on 2026-11-03 with 133 stars and 18 stars gained in roughly one hour, plus one fork; it has not been recorded on the all-language daily GitHub rank. With such a short observation window the numbers show early attention, not sustained adoption, and there is no long trend to read from this data. The technical substance is the draw: a paper with released distilled maps, ten native checkpoints, and an SGLang plugin path, but all of it assumes an H200-class GPU environment and an SGLang fork, so the primary audience will be attention researchers rather than general application teams.
Written by AI from the project's README, updated . The project's own docs are the source of truth.