From 27679c3339394dfe3fa834f169660c792624db08 Mon Sep 17 00:00:00 2001 From: Mick Date: Sat, 10 Oct 2026 22:16:27 +0800 Subject: [PATCH] [diffusion] fix: Qwen-Image BCG prompt padding misses every graph and attends the pad The Qwen-Image padder bucketed `freqs_cis` but not the complex RoPE cache `freqs_complex`, whose text half keeps the prompt length, so every request missed the captured graphs. A miss runs the padded kwargs eagerly, and the padder had also bucketed `txt_seq_lens`, which the eager forward uses to build the varlen text ranges: the zero-padded text rows were attended as real tokens, changing the image. Pad `freqs_complex` with `freqs_cis` and leave `txt_seq_lens` unset on padded inputs so the mask alone marks the valid text. Remove the now-unused `bucket_txt_seq_lens` helper. Co-Authored-By: Claude Opus 5.5 --- .../model_padders/qwen_image.py | 26 +++++++++---------- .../breakable_cuda_graph/prompt_padding.py | 14 ---------- .../test/unit/test_diffusion_bcg_padding.py | 18 ++++++++++--- 3 files changed, 26 insertions(+), 32 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/breakable_cuda_graph/model_padders/qwen_image.py b/python/sglang/multimodal_gen/runtime/breakable_cuda_graph/model_padders/qwen_image.py index e5b0efb383f0..c2e5fad1c424 100644 --- a/python/sglang/multimodal_gen/runtime/breakable_cuda_graph/model_padders/qwen_image.py +++ b/python/sglang/multimodal_gen/runtime/breakable_cuda_graph/model_padders/qwen_image.py @@ -81,21 +81,19 @@ def pad_qwen_prompt_kwargs( out["encoder_attention_mask"], dim=1, source=seq, target=bucket ) - freqs_cis = out.get("freqs_cis") - if isinstance(freqs_cis, tuple) and len(freqs_cis) == 2: - img_cache, txt_cache = freqs_cis - txt_cache = bcg_utils.pad_nested_dim( - txt_cache, dim=0, source=seq, target=bucket - ) - out["freqs_cis"] = (img_cache, txt_cache) - elif isinstance(freqs_cis, list) and len(freqs_cis) == 2: - img_cache, txt_cache = freqs_cis - txt_cache = bcg_utils.pad_nested_dim( - txt_cache, dim=0, source=seq, target=bucket - ) - out["freqs_cis"] = [img_cache, txt_cache] + # the real and the complex RoPE caches both carry the text length on dim 0 + for key in ("freqs_cis", "freqs_complex"): + freqs = out.get(key) + if isinstance(freqs, (tuple, list)) and len(freqs) == 2: + img_cache, txt_cache = freqs + txt_cache = bcg_utils.pad_nested_dim( + txt_cache, dim=0, source=seq, target=bucket + ) + out[key] = type(freqs)((img_cache, txt_cache)) - out["txt_seq_lens"] = bcg_utils.bucket_txt_seq_lens(out.get("txt_seq_lens"), bucket) + # the mask alone marks the valid text: a bucketed length would make the + # eager forward of a missed graph attend the pad rows + out["txt_seq_lens"] = None return out diff --git a/python/sglang/multimodal_gen/runtime/breakable_cuda_graph/prompt_padding.py b/python/sglang/multimodal_gen/runtime/breakable_cuda_graph/prompt_padding.py index b9721d41dd0f..345437a0e936 100644 --- a/python/sglang/multimodal_gen/runtime/breakable_cuda_graph/prompt_padding.py +++ b/python/sglang/multimodal_gen/runtime/breakable_cuda_graph/prompt_padding.py @@ -138,20 +138,6 @@ def pad_nested_dim( return obj -def bucket_txt_seq_lens(txt_seq_lens: Any, bucket: int) -> Any: - if txt_seq_lens is None: - return txt_seq_lens - if torch.is_tensor(txt_seq_lens): - return torch.full_like(txt_seq_lens, bucket) - if isinstance(txt_seq_lens, list): - return [bucket_txt_seq_lens(seq_len, bucket) for seq_len in txt_seq_lens] - if isinstance(txt_seq_lens, tuple): - return tuple(bucket_txt_seq_lens(seq_len, bucket) for seq_len in txt_seq_lens) - if isinstance(txt_seq_lens, int): - return bucket - return txt_seq_lens - - def prompt_seq_and_dim(call_kwargs: dict) -> tuple[int, int] | None: """Return ``(text_seq_len, seq_dim)`` inferred from the prompt embeddings or a prompt mask, or ``None`` when no text conditioning is present.""" diff --git a/python/sglang/multimodal_gen/test/unit/test_diffusion_bcg_padding.py b/python/sglang/multimodal_gen/test/unit/test_diffusion_bcg_padding.py index e0d051ca8866..b803b2d6f5b6 100644 --- a/python/sglang/multimodal_gen/test/unit/test_diffusion_bcg_padding.py +++ b/python/sglang/multimodal_gen/test/unit/test_diffusion_bcg_padding.py @@ -231,6 +231,10 @@ def _qwen_kwargs(self, seq_len: int, *, fill: float = 1.0): torch.zeros(4096, 128, dtype=torch.float32), torch.ones(seq_len, 128, dtype=torch.float32), ), + "freqs_complex": ( + torch.zeros(4096, 64, dtype=torch.complex64), + torch.ones(seq_len, 64, dtype=torch.complex64), + ), "img_shapes": [[(1, 64, 64)]], } @@ -251,8 +255,14 @@ def test_qwen_prompt_lengths_share_bucket_signature(self): self.assertTrue(longer["encoder_hidden_states_mask"][0, :47].all()) self.assertFalse(longer["encoder_hidden_states_mask"][0, 47:].any()) self.assertEqual(short["freqs_cis"][1].shape, (256, 128)) - self.assertEqual(short["txt_seq_lens"], [256]) - self.assertEqual(longer["txt_seq_lens"], [256]) + # the complex RoPE cache must follow the bucket too, or every prompt + # length misses the captured graph + self.assertEqual(short["freqs_complex"][1].shape, (256, 64)) + self.assertFalse(short["freqs_complex"][1][19:].any()) + # a bucketed length would tell a missed graph's eager forward to attend + # the pad rows; the mask alone marks the valid text + self.assertIsNone(short["txt_seq_lens"]) + self.assertIsNone(longer["txt_seq_lens"]) self.assertEqual(_signature_kwargs(short), _signature_kwargs(longer)) def test_qwen_prompt_content_changes_do_not_change_signature(self): @@ -331,8 +341,8 @@ def kwargs(valid_len: int): ) self.assertEqual(first["encoder_hidden_states"][0].shape[1], 64) - self.assertEqual(first["txt_seq_lens"], [64]) - self.assertEqual(second["txt_seq_lens"], [64]) + self.assertIsNone(first["txt_seq_lens"]) + self.assertIsNone(second["txt_seq_lens"]) self.assertTrue(first["encoder_hidden_states_mask"][0, :19].all()) self.assertFalse(first["encoder_hidden_states_mask"][0, 19:].any()) self.assertTrue(second["encoder_hidden_states_mask"][0, :47].all())