Add Cosmos3 Edge and Distilled checkpoints support by bastefaniak · Pull Request #5001 · vllm-project/vllm-omni

bastefaniak · 2026-07-09T21:58:04Z

Summary

This MR adds support for the upcoming Cosmos3 Edge and distilled models and
improves parity with the original implementation in cosmos-framework.

Cosmos3 Edge

Adds Cosmos3EdgeVFMTransformer, a Nemotron-based transformer architecture
used by Cosmos3 Edge models.
Selects the Edge transformer using the backbone_type field from
transformer/config.json.

Distilled models

Distilled checkpoints differ from regular one only by scheduler used, detects them from FlowMatchEulerDiscreteScheduler in the scheduler configuration.
Forces guidance_scale=1.0 for distilled inference.

Pipeline compatibility

Registers Cosmos3OmniPipeline as an alias for
Cosmos3OmniDiffusersPipeline to support the new checkpoint configuration.

Parity with `cosmos-framework`

Generates TimestepEmbedder frequencies in FP32.
Loads JSON prompts and updates their metadata from request parameters.
Replaces the previous UniPCMultistepScheduler approximation with the
native FlowUniPCMultistepScheduler implementation.

Note: The regression reported in
#4627 was not caused by
the scheduler change. It was caused by the video-generation default
flow_shift being 1.0; the default is now 10.0.

vLLM Version:
0.25

vLLM-Omni Commit:
Top of tree of this MR

Test Result

python3 -m pytest -q   tests/diffusion/models/cosmos3   tests/diffusion/cache/test_cache_dit.py   tests/config/test_pipeline_registry.py   tests/model_extras/test_model_extras.py
158 passed, 22 warnings in 1.74s

pytest -s -v -m "core_model and cpu"
4185 passed, 22 skipped, 2189 deselected, 72 warnings in 266.64s (0:04:26)

BEFORE SUBMITTING: read CONTRIBUTING.md and run the precheck-pr skill with the code agent for a self-check against project conventions.
(anything written below this line will be removed by GitHub Actions)

Add native Cosmos3 Edge image and video generation support across the transformer, diffusion pipeline, model registry, documentation, and tests. - add Cosmos3EdgeVFMTransformer and select it from the checkpoint's Edge backbone type - implement the Edge-specific ReLU2 feed-forward and Q/K normalization behavior - support Edge T2I, T2V, and I2V resolutions, with video guidance scale 5.0 and flow shift 3.0 - register Cosmos3OmniPipeline as an alias of Cosmos3OmniDiffusersPipeline for the new checkpoint model-index format - register the Edge transformer with Cosmos3 Cache-DiT integration - round requested video lengths to the causal VAE temporal grid - use FlowUniPCMultistepScheduler for regular checkpoints and apply the request flow shift while constructing timesteps - select FlowMatchEulerDiscreteScheduler from distilled checkpoint configuration and use its fixed sigma schedule - validate distilled checkpoint modalities and enforce guidance scale 1.0 - parse structured JSON prompts, update duration, FPS, resolution, and aspect-ratio metadata, and serialize them for model conditioning - distinguish explicitly provided guidance from request fallback values - seed PyTorch and CUDA consistently for request reproducibility - create timestep embedding frequencies in FP32 - add transformer, pipeline, registry, scheduler, prompt-formatting, and integration coverage - document Cosmos3 Edge registration, supported modes, defaults, and known limitations Signed-off-by: Bartosz Stefaniak <bstefaniak@nvidia.com>

Signed-off-by: Bartosz Stefaniak <bstefaniak@nvidia.com>

chatgpt-codex-connector · 2026-07-09T21:58:10Z

Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits.
Credits must be used to enable repository wide code reviews.

Signed-off-by: Bartosz Stefaniak <bstefaniak@nvidia.com>

hsliuustc0106 · 2026-07-13T12:39:42Z

+
+#### Notes
+
+- **Performance:** TBD. Checkpoint is not yet released


next PR for updates?

Would it be fine, if I removed docs changes and created followup PR when checkpoint is oficially released?

Removed Cosmos3-Edge from docs, will add in followup PR when checkpoint is released

hsliuustc0106 · 2026-07-13T12:40:48Z

        "Wan22VACEPipeline": enable_cache_for_wan22,
        "Wan22S2VPipeline": enable_cache_for_wan22_s2v,
        "Cosmos3OmniDiffusersPipeline": enable_cache_for_cosmos3,
+        "Cosmos3OmniPipeline": enable_cache_for_cosmos3,


this one is for edge? if so can we rename the pipelines?

The new name is required by diffusers. I wanted to have both names to remain backward compatible with previously downloaded checkpoints

