By David Li and Andy Luo

Bringing world-model and robot-policy serving to AMD GPUs: end-to-end latency work on DreamZero in vLLM-Omni, and a hands-ready path to evaluate robot policies in simulation on ROCm.

For the last year most of the inference-serving conversation has been about LLMs: tokens per second, KV cache, prefix caching, speculative decoding. That world is mature. What is not yet mature — and what is moving fastest right now — is serving world models and robot policies: models that do not emit text, but emit future video frames and robot actions in a closed control loop.

This post is a deep dive into the work of getting one of those models, DreamZero, running efficiently inside vLLM-Omni, and pushing that stack onto AMD MI300X / ROCm. It covers the end-to-end latency optimization, the serving path for robot policies, and a working experiment: a robot policy driving a simulated arm to grasp a mug.

TL;DR

  • DreamZero E2E latency in vLLM-Omni went from 564.7 ms → 398.3 ms per closed-loop step — roughly a 1.42× speedup — through host-side scheduler restructuring, FP8 quantization, and VAE Conv3D fusion.
  • Robot-policy serving is hands-ready in vLLM-Omni. The serving path (OpenPI-compatible frontend, a DreamZero recipe, and bit-wise reproducible runtime) works end-to-end against a simulator.
  • The same building blocks that made LLM serving portable to ROCm — clean backend fallbacks, FP8 kernels, fused ops — carry over to world models, with a few new problem areas (3D convolutions, diffusion decode, action heads).
  • Below is a success experiment: a robot policy served through vLLM-Omni, evaluated in simulation, cleanly grasping and lifting a mug off a cluttered table.

Why “World Models” Are a Different Serving Problem

An LLM server is, at its core, a token loop: prefill a prompt, then autoregressively decode tokens, reusing a KV cache. The compute is dominated by attention and MLP GEMMs, and the whole ecosystem — batching, paging, speculative decoding — is built around that shape.

A world model for robotics breaks most of those assumptions:

  • The output is not text. DreamZero is a World Action Model (WAM): it jointly predicts future world states (as video frames) and the robot actions that would produce them. It is built on an image-to-video diffusion backbone, so a single “step” involves diffusion denoising, not just a matmul over a vocabulary.
  • The loop is a control loop, not a chat. Inputs are observations (camera frames), proprioceptive state, and a task instruction. Outputs are action chunks that get executed on a robot (or a simulator), which then produces the next observation. Latency per step is a hard real-time constraint, not a nicety.
  • The compute mix is heterogeneous. A step touches a diffusion transformer, a 3D VAE (Conv3D encode/decode over the temporal dimension), and an action head. Each of these has a different kernel profile, and each is a different porting target on a non-CUDA backend.

DreamZero itself is a ~14B-parameter autoregressive video-diffusion model that targets real-time closed-loop control (on the order of a few Hz). The reason a serving framework matters here is exactly that tension: a 14B diffusion model is heavy, and a control loop wants millisecond-scale steps. You do not get there without systems work.

That is where vLLM-Omni comes in. It generalizes vLLM’s serving machinery beyond text into an any-to-any, disaggregated pipeline — text, image, video, audio, and actions — and explicitly supports robot-policy models as a first-class serving target.

The E2E Latency Work: 564.7 ms → 398.3 ms

