open-jev-fast

Open-Jev-27B inference in 20.1 ms instead of 257.0 ms on one local NVIDIA B300. Same model, bf16, a drop-in replacement for Open-Jev's local server.

Code
Latency of one inference as optimizations are added

Example request: 3 questions, 7 candidates, 539 tokens. Median latency of the forward pass plus scoring head.

TL;DR

An inference backend for Open-Jev[1]. It runs locally, keeps the model and the request format of Open-Jev's local server, and replaces the forward pass with fused CUDA kernels, a prefix tree that computes shared prompt text once, and CUDA Graphs.

12.8×
257.0 → 20.1 ms
vs. default PyTorch path
5.1×
102.9 → 20.1 ms
vs. PyTorch + FLA kernels
13.7×
703 → 51 ms
JevBench mean latency, HTTP
197 / 231
original: 198 / 231
JevBench accuracy; one coin flip

How it works

1

Fused CUDA kernels

Six hand-written kernels replace every non-matmul op in a layer. Kernels per request drop from ~5000 to 945, and the CPU launch bottleneck goes away.

2

Two-level prefix tree

The shared context is computed once, each question once, and each candidate only its tail, all in one packed forward pass.

3

Tuned GEMMs + CUDA Graphs

Merged weights, split-K with fp32 partial sums, per-shape cuBLASLt algorithms, and one CUDA Graph per request shape.

Two-level prefix tree for the example request

Prefix tree for the example request. Linear-attention layers still run each candidate over its full path, so results match the original; full-attention layers use an ancestor mask.

JevBench

Per-task latency distribution on JevBench

Cumulative share of the 231 public JevBench[11] tasks finished within a given latency, sent over HTTP to the local server, one request at a time.

Original jev.server (with FLA)open-jev-fast
Accuracy198 / 231197 / 231
Mean latency703 ms51.3 ms
P50 / P95171 / 1272 ms24.2 / 148 ms
Slowest task16.0 s0.58 s

Correctness

  • Kernels match the PyTorch/FLA reference bit for bit on 98–100% of elements.
  • Probabilities differ from the original (with FLA) by at most 0.0034 on the example request and 0.035 on a 10,722-token request; decisions are unchanged.
  • On JevBench, accuracy changes by one task, a coin flip (original: 0.504 vs 0.496).

Setup and limits

  • Latency is the forward pass plus scoring head for the example request, with prompts tokenized beforehand. All three numbers use the same script: median of 30 runs after 5 warm-ups.
  • Tested only on an NVIDIA B300. Needs about 96 GiB of GPU memory.
  • bf16 only; one request at a time; a new request shape adds 1–3 s once, for graph capture.

Full details, the step-by-step breakdown and failed attempts are in the README and report (PDF).

References
  1. Open-Jev contributors. Open-Jev. github.com/Zefan-Cai/Open-Jev, 2026.
  2. Open-Jev-27B-v1.1. huggingface.co/ZefanCai/Open-Jev-27B-v1.1.
  3. Qwen Team. Qwen3.8-27B. huggingface.co/Qwen/Qwen3.8-27B.
  4. S. Yang and Y. Zhang. FLA: A Triton-Based Library for Hardware-Efficient Implementations of Linear Attention Mechanism. github.com/fla-org/flash-linear-attention, 2024.
  5. S. Yang, J. Kautz and A. Hatamizadeh. Gated Delta Networks: Improving Mamba2 with Delta Rule. ICLR 2025. arXiv:2412.06464. S. Yang et al. Parallelizing Linear Transformers with the Delta Rule over Sequence Length. NeurIPS 2024. arXiv:2406.06484.
  6. T. Dao. causal-conv1d. github.com/Dao-AILab/causal-conv1d.
  7. T. Wolf et al. Transformers: State-of-the-Art Natural Language Processing. EMNLP 2020 Demos. aclanthology.org/2020.emnlp-demos.6.
  8. S. Mangrulkar et al. PEFT. github.com/huggingface/peft, 2022.
  9. J. Ansel et al. PyTorch 2: Faster Machine Learning Through Dynamic Python Bytecode Transformation and Graph Compilation. ASPLOS 2024.
  10. E. J. Hu et al. LoRA: Low-Rank Adaptation of Large Language Models. ICLR 2022. arXiv:2106.09685.
  11. F. Standhartinger and contributors. JevBench. github.com/fstandhartinger/jevbench, 2026.
  12. J. Juravsky et al. Hydragen: High-Throughput LLM Inference with Shared Prefixes. 2024. arXiv:2402.05099.
  13. L. Zheng et al. SGLang: Efficient Execution of Structured Language Model Programs. NeurIPS 2024. arXiv:2312.07104.
  14. X. Miao et al. SpecInfer: Accelerating Large Language Model Serving with Tree-based Speculative Inference and Verification. ASPLOS 2024. arXiv:2305.09781.
  15. T. Cai et al. Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads. ICML 2024. arXiv:2401.10774.
  16. NVIDIA. cuBLAS/cuBLASLt (docs); CUTLASS (github.com/NVIDIA/cutlass).