[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