The headline result is the per-step end-to-end latency on DreamZero: 564.7 ms down to 398.3 ms. That did not come from one trick. It came from attacking three different parts of the step. (Full perf breakdown is tracked in vllm-omni#4127.)

1. Host-side sched_step restructuring

The first target was not the GPU at all — it was the host. In a control loop, the scheduler runs every step, and every step is short, so any Python/host overhead in the scheduling path is paid over and over. In LLM serving you can amortize this over long decode sequences; in a robot loop each step is a fresh scheduling decision, so host-side latency shows up directly in the wall-clock budget.

The fix restructures the host-side scheduling step so the per-step CPU work overlaps with GPU execution instead of serializing in front of it. Tracked in PR #5971.

The general lesson, which applies just as much on ROCm: for short-step generative loops, profile the host, not just the device. It is very easy to leave 100+ ms on the table in Python-side scheduling and launch overhead when each step is only a few hundred ms total.

2. FP8 quantization

The second lever is the obvious one for a 14B model on a latency budget: FP8 quantization (PR #6203). Moving the heavy diffusion-transformer GEMMs to FP8 cuts both the memory-bandwidth pressure and the compute time of the dominant matmuls.

This is also where the AMD story is genuinely good news rather than a caveat. MI300X has native FP8 (OCP e4m3/e5m2) support, and the ROCm software stack (hipBLASLt, and AITER for fused paths) has been pushing FP8 GEMM performance hard for exactly the LLM-serving reasons. World models inherit that work for free: the same FP8 GEMM kernels that accelerate an LLM’s MLP also accelerate a diffusion transformer’s blocks. The activation ranges differ, so quantization scheme and scaling matter, but the kernel substrate is already there.

3. VAE Conv3D → fused DSL kernel

The third piece is the most world-model-specific. DreamZero’s video path runs through a 3D VAE: to turn latent tokens back into video frames (and to encode observations), it uses Conv3D over the temporal dimension. Conv3D is not part of the classic LLM kernel set, and a naive Conv3D can dominate the decode side of a diffusion step.

The optimization rewrites the VAE Conv3D as a fused kernel (via an internal fusion DSL). This is done and measured, but it is not upstreamed as a PR yet because of an upstream version conflict — a good reminder that “it works and it’s fast on my branch” and “it’s merged” are two different milestones, especially in a fast-moving project.

For ROCm specifically, Conv3D is one of those operations where you want to be careful about which backend actually runs: MIOpen has 3D convolution support, but for a fixed shape inside a hot loop a hand-fused kernel that avoids the generic conv path is usually the win. This mirrors a pattern I have hit repeatedly porting 3D / world-model repos to AMD — the durable problem is almost never the transformer; it is the “extra” spatial/temporal ops (sparse conv, Conv3D, rasterization) hanging off the side of the model.

Where the 166 ms went

Optimization Layer Mechanism
Host-side sched_step Scheduler (CPU) Overlap per-step host work with GPU execution
FP8 quantization Diffusion transformer GEMMs Lower-precision matmuls, less bandwidth + compute
VAE Conv3D fusion 3D VAE decode Fused kernel replaces generic Conv3D path

The combined effect is the 564.7 ms → 398.3 ms result. None of these are exotic; the point is that the same three categories of optimization — host overhead, precision, and fusing the odd non-transformer op — are exactly what you attack when you move any generative loop onto a new accelerator.

Serving a Robot Policy End-to-End

Making the model fast is half the story. The other half is turning “a model that outputs actions” into “a service a robot (or an evaluation harness) can talk to.” Several pieces had to line up.

OpenPI-compatible frontend

Robot policies in the open ecosystem increasingly speak an OpenPI-style API — a realtime endpoint where a client streams observations and gets back action chunks. vLLM-Omni exposes a robot/OpenPI-compatible frontend so existing robot clients and evaluation harnesses can drive the served policy without a bespoke protocol. A frontend fix for that path landed in PR #6051.

Bit-wise reproducible runtime

This one is easy to underrate and turns out to be essential for robotics. When you are evaluating a policy, you need runs to be reproducible — otherwise you cannot tell whether a change in success rate came from your code change or from nondeterministic kernels. Diffusion sampling, reduction order, and atomics can all inject run-to-run variance.

PR #6166 makes the runtime bit-wise reproducible, which means two evaluation runs of the same policy on the same seeds produce identical action streams. On ROCm this is a real constraint to hold onto: several fast kernels use atomic accumulation or nondeterministic reduction orders, and you have to explicitly opt into deterministic paths to preserve reproducibility. It is worth the cost during evaluation.

A DreamZero recipe

Finally, a DreamZero serving recipe (PR #6148) captures the launch configuration — which model, which precision, which endpoint, which loop parameters — so serving the policy is a documented, repeatable vllm serve-style invocation rather than tribal knowledge. Recipes are how this kind of work stops being a one-off.

The Experiment: Grasping a Mug in Simulation

Here is the payoff — a robot policy served through vLLM-Omni and evaluated in a simulator. The task: a tabletop scene with two apples, two books, a square plate, and a white mug. The policy has to identify the mug, approach it, close the gripper, and lift it clear of the table without disturbing the other objects.

The arm goes straight for the mug, dwells to settle the grasp, closes on it, and lifts it off the table — clean, with the apples, books, and plate left untouched. It is a simple task by human standards, but end-to-end it exercises the whole stack: observation encode → world/action prediction (the served DreamZero-style policy) → action chunk → simulator step → next observation, all through the OpenPI-compatible serving path, on a reproducible runtime.

That is the meaningful milestone. It is not “we generated a nice video.” It is “a robot policy, served by vLLM-Omni through a standard robot API, closed the loop against a simulator and completed a manipulation task.” The video is the artifact; the point is the plumbing behind it.

Where This Fits: Stage 1 and Stage 2

The plan for robot-policy serving on AMD GPUs is laid out in RFC-6168. The short version:

Stage 1 — hands-ready policy serving. This is what the work above delivers. Serving robot-policy / world-action models in vLLM-Omni is working end-to-end: a fast DreamZero step, an OpenPI-compatible frontend, a reproducible runtime, and a documented recipe, validated against a simulator. From here the natural expansion is more models — bringing additional open-source robot policies (GR00T, Lingbot, and others discussed in the RFC) onto the same serving path.

Stage 2 — RL rollout and session management. The harder, more interesting problem is next: treating policy serving as the inference engine underneath RL rollout. That pulls in session / request / cache management — how you keep per-episode state, how you batch many concurrent rollouts, how you cache and reuse computation across steps of the same episode. This is where the LLM-serving intuition (KV cache, paging, continuous batching) has to be rethought for episodic, stateful, action-emitting loops. It is the focus going forward.

The ROCm Takeaway

Zooming out, the encouraging result is how transferable the LLM-serving toolkit turned out to be. The three things that bought the DreamZero speedup — killing host-side scheduling overhead, FP8-quantizing the heavy GEMMs, and fusing the one weird non-transformer op — are the same three moves you make porting any generative loop to AMD. MI300X’s FP8 support and the maturing ROCm/AITER kernel stack mean the transformer core is largely a solved substrate.

The genuinely new surface area for world models is the non-LLM parts: 3D VAE Conv3D, diffusion decode, action heads, and the reproducibility requirements that come with evaluating policies rather than just chatting with a model. Those are tractable, and they are exactly the pieces this work has been chipping away at.

World models and robot policies are where multimodal serving gets physical. Getting them to run — and run fast, and run reproducibly — on AMD GPUs is now hands-ready for Stage 1. Stage 2 is where it gets really interesting.

Links: Perf summary #4127 · Host-side sched_step #5971 · FP8 quant #6203 · Reproducible runtime #6166 · DreamZero recipe #6148 · OpenPI frontend fix #6051 · RFC-6168