[Bugfix][Cosmos3] Resolve per-component transformer quantization config (#8101)

Signed-off-by: NumberWan <wantszkin2003@gmail.com>
Co-authored-by: Oliver7th <fangzhengyi@huawei.com>

tests/diffusion/models/cosmos3/test_cosmos3_transformer.py

@@ -75,6 +75,248 @@ def _tiny_cosmos3_edge_config(**overrides):
     return config
 
 
+@pytest.mark.parametrize(
+    ("config_kind", "expected_language", "expected_gen"),
+    [
+        ("flat", "transformer", "transformer"),
+        ("transformer", "transformer", "transformer"),
+        ("other_component", None, None),
+        ("language_model", "language", None),
+        ("gen_layers", None, "gen"),
+        ("transformer_language_override", "language", "transformer"),
+        ("transformer_gen_disabled", "transformer", None),
+        ("independent_subcomponents", "language", "gen"),
+    ],
+)
+def test_transformer_resolves_global_and_subcomponent_quant_configs(
+    monkeypatch: pytest.MonkeyPatch,
+    config_kind: str,
+    expected_language: str | None,
+    expected_gen: str | None,
+) -> None:
+    """Resolve pipeline-level defaults and Cosmos3 subcomponent overrides."""
+    from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3
+    from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3VFMTransformer
+    from vllm_omni.quantization.component_config import ComponentQuantizationConfig
+
+    received: list[tuple[str, object]] = []
+    configs = {
+        "transformer": object(),
+        "language": object(),
+        "gen": object(),
+    }
+
+    class _StubLanguageModel(nn.Module):
+        def __init__(self, *, quant_config, **kwargs) -> None:
+            del kwargs
+            super().__init__()
+            received.append(("language_model", quant_config))
+            self.layers = nn.ModuleList()
+
+    class _StubGenDecoderLayer(nn.Module):
+        def __init__(self, *, quant_config, **kwargs) -> None:
+            del kwargs
+            super().__init__()
+            received.append(("gen_layers", quant_config))
+
+    monkeypatch.setattr(Cosmos3VFMTransformer, "_language_model_cls", _StubLanguageModel)
+    monkeypatch.setattr(transformer_cosmos3, "Cosmos3GenDecoderLayer", _StubGenDecoderLayer)
+
+    if config_kind == "flat":
+        top_level_config = configs["transformer"]
+    elif config_kind == "transformer":
+        top_level_config = ComponentQuantizationConfig({"transformer": configs["transformer"]})
+    elif config_kind == "other_component":
+        top_level_config = ComponentQuantizationConfig({"vae": configs["transformer"]})
+    elif config_kind == "language_model":
+        top_level_config = ComponentQuantizationConfig({"language_model": configs["language"]})
+    elif config_kind == "gen_layers":
+        top_level_config = ComponentQuantizationConfig({"gen_layers": configs["gen"]})
+    elif config_kind == "transformer_language_override":
+        top_level_config = ComponentQuantizationConfig(
+            {"transformer": configs["transformer"], "language_model": configs["language"]}
+        )
+    elif config_kind == "transformer_gen_disabled":
+        top_level_config = ComponentQuantizationConfig({"transformer": configs["transformer"], "gen_layers": None})
+    else:
+        top_level_config = ComponentQuantizationConfig(
+            {"language_model": configs["language"], "gen_layers": configs["gen"]}
+        )
+
+    Cosmos3VFMTransformer(
+        SimpleNamespace(
+            tf_model_config=_tiny_cosmos3_config(num_hidden_layers=1),
+            dtype=torch.float32,
+            quantization_config=top_level_config,
+            custom_pipeline_args={},
+            model_config={},
+        )
+    )
+
+    assert received == [
+        ("language_model", None if expected_language is None else configs[expected_language]),
+        ("gen_layers", None if expected_gen is None else configs[expected_gen]),
+    ]
+
+
+def test_transformer_component_quant_reaches_real_linear(monkeypatch: pytest.MonkeyPatch) -> None:
+    """#7323 regression: unwrap ComponentQuantizationConfig before ColumnParallelLinear.
+
+    Passing the router through unchanged makes ``resolve(language_model.*)`` miss
+    a ``transformer``-only map and raise ``ValueError: All linear layers should
+    support quant method``. A leaf config must reach the real linear path after
+    resolution. Attention backends are stubbed so this stays CPU-safe; MLP and
+    QKV linears remain real ``ColumnParallelLinear`` instances.
+    """
+    from vllm.model_executor.layers.linear import UnquantizedLinearMethod
+    from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
+
+    from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3
+    from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import (
+        Cosmos3GatedMLP,
+        Cosmos3VFMTransformer,
+        _resolve_cosmos3_quant_configs,
+    )
+    from vllm_omni.quantization.component_config import ComponentQuantizationConfig
+
+    class _LeafQuantConfig(QuantizationConfig):
+        def get_name(self) -> str:
+            return "test_leaf"
+
+        def get_quant_method(self, layer, prefix):  # noqa: ANN001
+            del layer, prefix
+            return UnquantizedLinearMethod()
+
+        @classmethod
+        def get_supported_act_dtypes(cls):
+            return [torch.float32]
+
+        def get_min_capability(self) -> int:
+            return 0
+
+        @classmethod
+        def from_config(cls, config):  # noqa: ANN001
+            raise NotImplementedError
+
+        def get_config_filenames(self) -> list[str]:
+            return []
+
+    class _StubFrameworkAttention(nn.Module):
+        def __init__(self, *args, **kwargs) -> None:
+            del args, kwargs
+            super().__init__()
+
+    monkeypatch.setattr(transformer_cosmos3, "FrameworkAttention", _StubFrameworkAttention)
+
+    leaf = _LeafQuantConfig()
+    router = ComponentQuantizationConfig({"transformer": leaf})
+
+    with pytest.raises(ValueError, match="All linear layers should support quant method"):
+        Cosmos3GatedMLP(
+            hidden_size=8,
+            intermediate_size=16,
+            quant_config=router,
+            prefix="language_model.layers.0.mlp",
+        )
+
+    language_qc, gen_qc = _resolve_cosmos3_quant_configs(router)
+    assert language_qc is leaf
+    assert gen_qc is leaf
+
+    model = Cosmos3VFMTransformer(
+        SimpleNamespace(
+            tf_model_config=_tiny_cosmos3_config(num_hidden_layers=1),
+            dtype=torch.float32,
+            quantization_config=router,
+            custom_pipeline_args={},
+            model_config={},
+        )
+    )
+    mlp = model.language_model.layers[0].mlp
+    assert isinstance(mlp.gate_proj.quant_method, UnquantizedLinearMethod)
+    assert isinstance(model.gen_layers[0].mlp.gate_proj.quant_method, UnquantizedLinearMethod)
+    assert isinstance(
+        model.language_model.layers[0].self_attn.to_q.quant_method,
+        UnquantizedLinearMethod,
+    )
+
+
+def test_transformer_pathway_default_keeps_nested_layer_override() -> None:
+    """Pathway leaf plus a more-specific key must still longest-prefix match."""
+    from vllm.model_executor.layers.linear import UnquantizedLinearMethod
+    from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
+
+    from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import (
+        Cosmos3GatedMLP,
+        _resolve_cosmos3_quant_configs,
+    )
+    from vllm_omni.quantization.component_config import ComponentQuantizationConfig
+
+    class _TaggedLinearMethod(UnquantizedLinearMethod):
+        def __init__(self, tag: str) -> None:
+            super().__init__()
+            self.tag = tag
+
+    class _TaggedQuantConfig(QuantizationConfig):
+        def __init__(self, tag: str) -> None:
+            super().__init__()
+            self.tag = tag
+
+        def get_name(self) -> str:
+            return self.tag
+
+        def get_quant_method(self, layer, prefix):  # noqa: ANN001
+            del layer, prefix
+            return _TaggedLinearMethod(self.tag)
+
+        @classmethod
+        def get_supported_act_dtypes(cls):
+            return [torch.float32]
+
+        def get_min_capability(self) -> int:
+            return 0
+
+        @classmethod
+        def from_config(cls, config):  # noqa: ANN001
+            raise NotImplementedError
+
+        def get_config_filenames(self) -> list[str]:
+            return []
+
+    pathway = _TaggedQuantConfig("pathway")
+    mlp_override = _TaggedQuantConfig("mlp")
+    language_qc, gen_qc = _resolve_cosmos3_quant_configs(
+        ComponentQuantizationConfig(
+            {
+                "language_model": pathway,
+                "gen_layers": pathway,
+                "language_model.layers.0.mlp": mlp_override,
+            }
+        )
+    )
+
+    assert isinstance(language_qc, ComponentQuantizationConfig)
+    assert language_qc.resolve("language_model.layers.0.mlp.gate_proj") is mlp_override
+    assert language_qc.resolve("language_model.layers.0.self_attn.to_q") is pathway
+    assert language_qc.resolve("language_model.layers.1.mlp.gate_proj") is pathway
+    assert gen_qc is pathway
+
+    mlp0 = Cosmos3GatedMLP(
+        hidden_size=8,
+        intermediate_size=16,
+        quant_config=language_qc,
+        prefix="language_model.layers.0.mlp",
+    )
+    mlp1 = Cosmos3GatedMLP(
+        hidden_size=8,
+        intermediate_size=16,
+        quant_config=language_qc,
+        prefix="language_model.layers.1.mlp",
+    )
+    assert mlp0.gate_proj.quant_method.tag == "mlp"
+    assert mlp1.gate_proj.quant_method.tag == "pathway"
+
+
 def test_mrope_position_ids_cover_text_video_sound_and_action() -> None:
     from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import (
         compute_mrope_position_ids_action,

vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py

@@ -39,6 +39,10 @@ from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelInput, Seque
 from vllm_omni.diffusion.forward_context import get_forward_context, is_forward_context_available
 from vllm_omni.diffusion.layers.norm import RMSNorm as _VllmRMSNorm
 from vllm_omni.platforms import current_omni_platform
+from vllm_omni.quantization.component_config import (
+    ComponentQuantizationConfig,
+    resolve_component_quant_config,
+)
 
 from .mixed_precision import (
     Cosmos3MixedPrecisionConfig,
@@ -52,6 +56,53 @@ if TYPE_CHECKING:
 logger = init_logger(__name__)
 
 
+def _pathway_quant_config(
+    components: dict[str, QuantizationConfig | None],
+    pathway: str,
+    fallback: QuantizationConfig | None,
+) -> QuantizationConfig | None:
+    """Leaf overlay, or a scoped router when nested keys exist under ``pathway``.
+
+    Linear layers announce prefixes such as ``language_model.layers.0.mlp``.
+    Passing a single pathway leaf would drop more-specific keys the factory
+    still accepts (for example ``language_model.layers.0.mlp``). Nested keys
+    keep a ``ComponentQuantizationConfig`` so longest-prefix matching still
+    selects them; unmatched prefixes use the pathway root or ``fallback``.
+    """
+    nested = {key: value for key, value in components.items() if key.startswith(f"{pathway}.")}
+    has_root = pathway in components
+    root = components[pathway] if has_root else fallback
+    if not nested:
+        return root
+    scoped = dict(nested)
+    if has_root:
+        scoped[pathway] = root
+    return ComponentQuantizationConfig(scoped, default_config=root)
+
+
+def _resolve_cosmos3_quant_configs(
+    quant_config: QuantizationConfig | None,
+) -> tuple[QuantizationConfig | None, QuantizationConfig | None]:
+    """Resolve the Cosmos3 reasoner and generator quantization configs.
+
+    A pipeline-level ``transformer`` entry is the default for both internal
+    pathways (same leaf-unwrapping pattern as Flux2 / MiniMax / Boogu). The
+    historical ``language_model`` and ``gen_layers`` scopes remain supported as
+    exact-key overlays, including explicit ``None`` entries that leave one
+    pathway unquantized. Nested keys under those roots keep longest-prefix
+    routing via a pathway-scoped ``ComponentQuantizationConfig``.
+    """
+    if not isinstance(quant_config, ComponentQuantizationConfig):
+        return quant_config, quant_config
+
+    transformer_config = resolve_component_quant_config(quant_config, "transformer")
+    components = quant_config.component_configs
+    return (
+        _pathway_quant_config(components, "language_model", transformer_config),
+        _pathway_quant_config(components, "gen_layers", transformer_config),
+    )
+
+
 class RMSNorm(_VllmRMSNorm):
     """Cosmos3-local RMSNorm that uses the FP32 native implementation."""
 
@@ -1288,7 +1339,9 @@ class Cosmos3VFMTransformer(nn.Module):
         self.use_und_k_norm_for_gen = _tf_config_get(model_config, "use_und_k_norm_for_gen", None)
 
         dtype = od_config.dtype
-        quant_config = getattr(od_config, "quantization_config", None) if od_config else None
+        language_model_quant_config, gen_layers_quant_config = _resolve_cosmos3_quant_configs(
+            getattr(od_config, "quantization_config", 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":
@@ -1314,7 +1367,7 @@ class Cosmos3VFMTransformer(nn.Module):
             rms_norm_eps=self.rms_norm_eps,
             rope_theta=self.rope_theta,
             mrope_section=self.mrope_section,
-            quant_config=quant_config,
+            quant_config=language_model_quant_config,
             prefix="language_model",
             **self._language_model_kwargs(),
         )
@@ -1352,7 +1405,7 @@ class Cosmos3VFMTransformer(nn.Module):
                     num_key_value_heads=self.num_key_value_heads,
                     head_dim=self.head_dim,
                     rms_norm_eps=self.rms_norm_eps,
-                    quant_config=quant_config,
+                    quant_config=gen_layers_quant_config,
                     mlp_cls=self._gen_mlp_cls,
                     qk_norm=self.qk_norm_for_diffusion,
                     prefix=f"gen_layers.{i}",