Inference optimization for MolmoBot VLA: torch.compile + CUDA graphs + action-expert KV cache

[nguyenhoangthuan99]

Inference optimization for MolmoBot VLA: torch.compile + CUDA graphs, and an action-expert KV cache

Hi team! 👋

This PR bundles complementary, backward-compatible inference speedups for SynthManipMolmoInferenceWrapper / the flow-matching action expert:

1. torch.compile + CUDA graphs — CUDA-graph capture was being silently skipped due to three model-level issues; this fixes all three.
2. Action-expert cross-attention KV cache — the action expert re-projected the entire VLM context into cross-attention K/V on every flow-matching step; now it's projected once per query and reused.
3. Length-robust compiled serving (COMPILE_PAD_LEN) — fixed-length padding so the compiled graph is captured once and never recompiles per instruction length.
4. Optional fp8 (quantize='fp8') — per-tensor dynamic fp8 on the dense LLM blocks: ~1.16× and −3.6 GB weights on sm_120, at 0.4 % action delta.

All are opt-in-safe (every new param defaults to the previous behavior) and come with benchmark + correctness tests.

Changes overview

File	Δ	Change
olmo/models/molmobot/inference_wrapper.py	+187 / −5	New optional params compile_mode, compile_cudagraphs, compile_pad_len, compile_dynamic, quantize. Pre-move BufferCache to GPU, pre-build causal mask, warmup_cache at load. Embedded _Fp8DynamicLinear + fp8 conversion (Part 3).
olmo/models/video_olmo/video_olmo.py	+12 / −4	Rewrite x_flat[mask] += image_features → functional masked_scatter + add (bitwise identical, unit-verified)
olmo/torch_util.py	+6	Add BufferCache.to(device) — BufferCache is a plain dict, not registered buffers, so model.to() doesn't move it
olmo/nn/action_expert.py	+146 / −6	KV cache: precompute_kv() + optional kv_cache fast path in attention; inference-only prepare_cross_context() / denoise_step(). forward() untouched.
olmo/models/molmobot/molmobot.py	+22 / −6	KV cache: generate_actions() builds the cross-attention context once, then calls denoise_step() per flow step
scripts/benchmark_inference.py	new	3-way latency benchmark using random dummy inputs (no sim/hardware needed)
scripts/test_correctness.py	new	Unit test + replay stability + eager vs compiled action match

All new parameters default to the previous behavior. No existing caller is affected, and the training path is unchanged.

---

Part 1 — torch.compile + CUDA graphs + fixed-length padding

What was blocking CUDA graphs (and how I fixed it)

1. CPU-resident cache tensors — BufferCache is a plain dict, not registered nn.Buffers. model.to(device) never moved image_tokens; every forward paid a CPU→GPU copy and inductor skipped graph capture.
   → BufferCache.to(device) + explicit call at load.

2. Lazily-built cache tensors — casual_mask and RoPE sin/cos were created inside the first forward and stored in the cache. Under cudagraphs they'd be backed by graph-pool memory that later replays overwrite.
   → Pre-build causal mask at max_seq_len; call warmup_cache before compile.

3. In-place input mutation — x_flat[mask] += image_features makes x an input of the resumed subgraph; inductor refuses to capture graphs that mutate inputs.
   → masked_scatter into zeros + add (provably bitwise identical, unit-tested).

Length-robustness (COMPILE_PAD_LEN)

The model packs inputs to the actual seq_len, which varies per instruction (~1621–1634 = ~1600 fixed image tokens + a few instruction tokens). A whole-graph capture therefore recompiles ~365 s per distinct length — a net loss for varied-instruction eval sweeps. COMPILE_PAD_LEN=N (default 2048) pads every request to a fixed seq_len via the collator's to_max path, so the graph is captured exactly once and never recompiles — no mid-episode stalls. Padded positions are masked out, so actions are numerically identical to the unpadded call (worst |Δ| = 0.0033, bf16 noise). COMPILE_DYNAMIC=1 is a riskier fallback (marks the seq dim dynamic; can still recompile mid-run). All tables below use COMPILE_PAD_LEN=2048.

Measured results — full-query latency + GPU memory

