[Model] Avoid repeated attention and MLP work in Klein TP (#8553)

Signed-off-by: QianCyrus <101633534+QianCyrus@users.noreply.github.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

vllm_omni/diffusion/models/flux2_klein/flux2_klein_transformer.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 # Copyright 2025 Black Forest Labs, The HuggingFace Team and The InstantX Team. All rights reserved.
 #
@@ -29,6 +29,7 @@ from diffusers.models.embeddings import (
 )
 from diffusers.models.modeling_outputs import Transformer2DModelOutput
 from diffusers.models.normalization import AdaLayerNormContinuous
+from vllm.distributed import get_tensor_model_parallel_world_size, tensor_model_parallel_all_gather
 from vllm.logger import init_logger
 from vllm.model_executor.layers.layernorm import RMSNorm
 from vllm.model_executor.layers.linear import (
@@ -50,6 +51,7 @@ from vllm_omni.diffusion.distributed.sp_plan import (
 from vllm_omni.diffusion.forward_context import get_forward_context
 from vllm_omni.diffusion.layers.rope import RotaryEmbedding, apply_rope_to_qk
 from vllm_omni.diffusion.models.host_weight_contract import FinalLayoutModelContract
+from vllm_omni.diffusion.offloader.config import offload_enabled
 
 logger = init_logger(__name__)
 if TYPE_CHECKING:
@@ -70,6 +72,19 @@ def _join_prefix(prefix: str, name: str) -> str:
     return f"{prefix}.{name}" if prefix else name
 
 
+def _use_local_single_stream_tp(od_config: OmniDiffusionConfig) -> bool:
+    parallel = od_config.parallel_config
+    return (
+        od_config.dtype == torch.bfloat16
+        and parallel.tensor_parallel_size > 1
+        and parallel.sequence_parallel_size in (None, 1)
+        and not parallel.use_hsdp
+        and od_config.quantization_config is None
+        and od_config.lora_path is None
+        and not offload_enabled(od_config)
+    )
+
+
 class Flux2SwiGLU(nn.Module):
     """SwiGLU activation used by Flux2."""
 
@@ -356,6 +371,7 @@ class Flux2ParallelSelfAttention(nn.Module):
         mlp_mult_factor: int = 2,
         quant_config: "QuantizationConfig | None" = None,
         prefix: str = "",
+        local_tp: bool = False,
     ):
         super().__init__()
         self.parallel_config = parallel_config
@@ -370,14 +386,40 @@ class Flux2ParallelSelfAttention(nn.Module):
         self.mlp_hidden_dim = int(query_dim * self.mlp_ratio)
         self.mlp_mult_factor = mlp_mult_factor
 
-        self.to_qkv_mlp_proj = ColumnParallelLinear(
-            self.query_dim,
-            self.inner_dim * 3 + self.mlp_hidden_dim * self.mlp_mult_factor,
-            bias=bias,
-            gather_output=True,
-            quant_config=quant_config,
-            prefix=_join_prefix(prefix, "to_qkv_mlp_proj"),
+        self.local_tp = (
+            local_tp
+            and quant_config is None
+            and parallel_config.sequence_parallel_size in (None, 1)
+            and mlp_mult_factor == 2
         )
+        tp_size = get_tensor_model_parallel_world_size() if self.local_tp else 1
+        self.local_tp = (
+            self.local_tp and tp_size > 1 and self.heads % tp_size == 0 and self.mlp_hidden_dim % tp_size == 0
+        )
+        tp_size = tp_size if self.local_tp else 1
+        self.query_num_heads = self.heads // tp_size
+        self.local_inner_dim = self.inner_dim // tp_size
+        self.local_mlp_hidden_dim = self.mlp_hidden_dim // tp_size
+
+        if self.local_tp:
+            # Split each packed segment independently so Q/K/V heads and
+            # SwiGLU gate/up channels stay paired on each TP rank.
+            self.to_qkv_mlp_proj = MergedColumnParallelLinear(
+                self.query_dim,
+                [self.inner_dim] * 3 + [self.mlp_hidden_dim] * 2,
+                bias=bias,
+                gather_output=False,
+                prefix=_join_prefix(prefix, "to_qkv_mlp_proj"),
+            )
+        else:
+            self.to_qkv_mlp_proj = ColumnParallelLinear(
+                self.query_dim,
+                self.inner_dim * 3 + self.mlp_hidden_dim * self.mlp_mult_factor,
+                bias=bias,
+                gather_output=True,
+                quant_config=quant_config,
+                prefix=_join_prefix(prefix, "to_qkv_mlp_proj"),
+            )
         self.mlp_act_fn = Flux2SwiGLU()
 
         self.norm_q = RMSNorm(dim_head, eps=eps)
@@ -393,7 +435,7 @@ class Flux2ParallelSelfAttention(nn.Module):
         )
         self.rope = RotaryEmbedding(is_neox_style=False)
         self.attn = Attention(
-            num_heads=self.heads,
+            num_heads=self.query_num_heads,
             head_size=self.head_dim,
             softmax_scale=1.0 / (self.head_dim**0.5),
             causal=False,
@@ -409,14 +451,14 @@ class Flux2ParallelSelfAttention(nn.Module):
         hidden_states, _ = self.to_qkv_mlp_proj(hidden_states)
         qkv, mlp_hidden_states = torch.split(
             hidden_states,
-            [3 * self.inner_dim, self.mlp_hidden_dim * self.mlp_mult_factor],
+            [3 * self.local_inner_dim, self.local_mlp_hidden_dim * self.mlp_mult_factor],
             dim=-1,
         )
 
         query, key, value = qkv.chunk(3, dim=-1)
-        query = query.unflatten(-1, (self.heads, -1))
-        key = key.unflatten(-1, (self.heads, -1))
-        value = value.unflatten(-1, (self.heads, -1))
+        query = query.unflatten(-1, (self.query_num_heads, -1))
+        key = key.unflatten(-1, (self.query_num_heads, -1))
+        value = value.unflatten(-1, (self.query_num_heads, -1))
 
         query = self.norm_q(query)
         key = self.norm_k(key)
@@ -468,6 +510,10 @@ class Flux2ParallelSelfAttention(nn.Module):
             if attention_mask is not None:
                 if attention_mask.dim() == 3:
                     attention_mask = attention_mask.unsqueeze(1)
+                if self.local_tp and attention_mask.dim() == 4 and attention_mask.shape[1] != 1:
+                    attention_mask = attention_mask.narrow(
+                        1, self.to_qkv_mlp_proj.tp_rank * self.query_num_heads, self.query_num_heads
+                    )
                 attn_metadata = AttentionMetadata(attn_mask=attention_mask)
 
             attn_output = self.attn(query, key, value, attn_metadata)
@@ -475,6 +521,9 @@ class Flux2ParallelSelfAttention(nn.Module):
         attn_output = attn_output.flatten(2, 3).to(query.dtype)
 
         mlp_hidden_states = self.mlp_act_fn(mlp_hidden_states)
+        if self.local_tp:
+            attn_output = tensor_model_parallel_all_gather(attn_output.contiguous())
+            mlp_hidden_states = tensor_model_parallel_all_gather(mlp_hidden_states.contiguous())
         hidden_states = torch.cat([attn_output, mlp_hidden_states], dim=-1)
         hidden_states, _ = self.to_out(hidden_states)
         return hidden_states
@@ -492,6 +541,7 @@ class Flux2SingleTransformerBlock(nn.Module):
         bias: bool = False,
         quant_config: "QuantizationConfig | None" = None,
         prefix: str = "",
+        local_tp: bool = False,
     ):
         super().__init__()
         self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
@@ -508,6 +558,7 @@ class Flux2SingleTransformerBlock(nn.Module):
             mlp_mult_factor=2,
             quant_config=quant_config,
             prefix=_join_prefix(prefix, "attn"),
+            local_tp=local_tp,
         )
 
     def forward(
@@ -834,6 +885,7 @@ class Flux2Transformer2DModel(nn.Module):
         guidance_embeds: bool = True,
         od_config: OmniDiffusionConfig = None,
         quant_config: "QuantizationConfig | None" = None,
+        local_tp: bool = False,
     ):
         super().__init__()
         # Warm HWR restoration skips load_weights(), so loader consumers such
@@ -911,10 +963,18 @@ class Flux2Transformer2DModel(nn.Module):
                     bias=False,
                     quant_config=quant_config,
                     prefix=f"single_transformer_blocks.{i}",
+                    local_tp=local_tp,
                 )
                 for i in range(num_single_layers)
             ]
         )
