Running Cosmos 3 Robot Policies on a Single RTX 4090

Code: cosmos-lite v0.3.0
Models: Nano GenW8A8 · Edge GenW8A8

NVIDIA’s Cosmos 3¹ shipped with robot policies trained on DROID². Running them yourself is harder than it sounds: the 16B Nano policy needs about 33 GB just to load, and the reference workflow assumes a datacenter GPU like an A100 or H100. A single RTX 4090 could not run it at all.

Cosmos Lite fills that gap. It is an efficient inference runtime for Cosmos 3 policies, built on top of Cosmos Framework. It covers two use cases: local serving as a batch-one policy server for real-robot deployment, and parallel rollout as a shared server feeding multiple simulator lanes, exporting rollout trajectories for further training.

The benchmark covers all 120 RoboLab³ tasks, 10 episodes each, all with the same sampler settings (guidance 3, two UniPC denoise steps, shift 5), and latency is measured per action request on a single RTX 4090.

Both Nano rows require quantization because BF16 does not fit a 24 GB GPU. GenW8A8 holds task success while running Nano 2.5× faster than W8A16 and Edge 1.76× faster than BF16. One caveat, stated plainly: every rate here trails the NVIDIA-reported rows (31.7% vs. 36.8% for Nano, 19.3% vs. 22.9% for Edge). “Holds task success” means statistically indistinguishable, not free. Verify success on your own tasks before deploying.

Make it fit via quantization

The first problem is fitting a 16B policy into 24 GB, and the model’s structure determines how. A Cosmos 3 policy is a Mixture-of-Transformers (MoT): every decoder layer has an understanding branch for language, state, and visual conditioning, and a generation branch that turns diffusion latents into actions. The two branches differ in token counts, matrix shapes, and precision sensitivity, so they are quantized separately. We quantize the Linear modules inside the MoT decoder, while vision, embeddings, action adapters, and the VAE stay BF16.

Mixed GenW8A8 shows the best tradeoff among evaluated configurations, better than full-model W8A8, which I also built and benchmarked. On Nano, full W8A8 saved only ~10 ms over GenW8A8 while using 3.5 GB more memory; on Edge it was outright slower than the mixed scheme. The generation branch holds the large GEMMs, which is where FP8 tensor cores pay off, while the understanding branch stays accurate at W4A16 with AWQ-style calibration⁶ over 128 training episodes.

One caveat from this stage: open-loop action error does not predict closed-loop success. The best artifact by replay error lost 12 points of success rate against a simpler one in matched rollouts, and calibration, while it improves the error distribution, does not remove rare action spikes. Every precision decision in the release is therefore gated on rollout success rate.

Do less work via sampler tuning

After quantization, the next lever is how much computation each request performs. The upstream RoboLab server defaults to guidance 3 with four UniPC⁷ denoise steps, and v0.1.0, the first deployable release, inherited that setting. Halving the step count to two, a YAML override and nothing more, halves the denoiser forwards per request.

Two steps cut Nano request latency by 41%, and in the paired rollout gates above the success rate held or improved on both models. This is less surprising than it sounds: UniPC is a numerical solver, so more steps mean finer integration, not a better policy. When the learned vector field and the inference schedule are imperfectly matched, extra solver points can move actions away from good trajectories. Edge showed the same pattern, improving from 21/50 to 34/50 as steps went from four to two.

The table also explains why guidance stays at 3: dropping it to 1 removes the CFG⁸ double-forward and looks like free latency, but it loses enough closed-loop success to fail the gate. These are paired gates on one task; the full-suite results carry the usual uncertainty, and other environments or robots may prefer different settings.

Do it cheaper via kernels and compilation

The last lever is the cost of the computation that remains. Early profiling invalidated two common assumptions here: attention took only ~4% of GPU kernel time, and weight-only quantization measured roughly neutral on Nano’s real shapes and slightly negative on Edge. The changes that survived testing, in the order they landed:

About the Triton row: CUTLASS under-occupies Edge’s dominant SM89 shapes, and the shape-tuned Triton kernel is ~8% faster end-to-end there. The same kernels were also faster on Nano, but Nano’s paired rollout moved from 49/50 to 47/50. Two episodes is likely noise (needs further verification), yet with a 49/50 control there was no quality margin to trade for ~9% latency, so Nano kept CUTLASS.

What didn’t work

Most of the exploration was coding-agent-driven, so failed attempts were cheap to generate, and the dead-end list is long enough to be worth recording. Grouped by where they failed:

Failed at the operator or end-to-end level:

• Whole-model torch.compile: graph breaks around the custom GEMMs, no gain, 20 s first requests.
• Static-shape compile and whole-request CUDA Graphs: recompiles on new prompt lengths; capture overhead with no steady gain.
• One global attention backend: SageAttention loses on the short condition attention; FlashInfer never beat FlashAttention2.
• Concatenated projection weights: some isolated operator wins, 19 GB peak memory on Nano, inconsistent end-to-end; replaced by shared activation quantization.
• Edge condition K/V cache: numerically correct, no repeatable gain; the cache ships for Nano only.
• torchao INT8 weight-only: 3–4× slower than BF16 on the model’s real shapes.
• Batching the CFG branches: fewer but larger attention calls, net slower end to end.
• Shorter action chunks (32→16/8): lower nominal compute, worse measured latency, different control semantics.

Passed the kernels, failed the rollout:

• Guidance 1: faster requests, more closed-loop failures (40/50 vs. 45/50 on Nano).
• Nano Triton FP8 GEMM: faster kernels and replay, rollout point estimate dipped (49→47/50).

Deferred:

• TensorRT-LLM: credible, but exporting the full MoT-plus-vision-plus-VAE stack is a different-scale project.

The full decision map, with the measurements behind each drop, is in the repo’s optimization report.

What’s next

v0.3.0 is the first fully measured release, with the RTX 4090 as its primary tested target. Planned next:

• More NVIDIA GPU architectures and inference backends.
• More world-action models beyond the current Nano and Edge DROID policies.
• Real-robot evaluation and optimization.

If you are working on policy deployment, simulation data generation, or inference for MoT-style models, I’d love to hear from you: open an issue or email me.