• [2026.05.21] We released Jet-Long.

• 2. Installation
• 3. Quick Evaluation (Qwen3-1.7B-Base)
• 4. Shipped methods
• 5. How a method is composed on top of a Qwen3 Base
• 6. Hardware Requirements
• 7. Project Structure
• 8. Acknowledgements
• 9. Contact
• 10. Contributing
• 11. License
• 12. BibTeX

git clone https://github.com/jet-ai-projects/jet-long.git
cd jet-long

There are two installation profiles, both on Python 3.12 + torch 2.9.1+cu130:

• Default (FlashAttention-2) — Steps 1–3 below. Supports every method (base, yarn, dntk, selfext, dca, jetlong, jetlong_freq).
• Fused CuTe kernel (optional, SM90/H100 only) — the Fused CuTe kernel environment section below. A separate conda env for the jetlong_fused method.

The only out-of-band dep is the FlashAttention build, which must be linked against the same CUDA toolkit as torch.

conda create -n jtl python=3.12 -y
conda activate jtl

pip install -e . reads pyproject.toml and pulls torch 2.9.1, transformers 5.3.0, datasets 4.6.1, accelerate, wandb, nvidia-ml-py, lm_eval (jetlm-ai fork pinned to the jetai tag), and the RULER/PG-19 metric deps (jieba, rouge, etc.):

pip install -e .

If your system CUDA major version doesn't match the torch wheel (the default pip wheel ships cu12, but H100 / Blackwell hosts often expose CUDA 13), force a matching torch wheel before building flash-attn:

# Example: cu13 host
pip install torch==2.9.1 --index-url https://download.pytorch.org/whl/cu130 --force-reinstall

Install or build flash-attn 2.8.3 against your torch + CUDA stack — either a matching pre-built wheel from the Dao-AILab releases page, or a from-source build of the v2.8.3 tag. If you build from source, set FLASH_ATTENTION_FORCE_BUILD=TRUE so pip actually compiles instead of falling back to a stale pre-built wheel.

Note: flash-attn==2.8.3 is required. The bleeding-edge 2.8.4 HEAD build has a small-q NaN regression that silently corrupts decode. See requirements.txt for the pin.

The jetlong_fused method runs JetLong through a fused FlashAttention-4 / CuTe SM90 kernel (jetlm/kernels/) instead of the FA2 A+B−C merge. It needs a separate conda env with the CuTe-DSL runtime plus flash_attn_4 (which provides flash_attn.cute), on the same torch 2.9.1+cu130:

conda create -n jtl-fused python=3.12 -y
conda activate jtl-fused
pip install -e ".[fused]"     # base deps + nvidia-cutlass-dsl + quack-kernels + cuda-python + ...
pip install torch==2.9.1 --index-url https://download.pytorch.org/whl/cu130 --force-reinstall
# flash_attn_4 ships flash_attn.cute (beta) — install 4.0.0b10 (the block-sparse
# prefill API the kernel relies on changed in b11):
pip install flash_attn_4==4.0.0b10

⚠️ This environment runs only jetlong_fused. flash_attn_4 provides flash_attn.cute but not the FA2 forward kernels, so every other method — including the non-fused jetlong, yarn, dntk, selfext, dca, jetlong_freq — will not run here; they require the default flash-attn 2.8.3 environment above. Hard requirements: NVIDIA SM90/H100 + CUDA 13. jetlong_fused is long-context-only (decode context must exceed w_0); use plain jetlong for short contexts.

Validate the kernel against reference attention with pytest tests/test_fa4_cute_kernels.py and python scripts/throughput/check_cute_jetlong_parity.py --config-path model_cache/Qwen3-1.7B-Base-jetlong_fused (run scripts/model/setup_model_variants.sh first so the modeling file is symlinked adjacent to the config). Load jetlong_fused models with attn_implementation="sdpa" (the fused kernel is invoked explicitly on the extended path; SDPA covers the within-32K-window path).

Measure prefill throughput and CUDA-graphed decode-kernel latency. The jetlong_fused commands run in the fused env; the FA2 baselines run in the default env (the two environments are mutually exclusive), so each comparison is two processes:

# Prefill throughput (full-forward tok/s) — fused vs FA2 baseline, across lengths:
python scripts/throughput/bench_prefill_throughput.py --model_path model_cache/Qwen3-8B-Base-jetlong_fused --attn sdpa            # fused env
python scripts/throughput/bench_prefill_throughput.py --model_path model_cache/Qwen3-8B-Base --attn flash_attention_2            # default env

# Generation throughput (decode-kernel latency, CUDA-graphed) — fused vs single-FA2-decode baseline:
python scripts/throughput/bench_cute_decode_sm90_jetlong.py --cuda-graph --nheads 32 --nheads-kv 8 --head-dim 128                # fused env
python scripts/throughput/bench_fa2_decode_baseline.py --cuda-graph --nheads 32 --nheads-kv 8 --head-dim 128                     # default env

