diff --git a/python/sglang/kernels/jit/csrc/diffusion/qknorm_rope.cuh b/python/sglang/kernels/jit/csrc/diffusion/qknorm_rope.cuh index a514ea82c..40c2b4fe0 100644 --- a/python/sglang/kernels/jit/csrc/diffusion/qknorm_rope.cuh +++ b/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 __grid_constant__ params) { using namespace device; @@ -185,7 +186,8 @@ __global__ void fused_qknorm_rope_warp(const QKNormRopeParamsT __grid_c constexpr uint32_t kRotaryLanes = kRopeDim / kElemsPerThread; constexpr uint32_t kHalfRotaryLanes = kRotaryLanes / 2; constexpr uint32_t kActiveMask = active_mask(); - 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 __grid_c auto output_vec = norm::apply_norm_warp(input_vec, weight_vec, eps); const auto pos = static_cast(static_cast(positions)[token_id]); const auto cos_ptr = static_cast(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 __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 cos = load_cache_value(cos_ptr, half_idx); - const auto sin = load_cache_value(sin_ptr, half_idx); + const auto cache_idx = + (kCacheHasFullWidth ? lane_id : lane_id % kHalfRotaryLanes) * kElemsPerThread + 2 * j + i; + 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 __grid_c const auto pos = static_cast(static_cast(positions)[token_id]); const auto cos_ptr = static_cast(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 __grid_c if (lane_id < kHalfRotaryLanes) { swapped = -swapped; } - const auto half_idx = (lane_id % kHalfRotaryLanes) * kElemsPerThread + i; - const float cos = cast(load_cache_value(cos_ptr, half_idx)); - const float sin = cast(load_cache_value(sin_ptr, half_idx)); + const auto cache_idx = (kCacheHasFullWidth ? lane_id : lane_id % kHalfRotaryLanes) * kElemsPerThread + i; + const float cos = cast(load_cache_value(cos_ptr, cache_idx)); + const float sin = cast(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 __grid_c const auto pos = static_cast(static_cast(positions)[token_id]); const auto cos_ptr = static_cast(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 @@ -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().with_device(device).verify(q); TensorMatcher({N, K, D}).with_strides({Dk, Dd, 1}).with_dtype().with_device(device).verify(k); TensorMatcher({D}).with_dtype().with_device(device).verify(q_weight).verify(k_weight); - TensorMatcher({-1, R}).with_dtype().with_device(device).verify(cos_sin_cache); + TensorMatcher({-1, kCacheHasFullWidth ? 2 * kRopeDim : kRopeDim}) + .with_dtype() + .with_device(device) + .verify(cos_sin_cache); TensorMatcher({N}).with_dtype(id_type).with_device(device).verify(positions); const auto num_tokens = static_cast(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 static constexpr auto kernel = fused_qknorm_rope_warp< kHeadDim, @@ -510,6 +520,7 @@ struct QKNormRopePackKVKernel { CacheDType, kRoundNormBeforeRope, true, + kCacheHasFullWidth, IdType>; static void diff --git a/python/sglang/kernels/ops/diffusion/qknorm_rope.py b/python/sglang/kernels/ops/diffusion/qknorm_rope.py index 37b0ccb05..32cd8ac58 100644 --- a/python/sglang/kernels/ops/diffusion/qknorm_rope.py +++ b/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), diff --git a/python/sglang/multimodal_gen/runtime/layers/layernorm.py b/python/sglang/multimodal_gen/runtime/layers/layernorm.py index a2202bec4..c89a9ad3b 100755 --- a/python/sglang/multimodal_gen/runtime/layers/layernorm.py +++ b/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 diff --git a/python/sglang/multimodal_gen/runtime/models/dits/ernie_image.py b/python/sglang/multimodal_gen/runtime/models/dits/ernie_image.py index 3165efaae..721e1ffe3 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/ernie_image.py +++ b/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, ) - - q = _ernie_rope(q, rope_cos, rope_sin) - k = _ernie_rope(k, rope_cos, rope_sin) + else: + q = _ernie_rope(q, 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, diff --git a/test/registered/kernels/ops/diffusion/test_ernie_norm_scale_shift.py b/test/registered/kernels/ops/diffusion/test_ernie_norm_scale_shift.py index 703ed0794..747a4d807 100644 --- a/test/registered/kernels/ops/diffusion/test_ernie_norm_scale_shift.py +++ b/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__])) diff --git a/test/registered/kernels/ops/diffusion/test_qknorm_rope.py b/test/registered/kernels/ops/diffusion/test_qknorm_rope.py index 87e3cc462..be3b67d6e 100644 --- a/test/registered/kernels/ops/diffusion/test_qknorm_rope.py +++ b/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,