[diffusion] fix: Qwen-Image BCG prompt padding misses every graph and attends the pad

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
 
 

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."""

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())