From 055963b3fd22a0365d837a9631833fa43551346e Mon Sep 17 00:00:00 2001 From: Lancer Date: Fri, 27 Mar 2026 09:21:50 +0800 Subject: [PATCH] [Bugfix] Fix tp and Quantization incompatible for Flux (#2180) Signed-off-by: Lancer --- vllm_omni/diffusion/models/flux/flux_transformer.py | 5 +++-- vllm_omni/diffusion/models/flux2/flux2_transformer.py | 5 +++-- .../models/flux2_klein/flux2_klein_transformer.py | 8 +++++--- 3 files changed, 11 insertions(+), 7 deletions(-) diff --git a/vllm_omni/diffusion/models/flux/flux_transformer.py b/vllm_omni/diffusion/models/flux/flux_transformer.py index 8f81f3cd2..7e4f1dc24 100644 --- a/vllm_omni/diffusion/models/flux/flux_transformer.py +++ b/vllm_omni/diffusion/models/flux/flux_transformer.py @@ -248,9 +248,10 @@ class FluxAttention(torch.nn.Module): encoder_hidden_states, hidden_states = hidden_states.split_with_sizes( [encoder_hidden_states.shape[1], hidden_states.shape[1] - encoder_hidden_states.shape[1]], dim=1 ) - hidden_states = self.to_out[0](hidden_states) + # Contiguous for FP8 quantization in RowParallelLinear + hidden_states = self.to_out[0](hidden_states.contiguous()) hidden_states = self.to_out[1](hidden_states) - encoder_hidden_states = self.to_add_out(encoder_hidden_states) + encoder_hidden_states = self.to_add_out(encoder_hidden_states.contiguous()) return hidden_states, encoder_hidden_states else: diff --git a/vllm_omni/diffusion/models/flux2/flux2_transformer.py b/vllm_omni/diffusion/models/flux2/flux2_transformer.py index 62f415393..116e499b0 100644 --- a/vllm_omni/diffusion/models/flux2/flux2_transformer.py +++ b/vllm_omni/diffusion/models/flux2/flux2_transformer.py @@ -209,9 +209,10 @@ class Flux2Attention(nn.Module): [context_len, hidden_states.shape[1] - context_len], dim=1, ) - encoder_hidden_states = self.to_add_out(encoder_hidden_states) + # Contiguous for FP8 quantization in RowParallelLinear + encoder_hidden_states = self.to_add_out(encoder_hidden_states.contiguous()) - hidden_states = self.to_out[0](hidden_states) + hidden_states = self.to_out[0](hidden_states.contiguous()) hidden_states = self.to_out[1](hidden_states) if encoder_hidden_states is not None: diff --git a/vllm_omni/diffusion/models/flux2_klein/flux2_klein_transformer.py b/vllm_omni/diffusion/models/flux2_klein/flux2_klein_transformer.py index fa7bbc4fe..1d375ca8d 100644 --- a/vllm_omni/diffusion/models/flux2_klein/flux2_klein_transformer.py +++ b/vllm_omni/diffusion/models/flux2_klein/flux2_klein_transformer.py @@ -255,7 +255,8 @@ class Flux2Attention(nn.Module): [txt_len, hidden_states.shape[1] - txt_len], dim=1, ) - encoder_hidden_states = self.to_add_out(encoder_hidden_states) + # Contiguous for FP8 quantization in RowParallelLinear + encoder_hidden_states = self.to_add_out(encoder_hidden_states.contiguous()) else: query = torch.cat([encoder_query, query], dim=1) key = torch.cat([encoder_key, key], dim=1) @@ -282,7 +283,8 @@ class Flux2Attention(nn.Module): [context_len, hidden_states.shape[1] - context_len], dim=1, ) - encoder_hidden_states = self.to_add_out(encoder_hidden_states) + # Contiguous for FP8 quantization in RowParallelLinear + encoder_hidden_states = self.to_add_out(encoder_hidden_states.contiguous()) else: if image_rotary_emb is not None: cos, sin = image_rotary_emb @@ -300,7 +302,7 @@ class Flux2Attention(nn.Module): hidden_states = self.attn(query, key, value, attn_metadata) hidden_states = hidden_states.flatten(2, 3).to(query.dtype) - hidden_states = self.to_out[0](hidden_states) + hidden_states = self.to_out[0](hidden_states.contiguous()) hidden_states = self.to_out[1](hidden_states) if encoder_hidden_states is not None: