import pytest import torch from sgl_kernel import fused_qk_norm_rope as sgl_fused_qk_norm_rope from sglang.srt.layers.layernorm import RMSNorm from sglang.srt.layers.rotary_embedding import get_rope from sglang.srt.server_args import ( ServerArgs, get_global_server_args, set_global_server_args_for_scheduler, ) from sglang.srt.utils import ( cpu_has_amx_support, is_cpu, is_cuda, is_hip, is_npu, is_xpu, ) _is_cuda = is_cuda() _is_hip = is_hip() _is_cpu = is_cpu() _is_cpu_amx_available = cpu_has_amx_support() _is_npu = is_npu() _is_xpu = is_xpu() device = torch.device("cuda" if torch.cuda.is_available() else "cpu") @torch.inference_mode() def torch_ref_rms_norm_rope( qkv, num_heads_q, num_heads_k, num_heads_v, head_dim, eps, q_weight, k_weight, base, is_neox, position_ids, partial_rotary_factor, ): """ PyTorch reference implementation of RMSNorm+RoPE for verification. Uses SGLang's own RMSNorm and RotaryEmbedding modules to ensure consistency with the expected behavior of the fused kernel. Args: qkv: Combined QKV tensor of shape [num_tokens, hidden_size] num_heads_q: Number of query heads num_heads_k: Number of key heads num_heads_v: Number of value heads (unused for normalization/RoPE but needed for tensor splitting) head_dim: Dimension of each head eps: Epsilon value for RMS normalization q_weight: RMSNorm weights for query [head_dim] k_weight: RMSNorm weights for key [head_dim] base: Base value for RoPE calculations is_neox: Whether to use NeoX style RoPE position_ids: Position IDs for RoPE of shape [num_tokens] partial_rotary_factor: Partial rotary factor Returns: Combined tensor with Q and K parts normalized and RoPE applied """ # Get input shape information num_tokens = qkv.shape[0] hidden_size = qkv.shape[1] # Calculate dimensions for Q, K, V segments q_size = num_heads_q * head_dim k_size = num_heads_k * head_dim v_size = num_heads_v * head_dim # Verify dimensions match assert ( hidden_size == q_size + k_size + v_size ), f"Hidden size {hidden_size} doesn't match Q+K+V dimensions {q_size + k_size + v_size}" # Split the tensor into Q, K, V parts q = qkv[:, :q_size] k = qkv[:, q_size : q_size + k_size] v = qkv[:, q_size + k_size :] rotary_emb = get_rope( head_dim, rotary_dim=head_dim, max_position=8192, base=10000, is_neox_style=is_neox, rope_scaling=None, dual_chunk_attention_config=None, partial_rotary_factor=partial_rotary_factor, ) rotary_emb = rotary_emb.to(qkv.device) # Create and apply RMSNorm modules with custom weights q_norm = RMSNorm(hidden_size=head_dim, eps=eps).to(qkv.device).to(qkv.dtype) q_norm.weight.data.copy_(q_weight) k_norm = RMSNorm(hidden_size=head_dim, eps=eps).to(qkv.device).to(qkv.dtype) k_norm.weight.data.copy_(k_weight) q_by_head = q.reshape(-1, head_dim) q_by_head = q_norm(q_by_head) k_by_head = k.reshape(-1, head_dim) k_by_head = k_norm(k_by_head) q = q_by_head.view(q.shape) k = k_by_head.view(k.shape) [q_rope, k_rope] = rotary_emb( position_ids, q, k, fused_set_kv_buffer_arg=None, ) # Combine Q, K, V back together result = torch.cat([q_rope, k_rope, v], dim=1) return result head_dims = [64, 128] # (Q heads, K heads, V heads) num_heads_groups = [ (16, 8, 8), # Qwen3-0.6B, Qwen3-1.7B (32, 8, 8), # Qwen3-4B, Qwen3-8B, Qwen3-30B-A3B (40, 8, 8), # Qwen3-14B (64, 8, 8), # Qwen3-32B, Qwen3-235B-A22B (12, 1, 1), # GLM4.6 TP8 ] num_tokens_list = [1, 3, 8, 32, 256] is_neox_list = [False, True] dtypes = [torch.bfloat16] partial_rotary_factor_list = [1.0, 0.5] @pytest.mark.skipif(not _is_cuda, reason="Skipping CUDA/ROCm only tests.") @pytest.mark.parametrize("head_dim", head_dims) @pytest.mark.parametrize("num_heads_group", num_heads_groups) @pytest.mark.parametrize("num_tokens", num_tokens_list) @pytest.mark.parametrize("is_neox", is_neox_list) @pytest.mark.parametrize("dtype", dtypes) @pytest.mark.parametrize("partial_rotary_factor", partial_rotary_factor_list) def test_fused_qk_norm_rope( head_dim, num_heads_group, num_tokens, is_neox, dtype, partial_rotary_factor ): """ Test the fused QK RMSNorm + RoPE operation with various configurations. This test verifies that the fused kernel correctly applies: 1. RMSNorm to both query (Q) and key (K) portions of the QKV tensor 2. Rotary Position Embeddings (RoPE) to the normalized Q and K 3. Leaves the value (V) portion unchanged Args: head_dim: Dimension of each attention head num_heads_group: Tuple of (num_heads_q, num_heads_k, num_heads_v) num_tokens: Number of tokens to process dtype: Data type (float16 or bfloat16) """ set_global_server_args_for_scheduler(ServerArgs(model_path="dummy")) device = "cuda" torch_dtype = dtype # Unpack head counts num_heads_q, num_heads_k, num_heads_v = num_heads_group # Calculate total hidden dimension hidden_size = (num_heads_q + num_heads_k + num_heads_v) * head_dim # Generate random inputs directly as 2D [num_tokens, hidden_size] torch.random.manual_seed(0) qkv = torch.randn(num_tokens, hidden_size, dtype=torch_dtype, device=device) qkv_copy = qkv.clone() # Generate position IDs with +100 offset to test decoding scenarios position_ids = torch.arange(num_tokens, dtype=torch.int32, device=device) + 100 # Generate random weights for RMSNorm q_weight = torch.randn(head_dim, dtype=torch_dtype, device=device) * 5.0 k_weight = torch.randn(head_dim, dtype=torch_dtype, device=device) * 5.0 # Set RMSNorm and RoPE parameters eps = 1e-5 base = 10000.0 factor, low, high, attention_factor = 1.0, 0, 0, 1.0 # Run the custom fusedQKNormRope operation sgl_fused_qk_norm_rope( qkv, num_heads_q, num_heads_k, num_heads_v, head_dim, eps, q_weight, k_weight, base, is_neox, position_ids, factor, low, high, attention_factor, int(head_dim * partial_rotary_factor), ) output = qkv # This op is inplace # Compute reference output using TensorRT-LLM modules ref_output = torch_ref_rms_norm_rope( qkv_copy, num_heads_q, num_heads_k, num_heads_v, head_dim, eps, q_weight, k_weight, base, is_neox, position_ids, partial_rotary_factor, ) # Compare outputs from custom kernel vs reference implementation torch.testing.assert_close( output, ref_output, rtol=5e-2, atol=1e-1, )