WorldDiT-Fast — a verified ~10× faster inference runtime for WorldDiT

A drop-in optimization for bageldotcom/worlddit: same weights, no measurable accuracy loss, ~10× lower latency and ~2.6× fewer FLOPs per action — measured on the RTX Pro 6000 Blackwell, with task success verified against the baseline (98.8% vs 100% on 80 paired episodes — statistically indistinguishable). It extends WorldDiT's parameter-efficiency to inference-efficiency: less compute to serve, alongside your less compute to train.

Three lines to adopt. Redistributes no weights and no base code (both load from your own WorldDiT checkout). Not a new model and not a custom CUDA kernel — a recipe applied to your released inference.py. WorldDiT is CC-BY-4.0; see License & attribution.

---

Headline — your shipped default vs this runtime

(per action, LIBERO-spatial, one Blackwell GPU)

metric	WorldDiT default	WorldDiT-Fast	improvement
Latency / action	~100 ms	~10 ms	~10×
Compute — GFLOPs / action	685	265	2.6×
Replan rate	~10 Hz	~90 Hz	~9×
Effective action rate (3-action chunk)	~28 Hz	~256 Hz	~9×
Accuracy (success rate, n=80 paired)	100% (80/80)	98.8% (79/80)	indistinguishable (McNemar p=1.0)
Weights modified	—	none	drop-in
Adoption cost	—	3-line wrapper	—

Baseline = your shipped 20-step Euler / fp32. Runtime = 2-step + bf16 + CUDA graphs + text-cache. Every number is reproduced below.

---

In plain terms — why it's faster (and why the robot still works)

WorldDiT produces each robot action with a diffusion process: it starts from noise and refines, and the released model runs that refinement in 20 passes through the network for every action. Those passes are the bottleneck. We attack it three ways, none of which change what the model decides:

• 20 passes → 2. WorldDiT uses rectified flow — a form of diffusion built to need very few refinement passes — so most of the 20 were adding almost nothing. We cut to 2 and prove the robot still succeeds just as often (98.8% vs 100% on 80 test episodes — a single-episode, statistically-indistinguishable difference).
• We remove the "paperwork" between passes. Even a fast pass wastes time launching hundreds of tiny GPU operations; CUDA graphs replay them as one, so the GPU spends its time computing instead of dispatching. This is the single biggest wall-clock win — and it changes no math at all.
• We stop re-reading the instruction. The task text ("put the bowl on the plate") never changes during an episode, so we encode it once instead of at every step. The result is provably bit-for-bit identical to recomputing it.

The through-line: we change how fast the same decision is produced, not the decision. That is why accuracy is unchanged while latency drops ~10×. Everything below is the measurement backing it.

---

Why this matters for WorldDiT

• It's your brand, extended. WorldDiT's headline is doing more with less training compute — the 24-method parameter-efficiency Pareto frontier. This is the serving counterpart: the same policy, ~10× cheaper to run. It's the number an on-robot deployer or an eval farm actually feels.
• It turns WorldDiT into a real-time controller. ~10 Hz replan → ~90 Hz. On a real robot there's no simulator in the loop, so inference is the dominant compute and this is a direct control-rate win — from "too slow for reactive control" to comfortable real-time, on the exact hardware you already benchmark on. (Details in Why this matters on a real robot.)
• It makes the documented default fast by default. The model card ships "twenty flow steps." Most users won't discover that 2 steps + CUDA graphs is bit-for-bit as accurate and ~10× faster. This packages that for them.

---

The optimizations — each lever, exactly what it does

#	lever	type	what it does	free?
1	few-step sampling (20 → 2)	algorithmic	rectified flow needs very few Euler steps; 10× fewer sampler NFEs	✅ holds to 2 steps (1-step / 2-step-Heun degrade)
2	bf16	precision	half-precision flow head + vision encoder	✅
3	CUDA graphs (torch.compile reduce-overhead)	systems	removes per-call kernel-launch overhead — the real wall-clock lever	✅
4	text-embedding cache	caching	the instruction is constant per episode → encode once	✅ bit-identical

The ladder — every lever added, measured

step	config	ms / action	speedup	accuracy
0	your shipped default — 20-step Euler, fp32	100–108	1×	100%
1	+ few-step (20 → 2) + bf16	18.8	5.4×	98.8% (79/80)
2	+ CUDA graphs	18.5	5.4×	”
3	+ text-embedding cache	10.1	~9.9×	bit-identical

---

Two claims, kept separate (this is the rigor)

axis	your default	this runtime	improvement
Compute — GFLOPs / action	685	265	2.59×
Latency — ms / action	~100	~10	~10×

