[BugFix][Qwen-Image] align txt_seq_lens RoPE width with padded embeds (#4474)

Signed-off-by: samithuang <285365963@qq.com>

tests/diffusion/models/qwen_image/test_qwen_image_rope_table_failure_case.py

@@ -0,0 +1,44 @@
+"""Regression tests for #4443: padded prompt width vs RoPE table length.
+
+When ``max(txt_seq_lens) < prompt_embeds.shape[1]``, ``QwenEmbedRope`` builds a
+text frequency table that is too short for the padded encoder width.
+"""
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.models.qwen_image.qwen_image_transformer import QwenEmbedRope
+from vllm_omni.diffusion.models.qwen_image.rope_utils import txt_seq_lens_from_embeds
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+# Representative request-batch padding: valid tokens 10/20, shared width 32.
+PADDED_WIDTH = 32
+VALID_LENS = [10, 20]
+
+
+def test_pre_pr_txt_seq_lens_is_shorter_than_padded_width():
+    mask = torch.zeros(len(VALID_LENS), PADDED_WIDTH, dtype=torch.bool)
+    mask[0, : VALID_LENS[0]] = True
+    mask[1, : VALID_LENS[1]] = True
+
+    pre_pr_txt_seq_lens = mask.sum(dim=1).tolist()
+    assert pre_pr_txt_seq_lens == VALID_LENS
+    assert max(pre_pr_txt_seq_lens) < PADDED_WIDTH
+
+
+def test_post_pr_txt_seq_lens_matches_padded_width():
+    prompt_embeds = torch.zeros(len(VALID_LENS), PADDED_WIDTH, 8)
+    assert txt_seq_lens_from_embeds(prompt_embeds) == [PADDED_WIDTH, PADDED_WIDTH]
+
+
+def test_rope_table_length_matches_padded_width_after_fix():
+    rope = QwenEmbedRope(theta=10000, axes_dim=[16, 56, 56], scale_rope=True)
+    img_shapes = [[(1, 16, 16)]]
+
+    _, freqs_pre_pr = rope(img_shapes, VALID_LENS, device="cpu")
+    _, freqs_post_pr = rope(img_shapes, [PADDED_WIDTH, PADDED_WIDTH], device="cpu")
+
+    assert freqs_pre_pr.shape[0] == max(VALID_LENS)
+    assert freqs_pre_pr.shape[0] < PADDED_WIDTH
+    assert freqs_post_pr.shape[0] == PADDED_WIDTH

tests/diffusion/models/qwen_image/test_qwen_image_rope_utils.py

@@ -0,0 +1,30 @@
+import pytest
+import torch
+
+from vllm_omni.diffusion.models.qwen_image.rope_utils import txt_seq_lens_from_embeds
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_txt_seq_lens_from_embeds_uses_padded_width_not_valid_token_count():
+    prompt_embeds = torch.zeros(2, 64, 8)
+    prompt_embeds_mask = torch.zeros(2, 64, dtype=torch.bool)
+    prompt_embeds_mask[:, :10] = True
+
+    assert prompt_embeds_mask.sum(dim=1).tolist() == [10, 10]
+    assert txt_seq_lens_from_embeds(prompt_embeds) == [64, 64]
+
+
+def test_txt_seq_lens_from_embeds_supports_2d_embeds():
+    prompt_embeds = torch.zeros(48, 16)
+
+    assert txt_seq_lens_from_embeds(prompt_embeds) == [48]
+
+
+def test_txt_seq_lens_from_embeds_returns_none_for_missing_embeds():
+    assert txt_seq_lens_from_embeds(None) is None
+
+
+def test_txt_seq_lens_from_embeds_rejects_invalid_rank():
+    with pytest.raises(ValueError, match="prompt_embeds must be 2D or 3D"):
+        txt_seq_lens_from_embeds(torch.zeros(2, 3, 4, 5))

tests/diffusion/test_input_batch_txt_seq_lens.py

@@ -0,0 +1,31 @@
+import pytest
+import torch
+
+from vllm_omni.diffusion.worker.input_batch import _prepare_request_prompt_field
+from vllm_omni.diffusion.worker.utils import DiffusionRequestState
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+pytestmark = [pytest.mark.cpu]
+
+
+def test_prepare_request_prompt_field_refreshes_txt_seq_lens_after_padding():
+    state = DiffusionRequestState(
+        request_id="req-0",
+        sampling=OmniDiffusionSamplingParams(),
+        prompt_embeds=torch.randn(1, 10, 4),
+        prompt_embeds_mask=torch.ones(1, 10, dtype=torch.bool),
+        txt_seq_lens=[10],
+    )
+
+    embeds, mask = _prepare_request_prompt_field(
+        state,
+        embeds_attr="prompt_embeds",
+        mask_attr="prompt_embeds_mask",
+        seq_lens_attr="txt_seq_lens",
+        target_seq_len=32,
+    )
+
+    assert embeds.shape == (1, 32, 4)
+    assert mask is not None and mask.shape == (1, 32)
+    assert state.txt_seq_lens == [32]
+    assert mask.sum(dim=1).tolist() == [10]

vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py

@@ -34,6 +34,7 @@ from vllm_omni.diffusion.models.qwen_image.cfg_parallel import (
 from vllm_omni.diffusion.models.qwen_image.qwen_image_transformer import (
     QwenImageTransformer2DModel,
 )
+from vllm_omni.diffusion.models.qwen_image.rope_utils import txt_seq_lens_from_embeds
 from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.utils.prompt_utils import (
@@ -730,10 +731,8 @@ class QwenImagePipeline(nn.Module, QwenImageCFGParallelMixin, DiffusionPipelineP
         else:
             guidance = None
 
-        txt_seq_lens = prompt_embeds_mask.sum(dim=1).tolist() if prompt_embeds_mask is not None else None
-        negative_txt_seq_lens = (
-            negative_prompt_embeds_mask.sum(dim=1).tolist() if negative_prompt_embeds_mask is not None else None
-        )
+        txt_seq_lens = txt_seq_lens_from_embeds(prompt_embeds)
+        negative_txt_seq_lens = txt_seq_lens_from_embeds(negative_prompt_embeds)
 
         return {
             "prompt_embeds": prompt_embeds,

vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit.py

@@ -36,6 +36,7 @@ from vllm_omni.diffusion.models.qwen_image.pipeline_qwen_image import calculate_
 from vllm_omni.diffusion.models.qwen_image.qwen_image_transformer import (
     QwenImageTransformer2DModel,
 )
+from vllm_omni.diffusion.models.qwen_image.rope_utils import txt_seq_lens_from_embeds
 from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.utils.prompt_utils import (
@@ -840,10 +841,8 @@ class QwenImageEditPipeline(nn.Module, SupportImageInput, QwenImageCFGParallelMi
         if self.attention_kwargs is None:
             self._attention_kwargs = {}
 
-        txt_seq_lens = prompt_embeds_mask.sum(dim=1).tolist() if prompt_embeds_mask is not None else None
-        negative_txt_seq_lens = (
-            negative_prompt_embeds_mask.sum(dim=1).tolist() if negative_prompt_embeds_mask is not None else None
-        )
+        txt_seq_lens = txt_seq_lens_from_embeds(prompt_embeds)
+        negative_txt_seq_lens = txt_seq_lens_from_embeds(negative_prompt_embeds)
 
         latents = self.diffuse(
             prompt_embeds,

vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit_plus.py

@@ -40,6 +40,7 @@ from vllm_omni.diffusion.models.qwen_image.pipeline_qwen_image_edit import (
 from vllm_omni.diffusion.models.qwen_image.qwen_image_transformer import (
     QwenImageTransformer2DModel,
 )
+from vllm_omni.diffusion.models.qwen_image.rope_utils import txt_seq_lens_from_embeds
 from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.utils.prompt_utils import (
@@ -812,10 +813,8 @@ class QwenImageEditPlusPipeline(
         if self.attention_kwargs is None:
             self._attention_kwargs = {}
 
-        txt_seq_lens = prompt_embeds_mask.sum(dim=1).tolist() if prompt_embeds_mask is not None else None
-        negative_txt_seq_lens = (
-            negative_prompt_embeds_mask.sum(dim=1).tolist() if negative_prompt_embeds_mask is not None else None
-        )
+        txt_seq_lens = txt_seq_lens_from_embeds(prompt_embeds)
+        negative_txt_seq_lens = txt_seq_lens_from_embeds(negative_prompt_embeds)
 
         latents = self.diffuse(
             prompt_embeds,

vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_layered.py

@@ -35,6 +35,7 @@ from vllm_omni.diffusion.models.qwen_image.cfg_parallel import (
 from vllm_omni.diffusion.models.qwen_image.qwen_image_transformer import (
     QwenImageTransformer2DModel,
 )
+from vllm_omni.diffusion.models.qwen_image.rope_utils import txt_seq_lens_from_embeds
 from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.utils.prompt_utils import (
@@ -848,10 +849,8 @@ the image\n<|vision_start|><|image_pad|><|vision_end|><|im_end|>\n<|im_start|>as
         if self.attention_kwargs is None:
             self._attention_kwargs = {}
 
-        txt_seq_lens = prompt_embeds_mask.sum(dim=1).tolist() if prompt_embeds_mask is not None else None
-        negative_txt_seq_lens = (
-            negative_prompt_embeds_mask.sum(dim=1).tolist() if negative_prompt_embeds_mask is not None else None
-        )
+        txt_seq_lens = txt_seq_lens_from_embeds(prompt_embeds)
+        negative_txt_seq_lens = txt_seq_lens_from_embeds(negative_prompt_embeds)
         is_rgb = torch.tensor([0] * batch_size).to(device=self.device, dtype=torch.long)
 
         latents = self.diffuse(

vllm_omni/diffusion/models/qwen_image/rope_utils.py

@@ -0,0 +1,26 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+import torch
+
+
+def txt_seq_lens_from_embeds(prompt_embeds: torch.Tensor | None) -> list[int] | None:
+    """Return per-row RoPE text lengths from padded encoder hidden-state width.
+
+    Diffusers derives Qwen-Image text RoPE length from
+    ``encoder_hidden_states.shape[1]``, not from the count of valid (non-padding)
+    tokens in the attention mask. Callers that right-pad prompt embeddings for
+    batching must pass the padded width here so ``QwenEmbedRope`` builds a table
+    long enough for every text position.
+    """
+    if prompt_embeds is None:
+        return None
+    if prompt_embeds.ndim == 2:
+        prompt_embeds = prompt_embeds.unsqueeze(0)
+    if prompt_embeds.ndim != 3:
+        raise ValueError(f"prompt_embeds must be 2D or 3D, got shape={tuple(prompt_embeds.shape)}")
+    seq_len = int(prompt_embeds.shape[1])
+    batch_size = int(prompt_embeds.shape[0])
+    return [seq_len] * batch_size

vllm_omni/diffusion/worker/input_batch.py

@@ -253,8 +253,15 @@ def _prepare_request_prompt_field(
                 f"{embeds_attr} for request {state.request_id} requires seq_len "
                 f"{max_actual_seq_len}, got target {target_seq_len}."
             )
-        return embeds[:, :target_seq_len], None if mask is None else mask[:, :target_seq_len]
+        embeds = embeds[:, :target_seq_len]
+        mask = None if mask is None else mask[:, :target_seq_len]
+        setattr(state, embeds_attr, embeds)
+        if mask is not None:
+            setattr(state, mask_attr, mask)
+        setattr(state, seq_lens_attr, [int(embeds.shape[1])] * int(embeds.shape[0]))
+        return embeds, mask
 
+    setattr(state, seq_lens_attr, [int(embeds.shape[1])] * int(embeds.shape[0]))
     return embeds, mask