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

Co-authored-by: Mick Qian <mickqian@users.noreply.github.com>

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 kind = blockIdx.y;  // Q, K, and V have independent grids.
-  const uint32_t total_works = 3 * token_head_works;
+  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 joint_token = work / heads_per_token;
-    const uint32_t token_head = work % token_head_works;
+    const uint32_t head = (work % heads_per_token) * (is_value ? 2 : 1);
-    const uint32_t joint_token = token_head / params.num_heads;
-    const uint32_t head = token_head % params.num_heads;
     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) {
+      // 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 = i == 0 ? cos_pair.x : cos_pair.y;
-      const float cos = __ldg(cos_ptr + cache_idx);
+      const float sin = i == 0 ? sin_pair.x : sin_pair.y;
-      const float sin = __ldg(sin_ptr + cache_idx);
+      // Preserve the original QKNorm/RoPE contraction: round the sin product
-      elems[i] = x * cos - y * sin;
+      // before adding it to the fused cos product, including for vector loads.
-      elems[i + 1] = y * cos + x * sin;
+      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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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,
+    not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] < 9,
-    reason="Qwen-Image QKV epilogue requires SM100+",
+    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():
+@pytest.mark.parametrize(
-    heads = 4
+    "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,9 +126,16 @@ def test_qwen_qkv_epilogue_is_bit_exact():
     ]
     assert all(not tensor.is_contiguous() for tensor in (*img_views, *txt_views))
 
+    # 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_views,
+            *img_inputs,
-        *txt_views,
+            *txt_inputs,
             norms[0].weight,
             norms[1].weight,
             norms[2].weight,
@@ -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

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