[Feat] add helios fp8 quantization (#1916)

Signed-off-by: rongfu.leng <lenronfu@gmail.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

docs/user_guide/diffusion/quantization/fp8.md

@@ -65,6 +65,7 @@ The available `ignored_layers` names depend on the model architecture (e.g., `to
 | Flux | `black-forest-labs/FLUX.1-dev` | All layers | None |
 | HunyuanImage-3 | `tencent/HunyuanImage3` | All layers | None |
 | HunyuanVideo-1.5 | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v`, `720p_t2v`, `480p_i2v` | All layers | None |
+| Helios | `BestWishYsh/Helios-Base`, `BestWishYsh/Helios-Mid`, `BestWishYsh/Helios-Distilled` | All layers | None |
 
 ## Combining with Other Features
 

examples/offline_inference/helios/end2end.py

@@ -178,6 +178,12 @@ def parse_args() -> argparse.Namespace:
         action="store_true",
         help="Enable diffusion pipeline profiler to display stage durations.",
     )
+    parser.add_argument(
+        "--quantization",
+        type=str,
+        default=None,
+        help="Quantization method (e.g. 'fp8').",
+    )
 
     # Memory & parallelism
     parser.add_argument("--vae-use-slicing", action="store_true", help="Enable VAE slicing.")
@@ -213,6 +219,7 @@ def main():
         parallel_config=parallel_config,
         enforce_eager=args.enforce_eager,
         enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
+        quantization=args.quantization,
     )
 
     # Validate I2V / V2V arguments

vllm_omni/diffusion/models/helios/helios_transformer.py

@@ -5,7 +5,7 @@
 import math
 from collections.abc import Iterable
 from functools import lru_cache
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 import torch
 import torch.nn as nn
@@ -29,6 +29,11 @@ from vllm_omni.diffusion.distributed.sp_plan import (
     SequenceParallelOutput,
 )
 
+if TYPE_CHECKING:
+    from vllm.model_executor.layers.quantization.base_config import (
+        QuantizationConfig,
+    )
+
 logger = init_logger(__name__)
 
 
@@ -96,7 +101,15 @@ class DistributedRMSNorm(nn.Module):
 class ColumnParallelGELU(nn.Module):
     """Column parallel linear with GELU activation."""
 
-    def __init__(self, dim_in: int, dim_out: int, *, approximate: str = "tanh", bias: bool = True):
+    def __init__(
+        self,
+        dim_in: int,
+        dim_out: int,
+        *,
+        approximate: str = "tanh",
+        bias: bool = True,
+        quant_config: "QuantizationConfig | None" = None,
+    ):
         super().__init__()
         self.proj = ColumnParallelLinear(
             dim_in,
@@ -104,6 +117,7 @@ class ColumnParallelGELU(nn.Module):
             bias=bias,
             gather_output=False,
             return_bias=False,
+            quant_config=quant_config,
         )
         self.approximate = approximate
 
@@ -121,18 +135,15 @@ class HeliosFeedForward(nn.Module):
         inner_dim: int,
         dim_out: int | None = None,
         bias: bool = True,
+        quant_config: "QuantizationConfig | None" = None,
     ) -> None:
         super().__init__()
         dim_out = dim_out or dim
 
-        self.net_0 = ColumnParallelGELU(dim, inner_dim, approximate="tanh", bias=bias)
+        self.net_0 = ColumnParallelGELU(dim, inner_dim, approximate="tanh", bias=bias, quant_config=quant_config)
         self.net_1 = nn.Identity()
         self.net_2 = RowParallelLinear(
-            inner_dim,
-            dim_out,
-            bias=bias,
-            input_is_parallel=True,
-            return_bias=False,
+            inner_dim, dim_out, bias=bias, input_is_parallel=True, return_bias=False, quant_config=quant_config
         )
 
     def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
@@ -262,6 +273,7 @@ class HeliosSelfAttention(nn.Module):
         dropout: float = 0.0,
         is_amplify_history: bool = False,
         history_scale_mode: str = "per_head",
+        quant_config: "QuantizationConfig | None" = None,
     ):
         super().__init__()
 
@@ -275,6 +287,7 @@ class HeliosSelfAttention(nn.Module):
             head_size=head_dim,
             total_num_heads=num_heads,
             bias=True,
+            quant_config=quant_config,
         )
 
         self.num_heads = self.to_qkv.num_heads
@@ -290,6 +303,7 @@ class HeliosSelfAttention(nn.Module):
             bias=True,
             input_is_parallel=True,
             return_bias=False,
+            quant_config=quant_config,
         )
         self.dropout = nn.Dropout(dropout)
 
@@ -368,6 +382,7 @@ class HeliosCrossAttention(nn.Module):
         head_dim: int,
         eps: float = 1e-5,
         dropout: float = 0.0,
+        quant_config: "QuantizationConfig | None" = None,
     ):
         super().__init__()
 
