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

This reverts commit b63b36ded3.

docs/user_guide/quantization/fp8.md

@@ -52,17 +52,6 @@ guide. FP8 on Ampere may use a weight-only path where available.
 Audio encoder, vision encoder, talker, and code2wav stay in BF16 unless a
 model-specific guide says otherwise.
 
-### MoT Diffusion Model (SenseNova-U1)
-
-| Model | HF models | Scope | Status | Notes |
-|-------|-----------|-------|--------|-------|
-| SenseNova-U1 | `SenseNova/SenseNova-U1-8B-MoT` | Gen-path only (`*_mot_gen` layers) | Validated | Understanding-path layers stay in BF16; no `ignored_layers` needed |
-
-SenseNova-U1 uses Mixture-of-Tokenizers (MoT) with duplicate weights per
-layer: understanding path (`qkv_proj`, `mlp`) and generation path
-(`qkv_proj_mot_gen`, `mlp_mot_gen`).  When `quantization="fp8"` is set,
-the pipeline automatically applies FP8 only to gen-path GEMMs.
-
 ### Multi-Stage Diffusion Model (BAGEL, GLM-Image)
 
 | Model | Scope | Status | Notes |

tests/diffusion/quantization/test_quantization_fp8.py

@@ -6,13 +6,11 @@ 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
@@ -438,92 +436,3 @@ 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

@@ -1,175 +0,0 @@
-# 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

vllm_omni/diffusion/models/sensenova_u1/pipeline_sensenova_u1.py

@@ -31,7 +31,6 @@ import torchvision.transforms as T
 from PIL import Image
 from transformers import AutoTokenizer
 from vllm.logger import init_logger
-from vllm.model_executor.layers.quantization import QuantizationConfig
 
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.utils import get_local_device
@@ -484,39 +483,6 @@ def _optimized_scale(positive_flat, negative_flat):
     return dot / sq_norm
 
 
-# ---------------------------------------------------------------------------
-# Gen-only FP8 wrapper
-# ---------------------------------------------------------------------------
-
-
-class _GenOnlyQuantConfig:
-    """Wraps an inner QuantizationConfig so only ``*_mot_gen`` layers are quantized.
-
-    SenseNova-U1 uses Mixture-of-Tokenizers with duplicate weights per layer:
-    understanding path (``qkv_proj``, ``mlp``) and generation path
-    (``qkv_proj_mot_gen``, ``mlp_mot_gen``).  The RFC requires FP8 only on the
-    gen-path GEMMs during denoising; und-path layers stay in BF16.
-
-    This is *not* a full ``QuantizationConfig`` subclass — it only needs to
-    satisfy the duck-typed interface that ``LinearBase.__init__`` uses:
-    ``get_quant_method(layer, prefix)``.  The outer registry and weight-loader
-    still operate on the original config stored in ``od_config.quantization_config``.
-    """
-
-    def __init__(self, inner: QuantizationConfig):
-        self._inner = inner
-
-    def get_quant_method(self, layer, prefix: str = ""):
-        from vllm.model_executor.layers.linear import UnquantizedLinearMethod
-
-        if "mot_gen" in prefix:
-            return self._inner.get_quant_method(layer, prefix)
-        return UnquantizedLinearMethod()
-
-    def __getattr__(self, name):
-        return getattr(self._inner, name)
-
-
 # ---------------------------------------------------------------------------
 # Pipeline
 # ---------------------------------------------------------------------------
@@ -583,15 +549,9 @@ class SenseNovaU1Pipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipeli
         self.img_context_token_id = self.tokenizer.convert_tokens_to_ids(IMG_CONTEXT_TOKEN)
         self.img_start_token_id = self.tokenizer.convert_tokens_to_ids(IMG_START_TOKEN)
 
-        # Language model (TP-aware).
-        # Wrap quant_config so only gen-path (mot_gen) layers are quantized;
-        # understanding-path layers stay in BF16.
-        quant_config = od_config.quantization_config
-        if quant_config is not None:
-            quant_config = _GenOnlyQuantConfig(quant_config)
+        # Language model (TP-aware)
         self.language_model = SenseNovaU1ForCausalLM(
             self.llm_cfg,
-            quant_config=quant_config,
             prefix="language_model",
         )