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

[Diffusion] Optimize Qwen-Image-Edit attention on Hopper (#38584)

Browse Source 
Co-authored-by: Mick Qian <mickqian@users.noreply.github.com>
This commit is contained in:
Xiaoyu Zhang
2026-09-13 09:15:01 +08:00
committed by GitHub
co-authored by Mick Qian
parent fa663e7297
commit 3e035a3513
9 changed files with 362 additions and 48 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/kernels/ops/diffusion/test_model_fast_paths.py
+87
View File
@@ -66,6 +66,7 @@ from sglang.kernels.ops.diffusion.common.platform import is_cuda
from sglang.multimodal_gen.configs.models.vaes.stablediffusion3 import (
StableDiffusion3VAEConfig,
)
from sglang.multimodal_gen.runtime.layers.attention.backends import sdpa as sdpa_backend
from sglang.multimodal_gen.runtime.layers.layernorm import (
RMSNorm,
RMSNormNoWeight,
@@ -157,6 +158,92 @@ def _seed_cuda():
torch.cuda.manual_seed(0)
@pytest.mark.skipif(
not is_cuda()
or torch.cuda.get_device_capability()[0] != 9
or sdpa_backend.torch_varlen_attn is None,
reason="packed Flash SDPA requires Hopper and PyTorch varlen attention",
)
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
@pytest.mark.parametrize("causal", [False, True])
@pytest.mark.parametrize(
"lengths,head_dim",
[
(([64] * 7 + [32]) * 11 + [16] * 7 + [8], 80),
([0, 64, 64, 0, 32, 0], 80),
([32] * 8, 64),
([128, 64, 128], 128),
([2048], 80),
],
)
def test_packed_sdpa_uses_native_varlen_without_changing_values(
dtype, causal, lengths, head_dim
):
bounds = [0]
for length in lengths:
bounds.append(bounds[-1] + length)
packed = torch.randn(bounds[-1], 3, 16, head_dim, device="cuda", dtype=dtype)
q, k, v = packed.unbind(1)
q, k = q.contiguous(), k.contiguous()
attention = sdpa_backend.SDPAImpl(
16, head_dim, causal=causal, softmax_scale=head_dim**-0.5
)
cu = torch.tensor(bounds, device="cuda", dtype=torch.int32)
with (
torch.no_grad(),
torch.nn.attention.sdpa_kernel(torch.nn.attention.SDPBackend.FLASH_ATTENTION),
):
expected = torch.cat(
[
attention.forward(q[a:b][None], k[a:b][None], v[a:b][None], None)[0]
for a, b in zip(bounds[:-1], bounds[1:])
if a != b
]
)
with patch.object(attention, "forward", wraps=attention.forward) as forward:
actual = attention.forward_varlen(
q,
k,
v,
cu_seqlens=cu,
max_seqlen=max(lengths),
cu_seqlens_host=tuple(bounds),
)
assert forward.call_count == (1 if len(lengths) == 1 else 0)
assert torch.equal(actual, expected)
@pytest.mark.skipif(not is_cuda(), reason="CUDA packed SDPA fallback test")
def test_packed_sdpa_preserves_training_dropout_and_missing_api_fallbacks():
attention = sdpa_backend.SDPAImpl(4, 64, causal=True, softmax_scale=64**-0.5)
qkv = torch.randn(
64, 3, 4, 64, device="cuda", dtype=torch.bfloat16, requires_grad=True
)
q, k, v = qkv.unbind(1)
cu = torch.tensor([0, 32, 64], device="cuda", dtype=torch.int32)
kwargs = dict(cu_seqlens=cu, max_seqlen=32, cu_seqlens_host=(0, 32, 64))
with patch.object(attention, "forward", wraps=attention.forward) as forward:
result = attention.forward_varlen(q, k, v, **kwargs)
assert forward.call_count == 2
result.float().square().mean().backward()
assert qkv.grad is not None and torch.isfinite(qkv.grad).all()
attention.dropout = 0.1
with (
torch.no_grad(),
patch.object(attention, "forward", wraps=attention.forward) as forward,
):
attention.forward_varlen(q, k, v, **kwargs)
assert forward.call_count == 2
attention.dropout = 0.0
with (
torch.no_grad(),
patch.object(sdpa_backend, "torch_varlen_attn", None),
patch.object(attention, "forward", wraps=attention.forward) as forward,
):
attention.forward_varlen(q, k, v, **kwargs)
assert forward.call_count == 2
def test_bitexact_norm_guards_follow_platform():
# Runs on both lanes, with shapes inside every guard's contract so only the
# platform decides: engaged on CUDA, rejected on ROCm. A fatal LLVM error
test/registered/kernels/ops/diffusion/test_qwen_qkv_epilogue.py
+156 -23
View File
@@ -4,6 +4,7 @@ from unittest.mock import patch
import pytest
import torch
import sglang.multimodal_gen.runtime.models.dits.qwen_image as qwen_image
from sglang.kernels.ops.diffusion import try_fused_qwen_qkv_epilogue
from sglang.multimodal_gen.runtime.layers.layernorm import (
RMSNorm,
@@ -12,10 +13,11 @@ from sglang.multimodal_gen.runtime.layers.layernorm import (
from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(est_time=15, stage="base-b-kernel-unit", runner_config="4-gpu-b200")
register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="1-gpu-large")
pytestmark = pytest.mark.skipif(
not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] < 10,
reason="Qwen-Image QKV epilogue requires SM100+",
not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] < 9,
reason="Qwen-Image QKV epilogue requires SM90+",
)
@@ -24,11 +26,12 @@ def _seed_cuda():
torch.cuda.manual_seed(0)
def test_qwen_qkv_epilogue_is_bit_exact():
heads = 4
@pytest.mark.parametrize(
"img_tokens,txt_tokens,heads",
[(17, 7, 1), (17, 7, 3), (17, 7, 4), (8152, 1365, 24)],
)
def test_qwen_qkv_epilogue_is_bit_exact(img_tokens, txt_tokens, heads):
head_dim = 128
img_tokens = 17
txt_tokens = 7
img_qkv = [
torch.randn(
1,
@@ -55,6 +58,9 @@ def test_qwen_qkv_epilogue_is_bit_exact():
RMSNorm(head_dim, eps=1e-6).to(device="cuda", dtype=torch.bfloat16)
for _ in range(4)
]
with torch.no_grad():
for norm in norms:
norm.weight.copy_(torch.randn_like(norm.weight))
def cache(tokens):
angles = torch.randn(tokens, head_dim // 2, device="cuda")
@@ -120,23 +126,30 @@ def test_qwen_qkv_epilogue_is_bit_exact():
]
assert all(not tensor.is_contiguous() for tensor in (*img_views, *txt_views))
packed_actual = try_fused_qwen_qkv_epilogue(
*img_views,
*txt_views,
norms[0].weight,
norms[1].weight,
norms[2].weight,
norms[3].weight,
img_cache,
txt_cache,
1e-6,
1e-6,
)
assert packed_actual is not None
assert all(
torch.equal(result, reference)
for result, reference in zip(packed_actual, expected)
)
# Unquantized image projections and packed text projections can use
# different token strides; each family also works when both are packed.
for img_inputs, txt_inputs in (
(img_views, txt_views),
(img_qkv, txt_views),
(img_views, txt_qkv),
):
packed_actual = try_fused_qwen_qkv_epilogue(
*img_inputs,
*txt_inputs,
norms[0].weight,
norms[1].weight,
norms[2].weight,
norms[3].weight,
img_cache,
txt_cache,
1e-6,
1e-6,
)
assert packed_actual is not None
assert all(
torch.equal(result, reference)
for result, reference in zip(packed_actual, expected)
)
def test_qwen_qkv_epilogue_rejects_compile():
@@ -190,5 +203,125 @@ def test_qwen_qkv_epilogue_rejects_unsupported_head_dim():
)
@pytest.mark.parametrize("mode", ["dense", "attention_mask", "text_mask", "sharded"])
def test_qwen_attention_preserves_normalized_segments(mode):
heads, head_dim, img_tokens, txt_tokens = 4, 128, 17, 7
dim = heads * head_dim
img = torch.randn(1, img_tokens, dim, device="cuda", dtype=torch.bfloat16)
txt = torch.randn(1, txt_tokens, dim, device="cuda", dtype=torch.bfloat16)
projections = [torch.randn_like(img) for _ in range(3)] + [
torch.randn_like(txt) for _ in range(3)
]
class IdentityProjection(torch.nn.Module):
def forward(self, x):
return x, None
class CaptureAttention(torch.nn.Module):
sp_attention_mode = "kv_gather"
def forward(self, q, k, v, **kwargs):
tensors = [q, k, v]
if kwargs["q_prefix"] is not None:
tensors = [
torch.cat([kwargs[name], tensor], dim=1)
for name, tensor in zip(
("q_prefix", "k_prefix", "v_prefix"), tensors
)
]
self.inputs = tuple(tensor.clone() for tensor in tensors)
return tensors[0]
module = object.__new__(qwen_image.QwenImageCrossAttention)
torch.nn.Module.__init__(module)
module._unquantized_added_qkv_is_packed = False
module.local_num_heads = heads
module.head_dim = head_dim
module.qk_norm = True
for name in ("norm_q", "norm_k", "norm_added_q", "norm_added_k"):
setattr(module, name, RMSNorm(head_dim).to(device="cuda", dtype=img.dtype))
module.attn = CaptureAttention()
module.to_out = torch.nn.ModuleList([IdentityProjection()])
module.to_add_out = IdentityProjection()
def cache(tokens):
angle = torch.randn(tokens, head_dim // 2, device="cuda")
return torch.cat([angle.cos(), angle.sin()], dim=-1)
caches = (cache(img_tokens), cache(txt_tokens))
kwargs = {}
if mode == "attention_mask":
kwargs["attn_mask"] = torch.ones(1, img_tokens + txt_tokens, device="cuda")
elif mode == "text_mask":
kwargs["encoder_hidden_states_mask"] = torch.ones(1, txt_tokens, device="cuda")
elif mode == "sharded":
kwargs["sp_text_sharded"] = True
with patch.object(
qwen_image,
"_get_qkv_projections",
side_effect=lambda *a, **kw: tuple(t.clone() for t in projections),
):
module.use_fused_qkv_epilogue = False
expected = module(img, txt, image_rotary_emb=caches, **kwargs)
expected_inputs = module.attn.inputs
module.use_fused_qkv_epilogue = True
with patch.object(
qwen_image,
"try_fused_qwen_qkv_epilogue",
wraps=try_fused_qwen_qkv_epilogue,
) as fused:
actual = module(img, txt, image_rotary_emb=caches, **kwargs)
assert fused.call_count == (1 if mode == "dense" else 0)
assert all(torch.equal(a, b) for a, b in zip(actual[:2], expected[:2]))
assert all(torch.equal(a, b) for a, b in zip(module.attn.inputs, expected_inputs))
@pytest.mark.parametrize("misaligned_image", [True, False])
def test_qwen_qkv_epilogue_rejects_misaligned_token_stride(misaligned_image):
tensor = torch.empty(1, 2, 2, 128, device="cuda", dtype=torch.bfloat16)
row = torch.empty(128, device="cuda", dtype=torch.bfloat16)
cache = torch.empty(2, 128, device="cuda", dtype=torch.float32)
pitched = torch.empty_strided(
(1, 2, 2, 128), (520, 260, 128, 1), device="cuda", dtype=torch.bfloat16
)
assert pitched.data_ptr() % 32 == 0
assert pitched.stride(1) * pitched.element_size() % 16 == 8
img = pitched if misaligned_image else tensor
txt = tensor if misaligned_image else pitched
assert (
try_fused_qwen_qkv_epilogue(
*([img] * 3),
*([txt] * 3),
*([row] * 4),
cache,
cache,
1e-6,
1e-6,
)
is None
)
@pytest.mark.parametrize("misaligned_image", [True, False])
def test_qwen_qkv_epilogue_rejects_misaligned_cache(misaligned_image):
tensor = torch.empty(1, 1, 1, 128, device="cuda", dtype=torch.bfloat16)
row = torch.empty(128, device="cuda", dtype=torch.bfloat16)
cache = torch.empty(1, 128, device="cuda", dtype=torch.float32)
misaligned = torch.empty(129, device="cuda", dtype=torch.float32)[1:].view(1, 128)
assert misaligned.is_contiguous() and misaligned.data_ptr() % 8 != 0
assert (
try_fused_qwen_qkv_epilogue(
*([tensor] * 6),
*([row] * 4),
misaligned if misaligned_image else cache,
cache if misaligned_image else misaligned,
1e-6,
1e-6,
)
is None
)
if __name__ == "__main__":
sys.exit(pytest.main([__file__, "-v"]))
test/registered/kernels/ops/layernorm/test_kernels_namespace.py
+16
View File
@@ -45,6 +45,22 @@ def test_unknown_op_or_backend_raises():
K.select_kernel("gemm.fp8_scaled_mm", backend=KernelBackend.TRITON)
@pytest.mark.parametrize(
"platform, eligible",
[
(_SM90, True),
(_SM100, True),
(PlatformInfo(device_type="cuda", cuda_arch_major=8, cuda_arch_minor=0), False),
(_CPU, False),
(_HIP, False),
],
)
def test_qwen_qkv_registry_accepts_hopper(platform, eligible):
spec = K.select_kernel("diffusion.qwen_qkv_epilogue")
assert spec.backend is KernelBackend.JIT
assert K.capabilities_satisfied(spec.capabilities, platform) is eligible
def test_multi_backend_requires_explicit_backend(monkeypatch):
# Device is a hard eligibility filter, not a ranking: >1 usable backend on
# the current device means selection must name one.