minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 15 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
python/sglang/kernels/jit/csrc/diffusion/qwen_qkv_epilogue.cuh
+36 -20
View File
@@ -37,7 +37,8 @@ struct Params {
const void* txt_k_weight;
const void* img_cache;
const void* txt_cache;
int64_t input_token_stride_bytes;
int64_t img_token_stride_bytes;
int64_t txt_token_stride_bytes;
int64_t output_token_stride_bytes;
int64_t head_stride_bytes;
uint32_t img_tokens;
@@ -57,14 +58,15 @@ __global__ void qwen_qkv_epilogue_kernel(const Params __grid_constant__ params)
const uint32_t start = blockIdx.x * kWarps + warp;
const uint32_t workers = gridDim.x * kWarps;
const uint32_t total_tokens = params.txt_tokens + params.img_tokens;
const uint32_t token_head_works = total_tokens * params.num_heads;
const uint32_t total_works = 3 * token_head_works;
const uint32_t kind = blockIdx.y; // Q, K, and V have independent grids.
const bool is_value = kind == 2;
const uint32_t v_heads = div_ceil(params.num_heads, uint32_t(2));
const uint32_t heads_per_token = is_value ? v_heads : params.num_heads;
const uint32_t total_works = total_tokens * heads_per_token;
for (uint32_t work = start; work < total_works; work += workers) {
const uint32_t kind = work / token_head_works; // 0: Q, 1: K, 2: V.
const uint32_t token_head = work % token_head_works;
const uint32_t joint_token = token_head / params.num_heads;
const uint32_t head = token_head % params.num_heads;
const uint32_t joint_token = work / heads_per_token;
const uint32_t head = (work % heads_per_token) * (is_value ? 2 : 1);
const bool is_text = joint_token < params.txt_tokens;
const uint32_t source_token = is_text ? joint_token : joint_token - params.txt_tokens;
@@ -81,16 +83,26 @@ __global__ void qwen_qkv_epilogue_kernel(const Params __grid_constant__ params)
output_base = params.joint_v;
}
const auto input_token_stride_bytes = is_text ? params.txt_token_stride_bytes : params.img_token_stride_bytes;
const void* input =
pointer::offset(input_base, source_token * params.input_token_stride_bytes, head * params.head_stride_bytes);
pointer::offset(input_base, source_token * input_token_stride_bytes, head * params.head_stride_bytes);
void* output =
pointer::offset(output_base, joint_token * params.output_token_stride_bytes, head * params.head_stride_bytes);
auto input_vec = load_as<Storage>(input, lane);
if (kind == 2) {
store_as<Storage>(output, input_vec, lane);
// One warp copies two adjacent heads, keeping all V workers active while
// issuing twice as many bytes per memory instruction.
if (head + 1 < params.num_heads) {
using CopyStorage = AlignedVector<Packed, 2 * kVecSize>;
const auto input_vec = load_as<CopyStorage>(input, lane);
store_as<CopyStorage>(output, input_vec, lane);
} else {
const auto input_vec = load_as<Storage>(input, lane);
store_as<Storage>(output, input_vec, lane);
}
continue;
}
auto input_vec = load_as<Storage>(input, lane);
const void* weight_base;
if (kind == 0) {
@@ -123,15 +135,20 @@ __global__ void qwen_qkv_epilogue_kernel(const Params __grid_constant__ params)
const auto* cache = static_cast<const float*>(is_text ? params.txt_cache : params.img_cache);
const auto* cos_ptr = cache + source_token * kHeadDim;
const auto* sin_ptr = cos_ptr + kHeadDim / 2;
// Each lane consumes two adjacent cache entries. Loading them together
// avoids the half-used sectors from two stride-2 scalar loads.
const auto cos_pair = __ldg(reinterpret_cast<const float2*>(cos_ptr) + lane);
const auto sin_pair = __ldg(reinterpret_cast<const float2*>(sin_ptr) + lane);
#pragma unroll
for (uint32_t i = 0; i < kElemsPerThread; i += 2) {
const float x = elems[i];
const float y = elems[i + 1];
const uint32_t cache_idx = (lane * kElemsPerThread + i) / 2;
const float cos = __ldg(cos_ptr + cache_idx);
const float sin = __ldg(sin_ptr + cache_idx);
elems[i] = x * cos - y * sin;
elems[i + 1] = y * cos + x * sin;
const float cos = i == 0 ? cos_pair.x : cos_pair.y;
const float sin = i == 0 ? sin_pair.x : sin_pair.y;
// Preserve the original QKNorm/RoPE contraction: round the sin product
// before adding it to the fused cos product, including for vector loads.
elems[i] = __fmaf_rn(x, cos, -__fmul_rn(y, sin));
elems[i + 1] = __fmaf_rn(y, cos, __fmul_rn(x, sin));
}
#pragma unroll
@@ -204,7 +221,6 @@ struct QwenQKVEpilogueKernel {
RuntimeCheck(
txt_q.stride(0) == txt_k.stride(0) && txt_q.stride(0) == txt_v.stride(0),
"text QKV inputs must use the same token stride");
RuntimeCheck(img_q.stride(0) == txt_q.stride(0), "image/text QKV token strides must match");
RuntimeCheck(
img_q.stride(1) == kHeadDim && img_k.stride(1) == kHeadDim && img_v.stride(1) == kHeadDim,
"image QKV heads must be contiguous");
@@ -220,11 +236,10 @@ struct QwenQKVEpilogueKernel {
const uint32_t img_tokens = static_cast<uint32_t>(NI.unwrap());
const uint32_t txt_tokens = static_cast<uint32_t>(NT.unwrap());
const uint32_t num_heads = static_cast<uint32_t>(H.unwrap());
const uint32_t total_works = 3 * (img_tokens + txt_tokens) * num_heads;
const uint32_t total_works = (img_tokens + txt_tokens) * num_heads;
if (total_works == 0) return;
const int64_t head_stride_bytes = kHeadDim * sizeof(bf16_t);
const int64_t input_token_stride_bytes = img_q.stride(0) * sizeof(bf16_t);
const int64_t output_token_stride_bytes = num_heads * head_stride_bytes;
const auto params = Params{
.joint_q = joint_q.data_ptr(),
@@ -242,7 +257,8 @@ struct QwenQKVEpilogueKernel {
.txt_k_weight = txt_k_weight.data_ptr(),
.img_cache = img_cache.data_ptr(),
.txt_cache = txt_cache.data_ptr(),
.input_token_stride_bytes = input_token_stride_bytes,
.img_token_stride_bytes = img_q.stride(0) * sizeof(bf16_t),
.txt_token_stride_bytes = txt_q.stride(0) * sizeof(bf16_t),
.output_token_stride_bytes = output_token_stride_bytes,
.head_stride_bytes = head_stride_bytes,
.img_tokens = img_tokens,
@@ -256,7 +272,7 @@ struct QwenQKVEpilogueKernel {
static const uint32_t blocks_per_sm = runtime::get_blocks_per_sm(qwen_qkv_epilogue_kernel, kThreads);
const uint32_t needed_blocks = div_ceil(total_works, uint32_t(kWarps));
const uint32_t blocks = std::min(blocks_per_sm * sm_count, needed_blocks);
LaunchKernel(blocks, kThreads, device.unwrap())(qwen_qkv_epilogue_kernel, params);
LaunchKernel(dim3(blocks, 3), kThreads, device.unwrap())(qwen_qkv_epilogue_kernel, params);
}
};
python/sglang/kernels/ops/diffusion/README.md
+1 -1
View File
@@ -143,7 +143,7 @@ tensor copy per residual site.
| `fused_qknorm_rope_pack_kv` | JIT CUDA | as above, also packs prefix K/V |
| `fused_qknorm_rope_out_of_place` | JIT CUDA | as above, bit-equal to the in-place kernel; reads strided q/k and writes contiguous copies, inputs untouched (VDN-H3 keeps the raw q/k for its linear branch) |
| `try_fused_flux2_qkv_epilogue` | KDA (JIT CUDA) | bit-exact vs the selected BF16 chain | FLUX.2 QK RMSNorm + RoPE + joint QKV packing |
| `try_fused_qwen_qkv_epilogue` | JIT CUDA | bit-exact vs the selected BF16 chain | Qwen-Image QK RMSNorm + RoPE + joint QKV writes; SM100+ |
| `try_fused_qwen_qkv_epilogue` | JIT CUDA | bit-exact vs the selected BF16 chain | Qwen-Image QK RMSNorm + RoPE + joint QKV writes; SM90+ |
| `fused_rope_rotate_half_bitexact` | Triton | bit-exact (elementwise only) |
| `fused_interleaved_rope_fp64` | JIT CUDA | bit-exact vs paired SANA-Video fp64 RoPE |
| `fused_inplace_helios_qk_rope` | JIT CUDA | bit-exact paired in-place RoPE for Helios' transposed frequency layout |
python/sglang/kernels/ops/diffusion/__init__.py
+2 -1
View File
@@ -38,6 +38,7 @@ from sglang.kernels.spec import (
)
_CUDA = frozenset({CapabilityRequirement.CUDA})
_CUDA_SM90_PLUS = frozenset({CapabilityRequirement.cuda(min_sm=(9, 0))})
_CUDA_SM100_PLUS = frozenset({CapabilityRequirement.cuda(min_sm=(10, 0))})
_HIP = frozenset({CapabilityRequirement.HIP})
@@ -241,7 +242,7 @@ _SPECS: tuple[tuple[str, KernelBackend, str, frozenset, str], ...] = (
"diffusion.qwen_qkv_epilogue",
KernelBackend.JIT,
"rope.qwen_qkv_epilogue_jit:try_fused_qwen_qkv_epilogue",
_CUDA_SM100_PLUS,
_CUDA_SM90_PLUS,
"Qwen-Image QK RMS-norm, RoPE, and joint QKV writes.",
),
(
python/sglang/kernels/ops/diffusion/rope/qwen_qkv_epilogue_jit.py
+5 -1
View File
@@ -39,6 +39,8 @@ def _qkv_tensor(tensor: torch.Tensor, like: torch.Tensor | None = None) -> bool:
and tensor.numel() > 0
and tensor.stride(-1) == 1
and tensor.stride(-2) == _HEAD_DIM
# The paired-head V copy uses 16-byte loads at every token boundary.
and tensor.stride(1) * tensor.element_size() % 16 == 0
and tensor.data_ptr() % _ALIGN == 0
and (
like is None
@@ -85,7 +87,7 @@ def try_fused_qwen_qkv_epilogue(
and _qkv_tensor(txt_v, txt_q)
and img_q.shape[2] == txt_q.shape[2]
and torch.version.cuda is not None
and torch.cuda.get_device_capability(img_q.device)[0] >= 10
and torch.cuda.get_device_capability(img_q.device)[0] >= 9
):
return None
@@ -122,6 +124,8 @@ def try_fused_qwen_qkv_epilogue(
and txt_cache.shape[0] >= txt_q.shape[1]
and img_cache.is_contiguous()
and txt_cache.is_contiguous()
and img_cache.data_ptr() % 8 == 0
and txt_cache.data_ptr() % 8 == 0
):
return None
python/sglang/multimodal_gen/runtime/layers/attention/backends/sdpa.py
+45 -1
View File
@@ -7,6 +7,11 @@ from contextlib import nullcontext
import torch
from torch.nn.attention import SDPBackend, sdpa_kernel
try:
from torch.nn.attention.varlen import varlen_attn as torch_varlen_attn
except ImportError:
torch_varlen_attn = None
from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import ( # FlashAttentionMetadata,
AttentionBackend,
AttentionImpl,
@@ -118,7 +123,46 @@ class SDPAImpl(AttentionImpl):
max_seqlen: int,
cu_seqlens_host: tuple[int, ...] | None = None,
) -> torch.Tensor:
del max_seqlen
if (
type(self) is SDPAImpl
and torch_varlen_attn is not None
and not torch.compiler.is_compiling()
and not torch.is_grad_enabled()
and query.is_cuda
and torch.version.hip is None
and query.ndim == 3
and query.shape == key.shape == value.shape
and query.dtype in (torch.float16, torch.bfloat16)
and query.dtype == key.dtype == value.dtype
and query.device == key.device == value.device == cu_seqlens.device
and query.stride(-1) == key.stride(-1) == value.stride(-1) == 1
and query.numel() > 0
and query.shape[-1] <= 256
and query.shape[-1] % 8 == 0
and cu_seqlens.dtype == torch.int32
and cu_seqlens.ndim == 1
and cu_seqlens.is_contiguous()
and cu_seqlens.numel() > 2
and max_seqlen > 0
and self.dropout == 0.0
and not self.allow_cudnn_sdp
and torch.backends.cuda.flash_sdp_enabled()
and not torch.backends.cuda.cudnn_sdp_enabled()
and torch.cuda.get_device_capability(query.device)[0] == 9
):
# Keep the existing Flash SDPA arithmetic while consuming all
# packed windows in one call, including ragged and empty windows.
return torch_varlen_attn(
query,
key,
value,
cu_seqlens,
cu_seqlens,
max_seqlen,
max_seqlen,
scale=self.softmax_scale,
window_size=(-1, 0) if self.causal else (-1, -1),
)
bounds = (
cu_seqlens_host
if cu_seqlens_host is not None
python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py
+14 -1
View File
@@ -3,6 +3,7 @@
# SPDX-License-Identifier: Apache-2.0
import functools
import os
from dataclasses import dataclass
from typing import Any, Dict, List, Optional, Tuple, Union
@@ -750,10 +751,18 @@ class QwenImageCrossAttention(nn.Module):
self.prefix = prefix
self.defer_output_bias = _defer_modelopt_output_bias(quant_config)
quant_name = _modelopt_quant_name(quant_config)
capability = current_platform.get_device_capability()
self.use_fused_qkv_epilogue = quant_name in {
"modelopt_fp4",
"modelopt_fp8",
}
} or (
quant_config is None
and current_platform.is_cuda()
and capability is not None
and capability.major == 9
and os.getenv("SGLANG_ENABLE_FUSED_QKNORM_ROPE", "1").lower()
not in {"0", "false", "off", "no"}
)
self.use_fused_qkv = (
isinstance(quant_config, NunchakuConfig) or quant_name == "modelopt_fp8"
)
@@ -1010,6 +1019,10 @@ class QwenImageCrossAttention(nn.Module):
and txt_cache is not None
and not sp_text_sharded
and sp_txt_pad == 0
# Masked attention packs the image and text segments separately.
# Its prefix tensors must go through the ordinary normalization.
and attn_mask is None
and encoder_hidden_states_mask is None
):
joint_qkv = try_fused_qwen_qkv_epilogue(
img_query,