SinkRouter is a training-free inference framework for accelerating long-context decoding in large language models (LLMs) and large multimodal models (LMMs). It uses the attention-sink signal to route decoding work at KV-head-group granularity: sink groups can bypass the long KV-cache scan, while active groups use a hardware-aware Triton attention kernel.
The implementation is based on the accompanying paper, SinkRouter: Sink-Aware Routing for Efficient Long-Context Decoding in Large Language and Multimodal Models. The repository contains model patches and evaluation scripts. Model checkpoints and benchmark datasets are not included; provide their local paths in the commands below.
During prefill, SinkRouter keeps the standard attention path and records the position-0 key used for routing. During single-token decoding, the implementation:
- compares the RoPE query with the cached position-0 key;
- routes each GQA KV group before scanning the long cache;
- returns the position-0 value for routed sink groups;
- computes active groups with tiled QK/PV and online softmax;
- uses Split-K parallelism and an in-kernel reduction for long contexts.
The main implementation is in src/models/optimized.py and src/models/triton/triton_gqa_simple.py. The multimodal patch is in src/models/vlmoptimized.py.
SinkRouter keeps the standard prefill path, routes each decode-time head group from its similarity to the position-0 key, and bypasses long KV-cache loading for sink groups. Active groups follow the standard attention path before their outputs update the residual stream.
.
├── docs/figures/
│ └── sinkrouter_overview.jpg method overview figure
├── src/models/
│ ├── optimized.py Llama attention replacement
│ ├── vlmoptimized.py LLaVA patch
│ └── triton/triton_gqa_simple.py Triton GQA kernels
├── scripts/
│ ├── acc_eval/
│ │ ├── eval_milebench.py
│ │ ├── eval_mmvp_baseline.py
│ │ ├── eval_mmvp_sinkrouter.py
│ │ ├── eval_cvbench_baseline.py
│ │ └── eval_cvbench_sinkrouter.py
│ └── efficient_eval/
│ ├── decode_benchmark.py
│ ├── bench_128k.py
│ ├── bench_256k.py
│ └── bench_512k.py
└── results.py/ generated local evaluation outputs
The commands below assume Linux, Python 3.10+, a CUDA-enabled PyTorch installation, and a GPU supported by Triton.
git clone https://github.com/streamer-AP/SinkRouter.git
cd SinkRouter
python3 -m venv .venv
source .venv/bin/activate
python -m pip install --upgrade pip
python -m pip install torch torchvision transformers accelerate triton
python -m pip install numpy pandas pillow tqdm rouge pycocotools
python -m pip install flash-attn --no-build-isolationInstall a PyTorch build matching the CUDA driver on your machine before installing flash-attn. The multimodal scripts use the Hugging Face LLaVA processor; the text-only speed scripts use a Llama checkpoint.
The scripts accept local paths, so no checkpoint or dataset is downloaded by the repository. Typical paths are:
export LLAMA_MODEL=/path/to/llama-3.1-8b
export LLAVA_MODEL=/path/to/llava-1.5-7b
export MILEBENCH_DATA=/path/to/MileBench
export MMVP_DATA=/path/to/MMVP
export CVBENCH_DATA=/path/to/CV-BenchUse --help on any script to inspect all available arguments.
Run the decode benchmark at the three context lengths used by the release. Each command compares the baseline model with SinkRouter and writes a JSON result under results.py/efficiency/ unless --output is supplied.
python scripts/efficient_eval/bench_128k.py \
--model-path "$LLAMA_MODEL" \
--device cuda:0 \
--dtype bfloat16 \
--decode-steps 50 \
--threshold 0.5 \
--output results.py/efficiency/decode_128k.json
python scripts/efficient_eval/bench_256k.py \
--model-path "$LLAMA_MODEL" \
--device cuda:0 \
--dtype bfloat16 \
--decode-steps 50 \
--threshold 0.5 \
--output results.py/efficiency/decode_256k.json
python scripts/efficient_eval/bench_512k.py \
--model-path "$LLAMA_MODEL" \
--device cuda:0 \
--dtype bfloat16 \
--decode-steps 50 \
--threshold 0.5 \
--output results.py/efficiency/decode_512k.jsonThe lower-level runner is useful for a custom context length:
python scripts/efficient_eval/decode_benchmark.py \
--model-path "$LLAMA_MODEL" \
--context-length 131072 \
--device cuda:0 \
--dtype bfloat16 \
--decode-steps 50 \
--warmup-steps 1 \
--chunk-size 32768 \
--threshold 0.5 \
--output results.py/efficiency/custom_decode.jsonOptional Triton launch parameters can be swept without editing the kernel:
SINK_BLOCK_K=128 SINK_NUM_WARPS=4 SINK_NUM_STAGES=2 SINK_NUM_SPLITS=32 \
python scripts/efficient_eval/decode_benchmark.py \
--model-path "$LLAMA_MODEL" \
--context-length 524288 \
--device cuda:0 \
--dtype bfloat16The MileBench runner supports baseline generation and SinkRouter with --use_sink_router. The --dataset_name value must match a dataset directory/name available under the supplied MileBench data root.
python scripts/acc_eval/eval_milebench.py \
--model_path "$LLAVA_MODEL" \
--data_dir "$MILEBENCH_DATA" \
--dataset_name ActionLocalization \
--output_dir results.py/milebench/llava15_7b/sink_router \
--max_new_tokens 512 \
--use_sink_router \
--overwriteFor the baseline, omit --use_sink_router and write to a different output directory:
python scripts/acc_eval/eval_milebench.py \
--model_path "$LLAVA_MODEL" \
--data_dir "$MILEBENCH_DATA" \
--dataset_name ActionLocalization \
--output_dir results.py/milebench/llava15_7b/baseline \
--max_new_tokens 512 \
--overwriteThe MMVP data directory must contain Questions.csv and MMVP_Images/. Run the optimized and baseline scripts separately:
python scripts/acc_eval/eval_mmvp_sinkrouter.py \
--model-path "$LLAVA_MODEL" \
--directory "$MMVP_DATA" \
--answers-file results.py/mmvp/llava15_7b_sinkrouter.jsonl
python scripts/acc_eval/eval_mmvp_baseline.py \
--model-path "$LLAVA_MODEL" \
--directory "$MMVP_DATA" \
--answers-file results.py/mmvp/llava15_7b_baseline.jsonlThe CV-Bench data directory must contain test_2d.jsonl, test_3d.jsonl, and the corresponding img/2D/ and img/3D/ folders.
python scripts/acc_eval/eval_cvbench_sinkrouter.py \
--model-path "$LLAVA_MODEL" \
--data-dir "$CVBENCH_DATA" \
--subsets 2D 3D \
--answers-file results.py/cvbench/llava15_7b_sinkrouter.jsonl
python scripts/acc_eval/eval_cvbench_baseline.py \
--model-path "$LLAVA_MODEL" \
--data-dir "$CVBENCH_DATA" \
--subsets 2D 3D \
--answers-file results.py/cvbench/llava15_7b_baseline.jsonlThe main routing controls are configured on SinkRouterAttention and can be changed through environment variables or the benchmark arguments:
export SINK_THRESHOLD=0.5
export SINK_LAYER_START=0
export SINK_LAYER_END=31SINK_THRESHOLD controls the cosine routing threshold. SINK_BLOCK_K, SINK_NUM_WARPS, SINK_NUM_STAGES, and SINK_NUM_SPLITS control Triton launch configuration. Unsupported input shapes automatically fall back to the standard attention path.
Efficiency runs write JSON summaries containing context length, baseline and SinkRouter latency, throughput, sink ratio, and speedup. MMVP and CV-Bench write JSONL predictions plus a _summary.txt file. MileBench writes the per-dataset generation files under the selected output directory.
If you use this code, please cite the SinkRouter paper:
@inproceedings{liu2026sinkrouter,
title = {SinkRouter: Sink-Aware Routing for Efficient Long-Context Decoding in Large Language and Multimodal Models},
author = {Liu, Junnan and Liu, Xinyan and Gao, Peifeng and Qi, Zhaobo and Zhang, Beichen and Zhang, Weigang and Chan, Antoni B.},
booktitle = {Proceedings of the 34th ACM International Conference on Multimedia},
year = {2026}
}The code in this repository is released under the MIT License. See LICENSE for the full license text.
Benchmark datasets and model checkpoints remain subject to their original licenses and are not redistributed here.