[Diffusion][ERNIE] Fuse QKNorm with full-width RoPE (#34620)

python/sglang/kernels/jit/csrc/diffusion/qknorm_rope.cuh

@@ -172,6 +172,7 @@ template <
     typename CacheDType,
     bool kRoundNormBeforeRope,
     bool kPackKV,
+    bool kCacheHasFullWidth,
     typename IdType>
 __global__ void fused_qknorm_rope_warp(const QKNormRopeParamsT<kPackKV> __grid_constant__ params) {
   using namespace device;
@@ -185,7 +186,8 @@ __global__ void fused_qknorm_rope_warp(const QKNormRopeParamsT<kPackKV> __grid_c
   constexpr uint32_t kRotaryLanes = kRopeDim / kElemsPerThread;
   constexpr uint32_t kHalfRotaryLanes = kRotaryLanes / 2;
   constexpr uint32_t kActiveMask = active_mask<kRotaryLanes>();
-  constexpr int64_t kCosSinStrideBytes = kRopeDim * sizeof(CacheDType);
+  constexpr int64_t kCacheRotaryDim = kCacheHasFullWidth ? 2 * kRopeDim : kRopeDim;
+  constexpr int64_t kCosSinStrideBytes = kCacheRotaryDim * sizeof(CacheDType);
 
   static_assert(kElemsPerThread % 2 == 0, "Each lane must own an even number of elements");
   static_assert(kRopeDim > 0 && kRopeDim <= kHeadDim, "Invalid rope dimension");
@@ -285,7 +287,7 @@ __global__ void fused_qknorm_rope_warp(const QKNormRopeParamsT<kPackKV> __grid_c
       auto output_vec = norm::apply_norm_warp<kHeadDim>(input_vec, weight_vec, eps);
       const auto pos = static_cast<int64_t>(static_cast<const IdType*>(positions)[token_id]);
       const auto cos_ptr = static_cast<const CacheDType*>(pointer::offset(cos_sin_cache_ptr, pos * kCosSinStrideBytes));
-      const auto sin_ptr = cos_ptr + kRopeDim / 2;
+      const auto sin_ptr = cos_ptr + (kCacheHasFullWidth ? kRopeDim : kRopeDim / 2);
 
       if constexpr (kIsNeox) {
         if (lane_id < kRotaryLanes) {
@@ -301,9 +303,10 @@ __global__ void fused_qknorm_rope_warp(const QKNormRopeParamsT<kPackKV> __grid_c
             const auto& partner_values = unpack(partner_vec);
 #pragma unroll
             for (uint32_t i = 0; i < 2; ++i) {
-              const auto half_idx = (lane_id % kHalfRotaryLanes) * kElemsPerThread + 2 * j + i;
+              const auto cache_idx =
-              const auto cos = load_cache_value(cos_ptr, half_idx);
+                  (kCacheHasFullWidth ? lane_id : lane_id % kHalfRotaryLanes) * kElemsPerThread + 2 * j + i;
-              const auto sin = load_cache_value(sin_ptr, half_idx);
+              const auto cos = load_cache_value(cos_ptr, cache_idx);
+              const auto sin = load_cache_value(sin_ptr, cache_idx);
               values[i] = lane_id < kHalfRotaryLanes ? rotary_sub(values[i], cos, partner_values[i], sin)
                                                      : rotary_add(values[i], cos, partner_values[i], sin);
             }
@@ -354,7 +357,7 @@ __global__ void fused_qknorm_rope_warp(const QKNormRopeParamsT<kPackKV> __grid_c
         const auto pos = static_cast<int64_t>(static_cast<const IdType*>(positions)[token_id]);
         const auto cos_ptr =
             static_cast<const CacheDType*>(pointer::offset(cos_sin_cache_ptr, pos * kCosSinStrideBytes));
-        const auto sin_ptr = cos_ptr + kRopeDim / 2;
+        const auto sin_ptr = cos_ptr + (kCacheHasFullWidth ? kRopeDim : kRopeDim / 2);
         const auto partner_lane = lane_id < kHalfRotaryLanes ? lane_id + kHalfRotaryLanes : lane_id - kHalfRotaryLanes;
 
 #pragma unroll
@@ -363,9 +366,9 @@ __global__ void fused_qknorm_rope_warp(const QKNormRopeParamsT<kPackKV> __grid_c
           if (lane_id < kHalfRotaryLanes) {
             swapped = -swapped;
           }
-          const auto half_idx = (lane_id % kHalfRotaryLanes) * kElemsPerThread + i;
+          const auto cache_idx = (kCacheHasFullWidth ? lane_id : lane_id % kHalfRotaryLanes) * kElemsPerThread + i;
-          const float cos = cast<fp32_t>(load_cache_value(cos_ptr, half_idx));
+          const float cos = cast<fp32_t>(load_cache_value(cos_ptr, cache_idx));
-          const float sin = cast<fp32_t>(load_cache_value(sin_ptr, half_idx));
+          const float sin = cast<fp32_t>(load_cache_value(sin_ptr, cache_idx));
           elems[i] = elems[i] * cos + swapped * sin;
         }
       }
@@ -374,7 +377,7 @@ __global__ void fused_qknorm_rope_warp(const QKNormRopeParamsT<kPackKV> __grid_c
         const auto pos = static_cast<int64_t>(static_cast<const IdType*>(positions)[token_id]);
         const auto cos_ptr =
             static_cast<const CacheDType*>(pointer::offset(cos_sin_cache_ptr, pos * kCosSinStrideBytes));
-        const auto sin_ptr = cos_ptr + kRopeDim / 2;
+        const auto sin_ptr = cos_ptr + (kCacheHasFullWidth ? kRopeDim : kRopeDim / 2);
 
 #pragma unroll
         for (uint32_t i = 0; i < kElemsPerThread; i += 2) {
@@ -406,7 +409,8 @@ template <
     bool kUsePDL,
     typename DType,
     typename CacheDType,
-    bool kRoundNormBeforeRope>
+    bool kRoundNormBeforeRope,
+    bool kCacheHasFullWidth>
 struct QKNormRopeKernel {
   static_assert(kHeadDim <= 256, "Only head_dim <= 256 is supported");
   template <typename IdType>
@@ -419,6 +423,7 @@ struct QKNormRopeKernel {
       CacheDType,
       kRoundNormBeforeRope,
       false,
+      kCacheHasFullWidth,
       IdType>;
 
   static void
@@ -448,7 +453,10 @@ struct QKNormRopeKernel {
     TensorMatcher({N, Q, D}).with_strides({Dq, Dd, 1}).with_dtype<DType>().with_device(device).verify(q);
     TensorMatcher({N, K, D}).with_strides({Dk, Dd, 1}).with_dtype<DType>().with_device(device).verify(k);
     TensorMatcher({D}).with_dtype<DType>().with_device(device).verify(q_weight).verify(k_weight);
-    TensorMatcher({-1, R}).with_dtype<CacheDType>().with_device(device).verify(cos_sin_cache);
+    TensorMatcher({-1, kCacheHasFullWidth ? 2 * kRopeDim : kRopeDim})
+        .with_dtype<CacheDType>()
+        .with_device(device)
+        .verify(cos_sin_cache);
     TensorMatcher({N}).with_dtype<int32_t, int64_t>(id_type).with_device(device).verify(positions);
 
     const auto num_tokens = static_cast<uint32_t>(N.unwrap());
@@ -498,8 +506,10 @@ template <
     bool kUsePDL,
     typename DType,
     typename CacheDType,
-    bool kRoundNormBeforeRope>
+    bool kRoundNormBeforeRope,
+    bool kCacheHasFullWidth>
 struct QKNormRopePackKVKernel {
+  static_assert(!kCacheHasFullWidth, "KV packing does not support full-width cos/sin caches");
   template <typename IdType>
   static constexpr auto kernel = fused_qknorm_rope_warp<
       kHeadDim,
@@ -510,6 +520,7 @@ struct QKNormRopePackKVKernel {
       CacheDType,
       kRoundNormBeforeRope,
       true,
+      kCacheHasFullWidth,
       IdType>;
 
   static void

python/sglang/kernels/ops/diffusion/qknorm_rope.py

@@ -32,6 +32,7 @@ def _jit_qknorm_rope_module(
     cache_dtype: torch.dtype,
     round_norm_before_rope: bool,
     pack_kv: bool = False,
+    cache_has_full_width: bool = False,
 ) -> Module:
     args = make_cpp_args(
         head_dim,
@@ -41,6 +42,7 @@ def _jit_qknorm_rope_module(
         dtype,
         cache_dtype,
         round_norm_before_rope,
+        cache_has_full_width,
     )
     op_name = "qknorm_rope_pack_kv" if pack_kv else "qknorm_rope"
     kernel_name = "QKNormRopePackKVKernel" if pack_kv else "QKNormRopeKernel"
@@ -60,6 +62,7 @@ def _can_use_fused_qknorm_rope(
     cache_dtype: torch.dtype,
     round_norm_before_rope: bool,
     pack_kv: bool,
+    cache_has_full_width: bool,
 ) -> bool:
     if dtype not in _SUPPORTED_DTYPES or cache_dtype not in _SUPPORTED_CACHE_DTYPES:
         logger.warning(
@@ -93,6 +96,12 @@ def _can_use_fused_qknorm_rope(
                 rotary_lanes,
             )
             return False
+    elif cache_has_full_width:
+        logger.warning("Full-width cos/sin caches are only supported for NeoX RoPE")
+        return False
+    if pack_kv and cache_has_full_width:
+        logger.warning("KV packing does not support full-width cos/sin caches")
+        return False
     if round_norm_before_rope and cache_dtype != dtype:
         logger.warning(
             "Exact fused QKNorm+RoPE requires cache dtype %s to match activation dtype %s",
@@ -109,6 +118,7 @@ def _can_use_fused_qknorm_rope(
             cache_dtype,
             round_norm_before_rope,
             pack_kv,
+            cache_has_full_width,
         )
         return True
     except Exception as e:
@@ -127,6 +137,7 @@ def can_use_fused_inplace_qknorm_rope(
     cache_dtype: torch.dtype = torch.float32,
     round_norm_before_rope: bool = False,
     pack_kv: bool = False,
+    cache_has_full_width: bool = False,
 ) -> bool:
     return _can_use_fused_qknorm_rope(
         head_dim,
@@ -136,6 +147,7 @@ def can_use_fused_inplace_qknorm_rope(
         cache_dtype,
         round_norm_before_rope,
         pack_kv,
+        cache_has_full_width,
     )
 
 
@@ -153,9 +165,12 @@ def fused_inplace_qknorm_rope(
     head_dim: int = 0,
     rope_dim: int = 0,
     round_norm_before_rope: bool = False,
+    cache_has_full_width: bool = False,
 ) -> None:
     head_dim = head_dim or q.size(-1)
-    rope_dim = rope_dim or cos_sin_cache.size(-1)
+    if not rope_dim:
+        cache_width = cos_sin_cache.size(-1)
+        rope_dim = cache_width // 2 if cache_has_full_width else cache_width
     module = _jit_qknorm_rope_module(
         head_dim,
         rope_dim,
@@ -163,6 +178,8 @@ def fused_inplace_qknorm_rope(
         q.dtype,
         cos_sin_cache.dtype,
         round_norm_before_rope,
+        False,
+        cache_has_full_width,
     )
     module.qknorm_rope(q, k, q_weight, k_weight, cos_sin_cache, positions, eps)
 
@@ -198,6 +215,7 @@ def fused_qknorm_rope_pack_kv(
         cos_sin_cache.dtype,
         round_norm_before_rope,
         True,
+        False,
     )
     module.qknorm_rope_pack_kv(
         q.view(-1, q.shape[-2], head_dim),

python/sglang/multimodal_gen/runtime/layers/layernorm.py

@@ -973,11 +973,15 @@ def apply_qk_norm_rope(
     position_offset: int = 0,
     allow_inplace: bool = True,
     allow_strided_qk: bool = False,
+    round_norm_before_rope: bool = False,
+    cache_has_full_width: bool = False,
 ) -> Tuple[torch.Tensor, torch.Tensor]:
     """Apply QK RMSNorm followed by RoPE, fusing supported CUDA/XPU shapes.
 
     Strided packed-QKV views require an explicit opt-in because selecting the fused
     kernel changes the numerical path for models that historically used the fallback.
+    ``cache_has_full_width`` describes ``[full cos, full sin]`` cache rows and
+    requires the fused CUDA path; the ordinary cache stores half-width cos/sin.
     """
 
     from sglang.multimodal_gen.runtime.layers.rotary_embedding import (
@@ -1004,7 +1008,12 @@ def apply_qk_norm_rope(
     batch_size, seq_len, _, _ = q.shape
     q_eps = q_norm.variance_epsilon
     k_eps = k_norm.variance_epsilon
-    rope_dim = cos_sin_cache.size(-1)
+    cache_width = cos_sin_cache.size(-1)
+    if cache_has_full_width and cache_width % 2:
+        raise ValueError(
+            f"full-width cos/sin cache must have even width, got {cache_width}"
+        )
+    rope_dim = cache_width // 2 if cache_has_full_width else cache_width
     if rope_dim % 2 != 0 or rope_dim > head_dim:
         raise ValueError(
             f"cos_sin_cache width must be even and <= head_dim, got {rope_dim} vs {head_dim}"
@@ -1054,7 +1063,15 @@ def apply_qk_norm_rope(
         and k_norm.weight.dtype == k.dtype
         and q_has_supported_layout
         and k_has_supported_layout
-        and can_use_fused_inplace_qknorm_rope(head_dim, rope_dim, is_neox, q.dtype)
+        and can_use_fused_inplace_qknorm_rope(
+            head_dim=head_dim,
+            rope_dim=rope_dim,
+            is_neox=is_neox,
+            dtype=q.dtype,
+            cache_dtype=cos_sin_cache.dtype,
+            round_norm_before_rope=round_norm_before_rope,
+            cache_has_full_width=cache_has_full_width,
+        )
     ):
         fused_inplace_qknorm_rope(
             q=q.view(-1, q.shape[-2], head_dim),
@@ -1067,9 +1084,14 @@ def apply_qk_norm_rope(
             eps=q_eps,
             head_dim=head_dim,
             rope_dim=rope_dim,
+            round_norm_before_rope=round_norm_before_rope,
+            cache_has_full_width=cache_has_full_width,
         )
         return q, k
 
+    if cache_has_full_width:
+        raise RuntimeError("full-width cos/sin cache requires fused QKNorm+RoPE")
+
     if (
         _is_xpu
         and allow_inplace

python/sglang/multimodal_gen/runtime/models/dits/ernie_image.py

@@ -49,7 +49,11 @@ from sglang.multimodal_gen.runtime.layers.attention.layer import (
     USPAttention,
     build_varlen_mask_meta,
 )
-from sglang.multimodal_gen.runtime.layers.layernorm import RMSNorm, apply_qk_norm
+from sglang.multimodal_gen.runtime.layers.layernorm import (
+    RMSNorm,
+    apply_qk_norm,
+    apply_qk_norm_rope,
+)
 from sglang.multimodal_gen.runtime.layers.linear import (
     ColumnParallelLinear,
     MergedColumnParallelLinear,
@@ -68,6 +72,7 @@ logger = init_logger(__name__)
 _ERNIE_NORM = BitExactFusionGate("ERNIE fused-norm")
 _ERNIE_GATED_NORM = BitExactFusionGate("ERNIE fused gated-norm")
 _ERNIE_ROPE = BitExactFusionGate("ERNIE fused RoPE")
+_ERNIE_QKNORM_ROPE = BitExactFusionGate("ERNIE fused QKNorm+RoPE")
 _ERNIE_GEGLU = BitExactFusionGate("ERNIE fused GELU-mul")
 
 
@@ -284,6 +289,8 @@ class ErnieImageSelfAttention(nn.Module):
         x: torch.Tensor,
         rope_cos: torch.Tensor,
         rope_sin: torch.Tensor,
+        rope_cache: torch.Tensor,
+        rope_positions: torch.Tensor,
         attn_mask: torch.Tensor | None = None,
         attn_mask_meta: dict | None = None,
     ) -> torch.Tensor:
@@ -298,16 +305,20 @@ class ErnieImageSelfAttention(nn.Module):
         v = v.view(B, S, self.num_local_heads, self.head_dim)
 
         if self.qk_layernorm:
-            q, k = apply_qk_norm(
+            q, k = _ernie_qknorm_rope(
                 q,
                 k,
                 self.norm_q,
                 self.norm_k,
                 self.head_dim,
+                rope_cos,
+                rope_sin,
+                rope_cache,
+                rope_positions,
             )
-
+        else:
-        q = _ernie_rope(q, rope_cos, rope_sin)
+            q = _ernie_rope(q, rope_cos, rope_sin)
-        k = _ernie_rope(k, rope_cos, rope_sin)
+            k = _ernie_rope(k, rope_cos, rope_sin)
 
         attn_out = self.attn(
             q, k, v, attn_mask=attn_mask, attn_mask_meta=attn_mask_meta
@@ -378,6 +389,8 @@ class ErnieImageSharedAdaLNBlock(nn.Module):
         x: torch.Tensor,
         rope_cos: torch.Tensor,
         rope_sin: torch.Tensor,
+        rope_cache: torch.Tensor,
+        rope_positions: torch.Tensor,
         shift_msa: torch.Tensor,
         scale_msa: torch.Tensor,
         gate_msa: torch.Tensor,
@@ -393,6 +406,8 @@ class ErnieImageSharedAdaLNBlock(nn.Module):
             x,
             rope_cos,
             rope_sin,
+            rope_cache,
+            rope_positions,
             attn_mask=attn_mask,
             attn_mask_meta=attn_mask_meta,
         )
@@ -472,6 +487,89 @@ def _ernie_rope(
     return _apply_rotary_bshd_eager(x, cos_, sin_)
 
 
+def _ernie_qknorm_rope_reference(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    q_norm: RMSNorm,
+    k_norm: RMSNorm,
+    head_dim: int,
+    rope_cos: torch.Tensor,
+    rope_sin: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    q, k = apply_qk_norm(q, k, q_norm, k_norm, head_dim)
+    return _ernie_rope(q, rope_cos, rope_sin), _ernie_rope(k, rope_cos, rope_sin)
+
+
+def _ernie_qknorm_rope(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    q_norm: RMSNorm,
+    k_norm: RMSNorm,
+    head_dim: int,
+    rope_cos: torch.Tensor,
+    rope_sin: torch.Tensor,
+    rope_cache: torch.Tensor,
+    rope_positions: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Fuse ERNIE's QK RMSNorm and rotate-half RoPE without changing bits."""
+    verified = _ERNIE_QKNORM_ROPE.verified
+    if not _ERNIE_QKNORM_ROPE.disabled and (
+        verified or _ERNIE_QKNORM_ROPE.can_attempt_once()
+    ):
+        q_input = q.clone() if not verified else q
+        k_input = k.clone() if not verified else k
+        try:
+            out = apply_qk_norm_rope(
+                q=q,
+                k=k,
+                q_norm=q_norm,
+                k_norm=k_norm,
+                head_dim=head_dim,
+                cos_sin_cache=rope_cache,
+                is_neox=True,
+                positions=rope_positions,
+                round_norm_before_rope=True,
+                cache_has_full_width=True,
+            )
+        except Exception as exc:
+            _ERNIE_QKNORM_ROPE.on_exception(exc, logger=logger)
+            return _ernie_qknorm_rope_reference(
+                q_input,
+                k_input,
+                q_norm,
+                k_norm,
+                head_dim,
+                rope_cos,
+                rope_sin,
+            )
+        else:
+            if verified:
+                return out
+            ref = _ernie_qknorm_rope_reference(
+                q_input,
+                k_input,
+                q_norm,
+                k_norm,
+                head_dim,
+                rope_cos,
+                rope_sin,
+            )
+            return _ERNIE_QKNORM_ROPE.accept_or_fallback(
+                out,
+                ref,
+                equal=tensors_equal,
+                logger=logger,
+                mismatch_msg=(
+                    "ERNIE fused QKNorm+RoPE fast path is not bit-exact on "
+                    "this platform; falling back to split kernels"
+                ),
+            )
+
+    return _ernie_qknorm_rope_reference(
+        q, k, q_norm, k_norm, head_dim, rope_cos, rope_sin
+    )
+
+
 def _eager_geglu(gate_up: torch.Tensor) -> torch.Tensor:
     gate, up = gate_up.chunk(2, dim=-1)
     return up * F.gelu(gate)
@@ -687,6 +785,10 @@ class ErnieImageTransformer2DModel(CachableDiT, LayerwiseOffloadableModuleMixin)
         all_ids = torch.cat([image_ids, text_ids], dim=1)
         rotary_pos_emb = self.pos_embed(all_ids)
         rope_cos, rope_sin = _precompute_rope_cos_sin(rotary_pos_emb, dtype)
+        rope_cache = torch.cat((rope_cos, rope_sin), dim=-1).contiguous()
+        rope_positions = torch.arange(
+            rope_cache.shape[0], device=device, dtype=torch.long
+        )
 
         attn_mask = attn_mask_meta = None
         if encoder_hidden_states_mask is not None:
@@ -715,6 +817,8 @@ class ErnieImageTransformer2DModel(CachableDiT, LayerwiseOffloadableModuleMixin)
                 x,
                 rope_cos,
                 rope_sin,
+                rope_cache,
+                rope_positions,
                 shift_msa,
                 scale_msa,
                 gate_msa,

test/registered/kernels/ops/diffusion/test_ernie_norm_scale_shift.py

@@ -1,6 +1,7 @@
 """ERNIE fused norm/scale/shift fast paths must stay bit-exact vs eager."""
 
 import sys
+from unittest.mock import patch
 
 import pytest
 import torch
@@ -10,6 +11,8 @@ from sglang.multimodal_gen.runtime.layers.layernorm import RMSNorm
 from sglang.multimodal_gen.runtime.models.dits.ernie_image import (
     _ernie_gated_norm_scale_shift,
     _ernie_norm_scale_shift,
+    _ernie_qknorm_rope,
+    _ernie_qknorm_rope_reference,
 )
 from sglang.test.ci.ci_register import register_cuda_ci
 
@@ -53,5 +56,80 @@ def test_fused_norm_scale_shift_is_bit_exact(shape):
     assert not ernie_image._ERNIE_GATED_NORM.disabled
 
 
+def test_fused_qknorm_rope_is_bit_exact():
+    torch.manual_seed(1)
+    ernie_image._ERNIE_QKNORM_ROPE.disabled = False
+    ernie_image._ERNIE_QKNORM_ROPE.verified = False
+    batch, seq, heads, head_dim = 1, 257, 32, 128
+    q = torch.randn(batch, seq, heads, head_dim, device="cuda", dtype=torch.bfloat16)
+    k = torch.randn_like(q)
+    q_norm = RMSNorm(head_dim, eps=1e-6).to(device="cuda", dtype=torch.bfloat16)
+    k_norm = RMSNorm(head_dim, eps=1e-6).to(device="cuda", dtype=torch.bfloat16)
+    cos = torch.randn(seq, head_dim, device="cuda", dtype=torch.bfloat16)
+    sin = torch.randn_like(cos)
+    cache = torch.cat((cos, sin), dim=-1).contiguous()
+    positions = torch.arange(seq, device="cuda", dtype=torch.long)
+
+    q_ref, k_ref = _ernie_qknorm_rope_reference(
+        q.clone(), k.clone(), q_norm, k_norm, head_dim, cos, sin
+    )
+    q_out, k_out = _ernie_qknorm_rope(
+        q,
+        k,
+        q_norm,
+        k_norm,
+        head_dim,
+        cos,
+        sin,
+        cache,
+        positions,
+    )
+
+    assert torch.equal(q_out, q_ref)
+    assert torch.equal(k_out, k_ref)
+    assert ernie_image._ERNIE_QKNORM_ROPE.verified
+    assert not ernie_image._ERNIE_QKNORM_ROPE.disabled
+
+
+def test_qknorm_rope_first_attempt_exception_uses_pristine_inputs():
+    torch.manual_seed(2)
+    ernie_image._ERNIE_QKNORM_ROPE.disabled = False
+    ernie_image._ERNIE_QKNORM_ROPE.verified = False
+    batch, seq, heads, head_dim = 1, 17, 4, 128
+    q = torch.randn(batch, seq, heads, head_dim, device="cuda", dtype=torch.bfloat16)
+    k = torch.randn_like(q)
+    q_norm = RMSNorm(head_dim, eps=1e-6).to(device="cuda", dtype=torch.bfloat16)
+    k_norm = RMSNorm(head_dim, eps=1e-6).to(device="cuda", dtype=torch.bfloat16)
+    cos = torch.randn(seq, head_dim, device="cuda", dtype=torch.bfloat16)
+    sin = torch.randn_like(cos)
+    cache = torch.cat((cos, sin), dim=-1).contiguous()
+    positions = torch.arange(seq, device="cuda", dtype=torch.long)
+    q_ref, k_ref = _ernie_qknorm_rope_reference(
+        q.clone(), k.clone(), q_norm, k_norm, head_dim, cos, sin
+    )
+
+    def mutate_then_raise(**kwargs):
+        kwargs["q"].zero_()
+        kwargs["k"].zero_()
+        raise RuntimeError("synthetic kernel failure")
+
+    with patch.object(ernie_image, "apply_qk_norm_rope", mutate_then_raise):
+        q_out, k_out = _ernie_qknorm_rope(
+            q,
+            k,
+            q_norm,
+            k_norm,
+            head_dim,
+            cos,
+            sin,
+            cache,
+            positions,
+        )
+
+    assert torch.equal(q_out, q_ref)
+    assert torch.equal(k_out, k_ref)
+    assert ernie_image._ERNIE_QKNORM_ROPE.disabled
+
+
 if __name__ == "__main__":
     sys.exit(pytest.main([__file__]))

test/registered/kernels/ops/diffusion/test_qknorm_rope.py

@@ -216,6 +216,45 @@ def test_qknorm_rope_preserves_split_bf16_rounding() -> None:
     assert torch.equal(k_ref, k_fused)
 
 
+def test_qknorm_rope_preserves_full_width_neox_cache() -> None:
+    from sglang.kernels.ops.diffusion.qknorm_rope import fused_inplace_qknorm_rope
+    from sglang.kernels.ops.layernorm.norm import fused_inplace_qknorm
+
+    num_tokens, num_heads, head_dim = 257, 32, 128
+    q = torch.randn(num_tokens, num_heads, head_dim, device=DEVICE, dtype=DTYPE)
+    k = torch.randn_like(q)
+    q_weight = torch.randn(head_dim, device=DEVICE, dtype=DTYPE)
+    k_weight = torch.randn(head_dim, device=DEVICE, dtype=DTYPE)
+    positions = torch.arange(num_tokens, device=DEVICE, dtype=torch.int64)
+    cos = torch.randn(num_tokens, head_dim, device=DEVICE, dtype=DTYPE)
+    sin = torch.randn_like(cos)
+    cache = torch.cat((cos, sin), dim=-1).contiguous()
+
+    q_ref, k_ref = q.clone(), k.clone()
+    fused_inplace_qknorm(q_ref, k_ref, q_weight, k_weight, eps=1e-6)
+    half = head_dim // 2
+    q1, q2 = q_ref[..., :half], q_ref[..., half:]
+    k1, k2 = k_ref[..., :half], k_ref[..., half:]
+    q_ref = torch.cat((-q2, q1), dim=-1) * sin[:, None, :] + q_ref * cos[:, None, :]
+    k_ref = torch.cat((-k2, k1), dim=-1) * sin[:, None, :] + k_ref * cos[:, None, :]
+
+    fused_inplace_qknorm_rope(
+        q,
+        k,
+        q_weight,
+        k_weight,
+        cache,
+        positions,
+        is_neox=True,
+        eps=1e-6,
+        round_norm_before_rope=True,
+        cache_has_full_width=True,
+    )
+
+    assert torch.equal(q, q_ref)
+    assert torch.equal(k, k_ref)
+
+
 def test_qknorm_rope_requires_opt_in_for_strided_packed_gqa() -> None:
     from sglang.kernels.ops.diffusion.qknorm_rope import (
         fused_inplace_qknorm_rope,