[Bugfix] Fix tp and Quantization incompatible for Flux (#2180)

Signed-off-by: Lancer <maruixiang6688@gmail.com>

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:

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:

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: