From 28b43bf693a8d93d9241ef0a71a11eb245f6f3de Mon Sep 17 00:00:00 2001 From: Mick Date: Fri, 7 Aug 2026 17:55:15 +0800 Subject: [PATCH] [diffusion] perf: build qwen's masked varlen metadata host-side (#33954) Co-authored-by: Claude Fable 5 --- .../runtime/models/dits/qwen_image.py | 21 +++++++++++ .../test/unit/test_varlen_meta_host_build.py | 36 +++++++++++++++++++ 2 files changed, 57 insertions(+) create mode 100644 python/sglang/multimodal_gen/test/unit/test_varlen_meta_host_build.py diff --git a/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py b/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py index cbb2ea280..18107c1b6 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py @@ -39,6 +39,7 @@ from sglang.multimodal_gen.runtime.layers.attention import ( DynamicVarlenMaskMeta, USPAttention, build_varlen_mask_meta, + build_varlen_mask_meta_from_ranges, ) from sglang.multimodal_gen.runtime.layers.elementwise import MulAdd from sglang.multimodal_gen.runtime.layers.fused_scale_shift_gate import ( @@ -1565,6 +1566,26 @@ class QwenImageTransformer2DModel(CachableDiT, LayerwiseOffloadableModuleMixin): # once, so build varlen metadata replay-locally from the current # static mask instead of closing over stale cu_seqlens/indices. block_attention_kwargs["attn_mask_meta"] = DynamicVarlenMaskMeta() + elif ( + txt_seq_lens is not None + and len(txt_seq_lens) == batch_size + and all(0 <= n <= encoder_hidden_states.shape[1] for n in txt_seq_lens) + ): + # txt_seq_lens already carries each row's valid text prefix + # (the mask is built from it), so the varlen metadata can be + # assembled host-side; the mask-based builder costs a GPU + # nonzero plus a device sync on every denoising step. + txt_len = encoder_hidden_states.shape[1] + block_attention_kwargs["attn_mask_meta"] = ( + build_varlen_mask_meta_from_ranges( + [ + [(0, int(n)), (txt_len, txt_len + image_seq_len)] + for n in txt_seq_lens + ], + max_seqlen=txt_len + image_seq_len, + device=hidden_states.device, + ) + ) else: # Precompute varlen metadata once per request so every block # reuses the same cu_seqlens / indices instead of rebuilding. diff --git a/python/sglang/multimodal_gen/test/unit/test_varlen_meta_host_build.py b/python/sglang/multimodal_gen/test/unit/test_varlen_meta_host_build.py new file mode 100644 index 000000000..203082673 --- /dev/null +++ b/python/sglang/multimodal_gen/test/unit/test_varlen_meta_host_build.py @@ -0,0 +1,36 @@ +"""Host-built varlen metadata must match the mask-based (nonzero) builder.""" + +import unittest + +import torch + +from sglang.multimodal_gen.runtime.layers.attention.layer import ( + build_varlen_mask_meta, + build_varlen_mask_meta_from_ranges, +) + + +class TestHostVarlenMetaEquivalence(unittest.TestCase): + def test_prefix_text_plus_full_image_matches_nonzero_builder(self): + txt_len, img_len = 7, 5 + txt_seq_lens = [3, 7, 0] + bs = len(txt_seq_lens) + mask = torch.zeros(bs, txt_len + img_len, dtype=torch.bool) + for row, n in enumerate(txt_seq_lens): + mask[row, :n] = True + mask[row, txt_len:] = True + + ref = build_varlen_mask_meta(mask) + host = build_varlen_mask_meta_from_ranges( + [[(0, n), (txt_len, txt_len + img_len)] for n in txt_seq_lens], + max_seqlen=txt_len + img_len, + device=mask.device, + ) + + for key in ("cu_seqlens", "indices", "inv_indices"): + torch.testing.assert_close(host[key], ref[key], rtol=0, atol=0) + self.assertEqual(host["max_seqlen"], ref["max_seqlen"]) + + +if __name__ == "__main__": + unittest.main()