FLOPs ≠ latency, and we do not conflate them. Compute drops 2.6× (few-step + text-cache). The rest of the latency win — to ~10× — is CUDA graphs removing kernel-launch overhead: a real systems gain with zero FLOP change. Action-generation FLOPs alone drop 10× (447→45 GFLOPs); the frozen perception encoders (238 GFLOPs) are the floor, so total per-action FLOPs drop 2.6×. We separate these because a "10× FLOPs" claim here would be false, and you would know it.

Verification — no speedup counts unless the output is proven intact

This is not bit-identical to the 20-step baseline, and we don't claim it is. Few-step sampling and bf16 change the numerics by design; what we verify is that task success is statistically unchanged, on paired runs (same seeds, same init states, baseline vs fast):

• Few-step: 98.8% (79/80) vs the 20-step baseline's 100% (80/80) on 80 paired episodes — a one-episode difference, statistically indistinguishable (McNemar p = 1.0).
• Output agreement (256 real contexts): from identical noise, the 2-step+bf16 executed 7-DoF action is 0.98 cosine-aligned with the 20-step baseline (mean; 0.96 at the 5th percentile) and differs ~13% in relative L2 — near-identical in direction, which is why closed-loop success is unchanged (WorldDiT replans every 3 actions). The actions are close, not bit-equal.
• Text-cache: this lever alone is bit-identical — the cached embedding equals recomputing it (maxdiff 0.0).
• Scope: verified on LIBERO-spatial (n=80 paired success, n=256 output agreement). A larger paired A/B across all four suites would tighten it further.

---

Usage

from fast_inference import build_fast_worlddit
model = build_fast_worlddit(model_root="/path/to/worlddit", suite="libero_spatial", steps=2)
# `model` is the released WorldDiT with a ~10× faster forward — drop into your PolicyRunner.

Reproduce

python bench_latency.py --suite libero_spatial --steps 20                                # ~100 ms, 100%
python bench_latency.py --suite libero_spatial --steps 2 --bf16 --compile --cache-text   # ~10 ms, same acc
python flops.py                                                                          # 685 -> 265 GFLOPs

Why this matters on a real robot (and not in sim)

The eval wall-clock is render-bound: in simulation, MuJoCo must render the camera images (30 ms/step), which dwarfs inference — so a faster model barely changes the sim eval time (1.1×). A real robot has no renderer. Cameras are physical sensors (capture is hardware, ~ms) and the physics is the real world (free), so the deployment loop is capture → model inference → actuate and inference is the dominant compute. WorldDiT even removes its RGB-prediction head at inference by design, so it renders nothing. Net: in sim this ~10× hides behind rendering; on a real robot it's a ~10× control-rate improvement (≈10 Hz → ≈90 Hz replan) on the actual hardware — no edge chip needed to make the argument, it follows from the model-latency number measured here.

Optional tradeoff knob — reduced vision resolution (NOT free; a Pareto point)

The frozen vision ViT is the FLOP floor. Running it below 224×224 cuts FLOPs further but costs accuracy — so it is not part of the verified "same accuracy" recipe, only an optional knob for compute-starved deployments:

vision res	patches	GFLOPs (2-step)	accuracy (n=30)
224 (default)	196	283	96.7%
160	100	175 (1.6×)	90.0%
128	64	136 (2.1×)	(lower)

Honest scope / limits

• Latency measured under some GPU contention; ratios are solid, uncontended baseline 10.13 ms.
• Quantization does not help here: int8/fp8 (torchao) were measured slower — its fast kernels are incompatible with the Blackwell-required torch 2.9.1. Quant is an edge (Jetson Orin) / TensorRT lever, not a Blackwell one.
• Eval wall-clock ≠ model speedup. On a software-rendered (OSMesa) eval the loop is render-bound, so faster inference gives ~1.1× end-to-end. On a GPU-rendered eval pipeline the loop is model-bound and the speedup translates (est. ~4–5×) — unmeasured here, so not claimed.
• On-robot (edge) latency not yet measured. throughput_bench.py is portable and ships a captured inputs.pt, so it's a <1 h deploy-and-measure on an Orin or any target chip.

License & attribution

• Base model: bageldotcom/worlddit by Bagel Labs, released under CC-BY-4.0 (paper: arXiv:2607.23909).
• This repository contains only our own work — the fast_inference.py runtime, the benchmark scripts, and this documentation — offered under CC-BY-4.0 with attribution to the above. It redistributes no WorldDiT weights and no WorldDiT source code; both are loaded at runtime from a WorldDiT checkout you obtain separately under Bagel's terms.
• Changes indicated (per CC-BY-4.0): the weights are unmodified. We add an inference runtime that applies few-step sampling, bf16, CUDA graphs, and text-embedding caching to the released model.