From b63b36ded3638d035123fbb0a1558477207c8d3e Mon Sep 17 00:00:00 2001 From: "Guanzhe (Leo) Huang" <116847811+leohuang257@users.noreply.github.com> Date: Fri, 5 Jun 2026 10:51:11 +0800 Subject: [PATCH] [Feature] gen-only FP8 Quantization Support for SenseNova-U1 (#3943) Signed-off-by: leohuang257 --- docs/user_guide/quantization/fp8.md | 11 ++ .../quantization/test_quantization_fp8.py | 91 +++++++++ .../quantization/test_sensenova_u1_fp8.py | 175 ++++++++++++++++++ .../sensenova_u1/pipeline_sensenova_u1.py | 42 ++++- 4 files changed, 318 insertions(+), 1 deletion(-) create mode 100644 tests/diffusion/quantization/test_sensenova_u1_fp8.py diff --git a/docs/user_guide/quantization/fp8.md b/docs/user_guide/quantization/fp8.md index c8f14548c..69e161a60 100644 --- a/docs/user_guide/quantization/fp8.md +++ b/docs/user_guide/quantization/fp8.md @@ -52,6 +52,17 @@ 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 | diff --git a/tests/diffusion/quantization/test_quantization_fp8.py b/tests/diffusion/quantization/test_quantization_fp8.py index 58f91742d..8ff16e1f1 100644 --- a/tests/diffusion/quantization/test_quantization_fp8.py +++ b/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)" diff --git a/tests/diffusion/quantization/test_sensenova_u1_fp8.py b/tests/diffusion/quantization/test_sensenova_u1_fp8.py new file mode 100644 index 000000000..4ed30e9a6 --- /dev/null +++ b/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 diff --git a/vllm_omni/diffusion/models/sensenova_u1/pipeline_sensenova_u1.py b/vllm_omni/diffusion/models/sensenova_u1/pipeline_sensenova_u1.py index fe8695b07..917cea13e 100644 --- a/vllm_omni/diffusion/models/sensenova_u1/pipeline_sensenova_u1.py +++ b/vllm_omni/diffusion/models/sensenova_u1/pipeline_sensenova_u1.py @@ -31,6 +31,7 @@ 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 @@ -483,6 +484,39 @@ 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 # --------------------------------------------------------------------------- @@ -549,9 +583,15 @@ 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) + # 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) self.language_model = SenseNovaU1ForCausalLM( self.llm_cfg, + quant_config=quant_config, prefix="language_model", )