The Problem Training or fine‑tuning very long‑context language models quickly runs into quadratic attention cost and memory limits. Existing sparse‑attention tricks either require hand‑crafted chunking or break end‑to‑end learning, making it hard to get stable ultra‑long context performance.

What This Does The repository implements HiLS‑Attention – a learned, chunk‑wise sparse attention that estimates chunk importance with compressed keys and factorises the attention into inter‑chunk and intra‑chunk softmaxes. The core logic lives in

  • models/FlashHiLS/hils_attention.py – key class HiLSAttention and helper chunk_attn_pool_tilelang.
  • models/FlashHiLS/modeling_olmo_hils.py / modeling_qwen_hils.py – model wrappers that plug the attention into OLMo‑3 and Qwen‑3 families.

Configuration files under configs/hils_attention/ (e.g., config_hils_attn_8KA2K_HoPE_345M_prop3p1_qcal_r64.json) define chunk sizes, token‑landmark settings and training hyper‑parameters. Evaluation scripts in eval/ (e.g., eval_longbench_v1.py) load a checkpoint, run the model on the LongBench benchmark, and write per‑task results.

How It Is Wired Execution starts at the entry point eval/eval_longbench_v1.py:main (line 648). main parses CLI args, calls load_model (which imports models/FlashHiLS/modeling_olmo_hils.py), then iterates over benchmark datasets. For each example it invokes forward on the model; forward routes through apply_rotary_pos_emb → rotate_half, then to hils_attention.forward. Inside hils_attention.forward the most‑used internal calls are:

  • chunk_attn_pool_tilelang (12 call sites) – builds chunk‑wise attention tensors.
  • rms_norm (6 call sites) – normalises hidden states.
  • assert_close (13 call sites) – sanity‑checks intermediate tensors.

The call graph shows backward is invoked from 26 distinct locations, indicating that gradient‑flow code is widely shared. No circular imports were detected, but the modules models/FlashHiLS/modeling_olmo_hils.py and ops/flex_attn_tilelang.py each have >10 inbound edges, making them high‑impact change points.

External effects are limited to filesystem writes: main → append_summary_log creates output directories via os.makedirs. No network or database calls are present.

How To Use It

# Clone the original upstream repo
git clone https://github.com/moses-y/HiLS-Attention
cd HiLS-Attention

# Recommended environment (uv)
curl -LsSf https://astral.sh/uv/install.sh | sh
uv sync
source .venv/bin/activate

# Or pip/conda alternative (from README)
conda create -n hils python=3.11 -y && conda activate hils
pip install torch==2.8.0 torchvision==0.23.0 torchaudio==2.8.0 \
            -f https://download.pytorch.org/whl/cu128
pip install -r requirements.txt

To evaluate a pre‑trained checkpoint:

export MODEL_PATH=/path/to/hf_ckpt
export OUTPUT_DIR=./eval_results
python eval/eval_longbench_v1.py \
    --model_path $MODEL_PATH \
    --config configs/hils_attention/config_hils_attn_8KA2K_HoPE_345M_prop3p1_qcal_r64.json \
    --output_dir $OUTPUT_DIR

Training scripts are in scripts/pretrain/ and scripts/cpt/; they expect CORPUS_PATH and OUTPUT_DIR environment variables as shown in the README.

Real‑World Use A company that serves legal‑document search over millions of pages can replace the default dense transformer in its inference service with the HiLS‑enabled OLMo‑3 checkpoint. By loading the checkpoint (converted to HF format with scripts/ckpt_transfer/dcp_hf_transfer.sh) and calling the same forward API, they gain roughly 2× faster inference on 64 k token prompts while preserving accuracy on short‑context QA tasks.

Code Health & Issues

  • High – License missing – no LICENSE file; redistribution rights undefined.
  • High – No CI pipeline – 76 source files, no .github/workflows or similar.
  • Medium – Dependabot absent – only pyproject.toml/requirements.txt; no automated dependency updates.
  • Medium – Large binary in repo – configs/olmo3_vocab/tokenizer.json (6.8 MiB) should be moved to Git LFS or external storage.

Measured quality concerns (static analysis):

  • High – Deep nesting in ops/flex_attn_tilelang.py, models/FlashHiLS/modeling_olmo_hils.py, ops/hils_fwd_bwd_head.py (max depth 10).
  • High – Oversized files (ops/flex_attn_tilelang.py, models/FlashHiLS/modeling_olmo_hils.py, models/FlashHiLS/modeling_qwen_hils.py > 1700 LOC).
  • High – Duplicated code blocks across several eval/configs/datasets/*.py files.
  • Medium – Broad exception handling (bare except:) in models/FlashHiLS/hils_attention.py and others.
  • Medium – High branching density in several shell scripts under scripts/eval/.

Addressing the high‑severity items (license and CI) should be the first priority; refactoring large/duplicated modules will improve maintainability for future extensions.

The Bottom Line HiLS‑Attention provides a functional, end‑to‑end implementation of learned sparse attention with clear entry points for training and evaluation. The code works but suffers from maintainability problems (deep nesting, huge files) and lacks basic engineering hygiene (license, CI, dependency lock). It is suitable for research teams or early‑stage product prototypes that can allocate effort to clean‑up and add proper release processes.