[Diffusion][Bug-Fix]: Fix Ulysses SP Accurary Error (#377)

Signed-off-by: Didan Deng <33117903+wtomin@users.noreply.github.com>
Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>
Co-authored-by: Jiangyun Zhu <riverclouds.zhu@qq.com>

docs/user_guide/acceleration/parallelism_acceleration.md

@@ -6,7 +6,7 @@ This guide includes how to use parallelism methods in vLLM-Omni to speed up diff
 
 The following parallelism methods are currently supported in vLLM-Omni:
 
-1. DeepSpeed Ulysses Sequence Parallel (Ulysses-SP) ([paper](https://arxiv.org/pdf/2309.14509)): Ulysses-SP splits the input along the sequence dimension and uses all-to-all communication to allow each device to compute only a subset of attention heads.
+1. DeepSpeed Ulysses Sequence Parallel (DeepSpeed Ulysses-SP) ([arxiv paper](https://arxiv.org/pdf/2309.14509)): Ulysses-SP splits the input along the sequence dimension and uses all-to-all communication to allow each device to compute only a subset of attention heads.
 
 
 The following table shows which models are currently supported by parallelism method:
@@ -22,9 +22,9 @@ The following table shows which models are currently supported by parallelism me
 
 #### Ulysses-SP
 
-##### Quick Start
+##### Offline Inference
 
-An example of using Ulysses-SP is shown below:
+An example of offline inference script using [Ulysses-SP](https://arxiv.org/pdf/2309.14509) is shown below:
 ```python
 from vllm_omni import Omni
 from vllm_omni.diffusion.data import DiffusionParallelConfig
@@ -70,56 +70,84 @@ To measure the parallelism methods, we run benchmarks with **Qwen/Qwen-Image** m
 
 ##### How to parallelize a new model
 
-If a diffusion model has been deployed in vLLM-Omni and supports single-card inference, you can refer to the following instruction on how to parallelize this model with Ulysses-SP.
+If a diffusion model has been deployed in vLLM-Omni and supports single-card inference, you can refer to the following instructions to parallelize it with [Ulysses-SP](https://arxiv.org/pdf/2309.14509).
 
-First, please edit the `TransformerModel`'s `forward` function in the `xxx_model_transformer.py` to make the inputs (image hidden states, positional embeddings, etc.) as chunks separated at the sequence dimension. Taking `qwen_image_transformer.py` as an example:
+This section uses **Qwen-Image** (`QwenImageTransformer2DModel`) as the reference implementation. Qwen-Image is a **dual-stream** transformer (text + image) that performs **joint attention** across the concatenated sequences. Because of that, when enabling sequence parallel you typically:
+
+- Chunk **image tokens** (`hidden_states`) across SP ranks along the **sequence dimension**.
+- Keep **text embeddings** (`encoder_hidden_states`) **replicated** on all SP ranks for correctness (unless you implement full joint-SP semantics and explicitly split text too).
+- Chunk **image RoPE freqs** to match the chunked image tokens.
+- `all_gather` the final output back along the sequence dimension.
+
+First, add the sequence-parallel helpers and (for Qwen-Image) the forward-context flag. Then, in the transformer's `forward()`:
+
+- Chunk `hidden_states` by SP world size.
+- Set `get_forward_context().split_text_embed_in_sp = False` because Qwen-Image uses joint attention and we would use a full text embedding in each rank.
+- After `pos_embed`, chunk `img_freqs` on `dim=0` (token axis) to match the chunked image tokens; only chunk `txt_freqs` if you explicitly split text embeddings in SP.
+
+Taking `vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py` as an example:
 
 ```diff
-class QwenImageTransformer2DModel(nn.Module):
+from vllm_omni.diffusion.distributed.parallel_state import (
+    get_sequence_parallel_rank,
+    get_sequence_parallel_world_size,
+    get_sp_group,
+)
+from vllm_omni.diffusion.forward_context import get_forward_context
+
+class QwenImageTransformer2DModel(...):
     def forward(
         self,
         hidden_states: torch.Tensor,
         encoder_hidden_states: torch.Tensor = None,
+        encoder_hidden_states_mask: torch.Tensor = None,
+        timestep: torch.LongTensor = None,
+        img_shapes: list[tuple[int, int, int]] | None = None,
+        txt_seq_lens: list[int] | None = None,
         ...
     ):
-+   if self.parallel_config.sequence_parallel_size > 1:
-+       hidden_states = torch.chunk(hidden_states, get_sequence_parallel_world_size(), dim=-2)[
-+           get_sequence_parallel_rank()
-+      ]
++       if self.parallel_config.sequence_parallel_size > 1:
++           # Chunk image tokens along sequence dimension.
++           hidden_states = torch.chunk(
++               hidden_states,
++               get_sequence_parallel_world_size(),
++               dim=-2,
++           )[get_sequence_parallel_rank()]
++
++           # Qwen-Image uses *dual-stream* (text + image) and runs *joint attention*.
++           # Text embeddings should be replicated across SP ranks for correctness.
++           get_forward_context().split_text_embed_in_sp = False
 
-    hidden_states = self.img_in(hidden_states)
+        hidden_states = self.img_in(hidden_states)
 
-    ...
-    image_rotary_emb = self.pos_embed(img_shapes, txt_seq_lens, device=hidden_states.device)
+        ...
+        image_rotary_emb = self.pos_embed(img_shapes, txt_seq_lens, device=hidden_states.device)
 
-+   def get_rotary_emb_chunk(freqs):
-+       freqs = torch.chunk(freqs, get_sequence_parallel_world_size(), dim=0)[get_sequence_parallel_rank()]
-+       return freqs
-
-+   if self.parallel_config.sequence_parallel_size > 1:
-+       img_freqs, txt_freqs = image_rotary_emb
-+       img_freqs = get_rotary_emb_chunk(img_freqs)
-+       image_rotary_emb = (img_freqs, txt_freqs)
++       def get_rotary_emb_chunk(freqs):
++           return torch.chunk(freqs, get_sequence_parallel_world_size(), dim=0)[
++               get_sequence_parallel_rank()
++           ]
++
++       if self.parallel_config.sequence_parallel_size > 1:
++           img_freqs, txt_freqs = image_rotary_emb
++           img_freqs = get_rotary_emb_chunk(img_freqs)
++           if get_forward_context().split_text_embed_in_sp:
++               txt_freqs = get_rotary_emb_chunk(txt_freqs)
++           image_rotary_emb = (img_freqs, txt_freqs)
 ```
 
-Next, at the end of the `forward` function, please call `get_sp_group().all_gather` to gather the chunked outputs across devices, and concatenate them at the sequence dimension.
-
+Next, at the end of the `forward()` function, call `get_sp_group().all_gather` to gather the chunked outputs across devices and concatenate them along the sequence dimension (matching the earlier `dim=-2` chunking):
 
 ```diff
-class QwenImageTransformer2DModel(nn.Module):
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        encoder_hidden_states: torch.Tensor = None,
+class QwenImageTransformer2DModel(...):
+    def forward(...):
         ...
-    ):
-    # Use only the image part (hidden_states) from the dual-stream blocks
-    hidden_states = self.norm_out(hidden_states, temb)
-    output = self.proj_out(hidden_states)
+        hidden_states = self.norm_out(hidden_states, temb)
+        output = self.proj_out(hidden_states)
 
-+   if self.parallel_config.sequence_parallel_size > 1:
-+       output = get_sp_group().all_gather(output, dim=-2)
-    return Transformer2DModelOutput(sample=output)
++       if self.parallel_config.sequence_parallel_size > 1:
++           output = get_sp_group().all_gather(output, dim=-2)
+        return Transformer2DModelOutput(sample=output)
 ```
 
 Finally, you can set the parallel configuration and pass it to `Omni` and start parallel inference with:

tests/diffusion/attention/test_ulysses_sequence_parallel.py

@@ -8,6 +8,9 @@ import pytest
 import torch
 from vllm.platforms import current_platform
 
+from vllm_omni.diffusion.attention.backends.abstract import (
+    AttentionMetadata,
+)
 from vllm_omni.diffusion.attention.layer import Attention
 from vllm_omni.diffusion.data import (
     DiffusionParallelConfig,
@@ -16,9 +19,11 @@ from vllm_omni.diffusion.data import (
 )
 from vllm_omni.diffusion.distributed.parallel_state import (
     destroy_distributed_env,
+    get_sequence_parallel_world_size,
     init_distributed_environment,
     initialize_model_parallel,
 )
+from vllm_omni.diffusion.forward_context import set_forward_context
 from vllm_omni.utils.platform_utils import detect_device_type
 
 device_type = detect_device_type()
@@ -29,6 +34,9 @@ elif device_type == "npu":
 else:
     raise ValueError(f"Unsupported device type: {device_type} for this test script! Expected GPU or NPU.")
 
+global split_text_embed_in_sp
+split_text_embed_in_sp = False
+
 
 def update_environment_variables(envs_dict: dict[str, str]):
     """Update multiple environment variables with logging."""
@@ -70,26 +78,52 @@ class TestAttentionModel(torch.nn.Module):
         self.v_proj = torch.nn.Linear(hidden_size, (num_kv_heads or num_heads) * head_size)
         self.o_proj = torch.nn.Linear(num_heads * head_size, hidden_size)
 
-    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+    def forward(self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor | None = None) -> torch.Tensor:
         """Forward pass through attention layer."""
         batch_size, seq_len, _ = hidden_states.shape
 
+        # Combine hidden_states and encoder_hidden_states if provided
+        if encoder_hidden_states is not None:
+            # Concatenate along sequence dimension
+            combined_hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)
+        else:
+            combined_hidden_states = hidden_states
+
         # Project to Q, K, V
-        q = self.q_proj(hidden_states)
-        k = self.k_proj(hidden_states)
-        v = self.v_proj(hidden_states)
+        q = self.q_proj(combined_hidden_states)
+        k = self.k_proj(combined_hidden_states)
+        v = self.v_proj(combined_hidden_states)
 
-        # Reshape to (batch_size, seq_len, num_heads, head_size)
-        q = q.view(batch_size, seq_len, self.num_heads, self.head_size)
-        k = k.view(batch_size, seq_len, k.shape[-1] // self.head_size, self.head_size)
-        v = v.view(batch_size, seq_len, v.shape[-1] // self.head_size, self.head_size)
+        # Reshape to (batch_size, total_seq_len, num_heads, head_size)
+        total_seq_len = combined_hidden_states.shape[1]
+        q = q.view(batch_size, total_seq_len, self.num_heads, self.head_size)
+        k = k.view(batch_size, total_seq_len, k.shape[-1] // self.head_size, self.head_size)
+        v = v.view(batch_size, total_seq_len, v.shape[-1] // self.head_size, self.head_size)
 
-        # Apply attention
-        attn_output = self.attention(q, k, v)
+        # Apply attention with split logic
+        if get_sequence_parallel_world_size() > 1 and not split_text_embed_in_sp and encoder_hidden_states is not None:
+            q_encoder, q_hidden = torch.split(q, [total_seq_len - seq_len, seq_len], dim=1)
+            k_encoder, k_hidden = torch.split(k, [total_seq_len - seq_len, seq_len], dim=1)
+            v_encoder, v_hidden = torch.split(v, [total_seq_len - seq_len, seq_len], dim=1)
+
+            # Use hidden_states part as main attention, encoder part as joint
+            attn_output = self.attention(
+                q_hidden,
+                k_hidden,
+                v_hidden,
+                AttentionMetadata(
+                    joint_query=q_encoder,
+                    joint_key=k_encoder,
+                    joint_value=v_encoder,
+                ),
+            )
+        else:
+            attn_output = self.attention(q, k, v)
 
         # Reshape back and project
-        attn_output = attn_output.view(batch_size, seq_len, -1)
+        attn_output = attn_output.view(batch_size, total_seq_len, -1)
         output = self.o_proj(attn_output)
+        output = output[:, encoder_hidden_states.shape[1] :, :] if encoder_hidden_states is not None else output
 
         return output
 
@@ -127,10 +161,10 @@ class TestMultiLayerAttentionModel(torch.nn.Module):
             ]
         )
 
-    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+    def forward(self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor | None = None) -> torch.Tensor:
         """Forward pass through multiple attention layers."""
         for layer in self.layers:
-            hidden_states = hidden_states + layer(hidden_states)
+            hidden_states = hidden_states + layer(hidden_states, encoder_hidden_states)
         return hidden_states
 
 
@@ -141,15 +175,16 @@ class TestMultiLayerAttentionModel(torch.nn.Module):
         TestMultiLayerAttentionModel,
     ],
 )
-@pytest.mark.parametrize("ulysses_degree", [2, 4])
+@pytest.mark.parametrize("ulysses_degree", [2])
 @pytest.mark.parametrize("ring_degree", [1])
 @pytest.mark.parametrize("batch_size", [2])
 @pytest.mark.parametrize("seq_len", [16])
+@pytest.mark.parametrize("encoder_seq_len", [16, 13])  # Test both divisible and non-divisible
 @pytest.mark.parametrize("num_heads", [8])
 @pytest.mark.parametrize("head_size", [8])
-@pytest.mark.parametrize("causal", [True])
+@pytest.mark.parametrize("causal", [False])
 @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
-@pytest.mark.parametrize("use_sync", [True, False])
+@pytest.mark.parametrize("use_sync", [True])
 @pytest.mark.parametrize("dynamic", [False])
 @pytest.mark.parametrize("use_compile", [False])
 def test_ulysses_attention(
@@ -163,12 +198,18 @@ def test_ulysses_attention(
     use_compile: bool,
     batch_size: int,
     seq_len: int,
+    encoder_seq_len: int,
     num_heads: int,
     head_size: int,
 ):
     """Test Ulysses attention by comparing with and without SP enabled."""
     sequence_parallel_size = ulysses_degree * ring_degree
 
+    # Determine if we can split encoder_hidden_states in SP
+    can_split_encoder = (encoder_seq_len % sequence_parallel_size) == 0
+    print(f"\nEncoder sequence length: {encoder_seq_len}, SP size: {sequence_parallel_size}")
+    print(f"Can split encoder in SP: {can_split_encoder}")
+
     # Create temporary files to share results between processes
     with tempfile.NamedTemporaryFile(delete=False, suffix=".pkl") as f:
         baseline_output_file = f.name
@@ -178,6 +219,8 @@ def test_ulysses_attention(
         model_state_file = f.name
     with tempfile.NamedTemporaryFile(delete=False, suffix=".pkl") as f:
         input_data_file = f.name
+    with tempfile.NamedTemporaryFile(delete=False, suffix=".pkl") as f:
+        encoder_input_data_file = f.name
 
     try:
         # Step 1: Run without SP (baseline with ulysses_degree=1, ring_degree=1)
@@ -189,6 +232,7 @@ def test_ulysses_attention(
                 test_model_cls,
                 batch_size,
                 seq_len,
+                encoder_seq_len,
                 num_heads,
                 head_size,
                 dtype,
@@ -202,6 +246,7 @@ def test_ulysses_attention(
                 baseline_output_file,
                 model_state_file,
                 input_data_file,
+                encoder_input_data_file,
                 True,  # is_baseline
             ),
             nprocs=1,
@@ -216,6 +261,7 @@ def test_ulysses_attention(
                 test_model_cls,
                 batch_size,
                 seq_len,
+                encoder_seq_len,
                 num_heads,
                 head_size,
                 dtype,
@@ -229,6 +275,7 @@ def test_ulysses_attention(
                 sp_output_file,
                 model_state_file,
                 input_data_file,
+                encoder_input_data_file,
                 False,  # is_baseline
             ),
             nprocs=sequence_parallel_size,
@@ -299,7 +346,7 @@ def test_ulysses_attention(
 
     finally:
         # Clean up temporary files
-        for f in [baseline_output_file, sp_output_file, model_state_file, input_data_file]:
+        for f in [baseline_output_file, sp_output_file, model_state_file, input_data_file, encoder_input_data_file]:
             if os.path.exists(f):
                 os.remove(f)
 
@@ -310,6 +357,7 @@ def ulysses_attention_on_test_model(
     test_model_cls: type[torch.nn.Module],
     batch_size: int,
     seq_len: int,
+    encoder_seq_len: int,
     num_heads: int,
     head_size: int,
     dtype: torch.dtype,
@@ -323,6 +371,7 @@ def ulysses_attention_on_test_model(
     output_file: str,
     model_state_file: str,
     input_data_file: str,
+    encoder_input_data_file: str,
     is_baseline: bool,
 ):
     """Run Ulysses attention test on a test model and save results for comparison."""
@@ -330,8 +379,15 @@ def ulysses_attention_on_test_model(
     RANDOM_SEED = 42
     current_platform.seed_everything(RANDOM_SEED)
 
+    # Determine if we can split encoder_hidden_states based on divisibility
+    global split_text_embed_in_sp
+    split_text_embed_in_sp = (encoder_seq_len % sequence_parallel_size) == 0 and not is_baseline
+
     mode_str = "Baseline (no SP)" if is_baseline else f"SP (ulysses={ulysses_degree}, ring={ring_degree})"
     print(f"\n[{mode_str}] Rank {local_rank}/{world_size} - Random seed set to {RANDOM_SEED}")
+    print(
+        f"[{mode_str}] Rank {local_rank}/{world_size} - encoder_seq_len={encoder_seq_len}, split_text_embed_in_sp={split_text_embed_in_sp}"
+    )
 
     device = torch.device(f"{device_type}:{local_rank}")
     torch_device.set_device(device)
@@ -379,7 +435,7 @@ def ulysses_attention_on_test_model(
     )
 
     # Set the config so Attention can access it
-    with set_current_omni_diffusion_config(od_config):
+    with set_forward_context(omni_diffusion_config=od_config), set_current_omni_diffusion_config(od_config):
         # Create model
         hidden_size = num_heads * head_size
 
@@ -418,10 +474,17 @@ def ulysses_attention_on_test_model(
                 dtype=dtype,
                 device="cpu",
             )
+            full_encoder_hidden_states = torch.randn(
+                (batch_size, encoder_seq_len, hidden_size),
+                dtype=dtype,
+                device="cpu",
+            )
             with open(input_data_file, "wb") as f:
                 pickle.dump(full_hidden_states.detach().cpu().float().numpy(), f)
+            with open(encoder_input_data_file, "wb") as f:
+                pickle.dump(full_encoder_hidden_states.detach().cpu().float().numpy(), f)
 
-            print("[Baseline] Saved model state and input data")
+            print("[Baseline] Saved model state and input data (including encoder_hidden_states)")
 
         # Synchronize to ensure baseline has saved data before SP loads it
         if world_size > 1:
@@ -437,7 +500,13 @@ def ulysses_attention_on_test_model(
             full_hidden_states_np = pickle.load(f)
         full_hidden_states = torch.from_numpy(full_hidden_states_np).to(device).to(dtype)
 
-        print(f"[Rank {local_rank}] Loaded model state and full input data with shape {full_hidden_states.shape}")
+        with open(encoder_input_data_file, "rb") as f:
+            full_encoder_hidden_states_np = pickle.load(f)
+        full_encoder_hidden_states = torch.from_numpy(full_encoder_hidden_states_np).to(device).to(dtype)
+
+        print(
+            f"[Rank {local_rank}] Loaded model state and full input data with shape {full_hidden_states.shape}, encoder shape {full_encoder_hidden_states.shape}"
+        )
 
         # Split input sequence according to sequence parallel BEFORE model forward
         # Each rank gets a contiguous chunk of the sequence dimension
@@ -446,10 +515,28 @@ def ulysses_attention_on_test_model(
         end_idx = start_idx + local_seq_len
         hidden_states = full_hidden_states[:, start_idx:end_idx, :].contiguous()
 
-        print(
-            f"[Rank {local_rank}] Split input: local_seq_len={local_seq_len}, "
-            f"indices=[{start_idx}:{end_idx}], local_shape={hidden_states.shape}"
-        )
+        # Handle encoder_hidden_states splitting based on split_text_embed_in_sp
+        if get_sequence_parallel_world_size() > 1 and split_text_embed_in_sp:
+            # Split encoder_hidden_states in the same way as hidden_states
+            local_encoder_seq_len = encoder_seq_len // sequence_parallel_size
+            encoder_start_idx = local_rank * local_encoder_seq_len
+            encoder_end_idx = encoder_start_idx + local_encoder_seq_len
+            encoder_hidden_states = full_encoder_hidden_states[:, encoder_start_idx:encoder_end_idx, :].contiguous()
+            print(
+                f"[Rank {local_rank}] Split input: local_seq_len={local_seq_len}, "
+                f"indices=[{start_idx}:{end_idx}], hidden_states shape={hidden_states.shape}, "
+                f"encoder_hidden_states (split) shape={encoder_hidden_states.shape}, "
+                f"encoder_indices=[{encoder_start_idx}:{encoder_end_idx}]"
+            )
+        else:
+            # No splitting for encoder_hidden_states, use full sequence
+            encoder_hidden_states = full_encoder_hidden_states
+            print(
+                f"[Rank {local_rank}] Split input: local_seq_len={local_seq_len}, "
+                f"indices=[{start_idx}:{end_idx}], hidden_states shape={hidden_states.shape}, "
+                f"encoder_hidden_states (full) shape={encoder_hidden_states.shape}, "
+                f"split_text_embed_in_sp={split_text_embed_in_sp}"
+            )
 
         if dynamic:
             torch._dynamo.mark_dynamic(hidden_states, 0)
@@ -461,26 +548,13 @@ def ulysses_attention_on_test_model(
 
         # Run forward pass with local sequence chunk
         print(f"[Rank {local_rank}] Running forward pass...")
-        output = model(hidden_states)
+        output = model(hidden_states, encoder_hidden_states)
         print(f"[Rank {local_rank}] Forward pass completed, output shape: {output.shape}")
-
-        # Verify output shape
         assert output.shape == (batch_size, local_seq_len, hidden_size), (
             f"Output shape mismatch: expected {(batch_size, local_seq_len, hidden_size)}, got {output.shape}"
         )
 
-        # Verify SP usage for non-baseline runs
-        if not is_baseline:
-            if hasattr(model, "attention"):
-                assert hasattr(model.attention, "use_ulysses"), "Attention should have use_ulysses attribute"
-                assert model.attention.use_ulysses, "Attention should be using Ulysses"
-            elif hasattr(model, "layers"):
-                for i, layer in enumerate(model.layers):
-                    assert hasattr(layer.attention, "use_ulysses"), (
-                        f"Layer {i} attention should have use_ulysses attribute"
-                    )
-                    assert layer.attention.use_ulysses, f"Layer {i} attention should be using Ulysses"
-
+        output = output.contiguous()
         # Gather outputs from all ranks AFTER computation
         if world_size > 1:
             print(f"[Rank {local_rank}] Gathering outputs from all {world_size} ranks...")
@@ -491,7 +565,7 @@ def ulysses_attention_on_test_model(
                 # Concatenate along sequence dimension to reconstruct full sequence
                 full_output = torch.cat(gathered_outputs, dim=1)
                 print(f"[Rank 0] Gathered and concatenated outputs: {full_output.shape}")
-                # Verify the full output shape matches expected
+
                 assert full_output.shape == (batch_size, seq_len, hidden_size), (
                     f"Gathered output shape mismatch: expected {(batch_size, seq_len, hidden_size)}, "
                     f"got {full_output.shape}"
@@ -512,9 +586,10 @@ def ulysses_attention_on_test_model(
             mode_str = "baseline (no SP)" if is_baseline else f"SP (ulysses={ulysses_degree}, ring={ring_degree})"
             print(
                 f"\n[{mode_str}] ✓ Saved output with shape {full_output.shape}:\n"
-                f"  - batch_size={batch_size}, seq_len={seq_len}\n"
+                f"  - batch_size={batch_size}, seq_len={seq_len}, encoder_seq_len={encoder_seq_len}\n"
                 f"  - num_heads={num_heads}, head_size={head_size}\n"
                 f"  - dtype={dtype}, causal={causal}, use_sync={use_sync}\n"
+                f"  - split_text_embed_in_sp={split_text_embed_in_sp}\n"
             )
 
         destroy_distributed_env()

tests/e2e/offline_inference/test_sequence_parallel.py

@@ -13,8 +13,10 @@ import os
 import sys
 from pathlib import Path
 
+import numpy as np
 import pytest
 import torch
+from PIL import Image
 
 # ruff: noqa: E402
 REPO_ROOT = Path(__file__).resolve().parents[2]
@@ -31,42 +33,113 @@ os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1"
 # Use random weights model for testing
 models = ["riverclouds/qwen_image_random"]
 
+PROMPT = "a photo of a cat sitting on a laptop keyboard"
+
+
+def _pil_to_float_rgb_tensor(img: Image.Image) -> torch.Tensor:
+    """Convert PIL image to float32 RGB tensor in [0, 1] with shape [H, W, 3]."""
+    arr = np.asarray(img.convert("RGB"), dtype=np.float32) / 255.0
+    return torch.from_numpy(arr)
+
+
+def _diff_metrics(a: Image.Image, b: Image.Image) -> tuple[float, float]:
+    """Return (mean_abs_diff, max_abs_diff) over RGB pixels in [0, 1]."""
+    ta = _pil_to_float_rgb_tensor(a)
+    tb = _pil_to_float_rgb_tensor(b)
+    assert ta.shape == tb.shape, f"Image shapes differ: {ta.shape} vs {tb.shape}"
+    abs_diff = torch.abs(ta - tb)
+    return abs_diff.mean().item(), abs_diff.max().item()
+
 
 @pytest.mark.parametrize("model_name", models)
-@pytest.mark.parametrize("ulysses_degree", [1, 2])
+@pytest.mark.parametrize("ulysses_degree", [2])
 @pytest.mark.parametrize("ring_degree", [1])
-def test_sequence_parallel(model_name: str, ulysses_degree: int, ring_degree: int):
-    """Test SP (Ulysses-SP + Ring-SP) backend with diffusion model."""
-    # Skip if not enough GPUs available
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
+def test_sequence_parallel(model_name: str, ulysses_degree: int, ring_degree: int, dtype: torch.dtype):
+    """Compare baseline (ulysses_degree=1) vs SP (ulysses_degree>1) outputs."""
+    if ulysses_degree <= 1:
+        pytest.skip("This test compares ulysses_degree=1 vs ulysses_degree>1; provide ulysses_degree>1.")
+
+    # Skip if not enough GPUs available for SP run
     if device_count() < ulysses_degree:
         pytest.skip(f"Test requires {ulysses_degree} GPUs but only {device_count()} available")
 
-    # Configure sequence parallel with DiffusionParallelConfig
-    parallel_config = DiffusionParallelConfig(ulysses_degree=ulysses_degree, ring_degree=ring_degree)
-
-    m = Omni(
-        model=model_name,
-        parallel_config=parallel_config,
-    )
-
     # Use minimal settings for fast testing
     height = 256
     width = 256
     num_inference_steps = 4  # Minimal steps for fast test
+    seed = 42
 
-    images = m.generate(
-        "a photo of a cat sitting on a laptop keyboard",
-        height=height,
-        width=width,
-        num_inference_steps=num_inference_steps,
-        guidance_scale=0.0,
-        generator=torch.Generator(get_device_name()).manual_seed(42),
-        num_outputs_per_prompt=1,  # Single output for speed
+    # Step 1: Baseline (no Ulysses sequence parallel)
+    baseline_parallel_config = DiffusionParallelConfig(ulysses_degree=1, ring_degree=1)
+    baseline = Omni(
+        model=model_name,
+        parallel_config=baseline_parallel_config,
+        dtype=dtype,
+    )
+    try:
+        baseline_images = baseline.generate(
+            PROMPT,
+            height=height,
+            width=width,
+            num_inference_steps=num_inference_steps,
+            guidance_scale=0.0,
+            generator=torch.Generator(get_device_name()).manual_seed(seed),
+            num_outputs_per_prompt=1,
+        )
+    finally:
+        baseline.close()
+
+    assert baseline_images is not None
+    assert len(baseline_images) == 1
+    assert baseline_images[0].width == width
+    assert baseline_images[0].height == height
+
+    # Step 2: SP (Ulysses-SP + Ring-SP)
+    sp_parallel_config = DiffusionParallelConfig(ulysses_degree=ulysses_degree, ring_degree=ring_degree)
+    sp = Omni(
+        model=model_name,
+        parallel_config=sp_parallel_config,
+        dtype=dtype,
+    )
+    try:
+        sp_images = sp.generate(
+            PROMPT,
+            height=height,
+            width=width,
+            num_inference_steps=num_inference_steps,
+            guidance_scale=0.0,
+            generator=torch.Generator(get_device_name()).manual_seed(seed),
+            num_outputs_per_prompt=1,
+        )
+    finally:
+        sp.close()
+
+    assert sp_images is not None
+    assert len(sp_images) == 1
+    assert sp_images[0].width == width
+    assert sp_images[0].height == height
+
+    # Step 3: Compare outputs
+    mean_abs_diff, max_abs_diff = _diff_metrics(baseline_images[0], sp_images[0])
+
+    # FP16/BF16 may differ slightly due to different computation order under parallelism.
+    if dtype in (torch.float16, torch.bfloat16):
+        mean_threshold = 2e-2
+        max_threshold = 2e-1
+    else:
+        mean_threshold = 1e-2
+        max_threshold = 1e-1
+
+    print(
+        "Image diff stats (baseline ulysses_degree=1 vs SP): "
+        f"mean_abs_diff={mean_abs_diff:.6e}, max_abs_diff={max_abs_diff:.6e}; "
+        f"thresholds: mean<={mean_threshold:.6e}, max<={max_threshold:.6e}; "
+        f"ulysses_degree={ulysses_degree}, ring_degree={ring_degree}, dtype={dtype}"
     )
 
-    # Verify generation succeeded
-    assert images is not None
-    assert len(images) == 1
-    # Check image size
-    assert images[0].width == width
-    assert images[0].height == height
+    assert mean_abs_diff <= mean_threshold and max_abs_diff <= max_threshold, (
+        f"Image diff exceeded threshold: mean_abs_diff={mean_abs_diff:.6e}, max_abs_diff={max_abs_diff:.6e} "
+        f"(thresholds: mean<={mean_threshold:.6e}, max<={max_threshold:.6e}); "
+        f"ulysses_degree={ulysses_degree}, ring_degree={ring_degree}, dtype={dtype}"
+    )

vllm_omni/diffusion/attention/backends/abstract.py

@@ -48,6 +48,14 @@ class AttentionBackend(ABC):
 @dataclass
 class AttentionMetadata:
     attn_mask: torch.Tensor | None = None
+    joint_query: torch.Tensor | None = None
+    # a replicated tensor among processes appended to the front or rear of query, depends the joint_strategy
+    joint_key: torch.Tensor | None = None
+    # a replicated tensor among processes appended to the front or rear of key, depends the joint_strategy
+    joint_value: torch.Tensor | None = None
+    # a replicated tensor among processes appended to the front or rear of value, depends the joint_strategy
+    joint_strategy: str = "front"
+    # the strategy to joint the query, key, and value, can be "front" or "rear"
 
 
 T = TypeVar("T", bound=AttentionMetadata)

vllm_omni/diffusion/attention/layer.py

@@ -7,17 +7,11 @@
 # https://github.com/feifeibear/long-context-attention/blob/main/yunchang/attention/layer.py
 
 import torch
-import torch.distributed as dist
 import torch.nn as nn
-from torch import Tensor
 
-from vllm_omni.diffusion.attention.backends.abstract import (
-    AttentionMetadata,
-)
+from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
+from vllm_omni.diffusion.attention.parallel import build_parallel_attention_strategy
 from vllm_omni.diffusion.attention.selector import get_attn_backend
-from vllm_omni.diffusion.data import get_current_omni_diffusion_config
-from vllm_omni.diffusion.distributed.comm import SeqAllToAll4D
-from vllm_omni.diffusion.distributed.parallel_state import get_sequence_parallel_world_size, get_sp_group
 
 
 class Attention(nn.Module):
@@ -49,25 +43,13 @@ class Attention(nn.Module):
         self.scatter_idx = scatter_idx
         self.gather_idx = gather_idx
         self.use_sync = use_sync
-        self.ring_pg: dist.ProcessGroup | None = None
-        self.ulysses_pg: dist.ProcessGroup | None = None
-        self.use_ulysses = False
-
-        try:
-            config = get_current_omni_diffusion_config()
-            if config.parallel_config.ulysses_degree > 1:
-                self.use_ulysses = True
-                # Get sequence parallel process group
-                try:
-                    sp_group = get_sp_group()
-                    self.ring_pg = sp_group.ring_group
-                    self.ulysses_pg = sp_group.ulysses_group
-                    assert get_sequence_parallel_world_size() > 1, "Sequence parallel world size must be > 1"
-                except (AssertionError, RuntimeError):
-                    # If sequence parallel group is not initialized, disable Ulysses
-                    self.use_ulysses = False
-        except Exception:
-            self.use_ulysses = False
+        # Parallel attention (communication / resharding) is a pluggable strategy.
+        # This keeps the attention kernel backend selection orthogonal.
+        self.parallel = build_parallel_attention_strategy(
+            scatter_idx=scatter_idx,
+            gather_idx=gather_idx,
+            use_sync=use_sync,
+        )
 
     def forward(
         self,
@@ -76,43 +58,13 @@ class Attention(nn.Module):
         value: torch.Tensor,
         attn_metadata: AttentionMetadata = None,
     ) -> torch.Tensor:
-        if self.use_ulysses:
-            return self._forward_ulysses(query, key, value, attn_metadata)
-        else:
-            # shape: (batch_size, seq_len, num_heads, head_size)
-            attn_output = self.attention.forward(query, key, value, attn_metadata)
-            return attn_output
+        # Parallel strategy may reshard/communicate QKV before the attention kernel.
+        query, key, value, attn_metadata, ctx = self.parallel.pre_attention(query, key, value, attn_metadata)
 
-    def _forward_ulysses(
-        self,
-        query: Tensor,
-        key: Tensor,
-        value: Tensor,
-        attn_metadata: AttentionMetadata = None,
-    ) -> Tensor:
-        """Ulysses attention forward pass with sequence parallelism."""
-        # scatter 2, gather 1
-        # (bs, seq_len/N, head_cnt, head_size) -> (bs, seq_len, head_cnt/N, head_size)
-        q = SeqAllToAll4D.apply(self.ulysses_pg, query, self.scatter_idx, self.gather_idx, self.use_sync)
-        k = SeqAllToAll4D.apply(self.ulysses_pg, key, self.scatter_idx, self.gather_idx, self.use_sync)
-        v = SeqAllToAll4D.apply(self.ulysses_pg, value, self.scatter_idx, self.gather_idx, self.use_sync)
+        # shape: (batch_size, seq_len, num_heads, head_size)
+        attn_output = self.attention.forward(query, key, value, attn_metadata)
+        if isinstance(attn_output, tuple):
+            attn_output = attn_output[0]
 
-        softmax_scale = self.softmax_scale
-        if softmax_scale is None:
-            softmax_scale = q.shape[-1] ** -0.5
-
-        context_layer = self.attention.forward(
-            q,
-            k,
-            v,
-            attn_metadata=attn_metadata,
-        )
-
-        if isinstance(context_layer, tuple):
-            context_layer = context_layer[0]
-
-        # (bs, seq_len, head_cnt/N, head_size) -> (bs, seq_len/N, head_cnt, head_size)
-        # scatter 1, gather 2
-        output = SeqAllToAll4D.apply(self.ulysses_pg, context_layer, self.gather_idx, self.scatter_idx, self.use_sync)
-
-        return output
+        # Parallel strategy may need to reverse resharding after the kernel.
+        return self.parallel.post_attention(attn_output, ctx)

vllm_omni/diffusion/attention/parallel/__init__.py

@@ -0,0 +1,22 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Parallel attention strategies.
+
+This package provides **communication / resharding strategies** for attention,
+orthogonal to the **attention kernel backend** (SDPA/Flash/Sage).
+
+The goal is to keep `vllm_omni.diffusion.attention.layer.Attention` small and
+extensible: adding a new parallelism method should not require editing the core
+Attention module, only adding a new strategy and selecting it in the factory.
+"""
+
+from .base import NoParallelAttention, ParallelAttentionContext, ParallelAttentionStrategy
+from .factory import build_parallel_attention_strategy
+
+__all__ = [
+    "ParallelAttentionStrategy",
+    "ParallelAttentionContext",
+    "NoParallelAttention",
+    "build_parallel_attention_strategy",
+]

vllm_omni/diffusion/attention/parallel/base.py

@@ -0,0 +1,82 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Protocol
+
+import torch
+
+from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
+
+
+@dataclass(frozen=True, slots=True)
+class ParallelAttentionContext:
+    """Opaque per-forward context returned by a parallel strategy.
+
+    Strategies may stash whatever they need here to finish post-processing after
+    the attention kernel runs (e.g. reverse resharding, slicing metadata, etc.).
+    """
+
+    name: str
+
+
+class ParallelAttentionStrategy(Protocol):
+    """Pluggable strategy for parallel attention communication/resharding.
+
+    This is intentionally orthogonal to the attention *kernel* backend.
+    The kernel backend implements `AttentionImpl.forward()` for a given device,
+    while the parallel strategy implements how Q/K/V and outputs are sharded /
+    communicated across ranks.
+    """
+
+    @property
+    def enabled(self) -> bool: ...
+
+    @property
+    def name(self) -> str: ...
+
+    def pre_attention(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata | None,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, AttentionMetadata | None, ParallelAttentionContext | None]:
+        """Runs before the attention kernel.
+
+        Returns possibly transformed Q/K/V and metadata, and an optional context
+        for `post_attention`.
+        """
+
+    def post_attention(
+        self,
+        attn_output: torch.Tensor,
+        ctx: ParallelAttentionContext | None,
+    ) -> torch.Tensor:
+        """Runs after the attention kernel."""
+
+
+class NoParallelAttention:
+    """Default strategy: do nothing (single device / no SP)."""
+
+    @property
+    def enabled(self) -> bool:
+        return False
+
+    @property
+    def name(self) -> str:
+        return "none"
+
+    def pre_attention(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata | None,
+    ):
+        return query, key, value, attn_metadata, None
+
+    def post_attention(self, attn_output: torch.Tensor, ctx: ParallelAttentionContext | None) -> torch.Tensor:
+        return attn_output

vllm_omni/diffusion/attention/parallel/factory.py

@@ -0,0 +1,45 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+from vllm_omni.diffusion.attention.parallel.base import NoParallelAttention, ParallelAttentionStrategy
+from vllm_omni.diffusion.attention.parallel.ulysses import UlyssesParallelAttention
+from vllm_omni.diffusion.data import get_current_omni_diffusion_config
+from vllm_omni.diffusion.distributed.parallel_state import get_sp_group
+
+
+def build_parallel_attention_strategy(
+    *,
+    scatter_idx: int,
+    gather_idx: int,
+    use_sync: bool,
+) -> ParallelAttentionStrategy:
+    """Select a parallel attention strategy based on current diffusion config.
+
+    Design principle:
+    - Attention kernel backend selection remains in `attention/selector.py`.
+    - Parallel attention selection is handled here, based on distributed config
+      and initialized process groups.
+    """
+    try:
+        cfg = get_current_omni_diffusion_config()
+        p = cfg.parallel_config
+    except Exception:
+        return NoParallelAttention()
+
+    # Current implementation supports Ulysses sequence-parallel attention.
+    # We intentionally do NOT infer ring attention from ring_degree here yet.
+    if getattr(p, "ulysses_degree", 1) > 1:
+        try:
+            sp_group = get_sp_group()
+        except Exception:
+            return NoParallelAttention()
+        return UlyssesParallelAttention(
+            sp_group=sp_group,
+            scatter_idx=scatter_idx,
+            gather_idx=gather_idx,
+            use_sync=use_sync,
+        )
+
+    return NoParallelAttention()

vllm_omni/diffusion/attention/parallel/ulysses.py

@@ -0,0 +1,135 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import torch
+import torch.distributed as dist
+
+from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
+from vllm_omni.diffusion.attention.parallel.base import ParallelAttentionContext
+from vllm_omni.diffusion.distributed.comm import SeqAllToAll4D
+from vllm_omni.diffusion.distributed.group_coordinator import SequenceParallelGroupCoordinator
+
+
+@dataclass(frozen=True, slots=True)
+class _UlyssesCtx(ParallelAttentionContext):
+    """Per-forward context for Ulysses sequence-parallel attention."""
+
+    ulysses_pg: dist.ProcessGroup
+    scatter_idx: int
+    gather_idx: int
+    use_sync: bool
+
+
+class UlyssesParallelAttention:
+    """Ulysses sequence-parallel strategy (all-to-all over seq/head dims).
+
+    This preserves the semantics previously implemented in
+    `Attention._forward_ulysses`:
+    - If `AttentionMetadata.joint_*` is provided, joint_query is concatenated to
+      query *before* all-to-all; joint_key/value are concatenated *after* all-to-all.
+    - joint_key/value are assumed to be replicated across SP ranks and are sliced
+      by ulysses head rank before concatenation.
+    """
+
+    def __init__(
+        self,
+        sp_group: SequenceParallelGroupCoordinator,
+        scatter_idx: int,
+        gather_idx: int,
+        use_sync: bool,
+    ) -> None:
+        self._sp_group = sp_group
+        self._ulysses_pg = sp_group.ulysses_group
+        self._scatter_idx = scatter_idx
+        self._gather_idx = gather_idx
+        self._use_sync = use_sync
+
+    @property
+    def enabled(self) -> bool:
+        return True
+
+    @property
+    def name(self) -> str:
+        return "ulysses"
+
+    def pre_attention(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata | None,
+    ):
+        joint_tensor_query = joint_tensor_key = joint_tensor_value = None
+        joint_strategy = None
+
+        if attn_metadata is not None:
+            joint_tensor_query = attn_metadata.joint_query
+            joint_tensor_key = attn_metadata.joint_key
+            joint_tensor_value = attn_metadata.joint_value
+            joint_strategy = attn_metadata.joint_strategy
+
+        is_joint = False
+        if joint_tensor_query is not None and joint_tensor_key is not None and joint_tensor_value is not None:
+            supported_joint_strategy = ["front", "rear"]
+            if joint_strategy not in supported_joint_strategy:
+                raise ValueError(
+                    f"joint_strategy: {joint_strategy} not supported."
+                    f" supported joint strategy: {supported_joint_strategy}"
+                )
+            if joint_strategy == "rear":
+                query = torch.cat([query, joint_tensor_query], dim=1)
+            else:
+                query = torch.cat([joint_tensor_query, query], dim=1)
+            is_joint = True
+        elif joint_tensor_query is None and joint_tensor_key is None and joint_tensor_value is None:
+            pass
+        else:
+            raise ValueError("joint_query, joint_key, and joint_value should be None or not None simultaneously.")
+
+        if is_joint:
+            # Slice joint key/value heads for this ulysses rank.
+            ulysses_world_size = self._sp_group.ulysses_world_size
+            ulysses_rank = self._sp_group.ulysses_rank
+            attn_heads_per_ulysses_rank = joint_tensor_key.shape[-2] // ulysses_world_size
+            joint_tensor_key = joint_tensor_key[
+                ...,
+                attn_heads_per_ulysses_rank * ulysses_rank : attn_heads_per_ulysses_rank * (ulysses_rank + 1),
+                :,
+            ]
+            joint_tensor_value = joint_tensor_value[
+                ...,
+                attn_heads_per_ulysses_rank * ulysses_rank : attn_heads_per_ulysses_rank * (ulysses_rank + 1),
+                :,
+            ]
+
+        # (bs, seq_len/P, head_cnt, head_size) -> (bs, seq_len, head_cnt/P, head_size)
+        query = SeqAllToAll4D.apply(self._ulysses_pg, query, self._scatter_idx, self._gather_idx, self._use_sync)
+        key = SeqAllToAll4D.apply(self._ulysses_pg, key, self._scatter_idx, self._gather_idx, self._use_sync)
+        value = SeqAllToAll4D.apply(self._ulysses_pg, value, self._scatter_idx, self._gather_idx, self._use_sync)
+
+        if is_joint:
+            # Concatenate joint key/value after all-to-all (matches previous implementation).
+            if joint_strategy == "front":
+                key = torch.cat([joint_tensor_key, key], dim=1)
+                value = torch.cat([joint_tensor_value, value], dim=1)
+            else:  # "rear"
+                key = torch.cat([key, joint_tensor_key], dim=1)
+                value = torch.cat([value, joint_tensor_value], dim=1)
+
+        ctx = _UlyssesCtx(
+            name=self.name,
+            ulysses_pg=self._ulysses_pg,
+            scatter_idx=self._scatter_idx,
+            gather_idx=self._gather_idx,
+            use_sync=self._use_sync,
+        )
+        return query, key, value, attn_metadata, ctx
+
+    def post_attention(self, attn_output: torch.Tensor, ctx: ParallelAttentionContext | None) -> torch.Tensor:
+        assert isinstance(ctx, _UlyssesCtx), f"Unexpected ctx type: {type(ctx)!r}"
+        # Reverse: (bs, seq_len, head_cnt/P, head_size) -> (bs, seq_len/P, head_cnt, head_size)
+        return SeqAllToAll4D.apply(ctx.ulysses_pg, attn_output, ctx.gather_idx, ctx.scatter_idx, ctx.use_sync)

vllm_omni/diffusion/forward_context.py

@@ -0,0 +1,90 @@
+from contextlib import contextmanager
+from dataclasses import dataclass
+
+from vllm.config import VllmConfig
+
+from vllm_omni.diffusion.attention.backends.abstract import (
+    AttentionMetadata,
+)
+from vllm_omni.diffusion.data import OmniDiffusionConfig
+
+
+@dataclass
+class ForwardContext:
+    """
+    set forward context for diffusion models
+    """
+
+    vllm_config: VllmConfig | None = None
+    omni_diffusion_config: OmniDiffusionConfig | None = None
+    attn_metadata: dict[str, AttentionMetadata] | list[dict[str, AttentionMetadata]] | None = None
+    split_text_embed_in_sp: bool = False
+    # whether to split the text embed in sequence parallel, if True, the text embed will be split in sequence parallel
+
+    def __post_init__(self):
+        pass
+
+
+_forward_context: ForwardContext | None = None
+
+
+def get_forward_context() -> ForwardContext:
+    """Get the current forward context."""
+    assert _forward_context is not None, (
+        "Forward context is not set. Please use `set_forward_context` to set the forward context."
+    )
+    return _forward_context
+
+
+def is_forward_context_available() -> bool:
+    return _forward_context is not None
+
+
+def create_forward_context(
+    vllm_config: VllmConfig | None = None,
+    omni_diffusion_config: OmniDiffusionConfig | None = None,
+    attn_metadata: dict[str, AttentionMetadata] | list[dict[str, AttentionMetadata]] | None = None,
+    split_text_embed_in_sp: bool = False,
+):
+    return ForwardContext(
+        vllm_config=vllm_config,
+        omni_diffusion_config=omni_diffusion_config,
+        attn_metadata=attn_metadata,
+        split_text_embed_in_sp=split_text_embed_in_sp,
+    )
+
+
+@contextmanager
+def override_forward_context(forward_context: ForwardContext | None):
+    """A context manager that overrides the current forward context.
+    This is used to override the forward context for a specific
+    forward pass.
+    """
+    global _forward_context
+    prev_context = _forward_context
+    _forward_context = forward_context
+    try:
+        yield
+    finally:
+        _forward_context = prev_context
+
+
+@contextmanager
+def set_forward_context(
+    vllm_config: VllmConfig | None = None,
+    omni_diffusion_config: OmniDiffusionConfig | None = None,
+    attn_metadata: dict[str, AttentionMetadata] | list[dict[str, AttentionMetadata]] | None = None,
+    split_text_embed_in_sp: bool = False,
+):
+    """A context manager that stores the current forward context,
+    can be attention metadata, split_text_embed_in_sp, etc.
+    Here we can inject common logic for every model forward pass.
+    """
+    forward_context = create_forward_context(
+        vllm_config=vllm_config,
+        omni_diffusion_config=omni_diffusion_config,
+        attn_metadata=attn_metadata,
+        split_text_embed_in_sp=split_text_embed_in_sp,
+    )
+    with override_forward_context(forward_context):
+        yield

vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py

@@ -19,14 +19,18 @@ from vllm.model_executor.layers.layernorm import RMSNorm
 from vllm.model_executor.layers.linear import QKVParallelLinear, ReplicatedLinear
 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 
+from vllm_omni.diffusion.attention.backends.abstract import (
+    AttentionMetadata,
+)
 from vllm_omni.diffusion.attention.layer import Attention
 from vllm_omni.diffusion.cache.base import CachedTransformer
-from vllm_omni.diffusion.data import OmniDiffusionConfig
+from vllm_omni.diffusion.data import OmniDiffusionConfig, get_current_omni_diffusion_config
 from vllm_omni.diffusion.distributed.parallel_state import (
     get_sequence_parallel_rank,
     get_sequence_parallel_world_size,
     get_sp_group,
 )
+from vllm_omni.diffusion.forward_context import get_forward_context
 from vllm_omni.diffusion.layers.rope import RotaryEmbedding
 
 logger = init_logger(__name__)
@@ -356,6 +360,12 @@ class QwenImageCrossAttention(nn.Module):
         )
         self.rope = RotaryEmbedding(is_neox_style=False)
 
+        try:
+            config = get_current_omni_diffusion_config()
+            self.parallel_config = config.parallel_config
+        except Exception:
+            self.parallel_config = None
+
     def forward(
         self,
         hidden_states: torch.Tensor,
@@ -405,11 +415,31 @@ class QwenImageCrossAttention(nn.Module):
         joint_value = torch.cat([txt_value, img_value], dim=1)
 
         # Compute joint attention
-        joint_hidden_states = self.attn(
-            joint_query,
-            joint_key,
-            joint_value,
-        )
+
+        if (
+            self.parallel_config is not None
+            and self.parallel_config.sequence_parallel_size > 1
+            and not get_forward_context().split_text_embed_in_sp
+        ):
+            # if using sequence parallel, but not splitting text embed,
+            #  we need to pass text embedding to attention layer as joint qkv
+            joint_hidden_states = self.attn(
+                img_query,
+                img_key,
+                img_value,
+                AttentionMetadata(
+                    joint_query=txt_query,
+                    joint_key=txt_key,
+                    joint_value=txt_value,
+                    joint_strategy="front",
+                ),
+            )
+        else:
+            joint_hidden_states = self.attn(
+                joint_query,
+                joint_key,
+                joint_value,
+            )
         joint_hidden_states = joint_hidden_states.flatten(2, 3)
         joint_hidden_states = joint_hidden_states.to(joint_query.dtype)
 
@@ -709,7 +739,10 @@ class QwenImageTransformer2DModel(CachedTransformer):
             hidden_states = torch.chunk(hidden_states, get_sequence_parallel_world_size(), dim=-2)[
                 get_sequence_parallel_rank()
             ]
-
+            # NOTE:
+            # QwenImage uses *dual-stream* (text + image) and runs a *joint attention*.
+            # text embeddings to be replicated across SP ranks for correctness.
+            get_forward_context().split_text_embed_in_sp = False
         hidden_states = self.img_in(hidden_states)
 
         # Ensure timestep tensor is on the same device and dtype as hidden_states
@@ -746,6 +779,8 @@ class QwenImageTransformer2DModel(CachedTransformer):
         if self.parallel_config.sequence_parallel_size > 1:
             img_freqs, txt_freqs = image_rotary_emb
             img_freqs = get_rotary_emb_chunk(img_freqs)
+            if get_forward_context().split_text_embed_in_sp:
+                txt_freqs = get_rotary_emb_chunk(txt_freqs)
             image_rotary_emb = (img_freqs, txt_freqs)
 
         for index_block, block in enumerate(self.transformer_blocks):

vllm_omni/diffusion/worker/gpu_worker.py

@@ -23,6 +23,7 @@ from vllm_omni.diffusion.distributed.parallel_state import (
     init_distributed_environment,
     initialize_model_parallel,
 )
+from vllm_omni.diffusion.forward_context import set_forward_context
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 
@@ -65,31 +66,33 @@ class GPUWorker:
         vllm_config = VllmConfig()
         vllm_config.parallel_config.tensor_parallel_size = self.od_config.parallel_config.tensor_parallel_size
         vllm_config.parallel_config.data_parallel_size = self.od_config.parallel_config.data_parallel_size
+        self.vllm_config = vllm_config
+        with (
+            set_current_omni_diffusion_config(self.od_config),
+            set_current_vllm_config(vllm_config),
+        ):
+            init_distributed_environment(world_size=world_size, rank=rank)
+            logger.info(f"Worker {self.rank}: Initialized device and distributed environment.")
+            parallel_config = self.od_config.parallel_config
+            initialize_model_parallel(
+                data_parallel_size=parallel_config.data_parallel_size,
+                cfg_parallel_size=parallel_config.cfg_parallel_size,
+                sequence_parallel_size=parallel_config.sequence_parallel_size,
+                ulysses_degree=parallel_config.ulysses_degree,
+                ring_degree=parallel_config.ring_degree,
+                tensor_parallel_size=parallel_config.tensor_parallel_size,
+                pipeline_parallel_size=parallel_config.pipeline_parallel_size,
+            )
 
-        with set_current_omni_diffusion_config(self.od_config):
-            with set_current_vllm_config(vllm_config):
-                init_distributed_environment(world_size=world_size, rank=rank)
-                logger.info(f"Worker {self.rank}: Initialized device and distributed environment.")
-                parallel_config = self.od_config.parallel_config
-                initialize_model_parallel(
-                    data_parallel_size=parallel_config.data_parallel_size,
-                    cfg_parallel_size=parallel_config.cfg_parallel_size,
-                    sequence_parallel_size=parallel_config.sequence_parallel_size,
-                    ulysses_degree=parallel_config.ulysses_degree,
-                    ring_degree=parallel_config.ring_degree,
-                    tensor_parallel_size=parallel_config.tensor_parallel_size,
-                    pipeline_parallel_size=parallel_config.pipeline_parallel_size,
+            load_config = LoadConfig()
+            model_loader = DiffusersPipelineLoader(load_config)
+            time_before_load = time.perf_counter()
+            with DeviceMemoryProfiler() as m:
+                self.pipeline = model_loader.load_model(
+                    od_config=self.od_config,
+                    load_device=f"cuda:{rank}",
                 )
-
-                load_config = LoadConfig()
-                model_loader = DiffusersPipelineLoader(load_config)
-                time_before_load = time.perf_counter()
-                with DeviceMemoryProfiler() as m:
-                    self.pipeline = model_loader.load_model(
-                        od_config=self.od_config,
-                        load_device=f"cuda:{rank}",
-                    )
-                time_after_load = time.perf_counter()
+            time_after_load = time.perf_counter()
 
         logger.info(
             "Model loading took %.4f GiB and %.6f seconds",
@@ -116,8 +119,8 @@ class GPUWorker:
         # Refresh cache context if needed
         if self.cache_backend is not None and self.cache_backend.is_enabled():
             self.cache_backend.refresh(self.pipeline, req.num_inference_steps)
-
-        output = self.pipeline.forward(req)
+        with set_forward_context(vllm_config=self.vllm_config, omni_diffusion_config=self.od_config):
+            output = self.pipeline.forward(req)
         return output
 
     def shutdown(self) -> None: