[Cosmos3] Mixed W8A8/W8A16 denoising for ModelOpt FP8 checkpoints (#14664)

* Add opt-in mixed W8A8/W8A16 denoising for Cosmos3 ModelOpt FP8 checkpoints.

Keep native ModelOpt GEMM on middle steps and dequant-linear W8A16 on the first/last steps so CFG cond/uncond share one precision per scheduler step.

* Read Cosmos3 mixed-precision schedule from the checkpoint policy.

Enable first/last W8A16 only when transformer/config.json declares diffusion_step_policy, so distilled FP8 stays native W8A8 instead of inheriting a hardcoded 3+3 window.

* Harden Cosmos3 mixed-precision loading and document official Hub fp8 schedules.

Read the checkpoint runtime policy from on-disk transformer/config.json when the live ModelOpt config omits it, fail closed on incomplete policies, and allow FP32 activations on the W8A16 path.

* Address review: simplify FP8 mixed docs, no-op non-ModelOpt backends, drop unit tests.

W8A8/W8A16 is explained without listing call-site overrides; TorchAO and other
quantizers keep their native forwards; focused tests are removed until Hub fp8
usage is clearer.

* Document FP8 mixed default vs none trade-offs without ModelOpt restore boilerplate.

Keep serialized restore in the ModelOpt guide and spell out that all-W8A8 is faster but can flicker on multi-step video.

* Apply style fixes

* Regenerate Cosmos modular auto docstrings for mixed-precision inputs.

---------

Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>
Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>

Author: 3 people

docs/source/en/api/pipelines/cosmos3.md

@@ -43,6 +43,38 @@ Two checkpoints are released on the Hub — [`nvidia/Cosmos3-Nano`](https://hugg
 > [!TIP]
 > Make sure to check out the Schedulers [guide](../../using-diffusers/schedulers) to learn how to explore the tradeoff between scheduler speed and quality, and see the [reuse components across pipelines](../../using-diffusers/loading#reusing-models-in-multiple-pipelines) section to learn how to efficiently load the same components into multiple pipelines.
 
+## FP8 mixed W8A8/W8A16 denoising
+
+Official ModelOpt FP8 checkpoints live on the Hub `fp8` revision (for example [`nvidia/Cosmos3-Nano`](https://huggingface.co/nvidia/Cosmos3-Nano) with `revision="fp8"`).
+
+All of these checkpoints are quantized the same way. **W8A8** uses 8-bit weights and 8-bit activations (the restored ModelOpt GEMM). **W8A16** reuses those same 8-bit weights but skips activation quantization: the FP8 weight is dequantized and a standard linear runs on BF16/FP16/FP32 activations.
+
+Running W8A8 on every step can produce visible flickering in multi-step video generation. The video Nano / Super / Super-I2V FP8 checkpoints therefore declare a schedule in `transformer/config.json`: **W8A16 on the first 3 and last 3 steps**, **W8A8 in the middle**. Diffusers reads those counts from the checkpoint rather than hardcoding them. Precision is chosen once per scheduler step so classifier-free guidance cond/uncond calls match.
+
+Image generation and few-step distilled checkpoints do not show that flickering, so Super-T2I and the distilled 4-step FP8 repos declare no schedule and stay W8A8 on every step. The schedule is also **ModelOpt FP8 only**: other quantization backends (for example TorchAO) keep their native forwards.
+
+Load the `fp8` revision with the same restore path as the [ModelOpt guide](../../quantization/modelopt) (`revision="fp8"` already carries the quantization config). Mixed precision then follows the checkpoint automatically:
+
+```python
+import torch
+from diffusers import Cosmos3OmniPipeline
+
+pipe = Cosmos3OmniPipeline.from_pretrained(
+    "nvidia/Cosmos3-Nano",
+    revision="fp8",
+    dtype=torch.bfloat16,
+    device_map="cuda",
+)
+result = pipe(prompt="...", num_inference_steps=35)
+```
+
+Two generate-time choices:
+
+- **Default** (`mixed_precision_format=None`): if the checkpoint declares `diffusion_step_policy`, run W8A16 on the first/last N steps and native W8A8 in the middle. That is the intended recipe for multi-step **video** FP8 (less flickering than all-W8A8). Distilled 4-step and Super-T2I FP8 omit the policy, so the default is already all W8A8.
+- **`mixed_precision_format="none"`**: keep every step on native W8A8. Faster, because W8A16 is dequant + `torch.nn.functional.linear` rather than the restored FP8 GEMM, but multi-step video can flicker. Use this to A/B the schedule or to match a fully quantized baseline.
+
+On one Blackwell workstation, Cosmos3-Nano `@fp8` at 720×1280 / 35 steps was about **27% slower** (T2I) and **13% slower** (49-frame T2V) with the default mixed schedule than with `"none"`. Those numbers are not a throughput guarantee. Pass `"fp8"` only to force the first/last-N schedule on a ModelOpt FP8 checkpoint that has no policy.
+
 ## Prompt upsampling
 
 Cosmos 3 was trained on long, highly descriptive captions. For optimal quality, short text prompts should be **upsampled into a specific JSON structure** before they are passed to the pipeline. The upsampler lives in the [cosmos-framework](https://github.com/NVIDIA/cosmos-framework) package.
@@ -1117,6 +1149,9 @@ config (from the checkpoint's `modular_model_index.json`) and `guidance_scale` i
 1.0 since guidance is baked into the weights — passing any other value for either raises an error,
 and `negative_prompt` is warned about and ignored.
 
+FP8 distilled checkpoints (`revision="fp8"`) do not declare a mixed-precision policy, so every
+step stays native W8A8.
+
 Prompts follow the same descriptive JSON structure as the non-distilled models, so short text
 must be upsampled first — use `--mode text2image` (T2I) or `--mode image2video` (I2V) as
 described in [Prompt upsampling](#prompt-upsampling), then pass the JSON via `json.dumps(...)`.

src/diffusers/modular_pipelines/cosmos/denoise.py

@@ -3,6 +3,11 @@
 import torch
 
 from ...models.transformers.transformer_cosmos3 import Cosmos3OmniTransformer
+from ...pipelines.cosmos.mixed_precision import (
+    Cosmos3MixedPrecisionConfig,
+    apply_cosmos3_mixed_precision_step,
+    reset_cosmos3_mixed_precision,
+)
 from ...schedulers import FlowMatchEulerDiscreteScheduler, UniPCMultistepScheduler
 from ..modular_pipeline import (
     BlockState,
@@ -463,18 +468,61 @@ def loop_inputs(self) -> list[InputParam]:
             InputParam(
                 name="num_warmup_steps", type_hint=int, required=True, description="Number of scheduler warmup steps."
             ),
+            InputParam(
+                name="mixed_precision_format",
+                type_hint=str,
+                default=None,
+                description="None follows the ModelOpt FP8 checkpoint schedule; 'none' keeps the native quantized forward; 'fp8' is ModelOpt FP8 only.",
+            ),
+            InputParam(
+                name="mixed_precision_first_steps",
+                type_hint=int,
+                default=None,
+                description="Optional leading W8A16 step count.",
+            ),
+            InputParam(
+                name="mixed_precision_last_steps",
+                type_hint=int,
+                default=None,
+                description="Optional trailing W8A16 step count.",
+            ),
+            InputParam(
+                name="mixed_precision_reasoner_policy",
+                type_hint=str,
+                default=None,
+                description="Optional reasoner path: 'high_precision' (W8A16) or 'base_precision' (native W8A8).",
+            ),
         ]
 
     @torch.no_grad()
     def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState:
         block_state = self.get_block_state(state)
-        with self.progress_bar(total=block_state.num_inference_steps) as progress_bar:
-            for i, t in enumerate(block_state.timesteps):
-                components, block_state = self.loop_step(components, block_state, i=i, t=t)
-                if i == len(block_state.timesteps) - 1 or (
-                    (i + 1) > block_state.num_warmup_steps and (i + 1) % components.scheduler.order == 0
-                ):
-                    progress_bar.update()
+        mixed_precision = Cosmos3MixedPrecisionConfig.resolve(
+            components.transformer,
+            mixed_precision_format=getattr(block_state, "mixed_precision_format", None),
+            mixed_precision_first_steps=getattr(block_state, "mixed_precision_first_steps", None),
+            mixed_precision_last_steps=getattr(block_state, "mixed_precision_last_steps", None),
+            mixed_precision_reasoner_policy=getattr(block_state, "mixed_precision_reasoner_policy", None),
+        )
+        trace = []
+        try:
+            with self.progress_bar(total=block_state.num_inference_steps) as progress_bar:
+                for i, t in enumerate(block_state.timesteps):
+                    apply_cosmos3_mixed_precision_step(
+                        components.transformer,
+                        mixed_precision,
+                        i,
+                        len(block_state.timesteps),
+                        trace=trace,
+                    )
+                    components, block_state = self.loop_step(components, block_state, i=i, t=t)
+                    if i == len(block_state.timesteps) - 1 or (
+                        (i + 1) > block_state.num_warmup_steps and (i + 1) % components.scheduler.order == 0
+                    ):
+                        progress_bar.update()
+        finally:
+            reset_cosmos3_mixed_precision(components.transformer, mixed_precision)
+        components._mixed_precision_trace = trace
         self.set_block_state(state, block_state)
         return components, state
 
@@ -849,6 +897,15 @@ class Cosmos3TransferDenoiseStep(Cosmos3DenoiseLoopWrapper):
               The number of denoising steps.
           num_warmup_steps (`int`):
               Number of scheduler warmup steps.
+          mixed_precision_format (`str`, *optional*):
+              None follows the ModelOpt FP8 checkpoint schedule; 'none' keeps the native quantized forward; 'fp8' is
+              ModelOpt FP8 only.
+          mixed_precision_first_steps (`int`, *optional*):
+              Optional leading W8A16 step count.
+          mixed_precision_last_steps (`int`, *optional*):
+              Optional trailing W8A16 step count.
+          mixed_precision_reasoner_policy (`str`, *optional*):
+              Optional reasoner path: 'high_precision' (W8A16) or 'base_precision' (native W8A8).
           control_latents (`list`):
               Clean control latents for this chunk, one per hint in canonical order.
           latents (`Tensor`):

src/diffusers/modular_pipelines/cosmos/modular_blocks_cosmos3.py

@@ -396,6 +396,15 @@ class Cosmos3VisionCoreDenoiseStep(SequentialPipelineBlocks):
               Torch generator for deterministic generation.
           num_inference_steps (`int`):
               The number of denoising steps.
+          mixed_precision_format (`str`, *optional*):
+              None follows the ModelOpt FP8 checkpoint schedule; 'none' keeps the native quantized forward; 'fp8' is
+              ModelOpt FP8 only.
+          mixed_precision_first_steps (`int`, *optional*):
+              Optional leading W8A16 step count.
+          mixed_precision_last_steps (`int`, *optional*):
+              Optional trailing W8A16 step count.
+          mixed_precision_reasoner_policy (`str`, *optional*):
+              Optional reasoner path: 'high_precision' (W8A16) or 'base_precision' (native W8A8).
           **denoiser_input_fields (`None`, *optional*):
               conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
           guidance_scale (`float`, *optional*, defaults to 6.0):
@@ -469,6 +478,15 @@ class Cosmos3VisionSoundCoreDenoiseStep(SequentialPipelineBlocks):
               The number of denoising steps.
           sound_latents (`Tensor`, *optional*):
               Pre-generated noisy sound latents.
+          mixed_precision_format (`str`, *optional*):
+              None follows the ModelOpt FP8 checkpoint schedule; 'none' keeps the native quantized forward; 'fp8' is
+              ModelOpt FP8 only.
+          mixed_precision_first_steps (`int`, *optional*):
+              Optional leading W8A16 step count.
+          mixed_precision_last_steps (`int`, *optional*):
+              Optional trailing W8A16 step count.
+          mixed_precision_reasoner_policy (`str`, *optional*):
+              Optional reasoner path: 'high_precision' (W8A16) or 'base_precision' (native W8A8).
           **denoiser_input_fields (`None`, *optional*):
               conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
           guidance_scale (`float`, *optional*, defaults to 6.0):
@@ -557,6 +575,15 @@ class Cosmos3VisionActionCoreDenoiseStep(SequentialPipelineBlocks):
               Action-frame indexes fixed by action conditioning.
           action_latents (`Tensor`, *optional*):
               Pre-generated noisy action latents.
+          mixed_precision_format (`str`, *optional*):
+              None follows the ModelOpt FP8 checkpoint schedule; 'none' keeps the native quantized forward; 'fp8' is
+              ModelOpt FP8 only.
+          mixed_precision_first_steps (`int`, *optional*):
+              Optional leading W8A16 step count.
+          mixed_precision_last_steps (`int`, *optional*):
+              Optional trailing W8A16 step count.
+          mixed_precision_reasoner_policy (`str`, *optional*):
+              Optional reasoner path: 'high_precision' (W8A16) or 'base_precision' (native W8A8).
           **denoiser_input_fields (`None`, *optional*):
               conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
           guidance_scale (`float`, *optional*, defaults to 6.0):
@@ -647,6 +674,15 @@ class Cosmos3VisionSoundActionCoreDenoiseStep(SequentialPipelineBlocks):
               Action-frame indexes fixed by action conditioning.
           action_latents (`Tensor`, *optional*):
               Pre-generated noisy action latents.
+          mixed_precision_format (`str`, *optional*):
+              None follows the ModelOpt FP8 checkpoint schedule; 'none' keeps the native quantized forward; 'fp8' is
+              ModelOpt FP8 only.
+          mixed_precision_first_steps (`int`, *optional*):
+              Optional leading W8A16 step count.
+          mixed_precision_last_steps (`int`, *optional*):
+              Optional trailing W8A16 step count.
+          mixed_precision_reasoner_policy (`str`, *optional*):
+              Optional reasoner path: 'high_precision' (W8A16) or 'base_precision' (native W8A8).
           **denoiser_input_fields (`None`, *optional*):
               conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
           guidance_scale (`float`, *optional*, defaults to 6.0):
@@ -748,6 +784,15 @@ class Cosmos3TransferChunkDenoiseStep(SequentialPipelineBlocks):
               Frame rate of the generated video.
           num_inference_steps (`int`):
               The number of denoising steps.
+          mixed_precision_format (`str`, *optional*):
+              None follows the ModelOpt FP8 checkpoint schedule; 'none' keeps the native quantized forward; 'fp8' is
+              ModelOpt FP8 only.
+          mixed_precision_first_steps (`int`, *optional*):
+              Optional leading W8A16 step count.
+          mixed_precision_last_steps (`int`, *optional*):
+              Optional trailing W8A16 step count.
+          mixed_precision_reasoner_policy (`str`, *optional*):
+              Optional reasoner path: 'high_precision' (W8A16) or 'base_precision' (native W8A8).
           **denoiser_input_fields (`None`, *optional*):
               conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
           guidance_scale (`float`, *optional*, defaults to 6.0):
@@ -890,6 +935,15 @@ class Cosmos3TransferCoreDenoiseStep(SequentialPipelineBlocks):
               Frame rate of the generated video.
           num_inference_steps (`int`):
               The number of denoising steps.
+          mixed_precision_format (`str`, *optional*):
+              None follows the ModelOpt FP8 checkpoint schedule; 'none' keeps the native quantized forward; 'fp8' is
+              ModelOpt FP8 only.
+          mixed_precision_first_steps (`int`, *optional*):
+              Optional leading W8A16 step count.
+          mixed_precision_last_steps (`int`, *optional*):
+              Optional trailing W8A16 step count.
+          mixed_precision_reasoner_policy (`str`, *optional*):
+              Optional reasoner path: 'high_precision' (W8A16) or 'base_precision' (native W8A8).
           **denoiser_input_fields (`None`, *optional*):
               conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
           guidance_scale (`float`, *optional*, defaults to 6.0):
@@ -1012,6 +1066,15 @@ class Cosmos3AutoCoreDenoiseStep(ConditionalPipelineBlocks):
               Frame rate of the generated video.
           num_inference_steps (`int`):
               The number of denoising steps.
+          mixed_precision_format (`str`, *optional*):
+              None follows the ModelOpt FP8 checkpoint schedule; 'none' keeps the native quantized forward; 'fp8' is
+              ModelOpt FP8 only.
+          mixed_precision_first_steps (`int`, *optional*):
+              Optional leading W8A16 step count.
+          mixed_precision_last_steps (`int`, *optional*):
+              Optional trailing W8A16 step count.
+          mixed_precision_reasoner_policy (`str`, *optional*):
+              Optional reasoner path: 'high_precision' (W8A16) or 'base_precision' (native W8A8).
           **denoiser_input_fields (`None`, *optional*):
               conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
           guidance_scale (`float`, *optional*, defaults to 6.0):
@@ -1215,6 +1278,15 @@ class Cosmos3OmniBlocks(SequentialPipelineBlocks):
               Torch generator for deterministic generation.
           num_inference_steps (`int`):
               The number of denoising steps.
+          mixed_precision_format (`str`, *optional*):
+              None follows the ModelOpt FP8 checkpoint schedule; 'none' keeps the native quantized forward; 'fp8' is
+              ModelOpt FP8 only.
+          mixed_precision_first_steps (`int`, *optional*):
+              Optional leading W8A16 step count.
+          mixed_precision_last_steps (`int`, *optional*):
+              Optional trailing W8A16 step count.
+          mixed_precision_reasoner_policy (`str`, *optional*):
+              Optional reasoner path: 'high_precision' (W8A16) or 'base_precision' (native W8A8).
           **denoiser_input_fields (`None`, *optional*):
               conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
           guidance_scale (`float`, *optional*, defaults to 6.0):

src/diffusers/modular_pipelines/cosmos/modular_blocks_cosmos3_distilled.py

@@ -114,6 +114,15 @@ class Cosmos3DistilledVisionCoreDenoiseStep(SequentialPipelineBlocks):
           guidance_scale (`float`, *optional*):
               Unused for distilled checkpoints; classifier-free guidance is baked into the weights and the scale is
               forced to 1.0. Passing a value other than 1.0 raises an error.
+          mixed_precision_format (`str`, *optional*):
+              None follows the ModelOpt FP8 checkpoint schedule; 'none' keeps the native quantized forward; 'fp8' is
+              ModelOpt FP8 only.
+          mixed_precision_first_steps (`int`, *optional*):
+              Optional leading W8A16 step count.
+          mixed_precision_last_steps (`int`, *optional*):
+              Optional trailing W8A16 step count.
+          mixed_precision_reasoner_policy (`str`, *optional*):
+              Optional reasoner path: 'high_precision' (W8A16) or 'base_precision' (native W8A8).
           **denoiser_input_fields (`None`, *optional*):
               conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
 
@@ -206,6 +215,15 @@ class Cosmos3DistilledBlocks(SequentialPipelineBlocks):
           guidance_scale (`float`, *optional*):
               Unused for distilled checkpoints; classifier-free guidance is baked into the weights and the scale is
               forced to 1.0. Passing a value other than 1.0 raises an error.
+          mixed_precision_format (`str`, *optional*):
+              None follows the ModelOpt FP8 checkpoint schedule; 'none' keeps the native quantized forward; 'fp8' is
+              ModelOpt FP8 only.
+          mixed_precision_first_steps (`int`, *optional*):
+              Optional leading W8A16 step count.
+          mixed_precision_last_steps (`int`, *optional*):
+              Optional trailing W8A16 step count.
+          mixed_precision_reasoner_policy (`str`, *optional*):
+              Optional reasoner path: 'high_precision' (W8A16) or 'base_precision' (native W8A8).
           **denoiser_input_fields (`None`, *optional*):
               conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
           output_type (`str`, *optional*, defaults to pil):