[Cosmos3] Mixed W8A8/W8A16 denoising for ModelOpt FP8 checkpoints

[ConstBob]

What does this PR do?

Adds denoising-step-aware mixed W8A8/W8A16 for Cosmos3 ModelOpt FP8 checkpoints, so official Hub revision=fp8 is not W8A8 on every step.

This is the Diffusers counterpart of the FP8 schedule contract in vllm-omni#6560 (not a port of NVFP4, tensor parallel, caches, or vllm serve).

Behavior

• Checkpoint weights stay serialized static W8A8 (weight_only: false). Mixing is runtime.
• Video Nano / Super / Super-I2V fp8: first 3 + last 3 scheduler steps W8A16 (dequant + F.linear); middle steps native ModelOpt W8A8.
• Precision is chosen once per scheduler step (CFG cond/uncond match). Reasoner defaults to A16 when the checkpoint says reasoner: a16.
• Distilled 4-step and Super-T2I fp8 have runtime: null → mixed off, all steps native W8A8.
• mixed_precision_format=None (default) reads quantization_config.runtime.diffusion_step_policy from the transformer (including on-disk transformer/config.json when the live ModelOpt config omits runtime). "none" disables only the schedule. "fp8" forces it.

Code

• src/diffusers/pipelines/cosmos/mixed_precision.py — policy parse (fail-closed), linear discovery, W8A16 dispatch.
• pipeline_cosmos3_omni.py and modular denoise.py — apply/reset once per step.
• docs/source/en/api/pipelines/cosmos3.md — user-facing FP8 mixed section.
• Out of scope: NVFP4, custom kernels, dense-weight cache, cookbook/Hub cards (follow-up after this lands).

Tests

Unit (CPU, in-tree): tests/pipelines/cosmos/test_cosmos3_mixed_precision.py — 17 passed (schedule boundaries, 1-step stays native, reasoner vs generation, overlay of live runtime: null + on-disk policy, FP32 W8A16, malformed policy / NVFP4 reject).

pytest tests/pipelines/cosmos/test_cosmos3_mixed_precision.py
# 17 passed, 5 subtests passed

GPU, not in CI (official Hub revision=fp8):

• Cosmos3-Nano @fp8: auto-enable 3+3; 8-step T2I E2E with exact trace W8A16×3 / W8A8×2 / W8A16×3; 35-step latent W8A16×3 / W8A8×29 / W8A16×3.
• Super-Text2Image-4Step @fp8 transformer: auto-off; GEMM stays native W8A8.

Not claimed in this PR: Super 64B generate, 189-frame I2V, vs-BF16 hashes, or the vLLM-Omni H100 throughput table.

Depends on NVIDIA ModelOpt restore for serialized FP8 (same as existing Cosmos3 FP8 loading).

Self-review (final)

Ran the self-review skill on huggingface/main...HEAD against .ai/references/review-rules.md (plus pipelines.md, modular.md, testing.md, code_style.md, pitfalls.md).

Diff files: mixed_precision.py, pipeline_cosmos3_omni.py, modular_pipelines/cosmos/denoise.py, test_cosmos3_mixed_precision.py, docs/source/en/api/pipelines/cosmos3.md.

Blocking issues: none.

Non-blocking — left for review (deliberate):

1. Transformer forward wrappers stay installed after __call__ (mixed_precision.py). Inactive wrappers call the original ModelOpt forward. Per pipelines.md gotcha 7 we could uninstall on exit (PAG-style); we kept wrappers to avoid re-wrapping hundreds of linears each generate.
2. W8A16 reads layer.weight then F.linear. Per testing.md, that bypasses group-offload leaf hooks. Same dense A16 reference as vLLM-Omni; group offload + serialized FP8 was not a target.
3. Focused tests use fake ModelOpt linears + unittest.TestCase. Per testing.md, new pipeline tests prefer real tiny classes and pytest mixins. Existing Cosmos3 pipeline tests are still unittest; ModelOpt restore is impractical in CI.
4. mixed_precision_overlap is parsed from the checkpoint / resolve(), not exposed on Cosmos3OmniPipeline.__call__.
5. Docs revision="fp8" example does not repeat full ModelOpt from_pretrained kwargs; loading still follows the ModelOpt guide.

Dead code: FIRST_LAST_N_FP8_POLICY, from_kwargs, and quantization_config_from_module are used (tests + resolve). No unused runtime methods.

Verdict: READY.

Before submitting

• Did you use an AI agent (Claude Code, Codex, Cursor, etc.) to help with this PR? If so:
  • Did you read the Coding with AI agents guide?
  • Did you run the self-review skill on the diff?
  • Did you share the final self-review notes in the PR description or a comment?
• Did you read the contributor guideline?
• Did you read our philosophy doc? (important for complex PRs)
• Was this discussed/approved via a GitHub issue or the forum? Please add a link to it if that's the case.
• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines, and
  here are tips on formatting docstrings.
• Did you write any new necessary tests?
• Are you the author (or part of the team) of the model/pipeline (only applicable for model/pipeline related PRs)?

Who can review?

@yiyixuxu @asomoza

[sayakpaul]

Thanks for the PR. I think it's a bit intrusive in terms of what it is doing. Could you provide runtime numbers (latency, throughput) and also how does it compare in terms of visual results (with and without comparisons would be great).

[sayakpaul]

I don't think we need to mention things like these. If we want the users to control them, we could simplify the text.

[sayakpaul]

Users may not be familiar with the convention of W8A8. It would make sense to elaborate on that.

[sayakpaul]

We don't need these tests yet. Let's first gauge if this checkpoint is used much.

[sayakpaul]

How does it work out when people would use a different quantization backend like TorchAO, for example?