Optimus: Efficient Diffusion-LLM Serving with Chunked Decode and Elastic Scheduling

Optimus is a serving system for block-diffusion language models (DLLMs) that improves inference throughput by dynamically adapting decoding granularity to runtime load. Standard block-diffusion decoding always processes a full token block per forward pass, which wastes compute when most tokens in the block are not yet decodable. Optimus replaces this with chunked decode — processing a smaller subset of tokens per pass — and uses an elastic scheduler to pick the optimal chunk size based on the current batch size, without any retraining or quality loss.

Supported model families:

• SDAR — JetLM/SDAR-8B-Chat-b32, JetLM/SDAR-4B-Chat-b32, and other SDAR variants
• LLaDA 2.0 — inclusionAI/LLaDA2.0-mini, inclusionAI/LLaDA2.0-flash, and other LLaDA 2.0 variants

---

How It Works

Chunked Decode

In standard block-diffusion decoding, each forward pass operates over the entire diffusion block regardless of how many tokens are actually ready to be decoded. Optimus decodes only a small chunk of tokens per pass. This avoids wasted computation on tokens that are not yet decodable and allows the system to trade off the number of decoding steps against the cost of each step.

Chunk size is a first-class runtime parameter. Smaller chunks reduce per-step compute; larger chunks improve GPU utilization by processing more tokens at once.

Elastic Scheduling

The elastic scheduler selects the chunk size at each iteration based on the current online serving conditions.

---

Project Structure

optimus/
├── engine/      # Modified LMDeploy engine with chunked-decode support
├── scheduler/   # Elastic scheduler, estimator implementations, and offline profiling tools
├── scripts/     # Evaluation and visualization entry points
└── requirements.txt

---

Installation

Requirements: Python · CUDA 12 · Datacenter-level GPU (Tested on A100)

# 1. Clone the repository
git clone <repo-url>
cd optimus

# 2. Create and activate a conda environment
conda create -n optimus python=3.12
conda activate optimus

# 3. Install Python dependencies
pip install -r requirements.txt

# 4. Install the local LMDeploy engine (contains Optimus patches)
cd engine && DISABLE_TURBOMIND=1 pip install -e . && cd ..

requirements.txt installs the upstream lmdeploy package to pull in its declared dependencies. The pip install -e . in step 4 installs the local patched engine as lmdeploy, which takes precedence at import time.

---

Usage

All scripts are run from the repo root. Run any script with --help to see the full list of options.

---

1. Visualize Chunked Decoding

Runs a single inference with step tracing enabled and saves a heatmap showing which tokens are decoded at each diffusion step. Useful for understanding how chunked decode progresses through a sequence.

bash scripts/run_visualize_decode.sh JetLM/SDAR-8B-Chat-b32
bash scripts/run_visualize_decode.sh JetLM/SDAR-8B-Chat-b32 --chunk-size 8 --dataset humaneval
bash scripts/run_visualize_decode.sh inclusionAI/LLaDA2.0-mini --extended-window

---

2. Offline Throughput Evaluation

Measures output tokens-per-second under a fixed concurrency setting. Supports multiple datasets (gsm8k, humaneval, ifeval, mbpp, sharegpt) and chunk-size sweeps.

bash scripts/run_offline_eval.sh JetLM/SDAR-8B-Chat-b32 sharegpt
bash scripts/run_offline_eval.sh JetLM/SDAR-8B-Chat-b32 sharegpt --chunk-size 8 --concurrency 128

---

3. Build Scheduler Estimators

The elastic scheduler requires two profiled artifacts built once per model. Pre-built artifacts for SDAR and LLaDA 2.0 on A100 are already included in scheduler/estimator_configs/.

Forward-time estimator — profiles forward-pass time across (batch size, chunk size) combinations:

bash scheduler/build_fwd_estimator.sh JetLM/SDAR-8B-Chat-b32

Unmask estimator — profiles how many tokens are decoded per step at each chunk size using ShareGPT prompts:

bash scheduler/build_unmask_estimator.sh JetLM/SDAR-8B-Chat-b32

---

4. Online Serving Evaluation

Evaluates throughput and latency under open-loop (Poisson arrival rate) or closed-loop (fixed concurrency) traffic. Pass --use-scheduler to enable the elastic chunk-size scheduler.

bash scripts/run_online_eval.sh JetLM/SDAR-8B-Chat-b32 --request-rate 4
bash scripts/run_online_eval.sh JetLM/SDAR-8B-Chat-b32 --request-rate 4 --use-scheduler
bash scripts/run_online_eval.sh JetLM/SDAR-8B-Chat-b32 --concurrency 64

---

End-to-End Walkthrough

The following commands run the complete Optimus pipeline on SDAR 8B with ShareGPT:

MODEL=JetLM/SDAR-8B-Chat-b32

# Step 1 — Visualize chunked decoding behavior on one sample
bash scripts/run_visualize_decode.sh $MODEL --chunk-size 16

# Step 2 — Measure offline throughput across chunk sizes
for CS in 32 16 8 4 1; do
    bash scripts/run_offline_eval.sh $MODEL sharegpt --chunk-size $CS
done

# Step 3 — Build scheduler estimators (one-time, per model)
bash scheduler/build_fwd_estimator.sh $MODEL
bash scheduler/build_unmask_estimator.sh $MODEL

# Step 4 — Online serving with elastic scheduling (scheduler enabled by default)
bash scripts/run_online_eval.sh $MODEL --request-rate 4

---

Coming Soon

• Accuracy evaluation — scripts and instructions for running quality benchmarks (GSM8K, HumanEval, IFEval, MBPP) to verify that chunked decode preserves generation accuracy relative to standard full-block decoding.
• Full experiment reproduction — end-to-end scripts to reproduce all throughput and latency results reported in the manuscript.

---

Acknowledgements

This repository is built on top of the following open-source projects:

• LMDeploy
• FOCUS
• SDAR
• LLaDA 2.0

We thank the authors and contributors of these projects for their valuable work.