hsliuustc0106 · 2026-07-13T12:46:41Z

 | `WanPipeline` | Wan2.1-T2V, Wan2.2-T2V, Wan2.2-TI2V | `Wan-AI/Wan2.1-T2V-1.3B-Diffusers`, `Wan-AI/Wan2.1-T2V-14B-Diffusers`, `Wan-AI/Wan2.2-T2V-A14B-Diffusers`, `Wan-AI/Wan2.2-TI2V-5B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `WanImageToVideoPipeline` | Wan2.2-I2V | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
-| `Cosmos3OmniDiffusersPipeline` | Cosmos3 T2I, T2V, I2V, V2V, T2V with sound, action policy | `nvidia/Cosmos3-Nano`, `nvidia/Cosmos3-Super` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
+| `Cosmos3OmniDiffusersPipeline` | Cosmos3 T2I, T2V, I2V, V2V, T2V with sound, action policy; Cosmos3 Edge T2I/T2V/I2V/action | `nvidia/Cosmos3-Nano`, `nvidia/Cosmos3-Super`, `nvidia/Cosmos3-Edge` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |


This marks Cosmos3 Edge as fully supported, but the PR does not include Edge run evidence yet. Please add runnable Edge online/offline commands, sample output, latency, and peak VRAM, or keep the support row scoped until that evidence is available.

Would it be fine, if I removed docs changes and created followup PR when checkpoint is oficially released?

Removed Cosmos3-Edge from docs, will add in followup PR when checkpoint is released

hsliuustc0106 · 2026-07-13T12:46:48Z

+
+#### Notes
+
+- **Performance:** TBD. Checkpoint is not yet released


Need actual Edge latency and peak VRAM here before merge, with hardware, resolution, step count, and the command used to measure. TBD makes the new model support unverifiable.

Would it be fine, if I removed docs changes and created followup PR when checkpoint is oficially released?

Removed Cosmos3-Edge from docs, will add in followup PR when checkpoint is released

Gaohan123 · 2026-07-13T14:17:11Z

 | `WanPipeline` | Wan2.1-T2V, Wan2.2-T2V, Wan2.2-TI2V | `Wan-AI/Wan2.1-T2V-1.3B-Diffusers`, `Wan-AI/Wan2.1-T2V-14B-Diffusers`, `Wan-AI/Wan2.2-T2V-A14B-Diffusers`, `Wan-AI/Wan2.2-TI2V-5B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `WanImageToVideoPipeline` | Wan2.2-I2V | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
-| `Cosmos3OmniDiffusersPipeline` | Cosmos3 T2I, T2V, I2V, V2V, T2V with sound, action policy | `nvidia/Cosmos3-Nano`, `nvidia/Cosmos3-Super` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
+| `Cosmos3OmniDiffusersPipeline` | Cosmos3 T2I, T2V, I2V, V2V, T2V with sound, action policy; Cosmos3 Edge T2I/T2V/I2V/action | `nvidia/Cosmos3-Nano`, `nvidia/Cosmos3-Super`, `nvidia/Cosmos3-Edge` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |


Gaohan123 · 2026-07-13T14:31:21Z


        half = frequency_embedding_size // 2
-        freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=target_dtype) / half)
+        freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half)


Why do we change it? It seems to affect the accuracy

I've changed it for better parity with cosmos-framework that contains source of truth for code. I verified on internal benchmarks, that quality remains very similar, but it's easier to debug when we are as close to original code as possible

Signed-off-by: Bartosz Stefaniak <bstefaniak@nvidia.com>

david6666666 · 2026-07-14T02:00:50Z

+        return hidden_states, k, v
+
+
+class Cosmos3EdgeLanguageModel(nn.Module):


Cosmos3LanguageModel declares _layerwise_offload_blocks_attrs = ["layers"]; this sibling does not, so the Edge UND tower is skipped by layerwise CPU offload. Please add the same attribute and cover it in the offload contract test.

david6666666 · 2026-07-14T02:00:58Z

+
+        self.is_distilled_model = False
+        if scheduler_class_name == COSMOS3_DISTILLED_CHECKPOINT_SCHEDULER_CLASS:
+            self.scheduler = FlowMatchEulerDiscreteScheduler.from_config(scheduler_config)


FlowMatchEulerDiscreteScheduler defaults to deterministic Euler, while Cosmos3 distilled sampling re-noises at every fixed step. Please force or validate stochastic_sampling=True here and cover an actual scheduler step.

We store that in checkpoint, but added forcing, tests and validation of other parameters

david6666666 · 2026-07-14T02:01:03Z

+        if seed is None:
            return
+        seed = int(seed)
+        torch.manual_seed(seed)


This mutates process-global RNG state for every request, and overlapping requests can interfere. Please thread the request-local generator into scheduler.step() instead.

Fixed, note that FlowMatchEulerDiscreteScheduler doesn't respect generator in diffusers==0.38.0, it's fixed in diffusers==0.39.0 (but we should be fine for now).

david6666666 · 2026-07-14T02:01:11Z

+    def _resolve_guidance_scale(self, sp: OmniDiffusionSamplingParams, default: float) -> float:
+        if self.is_distilled_model:
+            return 1.0
+        if sp.guidance_scale_provided:


The new 0.0 case is not reachable through a real request: OmniDiffusionRequest.__post_init__ treats zero as omitted and overwrites it with 1.0. Please fix the sentinel handling or drop that contract.

Is this comment about unittests? If so updated them

Signed-off-by: Bartosz Stefaniak <bstefaniak@nvidia.com>

