[Diffusion] Optimize Qwen-Image TP collectives and attention (#36680)

Co-authored-by: Mick <mickjagger19@icloud.com>
Co-authored-by: Cursor <cursoragent@cursor.com>

test/registered/kernels/ops/diffusion/test_layout.py

@@ -31,6 +31,7 @@ from sglang.kernels.ops.diffusion import (
 from sglang.kernels.ops.diffusion import (
     fused_causal_conv3d_cat_pad_cuda,
     fused_pack_qkv,
+    fused_pack_segmented_qkv,
     fused_scatter_to_padded,
     pack_qkv_destination_major,
     usp_merge_heads,
@@ -211,6 +212,28 @@ def test_varlen_pack_matches_index_select(dtype, shape):
         assert torch.equal(got, want)
 
 
+@pytest.mark.parametrize("dtype", VARLEN_DTYPES)
+@pytest.mark.parametrize("shape", VARLEN_SHAPES, ids=lambda s: s[0])
+def test_varlen_segmented_pack_matches_materialized_joint(dtype, shape):
+    _, bs, s_txt, s_img, num_heads, head_dim, valid_txt_lens = shape
+    torch.manual_seed(42)
+    indices, _ = _build_meta(_build_mask(bs, s_txt, s_img, valid_txt_lens))
+    txt_qkv = tuple(
+        torch.randn(bs, s_txt, num_heads, head_dim, dtype=dtype, device=DEVICE)
+        for _ in range(3)
+    )
+    img_qkv = tuple(
+        torch.randn(bs, s_img, num_heads, head_dim, dtype=dtype, device=DEVICE)
+        for _ in range(3)
+    )
+
+    got = fused_pack_segmented_qkv(*txt_qkv, *img_qkv, indices)
+    for actual, txt, img in zip(got, txt_qkv, img_qkv, strict=True):
+        joint = torch.cat([txt, img], dim=1)
+        expected = joint.flatten(0, 1).index_select(0, indices)
+        assert torch.equal(actual, expected)
+
+
 @pytest.mark.parametrize("dtype", VARLEN_DTYPES)
 @pytest.mark.parametrize("shape", VARLEN_SHAPES, ids=lambda s: s[0])
 def test_varlen_scatter_matches_index_copy(dtype, shape):
@@ -344,6 +367,45 @@ def _varlen_path(q, k, v, key_mask, softmax_scale):
     return fused_scatter_to_padded(out_unpad, meta["inv_indices"], bs, seq)
 
 
+@pytest.mark.parametrize("dtype", VARLEN_DTYPES)
+def test_fa_dense_scheduler_matches_single_sequence_varlen(dtype):
+    torch.manual_seed(7)
+    batch_size, seq, num_heads, head_dim = 1, 256, 4, 128
+    q, k, v = (
+        torch.randn(
+            batch_size,
+            seq,
+            num_heads,
+            head_dim,
+            dtype=dtype,
+            device=DEVICE,
+        )
+        for _ in range(3)
+    )
+    cu_seqlens = torch.tensor([0, seq], dtype=torch.int32, device=DEVICE)
+    kwargs = dict(
+        max_seqlen_q=seq,
+        max_seqlen_k=seq,
+        softmax_scale=head_dim**-0.5,
+        causal=False,
+        ver=_fa_backend.fa_ver,
+    )
+    try:
+        varlen = flash_attn_varlen_func(
+            q.flatten(0, 1),
+            k.flatten(0, 1),
+            v.flatten(0, 1),
+            cu_seqlens,
+            cu_seqlens,
+            **kwargs,
+        ).view_as(q)
+        dense = flash_attn_varlen_func(q, k, v, None, None, **kwargs)
+    except ImportError as exc:  # pragma: no cover - image-dependent
+        pytest.skip(f"FlashAttention unavailable: {exc}")
+
+    torch.testing.assert_close(dense, varlen, rtol=1e-3, atol=1e-3)
+
+
 @pytest.mark.parametrize("dtype", VARLEN_DTYPES)
 @pytest.mark.parametrize("shape", VARLEN_SHAPES, ids=lambda s: s[0])
 def test_varlen_path_matches_sdpa_on_valid_rows(dtype, shape):

test/registered/kernels/ops/diffusion/test_model_fast_paths.py

@@ -48,12 +48,15 @@ from sglang.kernels.ops.diffusion import (
     mark_fused_ln_modulate_site,
     mark_hunyuan_qknorm_site,
     mark_ltx2_rms_norm_modulate_site,
+    mark_qwen_image_added_qkv_site,
     mount_fused_ln_modulate,
     mount_hunyuan_qknorm,
     mount_ltx2_rms_norm_modulate,
+    mount_qwen_image_added_qkv,
     try_flux2_token_cat_nvfp4,
     unmount_hunyuan_qknorm,
     unmount_ltx2_rms_norm_modulate,
+    unmount_qwen_image_added_qkv,
     wan_rmsnorm_silu,
 )
 from sglang.kernels.ops.diffusion.common.platform import is_cuda
@@ -98,9 +101,11 @@ from sglang.multimodal_gen.runtime.models.dits.longcat_image import (
 )
 from sglang.multimodal_gen.runtime.models.dits.ltx_2 import _ltx2_rms_norm_modulate
 from sglang.multimodal_gen.runtime.models.dits.qwen_image import (
+    QwenImageCrossAttention,
     QwenImageTransformerBlock,
     _qwen_modulation_cache_key,
     _qwen_norm_out,
+    _split_unquantized_merged_linear,
 )
 from sglang.multimodal_gen.runtime.models.dits.sana import (
     _eager_ln_modulate as _sana_eager_ln_modulate,
@@ -412,6 +417,64 @@ class TestFlux2EagerFusions(CustomTestCase):
 # -------------------------------------------------------------------------
 
 
+class _PackedAddedQKV(nn.Module):
+    def __init__(self, dim: int):
+        super().__init__()
+        self.output_partition_sizes = [dim, dim, dim]
+        self.quant_config = None
+        self.weight = nn.Parameter(
+            torch.randn(3 * dim, dim, device="cuda", dtype=torch.bfloat16),
+            requires_grad=False,
+        )
+        self.bias = nn.Parameter(
+            torch.randn(3 * dim, device="cuda", dtype=torch.bfloat16),
+            requires_grad=False,
+        )
+        self.calls = 0
+
+    def forward(self, x):
+        self.calls += 1
+        return F.linear(x, self.weight, self.bias), None
+
+
+def test_qwen_added_qkv_lossless_uses_three_reference_gemms():
+    torch.manual_seed(20260831)
+    dim = 64
+    x = torch.randn(1, 17, dim, device="cuda", dtype=torch.bfloat16)
+    packed = _PackedAddedQKV(dim)
+
+    attention = QwenImageCrossAttention.__new__(QwenImageCrossAttention)
+    nn.Module.__init__(attention)
+    attention.use_fused_added_qkv = True
+    attention._unquantized_added_qkv_is_packed = True
+    attention.to_added_qkv = packed
+    mark_qwen_image_added_qkv_site(attention)
+
+    expected_lossless = _split_unquantized_merged_linear(packed, x)
+    actual_lossless = attention._get_added_qkv_projections(x)
+    assert packed.calls == 0
+    assert all(
+        torch.equal(actual, expected)
+        for actual, expected in zip(actual_lossless, expected_lossless)
+    )
+
+    assert mount_qwen_image_added_qkv(attention)
+    actual_high = attention._get_added_qkv_projections(x)
+    expected_high = tuple(
+        tensor.contiguous()
+        for tensor in F.linear(x, packed.weight, packed.bias).chunk(3, dim=-1)
+    )
+    assert packed.calls == 1
+    assert all(
+        torch.equal(actual, expected)
+        for actual, expected in zip(actual_high, expected_high)
+    )
+
+    unmount_qwen_image_added_qkv(attention)
+    attention._get_added_qkv_projections(x)
+    assert packed.calls == 1
+
+
 class _CountingProjection(nn.Module):
     def __init__(self, offset: float):
         super().__init__()

test/registered/kernels/ops/diffusion/test_sites.py

@@ -28,6 +28,7 @@ import torch.nn.functional as F
 import sglang.kernels.ops.diffusion.sites.fused_gate_rmsnorm_site as gate_rmsnorm
 import sglang.kernels.ops.diffusion.sites.fused_linear_gelu_site as linear_gelu
 import sglang.kernels.ops.diffusion.sites.lingbot_video_rmsnorm_site as lingbot_video_rmsnorm
+import sglang.kernels.ops.diffusion.sites.qwen_image_added_qkv_site as qwen_image_added_qkv
 import sglang.kernels.ops.diffusion.sites.sana_video_linear_attention_site as sana_video_linear_attention
 from sglang.kernels.ops.diffusion import (
     BitExactFusionGate,
@@ -91,6 +92,20 @@ def test_quality_gate_rejection_is_all_or_nothing():
     assert not fusion.mount(nn.Module())
 
 
+def test_qwen_image_added_qkv_site_is_request_scoped():
+    site = nn.Module()
+    site.to_added_qkv = nn.Module()
+    site.to_added_qkv.quant_config = None
+    site.to_added_qkv.output_partition_sizes = [8, 8, 8]
+    qwen_image_added_qkv.mark_qwen_image_added_qkv_site(site)
+
+    assert not qwen_image_added_qkv.qwen_image_added_qkv_active(site)
+    assert qwen_image_added_qkv.mount_qwen_image_added_qkv(site)
+    assert qwen_image_added_qkv.qwen_image_added_qkv_active(site)
+    qwen_image_added_qkv.unmount_qwen_image_added_qkv(site)
+    assert not qwen_image_added_qkv.qwen_image_added_qkv_active(site)
+
+
 # -------------------------------------------------------------------------
 # BitExactFusionGate protocol (CPU)
 # -------------------------------------------------------------------------