Prefill should get ~1.24–1.31× FA2 speedup past 32K.

Generation throughput shows no loss vs FA2 at the attention-kernel level — the fused decode measures at parity with or faster than a single FA2 decode (6–15% faster across 32K–131K on H100), and attention is a very small portion of the per-token cost, so end-to-end is ≈1.0× FA2 (note that measuring full end-to-end generation throughput directly requires manually capturing a full CUDA graph of the decode loop).

Below is the complete from-scratch evaluation of the Qwen3-1.7B + jetlong row on all three benchmarks. All commands are run from the repo root with the python env activated.

bash scripts/model/setup_model_variants.sh

RULER-500 and PG-19-subsample are hosted on the HF Hub and downloaded into processed_data/ via a single dispatcher:

bash scripts/data/download_data.sh

# One launcher per (size, method) under scripts/eval/ruler/
bash scripts/eval/ruler/qwen1d7B_base_jetlong.sh

MODEL_PATH=model_cache/Qwen3-1.7B-Base-jetlong bash scripts/eval/ppl/pg19_ppl_local.sh

First clone the upstream HELMET repo + data:

bash scripts/eval/helmet/set_up_helmet.sh
bash scripts/eval/helmet/run_helmet.sh Qwen3-1.7B-Base-jetlong

jetlong_fused reuses the model_cache and data from Steps 1–2, but runs RULER from the fused conda env (jtl-fused, see Section 2) because the kernel needs flash_attn.cute:

conda activate jtl-fused
bash scripts/eval/ruler/qwen1d7B_base_jetlong_fused.sh

This mirrors Step 3 with three fused-specific settings: it loads with attn_implementation="sdpa" (auto-detected from the architecture; the CuTe kernel takes over once the context passes the 32K window), runs at --eval_batch_size 8 so the lengths that exercise the kernel (≥32K, where it rejects left-padded batches) fall to batch 1, and sets --prefill_chunk_size 16384 to keep the 131072 prefill within one H100. RULER scores match the non-fused jetlong method — identical at ≤32K (both reduce to plain causal attention there) and lossless at the 64K/128K lengths the kernel handles.

Both jetlong and jetlong_freq share a 3-call FlashAttention prefill (near base RoPE + distant correction + windowed correction, merged via stabilized fp32 LSE); they differ in how the distant phase is computed (integer position grouping vs YaRN frequency interpolation). See jetlm/modeling/jetlong/modeling_qwen3_jetlong.py and jetlm/modeling/jetlong_freq/modeling_qwen3_jetlong_freq.py for details.

Every variant is a standard HuggingFace model directory at model_cache/<variant>/. The trick is that we never copy a 4 B / 8 B checkpoint: each variant is materialized as a directory of symlinks by scripts/model/setup_model_variants.sh, which sources three things:

The split keeps the per-method implementation in one place (jetlm/modeling/<method>/) while letting each (size, method) pair carry its own small config.json for the architecture dimensions and method knobs that differ per size.

Config-only methods (yarn, dntk) use HF's stock rope_scaling field and the unmodified Qwen3ForCausalLM class — no custom modeling code:

{
  "architectures": ["Qwen3ForCausalLM"],
  "rope_scaling": {"rope_type": "yarn", "factor": 4.0,
                   "original_max_position_embeddings": 32768},
  "rope_theta": 1000000, "max_position_embeddings": 131072,
  // ...base dims unchanged...
}

Methods with custom modeling code (selfext, dca, jetlong, jetlong_freq) add auto_map so HF's from_pretrained(..., trust_remote_code=True) picks up the sibling modeling file at load time, plus a method-specific knob block:

{
  "architectures": ["Qwen3JetLongForCausalLM"],
  "auto_map": {
    "AutoModelForCausalLM": "modeling_qwen3_jetlong.Qwen3JetLongForCausalLM"
  },
  "jetlong": {"w": 32768, "w_0": 2048, "yarn_alpha": 1, "yarn_beta": 32},
  "rope_scaling": null,
  // ...base dims unchanged...
}

The string "modeling_qwen3_jetlong.Qwen3JetLongForCausalLM" is a bare file name plus class name — not a Python import path. HF expects the file adjacent to config.json at load time, which is exactly what the modeling symlink provides.

After bash scripts/model/setup_model_variants.sh, a variant like Qwen3-1.7B-Base-jetlong becomes a flat directory of symlinks:

