# modify the script to use your local path
streamlit run applications/jacobi_model_chat.py

• Introduction
• Installation
• Model Weights
• Usage
  • Inference
  • Training
  • Evaluation
• Citation

Introduction

Why faster decoding?

AR decoding is high-quality but serial: one forward pass per token. Diffusion language models can decode many tokens in parallel, but typically require non-causal objectives and often break KV-cache-friendly serving.

fig1: Side-by-side comparison between Jacobi forcing decoding and text diffusion decoding, where Jacobi forcing decoding comes with more efficient KV cache reuse and is trained to generate higher quality drafts over a long horizon.

Jacobi Forcing bridges this gap by training an AR model to behave like a diffusion-style decoder while staying causal:

• Causal, left-to-right generation with KV-cache reuse
• Parallel token updates within a block of size $n$ (via Jacobi decoding) and training makes such convergence faster
• Multiblock decoding and Rejection recycling to exploit higher-quality draft with higher GPU utilization

fig2: Illustration of higher quality drafts that emerge from Jacobi Forcing model.

Installation

This section is demonstrative with path placeholders: adjust to match your repo structure.

1. Environment setup:

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

2. Clone this repository and build from source:

git clone https://github.com/hao-ai-lab/JacobiForcing.git
cd JacobiForcing

3. Install dependencies:

pip install -r requirements.txt

Model Weights

Base Models

Jacobi Forcing Models

Usage

Training

Jacobi Forcing training involves the following steps:

Prepare training data

Choice A: download existing data from Huggingface.

git lfs clone https://huggingface.co/datasets/JacobiForcing/OpenCodeInstruct_training_data_n32w16

Choice B

• step 1: Collect Jacobi trajectories from a base AR model (intermediate states + fixed-point state for all $n$-token blocks).

# generate trajctories using customized models
bash generate_trajectory/generation/generate_trajectory_opencodeinstruct_greedy.sh

• step 2: training sequence packing and mapping noise schedule to training sequence.

python3 generate_trajectory/data/2_prepare_efficient_cllm_training_data_progressive_noise_window.py \
    --input_path {trajectory_data_path} \
    --output_path {output_training_seq_path} \
    --n_token_seq_length {block_size} \
    --window_size {window_size} \
    --min_noisy_ratio 0 \
    --max_noisy_ratio 1.0 \
    --strategy "progressive"

Note: if the target model is not Qwen2.5, first modify generate_trajectory/generation/generate_trajectory_opencodeinstruct_greedy.sh to customize model path, trajectory data destimation, and input data path (you can download our length-bucketed input data from this link for code and this link for math).

Then adapt from the script generate_trajectory/generation/qwen2_modeling_jacobi_forcing_greedy.py to make your target model compatible.

fig3: Illustration of the training sequence packing process with an example (linear progressive) noise schedule mapping.

Noise-conditioned training over long horizons

cd /home/lah003/workspace/CLLM2/JacobiForcing
bash scripts/train/train_jacobi_forcing_coder_n32.sh

fig4: Jacobi Forcing uses the attention implementation shown above. It allows logits from clean blocks and noisy blocks to be generated with single forward pass to calculate the progressive consistency loss and AR loss.

Inference

Inference Engine

A lightweight, self-contained inference engine lives in inference_engine/. It supports autoregressive and Jacobi decoding (greedy & non-greedy) with FlashAttention, paged KV cache, CUDA graph capture, and tensor parallelism implmetend on top of nano-vLLM. On a single GPU the engine reaches 800–1000 tokens/second with Jacobi Forcing models.

# greedy Jacobi correctness test
python inference_engine/tests/test_jacobi_decoding_greedy.py --model-path /path/to/model

# non-greedy distribution similarity test
python inference_engine/tests/test_jacobi_decoding_nongreedy.py --model-path /path/to/model

Multiblock Decoding

Jacobi Forcing decoding typically exposes knobs like:

• block size n (tokens updated in parallel)

• rejection recycling verification budget pool_size

• block count K (maximum blocks “in flight”)

• activation ratio r

fig5: Illustration of multiblock Jacobi decoding with rejection recycling. High-quality n-grams from earlier iterations are reused as drafts.

Recommended starting point (from our grid search):

n=64, K=2, pool_size=4, r=0.85

To run comprehensive grid search profiling for TPS speedup and TPF across different settings, run:

cd JacobiForcing
bash scripts/inference/scanning_hyperparameter_jacobi_decoding_mr.sh

To run a specific decoding setting with multiblock decoding and rejection recycling, run:

# vanilla Jacobi decoding
python3 JacobiForcing/jacobi_forcing_inference_humaneval.py

# with multiblock decoding and rejection recycling
python3 JacobiForcing/jacobi_forcing_inference_MR_humaneval.py

Evaluation

Generation Quality Evaluation

We evaluate baseline models' and Jacobi Forcing models' performance on HumanEval, MBPP, GSM8K and MATH following the settings in evalchemy.

Performance Comparison

*Here we report the strongest checkpoints released by the authors; in principle EAGLE-3 and HASS are lossless in comparison with greedy AR checkpoints if they were trained with the Qwen2.5-7B backbone. Note that SD has a worse acceptance length (TPF) to TPS conversion ratio due to other overheads in the algorithm like token drafting using draft head, tree-like verification overhead, feature merging from different layers etc.

Performance Summary

Overall, Jacobi Forcing model consistently delivers up to $3-4\times$ wall-clock speedup on coding and math tasks with only minor accuracy changes versus greedy AR, while significantly outperforming both dLLMs and prior consistency-based parallel decoders in the accuracy–throughput tradeoff.

On a single B200 GPU with much higher FLOPs, the same Jacobi Forcing model with multiblock + rejection recycling can achieve an even more significant speedup at around 330 tokens/s (vs. around 80 tokens/s using AR), showing that the design continues to scale on newer accelerators.

Citation

This is the official project repository for the following paper. If you find this repository helpful, Please kindly cite:

@misc{hu2025fastaccuratecausalparallel,
      title={Fast and Accurate Causal Parallel Decoding using Jacobi Forcing}, 
      author={Lanxiang Hu and Siqi Kou and Yichao Fu and Samyam Rajbhandari and Tajana Rosing and Yuxiong He and Zhijie Deng and Hao Zhang},
      year={2025},
      eprint={2512.14681},
      archivePrefix={arXiv},
      primaryClass={cs.CL},
      url={https://arxiv.org/abs/2512.14681}, 
}