+        if any(block.attn.local_tp for block in self.single_transformer_blocks):
+            # Rank-local shapes are unchanged, but packed weight rows belong
+            # to different ranks. Do not restore the old final-layout cache.
+            self.host_weight_restore_contract = FinalLayoutModelContract(
+                implementation_id="flux2-klein-dit",
+                version="2-local-single-stream-tp",
+            )
 
         self.norm_out = AdaLayerNormContinuous(
             self.inner_dim,

vllm_omni/diffusion/models/flux2_klein/pipeline_flux2_klein.py

@@ -43,6 +43,7 @@ from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineL
 from vllm_omni.diffusion.model_loader.hub_prefetch import from_pretrained_with_prefetch, prefetch_subfolders
 from vllm_omni.diffusion.models.flux2_klein.flux2_klein_transformer import (
     Flux2Transformer2DModel,
+    _use_local_single_stream_tp,
 )
 from vllm_omni.diffusion.models.interface import SupportImageInput, SupportsComponentDiscovery
 from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
@@ -262,7 +263,11 @@ class Flux2KleinPipeline(
         ).to(self._execution_device)
 
         transformer_kwargs = get_transformer_config_kwargs(od_config.tf_model_config, Flux2Transformer2DModel)
-        self.transformer = Flux2Transformer2DModel(quant_config=od_config.quantization_config, **transformer_kwargs)
+        self.transformer = Flux2Transformer2DModel(
+            quant_config=od_config.quantization_config,
+            local_tp=_use_local_single_stream_tp(od_config),
+            **transformer_kwargs,
+        )
 
         self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8
         self.latent_channels = self.vae.config.latent_channels if hasattr(self.vae, "config") else 16