[Model] Add Cosmos3 mixed W8A8/W8A16 and W4A4/W4A16 denoising (#6560)

Signed-off-by: Rahul Steiger <rsteiger@aws-cmh-slurm-1-vscode-04.cm.cluster>
Signed-off-by: Wojciech Kutak <wkutak@nvidia.com>
Co-authored-by: Rahul Steiger <rsteiger@nvidia.com>

docs/user_guide/quantization/modelopt.md

@@ -204,6 +204,91 @@ omni = Omni(
 | `--linear-backend cutlass` | str | auto | Select the validated CUTLASS linear backend for supported ModelOpt NVFP4 or mixed FP8/NVFP4 diffusion stages |
 | `--moe-backend cutlass` | str | auto | Select the validated CUTLASS MoE backend for supported ModelOpt mixed MoE checkpoints |
 
+## Cosmos3 Mixed-Precision Schedule
+
+The experimental Cosmos3 schedule can use the checkpoint-native quantized GEMM for middle denoising
+steps and dense 16-bit activation execution for selected first and last steps.
+The schedule supports serialized ModelOpt FP8 and NVFP4 checkpoints:
+
+| Detected weight format | Native middle steps | A16 boundary steps |
+|------------------------|---------------------|--------------------|
+| FP8 | W8A8 | W8A16 |
+| NVFP4 | W4A4 | W4A16 |
+
+The native and A16 paths share one quantized weight representation. The A16
+path dequantizes the live backend weight and calls `F.linear`; it does not keep
+a second checkpoint-weight snapshot. Scheduled FP8 currently requires
+serialized tensorwise scales and a backend that retains canonical FP8 weights.
+Scheduled NVFP4 requires a supported CUTLASS-compatible native or FlashInfer
+layout. SmoothQuant, Marlin-repacked FP8, and per-channel/per-token FP8
+schedules fail closed.
+
+The reasoner uses dense A16 execution by default when its weights are FP8 or
+NVFP4. Set the nested `reasoner` field to `native` to retain the checkpoint-native
+W8A8/W4A4 path. A BF16 reasoner remains BF16 in either mode. The schedule
+currently requires tensor parallel size 1 and one active request per worker,
+does not support HSDP, and does not support block-scaled FP8. Its live-weight
+design is compatible with model-level and standard layer-wise offload, but
+those paths still require GPU end-to-end validation. Distributed layer-wise
+offload is rejected because its direct loader bypasses ModelOpt post-load
+transformations.
+
+A checkpoint can opt in by adding a versioned policy to its authoritative
+`transformer/config.json`. The same policy schema is used for FP8 and NVFP4:
+
+```json
+{
+  "quantization_config": {
+    "quant_method": "modelopt",
+    "quant_algo": "FP8",
+    "runtime": {
+      "diffusion_step_policy": {
+        "schema_version": 1,
+        "type": "first_last_n",
+        "index_space": "denoising_loop_iteration",
+        "scope": ["transformer"],
+        "default_mode": "native",
+        "first_steps": {"count": 3, "mode": "a16"},
+        "last_steps": {"count": 3, "mode": "a16"},
+        "overlap": "a16",
+        "reasoner": "a16"
+      }
+    }
+  }
+}
+```
+
+The example preserves the checkpoint's existing `quant_method` and
+`quant_algo`; only `runtime` is added.
+
+When present and valid, the checkpoint policy is the default. An explicit
+runtime object overrides it. To retain ordinary checkpoint-native execution
+for every step, disable only the schedule:
+
+```bash
+vllm serve /path/to/Cosmos3-Nano-modelopt \
+  --omni \
+  --additional-config '{"cosmos3_mixed_precision":{"enabled":false}}'
+```
+
+This does not disable checkpoint quantization; it prevents installation of the
+mixed-step runtime.
+
+```bash
+vllm serve /path/to/Cosmos3-Nano-modelopt \
+  --omni \
+  --additional-config \
+  '{"cosmos3_mixed_precision":{"first_steps":3,"last_steps":3,"reasoner":"a16"}}'
+```
+
+The presence of `cosmos3_mixed_precision` supplies an explicit runtime
+override; an empty object uses the defaults shown above.
+The schedule uses vLLM's selected native backend and fails during loading if
+its live weight layout cannot support A16 dequantization.
+The runtime selects FP8 or NVFP4 from the checkpoint's ModelOpt configuration.
+BF16 linears are untouched. Mixed FP8/NVFP4 checkpoints are not supported by
+this schedule.
+
 ## Validation and Notes
 
 1. Compare the ModelOpt checkpoint against the BF16 baseline with the same

recipes/cosmos3/Cosmos3-Nano.md

@@ -115,6 +115,38 @@ shift at the same seed). The pipeline
 auto-resolves from `model_index.json`; pass
 `--model-class-name Cosmos3OmniDiffusersPipeline` to force it explicitly.
 
+For a serialized ModelOpt FP8 or NVFP4 checkpoint, an experimental
+mixed-precision schedule can use native W8A8/W4A4 in middle denoising steps
+and dense W8A16/W4A16 in the first and last steps:
+
+```bash
+vllm serve /path/to/Cosmos3-Nano-modelopt \
+  --omni \
+  --additional-config \
+  '{"cosmos3_mixed_precision":{"first_steps":3,"last_steps":3,"reasoner":"a16"}}'
+```
+
+The nested object's presence supplies an explicit runtime override; use an
+empty object for these defaults. A compatible checkpoint can instead declare
+the versioned `runtime.diffusion_step_policy` in its authoritative
+`transformer/config.json`. The checkpoint policy is used only when no runtime
+override is present. Use
+`--additional-config '{"cosmos3_mixed_precision":{"enabled":false}}'` to
+disable the checkpoint schedule without disabling checkpoint quantization.
+This path currently requires tensor parallel size 1 and does not support HSDP
+or block-scaled FP8. It keeps one live quantized weight representation:
+scheduled FP8 requires serialized tensorwise scales and a canonical backend,
+while scheduled NVFP4 requires a supported CUTLASS-compatible native or
+FlashInfer layout. Its dequantization is a correctness baseline; no speedup is
+implied. Quantized reasoner weights use A16 by default; set
+`"reasoner":"native"` in the nested object to keep W8A8/W4A4.
+The checkpoint's ModelOpt configuration selects FP8 or NVFP4; mixed FP8/NVFP4
+checkpoints are not supported by this schedule. BF16 linears are left unchanged.
+The schedule currently supports one active request per worker. Model-level and
+layer-wise offload are designed to work with this live-weight path but still
+require GPU validation. Distributed layer-wise offload and incompatible native
+weight layouts fail during loading.
+
 #### Verification
 
 Best quality uses the JSON-upsampled prompts from `assets/` (download with

tests/diffusion/quantization/test_cosmos3_mixed_precision.py

@@ -0,0 +1,690 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.models.cosmos3.mixed_precision import (
+    Cosmos3MixedPrecisionConfig,
+    Cosmos3MixedPrecisionRuntime,
+    resolve_mixed_precision_config,
+)
+from vllm_omni.diffusion.models.cosmos3.mixed_precision import runtime as runtime_impl
+from vllm_omni.diffusion.models.cosmos3.mixed_precision.checkpoint import read_checkpoint_policy
+from vllm_omni.diffusion.models.cosmos3.mixed_precision.runtime import (
+    Cosmos3MixedPrecisionLinearMethod,
+)
+from vllm_omni.diffusion.models.cosmos3.mixed_precision.strategy import (
+    Fp8W8A8W8A16Strategy,
+    Nvfp4W4A4W4A16Strategy,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+def _policy(
+    *,
+    first_steps: int = 3,
+    last_steps: int = 3,
+    reasoner: str = "a16",
+    scope: list[str] | None = None,
+) -> dict[str, object]:
+    return {
+        "schema_version": 1,
+        "type": "first_last_n",
+        "index_space": "denoising_loop_iteration",
+        "scope": scope or ["transformer"],
+        "default_mode": "native",
+        "first_steps": {"count": first_steps, "mode": "a16"},
+        "last_steps": {"count": last_steps, "mode": "a16"},
+        "overlap": "a16",
+        "reasoner": reasoner,
+    }
+
+
+def _fake_quant_config(name: str, *, serialized: bool = True):
+    if name == "modelopt":
+        return SimpleNamespace(
+            get_name=lambda: name,
+            is_checkpoint_fp8_serialized=serialized,
+        )
+    if name == "modelopt_fp4":
+        return SimpleNamespace(
+            get_name=lambda: name,
+            quant_method="NVFP4",
+            is_checkpoint_nvfp4_serialized=serialized,
+        )
+    return SimpleNamespace(get_name=lambda: name)
+
+
+def _checkpoint_od_config(
+    quant_config,
+    *,
+    policy: object | None = None,
+    additional_config: dict[str, object] | None = None,
+):
+    quant_config_name = quant_config.get_name()
+    quant_algo = {
+        "modelopt": "FP8",
+        "modelopt_fp4": "NVFP4",
+    }.get(quant_config_name, "FP8")
+    disk_quant_config: dict[str, object] = {
+        "quant_method": quant_config_name,
+        "quant_algo": quant_algo,
+    }
+    if policy is not None:
+        disk_quant_config["runtime"] = {"diffusion_step_policy": policy}
+    return SimpleNamespace(
+        quantization_config=quant_config,
+        tf_model_config=SimpleNamespace(
+            params={"quantization_config": disk_quant_config},
+            quant_config=quant_config,
+        ),
+        additional_config=additional_config,
+    )
+
+
+def _swizzle_blockscale_cpu(scale: torch.Tensor) -> torch.Tensor:
+    rows, cols = scale.shape
+    padded_rows = ((rows + 127) // 128) * 128
+    padded_cols = ((cols + 3) // 4) * 4
+    padded = torch.zeros(
+        (1, padded_rows, padded_cols),
+        dtype=scale.dtype,
+        device=scale.device,
+    )
+    padded[0, :rows, :cols] = scale
+    padded = padded.reshape(1, padded_rows // 128, 4, 32, padded_cols // 4, 4)
+    return (
+        padded.permute(0, 1, 4, 3, 2, 5)
+        .contiguous()
+        .reshape(
+            padded_rows,
+            padded_cols,
+        )
+    )
+
+
+def test_config_parses_nested_schedule_and_reasoner_policy() -> None:
+    config = Cosmos3MixedPrecisionConfig.from_additional_config(
+        {
+            "cosmos3_mixed_precision": {
+                "first_steps": 2,
+                "last_steps": 4,
+                "reasoner": "native",
+            }
+        }
+    )
+
+    assert config is not None
+    assert config.reasoner == "native"
+    assert [index for index in range(10) if config.use_high_precision(index, 10)] == [0, 1, 6, 7, 8, 9]
+
+
+@pytest.mark.parametrize(
+    ("first_steps", "last_steps", "selected"),
+    [
+        (0, 2, [5, 6]),
+        (2, 0, [0, 1]),
+        (0, 0, []),
+        (4, 4, list(range(7))),
+    ],
+)
+def test_schedule_boundaries_and_overlap(
+    first_steps: int,
+    last_steps: int,
+    selected: list[int],
+) -> None:
+    config = Cosmos3MixedPrecisionConfig(
+        first_steps=first_steps,
+        last_steps=last_steps,
+    )
+    assert [index for index in range(7) if config.use_high_precision(index, 7)] == selected
+
+
+def test_one_step_request_honors_boundary_precision() -> None:
+    config = Cosmos3MixedPrecisionConfig(first_steps=1, last_steps=1)
+    assert config.use_high_precision(0, 1)
+
+
+@pytest.mark.parametrize(
+    ("values", "message"),
+    [
+        ({"cosmos3_mixed_precision": True}, "must be a mapping"),
+        ({"cosmos3_mixed_precision": {"first_steps": -1}}, "non-negative"),
+        ({"cosmos3_mixed_precision": {"reasoner": "fp16"}}, "must be one of"),
+        ({"cosmos3_mixed_precision": {"enabled": 0}}, "must be a boolean"),
+        (
+            {"cosmos3_mixed_precision": {"enabled": False, "first_steps": 1}},
+            "cannot be combined",
+        ),
+        ({"cosmos3_mixed_precision": {"unknown": 1}}, "Unknown"),
+    ],
+)
+def test_config_rejects_invalid_values(values: dict, message: str) -> None:
+    with pytest.raises((TypeError, ValueError), match=message):
+        Cosmos3MixedPrecisionConfig.from_additional_config(values)
+
+
+def test_config_presence_enables_defaults() -> None:
+    assert Cosmos3MixedPrecisionConfig.from_additional_config({}) is None
+    config = Cosmos3MixedPrecisionConfig.from_additional_config({"cosmos3_mixed_precision": {}})
+    assert config == Cosmos3MixedPrecisionConfig()
+
+
+def test_noop_config_does_not_install_runtime() -> None:
+    assert (
+        Cosmos3MixedPrecisionConfig.from_additional_config(
+            {
+                "cosmos3_mixed_precision": {
+                    "first_steps": 0,
+                    "last_steps": 0,
+                    "reasoner": "native",
+                }
+            }
+        )
+        is None
+    )
+
+
+def test_additional_config_distinguishes_absent_and_disabled() -> None:
+    assert Cosmos3MixedPrecisionConfig.resolve_additional_config({}) == (False, None)
+    assert Cosmos3MixedPrecisionConfig.resolve_additional_config({"cosmos3_mixed_precision": {"enabled": False}}) == (
+        True,
+        None,
+    )
+    assert Cosmos3MixedPrecisionConfig.resolve_additional_config(
+        {
+            "cosmos3_mixed_precision": {
+                "first_steps": 0,
+                "last_steps": 0,
+                "reasoner": "native",
+            }
+        }
+    ) == (True, None)
+
+
+def test_checkpoint_policy_round_trip_through_transformer_config() -> None:
+    from vllm_omni.diffusion.data import TransformerConfig
+
+    disk = {
+        "_class_name": "Cosmos3VFMTransformer",
+        "quantization_config": {
+            "quant_method": "modelopt",
+            "quant_algo": "FP8",
+            "runtime": {"diffusion_step_policy": _policy(first_steps=2, last_steps=4)},
+        },
+    }
+
+    tf_config = TransformerConfig.from_dict(disk)
+    assert tf_config.quant_config is not None
+    assert tf_config.quant_config.get_name() == "modelopt"
+    assert tf_config.to_dict()["quantization_config"]["runtime"] == disk["quantization_config"]["runtime"]
+    assert read_checkpoint_policy(SimpleNamespace(tf_model_config=tf_config)) == Cosmos3MixedPrecisionConfig(
+        first_steps=2,
+        last_steps=4,
+        reasoner="a16",
+    )
+
+
+@pytest.mark.parametrize("name", ["modelopt", "modelopt_fp4"])
+def test_checkpoint_policy_supports_serialized_modelopt_formats(name: str) -> None:
+    od_config = _checkpoint_od_config(_fake_quant_config(name), policy=_policy())
+    assert read_checkpoint_policy(od_config) == Cosmos3MixedPrecisionConfig()
+
+
+@pytest.mark.parametrize("name", ["modelopt", "modelopt_fp4"])
+def test_checkpoint_policy_rejects_nonserialized_modelopt_formats(name: str) -> None:
+    od_config = _checkpoint_od_config(
+        _fake_quant_config(name, serialized=False),
+        policy=_policy(),
+    )
+    with pytest.raises(ValueError, match="serialized ModelOpt"):
+        read_checkpoint_policy(od_config)
+
+
+def test_checkpoint_policy_rejects_mixed_modelopt_format() -> None:
+    with pytest.raises(ValueError, match="serialized ModelOpt FP8 or NVFP4"):
+        read_checkpoint_policy(_checkpoint_od_config(_fake_quant_config("modelopt_mixed"), policy=_policy()))
+
+
+def test_checkpoint_policy_missing_metadata_preserves_ordinary_path() -> None:
+    assert read_checkpoint_policy(_checkpoint_od_config(_fake_quant_config("modelopt"))) is None
+
+
+def test_checkpoint_policy_rejects_malformed_runtime_container() -> None:
+    od_config = _checkpoint_od_config(_fake_quant_config("modelopt"))
+    od_config.tf_model_config.params["quantization_config"]["runtime"] = True
+    with pytest.raises(TypeError, match="runtime must be a mapping"):
+        read_checkpoint_policy(od_config)
+
+
+@pytest.mark.parametrize(
+    ("field", "value", "message"),
+    [
+        ("schema_version", 2, "schema_version"),
+        ("type", "sigma_ranges", "type"),
+        ("index_space", "scheduler_timestep", "index_space"),
+        ("default_mode", "a16", "default_mode"),
+        ("overlap", "native", "overlap"),
+        ("reasoner", "fp16", "reasoner"),
+    ],
+)
+def test_checkpoint_policy_rejects_invalid_schema(field: str, value: object, message: str) -> None:
+    policy = _policy()
+    policy[field] = value
+    with pytest.raises(ValueError, match=message):
+        read_checkpoint_policy(_checkpoint_od_config(_fake_quant_config("modelopt"), policy=policy))
+
+
+@pytest.mark.parametrize(
+    ("step_range", "message"),
+    [
+        ({"count": -1, "mode": "a16"}, "non-negative"),
+        ({"count": True, "mode": "a16"}, "non-negative"),
+        ({"count": 1, "mode": "native"}, "mode"),
+        ({"mode": "a16"}, "Missing"),
+        ({"count": 1, "mode": "a16", "sigma": 1.0}, "Unknown"),
+    ],
+)
+def test_checkpoint_policy_rejects_invalid_step_range(step_range: object, message: str) -> None:
+    policy = _policy()
+    policy["first_steps"] = step_range
+    with pytest.raises((TypeError, ValueError), match=message):
+        read_checkpoint_policy(_checkpoint_od_config(_fake_quant_config("modelopt"), policy=policy))
+
+
+def test_checkpoint_policy_rejects_unsupported_quantization_config() -> None:
+    with pytest.raises(ValueError, match="serialized ModelOpt FP8 or NVFP4"):
+        read_checkpoint_policy(_checkpoint_od_config(_fake_quant_config("compressed-tensors"), policy=_policy()))
+
+
+def test_checkpoint_policy_for_other_scope_is_not_applied() -> None:
+    od_config = _checkpoint_od_config(
+        _fake_quant_config("modelopt"),
+        policy=_policy(scope=["vae"]),
+    )
+    assert read_checkpoint_policy(od_config) is None
+
+
+def test_additional_config_overrides_and_disables_checkpoint_policy() -> None:
+    checkpoint = _checkpoint_od_config(
+        _fake_quant_config("modelopt"),
+        policy=_policy(first_steps=1, last_steps=1),
+    )
+    assert resolve_mixed_precision_config(checkpoint) == (
+        Cosmos3MixedPrecisionConfig(first_steps=1, last_steps=1),
+        "checkpoint",
+    )
+
+    checkpoint.additional_config = {
+        "cosmos3_mixed_precision": {
+            "first_steps": 2,
+            "last_steps": 0,
+            "reasoner": "native",
+        }
+    }
+    assert resolve_mixed_precision_config(checkpoint) == (
+        Cosmos3MixedPrecisionConfig(first_steps=2, last_steps=0, reasoner="native"),
+        "additional_config",
+    )
+
+    checkpoint.additional_config = {"cosmos3_mixed_precision": {"enabled": False}}
+    assert resolve_mixed_precision_config(checkpoint) == (None, "additional_config_disabled")
+
+
+@pytest.mark.parametrize("name", ["modelopt", "modelopt_fp4"])
+def test_additional_config_supports_serialized_modelopt_formats_without_checkpoint_policy(name: str) -> None:
+    od_config = _checkpoint_od_config(
+        _fake_quant_config(name),
+        additional_config={"cosmos3_mixed_precision": {}},
+    )
+    assert resolve_mixed_precision_config(od_config) == (Cosmos3MixedPrecisionConfig(), "additional_config")
+
+
+@pytest.mark.parametrize("has_checkpoint_policy", [False, True])
+@pytest.mark.parametrize(
+    ("name", "serialized"),
+    [
+        ("modelopt_mixed", True),
+        ("mxfp8", True),
+        ("compressed-tensors", True),
+        ("modelopt", False),
+        ("modelopt_fp4", False),
+    ],
+)
+def test_additional_config_rejects_incompatible_quantization(
+    name: str, serialized: bool, has_checkpoint_policy: bool
+) -> None:
+    od_config = _checkpoint_od_config(
+        _fake_quant_config(name, serialized=serialized),
+        policy=_policy() if has_checkpoint_policy else None,
+        additional_config={"cosmos3_mixed_precision": {}},
+    )
+    with pytest.raises(ValueError, match="serialized ModelOpt"):
+        resolve_mixed_precision_config(od_config)
+
+
+@pytest.mark.parametrize("additional_config", [None, {"cosmos3_mixed_precision": {}}])
+def test_policy_validates_effective_quantization_config(additional_config: dict[str, object] | None) -> None:
+    od_config = _checkpoint_od_config(
+        _fake_quant_config("modelopt"),
+        policy=_policy(),
+        additional_config=additional_config,
+    )
+    od_config.quantization_config = _fake_quant_config("modelopt_mixed")
+    with pytest.raises(ValueError, match="serialized ModelOpt FP8 or NVFP4"):
+        resolve_mixed_precision_config(od_config)
+
+
+def test_additional_config_rejects_missing_quantization_config() -> None:
+    od_config = SimpleNamespace(additional_config={"cosmos3_mixed_precision": {}})
+    with pytest.raises(ValueError, match="serialized ModelOpt FP8 or NVFP4"):
+        resolve_mixed_precision_config(od_config)
+
+
+@pytest.mark.parametrize("policy", [_policy(), True])
+def test_additional_config_disable_bypasses_incompatible_checkpoint(policy: object) -> None:
+    od_config = _checkpoint_od_config(
+        _fake_quant_config("modelopt_mixed"),
+        policy=policy,
+        additional_config={"cosmos3_mixed_precision": {"enabled": False}},
+    )
+    od_config.quantization_config = _fake_quant_config("modelopt_mixed")
+    assert resolve_mixed_precision_config(od_config) == (None, "additional_config_disabled")
+
+
+def test_runtime_allows_standard_offload_and_rejects_distributed(monkeypatch) -> None:
+    from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3
+
+    monkeypatch.setattr(
+        transformer_cosmos3,
+        "get_tensor_model_parallel_world_size",
+        lambda: 1,
+    )
+    od_config = SimpleNamespace(
+        enable_cpu_offload=True,
+        enable_layerwise_offload=True,
+        enable_distributed_layerwise_offload=False,
+        parallel_config=SimpleNamespace(use_hsdp=False),
+    )
+    transformer_cosmos3._validate_mixed_precision_runtime(
+        Cosmos3MixedPrecisionConfig(),
+        od_config,
+    )
+    od_config.enable_distributed_layerwise_offload = True
+    with pytest.raises(ValueError, match="distributed layer-wise offload"):
+        transformer_cosmos3._validate_mixed_precision_runtime(
+            Cosmos3MixedPrecisionConfig(),
+            od_config,
+        )
+    od_config.enable_distributed_layerwise_offload = False
+    od_config.max_num_seqs = 2
+    with pytest.raises(ValueError, match="one active request"):
+        transformer_cosmos3._validate_mixed_precision_runtime(
+            Cosmos3MixedPrecisionConfig(),
+            od_config,
+        )
+
+
+class _Layer(torch.nn.Module):
+    def __init__(self, output_size: int, input_size: int) -> None:
+        super().__init__()
+        self.output_size_per_partition = output_size
+        self.input_size_per_partition = input_size
+
+
+class _BaseMethod:
+    def __init__(self, *, nvfp4: bool = False) -> None:
+        self.processed = False
+        self.apply_calls = 0
+        self.fp8_linear = object()
+        if nvfp4:
+            from vllm.model_executor.kernels.linear.nvfp4.cutlass import (
+                CutlassNvFp4LinearKernel,
+            )
+
+            self.kernel = object.__new__(CutlassNvFp4LinearKernel)
+        else:
+            self.kernel = None
+
+    def create_weights(self, *args, **kwargs) -> None:
+        pass
+
+    def process_weights_after_loading(self, layer) -> None:
+        self.processed = True
+        if layer.weight.dtype == torch.float8_e4m3fn:
+            layer.weight = layer.weight.t().contiguous()
+            layer.weight_scale = layer.weight_scale.max().reshape(1)
+        elif layer.weight.dtype == torch.uint8:
+            layer.weight_scale = _swizzle_blockscale_cpu(layer.weight_scale)
+            layer.weight_global_scale = layer.weight_scale_2.max().float().reshape(1)
+            del layer.weight_scale_2
+
+    def apply(self, layer, x, bias=None):
+        del bias
+        self.apply_calls += 1
+        return torch.full(
+            (*x.shape[:-1], layer.output_size_per_partition),
+            17,
+            dtype=x.dtype,
+            device=x.device,
+        )
+
+
+def _fp8_layer(output_size: int = 2, input_size: int = 4) -> _Layer:
+    layer = _Layer(output_size, input_size)
+    layer.weight = torch.ones(output_size, input_size, dtype=torch.float8_e4m3fn)
+    layer.weight_scale = torch.tensor([0.5], dtype=torch.float32)
+    return layer
+
+
+def _nvfp4_layer(output_size: int = 128, input_size: int = 64) -> _Layer:
+    layer = _Layer(output_size, input_size)
+    layer.weight = torch.full((output_size, input_size // 2), 0x22, dtype=torch.uint8)
+    layer.weight_scale = torch.ones(output_size, input_size // 16, dtype=torch.float8_e4m3fn)
+    layer.weight_scale_2 = torch.ones(1)
+    return layer
+
+
+def _runtime_and_method(
+    strategy,
+    layer: _Layer,
+    *,
+    path: str = "generation",
+    reasoner: str = "native",
+):
+    config = Cosmos3MixedPrecisionConfig(
+        first_steps=1,
+        last_steps=1,
+        reasoner=reasoner,
+    )
+    runtime = Cosmos3MixedPrecisionRuntime(config)
+    base = _BaseMethod(nvfp4=isinstance(strategy, Nvfp4W4A4W4A16Strategy))
+    method = Cosmos3MixedPrecisionLinearMethod(
+        base,
+        strategy,
+        runtime,
+        f"{path}.linear",
+        path,
+    )
+    method.process_weights_after_loading(layer)
+    return runtime, base, method
+
+
+@pytest.mark.parametrize(
+    ("strategy", "layer"),
+    [
+        (Fp8W8A8W8A16Strategy(), _fp8_layer()),
+        (Nvfp4W4A4W4A16Strategy(), _nvfp4_layer()),
+    ],
+)
+def test_generation_dispatches_native_middle_and_a16_edges(strategy, layer) -> None:
+    runtime, base, method = _runtime_and_method(strategy, layer)
+    x = torch.ones(1, layer.input_size_per_partition, dtype=torch.bfloat16)
+
+    runtime.set_step(1, 3)
+    assert torch.equal(
+        method.apply(layer, x),
+        torch.full((1, layer.output_size_per_partition), 17, dtype=x.dtype),
+    )
+    runtime.set_step(0, 3)
+    assert torch.equal(
+        method.apply(layer, x),
+        torch.nn.functional.linear(x, strategy.materialize(layer)),
+    )
+    assert base.apply_calls == 1
+
+
+@pytest.mark.parametrize(
+    ("strategy", "layer"),
+    [
+        (Fp8W8A8W8A16Strategy(), _fp8_layer()),
+        (Nvfp4W4A4W4A16Strategy(), _nvfp4_layer()),
+    ],
+)
+def test_a16_uses_only_live_native_weights(strategy, layer) -> None:
+    runtime, base, method = _runtime_and_method(strategy, layer)
+    assert base.processed
+    assert not hasattr(layer, "_cosmos3_precision_weight")
+    assert not hasattr(layer, "_cosmos3_precision_weight_scale")
+
+    runtime.set_step(0, 3)
+    output = method.apply(
+        layer,
+        torch.ones(1, layer.input_size_per_partition, dtype=torch.bfloat16),
+    )
+    assert output.shape == (1, layer.output_size_per_partition)
+
+
+def test_a16_follows_live_weight_storage_after_offload_rebinding() -> None:
+    layer = _fp8_layer()
+    strategy = Fp8W8A8W8A16Strategy()
+    runtime, _, method = _runtime_and_method(strategy, layer)
+    runtime.set_step(0, 3)
+    x = torch.ones(1, 4, dtype=torch.bfloat16)
+    before = method.apply(layer, x)
+
+    layer.weight = torch.zeros_like(layer.weight)
+    after = method.apply(layer, x)
+    assert not torch.equal(before, after)
+    assert torch.count_nonzero(after) == 0
+
+
+def test_fp8_rejects_smoothquant() -> None:
+    layer = _fp8_layer()
+    layer.pre_quant_scale = torch.ones(1)
+    with pytest.raises(ValueError, match="SmoothQuant"):
+        Fp8W8A8W8A16Strategy().validate_before_processing(
+            _BaseMethod(),
+            layer,
+            "gen.linear",
+        )
+
+
+def test_fp8_rejects_non_tensorwise_scales() -> None:
+    layer = _fp8_layer(output_size=3, input_size=2)
+    layer.weight_scale = torch.tensor([[1.0], [2.0], [3.0]])
+    with pytest.raises(ValueError, match="tensorwise FP8"):
+        Fp8W8A8W8A16Strategy().validate_before_processing(
+            _BaseMethod(),
+            layer,
+            "gen.linear",
+        )
+
+
+def test_nvfp4_reference_materializes_live_cutlass_weights() -> None:
+    layer = _nvfp4_layer()
+    strategy = Nvfp4W4A4W4A16Strategy()
+    _runtime_and_method(strategy, layer)
+    output = strategy.materialize(layer)
+    assert output.shape == (128, 64)
+    assert torch.equal(output, torch.ones_like(output))
+
+
+def test_nvfp4_rejects_fused_global_scales() -> None:
+    layer = _nvfp4_layer()
+    layer.weight_scale_2 = torch.tensor([1.0, 2.0])
+    strategy = Nvfp4W4A4W4A16Strategy()
+    with pytest.raises(ValueError, match="NVFP4 scales"):
+        strategy.validate_before_processing(
+            _BaseMethod(nvfp4=True),
+            layer,
+            "gen.qkv",
+        )
+
+
+@pytest.mark.parametrize(
+    ("policy", "uses_native"),
+    [("a16", False), ("native", True)],
+)
+def test_reasoner_policy_is_independent_of_generation_step(
+    policy: str,
+    uses_native: bool,
+) -> None:
+    layer = _fp8_layer()
+    runtime, base, method = _runtime_and_method(
+        Fp8W8A8W8A16Strategy(),
+        layer,
+        path="reasoner",
+        reasoner=policy,
+    )
+    runtime.set_step(1, 3)
+    method.apply(layer, torch.ones(1, 4, dtype=torch.bfloat16))
+    assert (base.apply_calls == 1) is uses_native
+
+
+def test_reset_clears_generation_state() -> None:
+    runtime = Cosmos3MixedPrecisionRuntime(Cosmos3MixedPrecisionConfig(first_steps=1, last_steps=1))
+    runtime.set_step(0, 5)
+    assert runtime.use_high_precision("generation")
+    runtime.reset()
+    assert not runtime.use_high_precision("generation")
+
+
+def test_install_discovers_generation_and_opt_in_reasoner(monkeypatch) -> None:
+    class _FakeLinear(torch.nn.Module):
+        def __init__(self, prefix: str) -> None:
+            super().__init__()
+            self.prefix = prefix
+            self.input_size_per_partition = 4
+            self.output_size_per_partition = 2
+            self.quant_method = SimpleNamespace(name="fp8")
+
+    monkeypatch.setattr(runtime_impl, "LinearBase", _FakeLinear)
+    strategy = Fp8W8A8W8A16Strategy()
+    monkeypatch.setattr(strategy, "accepts", lambda method: getattr(method, "name", None) == "fp8")
+    monkeypatch.setattr(runtime_impl, "_STRATEGIES", (strategy,))
+    runtime = Cosmos3MixedPrecisionRuntime(Cosmos3MixedPrecisionConfig(reasoner="a16"))
+    reasoner = _FakeLinear("reasoner.q_proj")
+    generation = [
+        _FakeLinear("generation.q_proj"),
+        _FakeLinear("generation.out_proj"),
+    ]
+    transformer = SimpleNamespace(
+        language_model=SimpleNamespace(layers=torch.nn.Sequential(reasoner)),
+        gen_layers=torch.nn.Sequential(*generation),
+    )
+    runtime.install(transformer)
+    assert isinstance(reasoner.quant_method, Cosmos3MixedPrecisionLinearMethod)
+    assert all(isinstance(layer.quant_method, Cosmos3MixedPrecisionLinearMethod) for layer in generation)
+
+
+def test_pipeline_helpers_forward_and_reset_schedule() -> None:
+    from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import (
+        Cosmos3OmniDiffusersPipeline,
+    )
+
+    calls = []
+    pipeline = object.__new__(Cosmos3OmniDiffusersPipeline)
+    pipeline.transformer = SimpleNamespace(
+        set_mixed_precision_step=lambda step, count: calls.append((step, count)),
+        reset_mixed_precision=lambda: calls.append("reset"),
+    )
+    pipeline._set_mixed_precision_step(2, 7)
+    pipeline._reset_mixed_precision()
+    assert calls == [(2, 7), "reset"]

vllm_omni/diffusion/models/cosmos3/mixed_precision/__init__.py

@@ -0,0 +1,14 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Scheduled A16 execution for quantized Cosmos3 checkpoints."""
+
+from .checkpoint import resolve_mixed_precision_config
+from .config import Cosmos3MixedPrecisionConfig
+from .runtime import Cosmos3MixedPrecisionRuntime
+
+__all__ = [
+    "Cosmos3MixedPrecisionConfig",
+    "Cosmos3MixedPrecisionRuntime",
+    "resolve_mixed_precision_config",
+]

vllm_omni/diffusion/models/cosmos3/mixed_precision/checkpoint.py

@@ -0,0 +1,176 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Checkpoint discovery for the Cosmos3 diffusion-step precision policy."""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+from typing import Literal, cast
+
+from .config import Cosmos3MixedPrecisionConfig, ReasonerPolicy, _non_negative_int
+
+PolicySource = Literal[
+    "none",
+    "checkpoint",
+    "additional_config",
+    "additional_config_disabled",
+]
+
+_COMPONENT = "transformer"
+_POLICY_FIELDS = frozenset(
+    {
+        "schema_version",
+        "type",
+        "index_space",
+        "scope",
+        "default_mode",
+        "first_steps",
+        "last_steps",
+        "overlap",
+        "reasoner",
+    }
+)
+_STEP_RANGE_FIELDS = frozenset({"count", "mode"})
+_SUPPORTED_MODEL_OPT_CONFIGS = frozenset(
+    {
+        "modelopt",
+        "modelopt_fp4",
+    }
+)
+
+
+def resolve_mixed_precision_config(
+    od_config: object,
+) -> tuple[Cosmos3MixedPrecisionConfig | None, PolicySource]:
+    """Resolve an explicit runtime override before the checkpoint default."""
+    override_present, override = Cosmos3MixedPrecisionConfig.resolve_additional_config(
+        getattr(od_config, "additional_config", None)
+    )
+    if override_present:
+        if override is not None:
+            _validate_checkpoint_quantization(od_config)
+        source: PolicySource = "additional_config" if override is not None else "additional_config_disabled"
+        return override, source
+
+    checkpoint = read_checkpoint_policy(od_config)
+    if checkpoint is not None:
+        return checkpoint, "checkpoint"
+    return None, "none"
+
+
+def read_checkpoint_policy(od_config: object) -> Cosmos3MixedPrecisionConfig | None:
+    """Read and validate a policy from ``transformer/config.json``.
+
+    Missing metadata preserves ordinary checkpoint behavior.  Metadata that is
+    present but malformed or incompatible fails closed.
+    """
+    tf_config = getattr(od_config, "tf_model_config", None)
+    params = getattr(tf_config, "params", None)
+    if not isinstance(params, Mapping):
+        return None
+
+    disk_quant_config = params.get("quantization_config")
+    if not isinstance(disk_quant_config, Mapping):
+        return None
+
+    if "runtime" not in disk_quant_config:
+        return None
+    runtime_config = disk_quant_config["runtime"]
+    if not isinstance(runtime_config, Mapping):
+        raise TypeError("quantization_config.runtime must be a mapping")
+
+    if "diffusion_step_policy" not in runtime_config:
+        return None
+    raw_policy = runtime_config["diffusion_step_policy"]
+    if not isinstance(raw_policy, Mapping):
+        raise TypeError("quantization_config.runtime.diffusion_step_policy must be a mapping")
+
+    policy = _parse_policy(raw_policy)
+    if policy is None:
+        return None
+    _validate_checkpoint_quantization(od_config)
+    return policy
+
+
+def _parse_policy(policy: Mapping[str, object]) -> Cosmos3MixedPrecisionConfig | None:
+    fields = set(policy)
+    unknown = fields - _POLICY_FIELDS
+    if unknown:
+        raise ValueError(f"Unknown diffusion_step_policy fields: {sorted(unknown)}")
+    missing = _POLICY_FIELDS - fields
+    if missing:
+        raise ValueError(f"Missing diffusion_step_policy fields: {sorted(missing)}")
+
+    schema_version = policy["schema_version"]
+    if not isinstance(schema_version, int) or isinstance(schema_version, bool) or schema_version != 1:
+        raise ValueError("diffusion_step_policy.schema_version must be the integer 1")
+    if policy["type"] != "first_last_n":
+        raise ValueError("diffusion_step_policy.type must be 'first_last_n'")
+    if policy["index_space"] != "denoising_loop_iteration":
+        raise ValueError("diffusion_step_policy.index_space must be 'denoising_loop_iteration'")
+    if policy["default_mode"] != "native":
+        raise ValueError("diffusion_step_policy.default_mode must be 'native'")
+    if policy["overlap"] != "a16":
+        raise ValueError("diffusion_step_policy.overlap must be 'a16'")
+
+    scope = policy["scope"]
+    if not isinstance(scope, list) or not scope or not all(isinstance(item, str) for item in scope):
+        raise TypeError("diffusion_step_policy.scope must be a non-empty list of strings")
+
+    first_steps = _parse_step_range(policy["first_steps"], "first_steps")
+    last_steps = _parse_step_range(policy["last_steps"], "last_steps")
+
+    reasoner = policy["reasoner"]
+    if reasoner not in {"native", "a16"}:
+        raise ValueError("diffusion_step_policy.reasoner must be 'native' or 'a16'")
+
+    if _COMPONENT not in scope:
+        return None
+    return Cosmos3MixedPrecisionConfig(
+        first_steps=first_steps,
+        last_steps=last_steps,
+        reasoner=cast(ReasonerPolicy, reasoner),
+    )
+
+
+def _parse_step_range(value: object, name: str) -> int:
+    if not isinstance(value, Mapping):
+        raise TypeError(f"diffusion_step_policy.{name} must be a mapping")
+    fields = set(value)
+    unknown = fields - _STEP_RANGE_FIELDS
+    if unknown:
+        raise ValueError(f"Unknown diffusion_step_policy.{name} fields: {sorted(unknown)}")
+    missing = _STEP_RANGE_FIELDS - fields
+    if missing:
+        raise ValueError(f"Missing diffusion_step_policy.{name} fields: {sorted(missing)}")
+    if value["mode"] != "a16":
+        raise ValueError(f"diffusion_step_policy.{name}.mode must be 'a16'")
+    return _non_negative_int(value["count"], f"diffusion_step_policy.{name}.count")
+
+
+def _validate_checkpoint_quantization(od_config: object) -> None:
+    # Match the effective quantization used to construct the transformer linears.
+    quant_config = getattr(od_config, "quantization_config", None)
+    if quant_config is None:
+        tf_config = getattr(od_config, "tf_model_config", None)
+        quant_config = getattr(tf_config, "quant_config", None)
+
+    get_name = getattr(quant_config, "get_name", None)
+    name = get_name() if callable(get_name) else None
+    if name not in _SUPPORTED_MODEL_OPT_CONFIGS:
+        raise ValueError(f"diffusion_step_policy requires a serialized ModelOpt FP8 or NVFP4 checkpoint, got {name!r}")
+
+    if name == "modelopt":
+        _require_serialized(quant_config, "is_checkpoint_fp8_serialized", "FP8")
+        return
+    if name == "modelopt_fp4":
+        if getattr(quant_config, "quant_method", "NVFP4") != "NVFP4":
+            raise ValueError("diffusion_step_policy requires checkpoint-native NVFP4 W4A4 linears")
+        _require_serialized(quant_config, "is_checkpoint_nvfp4_serialized", "NVFP4")
+        return
+
+
+def _require_serialized(quant_config: object | None, flag: str, format_name: str) -> None:
+    if not bool(getattr(quant_config, flag, False)):
+        raise ValueError(f"diffusion_step_policy requires serialized ModelOpt {format_name} weights")

vllm_omni/diffusion/models/cosmos3/mixed_precision/config.py

@@ -0,0 +1,106 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Configuration parsing and denoising-step policy for mixed precision."""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+from dataclasses import dataclass
+from typing import Literal, cast
+
+ReasonerPolicy = Literal["native", "a16"]
+
+_REASONER_POLICIES = frozenset({"native", "a16"})
+_CONFIG_FIELDS = frozenset({"enabled", "first_steps", "last_steps", "reasoner"})
+
+
+@dataclass(frozen=True)
+class Cosmos3MixedPrecisionConfig:
+    """Validated precision policy shared by every quantization strategy."""
+
+    first_steps: int = 3
+    last_steps: int = 3
+    reasoner: ReasonerPolicy = "a16"
+
+    @classmethod
+    def from_additional_config(
+        cls,
+        additional_config: Mapping[str, object] | None,
+    ) -> Cosmos3MixedPrecisionConfig | None:
+        """Parse Cosmos3 fields from vLLM-Omni's additional configuration."""
+        _, config = cls.resolve_additional_config(additional_config)
+        return config
+
+    @classmethod
+    def resolve_additional_config(
+        cls,
+        additional_config: Mapping[str, object] | None,
+    ) -> tuple[bool, Cosmos3MixedPrecisionConfig | None]:
+        """Return whether an override was supplied and its effective policy.
+
+        The separate presence bit distinguishes an absent runtime override from
+        an explicit disable.  That distinction matters when a checkpoint owns a
+        default policy.
+        """
+        if additional_config is not None and not isinstance(additional_config, Mapping):
+            raise TypeError("additional_config must be a mapping")
+        values = additional_config or {}
+        if "cosmos3_mixed_precision" not in values:
+            return False, None
+        raw_config = values["cosmos3_mixed_precision"]
+        if not isinstance(raw_config, Mapping):
+            raise TypeError("cosmos3_mixed_precision must be a mapping")
+        unknown = set(raw_config) - _CONFIG_FIELDS
+        if unknown:
+            raise ValueError(f"Unknown cosmos3_mixed_precision fields: {sorted(unknown)}")
+
+        enabled = raw_config.get("enabled", True)
+        if not isinstance(enabled, bool):
+            raise TypeError("cosmos3_mixed_precision.enabled must be a boolean")
+        if not enabled:
+            conflicting = set(raw_config) - {"enabled"}
+            if conflicting:
+                raise ValueError(
+                    "cosmos3_mixed_precision.enabled=false cannot be combined with "
+                    f"schedule fields: {sorted(conflicting)}"
+                )
+            return True, None
+
+        first_steps = _non_negative_int(
+            raw_config.get("first_steps", 3),
+            "cosmos3_mixed_precision.first_steps",
+        )
+        last_steps = _non_negative_int(
+            raw_config.get("last_steps", 3),
+            "cosmos3_mixed_precision.last_steps",
+        )
+
+        reasoner = str(raw_config.get("reasoner", "a16")).lower()
+        if reasoner not in _REASONER_POLICIES:
+            raise ValueError(
+                f"cosmos3_mixed_precision.reasoner must be one of {sorted(_REASONER_POLICIES)}, got {reasoner!r}"
+            )
+        config = cls(
+            first_steps=first_steps,
+            last_steps=last_steps,
+            reasoner=cast(ReasonerPolicy, reasoner),
+        )
+        if first_steps == 0 and last_steps == 0 and reasoner == "native":
+            return True, None
+        return True, config
+
+    def use_high_precision(self, step_index: int, num_steps: int) -> bool:
+        """Select the high path for an actual scheduler-step index."""
+        if num_steps <= 0:
+            raise ValueError(f"num_steps must be positive, got {num_steps}")
+        if step_index < 0 or step_index >= num_steps:
+            raise IndexError(f"step_index must be in [0, {num_steps}), got {step_index}")
+        return step_index < self.first_steps or step_index >= num_steps - self.last_steps
+
+
+def _non_negative_int(value: object, name: str) -> int:
+    """Validate an integer configuration field without accepting booleans."""
+    if not isinstance(value, int) or isinstance(value, bool) or value < 0:
+        raise TypeError(f"{name} must be a non-negative integer")
+    return value

vllm_omni/diffusion/models/cosmos3/mixed_precision/runtime.py

@@ -0,0 +1,116 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Format-agnostic schedule state, linear discovery, and dispatch."""
+
+from __future__ import annotations
+
+from typing import Literal
+
+import torch
+from vllm.model_executor.layers.linear import LinearBase, LinearMethodBase
+
+from .config import Cosmos3MixedPrecisionConfig
+from .strategy import (
+    Cosmos3PrecisionStrategy,
+    Fp8W8A8W8A16Strategy,
+    Nvfp4W4A4W4A16Strategy,
+)
+
+PrecisionPath = Literal["reasoner", "generation"]
+_STRATEGIES: tuple[Cosmos3PrecisionStrategy, ...] = (
+    Fp8W8A8W8A16Strategy(),
+    Nvfp4W4A4W4A16Strategy(),
+)
+
+
+def _strategy_for(method: object | None) -> Cosmos3PrecisionStrategy | None:
+    return next((strategy for strategy in _STRATEGIES if strategy.accepts(method)), None)
+
+
+class Cosmos3MixedPrecisionLinearMethod(LinearMethodBase):
+    """Dispatch between a checkpoint-native method and dense A16."""
+
+    def __init__(
+        self,
+        base_method: LinearMethodBase,
+        strategy: Cosmos3PrecisionStrategy,
+        runtime: Cosmos3MixedPrecisionRuntime,
+        module_name: str,
+        path: PrecisionPath,
+    ) -> None:
+        self.base_method = base_method
+        self.strategy = strategy
+        self.runtime = runtime
+        self.module_name = module_name
+        self.path = path
+
+    def create_weights(self, *args, **kwargs) -> None:
+        self.base_method.create_weights(*args, **kwargs)
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        self.strategy.validate_before_processing(
+            self.base_method,
+            layer,
+            self.module_name,
+        )
+        self.base_method.process_weights_after_loading(layer)
+        self.strategy.validate_after_processing(layer, self.module_name)
+
+    def apply(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        if not self.runtime.use_high_precision(self.path):
+            return self.base_method.apply(layer, x, bias)
+        return self.strategy.apply_high(layer, x, bias)
+
+
+class Cosmos3MixedPrecisionRuntime:
+    """Own one transformer's schedule state and wrapped linear inventory."""
+
+    def __init__(self, config: Cosmos3MixedPrecisionConfig) -> None:
+        self.config = config
+        # TODO: Make this request-local before Cosmos3 supports interleaved
+        # denoising requests on one transformer instance.
+        self._generation_high_precision = False
+
+    def install(self, transformer: torch.nn.Module) -> None:
+        components: dict[PrecisionPath, torch.nn.Module] = {
+            "generation": transformer.gen_layers,
+        }
+        if self.config.reasoner == "a16":
+            components["reasoner"] = transformer.language_model.layers
+
+        wrapped_count = 0
+        for path, component in components.items():
+            for local_name, layer in component.named_modules():
+                if not isinstance(layer, LinearBase):
+                    continue
+                base_method = getattr(layer, "quant_method", None)
+                strategy = _strategy_for(base_method)
+                if strategy is None:
+                    continue
+                module_name = getattr(layer, "prefix", None) or f"{path}.{local_name}"
+                layer.quant_method = Cosmos3MixedPrecisionLinearMethod(
+                    base_method,
+                    strategy,
+                    self,
+                    module_name,
+                    path,
+                )
+                wrapped_count += 1
+
+        if not wrapped_count:
+            raise ValueError("Cosmos3 mixed precision found no compatible FP8 or NVFP4 ModelOpt linears")
+
+    def use_high_precision(self, path: PrecisionPath) -> bool:
+        return self.config.reasoner == "a16" if path == "reasoner" else self._generation_high_precision
+
+    def set_step(self, step_index: int, num_steps: int) -> None:
+        self._generation_high_precision = self.config.use_high_precision(step_index, num_steps)
+
+    def reset(self) -> None:
+        self._generation_high_precision = False

vllm_omni/diffusion/models/cosmos3/mixed_precision/strategy.py

@@ -0,0 +1,218 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Format-specific live-weight validation and reference A16 materialization."""
+
+from __future__ import annotations
+
+import torch
+import torch.nn.functional as F
+from vllm.model_executor.layers.quantization.modelopt import (
+    ModelOptFp8LinearMethod,
+    ModelOptNvFp4LinearMethod,
+)
+from vllm.model_executor.layers.quantization.utils import nvfp4_emulation_utils
+
+_NVFP4_BLOCK_SIZE = 16
+
+
+class Cosmos3PrecisionStrategy:
+    """Quantization-format boundary behind the shared precision schedule."""
+
+    def accepts(self, method: object | None) -> bool:
+        """Return whether this strategy owns a native linear method."""
+        raise NotImplementedError
+
+    def validate_before_processing(
+        self,
+        method: object,
+        layer: torch.nn.Module,
+        module_name: str,
+    ) -> None:
+        """Reject native transforms that cannot support live dequantization."""
+
+    def validate_after_processing(
+        self,
+        layer: torch.nn.Module,
+        module_name: str,
+    ) -> None:
+        """Validate the live native weight representation."""
+        raise NotImplementedError
+
+    def materialize(self, layer: torch.nn.Module) -> torch.Tensor:
+        """Materialize the live quantized weight as a dense BF16 matrix."""
+        raise NotImplementedError
+
+    def apply_high(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: torch.Tensor | None,
+    ) -> torch.Tensor:
+        """Apply the reference dense A16 path."""
+        input_size = int(layer.input_size_per_partition)
+        output_size = int(layer.output_size_per_partition)
+        weight = self.materialize(layer)
+        expected_shape = (output_size, input_size)
+        if tuple(weight.shape) != expected_shape:
+            raise RuntimeError(f"Dense weight shape {tuple(weight.shape)} does not match {expected_shape}")
+        x_2d = x.reshape(-1, x.shape[-1])
+        if x_2d.shape[1] != input_size:
+            raise ValueError(f"Expected activation width {input_size}, got {x_2d.shape[1]}")
+        weight = weight.to(dtype=x.dtype)
+        if bias is not None and bias.dtype != x.dtype:
+            bias = bias.to(dtype=x.dtype)
+        output = F.linear(x_2d, weight, bias)
+        return output.view(*x.shape[:-1], output_size)
+
+
+class Fp8W8A8W8A16Strategy(Cosmos3PrecisionStrategy):
+    """Use native ModelOpt W8A8 or reference dense W8A16."""
+
+    def accepts(self, method: object | None) -> bool:
+        return isinstance(method, ModelOptFp8LinearMethod)
+
+    def validate_before_processing(
+        self,
+        method: object,
+        layer: torch.nn.Module,
+        module_name: str,
+    ) -> None:
+        if hasattr(layer, "pre_quant_scale"):
+            raise ValueError(f"{module_name} uses unsupported SmoothQuant pre_quant_scale")
+        weight_scale = getattr(layer, "weight_scale", None)
+        if not isinstance(weight_scale, torch.Tensor) or weight_scale.numel() != 1:
+            raise ValueError(f"{module_name} requires serialized tensorwise FP8 weights")
+        fp8_kernel = getattr(method, "fp8_linear", None)
+        from vllm.model_executor.kernels.linear.scaled_mm.marlin import (
+            MarlinFP8ScaledMMLinearKernel,
+        )
+
+        if isinstance(fp8_kernel, MarlinFP8ScaledMMLinearKernel):
+            raise ValueError(
+                f"{module_name} selected Marlin FP8, which repacks the only weight copy; "
+                "use a backend that retains canonical FP8 weights"
+            )
+
+    def validate_after_processing(
+        self,
+        layer: torch.nn.Module,
+        module_name: str,
+    ) -> None:
+        weight = getattr(layer, "weight", None)
+        scale = getattr(layer, "weight_scale", None)
+        if not isinstance(weight, torch.Tensor) or weight.dtype != torch.float8_e4m3fn:
+            dtype = getattr(weight, "dtype", None)
+            raise TypeError(f"{module_name} requires a canonical FP8 weight, got {dtype}")
+        if weight.ndim != 2:
+            raise ValueError(f"{module_name} requires a rank-2 FP8 weight, got {tuple(weight.shape)}")
+        if not isinstance(scale, torch.Tensor):
+            raise ValueError(f"{module_name} has an FP8 weight but no weight_scale")
+        if scale.numel() != 1:
+            raise ValueError(f"{module_name} requires one live FP8 weight scale")
+        input_size = int(layer.input_size_per_partition)
+        output_size = int(layer.output_size_per_partition)
+        if weight.shape[0] < input_size or weight.shape[1] < output_size:
+            raise ValueError(
+                f"{module_name} live FP8 weight shape {tuple(weight.shape)} does not cover "
+                f"(K, N)=({input_size}, {output_size})"
+            )
+        _validate_positive_finite_scale(scale, module_name)
+
+    def materialize(self, layer: torch.nn.Module) -> torch.Tensor:
+        input_size = int(layer.input_size_per_partition)
+        output_size = int(layer.output_size_per_partition)
+        weight = layer.weight[:input_size, :output_size]
+        scale = layer.weight_scale.reshape(1).to(device=weight.device, dtype=torch.float32)
+        return (weight.to(torch.float32) * scale).t().to(torch.bfloat16)
+
+
+class Nvfp4W4A4W4A16Strategy(Cosmos3PrecisionStrategy):
+    """Use native ModelOpt W4A4 or reference dense W4A16."""
+
+    def accepts(self, method: object | None) -> bool:
+        return isinstance(method, ModelOptNvFp4LinearMethod)
+
+    def validate_before_processing(
+        self,
+        method: object,
+        layer: torch.nn.Module,
+        module_name: str,
+    ) -> None:
+        kernel = getattr(method, "kernel", None)
+        from vllm.model_executor.kernels.linear.nvfp4.cutlass import (
+            CutlassNvFp4LinearKernel,
+        )
+        from vllm.model_executor.kernels.linear.nvfp4.flashinfer import (
+            FlashInferCuteDslNvFp4LinearKernel,
+            FlashInferCutlassNvFp4LinearKernel,
+        )
+
+        compatible_kernels = (
+            CutlassNvFp4LinearKernel,
+            FlashInferCuteDslNvFp4LinearKernel,
+            FlashInferCutlassNvFp4LinearKernel,
+        )
+        if not isinstance(kernel, compatible_kernels):
+            kernel_name = type(kernel).__name__ if kernel is not None else "none"
+            raise ValueError(
+                f"{module_name} selected unsupported NVFP4 backend {kernel_name}; "
+                "live W4A16 dequantization requires a CUTLASS-compatible layout"
+            )
+        global_scale = _nvfp4_global_scale(layer)
+        if global_scale is None or global_scale.numel() != 1:
+            raise ValueError(f"{module_name} is missing ModelOpt NVFP4 scales")
+
+    def validate_after_processing(
+        self,
+        layer: torch.nn.Module,
+        module_name: str,
+    ) -> None:
+        packed = getattr(layer, "weight", None)
+        scale = getattr(layer, "weight_scale", None)
+        global_scale = getattr(layer, "weight_global_scale", None)
+        if not isinstance(packed, torch.Tensor) or packed.dtype != torch.uint8:
+            dtype = getattr(packed, "dtype", None)
+            raise TypeError(f"{module_name} requires a live packed NVFP4 weight, got {dtype}")
+        if not isinstance(scale, torch.Tensor) or not isinstance(global_scale, torch.Tensor):
+            raise ValueError(f"{module_name} is missing live NVFP4 scales")
+        output_size = int(layer.output_size_per_partition)
+        input_size = int(layer.input_size_per_partition)
+        if input_size % _NVFP4_BLOCK_SIZE != 0:
+            raise ValueError(f"{module_name} input size {input_size} is not divisible by {_NVFP4_BLOCK_SIZE}")
+        expected_weight = (output_size, input_size // 2)
+        if packed.shape[0] < expected_weight[0] or packed.shape[1] < expected_weight[1]:
+            raise ValueError(
+                f"{module_name} live NVFP4 weight shape {tuple(packed.shape)} does not cover {expected_weight}"
+            )
+        _validate_positive_finite_scale(global_scale, module_name)
+        lut_handle = nvfp4_emulation_utils.kE2M1ToFloat_handle
+        if lut_handle.val.device != packed.device:
+            lut_handle.val = lut_handle.val.to(packed.device)
+
+    def materialize(self, layer: torch.nn.Module) -> torch.Tensor:
+        output_size = int(layer.output_size_per_partition)
+        input_size = int(layer.input_size_per_partition)
+        packed = layer.weight[:output_size, : input_size // 2]
+        return nvfp4_emulation_utils.dequantize_to_dtype(
+            packed,
+            layer.weight_scale,
+            layer.weight_global_scale,
+            torch.bfloat16,
+            _NVFP4_BLOCK_SIZE,
+            True,
+        )
+
+
+def _validate_positive_finite_scale(scale: torch.Tensor, module_name: str) -> None:
+    values = scale.detach().float()
+    if not torch.isfinite(values).all() or not (values > 0).all():
+        raise ValueError(f"{module_name} has a non-finite or non-positive weight scale")
+
+
+def _nvfp4_global_scale(layer: torch.nn.Module) -> torch.Tensor | None:
+    for name in ("weight_scale_2", "weight_global_scale"):
+        value = getattr(layer, name, None)
+        if isinstance(value, torch.Tensor):
+            return value
+    return None

vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py

@@ -1822,6 +1822,16 @@ class Cosmos3OmniDiffusersPipeline(
     def num_timesteps(self):
         return self._num_timesteps
 
+    def _set_mixed_precision_step(self, step_index: int, num_steps: int) -> None:
+        setter = getattr(self.transformer, "set_mixed_precision_step", None)
+        if setter is not None:
+            setter(step_index, num_steps)
+
+    def _reset_mixed_precision(self) -> None:
+        resetter = getattr(self.transformer, "reset_mixed_precision", None)
+        if resetter is not None:
+            resetter()
+
     @staticmethod
     def _distilled_unsupported_error(detail: str) -> ValueError:
         return ValueError(
@@ -2819,7 +2829,8 @@ class Cosmos3OmniDiffusersPipeline(
                 # Each CFG-parallel rank runs exactly one branch (rank 0 -> cond,
                 # else uncond), so session keying loads/stores only this rank's branch.
                 cfg_rank_is_negative = get_classifier_free_guidance_rank() != 0
-                for t in self.progress_bar(timesteps):
+                for step_index, t in enumerate(self.progress_bar(timesteps)):
+                    self._set_mixed_precision_step(step_index, len(timesteps))
                     timestep = t.unsqueeze(0)
                     # Outside the interval, scale=1 makes the combined output equal
                     # the cond branch. Every rank remains on the same iteration and
@@ -2858,7 +2869,8 @@ class Cosmos3OmniDiffusersPipeline(
                 uncond_cache: tuple = (None, None)
                 keep_uncond_for_cache = self._cache_requires_paired_cfg()
 
-                for t in self.progress_bar(timesteps):
+                for step_index, t in enumerate(self.progress_bar(timesteps)):
+                    self._set_mixed_precision_step(step_index, len(timesteps))
                     timestep = t.unsqueeze(0)
                     cfg_active = _cfg_active_at(t)
 
@@ -2912,7 +2924,8 @@ class Cosmos3OmniDiffusersPipeline(
             else:
                 # No CFG: a single cond branch per step. Bespoke (state None) keeps
                 # using the transformer-instance cache exactly as before.
-                for t in self.progress_bar(timesteps):
+                for step_index, t in enumerate(self.progress_bar(timesteps)):
+                    self._set_mixed_precision_step(step_index, len(timesteps))
                     timestep = t.unsqueeze(0)
                     self._kv_load_und(kv_state, is_negative=False)
                     noise_pred = self.transformer(
@@ -2928,6 +2941,7 @@ class Cosmos3OmniDiffusersPipeline(
                         self._kv_capture_und(kv_state, is_negative=False)
                     _assign_step_out(_step(noise_pred, t, latents, action_latents, sound_latents))
         finally:
+            self._reset_mixed_precision()
             # Cosmos3 currently receives a unique request_id rather than a
             # reusable rollout session id. Retaining its state would only pin
             # K/V buffers on device after this generation finishes.
@@ -3112,7 +3126,8 @@ class Cosmos3OmniDiffusersPipeline(
         self.transformer.reset_cache()
         self._cosmos3_branch_caches = {}
         try:
-            for t in self.progress_bar(timesteps):
+            for step_index, t in enumerate(self.progress_bar(timesteps)):
+                self._set_mixed_precision_step(step_index, len(timesteps))
                 timestep = t.unsqueeze(0)
                 step_guidance = guidance_scale if _active_at(t, guidance_interval) else 1.0
                 step_control = control_guidance if _active_at(t, control_guidance_interval) else 1.0
@@ -3224,6 +3239,7 @@ class Cosmos3OmniDiffusersPipeline(
                 )[0]
                 latents = velocity_mask * latents + (1.0 - velocity_mask) * condition_latents
         finally:
+            self._reset_mixed_precision()
             self._cosmos3_branch_caches = None
             self.transformer.reset_cache()
         return latents

vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py

@@ -40,6 +40,12 @@ from vllm_omni.diffusion.forward_context import get_forward_context, is_forward_
 from vllm_omni.diffusion.layers.norm import RMSNorm as _VllmRMSNorm
 from vllm_omni.platforms import current_omni_platform
 
+from .mixed_precision import (
+    Cosmos3MixedPrecisionConfig,
+    Cosmos3MixedPrecisionRuntime,
+    resolve_mixed_precision_config,
+)
+
 if TYPE_CHECKING:
     from vllm_omni.diffusion.offloader.sequential_backend import SequentialOffloadHook
 
@@ -136,6 +142,26 @@ def _od_config_get(od_config: Any, key: str, default: Any = None) -> Any:
     return default if value is None else value
 
 
+def _validate_mixed_precision_runtime(
+    config: Cosmos3MixedPrecisionConfig | None,
+    od_config: OmniDiffusionConfig,
+) -> None:
+    if config is None:
+        return
+    if get_tensor_model_parallel_world_size() != 1:
+        raise ValueError("Cosmos3 mixed precision currently supports tensor parallel size 1 only")
+    if int(getattr(od_config, "max_num_seqs", 1)) != 1:
+        raise ValueError("Cosmos3 mixed precision currently supports one active request per worker")
+    if bool(getattr(od_config, "enable_distributed_layerwise_offload", False)):
+        raise ValueError(
+            "Cosmos3 mixed precision does not support distributed layer-wise offload "
+            "because its direct loader bypasses ModelOpt post-load transformations"
+        )
+    parallel_config = getattr(od_config, "parallel_config", None)
+    if bool(getattr(parallel_config, "use_hsdp", False)):
+        raise ValueError("Cosmos3 mixed precision has not validated live backend weights under HSDP")
+
+
 def _as_bool(value: Any) -> bool:
     if isinstance(value, str):
         return value.strip().lower() in {"1", "true", "yes", "on"}
@@ -1241,6 +1267,19 @@ class Cosmos3VFMTransformer(nn.Module):
 
         dtype = od_config.dtype
         quant_config = getattr(od_config, "quantization_config", None) if od_config else None
+        mixed_precision_config, mixed_precision_source = resolve_mixed_precision_config(od_config)
+        if mixed_precision_config is None:
+            if mixed_precision_source == "additional_config_disabled":
+                logger.info("Cosmos3 checkpoint mixed-precision policy disabled by additional_config")
+        else:
+            logger.info(
+                "Cosmos3 mixed precision active (source=%s): first_steps=%d last_steps=%d reasoner=%s",
+                mixed_precision_source,
+                mixed_precision_config.first_steps,
+                mixed_precision_config.last_steps,
+                mixed_precision_config.reasoner,
+            )
+        _validate_mixed_precision_runtime(mixed_precision_config, od_config)
 
         self.language_model = self._language_model_cls(
             hidden_size=self.hidden_size,
@@ -1300,6 +1339,11 @@ class Cosmos3VFMTransformer(nn.Module):
             ]
         )
 
+        self.mixed_precision_runtime: Cosmos3MixedPrecisionRuntime | None = None
+        if mixed_precision_config is not None:
+            self.mixed_precision_runtime = Cosmos3MixedPrecisionRuntime(mixed_precision_config)
+            self.mixed_precision_runtime.install(self)
+
         self.norm_moe_gen = RMSNorm(self.hidden_size, eps=self.rms_norm_eps)
         self.gen_sp_prepare = Cosmos3GenSPPrepare()
         self.gen_sp_gather = nn.Identity()
@@ -1958,3 +2002,11 @@ class Cosmos3VFMTransformer(nn.Module):
     def post_load_weights(self) -> None:
         """Post-load processing: ensure correct dtypes."""
         self.time_embedder.to(torch.float32)
+
+    def set_mixed_precision_step(self, step_index: int, num_steps: int) -> None:
+        if self.mixed_precision_runtime is not None:
+            self.mixed_precision_runtime.set_step(step_index, num_steps)
+
+    def reset_mixed_precision(self) -> None:
+        if self.mixed_precision_runtime is not None:
+            self.mixed_precision_runtime.reset()