From 3e035a3513ba557c0da54ea1dac986b53e5f87d4 Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang <1182563586@qq.com> Date: Sun, 13 Sep 2026 09:15:01 +0800 Subject: [PATCH] [Diffusion] Optimize Qwen-Image-Edit attention on Hopper (#38584) Co-authored-by: Mick Qian --- .../jit/csrc/diffusion/qwen_qkv_epilogue.cuh | 56 ++++-- python/sglang/kernels/ops/diffusion/README.md | 2 +- .../sglang/kernels/ops/diffusion/__init__.py | 3 +- .../diffusion/rope/qwen_qkv_epilogue_jit.py | 6 +- .../runtime/layers/attention/backends/sdpa.py | 46 ++++- .../runtime/models/dits/qwen_image.py | 15 +- .../ops/diffusion/test_model_fast_paths.py | 87 +++++++++ .../ops/diffusion/test_qwen_qkv_epilogue.py | 179 +++++++++++++++--- .../ops/layernorm/test_kernels_namespace.py | 16 ++ 9 files changed, 362 insertions(+), 48 deletions(-) diff --git a/python/sglang/kernels/jit/csrc/diffusion/qwen_qkv_epilogue.cuh b/python/sglang/kernels/jit/csrc/diffusion/qwen_qkv_epilogue.cuh index 35dc8eacf..c5dce25de 100644 --- a/python/sglang/kernels/jit/csrc/diffusion/qwen_qkv_epilogue.cuh +++ b/python/sglang/kernels/jit/csrc/diffusion/qwen_qkv_epilogue.cuh @@ -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(input, lane); if (kind == 2) { - store_as(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; + const auto input_vec = load_as(input, lane); + store_as(output, input_vec, lane); + } else { + const auto input_vec = load_as(input, lane); + store_as(output, input_vec, lane); + } continue; } + auto input_vec = load_as(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(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(cos_ptr) + lane); + const auto sin_pair = __ldg(reinterpret_cast(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(NI.unwrap()); const uint32_t txt_tokens = static_cast(NT.unwrap()); const uint32_t num_heads = static_cast(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); } }; diff --git a/python/sglang/kernels/ops/diffusion/README.md b/python/sglang/kernels/ops/diffusion/README.md index 4aeb0c87e..5a8e811c6 100644 --- a/python/sglang/kernels/ops/diffusion/README.md +++ b/python/sglang/kernels/ops/diffusion/README.md @@ -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 | diff --git a/python/sglang/kernels/ops/diffusion/__init__.py b/python/sglang/kernels/ops/diffusion/__init__.py index f584fd422..53ddd3c5e 100644 --- a/python/sglang/kernels/ops/diffusion/__init__.py +++ b/python/sglang/kernels/ops/diffusion/__init__.py @@ -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.", ), ( diff --git a/python/sglang/kernels/ops/diffusion/rope/qwen_qkv_epilogue_jit.py b/python/sglang/kernels/ops/diffusion/rope/qwen_qkv_epilogue_jit.py index 5c1c48596..42acc0b3f 100644 --- a/python/sglang/kernels/ops/diffusion/rope/qwen_qkv_epilogue_jit.py +++ b/python/sglang/kernels/ops/diffusion/rope/qwen_qkv_epilogue_jit.py @@ -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 diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/backends/sdpa.py b/python/sglang/multimodal_gen/runtime/layers/attention/backends/sdpa.py index 5552cfa88..dc4fbca1b 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/backends/sdpa.py +++ b/python/sglang/multimodal_gen/runtime/layers/attention/backends/sdpa.py @@ -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 diff --git a/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py b/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py index cd899bd6e..cea45be71 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py @@ -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, diff --git a/test/registered/kernels/ops/diffusion/test_model_fast_paths.py b/test/registered/kernels/ops/diffusion/test_model_fast_paths.py index 8eb00e3d1..f3f9a6bb7 100644 --- a/test/registered/kernels/ops/diffusion/test_model_fast_paths.py +++ b/test/registered/kernels/ops/diffusion/test_model_fast_paths.py @@ -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 diff --git a/test/registered/kernels/ops/diffusion/test_qwen_qkv_epilogue.py b/test/registered/kernels/ops/diffusion/test_qwen_qkv_epilogue.py index ca1e356f9..48eeeddf1 100644 --- a/test/registered/kernels/ops/diffusion/test_qwen_qkv_epilogue.py +++ b/test/registered/kernels/ops/diffusion/test_qwen_qkv_epilogue.py @@ -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"])) diff --git a/test/registered/kernels/ops/layernorm/test_kernels_namespace.py b/test/registered/kernels/ops/layernorm/test_kernels_namespace.py index e715a0066..ae0523631 100644 --- a/test/registered/kernels/ops/layernorm/test_kernels_namespace.py +++ b/test/registered/kernels/ops/layernorm/test_kernels_namespace.py @@ -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.