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
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,
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.