Hardware: RTX PRO 6000 (Blackwell, sm_120) • Checkpoint: MolmoBot VLA (dense Qwen3-4B LLM + SigLIP ViT + flow-matching action expert, 10 flow steps) • Method: 40 warm queries (cold discarded), each = 8 images (480×640) + 14-D state + instruction, at COMPILE_PAD_LEN=2048. To isolate the compile/cudagraph/padding contribution the KV cache is held off here (Part 2 adds it). peak = max allocated after warmup; reserved = allocator reservation (incl. cudagraph pool).

Config	Mean	Jitter (σ)	Max	Peak GPU	Reserved	Speedup
Baseline (eager bf16)	303.5 ms	±3.4 ms	315.1 ms	11.9 GB	12.0 GB	1.00×
COMPILE_MODEL=1	222.2 ms	±14.0 ms	277.6 ms	12.3 GB	12.8 GB	1.37×
+COMPILE_CUDAGRAPHS=1	204.6 ms	±3.9 ms	221.4 ms	11.3 GB	11.7 GB	1.48×

• 1.48× end-to-end at the deployed default COMPILE_PAD_LEN=2048. The ratio is lower than an unpadded run because fixed-length padding adds constant GEMM work to every config, diluting the launch-overhead win.
• CUDA graphs give the lowest peak/reserved memory (11.3 / 11.7 GB — graph replay reuses captured memory) and the tightest worst case (max 221 ms vs 278 / 315).
• mode="default" matches "max-autotune" without the long autotuning.

---

Part 2 — Action-expert cross-attention KV cache

During flow-matching, the VLM backbone runs once and produces the per-layer LLM hidden states. The action expert is then integrated for num_steps Euler steps, and on every step it re-projected the full VLM context — the per-layer hidden states (≈1.6k tokens each), the robot state, and the attention mask — into cross-attention keys/values, across all layers. But those inputs are identical across steps; only the (few) action-token queries and the flow timestep change. So the large K/V projections were recomputed num_steps times over.

Fix: compute the cross-attention K/V (plus context projection, mask, and encoded state) once per query and reuse them across every step. Per step, the action expert only projects the queries from the action tokens.

• ActionExpertAttention: precompute_kv() + optional kv_cache fast path in forward() (defaults to None → original path, bitwise identical).
• ActionExpertBlock.forward(): threads an optional cross_kv_cache into attn2.
• ActionExpert: inference-only prepare_cross_context() (per-layer K/V + mask + encoded state, once) and denoise_step() (light per-step forward). forward() is left unchanged, so the training path is identical.
• MolmoBot.generate_actions(): builds the context cache once, then calls denoise_step() per flow step. The cache is a per-call local, so it stays a graph intermediate and remains compatible with torch.compile + CUDA graphs.

Measured results (action-model compute)

Hardware: RTX PRO 6000 (Blackwell)  •  same checkpoint, 10 flow steps.
Method: timing generate_actions only (image preprocessing done once, outside the loop), min latency over warm runs, at COMPILE_PAD_LEN=2048. KV-off is the faithful counterfactual on this branch — re-project the cross-context every flow step (the pre-cache behavior); KV-on is this branch. Everything else identical.

Mode	KV cache off	KV cache on	Speedup
eager	269.9 ms	252.8 ms	1.07×
torch.compile	197.2 ms	156.7 ms	1.26×
compile + cudagraphs	185.4 ms	144.4 ms	1.28×

• The cache stacks on top of compile + CUDA graphs and helps more there: once launch overhead is fused away, the redundant K/V GEMMs are a larger share of the remaining time.
• Combined — KV cache on + compile + cudagraphs (144.4 ms) vs KV cache off + eager (269.9 ms) = 1.87× on the action-model compute.
• The benefit grows with num_steps (more steps reusing the same cached context) and with context length.

---

Part 3 — Optional fp8 on the dense LLM blocks (quantize='fp8')

The served model is a dense Qwen3-4B (36 OLMoSequentialBlock; the moe_* config fields are inert). Its per-query cost, once compiled, is dominated by the LLM-backbone GEMMs. quantize='fp8' casts the transformer.blocks.* Linears (fused-QKV att_proj, attn_out, SwiGLU ff_proj, ff_out) to per-tensor dynamic fp8 (e4m3), before torch.compile so inductor fuses the per-call activation quant (amax → cast) into the fp8 _scaled_mm. ViT, the action-expert DiT, and embeddings stay bf16; RMSNorm / softmax / RoPE are never nn.Linear so they stay bf16 automatically.

