Jet-Long: Efficient Long-Context Extension with Dynamic Bifocal RoPE

Modern LLMs are increasingly deployed in long-context applications such as retrieval-augmented generation, repository-level coding, and agentic workflows whose accumulated reasoning and tool traces routinely push the input an order of magnitude past the pretraining window, making zero-shot context extension the dominant deployment path for open-weight checkpoints. Most existing zero-shot methods fix a single rescaling factor up front, so an aggressive factor sacrifices short-context fidelity while a conservative one breaks down at long contexts. We propose Jet-Long, a tuning-free zero-shot method that pairs a local RoPE-faithful window with a long-range window whose rescaling factor adapts dynamically to the current sequence length, recovering the base model exactly at short inputs while extrapolating cleanly at long ones. An inclusion-exclusion attention merge and an on-the-fly RoPE correction rotation make the bifocal construction essentially free at inference; fused into a single CuTe kernel, long-context prefill reaches up to $1.39\times$ FA2 throughput on H100 (approaching the Hopper-only FA4), and single-batch generation incurs $\le 4\%$ overhead at every length. On Qwen3-1.7B/4B/8B up to 128K context, Jet-Long leads RULER by $+4.79$/$+2.18$/$+2.03$~pp over the strongest baseline at 1.7B/4B/8B, achieves the best overall accuracy on HELMET-RAG (a benchmark identified by HELMET as the most efficient predictor of downstream long-context performance) and attains the lowest PG-19 perplexity. Jet-Long also generalizes to hybrid attention architectures such as Jet-Nemotron for further long-context improvement without retraining, and remains hyperparameter-resilient for ease of deployment.

한 줄 요약

Jet-Long은 RoPE 기반의 동적 이중 윈도우 구조를 통해 128K 컨텍스트까지 정확도와 추론 효율성을 동시에 확보하는 제로샷 컨텍스트 확장 기법이다.

핵심 기여도

**Jet-Long**이라는 튜닝이 필요 없는 제로샷 컨텍스트 확장 방법을 제안.
**로컬 RoPE-신실 윈도우**와 **동적 리스케일링을 적용한 원격 윈도우**를 결합하여 짧은 입력에서는 기존 모델과 동일한 성능, 긴 입력에서는 정확한 외삽을 달성.
**Inclusion-Exclusion Attention Merge**와 **On-the-fly RoPE Correction Rotation**을 통해 추론 비용을 최소화.
**CuTe 커널**로 통합하여 H100에서 최대 1.39× FA2 처리량 달성, 생성 단계에서 최대 4% 오버헤드만 발생.

핵심 아이디어

Jet-Long은 기존 제로샷 방법들이 고정된 리스케일링 인수를 사용하여 짧은 컨텍스트 정확도와 긴 컨텍스트 확장 사이에 트레이드오프를 강요하는 문제를 해결한다. 이에 반해 Jet-Long은 **현재 시퀀스 길이에 따라 리스케일링 인수가 동적으로 조정되는 원격 윈도우**를 도입한다. 이는 RoPE 기반의 위치 인코딩을 유지하면서도 긴 컨텍스트에서도 모델의 예측력을 보존한다. 또한, **Inclusion-Exclusion Attention Merge**를 통해 로컬과 원격 윈도우의 어텐션을 병합하고, **On-the-fly RoPE Correction Rotation**을 통해 KV 캐시를 변경하지 않으면서 정확한 위치 회전을 보장한다. 이 구조는 추론 시 비용을 거의 증가시키지 않으면서도 정확도를 향상시킨다.

기술적 접근법

**로컬 윈도우**: RoPE-신실한 윈도우로, 기존 훈련 시 컨텍스트 내의 위치 정보를 그대로 유지.
**원격 윈도우**: 리스케일링 함수 $ f(\cdot) $를 통해 위치를 훈련 범위 내로 매핑. 시퀀스 길이 $ L $이 훈련 윈도우 $ w_{\text{pretrained}} $ 이하일 경우 $ f(x) = x $로 기존 모델과 동일한 결과를 보장.
**Inclusion-Exclusion Attention Merge**: 로컬과 원격 윈도우의 어텐션을 세 번의 FlashAttention 패스로 병합.
**On-the-fly RoPE Correction Rotation**: KV 캐시를 변경하지 않고 실시간으로 RoPE 회전을 보정.
**CuTe 커널**: 위 기법을 단일 커널로 통합하여 H100에서 최대 1.39× FA2 처리량 달성.

주요 결과

**Qwen3-1.7B/4B/8B**에서 128K 컨텍스트까지 테스트:
RULER에서 기존 최고 기준 대비 +4.79 / +2.18 / +2.03 pp 개선.
HELMET-RAG에서 최고 정확도 달성.
PG-19 퍼플렉시티에서 최저 기록.
**H100 GPU**에서:
Long-context prefill 처리량 1.39× FA2 (Hopper FA4에 근접).
생성 단계에서 모든 길이에서 최대 4% 오버헤드 발생.

의의 및 한계

Jet-Long은 제로샷 컨텍스트 확장에서 **튜닝 없이도 짧은 컨텍스트 정확도와 긴 컨텍스트 확장력을 동시에 유지**할 수 있는 점에서 혁신적이다. 특히, **하이퍼파라미터 $ w_0 $**에 대한 내성이 높아 배포 용이성을 높인다. 또한, **Jet-Nemotron**과 같은 하이브리드 어텐션 아키텍처에도 일반화 가능하다는 점에서 실용적 가치가 있다. 그러나, **RoPE 기반 모델에만 적용 가능**하며, **softmax-attention 확산 문제는 다른 아키텍처와의 결합이 필요**하다는 한계가 있다.

실용적 활용

Jet-Long은 **소프트웨어 레포지토리 분석**, **대규모 문서 QA**, **агент 기반 워크플로우** 등 긴 컨텍스트를 요구하는 LLM 응용 분야에 즉시 적용 가능하다. 특히, **Qwen3-1.7B 이상의 모델**에서 뛰어난 성능을 보이므로, **산업 현장의 대규모 컨텍스트 처리 시스템**에 유용하게 활용될 수 있다.