@@ -382,6 +397,7 @@ class HeliosCrossAttention(nn.Module):
             bias=True,
             gather_output=False,
             return_bias=False,
+            quant_config=quant_config,
         )
         self.to_k = ColumnParallelLinear(
             dim,
@@ -389,6 +405,7 @@ class HeliosCrossAttention(nn.Module):
             bias=True,
             gather_output=False,
             return_bias=False,
+            quant_config=quant_config,
         )
         self.to_v = ColumnParallelLinear(
             dim,
@@ -396,6 +413,7 @@ class HeliosCrossAttention(nn.Module):
             bias=True,
             gather_output=False,
             return_bias=False,
+            quant_config=quant_config,
         )
 
         tp_size = get_tensor_model_parallel_world_size()
@@ -411,6 +429,7 @@ class HeliosCrossAttention(nn.Module):
             bias=True,
             input_is_parallel=True,
             return_bias=False,
+            quant_config=quant_config,
         )
         self.dropout = nn.Dropout(dropout)
 
@@ -461,6 +480,7 @@ class HeliosTransformerBlock(nn.Module):
         guidance_cross_attn: bool = False,
         is_amplify_history: bool = False,
         history_scale_mode: str = "per_head",
+        quant_config: "QuantizationConfig | None" = None,
     ):
         super().__init__()
 
@@ -475,19 +495,17 @@ class HeliosTransformerBlock(nn.Module):
             eps=eps,
             is_amplify_history=is_amplify_history,
             history_scale_mode=history_scale_mode,
+            quant_config=quant_config,
         )
 
         # 2. Cross-attention
         self.attn2 = HeliosCrossAttention(
-            dim=dim,
-            num_heads=num_heads,
-            head_dim=head_dim,
-            eps=eps,
+            dim=dim, num_heads=num_heads, head_dim=head_dim, eps=eps, quant_config=quant_config
         )
         self.norm2 = FP32LayerNorm(dim, eps, elementwise_affine=True) if cross_attn_norm else nn.Identity()
 
         # 3. Feed-forward
-        self.ffn = HeliosFeedForward(dim=dim, inner_dim=ffn_dim, dim_out=dim)
+        self.ffn = HeliosFeedForward(dim=dim, inner_dim=ffn_dim, dim_out=dim, quant_config=quant_config)
         self.norm3 = FP32LayerNorm(dim, eps, elementwise_affine=False)
 
         self.scale_shift_table = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5)
@@ -602,6 +620,7 @@ class HeliosTransformer3DModel(nn.Module):
         has_multi_term_memory_patch: bool = True,
         is_amplify_history: bool = False,
         history_scale_mode: str = "per_head",
+        quant_config: "QuantizationConfig | None" = None,
     ) -> None:
         super().__init__()
 
@@ -689,6 +708,7 @@ class HeliosTransformer3DModel(nn.Module):
                     guidance_cross_attn=guidance_cross_attn,
                     is_amplify_history=is_amplify_history,
                     history_scale_mode=history_scale_mode,
+                    quant_config=quant_config,
                 )
                 for _ in range(num_layers)
             ]

vllm_omni/diffusion/models/helios/pipeline_helios.py

@@ -9,7 +9,7 @@ import logging
 import math
 import os
 from collections.abc import Iterable
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 import numpy as np
 import torch
@@ -31,6 +31,11 @@ from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPi
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.platforms import current_omni_platform
 
+if TYPE_CHECKING:
+    from vllm.model_executor.layers.quantization.base_config import (
+        QuantizationConfig,
+    )
+
 logger = logging.getLogger(__name__)
 
 
@@ -80,7 +85,9 @@ def load_transformer_config(model_path: str, subfolder: str = "transformer", loc
     return load_json_config(model_path, subfolder, "config.json", local_files_only)
 
 
-def create_transformer_from_config(config: dict) -> HeliosTransformer3DModel:
+def create_transformer_from_config(
+    config: dict, quant_config: QuantizationConfig | None = None
+) -> HeliosTransformer3DModel:
     kwargs = {}
 
     key_map = [
@@ -112,7 +119,7 @@ def create_transformer_from_config(config: dict) -> HeliosTransformer3DModel:
                 val = tuple(val)
             kwargs[key] = val
 
-    return HeliosTransformer3DModel(**kwargs)
+    return HeliosTransformer3DModel(quant_config=quant_config, **kwargs)
 
 
 def get_helios_post_process_func(
@@ -192,7 +199,9 @@ class HeliosPipeline(nn.Module, CFGParallelMixin, ProgressBarMixin, DiffusionPip
         ).to(self.device)
 
         transformer_config = load_transformer_config(model, "transformer", local_files_only)
-        self.transformer = create_transformer_from_config(transformer_config)
+        self.transformer = create_transformer_from_config(
+            transformer_config, quant_config=od_config.quantization_config
+        )
 
         # Read scheduler config to determine scheduler type
         sched_cfg = load_json_config(model, "scheduler", "scheduler_config.json", local_files_only)