Signed-off-by: Bartosz Stefaniak <bstefaniak@nvidia.com> Signed-off-by: y_null <y_null@qq.com>

Signed-off-by: Bartosz Stefaniak <bstefaniak@nvidia.com>

bastefaniak added 3 commits July 9, 2026 14:22

Merge branch 'main' into bstefaniak/add-cosmos3-edge

d561db4

Signed-off-by: Bartosz Stefaniak <bstefaniak@nvidia.com>

Fix indent bug

bedc873

Signed-off-by: Bartosz Stefaniak <bstefaniak@nvidia.com>

bastefaniak requested review from Gaohan123, Isotr0py, RuixiangMa, SamitHuang, ZJY0516, david6666666, hsliuustc0106, princepride, wtomin, xuechendi, yenuo26 and ywang96 as code owners July 9, 2026 21:58

bastefaniak added 3 commits July 9, 2026 23:59

Merge branch 'main' into bstefaniak/add-cosmos3-edge

3702152

Merge branch 'main' into bstefaniak/add-cosmos3-edge

76c08db

Merge branch 'main' into bstefaniak/add-cosmos3-edge

e296526

Signed-off-by: Bartosz Stefaniak <bstefaniak@nvidia.com>

hsliuustc0106 reviewed Jul 13, 2026

View reviewed changes

Gaohan123 added this to the v0.26.0 milestone Jul 13, 2026

Gaohan123 reviewed Jul 13, 2026

View reviewed changes

bastefaniak added 2 commits July 13, 2026 16:04

Remove Cosmos3-Edge from docs until it's released

7d52094

Signed-off-by: Bartosz Stefaniak <bstefaniak@nvidia.com>

Merge branch 'main' into bstefaniak/add-cosmos3-edge

371cc41

david6666666 reviewed Jul 14, 2026

View reviewed changes

bastefaniak added 3 commits July 14, 2026 10:01

Merge branch 'main' into bstefaniak/add-cosmos3-edge

0fcdfa1

Signed-off-by: Bartosz Stefaniak <bstefaniak@nvidia.com>

Answer comments

bce4727

Signed-off-by: Bartosz Stefaniak <bstefaniak@nvidia.com>

Merge branch 'main' into bstefaniak/add-cosmos3-edge

fd5831f

bastefaniak requested review from Gaohan123, david6666666 and hsliuustc0106 July 14, 2026 19:33

Merge branch 'main' into bstefaniak/add-cosmos3-edge

c8f01ad

hsliuustc0106 added the ready label to trigger buildkite CI label Jul 16, 2026

Merge branch 'main' into bstefaniak/add-cosmos3-edge

5c4baf2

ywang96 approved these changes Jul 16, 2026

View reviewed changes

ywang96 merged commit 85f7a5b into vllm-project:main Jul 16, 2026
6 checks passed

nv-dmajchrowski mentioned this pull request Jul 17, 2026

Add Cosmos3 Edge and Distilled checkpoints support sgl-project/sglang#31590

Merged

5 tasks

y-null pushed a commit to y-null/vllm-omni that referenced this pull request Jul 21, 2026

Add Cosmos3 Edge and Distilled checkpoints support (vllm-project#5001)

c20ec29

Signed-off-by: Bartosz Stefaniak <bstefaniak@nvidia.com> Signed-off-by: y_null <y_null@qq.com>

bowieshi pushed a commit to bowieshi/vllm-omni that referenced this pull request Jul 22, 2026

Add Cosmos3 Edge and Distilled checkpoints support (vllm-project#5001)

438415a

Signed-off-by: Bartosz Stefaniak <bstefaniak@nvidia.com>

lishunyang12 mentioned this pull request Jul 22, 2026

[Recipe] Add Cosmos3-Edge #5313

Merged

yuanheng-zhao mentioned this pull request Jul 22, 2026

[Refactor][1/N]diffusion/cache: cache-dit #5226

Merged

nagisa-kunhah pushed a commit to nagisa-kunhah/vllm-omni that referenced this pull request Jul 30, 2026

Add Cosmos3 Edge and Distilled checkpoints support (vllm-project#5001)

f9f8901

Signed-off-by: Bartosz Stefaniak <bstefaniak@nvidia.com>

khairulkabir1661 pushed a commit to khairulkabir1661/vllm-omni that referenced this pull request Sep 25, 2026

Add Cosmos3 Edge and Distilled checkpoints support (vllm-project#5001)

119e1ff

Signed-off-by: Bartosz Stefaniak <bstefaniak@nvidia.com>


		#### Notes

		- Performance: TBD. Checkpoint is not yet released

		return hidden_states, k, v


		class Cosmos3EdgeLanguageModel(nn.Module):

Conversation

bastefaniak commented Jul 9, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Cosmos3 Edge

Distilled models

Pipeline compatibility

Parity with cosmos-framework

Test Result

Uh oh!

chatgpt-codex-connector Bot commented Jul 9, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

bastefaniak Jul 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

bastefaniak commented Jul 9, 2026 •

edited

Loading

Parity with `cosmos-framework`

bastefaniak Jul 13, 2026 •

edited

Loading