Why a small hand-rolled linear rather than a library: on sm_120 (RTX PRO 6000 Blackwell, torch 2.7.1) rowwise fp8 scaling is unsupported (per-tensor only), and torchao's tensor-subclass path does not fuse the activation quant there — measured a net 1.04× (a loss after overhead). A plain-torch amax → cast → _scaled_mm that torch.compile fuses recovers 1.5–1.9× on the individual GEMMs.

Measured results

Hardware: RTX PRO 6000 (Blackwell, sm_120) • Checkpoint: dense Qwen3-4B MolmoBot (5.0 B params, 10 flow steps) • Method: generate_actions only, min latency over warm runs, COMPILE_PAD_LEN=2048. Memory via torch.cuda: weights = allocated after load; peak = max allocated after warmup; reserved = allocator reservation (incl. cudagraph pool).

Config	generate_actions	Weights	Peak GPU	Reserved	vs compiled bf16
eager bf16	230 ms	10.2 GB	11.4 GB	11.6 GB	—
compile + cudagraphs (bf16)	145 ms	10.2 GB	11.4 GB	11.7 GB	1.00×
+ quantize='fp8'	124 ms	6.5 GB	10.6 GB	8.4 GB	1.16×

• −3.6 GB weights (10.2 → 6.5 GB, −36%): the ~3.6 B LLM-linear params drop from bf16 to fp8; reserved memory falls ~3.3 GB.
• Fidelity: 0.39 % max action-chunk |Δ| vs bf16 (mean 0.11 %), same seed/inputs.
• Only helps with compile_model=1 (eager fp8 is slower — the fusion is the point). Stacks with COMPILE_PAD_LEN.
• Caveat: this shifts the numerics slightly; validate task success rate (sim A/B) before trusting fp8 in eval. Default is bf16 (quantize=None).

Combined — cumulative speedup vs original eager

All at COMPILE_PAD_LEN=2048, generate_actions min latency, same dense Qwen3-4B checkpoint. Each row adds one optimization on top of the row above:

Stack	generate_actions	vs eager
eager (no compile, no KV cache, bf16)	269.9 ms	1.00×
+ torch.compile + CUDA graphs	185.4 ms	1.46×
+ action-expert KV cache	144.4 ms	1.87×
+ fp8 (quantize='fp8')	124.3 ms	2.17×

~2.17× faster than eager on the action-model compute, plus −3.6 GB weights (10.2 → 6.5 GB from fp8). Padding (COMPILE_PAD_LEN=2048) is held constant across the ladder — it's a serving-robustness feature (one compile, zero per-instruction recompiles), not itself a speedup. The compile + CUDA graphs + KV cache portion (1.87×) is numerically ~bit-exact and on by default; fp8 (the last 1.16× step) is opt-in and pending a sim success-rate A/B.

Correctness

Test	What it checks	Result
masked_scatter unit	Functional masked_scatter + add matches original x_flat[mask] += feats bitwise (fp32 + bf16)	✅ Bitwise identical
KV-cache equivalence	prepare_cross_context + denoise_step matches ActionExpert.forward bitwise (fp32 + bf16, cross_attn & self_attn, ±attention mask, single-step and full loop)	✅ Bitwise identical
Replay stability	Repeated same-seed queries produce identical output under cudagraphs	✅ Bit-stable
A-B-A stale-memory	Interleave different-shaped inputs → subsequent A queries stay bit-stable	✅ No graph-pool corruption
Eager vs compiled+cudagraphs	Full pipeline (compile + cudagraphs + KV cache) vs eager	✅ < 0.8% relative (< 2% tolerance)

scripts/test_correctness.py exercises the full pipeline, so with the KV cache in this branch it now validates the combined path (eager vs compiled+cudagraphs, replay stability) end-to-end.

---

How to reproduce

# 3-way latency benchmark: baseline vs compile vs compile+cudagraphs (KV cache always active)
CKPT=/path/to/checkpoint_unsharded python scripts/benchmark_inference.py

# Full correctness suite (unit + replay stability + eager vs compiled)
CKPT=/path/to/checkpoint_unsharded python scripts/test_correctness.py

Both scripts use random dummy inputs matching the checkpoint's obs contract (image resolution, camera count, state dim, instruction text). No simulation or hardware setup required. The KV-off vs KV-on numbers above were produced by running the same measurement against the parent (pre-cache) commit.