From 64629cc174d6e0ae8283c9407ffa5f64acec186f Mon Sep 17 00:00:00 2001 From: Lancer Date: Sat, 28 Mar 2026 18:28:10 +0800 Subject: [PATCH] [Feat] support quantization for Flux Kontext (#2184) Signed-off-by: Lancer --- vllm_omni/diffusion/models/flux/flux_transformer.py | 5 +++++ vllm_omni/diffusion/models/flux/pipeline_flux_kontext.py | 2 ++ 2 files changed, 7 insertions(+) diff --git a/vllm_omni/diffusion/models/flux/flux_transformer.py b/vllm_omni/diffusion/models/flux/flux_transformer.py index 7e4f1dc24..df3a26742 100644 --- a/vllm_omni/diffusion/models/flux/flux_transformer.py +++ b/vllm_omni/diffusion/models/flux/flux_transformer.py @@ -196,6 +196,8 @@ class FluxAttention(torch.nn.Module): attention_mask: torch.Tensor | None = None, **kwargs, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + # Ensure contiguous for FP8 quantized linear layers + hidden_states = hidden_states.contiguous() qkv, _ = self.to_qkv(hidden_states) q_size = self.to_qkv.num_heads * self.head_dim kv_size = self.to_qkv.num_kv_heads * self.head_dim @@ -209,6 +211,7 @@ class FluxAttention(torch.nn.Module): key = self.norm_k(key) if self.added_kv_proj_dim is not None: + encoder_hidden_states = encoder_hidden_states.contiguous() encoder_qkv, _ = self.add_kv_proj(encoder_hidden_states) add_q_size = self.add_kv_proj.num_heads * self.head_dim add_kv_size = self.add_kv_proj.num_kv_heads * self.head_dim @@ -720,6 +723,7 @@ class FluxKontextTransformer2DModel(FluxTransformer2DModel): guidance_embeds: bool = True, axes_dims_rope: tuple[int, int, int] = (16, 56, 56), theta: float = 10000.0, + quant_config: "QuantizationConfig | None" = None, ): super().__init__( od_config=od_config, @@ -735,6 +739,7 @@ class FluxKontextTransformer2DModel(FluxTransformer2DModel): guidance_embeds=guidance_embeds, axes_dims_rope=axes_dims_rope, theta=theta, + quant_config=quant_config, ) def forward( diff --git a/vllm_omni/diffusion/models/flux/pipeline_flux_kontext.py b/vllm_omni/diffusion/models/flux/pipeline_flux_kontext.py index bf741a64e..3232b436d 100644 --- a/vllm_omni/diffusion/models/flux/pipeline_flux_kontext.py +++ b/vllm_omni/diffusion/models/flux/pipeline_flux_kontext.py @@ -130,6 +130,8 @@ class FluxKontextPipeline(nn.Module, FluxPipelineMixin, SupportImageInput): ).to(self._execution_device) transformer_kwargs = get_transformer_config_kwargs(od_config.tf_model_config, FluxKontextTransformer2DModel) + transformer_kwargs["od_config"] = od_config + transformer_kwargs["quant_config"] = od_config.quantization_config self.transformer = FluxKontextTransformer2DModel(**transformer_kwargs) self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8