[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>

tests/diffusion/models/flux2/test_flux2_klein_local_tp.py

@@ -0,0 +1,331 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""CPU regressions for Klein's packed TP layout and local-head execution."""
+
+from contextlib import contextmanager
+
+import pytest
+import torch
+import torch.nn.functional as F
+from vllm.model_executor import parameter
+from vllm.model_executor.layers import linear
+from vllm.model_executor.layers.utils import default_unquantized_gemm
+from vllm.model_executor.model_loader.weight_utils import default_weight_loader
+
+from vllm_omni.diffusion.data import DiffusionParallelConfig, OmniDiffusionConfig
+from vllm_omni.diffusion.forward_context import ForwardContext
+from vllm_omni.diffusion.models.flux2_klein import flux2_klein_transformer as klein
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+@contextmanager
+def _default_dtype(dtype):
+    previous = torch.get_default_dtype()
+    torch.set_default_dtype(dtype)
+    try:
+        yield
+    finally:
+        torch.set_default_dtype(previous)
+
+
+class _CPUAttention(torch.nn.Module):
+    """Use SDPA while retaining the model's actual Q/K norm and RoPE."""
+
+    def __init__(self, num_heads, head_size, softmax_scale, causal=False, **kwargs):
+        super().__init__()
+        self.num_heads = num_heads
+        self.scale = softmax_scale
+
+    def forward(self, query, key, value, metadata=None):
+        assert query.shape[2] == self.num_heads
+        return F.scaled_dot_product_attention(
+            query.transpose(1, 2),
+            key.transpose(1, 2),
+            value.transpose(1, 2),
+            attn_mask=None if metadata is None else metadata.attn_mask,
+            scale=self.scale,
+        ).transpose(1, 2)
+
+
+@pytest.fixture(autouse=True)
+def _cpu_dispatch(monkeypatch):
+    monkeypatch.setattr(linear, "dispatch_unquantized_gemm", lambda *a, **kw: default_unquantized_gemm)
+    monkeypatch.setattr(klein, "Attention", _CPUAttention)
+    monkeypatch.setattr(klein, "get_forward_context", ForwardContext)
+
+
+def _set_tp(monkeypatch, tp, rank):
+    # Keep real vLLM parameters, constructors and checkpoint loaders. Only the
+    # distributed process metadata is substituted for this CPU unit test.
+    for module in (linear, parameter):
+        monkeypatch.setattr(module, "get_tensor_model_parallel_world_size", lambda: tp)
+        monkeypatch.setattr(module, "get_tensor_model_parallel_rank", lambda: rank)
+    monkeypatch.setattr(klein, "get_tensor_model_parallel_world_size", lambda: tp)
+
+
+def _attention(monkeypatch, tp, rank, *, local_tp=True, heads=8, dim_head=8, mlp_ratio=3, sp=1):
+    _set_tp(monkeypatch, tp, rank)
+    with _default_dtype(torch.bfloat16):
+        attention = klein.Flux2ParallelSelfAttention(
+            parallel_config=DiffusionParallelConfig(tensor_parallel_size=tp, ulysses_degree=sp),
+            query_dim=heads * dim_head,
+            heads=heads,
+            dim_head=dim_head,
+            mlp_ratio=mlp_ratio,
+            bias=True,
+            local_tp=local_tp,
+        )
+    attention.norm_q.forward = attention.norm_q.forward_native
+    attention.norm_k.forward = attention.norm_k.forward_native
+    attention.rope.forward = attention.rope.forward_native
+    return attention
+
+
+def _rank_rows(tensor, sizes, tp, rank):
+    return torch.cat([part.chunk(tp, dim=0)[rank] for part in tensor.split(sizes, dim=0)])
+
+
+@pytest.mark.parametrize("tp", [1, 2, 4, 8])
+@pytest.mark.parametrize("name", ["weight", "bias"])
+def test_packed_checkpoint_reconstruction_and_reload(monkeypatch, tp, name):
+    sizes = [64, 64, 64, 192, 192]
+    shape = (sum(sizes), 64) if name == "weight" else (sum(sizes),)
+    canonical = torch.arange(torch.tensor(shape).prod().item(), dtype=torch.float32).reshape(shape)
+    shards = []
+    for rank in range(tp):
+        attention = _attention(monkeypatch, tp, rank)
+        assert attention.local_tp == (tp > 1)
+        assert attention.query_num_heads == 8 // tp
+        projection = attention.to_qkv_mlp_proj
+        param = getattr(projection, name)
+        # Exact row markers verify the installed loader, without BF16 rounding
+        # hiding a channel-order error. The constructor still uses BF16.
+        param.data = param.data.float()
+        for weights in (canonical + 1, canonical):
+            param.weight_loader(param, weights)
+            torch.testing.assert_close(param, _rank_rows(weights, sizes, tp, rank), rtol=0, atol=0)
+        shards.append(param.detach().clone())
+    local_sizes = [size // tp for size in sizes]
+    reconstructed = torch.cat(
+        [torch.cat([shard.split(local_sizes, dim=0)[i] for shard in shards]) for i in range(len(sizes))]
+    )
+    torch.testing.assert_close(reconstructed, canonical, rtol=0, atol=0)
+
+
+@pytest.mark.parametrize("tp,heads", [(2, 8), (4, 24), (8, 24)])
+@pytest.mark.parametrize("mask_kind", ["none", "3d", "broadcast", "per_head"])
+@torch.inference_mode()
+def test_local_heads_and_gather_order_match_dense_attention(monkeypatch, tp, heads, mask_kind):
+    torch.manual_seed(17)
+    dim = heads * 8
+    baseline = _attention(monkeypatch, 1, 0, local_tp=False, heads=heads)
+    for name, param in baseline.named_parameters():
+        if name.startswith("norm_"):
+            param.fill_(1)
+        else:
+            param.normal_(std=0.03)
+    # A non-contiguous input and batch > 1 exercise the reshape boundaries.
+    hidden = torch.randn(2, dim, 7, dtype=torch.bfloat16).transpose(1, 2)
+    angles = torch.randn(7, 4)
+    rotary = angles.cos(), angles.sin()
+    mask = None
+    if mask_kind != "none":
+        mask_heads = heads if mask_kind == "per_head" else 1
+        mask = torch.rand(2, mask_heads, 7, 7) > 0.25
+        mask[..., 0] = True
+        if mask_kind == "3d":
+            mask = mask[:, 0]
+
+    captured = []
+    hook = baseline.to_out.register_forward_pre_hook(lambda module, args: captured.append(args[0].clone()))
+    expected = baseline(hidden, attention_mask=mask, image_rotary_emb=rotary)
+    hook.remove()
+    full_attention, full_mlp = captured[0].split([dim, 3 * dim], dim=-1)
+    weights = {name: param.detach().clone() for name, param in baseline.named_parameters()}
+
+    for rank in range(tp):
+        local = _attention(monkeypatch, tp, rank, heads=heads)
+        for name, param in local.named_parameters():
+            getattr(param, "weight_loader", default_weight_loader)(param, weights[name])
+        feature_calls: list[int] = []
+        output_calls: list[int] = []
+
+        def gather_features(value, dim=-1):
+            reference = (full_attention, full_mlp)[len(feature_calls)]
+            torch.testing.assert_close(value, reference.chunk(tp, dim=-1)[rank], rtol=0.02, atol=0.002)
+            feature_calls.append(value.shape[-1])
+            return reference
+
+        def gather_output(value, dim=-1):
+            torch.testing.assert_close(value, expected.chunk(tp, dim=-1)[rank], rtol=0.02, atol=0.002)
+            output_calls.append(value.shape[-1])
+            return expected
+
+        # Each outgoing tensor must match the independently computed dense
+        # result before the collective supplies the other ranks' slices.
+        monkeypatch.setattr(klein, "tensor_model_parallel_all_gather", gather_features)
+        monkeypatch.setattr(linear, "tensor_model_parallel_all_gather", gather_output)
+        actual = local(hidden, attention_mask=mask, image_rotary_emb=rotary)
+        torch.testing.assert_close(actual, expected)
+        assert feature_calls == [dim // tp, 3 * dim // tp]
+        assert output_calls == [dim // tp]
+
+
+@pytest.mark.parametrize(
+    "kwargs",
+    [
+        {"tp": 1},
+        {"tp": 4, "local_tp": False},
+        {"tp": 4, "sp": 2},
+        {"tp": 4, "heads": 6},
+        {"tp": 4, "mlp_ratio": 2.1},
+    ],
+)
+def test_ineligible_shapes_keep_original_projection(monkeypatch, kwargs):
+    attention = _attention(monkeypatch, rank=0, **kwargs)
+    assert not attention.local_tp
+    assert type(attention.to_qkv_mlp_proj) is linear.ColumnParallelLinear
+    assert attention.to_qkv_mlp_proj.gather_output
+    assert attention.query_num_heads == attention.heads
+
+
+@pytest.mark.parametrize(
+    "overrides,parallel_overrides,eligible",
+    [
+        ({}, {}, True),
+        ({"dtype": torch.float32}, {}, False),
+        ({"dtype": torch.float16}, {}, False),
+        ({}, {"tensor_parallel_size": 1}, False),
+        ({}, {"ulysses_degree": 2}, False),
+        ({"quantization_config": "fp8"}, {}, False),
+        ({"lora_path": "local-adapter"}, {}, False),
+        ({"enable_cpu_offload": True}, {}, False),
+        ({"enable_layerwise_offload": True}, {}, False),
+        ({"diffusion_offload_config": {"mode": "module", "components": ["dit"]}}, {}, False),
+        ({}, {"tensor_parallel_size": 1, "use_hsdp": True, "hsdp_shard_size": 4}, False),
+    ],
+)
+def test_pipeline_eligibility_uses_resolved_config(tmp_path, overrides, parallel_overrides, eligible):
+    parallel = DiffusionParallelConfig(**({"tensor_parallel_size": 4} | parallel_overrides))
+    config = OmniDiffusionConfig(
+        **({"model": str(tmp_path), "dtype": torch.bfloat16} | overrides),
+        parallel_config=parallel,
+    )
+    assert klein._use_local_single_stream_tp(config) is eligible
+
+
+@pytest.mark.parametrize("tp", [1, 2, 4, 8])
+def test_host_restore_contract_and_model_weight_reload(monkeypatch, tp):
+    _set_tp(monkeypatch, tp, 0)
+    kwargs = dict(
+        in_channels=8,
+        num_layers=1,
+        num_single_layers=1,
+        attention_head_dim=8,
+        num_attention_heads=8,
+        joint_attention_dim=8,
+        timestep_guidance_channels=8,
+        axes_dims_rope=(2, 2, 2, 2),
+        guidance_embeds=False,
+    )
+    with _default_dtype(torch.bfloat16):
+        baseline = klein.Flux2Transformer2DModel(**kwargs)
+        model = klein.Flux2Transformer2DModel(**kwargs, local_tp=True)
+    assert baseline.host_weight_restore_contract.version == "1"
+    assert model.host_weight_restore_contract.version == ("2-local-single-stream-tp" if tp > 1 else "1")
+    model.validate_restored_host_weights()
+    name = "single_transformer_blocks.0.attn.to_qkv_mlp_proj.weight"
+    weights = torch.linspace(-1, 1, 576 * 64).reshape(576, 64).to(torch.bfloat16)
+    for canonical in (weights, weights.neg()):
+        assert name in model.load_weights([(name, canonical)])
+        param = dict(model.named_parameters())[name]
+        torch.testing.assert_close(param, _rank_rows(canonical, [64, 64, 64, 192, 192], tp, 0), rtol=0, atol=0)
+    with _default_dtype(torch.bfloat16):
+        restored = klein.Flux2Transformer2DModel(**kwargs, local_tp=True)
+    restored.load_state_dict(model.state_dict(), strict=True)
+    restored.validate_restored_host_weights()
+    torch.testing.assert_close(dict(restored.named_parameters())[name], param, rtol=0, atol=0)
+
+
+@pytest.mark.parametrize("sizes", [(8, 24), (8, 8, 8, 24, 24)])
+@pytest.mark.parametrize("tp", [1, 2, 4, 8])
+@torch.inference_mode()
+def test_fused_lora_global_rows_dynamic_activation(monkeypatch, sizes, tp):
+    from vllm.lora.lora_model import LoRAModel
+    from vllm.lora.lora_weights import LoRALayerWeights
+    from vllm.lora.peft_helper import PEFTHelper
+    from vllm.lora.request import LoRARequest
+
+    from vllm_omni.diffusion.lora.layers.column_parallel_linear import DiffusionMergedColumnParallelLinearWithLoRA
+    from vllm_omni.diffusion.lora.manager import DiffusionLoRAManager
+
+    torch.manual_seed(29)
+    weights = torch.randn(sum(sizes), 8)
+    lora_a = torch.randn(2, 8)
+    lora_b = torch.randn(sum(sizes), 2)
+    hidden = torch.randn(2, 3, 8)
+    name = "transformer.to_qkv_mlp_proj"
+    peft = PEFTHelper(r=2, lora_alpha=2, target_modules=["to_qkv_mlp_proj"])
+    request = LoRARequest(lora_name="test", lora_int_id=7, lora_path="unused-by-memory-loader")
+    for rank in range(tp):
+        _set_tp(monkeypatch, tp, rank)
+        projection = linear.MergedColumnParallelLinear(8, list(sizes), bias=False, gather_output=False)
+        projection.weight.weight_loader(projection.weight, weights)
+        pipeline = torch.nn.Module()
+        pipeline.transformer = torch.nn.Module()
+        pipeline.transformer.to_qkv_mlp_proj = projection
+        manager = DiffusionLoRAManager(pipeline, device=torch.device("cpu"), dtype=torch.float32)
+        adapter = LoRAModel(7, 2, {name: LoRALayerWeights(name, 2, 2, lora_a, lora_b)})
+        # Adapter disk I/O is outside this regression. Actual manager injection,
+        # five-/two-slice wrapper, rank slicing and activation all run below.
+        monkeypatch.setattr(manager, "_load_adapter", lambda request: (adapter, peft))
+        base_output = F.linear(hidden, _rank_rows(weights, sizes, tp, rank))
+        delta = F.linear(F.linear(hidden, lora_a), _rank_rows(lora_b, sizes, tp, rank))
+        for scale in (0.5, 1.0):
+            manager.set_active_adapter(request, lora_scale=scale)
+            wrapped = pipeline.transformer.to_qkv_mlp_proj
+            assert isinstance(wrapped, DiffusionMergedColumnParallelLinearWithLoRA)
+            assert wrapped.n_slices == len(sizes)
+            actual, _ = wrapped(hidden)
+            torch.testing.assert_close(actual, base_output + scale * delta, rtol=1e-5, atol=1e-5)
+            manager.set_active_adapter(None)
+            actual, _ = wrapped(hidden)
+            torch.testing.assert_close(actual, base_output, rtol=1e-5, atol=1e-5)
+
+
+@pytest.mark.parametrize("image_count", [1, 2])
+def test_native_editing_reaches_normal_input_validation(tmp_path, monkeypatch, image_count):
+    from PIL import Image
+
+    from vllm_omni.diffusion.models.flux2_klein.pipeline_flux2_klein import Flux2KleinPipeline
+    from vllm_omni.diffusion.request import OmniDiffusionRequest
+    from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
+    from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+    class InputsReachedError(Exception):
+        pass
+
+    pipeline = object.__new__(Flux2KleinPipeline)
+    pipeline.od_config = OmniDiffusionConfig(
+        model=str(tmp_path),
+        dtype=torch.bfloat16,
+        parallel_config=DiffusionParallelConfig(tensor_parallel_size=4),
+    )
+    assert klein._use_local_single_stream_tp(pipeline.od_config)
+    images = [Image.new("RGB", (16, 16)) for _ in range(image_count)]
+    image_input = images[0] if image_count == 1 else images
+    request = OmniDiffusionRequest(
+        prompt={"prompt": "paint the boat red", "multi_modal_data": {"image": image_input}},
+        request_id="native-editing",
+        sampling_params=OmniDiffusionSamplingParams(height=512, width=512, num_inference_steps=4),
+    )
+
+    def check_inputs(**kwargs):
+        assert kwargs["prompt"] == "paint the boat red"
+        raise InputsReachedError
+
+    monkeypatch.setattr(pipeline, "check_inputs", check_inputs)
+    with pytest.raises(InputsReachedError):
+        pipeline.forward(DiffusionRequestBatch(requests=[request]))
+    assert request.prompt["multi_modal_data"]["image"] is image_input

vllm_omni/diffusion/lora/manager.py

@@ -693,6 +693,11 @@ class DiffusionLoRAManager:
                         )
                     b_splits = list(torch.split(deinterleaved, list(output_sizes), dim=0))
                 else:
+                    base = getattr(lora_layer, "base_layer", None)
+                    if isinstance(base, MergedColumnParallelLinear):
+                        # Adapter B contains global rows. The LoRA layer
+                        # slices each packed segment for its TP rank below.
+                        output_slices = base.output_sizes
                     total = sum(output_slices)
                     if lora_weights.lora_b.shape[0] != total:
                         raise ValueError(

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