Cosmos3 droid

[merceod]

What does this PR do?

• New registry family cosmos3_droid -> nvidia/Cosmos3-Nano-Policy-DROID (same Cosmos3Model; dims + capability flags load from the checkpoint, which differs from Nano only in sound_gen:false - the audio heads are already construction-gated on that flag, and the sound walk self-disables).
• configs/cosmos3_droid.yaml (Nano-shaped; no audio_decoder node; num_inference_steps_action: 4, guidance_scale_action: 3.0 i.e.the released DROID RoboLab policy-server sampling params), mstar serve cosmos3_droid, OpenAI image/video adapter registration, README + docs rows.
• Action-request serving defaults aligned with the reference action recipe for every cosmos3 checkpoint (previously they fell to the generic image/video defaults):
  • steps 30 / guidance 1.0 / flow_shift 5.0 / 832x480, config-overridable (num_inference_steps_action, guidance_scale_action, flow_shift_action).
  • action_chunk_size <-> num_frames default off each other (chunk 16 when neither is sent; num_frames ∈ {chunk, chunk+1} validated at submission. Previously an action request without num_frames failed mid-denoise).
  • The embodiment domain is now required and domain_name strings resolve through the published table (droid_lerobot=8); previously a missing domain_id silently became 0 (= the no_action domain).
  • raw_action_dim required + bounded by the model's padded action dim (forward_dynamics infers it from the conditioning action array).

How was it tested?

• CPU: cosmos3 serving/loader/droid tests 27 passed (new: an action-defaults battery in test_serving.py; test_droid.py
• GPU suite (both checkpoint dirs): 8F/50P — failure set identical to the documented coriander bf16-SDPA set; +6 new tests pass.
• Serving (coriander H200, tp1): policy (bare-defaults + pinned fixed-seed byte-determinism), inverse_dynamics (video obs -> action), forward_dynamics (image + action -> rollout video, visually inspected), t2v/i2v (eyeballed), /v1/videos with model=cosmos3_droid.
• Compiled bf16 VAE decode: a process's first execution of the compiled bf16 decode returns corrupted frames on some torch/cuDNN stacks (found on torch 2.11/cu130/cuDNN 9.19; deterministic, inductor-side, bf16-only, unaffected by autotune/layout/naming options, present even with a warm on-disk cache; fp32 is clean and aot_eager is bit-exact). The decoder now warms each new latent shape with one discarded decode before serving (test_vae_decode_bf16_warmup); first served outputs then match eager to <=5/255 on a fresh cache, and the fp32 path is byte-unchanged (nano t2i bit cells == mint on the fix commit).

Checklist

• ruff check . passes
• Added or updated tests / docs where relevant

[NSagan271]

LGTM, just a clarifying question on the decode issue.

[NSagan271]

Just checking---is this confirmed to not be, e.g., a race condition somewhere?

[merceod]

Yes, I checked this exactly. It reproduces in a bare bones single-threaded script (no serving stack, one thread, and default stream).Basically, compile vae.decode with the same flags, decode a fixed latent repeatedly and the result is that the first call is corrupted (mean abs err 0.23, max 2.0 vs eager) while every later call is clean and bit-identical. The corruption itself is bit-identical across fresh processes, unchanged under CUDA_LAUNCH_BLOCKING=1 (fully synchronous launches, so no async overlap at all), unchanged against a pre-populated inductor disk cache (no compilation happening in-process), and still present with TORCHINDUCTOR_COMPILE_THREADS=1. A race can't be timing-invariant and byte-reproducible under all of those. Since aot_eager is bit-exact vs. eager while inductor diverges, and fp32-compiled / bf16-eager are both clean, it's a deterministic first-execution miscompute in inductor's bf16 codegen (known-open upstream family, pytorch#164124).