minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 13 Packages Projects Releases Wiki Activity

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

Browse Source 
Co-authored-by: Mick <mickjagger19@icloud.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
Xiaoyu Zhang
2026-09-01 10:28:37 +08:00
committed by GitHub
co-authored by Mick Cursor
parent 562b661e0e
commit 71cee04ebe
19 changed files with 793 additions and 61 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/kernels/benchmark/diffusion/bench_varlen_segmented_pack.py
+95
View File
@@ -0,0 +1,95 @@
import torch
from sglang.kernels.jit.benchmark import marker
from sglang.kernels.ops.diffusion import fused_pack_qkv, fused_pack_segmented_qkv
from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(
est_time=12, stage="base-b-kernel-benchmark", runner_config="1-gpu-large"
)
DEVICE = "cuda"
DTYPE = torch.bfloat16
def _materialized_pack(q_prefix, k_prefix, v_prefix, q_main, k_main, v_main, indices):
return fused_pack_qkv(
torch.cat([q_prefix, q_main], dim=1),
torch.cat([k_prefix, k_main], dim=1),
torch.cat([v_prefix, v_main], dim=1),
indices,
)
def _segmented_pack(q_prefix, k_prefix, v_prefix, q_main, k_main, v_main, indices):
return fused_pack_segmented_qkv(
q_prefix, k_prefix, v_prefix, q_main, k_main, v_main, indices
)
@marker.parametrize(
"batch,prefix_rows,main_rows,heads,head_dim,valid_prefix_rows",
[
(1, 64, 4096, 12, 128, 24),
(2, 64, 1024, 8, 128, 40),
],
ci_vals=[(1, 64, 4096, 12, 128, 24)],
)
@marker.benchmark("provider", ["materialized", "segmented"])
def benchmark(
batch: int,
prefix_rows: int,
main_rows: int,
heads: int,
head_dim: int,
valid_prefix_rows: int,
provider: str,
) -> marker.BenchResult:
generator = torch.Generator(device=DEVICE).manual_seed(42)
prefixes = tuple(
torch.randn(
batch,
prefix_rows,
heads,
head_dim,
dtype=DTYPE,
device=DEVICE,
generator=generator,
)
for _ in range(3)
)
mains = tuple(
torch.randn(
batch,
main_rows,
heads,
head_dim,
dtype=DTYPE,
device=DEVICE,
generator=generator,
)
for _ in range(3)
)
mask = torch.zeros(batch, prefix_rows + main_rows, dtype=torch.bool, device=DEVICE)
mask[:, :valid_prefix_rows] = True
mask[:, prefix_rows:] = True
indices = mask.flatten().nonzero(as_tuple=False).flatten()
args = (*prefixes, *mains, indices)
expected = _materialized_pack(*args)
actual = _segmented_pack(*args)
assert all(
torch.equal(got, want) for got, want in zip(actual, expected, strict=True)
)
fn = _materialized_pack if provider == "materialized" else _segmented_pack
return marker.do_bench(
fn,
input_args=args,
graph_clone_args=tuple(range(len(args))),
disable_log_bandwidth=True,
)
if __name__ == "__main__":
benchmark.run()
test/registered/kernels/ops/diffusion/test_layout.py
+62
View File
@@ -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
+63
View File
@@ -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
+15
View File
@@ -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)
# -------------------------------------------------------------------------