[Feature] gen-only FP8 Quantization Support for SenseNova-U1 (#3943)

Signed-off-by: leohuang257 <masugchds@gmail.com>

tests/diffusion/quantization/test_quantization_fp8.py

@@ -6,11 +6,13 @@ End-to-end tests for the unified quantization framework (PR #1764).
 Validates FP8 quantization works correctly for all supported model types:
   - Single-stage diffusion models (FLUX.1-dev, Qwen-Image, Z-Image-Turbo)
   - Multi-stage models (BAGEL: LLM + Diffusion)
+  - MoT gen-only models (SenseNova-U1: gen-path FP8, und-path BF16)
 
 Tests verify:
   1. FP8 quantization produces valid images
   2. Memory usage is lower than BF16 baseline
   3. Multi-stage models only quantize the diffusion stage (not the LLM stage)
+  4. MoT models only quantize gen-path layers (mot_gen)
 
 Usage:
     # Run all FP8 quantization tests
@@ -436,3 +438,92 @@ def test_single_stage_quantization_config_key():
         quantization_config="fp8",
     )
     assert len(images) >= 1
+
+
+# ─── SenseNova-U1 gen-only FP8 tests ─────────────────────────────────────────
+
+_SENSENOVA_MODEL = "SenseNova/SenseNova-U1-8B-MoT"
+
+
+def _generate_sensenova_u1_image(
+    quantization_config: str | None = None,
+    num_inference_steps: int = 20,
+) -> tuple[Any, float]:
+    """Generate an image with SenseNova-U1 (single-stage MoT, gen-only FP8).
+
+    Returns (generated_image, peak_memory_gib).
+    """
+    omni_kwargs: dict[str, Any] = {}
+    if quantization_config:
+        omni_kwargs["quantization_config"] = quantization_config
+
+    with OmniRunner(_SENSENOVA_MODEL, **omni_kwargs) as runner:
+        torch.accelerator.reset_peak_memory_stats()
+
+        sampling_params = OmniDiffusionSamplingParams(
+            height=1024,
+            width=1024,
+            seed=42,
+            num_inference_steps=num_inference_steps,
+            extra_args={
+                "cfg_scale": 4.0,
+                "cfg_norm": "none",
+                "timestep_shift": 3.0,
+            },
+        )
+        omni_outputs = list(
+            runner.omni.generate(
+                prompts={"prompt": "A cat sitting on a windowsill", "modalities": ["image"]},
+                sampling_params_list=sampling_params,
+            )
+        )
+
+        peak_mem = torch.accelerator.max_memory_allocated() / (1024**3)
+
+        generated_image = None
+        for req_output in omni_outputs:
+            if images := getattr(req_output, "images", None):
+                generated_image = images[0]
+                break
+
+        assert generated_image is not None, "No images generated from SenseNova-U1"
+        assert generated_image.size == (1024, 1024), f"Expected 1024x1024, got {generated_image.size}"
+
+        return generated_image, peak_mem
+
+
+@hardware_test(res={"cuda": "H100"})
+def test_sensenova_u1_fp8_generates_image():
+    """SenseNova-U1 with gen-only FP8 generates a valid image.
+
+    FP8 is applied only to gen-path (mot_gen) layers; understanding-path
+    layers stay in BF16.
+    """
+    image, _ = _generate_sensenova_u1_image(quantization_config="fp8")
+    image.save("test_sensenova_u1_fp8.png")
+
+
+@hardware_test(res={"cuda": "H100"})
+def test_sensenova_u1_bf16_generates_image():
+    """SenseNova-U1 without quantization generates a valid image (baseline)."""
+    image, _ = _generate_sensenova_u1_image(quantization_config=None)
+    image.save("test_sensenova_u1_bf16.png")
+
+
+@hardware_test(res={"cuda": "H100"})
+def test_sensenova_u1_fp8_uses_less_memory():
+    """Gen-only FP8 should use less peak memory than BF16 for SenseNova-U1."""
+    _, mem_bf16 = _generate_sensenova_u1_image(
+        quantization_config=None,
+        num_inference_steps=10,
+    )
+    torch.accelerator.empty_cache()
+
+    _, mem_fp8 = _generate_sensenova_u1_image(
+        quantization_config="fp8",
+        num_inference_steps=10,
+    )
+
+    print(f"SenseNova-U1 BF16 peak memory: {mem_bf16:.2f} GiB")
+    print(f"SenseNova-U1 FP8  peak memory: {mem_fp8:.2f} GiB")
+    assert mem_fp8 < mem_bf16, f"FP8 ({mem_fp8:.2f} GiB) should use less memory than BF16 ({mem_bf16:.2f} GiB)"

tests/diffusion/quantization/test_sensenova_u1_fp8.py

@@ -0,0 +1,175 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""CPU tests for SenseNova-U1 gen-only FP8 quantization.
+
+Validates that the ``_GenOnlyQuantConfig`` wrapper correctly routes FP8
+quantization to only the gen-path (``*_mot_gen``) layers while leaving
+understanding-path layers in BF16.
+"""
+
+from types import SimpleNamespace
+from unittest.mock import MagicMock
+
+import pytest
+import torch
+from vllm.distributed import parallel_state
+from vllm.model_executor.layers.linear import LinearBase
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+@pytest.fixture(autouse=True)
+def _fake_tp_group():
+    """Provide a fake TP group so vLLM parallel layers can be instantiated.
+
+    Note: a shared ``init_fake_tp_group`` fixture exists in
+    ``tests/helpers/fixtures/env.py`` but requires pytest-mock.
+    """
+    mock_tp = MagicMock()
+    mock_tp.world_size = 1
+    mock_tp.rank_in_group = 0
+    old = parallel_state._TP
+    parallel_state._TP = mock_tp
+    yield
+    parallel_state._TP = old
+
+
+def _make_minimal_llm_config(num_layers: int = 2) -> SimpleNamespace:
+    """Create a minimal SenseNova-U1 LLM config for unit testing."""
+    return SimpleNamespace(
+        hidden_size=256,
+        intermediate_size=512,
+        num_hidden_layers=num_layers,
+        num_attention_heads=4,
+        num_key_value_heads=4,
+        head_dim=64,
+        hidden_act="silu",
+        rms_norm_eps=1e-6,
+        vocab_size=1024,
+        rope_theta=1000000.0,
+        rope_theta_hw=10000.0,
+        max_position_embeddings=4096,
+        max_position_embeddings_hw=1024,
+        attention_bias=True,
+        layer_types=["full_attention"] * num_layers,
+        tie_word_embeddings=False,
+    )
+
+
+def _build_model_with_gen_only_quant(quant_method: str | None = "fp8"):
+    """Construct SenseNovaU1ForCausalLM with gen-only FP8 wrapping."""
+    from unittest.mock import MagicMock, patch
+
+    from vllm_omni.diffusion.models.sensenova_u1.pipeline_sensenova_u1 import (
+        _GenOnlyQuantConfig,
+    )
+    from vllm_omni.diffusion.models.sensenova_u1.sensenova_u1_transformer import (
+        SenseNovaU1ForCausalLM,
+    )
+    from vllm_omni.quantization import build_quant_config
+
+    inner_config = build_quant_config(quant_method)
+    if inner_config is not None:
+        quant_config = _GenOnlyQuantConfig(inner_config)
+    else:
+        quant_config = None
+
+    llm_cfg = _make_minimal_llm_config(num_layers=2)
+
+    mock_vllm_config = MagicMock()
+    mock_vllm_config.model_config.dtype = torch.bfloat16
+    mock_vllm_config.model_config.quantization = quant_method
+
+    with (
+        torch.device("meta"),
+        patch("vllm.model_executor.layers.quantization.fp8.get_current_vllm_config", return_value=mock_vllm_config),
+    ):
+        model = SenseNovaU1ForCausalLM(
+            llm_cfg,
+            quant_config=quant_config,
+            prefix="language_model",
+        )
+    return model
+
+
+def test_mot_gen_layers_have_fp8_quant_method():
+    """Gen-path (mot_gen) Linear layers should have a real FP8 quant_method."""
+    from vllm.model_executor.layers.linear import UnquantizedLinearMethod
+
+    model = _build_model_with_gen_only_quant("fp8")
+
+    mot_gen_linears = [
+        (name, mod) for name, mod in model.named_modules() if isinstance(mod, LinearBase) and "mot_gen" in name
+    ]
+
+    assert len(mot_gen_linears) > 0, "No mot_gen LinearBase modules found"
+
+    for name, mod in mot_gen_linears:
+        assert mod.quant_method is not None, f"mot_gen layer {name!r} should have quant_method with FP8"
+        assert not isinstance(mod.quant_method, UnquantizedLinearMethod), (
+            f"mot_gen layer {name!r} should NOT use UnquantizedLinearMethod"
+        )
+
+
+def test_und_layers_have_no_quant_method():
+    """Understanding-path Linear layers should use UnquantizedLinearMethod (stay BF16)."""
+    from vllm.model_executor.layers.linear import UnquantizedLinearMethod
+
+    model = _build_model_with_gen_only_quant("fp8")
+
+    und_linears = [
+        (name, mod) for name, mod in model.named_modules() if isinstance(mod, LinearBase) and "mot_gen" not in name
+    ]
+
+    assert len(und_linears) > 0, "No und LinearBase modules found"
+
+    for name, mod in und_linears:
+        assert isinstance(mod.quant_method, UnquantizedLinearMethod), (
+            f"und layer {name!r} should use UnquantizedLinearMethod, got {type(mod.quant_method).__name__}"
+        )
+
+
+def test_no_quant_method_without_quantization():
+    """Without quantization, all layers should use UnquantizedLinearMethod."""
+    from vllm.model_executor.layers.linear import UnquantizedLinearMethod
+
+    model = _build_model_with_gen_only_quant(None)
+
+    for name, mod in model.named_modules():
+        if isinstance(mod, LinearBase):
+            assert isinstance(mod.quant_method, UnquantizedLinearMethod), (
+                f"LinearBase module {name!r} has quant_method="
+                f"{type(mod.quant_method).__name__} "
+                "but no quantization was configured"
+            )
+
+
+def test_und_and_gen_layers_both_present():
+    """Sanity check: the model has both und and gen layers."""
+    model = _build_model_with_gen_only_quant(None)
+
+    layer_names = [name for name, _ in model.named_modules()]
+
+    und_qkv = [n for n in layer_names if "qkv_proj" in n and "mot_gen" not in n]
+    gen_qkv = [n for n in layer_names if "qkv_proj_mot_gen" in n]
+    und_mlp = [n for n in layer_names if ".mlp." in n and "mot_gen" not in n]
+    gen_mlp = [n for n in layer_names if "mlp_mot_gen" in n]
+
+    assert len(und_qkv) > 0, "No und qkv_proj layers found"
+    assert len(gen_qkv) > 0, "No gen qkv_proj_mot_gen layers found"
+    assert len(und_mlp) > 0, "No und mlp layers found"
+    assert len(gen_mlp) > 0, "No gen mlp_mot_gen layers found"
+
+
+def test_gen_only_wrapper_delegates_attributes():
+    """_GenOnlyQuantConfig should delegate attribute access to the inner config."""
+    from vllm_omni.diffusion.models.sensenova_u1.pipeline_sensenova_u1 import (
+        _GenOnlyQuantConfig,
+    )
+    from vllm_omni.quantization import build_quant_config
+
+    inner = build_quant_config("fp8")
+    wrapper = _GenOnlyQuantConfig(inner)
+
+    assert wrapper.get_name() == "fp8"
+    assert wrapper.activation_scheme == inner.activation_scheme