model_cache/Qwen3-1.7B-Base-jetlong/
├── config.json                → ../../model_configs/Qwen3-1.7B-Base-jetlong/config.json
├── modeling_qwen3_jetlong.py  → ../../jetlm/modeling/jetlong/modeling_qwen3_jetlong.py
├── model.safetensors          → ../Qwen3-1.7B-Base/model.safetensors
├── tokenizer.json             → ../Qwen3-1.7B-Base/tokenizer.json
└── ... (other tokenizer + generation_config + LICENSE files)

From any caller — our eval scripts, the upstream HELMET driver, or a plain Python REPL — loading is just:

from transformers import AutoModelForCausalLM, AutoTokenizer
model = AutoModelForCausalLM.from_pretrained(
    "model_cache/Qwen3-1.7B-Base-jetlong",
    trust_remote_code=True,
    attn_implementation="flash_attention_2",
    torch_dtype="bfloat16",
)
tokenizer = AutoTokenizer.from_pretrained("model_cache/Qwen3-1.7B-Base-jetlong")

HF reads config.json, sees auto_map, copies the sibling modeling_qwen3_jetlong.py into its trust-remote-code module cache, imports Qwen3JetLongForCausalLM, and loads the base weights into it. Because we only overrode rotary embedding + attention forward and kept every other module verbatim, the pretrained safetensors load cleanly with no key mismatches.

Each jetlm/modeling/<method>/ subpackage subclasses HF's vanilla Qwen3 modules and overrides only what the method needs. For example jetlm/modeling/jetlong/modeling_qwen3_jetlong.py exports 5 classes (Qwen3JetLongAttention, Qwen3JetLongDecoderLayer, Qwen3JetLongForCausalLM, Qwen3JetLongModel, Qwen3JetLongRotaryEmbedding) — every other building block (Qwen3MLP, Qwen3RMSNorm, etc.) is imported unchanged from transformers.models.qwen3.modeling_qwen3. That's why every method is training-free: the weight shapes are byte-identical to the upstream base checkpoint.

Jet-Long is inference-only (no fine-tuning), and we run all evaluation on NVIDIA H100 GPUs. The default experiments in this repo are launched on a single 8×H100 node for fast iteration; the launchers under scripts/eval/ discover the visible GPU count via nvidia-smi (or CUDA_VISIBLE_DEVICES) and shard across whatever is present, so the same scripts work on smaller node counts at the cost of wallclock time.

jetlm/
  constant.py               # repo-wide version-string + W&B result-log suffix
  evaluation/
    meta_eval.py            # entry for RULER + lm-eval-harness configs
    ruler/                  # RULER evaluator (custom)
    ppl/                    # PG-19 anchored growing-window PPL
    helmet/                 # HELMET sharded driver + result aggregator + compat patches
    lm_eval_harness/        # wrapper around lm-eval harness
    configs/                # per-task YAMLs (ruler.yaml + pg19_ppl.yaml)
  kernels/                  # fused FA4 / CuTe SM90 kernels for jetlong_fused (+ csrc/ JIT-built merge)
  modeling/
    wrapper.py              # thin HFModel generate wrapper
    {dca,selfext,jetlong,jetlong_freq,jetlong_fused}/  # one subpackage per method with custom modeling
  utils/                    # distributed, atomic I/O, complete-markers
model_configs/              # per-(size, method) config.json (no modeling code)
scripts/
  data/                     # PG-19 + RULER data prep (download_data.{py,sh}, upload_*_hf.py, subsample/generate)
  model/                    # variant setup + base-model download
  eval/
    ruler/                  # per-(size, method) RULER launcher
    ppl/                    # PG-19 PPL launchers
    helmet/                 # HELMET launcher, GPT-4 judges, set_up_helmet.sh
  throughput/               # fused-vs-FA2 throughput benches + parity check (prefill, decode-kernel latency, CuTe-vs-ref parity)
tests/                      # FA4 CuTe kernel parity tests (test_fa4_cute_kernels.py)

We thank the authors of RULER, PG-19, and HELMET for the long-context benchmarks Jet-Long is evaluated on, and the Qwen3 team for the open-weights base models we extend.

• Han Cai
• Haozhan Tang
• Song Han

Please see CONTRIBUTING.md for contribution guidelines and Developer Certificate of Origin (DCO) sign-off requirements.

This codebase is released under the Apache License 2.0. Third-party dependency notices are listed in THIRD_PARTY_NOTICES.md.

If you find Jet-Long useful, please cite:

@misc{jetlong2026,
  title={Jet-Long: Efficient Long-Context Extension with Dynamic Bifocal RoPE},
  author={Haozhan Tang and Zerui Wang and Yuxian Gu and Song Han and Han Cai},
  year={2026},
  eprint={2607.07740},
  archivePrefix={arXiv},
  primaryClass={cs.LG},
  url={https://arxiv.org/abs/2607.07740},
}