From 5f5b77aa9cef5033c015eb40c9c228c5f68244c6 Mon Sep 17 00:00:00 2001 From: Samit <285365963@qq.com> Date: Tue, 16 Jun 2026 19:09:25 +0800 Subject: [PATCH] [BugFix][Qwen-Image] align txt_seq_lens RoPE width with padded embeds (#4474) Signed-off-by: samithuang <285365963@qq.com> --- ...test_qwen_image_rope_table_failure_case.py | 44 +++++++++++++++++++ .../qwen_image/test_qwen_image_rope_utils.py | 30 +++++++++++++ .../test_input_batch_txt_seq_lens.py | 31 +++++++++++++ .../models/qwen_image/pipeline_qwen_image.py | 7 ++- .../qwen_image/pipeline_qwen_image_edit.py | 7 ++- .../pipeline_qwen_image_edit_plus.py | 7 ++- .../qwen_image/pipeline_qwen_image_layered.py | 7 ++- .../diffusion/models/qwen_image/rope_utils.py | 26 +++++++++++ vllm_omni/diffusion/worker/input_batch.py | 9 +++- 9 files changed, 151 insertions(+), 17 deletions(-) create mode 100644 tests/diffusion/models/qwen_image/test_qwen_image_rope_table_failure_case.py create mode 100644 tests/diffusion/models/qwen_image/test_qwen_image_rope_utils.py create mode 100644 tests/diffusion/test_input_batch_txt_seq_lens.py create mode 100644 vllm_omni/diffusion/models/qwen_image/rope_utils.py diff --git a/tests/diffusion/models/qwen_image/test_qwen_image_rope_table_failure_case.py b/tests/diffusion/models/qwen_image/test_qwen_image_rope_table_failure_case.py new file mode 100644 index 000000000..2218db8d6 --- /dev/null +++ b/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 diff --git a/tests/diffusion/models/qwen_image/test_qwen_image_rope_utils.py b/tests/diffusion/models/qwen_image/test_qwen_image_rope_utils.py new file mode 100644 index 000000000..336ae2ed4 --- /dev/null +++ b/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)) diff --git a/tests/diffusion/test_input_batch_txt_seq_lens.py b/tests/diffusion/test_input_batch_txt_seq_lens.py new file mode 100644 index 000000000..315f0f703 --- /dev/null +++ b/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] diff --git a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py index cdc6de998..655641fb4 100644 --- a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py +++ b/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, diff --git a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit.py b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit.py index 35022e5eb..a83aa83ad 100644 --- a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit.py +++ b/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, diff --git a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit_plus.py b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit_plus.py index e053f85dc..7d2170e09 100644 --- a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit_plus.py +++ b/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, diff --git a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_layered.py b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_layered.py index a1fcff60d..19711cd8c 100644 --- a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_layered.py +++ b/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( diff --git a/vllm_omni/diffusion/models/qwen_image/rope_utils.py b/vllm_omni/diffusion/models/qwen_image/rope_utils.py new file mode 100644 index 000000000..68dbccb59 --- /dev/null +++ b/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 diff --git a/vllm_omni/diffusion/worker/input_batch.py b/vllm_omni/diffusion/worker/input_batch.py index ef502d0d4..375546774 100644 --- a/vllm_omni/diffusion/worker/input_batch.py +++ b/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