From 9081420cade5ded33664e9585160fb6685d2191b Mon Sep 17 00:00:00 2001
From: "Ruby Chen (Engrg-Hardware 1)" <rubchen@eos0123.eos.clusters.nvidia.com>
Date: Tue, 2 Dec 2025 07:02:26 -0800
Subject: [PATCH 01/37] Integrate rotary embedding CUDA kernel with
 comprehensive testing and benchmarking infrastructure.

---
 sgl-kernel/CMakeLists.txt                     |   1 +
 .../benchmark/bench_mm_rotary_embedding.py    | 197 ++++++++++++
 sgl-kernel/csrc/common_extension.cc           |   4 +
 .../csrc/multimodal/rotary_embedding.cu       | 283 ++++++++++++++++++
 sgl-kernel/include/sgl_kernel_ops.h           |  11 +
 .../python/sgl_kernel/rotary_embedding.py     |  23 ++
 sgl-kernel/tests/test_mm_rotary_embedding.py  | 189 ++++++++++++
 7 files changed, 708 insertions(+)
 create mode 100644 sgl-kernel/benchmark/bench_mm_rotary_embedding.py
 create mode 100644 sgl-kernel/csrc/multimodal/rotary_embedding.cu
 create mode 100644 sgl-kernel/python/sgl_kernel/rotary_embedding.py
 create mode 100644 sgl-kernel/tests/test_mm_rotary_embedding.py

diff --git a/sgl-kernel/CMakeLists.txt b/sgl-kernel/CMakeLists.txt
index 7b5d857a459b..6815479f1fef 100644
--- a/sgl-kernel/CMakeLists.txt
+++ b/sgl-kernel/CMakeLists.txt
@@ -310,6 +310,7 @@ set(SOURCES
     "csrc/kvcacheio/transfer.cu"
     "csrc/mamba/causal_conv1d.cu"
     "csrc/memory/store.cu"
+    "csrc/multimodal/rotary_embedding.cu"
     "csrc/memory/weak_ref_tensor.cpp"
 
     "csrc/moe/cutlass_moe/w4a8/scaled_mm_entry.cu"
diff --git a/sgl-kernel/benchmark/bench_mm_rotary_embedding.py b/sgl-kernel/benchmark/bench_mm_rotary_embedding.py
new file mode 100644
index 000000000000..ce36543e22a9
--- /dev/null
+++ b/sgl-kernel/benchmark/bench_mm_rotary_embedding.py
@@ -0,0 +1,197 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+
+from typing import List, Tuple
+
+import numpy as np
+import torch
+import triton
+
+from sgl_kernel.rotary_embedding import rotary_embedding as sgl_rotary_embedding
+
+def compute_cos_sin_cache(
+    max_seq_len: int,
+    rotary_dim: int,
+    base: float = 10000.0,
+    dtype: torch.dtype = torch.float32,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """
+    Compute separate cos and sin caches.
+
+    Returns:
+        cos, sin: shape (max_seq_len, rotary_dim / 2)
+    """
+    inv_freq = 1.0 / (
+        base ** (torch.arange(0, rotary_dim, 2, dtype=torch.float32) / rotary_dim)
+    )
+    t = torch.arange(max_seq_len, dtype=torch.float32)
+    freqs = torch.einsum("i,j->ij", t, inv_freq)
+    cos = freqs.cos().to(dtype)
+    sin = freqs.sin().to(dtype)
+    return cos, sin
+
+
+def benchmark_mm_rotary_embedding() -> None:
+    """Benchmark sgl_kernel.rotary_embedding vs vLLM & flash_attn (if available)."""
+    try:
+        from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding as vLLMRotaryEmbedding
+
+        HAS_VLLM = True
+    except ImportError:
+        vLLMRotaryEmbedding = None
+        HAS_VLLM = False
+        print("vLLM not available")
+
+    try:
+        from flash_attn.layers.rotary import RotaryEmbedding as FlashRotaryEmbedding
+
+        HAS_FLASH_ATTN = True
+    except ImportError:
+        FlashRotaryEmbedding = None
+        HAS_FLASH_ATTN = False
+        print("flash_attn not available")
+
+    device = "cuda"
+    dtype = torch.bfloat16
+    max_seq_len = 65536
+
+    seq_lens = [1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192]
+
+    # Test configurations: (batch_size, num_heads, num_kv_heads, head_size)
+    configs = [
+        # Decoding scenarios (batch_size=1)
+        (1, 32, 8, 128),
+        (1, 64, 8, 128),
+        (1, 32, 8, 64),
+        (1, 32, 8, 256),
+        (32, 32, 8, 64),
+        # # Edge cases
+        (1, 32, 1, 128),   # MQA
+        (32, 8, 8, 128),   # MHA
+        (1, 32, 8, 80),    # Non-standard head_size
+    ]
+
+    for batch_size, num_heads, num_kv_heads, head_size in configs:
+        try:
+            torch.cuda.synchronize()
+        except torch.AcceleratorError:
+            torch.cuda.empty_cache()
+            pass
+        
+        try:
+            rotary_dim = head_size
+            cos_cache, sin_cache = compute_cos_sin_cache(max_seq_len, rotary_dim, dtype=dtype)
+            cos_cache = cos_cache.to(device)
+            sin_cache = sin_cache.to(device)
+
+            print(f"\nConfig: batch_size={batch_size}, heads={num_heads}/{num_kv_heads}, head_size={head_size}, dtype={dtype}")
+            print("-" * 100)
+        except Exception as e:
+            print(f"\nSkipping config (batch_size={batch_size}, heads={num_heads}/{num_kv_heads}, head_size={head_size}): {e}")
+            continue
+
+        header = f"{'seq_len':>8}"
+        header += f" | {'ours (ms)':>10}"
+        if HAS_VLLM:
+            header += f" | {'vLLM (ms)':>10}"
+        if HAS_FLASH_ATTN:
+            header += f" | {'flash_attn (ms)':>14}"
+        header += f" | {'speedup':>9}"
+        print(header)
+        print("-" * 100)
+
+        results = []
+
+        for seq_len in seq_lens:
+            try:
+                num_tokens = batch_size * seq_len
+                query = torch.randn(num_tokens, num_heads * head_size, dtype=dtype, device=device)  
+                key = torch.randn(num_tokens, num_kv_heads * head_size, dtype=dtype, device=device)
+                positions = torch.arange(seq_len, device=device, dtype=torch.int64).repeat(batch_size)
+                cos = cos_cache[positions]
+                sin = sin_cache[positions]
+
+                row_str = f"{seq_len:8d}"
+                ours_time = None
+                vllm_time = None
+                fa_time = None
+            except (RuntimeError, torch.AcceleratorError) as e:
+                print(f"{seq_len:8d} | SKIP (CUDA error from previous iteration)")
+                torch.cuda.synchronize()
+                continue
+
+            # Ours: sgl_kernel.rotary_embedding (triton.testing.do_bench style)
+            def fn_ours() -> None:
+                q = query.clone()
+                k = key.clone()
+                sgl_rotary_embedding(cos, sin, q, k, head_size, True)
+
+            ms, _, _ = triton.testing.do_bench(fn_ours, quantiles=[0.5, 0.2, 0.8])
+            ours_time = 1000 * ms
+            row_str += f" | {ours_time:10.4f}"
+
+            # vLLM
+            if HAS_VLLM:
+                vllm_rope = vLLMRotaryEmbedding(
+                    head_size=head_size,
+                    rotary_dim=rotary_dim,
+                    max_position_embeddings=max_seq_len,
+                    base=10000,
+                    is_neox_style=True,
+                    dtype=dtype,
+                ).cuda()
+                
+                def fn_vllm() -> None:
+                    q = query.clone()
+                    k = key.clone()
+                    vllm_rope.forward_cuda(positions, q, k)
+
+                ms, _, _ = triton.testing.do_bench(fn_vllm, quantiles=[0.5, 0.2, 0.8])
+                vllm_time = 1000 * ms
+                row_str += f" | {vllm_time:10.4f}"
+            else:
+                row_str += f" | {'N/A':>10}"
+
+            # FlashAttention RotaryEmbedding (NeoX-style, QKV layout)
+            if HAS_FLASH_ATTN:
+                try:
+                    flash_rotary = FlashRotaryEmbedding(rotary_dim, device=device)
+                    qkv = torch.randn(batch_size, seq_len, 3, num_heads, head_size, dtype=dtype, device=device)
+
+                    def fn_flash_attn() -> None:
+                        qkv_fa = qkv.clone()
+                        flash_rotary(qkv_fa, seqlen_offset=0)
+
+                    ms, _, _ = triton.testing.do_bench(fn_flash_attn, quantiles=[0.5, 0.2, 0.8])
+                    fa_time = 1000 * ms
+                    row_str += f" | {fa_time:14.4f}"
+                except Exception:
+                    row_str += f" | {'ERROR':>14}"
+                    fa_time = None
+            else:
+                row_str += f" | {'N/A':>14}"
+
+            if HAS_VLLM and ours_time is not None and vllm_time is not None:
+                speedup = vllm_time / ours_time
+                row_str += f" | {speedup:9.2f}x"
+
+            print(row_str)
+            results.append({"seq_len": seq_len, "ours": ours_time, "vllm": vllm_time, "flash_attn": fa_time})
+
+        print("-" * 100)
+        print("\nAverage time:")
+        ours_vals = [r["ours"] for r in results if r["ours"] is not None]
+        if ours_vals:
+            print(f"  SGLang: {np.mean(ours_vals):.4f} ms")
+        if HAS_VLLM:
+            vllm_vals = [r["vllm"] for r in results if r["vllm"] is not None]
+            if vllm_vals:
+                print(f"  vLLM: {np.mean(vllm_vals):.4f} ms")
+        if HAS_FLASH_ATTN:
+            fa_vals = [r["flash_attn"] for r in results if r["flash_attn"] is not None]
+            if fa_vals:
+                print(f"  flash_attn: {np.mean(fa_vals):.4f} ms")
+   
+
+if __name__ == "__main__":
+    benchmark_mm_rotary_embedding()
\ No newline at end of file
diff --git a/sgl-kernel/csrc/common_extension.cc b/sgl-kernel/csrc/common_extension.cc
index 1344675b3c80..a92a66b497a9 100644
--- a/sgl-kernel/csrc/common_extension.cc
+++ b/sgl-kernel/csrc/common_extension.cc
@@ -90,6 +90,10 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) {
       "Tensor? v, Tensor!? k_buffer, Tensor!? v_buffer, Tensor? kv_cache_loc) -> ()");
   m.impl("apply_rope_pos_ids_cos_sin_cache", torch::kCUDA, &apply_rope_pos_ids_cos_sin_cache);
 
+  m.def(
+      "rotary_embedding(Tensor cos, Tensor sin, Tensor(a!) query, Tensor? key, int head_size, bool is_neox) -> ()");
+  m.impl("rotary_embedding", torch::kCUDA, &rotary_embedding);
+
   m.def(
       "downcast_fp8(Tensor k, Tensor v, Tensor k_out, Tensor v_out, Tensor k_scale, Tensor v_scale, Tensor loc, "
       "int mult, int offset) -> ()");
diff --git a/sgl-kernel/csrc/multimodal/rotary_embedding.cu b/sgl-kernel/csrc/multimodal/rotary_embedding.cu
new file mode 100644
index 000000000000..2356765742ac
--- /dev/null
+++ b/sgl-kernel/csrc/multimodal/rotary_embedding.cu
@@ -0,0 +1,283 @@
+/*
+ * Copyright (c) 2025 by SGLang team.
+ * Copyright (c) 2025 by FlashInfer team.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAGuard.h>
+#include <c10/cuda/CUDAException.h>
+#include <c10/cuda/CUDAStream.h>
+#include <torch/all.h>
+
+#include <cmath>
+#include "utils.h"
+
+template <typename scalar_t, bool IS_NEOX>
+inline __device__ void apply_token_rotary_embedding(
+    scalar_t* __restrict__ arr,
+    const scalar_t* __restrict__ cos_ptr,
+    const scalar_t* __restrict__ sin_ptr,
+    int rot_offset,
+    int embed_dim) {
+  int x_index, y_index;
+
+  if (IS_NEOX) {
+    // NeoX-style: interleaved layout [x0, y0, x1, y1, ...].
+    // For NeoX, cos/sin have shape [..., rotary_dim/2]; each index corresponds
+    // to one (x,y) pair.
+    x_index = 2 * rot_offset;
+    y_index = x_index + 1;
+
+    const float cos_val = static_cast<float>(SGLANG_LDG(cos_ptr + rot_offset));
+    const float sin_val = static_cast<float>(SGLANG_LDG(sin_ptr + rot_offset));
+
+    const float x = static_cast<float>(arr[x_index]);
+    const float y = static_cast<float>(arr[y_index]);
+    arr[x_index] = static_cast<scalar_t>(x * cos_val - y * sin_val);
+    arr[y_index] = static_cast<scalar_t>(y * cos_val + x * sin_val);
+
+  } else {
+    // GPT-J / LLaMA style when cos/sin are [..., rotary_dim], i.e. non-interleaved
+    // [x0, x1, ..., y0, y1, ...]. Here embed_dim is the "half" dimension and
+    // cos/sin have length 2 * embed_dim.
+    x_index = rot_offset; 
+    y_index = rot_offset + embed_dim;
+
+    const float cos_val_x =
+        static_cast<float>(SGLANG_LDG(cos_ptr + rot_offset));
+    const float sin_val_x =
+        static_cast<float>(SGLANG_LDG(sin_ptr + rot_offset));
+    const float cos_val_y =
+        static_cast<float>(SGLANG_LDG(cos_ptr + rot_offset + embed_dim));
+    const float sin_val_y =
+        static_cast<float>(SGLANG_LDG(sin_ptr + rot_offset + embed_dim));
+
+    const float x = static_cast<float>(arr[x_index]);
+    const float y = static_cast<float>(arr[y_index]);
+    arr[x_index] = static_cast<scalar_t>(x * cos_val_x - y * sin_val_x);
+    arr[y_index] = static_cast<scalar_t>(y * cos_val_y + x * sin_val_y);
+  }
+}
+
+template <typename scalar_t, bool IS_NEOX>
+inline __device__ void apply_rotary_embedding(
+    scalar_t* __restrict__ query,                        // [num_heads, head_size]
+    scalar_t* __restrict__ key,                          // [num_kv_heads, head_size]
+    const scalar_t* __restrict__ current_token_cos_ptr,  // [rot_dim]
+    const scalar_t* __restrict__ current_token_sin_ptr,  // [rot_dim]
+    const int head_size,
+    const int num_heads,
+    const int num_kv_heads,
+    const int rot_dim,
+    const int64_t head_stride_query,
+    const int64_t head_stride_key) {
+  const int embed_dim_for_rotation = IS_NEOX ? rot_dim : (rot_dim / 2);
+
+  const int nq_pairs = num_heads * embed_dim_for_rotation;
+  for (int i = threadIdx.x; i < nq_pairs; i += blockDim.x) {
+    const int head_idx = i / embed_dim_for_rotation;
+    const int rot_offset = i % embed_dim_for_rotation;
+
+    scalar_t* query_for_token_head = query + head_idx * (int)head_stride_query;
+
+    apply_token_rotary_embedding<scalar_t, IS_NEOX>(query_for_token_head, current_token_cos_ptr, current_token_sin_ptr, rot_offset, embed_dim_for_rotation);
+  }
+
+  if (key != nullptr) {
+    const int nk_pairs = num_kv_heads * embed_dim_for_rotation;
+    for (int i = threadIdx.x; i < nk_pairs; i += blockDim.x) {
+      const int head_idx = i / embed_dim_for_rotation;
+      const int rot_offset = i % embed_dim_for_rotation;
+
+      scalar_t* key_for_token_head = key + head_idx * (int)head_stride_key;
+
+      apply_token_rotary_embedding<scalar_t, IS_NEOX>(key_for_token_head, current_token_cos_ptr, current_token_sin_ptr, rot_offset, embed_dim_for_rotation);
+    }
+  }
+}
+
+template <typename scalar_t, bool IS_NEOX>
+__global__ void rotary_embedding_kernel(
+    const scalar_t* __restrict__ cos_data,  // [num_tokens, rot_dim_arg]
+    const scalar_t* __restrict__ sin_data,  // [num_tokens, rot_dim_arg]
+    scalar_t* __restrict__ query_total,
+    scalar_t* __restrict__ key_total,
+    const int rot_dim_arg,
+    const int64_t query_token_stride,
+    const int64_t key_token_stride,
+    const int64_t head_stride_query,
+    const int64_t head_stride_key,
+    const int num_heads,
+    const int num_kv_heads,
+    const int head_size) {
+  const int token_idx = blockIdx.x;
+  const scalar_t* current_token_cos_ptr = cos_data + token_idx * rot_dim_arg;
+  const scalar_t* current_token_sin_ptr = sin_data + token_idx * rot_dim_arg;
+
+  scalar_t* query_for_token = query_total + token_idx * (int)query_token_stride;
+  scalar_t* key_for_token = (key_total != nullptr) ? (key_total + token_idx * (int)key_token_stride) : nullptr;
+
+  apply_rotary_embedding<scalar_t, IS_NEOX>(
+      query_for_token,
+      key_for_token,
+      current_token_cos_ptr,
+      current_token_sin_ptr,
+      head_size,
+      num_heads,
+      num_kv_heads,
+      rot_dim_arg,
+      head_stride_query,
+      head_stride_key);
+}
+
+void rotary_embedding(
+    at::Tensor& cos,
+    at::Tensor& sin,
+    at::Tensor& query,
+    const std::optional<at::Tensor>& key,
+    int64_t head_size,
+    bool is_neox) {
+  TORCH_CHECK(
+      query.dim() == 2 || query.dim() == 3,
+      "query must be in  shape [num_tokens, hidden_size] or [num_tokens, num_heads, head_size]");
+  if (key.has_value()) {
+    TORCH_CHECK(
+        key->dim() == 2 || key->dim() == 3,
+        "key must be in  shape [num_tokens, hidden_size] or [num_tokens, num_kv_heads, head_size]");
+  }
+
+  int64_t num_tokens = query.size(0);
+
+  TORCH_CHECK(cos.dim() == 2, "cos must be in shape [num_tokens, D_cos]");
+  TORCH_CHECK(sin.dim() == 2, "sin must be in  shape [num_tokens, D_sin]");
+  TORCH_CHECK(cos.size(0) == num_tokens, "cos num_tokens mismatch with query");
+  TORCH_CHECK(sin.size(0) == num_tokens, "sin num_tokens mismatch with query");
+  TORCH_CHECK(cos.size(1) == sin.size(1), "cos and sin D_cos/D_sin mismatch");
+
+  TORCH_CHECK(cos.scalar_type() == query.scalar_type(), "cos dtype mismatch");
+  TORCH_CHECK(sin.scalar_type() == query.scalar_type(), "sin dtype mismatch");
+  TORCH_CHECK(cos.is_cuda() && sin.is_cuda() && query.is_cuda(), "All tensors must be on CUDA");
+  TORCH_CHECK(
+      query.is_contiguous(),
+      "query must be contiguous; got non-contiguous tensor");
+  if (key.has_value()) {
+    TORCH_CHECK(key->is_cuda(), "Key tensor must be on CUDA if provided");
+    TORCH_CHECK(key->scalar_type() == query.scalar_type(), "Key dtype mismatch");
+    TORCH_CHECK(
+        key->is_contiguous(),
+        "key must be contiguous when provided; got non-contiguous tensor");
+  }
+
+  int query_hidden_size_calculated;
+  if (query.dim() == 2) {
+    query_hidden_size_calculated = (int)query.size(1);
+  } else {
+    query_hidden_size_calculated = (int)query.size(1) * (int)query.size(2);
+    TORCH_CHECK(query.size(2) == head_size, "Query head_size mismatch in 3D tensor");
+  }
+  TORCH_CHECK(query_hidden_size_calculated % head_size == 0, "query_hidden_size not divisible by head_size");
+  int num_heads = (int)query_hidden_size_calculated / (int)head_size;
+
+  int key_hidden_size_calculated = 0;
+  int num_kv_heads = num_heads;
+  if (key.has_value()) {
+    TORCH_CHECK((int)key->size(0) == num_tokens, "Key num_tokens mismatch");
+    if (key->dim() == 2) {
+      key_hidden_size_calculated = (int)key->size(1);
+    } else {
+      key_hidden_size_calculated = (int)key->size(1) * (int)key->size(2);
+      TORCH_CHECK((int)key->size(2) == head_size, "Key head_size mismatch in 3D tensor");
+    }
+    TORCH_CHECK(key_hidden_size_calculated % head_size == 0, "key_hidden_size not divisible by head_size");
+    num_kv_heads = key_hidden_size_calculated / (int)head_size;
+  }
+  TORCH_CHECK(num_heads % num_kv_heads == 0, "num_heads must be divisible by num_kv_heads");
+
+  int rot_dim_from_cache = (int)cos.size(1);
+
+  int64_t query_token_stride = query_hidden_size_calculated;
+  int64_t key_token_stride = key.has_value() ? key_hidden_size_calculated : 0;
+
+  int64_t head_stride_query;
+  if (query.dim() == 3 && query.size(1) == num_heads && query.size(2) == head_size) {
+    head_stride_query = query.stride(1);
+  } else {
+    head_stride_query = head_size;
+  }
+
+  int64_t head_stride_key = head_size;
+  if (key.has_value()) {
+    if (key->dim() == 3 && key->size(1) == num_kv_heads && key->size(2) == head_size) {
+      head_stride_key = key->stride(1);
+    } else {
+      head_stride_key = head_size;
+    }
+  }
+
+  dim3 grid((int)num_tokens);
+
+  // Number of (x,y) pairs rotated per head:
+  //  - NeoX: cos.size(1) = rotary_dim/2  => pairs = rot_dim_from_cache
+  //  - GPT-J: cos.size(1) = rotary_dim   => pairs = rot_dim_from_cache / 2
+  int embed_dim_for_block_calc = is_neox ? rot_dim_from_cache : (rot_dim_from_cache / 2);
+  int max_pairs_to_rotate_per_token = std::max(num_heads * embed_dim_for_block_calc, num_kv_heads * embed_dim_for_block_calc);
+  dim3 block(std::min<int64_t>(max_pairs_to_rotate_per_token, 512L));
+
+  if (block.x == 0 && num_tokens > 0) block.x = 1;
+
+  const at::cuda::OptionalCUDAGuard device_guard(device_of(query));
+  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+
+  AT_DISPATCH_FLOATING_TYPES_AND2(
+      at::ScalarType::Half, at::ScalarType::BFloat16,
+      query.scalar_type(), "rotary_embedding", [&] {
+    using cuda_scalar_t = typename std::conditional<
+        std::is_same<scalar_t, at::Half>::value, nv_half,
+        typename std::conditional<
+            std::is_same<scalar_t, at::BFloat16>::value, nv_bfloat16,
+            scalar_t>::type>::type;
+    
+    if (is_neox) {
+      rotary_embedding_kernel<cuda_scalar_t, true><<<grid, block, 0, stream>>>(
+          reinterpret_cast<cuda_scalar_t*>(cos.data_ptr<scalar_t>()),
+          reinterpret_cast<cuda_scalar_t*>(sin.data_ptr<scalar_t>()),
+          reinterpret_cast<cuda_scalar_t*>(query.data_ptr<scalar_t>()),
+          key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<scalar_t>()) : nullptr,
+          rot_dim_from_cache,
+          query_token_stride,
+          key_token_stride,
+          head_stride_query,
+          head_stride_key,
+          num_heads,
+          num_kv_heads,
+          (int)head_size);
+    } else {
+      rotary_embedding_kernel<cuda_scalar_t, false><<<grid, block, 0, stream>>>(
+          reinterpret_cast<cuda_scalar_t*>(cos.data_ptr<scalar_t>()),
+          reinterpret_cast<cuda_scalar_t*>(sin.data_ptr<scalar_t>()),
+          reinterpret_cast<cuda_scalar_t*>(query.data_ptr<scalar_t>()),
+          key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<scalar_t>()) : nullptr,
+          rot_dim_from_cache,
+          query_token_stride,
+          key_token_stride,
+          head_stride_query,
+          head_stride_key,
+          num_heads,
+          num_kv_heads,
+          (int)head_size);
+    }
+  });
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+}
\ No newline at end of file
diff --git a/sgl-kernel/include/sgl_kernel_ops.h b/sgl-kernel/include/sgl_kernel_ops.h
index 552519fa49f0..0c8438135d4e 100644
--- a/sgl-kernel/include/sgl_kernel_ops.h
+++ b/sgl-kernel/include/sgl_kernel_ops.h
@@ -395,6 +395,17 @@ void fused_qk_norm_rope(
     double high,
     double attention_factor);
 
+/*
+ * From csrc/multimodal/rotary_embedding.cu
+ */
+void rotary_embedding(
+    at::Tensor& cos,
+    at::Tensor& sin,
+    at::Tensor& query,
+    const std::optional<at::Tensor>& key,
+    int64_t head_size,
+    bool is_neox);
+
 void cutlass_fp4_group_mm(
     torch::Tensor& output,
     const torch::Tensor& a,
diff --git a/sgl-kernel/python/sgl_kernel/rotary_embedding.py b/sgl-kernel/python/sgl_kernel/rotary_embedding.py
new file mode 100644
index 000000000000..6d5072e68461
--- /dev/null
+++ b/sgl-kernel/python/sgl_kernel/rotary_embedding.py
@@ -0,0 +1,23 @@
+from typing import Optional
+
+import torch
+
+
+# Adapted from https://github.com/vllm-project/vllm/blob/9214e60631a79506e7669650de87806a123e0b0b/vllm/_custom_ops.py#L249
+# pos encoding ops
+def rotary_embedding(
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    query: torch.Tensor,
+    key: Optional[torch.Tensor],
+    head_size: int,
+    is_neox: bool,
+) -> None:
+    torch.ops.sgl_kernel.rotary_embedding.default(
+        cos,
+        sin,
+        query,
+        key,
+        head_size,
+        is_neox,
+    )
diff --git a/sgl-kernel/tests/test_mm_rotary_embedding.py b/sgl-kernel/tests/test_mm_rotary_embedding.py
new file mode 100644
index 000000000000..3fb0a5bb2297
--- /dev/null
+++ b/sgl-kernel/tests/test_mm_rotary_embedding.py
@@ -0,0 +1,189 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+
+import os
+import sys
+from dataclasses import dataclass
+from typing import Tuple, List
+
+import torch
+
+try:
+    import pytest
+    HAS_PYTEST = True
+except ImportError:
+    HAS_PYTEST = False
+
+try:
+    from sgl_kernel.rotary_embedding import rotary_embedding as rotary_emb_module
+    HAS_ROTARY_EMBEDDING = True
+except ImportError:
+    rotary_emb_module = None
+    HAS_ROTARY_EMBEDDING = False
+    print("sgl_kernel.rotary_embedding not available")
+
+
+@dataclass
+class RotaryTestResult:
+    name: str
+    passed: bool
+    q_diff: float = 0.0
+    k_diff: float = 0.0
+    details: str = ""
+
+
+def compute_cos_sin_cache(
+    max_seq_len: int,
+    rotary_dim: int,
+    base: float = 10000.0,
+    dtype: torch.dtype = torch.float32,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """
+    Compute separate cos and sin caches.
+
+    Returns:
+        cos, sin: shape (max_seq_len, rotary_dim / 2)
+    """
+    inv_freq = 1.0 / (
+        base ** (torch.arange(0, rotary_dim, 2, dtype=torch.float32) / rotary_dim)
+    )
+    t = torch.arange(max_seq_len, dtype=torch.float32)
+    freqs = torch.einsum("i,j->ij", t, inv_freq)
+    cos = freqs.cos().to(dtype)
+    sin = freqs.sin().to(dtype)
+    return cos, sin
+
+
+def reference_rotary_neox(
+    query: torch.Tensor,  # (num_tokens, num_heads * head_size)
+    key: torch.Tensor,  # (num_tokens, num_kv_heads * head_size)
+    cos: torch.Tensor,  # (num_tokens, rotary_dim / 2)
+    sin: torch.Tensor,  # (num_tokens, rotary_dim / 2)
+    head_size: int,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """
+    Reference implementation of NeoX-style rotary embedding, using explicit
+    cos/sin caches and Q/K in flattened (num_tokens, heads * dim) layout.
+    """
+    num_tokens = query.size(0)
+    num_heads = query.size(1) // head_size
+    num_kv_heads = key.size(1) // head_size
+    rotary_dim = cos.size(1) * 2
+
+    q = query.view(num_tokens, num_heads, head_size).float()
+    k = key.view(num_tokens, num_kv_heads, head_size).float()
+
+    q_rot = q[..., :rotary_dim]
+    q_pass = q[..., rotary_dim:]
+    k_rot = k[..., :rotary_dim]
+    k_pass = k[..., rotary_dim:]
+
+    q_rot = q_rot.view(num_tokens, num_heads, rotary_dim // 2, 2)
+    k_rot = k_rot.view(num_tokens, num_kv_heads, rotary_dim // 2, 2)
+
+    cos_expanded = cos.float().unsqueeze(1)  # (tokens, 1, rotary_dim/2)
+    sin_expanded = sin.float().unsqueeze(1)
+
+    # Apply rotary: x' = x*cos - y*sin, y' = y*cos + x*sin
+    q_x = q_rot[..., 0]
+    q_y = q_rot[..., 1]
+    q_rot_out = torch.stack(
+        [
+            q_x * cos_expanded - q_y * sin_expanded,
+            q_y * cos_expanded + q_x * sin_expanded,
+        ],
+        dim=-1,
+    )
+
+    k_x = k_rot[..., 0]
+    k_y = k_rot[..., 1]
+    k_rot_out = torch.stack(
+        [
+            k_x * cos_expanded - k_y * sin_expanded,
+            k_y * cos_expanded + k_x * sin_expanded,
+        ],
+        dim=-1,
+    )
+
+    q_rot_out = q_rot_out.view(num_tokens, num_heads, rotary_dim)
+    k_rot_out = k_rot_out.view(num_tokens, num_kv_heads, rotary_dim)
+
+    q_out = torch.cat([q_rot_out, q_pass], dim=-1)
+    k_out = torch.cat([k_rot_out, k_pass], dim=-1)
+
+    q_out = q_out.view(num_tokens, num_heads * head_size).to(query.dtype)
+    k_out = k_out.view(num_tokens, num_kv_heads * head_size).to(key.dtype)
+
+    return q_out, k_out
+
+
+def _check_rotary_correctness(
+    batch_size: int = 2,
+    seq_len: int = 128,
+    num_heads: int = 32,
+    num_kv_heads: int = 8,
+    head_size: int = 128,
+    dtype: torch.dtype = torch.bfloat16,
+    tol: float = 1e-2,
+    device: str = "cuda",
+) -> RotaryTestResult:
+    """Helper function for rotary correctness checking."""
+    if not HAS_ROTARY_EMBEDDING:
+        return RotaryTestResult("basic (skipped: rotary_embedding not built)", True, details="rotary_embedding extension not found")
+
+    rotary_dim = head_size
+    max_seq_len = 8192
+    num_tokens = batch_size * seq_len
+
+    query = torch.randn(num_tokens, num_heads * head_size, dtype=dtype, device=device)
+    key = torch.randn(num_tokens, num_kv_heads * head_size, dtype=dtype, device=device)
+    cos_cache, sin_cache = compute_cos_sin_cache(max_seq_len, rotary_dim, dtype=dtype)
+    cos_cache = cos_cache.to(device)
+    sin_cache = sin_cache.to(device)
+    positions = torch.arange(seq_len, device=device).repeat(batch_size)
+    cos = cos_cache[positions]
+    sin = sin_cache[positions]
+
+    q_ref, k_ref = reference_rotary_neox(query.clone(), key.clone(), cos, sin, head_size)
+
+    q_out = query.clone()
+    k_out = key.clone()
+    rotary_emb_module(cos, sin, q_out, k_out, head_size, True)  # is_neox=True
+
+    q_diff = (q_out - q_ref).abs().max().item()
+    k_diff = (k_out - k_ref).abs().max().item()
+    passed = q_diff < tol and k_diff < tol
+
+    name = f"basic [bs={batch_size}, seq={seq_len}, heads={num_heads}/{num_kv_heads}, dim={head_size}, {dtype}]"
+    return RotaryTestResult(name, passed, q_diff, k_diff)
+
+
+@pytest.mark.parametrize("batch_size", [2, 32, 1])
+@pytest.mark.parametrize("seq_len", [1, 128, 512, 2048])
+@pytest.mark.parametrize("num_heads, num_kv_heads", [(32, 8), (64, 8), (8, 8), (32, 1)])
+@pytest.mark.parametrize("head_size", [64, 128, 256, 80, 320])
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
+def test_rotary_embedding_correctness(
+    batch_size: int,
+    seq_len: int,
+    num_heads: int,
+    num_kv_heads: int,
+    head_size: int,
+    dtype: torch.dtype,
+) -> None:
+    if not HAS_ROTARY_EMBEDDING:
+        pytest.skip("sgl_kernel.rotary_embedding not available")
+    
+    result = _check_rotary_correctness(
+        batch_size=batch_size,
+        seq_len=seq_len,
+        num_heads=num_heads,
+        num_kv_heads=num_kv_heads,
+        head_size=head_size,
+        dtype=dtype,
+    )
+    assert result.passed, f"{result.name} failed: Q={result.q_diff:.2e}, K={result.k_diff:.2e}"
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])

From 2aff99b60004f7d8171c1c3624488f838617759e Mon Sep 17 00:00:00 2001
From: RubiaCx <1084281732@qq.com>
Date: Tue, 2 Dec 2025 07:35:44 -0800
Subject: [PATCH 02/37] clean and format

---
 .../csrc/multimodal/rotary_embedding.cu       | 115 ++++++++----------
 sgl-kernel/include/sgl_kernel_ops.h           |   2 +-
 .../python/sgl_kernel/rotary_embedding.py     |   3 -
 sgl-kernel/tests/test_mm_rotary_embedding.py  |  12 +-
 4 files changed, 58 insertions(+), 74 deletions(-)

diff --git a/sgl-kernel/csrc/multimodal/rotary_embedding.cu b/sgl-kernel/csrc/multimodal/rotary_embedding.cu
index 2356765742ac..803995960b47 100644
--- a/sgl-kernel/csrc/multimodal/rotary_embedding.cu
+++ b/sgl-kernel/csrc/multimodal/rotary_embedding.cu
@@ -16,12 +16,13 @@
  */
 
 #include <ATen/cuda/CUDAContext.h>
-#include <c10/cuda/CUDAGuard.h>
 #include <c10/cuda/CUDAException.h>
+#include <c10/cuda/CUDAGuard.h>
 #include <c10/cuda/CUDAStream.h>
 #include <torch/all.h>
 
 #include <cmath>
+
 #include "utils.h"
 
 template <typename scalar_t, bool IS_NEOX>
@@ -35,8 +36,8 @@ inline __device__ void apply_token_rotary_embedding(
 
   if (IS_NEOX) {
     // NeoX-style: interleaved layout [x0, y0, x1, y1, ...].
-    // For NeoX, cos/sin have shape [..., rotary_dim/2]; each index corresponds
-    // to one (x,y) pair.
+    // For NeoX, cos/sin have shape [..., rotary_dim/2]; 
+    // each index corresponds to one (x,y) pair.
     x_index = 2 * rot_offset;
     y_index = x_index + 1;
 
@@ -50,19 +51,14 @@ inline __device__ void apply_token_rotary_embedding(
 
   } else {
     // GPT-J / LLaMA style when cos/sin are [..., rotary_dim], i.e. non-interleaved
-    // [x0, x1, ..., y0, y1, ...]. Here embed_dim is the "half" dimension and
-    // cos/sin have length 2 * embed_dim.
-    x_index = rot_offset; 
+    // [x0, x1, ..., y0, y1, ...]. Here embed_dim is the "half" dimension and cos/sin have length 2 * embed_dim.
+    x_index = rot_offset;
     y_index = rot_offset + embed_dim;
 
-    const float cos_val_x =
-        static_cast<float>(SGLANG_LDG(cos_ptr + rot_offset));
-    const float sin_val_x =
-        static_cast<float>(SGLANG_LDG(sin_ptr + rot_offset));
-    const float cos_val_y =
-        static_cast<float>(SGLANG_LDG(cos_ptr + rot_offset + embed_dim));
-    const float sin_val_y =
-        static_cast<float>(SGLANG_LDG(sin_ptr + rot_offset + embed_dim));
+    const float cos_val_x = static_cast<float>(SGLANG_LDG(cos_ptr + rot_offset));
+    const float sin_val_x = static_cast<float>(SGLANG_LDG(sin_ptr + rot_offset));
+    const float cos_val_y = static_cast<float>(SGLANG_LDG(cos_ptr + rot_offset + embed_dim));
+    const float sin_val_y = static_cast<float>(SGLANG_LDG(sin_ptr + rot_offset + embed_dim));
 
     const float x = static_cast<float>(arr[x_index]);
     const float y = static_cast<float>(arr[y_index]);
@@ -92,7 +88,8 @@ inline __device__ void apply_rotary_embedding(
 
     scalar_t* query_for_token_head = query + head_idx * (int)head_stride_query;
 
-    apply_token_rotary_embedding<scalar_t, IS_NEOX>(query_for_token_head, current_token_cos_ptr, current_token_sin_ptr, rot_offset, embed_dim_for_rotation);
+    apply_token_rotary_embedding<scalar_t, IS_NEOX>(
+        query_for_token_head, current_token_cos_ptr, current_token_sin_ptr, rot_offset, embed_dim_for_rotation);
   }
 
   if (key != nullptr) {
@@ -103,7 +100,8 @@ inline __device__ void apply_rotary_embedding(
 
       scalar_t* key_for_token_head = key + head_idx * (int)head_stride_key;
 
-      apply_token_rotary_embedding<scalar_t, IS_NEOX>(key_for_token_head, current_token_cos_ptr, current_token_sin_ptr, rot_offset, embed_dim_for_rotation);
+      apply_token_rotary_embedding<scalar_t, IS_NEOX>(
+          key_for_token_head, current_token_cos_ptr, current_token_sin_ptr, rot_offset, embed_dim_for_rotation);
     }
   }
 }
@@ -169,15 +167,11 @@ void rotary_embedding(
   TORCH_CHECK(cos.scalar_type() == query.scalar_type(), "cos dtype mismatch");
   TORCH_CHECK(sin.scalar_type() == query.scalar_type(), "sin dtype mismatch");
   TORCH_CHECK(cos.is_cuda() && sin.is_cuda() && query.is_cuda(), "All tensors must be on CUDA");
-  TORCH_CHECK(
-      query.is_contiguous(),
-      "query must be contiguous; got non-contiguous tensor");
+  TORCH_CHECK(query.is_contiguous(), "query must be contiguous; got non-contiguous tensor");
   if (key.has_value()) {
     TORCH_CHECK(key->is_cuda(), "Key tensor must be on CUDA if provided");
     TORCH_CHECK(key->scalar_type() == query.scalar_type(), "Key dtype mismatch");
-    TORCH_CHECK(
-        key->is_contiguous(),
-        "key must be contiguous when provided; got non-contiguous tensor");
+    TORCH_CHECK(key->is_contiguous(), "key must be contiguous when provided; got non-contiguous tensor");
   }
 
   int query_hidden_size_calculated;
@@ -232,7 +226,8 @@ void rotary_embedding(
   //  - NeoX: cos.size(1) = rotary_dim/2  => pairs = rot_dim_from_cache
   //  - GPT-J: cos.size(1) = rotary_dim   => pairs = rot_dim_from_cache / 2
   int embed_dim_for_block_calc = is_neox ? rot_dim_from_cache : (rot_dim_from_cache / 2);
-  int max_pairs_to_rotate_per_token = std::max(num_heads * embed_dim_for_block_calc, num_kv_heads * embed_dim_for_block_calc);
+  int max_pairs_to_rotate_per_token =
+      std::max(num_heads * embed_dim_for_block_calc, num_kv_heads * embed_dim_for_block_calc);
   dim3 block(std::min<int64_t>(max_pairs_to_rotate_per_token, 512L));
 
   if (block.x == 0 && num_tokens > 0) block.x = 1;
@@ -241,43 +236,41 @@ void rotary_embedding(
   const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
 
   AT_DISPATCH_FLOATING_TYPES_AND2(
-      at::ScalarType::Half, at::ScalarType::BFloat16,
-      query.scalar_type(), "rotary_embedding", [&] {
-    using cuda_scalar_t = typename std::conditional<
-        std::is_same<scalar_t, at::Half>::value, nv_half,
-        typename std::conditional<
-            std::is_same<scalar_t, at::BFloat16>::value, nv_bfloat16,
-            scalar_t>::type>::type;
-    
-    if (is_neox) {
-      rotary_embedding_kernel<cuda_scalar_t, true><<<grid, block, 0, stream>>>(
-          reinterpret_cast<cuda_scalar_t*>(cos.data_ptr<scalar_t>()),
-          reinterpret_cast<cuda_scalar_t*>(sin.data_ptr<scalar_t>()),
-          reinterpret_cast<cuda_scalar_t*>(query.data_ptr<scalar_t>()),
-          key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<scalar_t>()) : nullptr,
-          rot_dim_from_cache,
-          query_token_stride,
-          key_token_stride,
-          head_stride_query,
-          head_stride_key,
-          num_heads,
-          num_kv_heads,
-          (int)head_size);
-    } else {
-      rotary_embedding_kernel<cuda_scalar_t, false><<<grid, block, 0, stream>>>(
-          reinterpret_cast<cuda_scalar_t*>(cos.data_ptr<scalar_t>()),
-          reinterpret_cast<cuda_scalar_t*>(sin.data_ptr<scalar_t>()),
-          reinterpret_cast<cuda_scalar_t*>(query.data_ptr<scalar_t>()),
-          key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<scalar_t>()) : nullptr,
-          rot_dim_from_cache,
-          query_token_stride,
-          key_token_stride,
-          head_stride_query,
-          head_stride_key,
-          num_heads,
-          num_kv_heads,
-          (int)head_size);
-    }
-  });
+      at::ScalarType::Half, at::ScalarType::BFloat16, query.scalar_type(), "rotary_embedding", [&] {
+        using cuda_scalar_t = typename std::conditional<
+            std::is_same<scalar_t, at::Half>::value,
+            nv_half,
+            typename std::conditional<std::is_same<scalar_t, at::BFloat16>::value, nv_bfloat16, scalar_t>::type>::type;
+
+        if (is_neox) {
+          rotary_embedding_kernel<cuda_scalar_t, true><<<grid, block, 0, stream>>>(
+              reinterpret_cast<cuda_scalar_t*>(cos.data_ptr<scalar_t>()),
+              reinterpret_cast<cuda_scalar_t*>(sin.data_ptr<scalar_t>()),
+              reinterpret_cast<cuda_scalar_t*>(query.data_ptr<scalar_t>()),
+              key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<scalar_t>()) : nullptr,
+              rot_dim_from_cache,
+              query_token_stride,
+              key_token_stride,
+              head_stride_query,
+              head_stride_key,
+              num_heads,
+              num_kv_heads,
+              (int)head_size);
+        } else {
+          rotary_embedding_kernel<cuda_scalar_t, false><<<grid, block, 0, stream>>>(
+              reinterpret_cast<cuda_scalar_t*>(cos.data_ptr<scalar_t>()),
+              reinterpret_cast<cuda_scalar_t*>(sin.data_ptr<scalar_t>()),
+              reinterpret_cast<cuda_scalar_t*>(query.data_ptr<scalar_t>()),
+              key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<scalar_t>()) : nullptr,
+              rot_dim_from_cache,
+              query_token_stride,
+              key_token_stride,
+              head_stride_query,
+              head_stride_key,
+              num_heads,
+              num_kv_heads,
+              (int)head_size);
+        }
+      });
   C10_CUDA_KERNEL_LAUNCH_CHECK();
 }
\ No newline at end of file
diff --git a/sgl-kernel/include/sgl_kernel_ops.h b/sgl-kernel/include/sgl_kernel_ops.h
index 0c8438135d4e..38989d2d0ecb 100644
--- a/sgl-kernel/include/sgl_kernel_ops.h
+++ b/sgl-kernel/include/sgl_kernel_ops.h
@@ -396,7 +396,7 @@ void fused_qk_norm_rope(
     double attention_factor);
 
 /*
- * From csrc/multimodal/rotary_embedding.cu
+ * From csrc/multimodal/rotary_embedding
  */
 void rotary_embedding(
     at::Tensor& cos,
diff --git a/sgl-kernel/python/sgl_kernel/rotary_embedding.py b/sgl-kernel/python/sgl_kernel/rotary_embedding.py
index 6d5072e68461..70d34cb402b2 100644
--- a/sgl-kernel/python/sgl_kernel/rotary_embedding.py
+++ b/sgl-kernel/python/sgl_kernel/rotary_embedding.py
@@ -1,9 +1,6 @@
 from typing import Optional
-
 import torch
 
-
-# Adapted from https://github.com/vllm-project/vllm/blob/9214e60631a79506e7669650de87806a123e0b0b/vllm/_custom_ops.py#L249
 # pos encoding ops
 def rotary_embedding(
     cos: torch.Tensor,
diff --git a/sgl-kernel/tests/test_mm_rotary_embedding.py b/sgl-kernel/tests/test_mm_rotary_embedding.py
index 3fb0a5bb2297..dd24cf63ee49 100644
--- a/sgl-kernel/tests/test_mm_rotary_embedding.py
+++ b/sgl-kernel/tests/test_mm_rotary_embedding.py
@@ -40,13 +40,8 @@ def compute_cos_sin_cache(
 ) -> Tuple[torch.Tensor, torch.Tensor]:
     """
     Compute separate cos and sin caches.
-
-    Returns:
-        cos, sin: shape (max_seq_len, rotary_dim / 2)
     """
-    inv_freq = 1.0 / (
-        base ** (torch.arange(0, rotary_dim, 2, dtype=torch.float32) / rotary_dim)
-    )
+    inv_freq = 1.0 / (base ** (torch.arange(0, rotary_dim, 2, dtype=torch.float32) / rotary_dim))
     t = torch.arange(max_seq_len, dtype=torch.float32)
     freqs = torch.einsum("i,j->ij", t, inv_freq)
     cos = freqs.cos().to(dtype)
@@ -62,8 +57,8 @@ def reference_rotary_neox(
     head_size: int,
 ) -> Tuple[torch.Tensor, torch.Tensor]:
     """
-    Reference implementation of NeoX-style rotary embedding, using explicit
-    cos/sin caches and Q/K in flattened (num_tokens, heads * dim) layout.
+    Reference implementation of NeoX-style rotary embedding, 
+    using explicit cos/sin caches and Q/K in flattened (num_tokens, heads * dim) layout.
     """
     num_tokens = query.size(0)
     num_heads = query.size(1) // head_size
@@ -127,7 +122,6 @@ def _check_rotary_correctness(
     tol: float = 1e-2,
     device: str = "cuda",
 ) -> RotaryTestResult:
-    """Helper function for rotary correctness checking."""
     if not HAS_ROTARY_EMBEDDING:
         return RotaryTestResult("basic (skipped: rotary_embedding not built)", True, details="rotary_embedding extension not found")
 

From 1fc12596eae6d38ca0b056cfe9f4203333caf166 Mon Sep 17 00:00:00 2001
From: RubiaCx <1084281732@qq.com>
Date: Thu, 11 Dec 2025 08:06:45 -0800
Subject: [PATCH 03/37] WIP

---
 .gitignore                                    |   1 +
 .../runtime/models/dits/qwen_image.py         |  45 +++++--
 sgl-kernel/CMakeLists.txt                     |   2 +-
 sgl-kernel/csrc/common_extension.cc           |  32 ++++-
 .../rope}/rotary_embedding.cu                 |  35 +++--
 sgl-kernel/include/sgl_kernel_ops.h           |  10 ++
 sgl-kernel/python/sgl_kernel/__init__.py      |   1 +
 .../python/sgl_kernel/rotary_embedding.py     | 123 ++++++++++++++++--
 sgl-kernel/tests/test_mm_rotary_embedding.py  |   5 +-
 9 files changed, 215 insertions(+), 39 deletions(-)
 rename sgl-kernel/csrc/{multimodal => sgl_diffusion/rope}/rotary_embedding.cu (91%)

diff --git a/.gitignore b/.gitignore
index 3326fca2cec6..3039bf31d9bb 100644
--- a/.gitignore
+++ b/.gitignore
@@ -252,3 +252,4 @@ outputs/
 
 # Eval Cache
 .longbench_cache/
+logs/
\ No newline at end of file
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 1dd0781a8464..9491f8d67378 100644
--- a/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py
+++ b/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py
@@ -21,6 +21,11 @@
     apply_rotary_embedding,
     fuse_scale_shift_kernel,
 )
+
+try:
+    from sgl_kernel.rotary_embedding import rotary_embedding_cos_sin as sgl_rotary_embedding  # type: ignore
+except Exception:
+    sgl_rotary_embedding = None
 from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT
 from sglang.multimodal_gen.runtime.models.dits.utils import (
     delete_projection_layers,
@@ -397,18 +402,34 @@ def forward(
         # Apply RoPE
         if image_rotary_emb is not None:
             (img_cos, img_sin), (txt_cos, txt_sin) = image_rotary_emb
-            img_query = apply_rotary_embedding(
-                img_query, img_cos, img_sin, interleaved=True
-            )
-            img_key = apply_rotary_embedding(
-                img_key, img_cos, img_sin, interleaved=True
-            )
-            txt_query = apply_rotary_embedding(
-                txt_query, txt_cos, txt_sin, interleaved=True
-            )
-            txt_key = apply_rotary_embedding(
-                txt_key, txt_cos, txt_sin, interleaved=True
-            )
+
+            if sgl_rotary_embedding is not None:
+                def _apply_sgl_rope(q: torch.Tensor, k: Optional[torch.Tensor], cos, sin):
+                    # sgl_kernel expects contiguous [num_tokens, num_heads, head_size]
+                    q_shape = q.shape
+                    q_view = q.contiguous().view(-1, q_shape[-2], q_shape[-1])
+                    k_view = k.contiguous().view(-1, k.shape[-2], k.shape[-1]) if k is not None else None
+                    # interleaved=True; kernel内部会根据 cos/sin 形状处理 Neox / GPT-J 布局
+                    sgl_rotary_embedding(cos, sin, q_view, k_view, q_shape[-1], True)
+                    q_out = q_view.view(q_shape)
+                    k_out = k_view.view(k.shape) if k_view is not None else None
+                    return q_out, k_out
+
+                img_query, img_key = _apply_sgl_rope(img_query, img_key, img_cos, img_sin)
+                txt_query, txt_key = _apply_sgl_rope(txt_query, txt_key, txt_cos, txt_sin)
+            else:
+                img_query = apply_rotary_embedding(
+                    img_query, img_cos, img_sin, interleaved=True
+                )
+                img_key = apply_rotary_embedding(
+                    img_key, img_cos, img_sin, interleaved=True
+                )
+                txt_query = apply_rotary_embedding(
+                    txt_query, txt_cos, txt_sin, interleaved=True
+                )
+                txt_key = apply_rotary_embedding(
+                    txt_key, txt_cos, txt_sin, interleaved=True
+                )
 
         # Concatenate for joint attention
         # Order: [text, image]
diff --git a/sgl-kernel/CMakeLists.txt b/sgl-kernel/CMakeLists.txt
index 684c9f150361..0aaaea8f26eb 100644
--- a/sgl-kernel/CMakeLists.txt
+++ b/sgl-kernel/CMakeLists.txt
@@ -313,7 +313,7 @@ set(SOURCES
     "csrc/kvcacheio/transfer.cu"
     "csrc/mamba/causal_conv1d.cu"
     "csrc/memory/store.cu"
-    "csrc/multimodal/rotary_embedding.cu"
+    "csrc/sgl_diffusion/rope/rotary_embedding.cu"
     "csrc/memory/weak_ref_tensor.cpp"
 
     "csrc/moe/cutlass_moe/w4a8/scaled_mm_entry.cu"
diff --git a/sgl-kernel/csrc/common_extension.cc b/sgl-kernel/csrc/common_extension.cc
index faced414d1de..2c4ffad61255 100644
--- a/sgl-kernel/csrc/common_extension.cc
+++ b/sgl-kernel/csrc/common_extension.cc
@@ -91,8 +91,36 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) {
   m.impl("apply_rope_pos_ids_cos_sin_cache", torch::kCUDA, &apply_rope_pos_ids_cos_sin_cache);
 
   m.def(
-      "rotary_embedding(Tensor cos, Tensor sin, Tensor(a!) query, Tensor? key, int head_size, bool is_neox) -> ()");
-  m.impl("rotary_embedding", torch::kCUDA, &rotary_embedding);
+      "rotary_embedding(Tensor positions, Tensor! query,Tensor!? key, int head_size, Tensor cos_sin_cache, bool is_neox) -> ()");
+  m.impl(
+      "rotary_embedding",
+      torch::kCUDA,
+      static_cast<void (*)(torch::Tensor&, torch::Tensor&, std::optional<torch::Tensor>, int64_t, torch::Tensor&, bool)>(
+          &rotary_embedding));
+
+  m.def(
+      "apply_rotary_embedding_cached(Tensor positions, Tensor! query,Tensor!? key, int head_size, Tensor cos_sin_cache, bool is_neox) -> ()");
+  m.impl(
+      "apply_rotary_embedding_cached",
+      torch::kCUDA,
+      static_cast<void (*)(torch::Tensor&, torch::Tensor&, std::optional<torch::Tensor>, int64_t, torch::Tensor&, bool)>(
+          &rotary_embedding));
+
+  m.def(
+      "rotary_embedding_cos_sin(Tensor cos, Tensor sin, Tensor(a!) query, Tensor? key, int head_size, bool interleaved) -> ()");
+  m.impl(
+      "rotary_embedding_cos_sin",
+      torch::kCUDA,
+      static_cast<void (*)(at::Tensor&, at::Tensor&, at::Tensor&, const std::optional<at::Tensor>&, int64_t, bool)>(
+          &rotary_embedding_cos_sin));
+
+  m.def(
+      "apply_rotary_embedding(Tensor cos, Tensor sin, Tensor(a!) query, Tensor? key, int head_size, bool interleaved) -> ()");
+  m.impl(
+      "apply_rotary_embedding",
+      torch::kCUDA,
+      static_cast<void (*)(at::Tensor&, at::Tensor&, at::Tensor&, const std::optional<at::Tensor>&, int64_t, bool)>(
+          &rotary_embedding_cos_sin));
 
   m.def(
       "downcast_fp8(Tensor k, Tensor v, Tensor k_out, Tensor v_out, Tensor k_scale, Tensor v_scale, Tensor loc, "
diff --git a/sgl-kernel/csrc/multimodal/rotary_embedding.cu b/sgl-kernel/csrc/sgl_diffusion/rope/rotary_embedding.cu
similarity index 91%
rename from sgl-kernel/csrc/multimodal/rotary_embedding.cu
rename to sgl-kernel/csrc/sgl_diffusion/rope/rotary_embedding.cu
index 803995960b47..f7eb3b3847e4 100644
--- a/sgl-kernel/csrc/multimodal/rotary_embedding.cu
+++ b/sgl-kernel/csrc/sgl_diffusion/rope/rotary_embedding.cu
@@ -22,10 +22,11 @@
 #include <torch/all.h>
 
 #include <cmath>
+#include <vector>
 
 #include "utils.h"
 
-template <typename scalar_t, bool IS_NEOX>
+template <typename scalar_t, bool interleaved>
 inline __device__ void apply_token_rotary_embedding(
     scalar_t* __restrict__ arr,
     const scalar_t* __restrict__ cos_ptr,
@@ -34,7 +35,7 @@ inline __device__ void apply_token_rotary_embedding(
     int embed_dim) {
   int x_index, y_index;
 
-  if (IS_NEOX) {
+  if (interleaved) {
     // NeoX-style: interleaved layout [x0, y0, x1, y1, ...].
     // For NeoX, cos/sin have shape [..., rotary_dim/2]; 
     // each index corresponds to one (x,y) pair.
@@ -67,7 +68,7 @@ inline __device__ void apply_token_rotary_embedding(
   }
 }
 
-template <typename scalar_t, bool IS_NEOX>
+template <typename scalar_t, bool interleaved>
 inline __device__ void apply_rotary_embedding(
     scalar_t* __restrict__ query,                        // [num_heads, head_size]
     scalar_t* __restrict__ key,                          // [num_kv_heads, head_size]
@@ -79,7 +80,7 @@ inline __device__ void apply_rotary_embedding(
     const int rot_dim,
     const int64_t head_stride_query,
     const int64_t head_stride_key) {
-  const int embed_dim_for_rotation = IS_NEOX ? rot_dim : (rot_dim / 2);
+  const int embed_dim_for_rotation = interleaved ? rot_dim : (rot_dim / 2);
 
   const int nq_pairs = num_heads * embed_dim_for_rotation;
   for (int i = threadIdx.x; i < nq_pairs; i += blockDim.x) {
@@ -88,7 +89,7 @@ inline __device__ void apply_rotary_embedding(
 
     scalar_t* query_for_token_head = query + head_idx * (int)head_stride_query;
 
-    apply_token_rotary_embedding<scalar_t, IS_NEOX>(
+    apply_token_rotary_embedding<scalar_t, interleaved>(
         query_for_token_head, current_token_cos_ptr, current_token_sin_ptr, rot_offset, embed_dim_for_rotation);
   }
 
@@ -100,13 +101,13 @@ inline __device__ void apply_rotary_embedding(
 
       scalar_t* key_for_token_head = key + head_idx * (int)head_stride_key;
 
-      apply_token_rotary_embedding<scalar_t, IS_NEOX>(
+      apply_token_rotary_embedding<scalar_t, interleaved>(
           key_for_token_head, current_token_cos_ptr, current_token_sin_ptr, rot_offset, embed_dim_for_rotation);
     }
   }
 }
 
-template <typename scalar_t, bool IS_NEOX>
+template <typename scalar_t, bool interleaved>
 __global__ void rotary_embedding_kernel(
     const scalar_t* __restrict__ cos_data,  // [num_tokens, rot_dim_arg]
     const scalar_t* __restrict__ sin_data,  // [num_tokens, rot_dim_arg]
@@ -127,7 +128,7 @@ __global__ void rotary_embedding_kernel(
   scalar_t* query_for_token = query_total + token_idx * (int)query_token_stride;
   scalar_t* key_for_token = (key_total != nullptr) ? (key_total + token_idx * (int)key_token_stride) : nullptr;
 
-  apply_rotary_embedding<scalar_t, IS_NEOX>(
+  apply_rotary_embedding<scalar_t, interleaved>(
       query_for_token,
       key_for_token,
       current_token_cos_ptr,
@@ -140,13 +141,13 @@ __global__ void rotary_embedding_kernel(
       head_stride_key);
 }
 
-void rotary_embedding(
+void rotary_embedding_cos_sin(
     at::Tensor& cos,
     at::Tensor& sin,
     at::Tensor& query,
     const std::optional<at::Tensor>& key,
     int64_t head_size,
-    bool is_neox) {
+    bool interleaved) {
   TORCH_CHECK(
       query.dim() == 2 || query.dim() == 3,
       "query must be in  shape [num_tokens, hidden_size] or [num_tokens, num_heads, head_size]");
@@ -199,6 +200,14 @@ void rotary_embedding(
   }
   TORCH_CHECK(num_heads % num_kv_heads == 0, "num_heads must be divisible by num_kv_heads");
 
+  // If caller passed full-dim cos/sin for interleaved layout, downsample to half-dim once here.
+  if (interleaved && cos.size(1) == head_size) {
+    const int64_t half = head_size / 2;
+    std::vector<int64_t> new_shape = {cos.size(0), half, 2};
+    cos = cos.view(new_shape).select(2, 0).contiguous();  // even positions
+    sin = sin.view(new_shape).select(2, 1).contiguous();
+  }
+
   int rot_dim_from_cache = (int)cos.size(1);
 
   int64_t query_token_stride = query_hidden_size_calculated;
@@ -225,7 +234,7 @@ void rotary_embedding(
   // Number of (x,y) pairs rotated per head:
   //  - NeoX: cos.size(1) = rotary_dim/2  => pairs = rot_dim_from_cache
   //  - GPT-J: cos.size(1) = rotary_dim   => pairs = rot_dim_from_cache / 2
-  int embed_dim_for_block_calc = is_neox ? rot_dim_from_cache : (rot_dim_from_cache / 2);
+  int embed_dim_for_block_calc = interleaved ? rot_dim_from_cache : (rot_dim_from_cache / 2);
   int max_pairs_to_rotate_per_token =
       std::max(num_heads * embed_dim_for_block_calc, num_kv_heads * embed_dim_for_block_calc);
   dim3 block(std::min<int64_t>(max_pairs_to_rotate_per_token, 512L));
@@ -236,13 +245,13 @@ void rotary_embedding(
   const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
 
   AT_DISPATCH_FLOATING_TYPES_AND2(
-      at::ScalarType::Half, at::ScalarType::BFloat16, query.scalar_type(), "rotary_embedding", [&] {
+      at::ScalarType::Half, at::ScalarType::BFloat16, query.scalar_type(), "rotary_embedding_cos_sin", [&] {
         using cuda_scalar_t = typename std::conditional<
             std::is_same<scalar_t, at::Half>::value,
             nv_half,
             typename std::conditional<std::is_same<scalar_t, at::BFloat16>::value, nv_bfloat16, scalar_t>::type>::type;
 
-        if (is_neox) {
+        if (interleaved) {
           rotary_embedding_kernel<cuda_scalar_t, true><<<grid, block, 0, stream>>>(
               reinterpret_cast<cuda_scalar_t*>(cos.data_ptr<scalar_t>()),
               reinterpret_cast<cuda_scalar_t*>(sin.data_ptr<scalar_t>()),
diff --git a/sgl-kernel/include/sgl_kernel_ops.h b/sgl-kernel/include/sgl_kernel_ops.h
index 4f5aa38dc54d..075a9505797f 100644
--- a/sgl-kernel/include/sgl_kernel_ops.h
+++ b/sgl-kernel/include/sgl_kernel_ops.h
@@ -149,6 +149,7 @@ void apply_rope_pos_ids_cos_sin_cache(
     const std::optional<at::Tensor>& v_buffer,
     const std::optional<at::Tensor>& kv_cache_loc);
 
+// RoPE with position ids + cos/sin cache (used by SRT)
 void rotary_embedding(
     torch::Tensor& positions,
     torch::Tensor& query,
@@ -157,6 +158,15 @@ void rotary_embedding(
     torch::Tensor& cos_sin_cache,
     bool is_neox);
 
+// RoPE with explicit cos/sin tensors (used by diffusion / Qwen image)
+void rotary_embedding_cos_sin(
+    at::Tensor& cos,
+    at::Tensor& sin,
+    at::Tensor& query,
+    const std::optional<at::Tensor>& key,
+    int64_t head_size,
+    bool interleaved);
+
 void downcast_fp8(
     at::Tensor& k,
     at::Tensor& v,
diff --git a/sgl-kernel/python/sgl_kernel/__init__.py b/sgl-kernel/python/sgl_kernel/__init__.py
index 8e8994e04c95..dfe637cc46be 100644
--- a/sgl-kernel/python/sgl_kernel/__init__.py
+++ b/sgl-kernel/python/sgl_kernel/__init__.py
@@ -33,6 +33,7 @@
     rotary_embedding,
     silu_and_mul,
 )
+from sgl_kernel.rotary_embedding import apply_rotary_embedding, rotary_embedding_cos_sin
 from sgl_kernel.expert_specialization import (
     es_fp8_blockwise_scaled_grouped_mm,
     es_sm100_mxfp8_blockscaled_grouped_mm,
diff --git a/sgl-kernel/python/sgl_kernel/rotary_embedding.py b/sgl-kernel/python/sgl_kernel/rotary_embedding.py
index 70d34cb402b2..896aae937306 100644
--- a/sgl-kernel/python/sgl_kernel/rotary_embedding.py
+++ b/sgl-kernel/python/sgl_kernel/rotary_embedding.py
@@ -1,20 +1,125 @@
 from typing import Optional
 import torch
 
-# pos encoding ops
-def rotary_embedding(
+# Detect available kernels (old and new names).
+_HAS_COS_SIN = hasattr(torch.ops.sgl_kernel, "rotary_embedding_cos_sin")
+_HAS_GENERIC = hasattr(torch.ops.sgl_kernel, "rotary_embedding")
+
+print(f"HAS_COS_SIN: {_HAS_COS_SIN}")
+print(f"HAS_GENERIC: {_HAS_GENERIC}")
+
+def apply_rotary_embedding(
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    query: torch.Tensor,
+    key: Optional[torch.Tensor] = None,
+    head_size: int = 0,
+    interleaved: Optional[bool] = None,
+    *,
+    is_neox: Optional[bool] = None,
+) -> None:
+    """Apply rotary embedding with precomputed cos/sin.
+
+    - `interleaved=True` (NeoX) layout: `[x0, y0, x1, y1, ...]`
+    - `interleaved=False` (GPT-J/LLaMA) layout: `[x0, x1, ..., y0, y1, ...]`
+    - `is_neox` is kept as alias to stay consistent with other call-sites.
+    """
+    if interleaved is not None and is_neox is not None and interleaved != is_neox:
+        raise ValueError(
+            f"is_neox({is_neox}) and interleaved({interleaved}) mismatch; keep only one or make them equal."
+        )
+
+    if is_neox is not None:
+        effective_interleaved = is_neox
+    elif interleaved is not None:
+        effective_interleaved = interleaved
+    else:
+        effective_interleaved = True  # default NeoX for backward compatibility
+
+
+    # Legacy explicit name.
+    if _HAS_COS_SIN:
+        torch.ops.sgl_kernel.rotary_embedding_cos_sin(
+            cos,
+            sin,
+            query,
+            key if key is not None else None,
+            head_size,
+            effective_interleaved,
+        )
+        return
+
+    # Some older builds overload `rotary_embedding` directly with cos/sin signature.
+    if _HAS_GENERIC:
+        try:
+            torch.ops.sgl_kernel.rotary_embedding(
+                cos,
+                sin,
+                query,
+                key if key is not None else None,
+                head_size,
+                effective_interleaved,
+            )
+            return
+        except Exception:
+            pass
+
+    # Fallback: pack cos/sin into cache and call positions+cache variant.
+    positions = torch.arange(cos.size(0), device=cos.device, dtype=torch.int64)
+    cos_sin_cache = torch.cat([cos, sin], dim=1)
+    # Final fallback to legacy positions kernel name.
+    if _HAS_GENERIC:
+        torch.ops.sgl_kernel.rotary_embedding(
+            positions=positions,
+            query=query,
+            key=key if key is not None else None,
+            head_size=head_size,
+            cos_sin_cache=cos_sin_cache,
+            is_neox=effective_interleaved,
+        )
+        return
+
+    raise RuntimeError("No rotary embedding kernel is available in torch.ops.sgl_kernel")
+
+
+# Backward-compatible aliases
+def rotary_embedding_cos_sin(
     cos: torch.Tensor,
     sin: torch.Tensor,
     query: torch.Tensor,
-    key: Optional[torch.Tensor],
-    head_size: int,
-    is_neox: bool,
+    key: Optional[torch.Tensor] = None,
+    head_size: int = 0,
+    interleaved: Optional[bool] = None,
+    *,
+    is_neox: Optional[bool] = None,
 ) -> None:
-    torch.ops.sgl_kernel.rotary_embedding.default(
+    apply_rotary_embedding(
         cos,
         sin,
         query,
-        key,
-        head_size,
-        is_neox,
+        key=key,
+        head_size=head_size,
+        interleaved=interleaved,
+        is_neox=is_neox,
     )
+
+
+def rotary_embedding(
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    query: torch.Tensor,
+    key: Optional[torch.Tensor] = None,
+    head_size: int = 0,
+    interleaved: Optional[bool] = None,
+    *,
+    is_neox: Optional[bool] = None,
+) -> None:
+    apply_rotary_embedding(
+        cos,
+        sin,
+        query,
+        key=key,
+        head_size=head_size,
+        interleaved=interleaved,
+        is_neox=is_neox,
+    )
\ No newline at end of file
diff --git a/sgl-kernel/tests/test_mm_rotary_embedding.py b/sgl-kernel/tests/test_mm_rotary_embedding.py
index dd24cf63ee49..2a7ee8155361 100644
--- a/sgl-kernel/tests/test_mm_rotary_embedding.py
+++ b/sgl-kernel/tests/test_mm_rotary_embedding.py
@@ -15,12 +15,13 @@
     HAS_PYTEST = False
 
 try:
-    from sgl_kernel.rotary_embedding import rotary_embedding as rotary_emb_module
+    # Cos/sin-based rotary kernel from sgl_diffusion
+    from sgl_kernel.rotary_embedding import rotary_embedding_cos_sin as rotary_emb_module
     HAS_ROTARY_EMBEDDING = True
 except ImportError:
     rotary_emb_module = None
     HAS_ROTARY_EMBEDDING = False
-    print("sgl_kernel.rotary_embedding not available")
+    print("sgl_kernel.rotary_embedding_cos_sin not available")
 
 
 @dataclass

From ce00ef61751274807d4c462a68d242bb7bf56f16 Mon Sep 17 00:00:00 2001
From: RubiaCx <1084281732@qq.com>
Date: Fri, 12 Dec 2025 06:24:15 -0800
Subject: [PATCH 04/37] Add diffusion rotary kernel path and unify op naming

---
 .../benchmark/bench_mm_rotary_embedding.py    | 141 ++++++---
 sgl-kernel/csrc/common_extension.cc           |  16 -
 .../sgl_diffusion/rope/rotary_embedding.cu    | 292 ++++++++++++------
 sgl-kernel/python/sgl_kernel/__init__.py      |   2 +-
 .../python/sgl_kernel/rotary_embedding.py     | 148 +++++----
 5 files changed, 380 insertions(+), 219 deletions(-)

diff --git a/sgl-kernel/benchmark/bench_mm_rotary_embedding.py b/sgl-kernel/benchmark/bench_mm_rotary_embedding.py
index ce36543e22a9..8c4a2ec30a68 100644
--- a/sgl-kernel/benchmark/bench_mm_rotary_embedding.py
+++ b/sgl-kernel/benchmark/bench_mm_rotary_embedding.py
@@ -7,7 +7,8 @@
 import torch
 import triton
 
-from sgl_kernel.rotary_embedding import rotary_embedding as sgl_rotary_embedding
+from sgl_kernel.rotary_embedding import rotary_embedding_cos_sin as sgl_rotary_cos_sin
+from sgl_kernel.testing.rotary_embedding import RotaryEmbedding as NativeRotaryEmbedding
 
 def compute_cos_sin_cache(
     max_seq_len: int,
@@ -15,12 +16,6 @@ def compute_cos_sin_cache(
     base: float = 10000.0,
     dtype: torch.dtype = torch.float32,
 ) -> Tuple[torch.Tensor, torch.Tensor]:
-    """
-    Compute separate cos and sin caches.
-
-    Returns:
-        cos, sin: shape (max_seq_len, rotary_dim / 2)
-    """
     inv_freq = 1.0 / (
         base ** (torch.arange(0, rotary_dim, 2, dtype=torch.float32) / rotary_dim)
     )
@@ -32,10 +27,8 @@ def compute_cos_sin_cache(
 
 
 def benchmark_mm_rotary_embedding() -> None:
-    """Benchmark sgl_kernel.rotary_embedding vs vLLM & flash_attn (if available)."""
     try:
         from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding as vLLMRotaryEmbedding
-
         HAS_VLLM = True
     except ImportError:
         vLLMRotaryEmbedding = None
@@ -59,16 +52,14 @@ def benchmark_mm_rotary_embedding() -> None:
 
     # Test configurations: (batch_size, num_heads, num_kv_heads, head_size)
     configs = [
-        # Decoding scenarios (batch_size=1)
         (1, 32, 8, 128),
         (1, 64, 8, 128),
         (1, 32, 8, 64),
         (1, 32, 8, 256),
         (32, 32, 8, 64),
-        # # Edge cases
-        (1, 32, 1, 128),   # MQA
-        (32, 8, 8, 128),   # MHA
-        (1, 32, 8, 80),    # Non-standard head_size
+        (1, 32, 1, 128), 
+        (32, 8, 8, 128), 
+        (1, 32, 8, 80),  
     ]
 
     for batch_size, num_heads, num_kv_heads, head_size in configs:
@@ -77,13 +68,28 @@ def benchmark_mm_rotary_embedding() -> None:
         except torch.AcceleratorError:
             torch.cuda.empty_cache()
             pass
-        
+
         try:
             rotary_dim = head_size
-            cos_cache, sin_cache = compute_cos_sin_cache(max_seq_len, rotary_dim, dtype=dtype)
+            cos_cache, sin_cache = compute_cos_sin_cache(
+                max_seq_len, rotary_dim, dtype=dtype
+            )
             cos_cache = cos_cache.to(device)
             sin_cache = sin_cache.to(device)
 
+            cos_sin_cache = torch.cat([cos_cache, sin_cache], dim=-1).to(
+                device=device, dtype=dtype
+            )
+
+            native_rope = NativeRotaryEmbedding(
+                head_size=head_size,
+                rotary_dim=rotary_dim,
+                max_position_embeddings=max_seq_len,
+                base=10000,
+                is_neox_style=True,
+                dtype=dtype,
+            ).to(device)
+
             print(f"\nConfig: batch_size={batch_size}, heads={num_heads}/{num_kv_heads}, head_size={head_size}, dtype={dtype}")
             print("-" * 100)
         except Exception as e:
@@ -91,7 +97,9 @@ def benchmark_mm_rotary_embedding() -> None:
             continue
 
         header = f"{'seq_len':>8}"
-        header += f" | {'ours (ms)':>10}"
+        header += f" | {'native (ms)':>12}"
+        header += f" | {'sgl_cos_sin (ms)':>16}"
+        header += f" | {'sgl_pos (ms)':>12}"
         if HAS_VLLM:
             header += f" | {'vLLM (ms)':>10}"
         if HAS_FLASH_ATTN:
@@ -112,7 +120,9 @@ def benchmark_mm_rotary_embedding() -> None:
                 sin = sin_cache[positions]
 
                 row_str = f"{seq_len:8d}"
-                ours_time = None
+                native_time = None
+                sgl_cos_sin_time = None
+                sgl_pos_time = None
                 vllm_time = None
                 fa_time = None
             except (RuntimeError, torch.AcceleratorError) as e:
@@ -120,17 +130,40 @@ def benchmark_mm_rotary_embedding() -> None:
                 torch.cuda.synchronize()
                 continue
 
-            # Ours: sgl_kernel.rotary_embedding (triton.testing.do_bench style)
-            def fn_ours() -> None:
+            def fn_native() -> None:
+                q = query.clone()
+                k = key.clone()
+                native_rope.forward_native(positions, q, k)
+
+            ms, _, _ = triton.testing.do_bench(fn_native, quantiles=[0.5, 0.2, 0.8])
+            native_time = 1000 * ms
+            row_str += f" | {native_time:12.4f}"
+
+            def fn_sgl_cos_sin() -> None:
+                q = query.clone()
+                k = key.clone()
+                sgl_rotary_cos_sin(cos, sin, q, k, head_size, True)
+
+            ms, _, _ = triton.testing.do_bench(fn_sgl_cos_sin, quantiles=[0.5, 0.2, 0.8])
+            sgl_cos_sin_time = 1000 * ms
+            row_str += f" | {sgl_cos_sin_time:16.4f}"
+
+            def fn_sgl_pos() -> None:
                 q = query.clone()
                 k = key.clone()
-                sgl_rotary_embedding(cos, sin, q, k, head_size, True)
+                torch.ops.sgl_kernel.rotary_embedding(
+                    positions,
+                    q,
+                    k,
+                    head_size,
+                    cos_sin_cache,
+                    True,
+                )
 
-            ms, _, _ = triton.testing.do_bench(fn_ours, quantiles=[0.5, 0.2, 0.8])
-            ours_time = 1000 * ms
-            row_str += f" | {ours_time:10.4f}"
+            ms, _, _ = triton.testing.do_bench(fn_sgl_pos, quantiles=[0.5, 0.2, 0.8])
+            sgl_pos_time = 1000 * ms
+            row_str += f" | {sgl_pos_time:12.4f}"
 
-            # vLLM
             if HAS_VLLM:
                 vllm_rope = vLLMRotaryEmbedding(
                     head_size=head_size,
@@ -140,7 +173,7 @@ def fn_ours() -> None:
                     is_neox_style=True,
                     dtype=dtype,
                 ).cuda()
-                
+
                 def fn_vllm() -> None:
                     q = query.clone()
                     k = key.clone()
@@ -149,40 +182,68 @@ def fn_vllm() -> None:
                 ms, _, _ = triton.testing.do_bench(fn_vllm, quantiles=[0.5, 0.2, 0.8])
                 vllm_time = 1000 * ms
                 row_str += f" | {vllm_time:10.4f}"
-            else:
-                row_str += f" | {'N/A':>10}"
 
-            # FlashAttention RotaryEmbedding (NeoX-style, QKV layout)
             if HAS_FLASH_ATTN:
                 try:
                     flash_rotary = FlashRotaryEmbedding(rotary_dim, device=device)
-                    qkv = torch.randn(batch_size, seq_len, 3, num_heads, head_size, dtype=dtype, device=device)
+                    qkv = torch.randn(
+                        batch_size,
+                        seq_len,
+                        3,
+                        num_heads,
+                        head_size,
+                        dtype=dtype,
+                        device=device,
+                    )
 
                     def fn_flash_attn() -> None:
                         qkv_fa = qkv.clone()
                         flash_rotary(qkv_fa, seqlen_offset=0)
 
-                    ms, _, _ = triton.testing.do_bench(fn_flash_attn, quantiles=[0.5, 0.2, 0.8])
+                    ms, _, _ = triton.testing.do_bench(
+                        fn_flash_attn, quantiles=[0.5, 0.2, 0.8]
+                    )
                     fa_time = 1000 * ms
                     row_str += f" | {fa_time:14.4f}"
                 except Exception:
                     row_str += f" | {'ERROR':>14}"
                     fa_time = None
-            else:
-                row_str += f" | {'N/A':>14}"
 
-            if HAS_VLLM and ours_time is not None and vllm_time is not None:
-                speedup = vllm_time / ours_time
-                row_str += f" | {speedup:9.2f}x"
+            if sgl_cos_sin_time is not None:
+                if HAS_VLLM and vllm_time is not None:
+                    speedup = vllm_time / sgl_cos_sin_time
+                    row_str += f" | {speedup:9.2f}x"
+                elif (not HAS_VLLM) and native_time is not None:
+                    speedup = native_time / sgl_cos_sin_time
+                    row_str += f" | {speedup:9.2f}x"
+                else:
+                    row_str += f" | {'N/A':>9}"
 
             print(row_str)
-            results.append({"seq_len": seq_len, "ours": ours_time, "vllm": vllm_time, "flash_attn": fa_time})
+            results.append(
+                {
+                    "seq_len": seq_len,
+                    "native": native_time,
+                    "sgl_cos_sin": sgl_cos_sin_time,
+                    "sgl_pos": sgl_pos_time,
+                    "vllm": vllm_time,
+                    "flash_attn": fa_time,
+                }
+            )
 
         print("-" * 100)
         print("\nAverage time:")
-        ours_vals = [r["ours"] for r in results if r["ours"] is not None]
-        if ours_vals:
-            print(f"  SGLang: {np.mean(ours_vals):.4f} ms")
+        native_vals = [r["native"] for r in results if r["native"] is not None]
+        if native_vals:
+            print(f"  Native: {np.mean(native_vals):.4f} ms")
+
+        sgl_cos_sin_vals = [r["sgl_cos_sin"] for r in results if r["sgl_cos_sin"] is not None]
+        if sgl_cos_sin_vals:
+            print(f"  SGLang (cos/sin API): {np.mean(sgl_cos_sin_vals):.4f} ms")
+
+        sgl_pos_vals = [r["sgl_pos"] for r in results if r["sgl_pos"] is not None]
+        if sgl_pos_vals:
+            print(f"  SGLang (positions API): {np.mean(sgl_pos_vals):.4f} ms")
         if HAS_VLLM:
             vllm_vals = [r["vllm"] for r in results if r["vllm"] is not None]
             if vllm_vals:
diff --git a/sgl-kernel/csrc/common_extension.cc b/sgl-kernel/csrc/common_extension.cc
index 2c4ffad61255..f82d15073cd9 100644
--- a/sgl-kernel/csrc/common_extension.cc
+++ b/sgl-kernel/csrc/common_extension.cc
@@ -98,14 +98,6 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) {
       static_cast<void (*)(torch::Tensor&, torch::Tensor&, std::optional<torch::Tensor>, int64_t, torch::Tensor&, bool)>(
           &rotary_embedding));
 
-  m.def(
-      "apply_rotary_embedding_cached(Tensor positions, Tensor! query,Tensor!? key, int head_size, Tensor cos_sin_cache, bool is_neox) -> ()");
-  m.impl(
-      "apply_rotary_embedding_cached",
-      torch::kCUDA,
-      static_cast<void (*)(torch::Tensor&, torch::Tensor&, std::optional<torch::Tensor>, int64_t, torch::Tensor&, bool)>(
-          &rotary_embedding));
-
   m.def(
       "rotary_embedding_cos_sin(Tensor cos, Tensor sin, Tensor(a!) query, Tensor? key, int head_size, bool interleaved) -> ()");
   m.impl(
@@ -114,14 +106,6 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) {
       static_cast<void (*)(at::Tensor&, at::Tensor&, at::Tensor&, const std::optional<at::Tensor>&, int64_t, bool)>(
           &rotary_embedding_cos_sin));
 
-  m.def(
-      "apply_rotary_embedding(Tensor cos, Tensor sin, Tensor(a!) query, Tensor? key, int head_size, bool interleaved) -> ()");
-  m.impl(
-      "apply_rotary_embedding",
-      torch::kCUDA,
-      static_cast<void (*)(at::Tensor&, at::Tensor&, at::Tensor&, const std::optional<at::Tensor>&, int64_t, bool)>(
-          &rotary_embedding_cos_sin));
-
   m.def(
       "downcast_fp8(Tensor k, Tensor v, Tensor k_out, Tensor v_out, Tensor k_scale, Tensor v_scale, Tensor loc, "
       "int mult, int offset) -> ()");
diff --git a/sgl-kernel/csrc/sgl_diffusion/rope/rotary_embedding.cu b/sgl-kernel/csrc/sgl_diffusion/rope/rotary_embedding.cu
index f7eb3b3847e4..ac89d8b20333 100644
--- a/sgl-kernel/csrc/sgl_diffusion/rope/rotary_embedding.cu
+++ b/sgl-kernel/csrc/sgl_diffusion/rope/rotary_embedding.cu
@@ -28,19 +28,16 @@
 
 template <typename scalar_t, bool interleaved>
 inline __device__ void apply_token_rotary_embedding(
-    scalar_t* __restrict__ arr,
-    const scalar_t* __restrict__ cos_ptr,
-    const scalar_t* __restrict__ sin_ptr,
+    scalar_t* __restrict__ arr,            // [head_size]
+    const scalar_t* __restrict__ cos_ptr,  // [rot_dim]
+    const scalar_t* __restrict__ sin_ptr,  // [rot_dim]
     int rot_offset,
-    int embed_dim) {
-  int x_index, y_index;
-
-  if (interleaved) {
+    int embed_dim) {  // for non-interleaved: half dim
+  if constexpr (interleaved) {
     // NeoX-style: interleaved layout [x0, y0, x1, y1, ...].
-    // For NeoX, cos/sin have shape [..., rotary_dim/2]; 
-    // each index corresponds to one (x,y) pair.
-    x_index = 2 * rot_offset;
-    y_index = x_index + 1;
+    // cos/sin: [..., rotary_dim/2], one entry per (x, y) pair.
+    const int x_index = 2 * rot_offset;
+    const int y_index = x_index + 1;
 
     const float cos_val = static_cast<float>(SGLANG_LDG(cos_ptr + rot_offset));
     const float sin_val = static_cast<float>(SGLANG_LDG(sin_ptr + rot_offset));
@@ -49,12 +46,11 @@ inline __device__ void apply_token_rotary_embedding(
     const float y = static_cast<float>(arr[y_index]);
     arr[x_index] = static_cast<scalar_t>(x * cos_val - y * sin_val);
     arr[y_index] = static_cast<scalar_t>(y * cos_val + x * sin_val);
-
   } else {
-    // GPT-J / LLaMA style when cos/sin are [..., rotary_dim], i.e. non-interleaved
-    // [x0, x1, ..., y0, y1, ...]. Here embed_dim is the "half" dimension and cos/sin have length 2 * embed_dim.
-    x_index = rot_offset;
-    y_index = rot_offset + embed_dim;
+    // GPT-J / LLaMA style: layout [x0, x1, ..., y0, y1, ...]
+    // cos/sin: [..., rotary_dim], one entry per (x, y) pair.
+    const int x_index = rot_offset;
+    const int y_index = rot_offset + embed_dim;
 
     const float cos_val_x = static_cast<float>(SGLANG_LDG(cos_ptr + rot_offset));
     const float sin_val_x = static_cast<float>(SGLANG_LDG(sin_ptr + rot_offset));
@@ -68,52 +64,88 @@ inline __device__ void apply_token_rotary_embedding(
   }
 }
 
+template <>
+inline __device__ void apply_token_rotary_embedding<float, true>(
+    float* __restrict__ arr,
+    const float* __restrict__ cos_ptr,
+    const float* __restrict__ sin_ptr,
+    int rot_offset,
+    int /*embed_dim*/) {
+  float2 xy = *reinterpret_cast<const float2*>(arr + rot_offset * 2);
+  const float cos_val = static_cast<float>(SGLANG_LDG(cos_ptr + rot_offset));
+  const float sin_val = static_cast<float>(SGLANG_LDG(sin_ptr + rot_offset));
+
+  float2 out;
+  out.x = xy.x * cos_val - xy.y * sin_val;
+  out.y = xy.y * cos_val + xy.x * sin_val;
+
+  *reinterpret_cast<float2*>(arr + rot_offset * 2) = out;
+}
+
+// 2D grid kernel: parallel over tokens (grid.x) and pair-tiles (grid.y)
 template <typename scalar_t, bool interleaved>
-inline __device__ void apply_rotary_embedding(
-    scalar_t* __restrict__ query,                        // [num_heads, head_size]
-    scalar_t* __restrict__ key,                          // [num_kv_heads, head_size]
-    const scalar_t* __restrict__ current_token_cos_ptr,  // [rot_dim]
-    const scalar_t* __restrict__ current_token_sin_ptr,  // [rot_dim]
-    const int head_size,
+__global__ void rotary_embedding_kernel_2d(
+    const scalar_t* __restrict__ cos_data,  // [num_tokens, rot_dim_arg]
+    const scalar_t* __restrict__ sin_data,  // [num_tokens, rot_dim_arg]
+    scalar_t* __restrict__ query_total,
+    scalar_t* __restrict__ key_total,
+    const int rot_dim_arg,
+    const int embed_dim_for_rotation,
+    const int64_t query_token_stride,
+    const int64_t key_token_stride,
+    const int64_t head_stride_query,
+    const int64_t head_stride_key,
     const int num_heads,
     const int num_kv_heads,
-    const int rot_dim,
-    const int64_t head_stride_query,
-    const int64_t head_stride_key) {
-  const int embed_dim_for_rotation = interleaved ? rot_dim : (rot_dim / 2);
+    const int head_size,
+    const int blocks_per_token) {
+  const int token_idx = blockIdx.x;
+  if (token_idx >= gridDim.x) {
+    return;
+  }
+
+  const scalar_t* current_token_cos_ptr = cos_data + token_idx * rot_dim_arg;
+  const scalar_t* current_token_sin_ptr = sin_data + token_idx * rot_dim_arg;
+
+  scalar_t* query_for_token = query_total + token_idx * (int)query_token_stride;
+  scalar_t* key_for_token = (key_total != nullptr) ? (key_total + token_idx * (int)key_token_stride) : nullptr;
+
+  const int local_block_idx = blockIdx.y; 
+  const int pair_stride = blockDim.x * blocks_per_token;
+  const int thread_pair_offset = local_block_idx * blockDim.x + threadIdx.x;
 
   const int nq_pairs = num_heads * embed_dim_for_rotation;
-  for (int i = threadIdx.x; i < nq_pairs; i += blockDim.x) {
+  for (int i = thread_pair_offset; i < nq_pairs; i += pair_stride) {
     const int head_idx = i / embed_dim_for_rotation;
     const int rot_offset = i % embed_dim_for_rotation;
 
-    scalar_t* query_for_token_head = query + head_idx * (int)head_stride_query;
-
+    scalar_t* query_for_token_head = query_for_token + head_idx * (int)head_stride_query;
     apply_token_rotary_embedding<scalar_t, interleaved>(
         query_for_token_head, current_token_cos_ptr, current_token_sin_ptr, rot_offset, embed_dim_for_rotation);
   }
 
-  if (key != nullptr) {
+  if (key_for_token != nullptr) {
     const int nk_pairs = num_kv_heads * embed_dim_for_rotation;
-    for (int i = threadIdx.x; i < nk_pairs; i += blockDim.x) {
+    for (int i = thread_pair_offset; i < nk_pairs; i += pair_stride) {
       const int head_idx = i / embed_dim_for_rotation;
       const int rot_offset = i % embed_dim_for_rotation;
 
-      scalar_t* key_for_token_head = key + head_idx * (int)head_stride_key;
-
+      scalar_t* key_for_token_head = key_for_token + head_idx * (int)head_stride_key;
       apply_token_rotary_embedding<scalar_t, interleaved>(
           key_for_token_head, current_token_cos_ptr, current_token_sin_ptr, rot_offset, embed_dim_for_rotation);
     }
   }
 }
 
+// 1D grid kernel: each block handles one token
 template <typename scalar_t, bool interleaved>
-__global__ void rotary_embedding_kernel(
+__global__ void rotary_embedding_kernel_1d(
     const scalar_t* __restrict__ cos_data,  // [num_tokens, rot_dim_arg]
     const scalar_t* __restrict__ sin_data,  // [num_tokens, rot_dim_arg]
     scalar_t* __restrict__ query_total,
     scalar_t* __restrict__ key_total,
     const int rot_dim_arg,
+    const int embed_dim_for_rotation,
     const int64_t query_token_stride,
     const int64_t key_token_stride,
     const int64_t head_stride_query,
@@ -122,23 +154,33 @@ __global__ void rotary_embedding_kernel(
     const int num_kv_heads,
     const int head_size) {
   const int token_idx = blockIdx.x;
+  if (token_idx >= gridDim.x) return;
+
   const scalar_t* current_token_cos_ptr = cos_data + token_idx * rot_dim_arg;
   const scalar_t* current_token_sin_ptr = sin_data + token_idx * rot_dim_arg;
 
   scalar_t* query_for_token = query_total + token_idx * (int)query_token_stride;
   scalar_t* key_for_token = (key_total != nullptr) ? (key_total + token_idx * (int)key_token_stride) : nullptr;
 
-  apply_rotary_embedding<scalar_t, interleaved>(
-      query_for_token,
-      key_for_token,
-      current_token_cos_ptr,
-      current_token_sin_ptr,
-      head_size,
-      num_heads,
-      num_kv_heads,
-      rot_dim_arg,
-      head_stride_query,
-      head_stride_key);
+  const int nq_pairs = num_heads * embed_dim_for_rotation;
+  for (int i = threadIdx.x; i < nq_pairs; i += blockDim.x) {
+    const int head_idx = i / embed_dim_for_rotation;
+    const int rot_offset = i % embed_dim_for_rotation;
+    scalar_t* query_for_token_head = query_for_token + head_idx * (int)head_stride_query;
+    apply_token_rotary_embedding<scalar_t, interleaved>(
+        query_for_token_head, current_token_cos_ptr, current_token_sin_ptr, rot_offset, embed_dim_for_rotation);
+  }
+
+  if (key_for_token != nullptr) {
+    const int nk_pairs = num_kv_heads * embed_dim_for_rotation;
+    for (int i = threadIdx.x; i < nk_pairs; i += blockDim.x) {
+      const int head_idx = i / embed_dim_for_rotation;
+      const int rot_offset = i % embed_dim_for_rotation;
+      scalar_t* key_for_token_head = key_for_token + head_idx * (int)head_stride_key;
+      apply_token_rotary_embedding<scalar_t, interleaved>(
+          key_for_token_head, current_token_cos_ptr, current_token_sin_ptr, rot_offset, embed_dim_for_rotation);
+    }
+  }
 }
 
 void rotary_embedding_cos_sin(
@@ -150,29 +192,29 @@ void rotary_embedding_cos_sin(
     bool interleaved) {
   TORCH_CHECK(
       query.dim() == 2 || query.dim() == 3,
-      "query must be in  shape [num_tokens, hidden_size] or [num_tokens, num_heads, head_size]");
+      "query must be in shape [num_tokens, hidden_size] or [num_tokens, num_heads, head_size]");
   if (key.has_value()) {
     TORCH_CHECK(
         key->dim() == 2 || key->dim() == 3,
-        "key must be in  shape [num_tokens, hidden_size] or [num_tokens, num_kv_heads, head_size]");
+        "key must be in shape [num_tokens, hidden_size] or [num_tokens, num_kv_heads, head_size]");
   }
 
-  int64_t num_tokens = query.size(0);
+  const int64_t num_tokens = query.size(0);
 
   TORCH_CHECK(cos.dim() == 2, "cos must be in shape [num_tokens, D_cos]");
-  TORCH_CHECK(sin.dim() == 2, "sin must be in  shape [num_tokens, D_sin]");
+  TORCH_CHECK(sin.dim() == 2, "sin must be in shape [num_tokens, D_sin]");
   TORCH_CHECK(cos.size(0) == num_tokens, "cos num_tokens mismatch with query");
   TORCH_CHECK(sin.size(0) == num_tokens, "sin num_tokens mismatch with query");
   TORCH_CHECK(cos.size(1) == sin.size(1), "cos and sin D_cos/D_sin mismatch");
 
-  TORCH_CHECK(cos.scalar_type() == query.scalar_type(), "cos dtype mismatch");
-  TORCH_CHECK(sin.scalar_type() == query.scalar_type(), "sin dtype mismatch");
-  TORCH_CHECK(cos.is_cuda() && sin.is_cuda() && query.is_cuda(), "All tensors must be on CUDA");
-  TORCH_CHECK(query.is_contiguous(), "query must be contiguous; got non-contiguous tensor");
+  TORCH_CHECK(cos.scalar_type() == query.scalar_type(), "cos dtype mismatch with query");
+  TORCH_CHECK(sin.scalar_type() == query.scalar_type(), "sin dtype mismatch with query");
+  TORCH_CHECK(cos.is_cuda() && sin.is_cuda() && query.is_cuda(), "cos/sin/query must be CUDA tensors");
+  TORCH_CHECK(query.is_contiguous(), "query must be contiguous");
   if (key.has_value()) {
-    TORCH_CHECK(key->is_cuda(), "Key tensor must be on CUDA if provided");
-    TORCH_CHECK(key->scalar_type() == query.scalar_type(), "Key dtype mismatch");
-    TORCH_CHECK(key->is_contiguous(), "key must be contiguous when provided; got non-contiguous tensor");
+    TORCH_CHECK(key->is_cuda(), "key must be CUDA tensor if provided");
+    TORCH_CHECK(key->scalar_type() == query.scalar_type(), "key dtype mismatch with query");
+    TORCH_CHECK(key->is_contiguous(), "key must be contiguous");
   }
 
   int query_hidden_size_calculated;
@@ -180,38 +222,39 @@ void rotary_embedding_cos_sin(
     query_hidden_size_calculated = (int)query.size(1);
   } else {
     query_hidden_size_calculated = (int)query.size(1) * (int)query.size(2);
-    TORCH_CHECK(query.size(2) == head_size, "Query head_size mismatch in 3D tensor");
+    TORCH_CHECK(query.size(2) == head_size, "query head_size mismatch in 3D tensor");
   }
   TORCH_CHECK(query_hidden_size_calculated % head_size == 0, "query_hidden_size not divisible by head_size");
-  int num_heads = (int)query_hidden_size_calculated / (int)head_size;
+  int num_heads = query_hidden_size_calculated / (int)head_size;
 
   int key_hidden_size_calculated = 0;
   int num_kv_heads = num_heads;
   if (key.has_value()) {
-    TORCH_CHECK((int)key->size(0) == num_tokens, "Key num_tokens mismatch");
+    TORCH_CHECK((int)key->size(0) == num_tokens, "key num_tokens mismatch");
     if (key->dim() == 2) {
       key_hidden_size_calculated = (int)key->size(1);
     } else {
       key_hidden_size_calculated = (int)key->size(1) * (int)key->size(2);
-      TORCH_CHECK((int)key->size(2) == head_size, "Key head_size mismatch in 3D tensor");
+      TORCH_CHECK((int)key->size(2) == head_size, "key head_size mismatch in 3D tensor");
     }
     TORCH_CHECK(key_hidden_size_calculated % head_size == 0, "key_hidden_size not divisible by head_size");
     num_kv_heads = key_hidden_size_calculated / (int)head_size;
   }
   TORCH_CHECK(num_heads % num_kv_heads == 0, "num_heads must be divisible by num_kv_heads");
 
-  // If caller passed full-dim cos/sin for interleaved layout, downsample to half-dim once here.
+  // NeoX interleaved: if cos/sin are full-dim (head_size), downsample once.
   if (interleaved && cos.size(1) == head_size) {
+    TORCH_CHECK(head_size % 2 == 0, "interleaved layout requires even head_size");
     const int64_t half = head_size / 2;
     std::vector<int64_t> new_shape = {cos.size(0), half, 2};
-    cos = cos.view(new_shape).select(2, 0).contiguous();  // even positions
+    cos = cos.view(new_shape).select(2, 0).contiguous();
     sin = sin.view(new_shape).select(2, 1).contiguous();
   }
 
-  int rot_dim_from_cache = (int)cos.size(1);
+  const int rot_dim_from_cache = (int)cos.size(1);
 
-  int64_t query_token_stride = query_hidden_size_calculated;
-  int64_t key_token_stride = key.has_value() ? key_hidden_size_calculated : 0;
+  const int64_t query_token_stride = query_hidden_size_calculated;
+  const int64_t key_token_stride = key.has_value() ? key_hidden_size_calculated : 0;
 
   int64_t head_stride_query;
   if (query.dim() == 3 && query.size(1) == num_heads && query.size(2) == head_size) {
@@ -229,57 +272,100 @@ void rotary_embedding_cos_sin(
     }
   }
 
-  dim3 grid((int)num_tokens);
+  const int embed_dim_for_rotation = interleaved ? rot_dim_from_cache : (rot_dim_from_cache / 2);
+  TORCH_CHECK(embed_dim_for_rotation > 0, "embed_dim_for_rotation must be > 0");
 
-  // Number of (x,y) pairs rotated per head:
-  //  - NeoX: cos.size(1) = rotary_dim/2  => pairs = rot_dim_from_cache
-  //  - GPT-J: cos.size(1) = rotary_dim   => pairs = rot_dim_from_cache / 2
-  int embed_dim_for_block_calc = interleaved ? rot_dim_from_cache : (rot_dim_from_cache / 2);
-  int max_pairs_to_rotate_per_token =
-      std::max(num_heads * embed_dim_for_block_calc, num_kv_heads * embed_dim_for_block_calc);
-  dim3 block(std::min<int64_t>(max_pairs_to_rotate_per_token, 512L));
+  const int max_pairs_to_rotate_per_token =
+      std::max(num_heads * embed_dim_for_rotation, num_kv_heads * embed_dim_for_rotation);
 
-  if (block.x == 0 && num_tokens > 0) block.x = 1;
+  const int threads_per_block = std::min<int>(256, std::max(128, embed_dim_for_rotation));
+  const int blocks_per_token = (max_pairs_to_rotate_per_token + threads_per_block - 1) / threads_per_block;
+
+  const bool use_grid_2d = (num_tokens <= 4) && (blocks_per_token > 1);
+  dim3 block(threads_per_block);
+  dim3 grid_2d((int)num_tokens, std::max(1, blocks_per_token));
+  dim3 grid_1d((int)num_tokens);
 
   const at::cuda::OptionalCUDAGuard device_guard(device_of(query));
   const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
 
   AT_DISPATCH_FLOATING_TYPES_AND2(
       at::ScalarType::Half, at::ScalarType::BFloat16, query.scalar_type(), "rotary_embedding_cos_sin", [&] {
+        using torch_scalar_t = scalar_t;
         using cuda_scalar_t = typename std::conditional<
-            std::is_same<scalar_t, at::Half>::value,
+            std::is_same<torch_scalar_t, at::Half>::value,
             nv_half,
-            typename std::conditional<std::is_same<scalar_t, at::BFloat16>::value, nv_bfloat16, scalar_t>::type>::type;
+            typename std::conditional<std::is_same<torch_scalar_t, at::BFloat16>::value, nv_bfloat16, torch_scalar_t>::
+                type>::type;
 
         if (interleaved) {
-          rotary_embedding_kernel<cuda_scalar_t, true><<<grid, block, 0, stream>>>(
-              reinterpret_cast<cuda_scalar_t*>(cos.data_ptr<scalar_t>()),
-              reinterpret_cast<cuda_scalar_t*>(sin.data_ptr<scalar_t>()),
-              reinterpret_cast<cuda_scalar_t*>(query.data_ptr<scalar_t>()),
-              key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<scalar_t>()) : nullptr,
-              rot_dim_from_cache,
-              query_token_stride,
-              key_token_stride,
-              head_stride_query,
-              head_stride_key,
-              num_heads,
-              num_kv_heads,
-              (int)head_size);
+          if (use_grid_2d) {
+            rotary_embedding_kernel_2d<cuda_scalar_t, true><<<grid_2d, block, 0, stream>>>(
+                reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
+                reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
+                reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
+                key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>()) : nullptr,
+                rot_dim_from_cache,
+                embed_dim_for_rotation,
+                query_token_stride,
+                key_token_stride,
+                head_stride_query,
+                head_stride_key,
+                num_heads,
+                num_kv_heads,
+                (int)head_size,
+                blocks_per_token);
+          } else {
+            rotary_embedding_kernel_1d<cuda_scalar_t, true><<<grid_1d, block, 0, stream>>>(
+                reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
+                reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
+                reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
+                key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>()) : nullptr,
+                rot_dim_from_cache,
+                embed_dim_for_rotation,
+                query_token_stride,
+                key_token_stride,
+                head_stride_query,
+                head_stride_key,
+                num_heads,
+                num_kv_heads,
+                (int)head_size);
+          }
         } else {
-          rotary_embedding_kernel<cuda_scalar_t, false><<<grid, block, 0, stream>>>(
-              reinterpret_cast<cuda_scalar_t*>(cos.data_ptr<scalar_t>()),
-              reinterpret_cast<cuda_scalar_t*>(sin.data_ptr<scalar_t>()),
-              reinterpret_cast<cuda_scalar_t*>(query.data_ptr<scalar_t>()),
-              key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<scalar_t>()) : nullptr,
-              rot_dim_from_cache,
-              query_token_stride,
-              key_token_stride,
-              head_stride_query,
-              head_stride_key,
-              num_heads,
-              num_kv_heads,
-              (int)head_size);
+          if (use_grid_2d) {
+            rotary_embedding_kernel_2d<cuda_scalar_t, false><<<grid_2d, block, 0, stream>>>(
+                reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
+                reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
+                reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
+                key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>()) : nullptr,
+                rot_dim_from_cache,
+                embed_dim_for_rotation,
+                query_token_stride,
+                key_token_stride,
+                head_stride_query,
+                head_stride_key,
+                num_heads,
+                num_kv_heads,
+                (int)head_size,
+                blocks_per_token);
+          } else {
+            rotary_embedding_kernel_1d<cuda_scalar_t, false><<<grid_1d, block, 0, stream>>>(
+                reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
+                reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
+                reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
+                key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>()) : nullptr,
+                rot_dim_from_cache,
+                embed_dim_for_rotation,
+                query_token_stride,
+                key_token_stride,
+                head_stride_query,
+                head_stride_key,
+                num_heads,
+                num_kv_heads,
+                (int)head_size);
+          }
         }
       });
+
   C10_CUDA_KERNEL_LAUNCH_CHECK();
-}
\ No newline at end of file
+}
diff --git a/sgl-kernel/python/sgl_kernel/__init__.py b/sgl-kernel/python/sgl_kernel/__init__.py
index dfe637cc46be..191360421816 100644
--- a/sgl-kernel/python/sgl_kernel/__init__.py
+++ b/sgl-kernel/python/sgl_kernel/__init__.py
@@ -33,7 +33,7 @@
     rotary_embedding,
     silu_and_mul,
 )
-from sgl_kernel.rotary_embedding import apply_rotary_embedding, rotary_embedding_cos_sin
+from sgl_kernel.rotary_embedding import rotary_embedding_cos_sin
 from sgl_kernel.expert_specialization import (
     es_fp8_blockwise_scaled_grouped_mm,
     es_sm100_mxfp8_blockscaled_grouped_mm,
diff --git a/sgl-kernel/python/sgl_kernel/rotary_embedding.py b/sgl-kernel/python/sgl_kernel/rotary_embedding.py
index 896aae937306..0c187e2ae4c1 100644
--- a/sgl-kernel/python/sgl_kernel/rotary_embedding.py
+++ b/sgl-kernel/python/sgl_kernel/rotary_embedding.py
@@ -1,57 +1,69 @@
-from typing import Optional
+from __future__ import annotations
+
+from typing import Literal, Optional
 import torch
 
-# Detect available kernels (old and new names).
 _HAS_COS_SIN = hasattr(torch.ops.sgl_kernel, "rotary_embedding_cos_sin")
 _HAS_GENERIC = hasattr(torch.ops.sgl_kernel, "rotary_embedding")
 
-print(f"HAS_COS_SIN: {_HAS_COS_SIN}")
-print(f"HAS_GENERIC: {_HAS_GENERIC}")
+
+def _resolve_interleaved(interleaved: Optional[bool], is_neox: Optional[bool]) -> bool:
+    if interleaved is not None and is_neox is not None and interleaved != is_neox:
+        raise ValueError(
+            f"is_neox({is_neox}) and interleaved({interleaved}) mismatch; keep only one or make them equal."
+        )
+    if is_neox is not None:
+        return is_neox
+    if interleaved is not None:
+        return interleaved
+    # Historical default in this wrapper: assume NeoX-style if not specified.
+    return True
+
 
 def apply_rotary_embedding(
-    cos: torch.Tensor,
-    sin: torch.Tensor,
+    *,
+    mode: Literal["cos_sin", "positions"],
     query: torch.Tensor,
     key: Optional[torch.Tensor] = None,
     head_size: int = 0,
     interleaved: Optional[bool] = None,
-    *,
     is_neox: Optional[bool] = None,
+    # cos/sin mode
+    cos: Optional[torch.Tensor] = None,
+    sin: Optional[torch.Tensor] = None,
+    # positions mode
+    positions: Optional[torch.Tensor] = None,
+    cos_sin_cache: Optional[torch.Tensor] = None,
 ) -> None:
-    """Apply rotary embedding with precomputed cos/sin.
+    """内部统一的 rotary dispatch：
+
+    - mode="cos_sin": 使用 (cos, sin) 直接旋转
+    - mode="positions": 使用 (positions, cos_sin_cache) 旋转
 
-    - `interleaved=True` (NeoX) layout: `[x0, y0, x1, y1, ...]`
-    - `interleaved=False` (GPT-J/LLaMA) layout: `[x0, x1, ..., y0, y1, ...]`
-    - `is_neox` is kept as alias to stay consistent with other call-sites.
+    interleaved/is_neox:
+      - True (NeoX): [x0, y0, x1, y1, ...]
+      - False (GPT-J/LLaMA): [x0, x1, ..., y0, y1, ...]
     """
-    if interleaved is not None and is_neox is not None and interleaved != is_neox:
-        raise ValueError(
-            f"is_neox({is_neox}) and interleaved({interleaved}) mismatch; keep only one or make them equal."
-        )
 
-    if is_neox is not None:
-        effective_interleaved = is_neox
-    elif interleaved is not None:
-        effective_interleaved = interleaved
-    else:
-        effective_interleaved = True  # default NeoX for backward compatibility
-
-
-    # Legacy explicit name.
-    if _HAS_COS_SIN:
-        torch.ops.sgl_kernel.rotary_embedding_cos_sin(
-            cos,
-            sin,
-            query,
-            key if key is not None else None,
-            head_size,
-            effective_interleaved,
-        )
-        return
+    effective_interleaved = _resolve_interleaved(interleaved, is_neox)
 
-    # Some older builds overload `rotary_embedding` directly with cos/sin signature.
-    if _HAS_GENERIC:
-        try:
+    if mode == "cos_sin":
+        if cos is None or sin is None:
+            raise ValueError("mode='cos_sin' requires cos and sin")
+
+        if _HAS_COS_SIN:
+            torch.ops.sgl_kernel.rotary_embedding_cos_sin(
+                cos,
+                sin,
+                query,
+                key if key is not None else None,
+                head_size,
+                effective_interleaved,
+            )
+            return
+
+        # Some older builds overload `rotary_embedding` directly with cos/sin signature.
+        if _HAS_GENERIC:
             torch.ops.sgl_kernel.rotary_embedding(
                 cos,
                 sin,
@@ -61,28 +73,45 @@ def apply_rotary_embedding(
                 effective_interleaved,
             )
             return
+
+        raise RuntimeError("No cos/sin rotary embedding kernel is available in torch.ops.sgl_kernel")
+
+    if mode == "positions":
+        if positions is None or cos_sin_cache is None:
+            raise ValueError("mode='positions' requires positions and cos_sin_cache")
+
+        if not _HAS_GENERIC:
+            raise RuntimeError("No positions rotary embedding kernel is available in torch.ops.sgl_kernel")
+
+        # Prefer keyword call if supported.
+        try:
+            torch.ops.sgl_kernel.rotary_embedding(
+                positions=positions,
+                query=query,
+                key=key if key is not None else None,
+                head_size=head_size,
+                cos_sin_cache=cos_sin_cache,
+                is_neox=effective_interleaved,
+            )
+            return
         except Exception:
             pass
 
-    # Fallback: pack cos/sin into cache and call positions+cache variant.
-    positions = torch.arange(cos.size(0), device=cos.device, dtype=torch.int64)
-    cos_sin_cache = torch.cat([cos, sin], dim=1)
-    # Final fallback to legacy positions kernel name.
-    if _HAS_GENERIC:
+        # Fallback to positional signature used by some builds/tests:
+        # rotary_embedding(positions, query, key, head_size, cos_sin_cache, is_neox)
         torch.ops.sgl_kernel.rotary_embedding(
-            positions=positions,
-            query=query,
-            key=key if key is not None else None,
-            head_size=head_size,
-            cos_sin_cache=cos_sin_cache,
-            is_neox=effective_interleaved,
+            positions,
+            query,
+            key if key is not None else None,
+            head_size,
+            cos_sin_cache,
+            effective_interleaved,
         )
         return
 
-    raise RuntimeError("No rotary embedding kernel is available in torch.ops.sgl_kernel")
+    raise ValueError(f"Unknown mode: {mode}")
 
 
-# Backward-compatible aliases
 def rotary_embedding_cos_sin(
     cos: torch.Tensor,
     sin: torch.Tensor,
@@ -94,30 +123,31 @@ def rotary_embedding_cos_sin(
     is_neox: Optional[bool] = None,
 ) -> None:
     apply_rotary_embedding(
-        cos,
-        sin,
-        query,
+        mode="cos_sin",
+        cos=cos,
+        sin=sin,
+        query=query,
         key=key,
         head_size=head_size,
         interleaved=interleaved,
         is_neox=is_neox,
     )
 
-
 def rotary_embedding(
-    cos: torch.Tensor,
-    sin: torch.Tensor,
+    positions: torch.Tensor,
     query: torch.Tensor,
     key: Optional[torch.Tensor] = None,
     head_size: int = 0,
     interleaved: Optional[bool] = None,
     *,
     is_neox: Optional[bool] = None,
+    cos_sin_cache: torch.Tensor,
 ) -> None:
     apply_rotary_embedding(
-        cos,
-        sin,
-        query,
+        mode="positions",
+        positions=positions,
+        cos_sin_cache=cos_sin_cache,
+        query=query,
         key=key,
         head_size=head_size,
         interleaved=interleaved,

From db3cf3329da7f6e11b2d85257e4187523a6f72b4 Mon Sep 17 00:00:00 2001
From: RubiaCx <1084281732@qq.com>
Date: Fri, 12 Dec 2025 07:35:45 -0800
Subject: [PATCH 05/37] typo

---
 .../runtime/models/dits/qwen_image.py          | 18 ++++--------------
 1 file changed, 4 insertions(+), 14 deletions(-)

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 9491f8d67378..94b32b36a5a9 100644
--- a/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py
+++ b/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py
@@ -399,7 +399,6 @@ def forward(
         if self.norm_added_k is not None:
             txt_key = self.norm_added_k(txt_key)
 
-        # Apply RoPE
         if image_rotary_emb is not None:
             (img_cos, img_sin), (txt_cos, txt_sin) = image_rotary_emb
 
@@ -409,7 +408,6 @@ def _apply_sgl_rope(q: torch.Tensor, k: Optional[torch.Tensor], cos, sin):
                     q_shape = q.shape
                     q_view = q.contiguous().view(-1, q_shape[-2], q_shape[-1])
                     k_view = k.contiguous().view(-1, k.shape[-2], k.shape[-1]) if k is not None else None
-                    # interleaved=True; kernel内部会根据 cos/sin 形状处理 Neox / GPT-J 布局
                     sgl_rotary_embedding(cos, sin, q_view, k_view, q_shape[-1], True)
                     q_out = q_view.view(q_shape)
                     k_out = k_view.view(k.shape) if k_view is not None else None
@@ -418,18 +416,10 @@ def _apply_sgl_rope(q: torch.Tensor, k: Optional[torch.Tensor], cos, sin):
                 img_query, img_key = _apply_sgl_rope(img_query, img_key, img_cos, img_sin)
                 txt_query, txt_key = _apply_sgl_rope(txt_query, txt_key, txt_cos, txt_sin)
             else:
-                img_query = apply_rotary_embedding(
-                    img_query, img_cos, img_sin, interleaved=True
-                )
-                img_key = apply_rotary_embedding(
-                    img_key, img_cos, img_sin, interleaved=True
-                )
-                txt_query = apply_rotary_embedding(
-                    txt_query, txt_cos, txt_sin, interleaved=True
-                )
-                txt_key = apply_rotary_embedding(
-                    txt_key, txt_cos, txt_sin, interleaved=True
-                )
+                img_query = apply_rotary_embedding(img_query, img_cos, img_sin, interleaved=True)
+                img_key = apply_rotary_embedding(img_key, img_cos, img_sin, interleaved=True)
+                txt_query = apply_rotary_embedding(txt_query, txt_cos, txt_sin, interleaved=True)
+                txt_key = apply_rotary_embedding(txt_key, txt_cos, txt_sin, interleaved=True)
 
         # Concatenate for joint attention
         # Order: [text, image]

From e4bc7c807aa0d86e817cdad18884d25c63843316 Mon Sep 17 00:00:00 2001
From: Ther-LF <2639852836@qq.com>
Date: Tue, 16 Dec 2025 16:17:07 +0800
Subject: [PATCH 06/37] perf: optimize rotary embedding via smem caching and
 vectorized loads

---
 .../sgl_diffusion/rope/rotary_embedding.cu    | 748 ++++++++++--------
 1 file changed, 412 insertions(+), 336 deletions(-)

diff --git a/sgl-kernel/csrc/sgl_diffusion/rope/rotary_embedding.cu b/sgl-kernel/csrc/sgl_diffusion/rope/rotary_embedding.cu
index ac89d8b20333..e7f69ac76e99 100644
--- a/sgl-kernel/csrc/sgl_diffusion/rope/rotary_embedding.cu
+++ b/sgl-kernel/csrc/sgl_diffusion/rope/rotary_embedding.cu
@@ -15,357 +15,433 @@
  * limitations under the License.
  */
 
-#include <ATen/cuda/CUDAContext.h>
-#include <c10/cuda/CUDAException.h>
-#include <c10/cuda/CUDAGuard.h>
-#include <c10/cuda/CUDAStream.h>
-#include <torch/all.h>
-
-#include <cmath>
-#include <vector>
-
-#include "utils.h"
-
-template <typename scalar_t, bool interleaved>
-inline __device__ void apply_token_rotary_embedding(
-    scalar_t* __restrict__ arr,            // [head_size]
-    const scalar_t* __restrict__ cos_ptr,  // [rot_dim]
-    const scalar_t* __restrict__ sin_ptr,  // [rot_dim]
-    int rot_offset,
-    int embed_dim) {  // for non-interleaved: half dim
-  if constexpr (interleaved) {
-    // NeoX-style: interleaved layout [x0, y0, x1, y1, ...].
-    // cos/sin: [..., rotary_dim/2], one entry per (x, y) pair.
-    const int x_index = 2 * rot_offset;
-    const int y_index = x_index + 1;
-
-    const float cos_val = static_cast<float>(SGLANG_LDG(cos_ptr + rot_offset));
-    const float sin_val = static_cast<float>(SGLANG_LDG(sin_ptr + rot_offset));
-
-    const float x = static_cast<float>(arr[x_index]);
-    const float y = static_cast<float>(arr[y_index]);
-    arr[x_index] = static_cast<scalar_t>(x * cos_val - y * sin_val);
-    arr[y_index] = static_cast<scalar_t>(y * cos_val + x * sin_val);
-  } else {
-    // GPT-J / LLaMA style: layout [x0, x1, ..., y0, y1, ...]
-    // cos/sin: [..., rotary_dim], one entry per (x, y) pair.
-    const int x_index = rot_offset;
-    const int y_index = rot_offset + embed_dim;
-
-    const float cos_val_x = static_cast<float>(SGLANG_LDG(cos_ptr + rot_offset));
-    const float sin_val_x = static_cast<float>(SGLANG_LDG(sin_ptr + rot_offset));
-    const float cos_val_y = static_cast<float>(SGLANG_LDG(cos_ptr + rot_offset + embed_dim));
-    const float sin_val_y = static_cast<float>(SGLANG_LDG(sin_ptr + rot_offset + embed_dim));
-
-    const float x = static_cast<float>(arr[x_index]);
-    const float y = static_cast<float>(arr[y_index]);
-    arr[x_index] = static_cast<scalar_t>(x * cos_val_x - y * sin_val_x);
-    arr[y_index] = static_cast<scalar_t>(y * cos_val_y + x * sin_val_y);
-  }
-}
-
-template <>
-inline __device__ void apply_token_rotary_embedding<float, true>(
-    float* __restrict__ arr,
-    const float* __restrict__ cos_ptr,
-    const float* __restrict__ sin_ptr,
-    int rot_offset,
-    int /*embed_dim*/) {
-  float2 xy = *reinterpret_cast<const float2*>(arr + rot_offset * 2);
-  const float cos_val = static_cast<float>(SGLANG_LDG(cos_ptr + rot_offset));
-  const float sin_val = static_cast<float>(SGLANG_LDG(sin_ptr + rot_offset));
-
-  float2 out;
-  out.x = xy.x * cos_val - xy.y * sin_val;
-  out.y = xy.y * cos_val + xy.x * sin_val;
-
-  *reinterpret_cast<float2*>(arr + rot_offset * 2) = out;
-}
-
-// 2D grid kernel: parallel over tokens (grid.x) and pair-tiles (grid.y)
-template <typename scalar_t, bool interleaved>
-__global__ void rotary_embedding_kernel_2d(
-    const scalar_t* __restrict__ cos_data,  // [num_tokens, rot_dim_arg]
-    const scalar_t* __restrict__ sin_data,  // [num_tokens, rot_dim_arg]
-    scalar_t* __restrict__ query_total,
-    scalar_t* __restrict__ key_total,
-    const int rot_dim_arg,
-    const int embed_dim_for_rotation,
-    const int64_t query_token_stride,
-    const int64_t key_token_stride,
-    const int64_t head_stride_query,
-    const int64_t head_stride_key,
-    const int num_heads,
-    const int num_kv_heads,
-    const int head_size,
-    const int blocks_per_token) {
-  const int token_idx = blockIdx.x;
-  if (token_idx >= gridDim.x) {
-    return;
-  }
-
+ #include <ATen/cuda/CUDAContext.h>
+ #include <c10/cuda/CUDAException.h>
+ #include <c10/cuda/CUDAGuard.h>
+ #include <c10/cuda/CUDAStream.h>
+ #include <torch/all.h>
+ 
+ #include <cmath>
+ #include <vector>
+ 
+ #include "utils.h"
+ 
+ template <typename scalar_t, bool interleaved>
+ inline __device__ void apply_token_rotary_embedding(
+     scalar_t* __restrict__ arr,            // [head_size]
+     const scalar_t* __restrict__ cos_ptr,  // [rot_dim]
+     const scalar_t* __restrict__ sin_ptr,  // [rot_dim]
+     int rot_offset,
+     int embed_dim) {  // for non-interleaved: half dim
+   if constexpr (interleaved) {
+     // NeoX-style: interleaved layout [x0, y0, x1, y1, ...].
+     // cos/sin: [..., rotary_dim/2], one entry per (x, y) pair.
+     const int x_index = 2 * rot_offset;
+     const int y_index = x_index + 1;
+ 
+     // Directly access SMEM
+     const float cos_val = static_cast<float>(cos_ptr[rot_offset]);
+     const float sin_val = static_cast<float>(sin_ptr[rot_offset]);
+ 
+     const float x = static_cast<float>(arr[x_index]);
+     const float y = static_cast<float>(arr[y_index]);
+     arr[x_index] = static_cast<scalar_t>(x * cos_val - y * sin_val);
+     arr[y_index] = static_cast<scalar_t>(y * cos_val + x * sin_val);
+   } else {
+     // GPT-J / LLaMA style: layout [x0, x1, ..., y0, y1, ...]
+     // cos/sin: [..., rotary_dim], one entry per (x, y) pair.
+     const int x_index = rot_offset;
+     const int y_index = rot_offset + embed_dim;
+ 
+     // Directly access SMEM
+     const float cos_val_x = static_cast<float>(cos_ptr[rot_offset]);
+     const float sin_val_x = static_cast<float>(sin_ptr[rot_offset]);
+     const float cos_val_y = static_cast<float>(cos_ptr[rot_offset + embed_dim]);
+     const float sin_val_y = static_cast<float>(sin_ptr[rot_offset + embed_dim]);
+ 
+     const float x = static_cast<float>(arr[x_index]);
+     const float y = static_cast<float>(arr[y_index]);
+     arr[x_index] = static_cast<scalar_t>(x * cos_val_x - y * sin_val_x);
+     arr[y_index] = static_cast<scalar_t>(y * cos_val_y + x * sin_val_y);
+   }
+ }
+ 
+ template <>
+ inline __device__ void apply_token_rotary_embedding<float, true>(
+     float* __restrict__ arr,
+     const float* __restrict__ cos_ptr,
+     const float* __restrict__ sin_ptr,
+     int rot_offset,
+     int /*embed_dim*/) {
+   float2 xy = *reinterpret_cast<const float2*>(arr + rot_offset * 2);
+   
+   // Directly access SMEM
+   const float cos_val = static_cast<float>(cos_ptr[rot_offset]);
+   const float sin_val = static_cast<float>(sin_ptr[rot_offset]);
+ 
+   float2 out;
+   out.x = xy.x * cos_val - xy.y * sin_val;
+   out.y = xy.y * cos_val + xy.x * sin_val;
+ 
+   *reinterpret_cast<float2*>(arr + rot_offset * 2) = out;
+ }
+ 
+ // 2D grid kernel: parallel over tokens (grid.x) and pair-tiles (grid.y)
+ template <typename scalar_t, bool interleaved>
+ __global__ void rotary_embedding_kernel_2d(
+     const scalar_t* __restrict__ cos_data,  // [num_tokens, rot_dim_arg]
+     const scalar_t* __restrict__ sin_data,  // [num_tokens, rot_dim_arg]
+     scalar_t* __restrict__ query_total,
+     scalar_t* __restrict__ key_total,
+     const int rot_dim_arg,
+     const int embed_dim_for_rotation,
+     const int64_t query_token_stride,
+     const int64_t key_token_stride,
+     const int64_t head_stride_query,
+     const int64_t head_stride_key,
+     const int num_heads,
+     const int num_kv_heads,
+     const int head_size,
+     const int blocks_per_token) {
+   
+   extern __shared__ char smem_[];
+   scalar_t* s_cos = reinterpret_cast<scalar_t*>(smem_);
+   scalar_t* s_sin = s_cos + rot_dim_arg;
+ 
+   const int token_idx = blockIdx.x;
+   if (token_idx >= gridDim.x) {
+     return;
+   }
+ 
+   // Pointers to Global Memory for the current token
   const scalar_t* current_token_cos_ptr = cos_data + token_idx * rot_dim_arg;
   const scalar_t* current_token_sin_ptr = sin_data + token_idx * rot_dim_arg;
 
-  scalar_t* query_for_token = query_total + token_idx * (int)query_token_stride;
-  scalar_t* key_for_token = (key_total != nullptr) ? (key_total + token_idx * (int)key_token_stride) : nullptr;
-
-  const int local_block_idx = blockIdx.y; 
-  const int pair_stride = blockDim.x * blocks_per_token;
-  const int thread_pair_offset = local_block_idx * blockDim.x + threadIdx.x;
+  constexpr int kVecBytes = 16;
+  const int scalar_size = sizeof(scalar_t);
+  const int vec_size = kVecBytes / scalar_size;
 
-  const int nq_pairs = num_heads * embed_dim_for_rotation;
-  for (int i = thread_pair_offset; i < nq_pairs; i += pair_stride) {
-    const int head_idx = i / embed_dim_for_rotation;
-    const int rot_offset = i % embed_dim_for_rotation;
+  bool is_aligned = (rot_dim_arg % vec_size == 0) &&
+                    (reinterpret_cast<uintptr_t>(current_token_cos_ptr) % kVecBytes == 0) &&
+                    (reinterpret_cast<uintptr_t>(current_token_sin_ptr) % kVecBytes == 0);
 
-    scalar_t* query_for_token_head = query_for_token + head_idx * (int)head_stride_query;
-    apply_token_rotary_embedding<scalar_t, interleaved>(
-        query_for_token_head, current_token_cos_ptr, current_token_sin_ptr, rot_offset, embed_dim_for_rotation);
-  }
-
-  if (key_for_token != nullptr) {
-    const int nk_pairs = num_kv_heads * embed_dim_for_rotation;
-    for (int i = thread_pair_offset; i < nk_pairs; i += pair_stride) {
-      const int head_idx = i / embed_dim_for_rotation;
-      const int rot_offset = i % embed_dim_for_rotation;
+  if (is_aligned) {
+    using vec_t = float4;
+    const vec_t* cos_vec_ptr = reinterpret_cast<const vec_t*>(current_token_cos_ptr);
+    const vec_t* sin_vec_ptr = reinterpret_cast<const vec_t*>(current_token_sin_ptr);
+    vec_t* s_cos_vec = reinterpret_cast<vec_t*>(s_cos);
+    vec_t* s_sin_vec = reinterpret_cast<vec_t*>(s_sin);
 
-      scalar_t* key_for_token_head = key_for_token + head_idx * (int)head_stride_key;
-      apply_token_rotary_embedding<scalar_t, interleaved>(
-          key_for_token_head, current_token_cos_ptr, current_token_sin_ptr, rot_offset, embed_dim_for_rotation);
+    for (int i = threadIdx.x; i < rot_dim_arg / vec_size; i += blockDim.x) {
+      s_cos_vec[i] = __ldg(cos_vec_ptr + i);
+      s_sin_vec[i] = __ldg(sin_vec_ptr + i);
+    }
+  } else {
+    for (int i = threadIdx.x; i < rot_dim_arg; i += blockDim.x) {
+      s_cos[i] = SGLANG_LDG(current_token_cos_ptr + i);
+      s_sin[i] = SGLANG_LDG(current_token_sin_ptr + i);
     }
   }
-}
-
-// 1D grid kernel: each block handles one token
-template <typename scalar_t, bool interleaved>
-__global__ void rotary_embedding_kernel_1d(
-    const scalar_t* __restrict__ cos_data,  // [num_tokens, rot_dim_arg]
-    const scalar_t* __restrict__ sin_data,  // [num_tokens, rot_dim_arg]
-    scalar_t* __restrict__ query_total,
-    scalar_t* __restrict__ key_total,
-    const int rot_dim_arg,
-    const int embed_dim_for_rotation,
-    const int64_t query_token_stride,
-    const int64_t key_token_stride,
-    const int64_t head_stride_query,
-    const int64_t head_stride_key,
-    const int num_heads,
-    const int num_kv_heads,
-    const int head_size) {
-  const int token_idx = blockIdx.x;
-  if (token_idx >= gridDim.x) return;
-
+   // Essential synchronization
+   __syncthreads();
+ 
+   scalar_t* query_for_token = query_total + token_idx * (int)query_token_stride;
+   scalar_t* key_for_token = (key_total != nullptr) ? (key_total + token_idx * (int)key_token_stride) : nullptr;
+ 
+   const int local_block_idx = blockIdx.y; 
+   const int pair_stride = blockDim.x * blocks_per_token;
+   const int thread_pair_offset = local_block_idx * blockDim.x + threadIdx.x;
+ 
+   const int nq_pairs = num_heads * embed_dim_for_rotation;
+   for (int i = thread_pair_offset; i < nq_pairs; i += pair_stride) {
+     const int head_idx = i / embed_dim_for_rotation;
+     const int rot_offset = i % embed_dim_for_rotation;
+ 
+     scalar_t* query_for_token_head = query_for_token + head_idx * (int)head_stride_query;
+     apply_token_rotary_embedding<scalar_t, interleaved>(
+         query_for_token_head, s_cos, s_sin, rot_offset, embed_dim_for_rotation);
+   }
+ 
+   if (key_for_token != nullptr) {
+     const int nk_pairs = num_kv_heads * embed_dim_for_rotation;
+     for (int i = thread_pair_offset; i < nk_pairs; i += pair_stride) {
+       const int head_idx = i / embed_dim_for_rotation;
+       const int rot_offset = i % embed_dim_for_rotation;
+ 
+       scalar_t* key_for_token_head = key_for_token + head_idx * (int)head_stride_key;
+       // Pass Shared Memory pointers
+       apply_token_rotary_embedding<scalar_t, interleaved>(
+           key_for_token_head, s_cos, s_sin, rot_offset, embed_dim_for_rotation);
+     }
+   }
+ }
+ 
+ // 1D grid kernel: each block handles one token
+ template <typename scalar_t, bool interleaved>
+ __global__ void rotary_embedding_kernel_1d(
+     const scalar_t* __restrict__ cos_data,  // [num_tokens, rot_dim_arg]
+     const scalar_t* __restrict__ sin_data,  // [num_tokens, rot_dim_arg]
+     scalar_t* __restrict__ query_total,
+     scalar_t* __restrict__ key_total,
+     const int rot_dim_arg,
+     const int embed_dim_for_rotation,
+     const int64_t query_token_stride,
+     const int64_t key_token_stride,
+     const int64_t head_stride_query,
+     const int64_t head_stride_key,
+     const int num_heads,
+     const int num_kv_heads,
+     const int head_size) {
+   
+   extern __shared__ char smem_[];
+   scalar_t* s_cos = reinterpret_cast<scalar_t*>(smem_);
+   scalar_t* s_sin = s_cos + rot_dim_arg;
+ 
+   const int token_idx = blockIdx.x;
+   if (token_idx >= gridDim.x) return;
+ 
   const scalar_t* current_token_cos_ptr = cos_data + token_idx * rot_dim_arg;
   const scalar_t* current_token_sin_ptr = sin_data + token_idx * rot_dim_arg;
 
-  scalar_t* query_for_token = query_total + token_idx * (int)query_token_stride;
-  scalar_t* key_for_token = (key_total != nullptr) ? (key_total + token_idx * (int)key_token_stride) : nullptr;
-
-  const int nq_pairs = num_heads * embed_dim_for_rotation;
-  for (int i = threadIdx.x; i < nq_pairs; i += blockDim.x) {
-    const int head_idx = i / embed_dim_for_rotation;
-    const int rot_offset = i % embed_dim_for_rotation;
-    scalar_t* query_for_token_head = query_for_token + head_idx * (int)head_stride_query;
-    apply_token_rotary_embedding<scalar_t, interleaved>(
-        query_for_token_head, current_token_cos_ptr, current_token_sin_ptr, rot_offset, embed_dim_for_rotation);
-  }
-
-  if (key_for_token != nullptr) {
-    const int nk_pairs = num_kv_heads * embed_dim_for_rotation;
-    for (int i = threadIdx.x; i < nk_pairs; i += blockDim.x) {
-      const int head_idx = i / embed_dim_for_rotation;
-      const int rot_offset = i % embed_dim_for_rotation;
-      scalar_t* key_for_token_head = key_for_token + head_idx * (int)head_stride_key;
-      apply_token_rotary_embedding<scalar_t, interleaved>(
-          key_for_token_head, current_token_cos_ptr, current_token_sin_ptr, rot_offset, embed_dim_for_rotation);
-    }
-  }
-}
-
-void rotary_embedding_cos_sin(
-    at::Tensor& cos,
-    at::Tensor& sin,
-    at::Tensor& query,
-    const std::optional<at::Tensor>& key,
-    int64_t head_size,
-    bool interleaved) {
-  TORCH_CHECK(
-      query.dim() == 2 || query.dim() == 3,
-      "query must be in shape [num_tokens, hidden_size] or [num_tokens, num_heads, head_size]");
-  if (key.has_value()) {
-    TORCH_CHECK(
-        key->dim() == 2 || key->dim() == 3,
-        "key must be in shape [num_tokens, hidden_size] or [num_tokens, num_kv_heads, head_size]");
-  }
+  constexpr int kVecBytes = 16;
+  const int scalar_size = sizeof(scalar_t);
+  const int vec_size = kVecBytes / scalar_size;
 
-  const int64_t num_tokens = query.size(0);
+  bool is_aligned = (rot_dim_arg % vec_size == 0) &&
+                    (reinterpret_cast<uintptr_t>(current_token_cos_ptr) % kVecBytes == 0) &&
+                    (reinterpret_cast<uintptr_t>(current_token_sin_ptr) % kVecBytes == 0);
 
-  TORCH_CHECK(cos.dim() == 2, "cos must be in shape [num_tokens, D_cos]");
-  TORCH_CHECK(sin.dim() == 2, "sin must be in shape [num_tokens, D_sin]");
-  TORCH_CHECK(cos.size(0) == num_tokens, "cos num_tokens mismatch with query");
-  TORCH_CHECK(sin.size(0) == num_tokens, "sin num_tokens mismatch with query");
-  TORCH_CHECK(cos.size(1) == sin.size(1), "cos and sin D_cos/D_sin mismatch");
+  if (is_aligned) {
+    using vec_t = float4;
+    const vec_t* cos_vec_ptr = reinterpret_cast<const vec_t*>(current_token_cos_ptr);
+    const vec_t* sin_vec_ptr = reinterpret_cast<const vec_t*>(current_token_sin_ptr);
+    vec_t* s_cos_vec = reinterpret_cast<vec_t*>(s_cos);
+    vec_t* s_sin_vec = reinterpret_cast<vec_t*>(s_sin);
 
-  TORCH_CHECK(cos.scalar_type() == query.scalar_type(), "cos dtype mismatch with query");
-  TORCH_CHECK(sin.scalar_type() == query.scalar_type(), "sin dtype mismatch with query");
-  TORCH_CHECK(cos.is_cuda() && sin.is_cuda() && query.is_cuda(), "cos/sin/query must be CUDA tensors");
-  TORCH_CHECK(query.is_contiguous(), "query must be contiguous");
-  if (key.has_value()) {
-    TORCH_CHECK(key->is_cuda(), "key must be CUDA tensor if provided");
-    TORCH_CHECK(key->scalar_type() == query.scalar_type(), "key dtype mismatch with query");
-    TORCH_CHECK(key->is_contiguous(), "key must be contiguous");
-  }
-
-  int query_hidden_size_calculated;
-  if (query.dim() == 2) {
-    query_hidden_size_calculated = (int)query.size(1);
-  } else {
-    query_hidden_size_calculated = (int)query.size(1) * (int)query.size(2);
-    TORCH_CHECK(query.size(2) == head_size, "query head_size mismatch in 3D tensor");
-  }
-  TORCH_CHECK(query_hidden_size_calculated % head_size == 0, "query_hidden_size not divisible by head_size");
-  int num_heads = query_hidden_size_calculated / (int)head_size;
-
-  int key_hidden_size_calculated = 0;
-  int num_kv_heads = num_heads;
-  if (key.has_value()) {
-    TORCH_CHECK((int)key->size(0) == num_tokens, "key num_tokens mismatch");
-    if (key->dim() == 2) {
-      key_hidden_size_calculated = (int)key->size(1);
-    } else {
-      key_hidden_size_calculated = (int)key->size(1) * (int)key->size(2);
-      TORCH_CHECK((int)key->size(2) == head_size, "key head_size mismatch in 3D tensor");
+    for (int i = threadIdx.x; i < rot_dim_arg / vec_size; i += blockDim.x) {
+      s_cos_vec[i] = __ldg(cos_vec_ptr + i);
+      s_sin_vec[i] = __ldg(sin_vec_ptr + i);
     }
-    TORCH_CHECK(key_hidden_size_calculated % head_size == 0, "key_hidden_size not divisible by head_size");
-    num_kv_heads = key_hidden_size_calculated / (int)head_size;
-  }
-  TORCH_CHECK(num_heads % num_kv_heads == 0, "num_heads must be divisible by num_kv_heads");
-
-  // NeoX interleaved: if cos/sin are full-dim (head_size), downsample once.
-  if (interleaved && cos.size(1) == head_size) {
-    TORCH_CHECK(head_size % 2 == 0, "interleaved layout requires even head_size");
-    const int64_t half = head_size / 2;
-    std::vector<int64_t> new_shape = {cos.size(0), half, 2};
-    cos = cos.view(new_shape).select(2, 0).contiguous();
-    sin = sin.view(new_shape).select(2, 1).contiguous();
-  }
-
-  const int rot_dim_from_cache = (int)cos.size(1);
-
-  const int64_t query_token_stride = query_hidden_size_calculated;
-  const int64_t key_token_stride = key.has_value() ? key_hidden_size_calculated : 0;
-
-  int64_t head_stride_query;
-  if (query.dim() == 3 && query.size(1) == num_heads && query.size(2) == head_size) {
-    head_stride_query = query.stride(1);
   } else {
-    head_stride_query = head_size;
-  }
-
-  int64_t head_stride_key = head_size;
-  if (key.has_value()) {
-    if (key->dim() == 3 && key->size(1) == num_kv_heads && key->size(2) == head_size) {
-      head_stride_key = key->stride(1);
-    } else {
-      head_stride_key = head_size;
+    for (int i = threadIdx.x; i < rot_dim_arg; i += blockDim.x) {
+      s_cos[i] = SGLANG_LDG(current_token_cos_ptr + i);
+      s_sin[i] = SGLANG_LDG(current_token_sin_ptr + i);
     }
   }
-
-  const int embed_dim_for_rotation = interleaved ? rot_dim_from_cache : (rot_dim_from_cache / 2);
-  TORCH_CHECK(embed_dim_for_rotation > 0, "embed_dim_for_rotation must be > 0");
-
-  const int max_pairs_to_rotate_per_token =
-      std::max(num_heads * embed_dim_for_rotation, num_kv_heads * embed_dim_for_rotation);
-
-  const int threads_per_block = std::min<int>(256, std::max(128, embed_dim_for_rotation));
-  const int blocks_per_token = (max_pairs_to_rotate_per_token + threads_per_block - 1) / threads_per_block;
-
-  const bool use_grid_2d = (num_tokens <= 4) && (blocks_per_token > 1);
-  dim3 block(threads_per_block);
-  dim3 grid_2d((int)num_tokens, std::max(1, blocks_per_token));
-  dim3 grid_1d((int)num_tokens);
-
-  const at::cuda::OptionalCUDAGuard device_guard(device_of(query));
-  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
-
-  AT_DISPATCH_FLOATING_TYPES_AND2(
-      at::ScalarType::Half, at::ScalarType::BFloat16, query.scalar_type(), "rotary_embedding_cos_sin", [&] {
-        using torch_scalar_t = scalar_t;
-        using cuda_scalar_t = typename std::conditional<
-            std::is_same<torch_scalar_t, at::Half>::value,
-            nv_half,
-            typename std::conditional<std::is_same<torch_scalar_t, at::BFloat16>::value, nv_bfloat16, torch_scalar_t>::
-                type>::type;
-
-        if (interleaved) {
-          if (use_grid_2d) {
-            rotary_embedding_kernel_2d<cuda_scalar_t, true><<<grid_2d, block, 0, stream>>>(
-                reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
-                reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
-                reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
-                key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>()) : nullptr,
-                rot_dim_from_cache,
-                embed_dim_for_rotation,
-                query_token_stride,
-                key_token_stride,
-                head_stride_query,
-                head_stride_key,
-                num_heads,
-                num_kv_heads,
-                (int)head_size,
-                blocks_per_token);
-          } else {
-            rotary_embedding_kernel_1d<cuda_scalar_t, true><<<grid_1d, block, 0, stream>>>(
-                reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
-                reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
-                reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
-                key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>()) : nullptr,
-                rot_dim_from_cache,
-                embed_dim_for_rotation,
-                query_token_stride,
-                key_token_stride,
-                head_stride_query,
-                head_stride_key,
-                num_heads,
-                num_kv_heads,
-                (int)head_size);
-          }
-        } else {
-          if (use_grid_2d) {
-            rotary_embedding_kernel_2d<cuda_scalar_t, false><<<grid_2d, block, 0, stream>>>(
-                reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
-                reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
-                reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
-                key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>()) : nullptr,
-                rot_dim_from_cache,
-                embed_dim_for_rotation,
-                query_token_stride,
-                key_token_stride,
-                head_stride_query,
-                head_stride_key,
-                num_heads,
-                num_kv_heads,
-                (int)head_size,
-                blocks_per_token);
-          } else {
-            rotary_embedding_kernel_1d<cuda_scalar_t, false><<<grid_1d, block, 0, stream>>>(
-                reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
-                reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
-                reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
-                key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>()) : nullptr,
-                rot_dim_from_cache,
-                embed_dim_for_rotation,
-                query_token_stride,
-                key_token_stride,
-                head_stride_query,
-                head_stride_key,
-                num_heads,
-                num_kv_heads,
-                (int)head_size);
-          }
-        }
-      });
-
-  C10_CUDA_KERNEL_LAUNCH_CHECK();
-}
+   __syncthreads();
+ 
+   scalar_t* query_for_token = query_total + token_idx * (int)query_token_stride;
+   scalar_t* key_for_token = (key_total != nullptr) ? (key_total + token_idx * (int)key_token_stride) : nullptr;
+ 
+   const int nq_pairs = num_heads * embed_dim_for_rotation;
+   for (int i = threadIdx.x; i < nq_pairs; i += blockDim.x) {
+     const int head_idx = i / embed_dim_for_rotation;
+     const int rot_offset = i % embed_dim_for_rotation;
+     scalar_t* query_for_token_head = query_for_token + head_idx * (int)head_stride_query;
+      
+     apply_token_rotary_embedding<scalar_t, interleaved>(
+         query_for_token_head, s_cos, s_sin, rot_offset, embed_dim_for_rotation);
+   }
+ 
+   if (key_for_token != nullptr) {
+     const int nk_pairs = num_kv_heads * embed_dim_for_rotation;
+     for (int i = threadIdx.x; i < nk_pairs; i += blockDim.x) {
+       const int head_idx = i / embed_dim_for_rotation;
+       const int rot_offset = i % embed_dim_for_rotation;
+       
+       scalar_t* key_for_token_head = key_for_token + head_idx * (int)head_stride_key;
+       apply_token_rotary_embedding<scalar_t, interleaved>(
+           key_for_token_head, s_cos, s_sin, rot_offset, embed_dim_for_rotation);
+     }
+   }
+ }
+ 
+ void rotary_embedding_cos_sin(
+     at::Tensor& cos,
+     at::Tensor& sin,
+     at::Tensor& query,
+     const std::optional<at::Tensor>& key,
+     int64_t head_size,
+     bool interleaved) {
+   TORCH_CHECK(
+       query.dim() == 2 || query.dim() == 3,
+       "query must be in shape [num_tokens, hidden_size] or [num_tokens, num_heads, head_size]");
+   if (key.has_value()) {
+     TORCH_CHECK(
+         key->dim() == 2 || key->dim() == 3,
+         "key must be in shape [num_tokens, hidden_size] or [num_tokens, num_kv_heads, head_size]");
+   }
+ 
+   const int64_t num_tokens = query.size(0);
+ 
+   TORCH_CHECK(cos.dim() == 2, "cos must be in shape [num_tokens, D_cos]");
+   TORCH_CHECK(sin.dim() == 2, "sin must be in shape [num_tokens, D_sin]");
+   TORCH_CHECK(cos.size(0) == num_tokens, "cos num_tokens mismatch with query");
+   TORCH_CHECK(sin.size(0) == num_tokens, "sin num_tokens mismatch with query");
+   TORCH_CHECK(cos.size(1) == sin.size(1), "cos and sin D_cos/D_sin mismatch");
+ 
+   TORCH_CHECK(cos.scalar_type() == query.scalar_type(), "cos dtype mismatch with query");
+   TORCH_CHECK(sin.scalar_type() == query.scalar_type(), "sin dtype mismatch with query");
+   TORCH_CHECK(cos.is_cuda() && sin.is_cuda() && query.is_cuda(), "cos/sin/query must be CUDA tensors");
+   TORCH_CHECK(query.is_contiguous(), "query must be contiguous");
+   if (key.has_value()) {
+     TORCH_CHECK(key->is_cuda(), "key must be CUDA tensor if provided");
+     TORCH_CHECK(key->scalar_type() == query.scalar_type(), "key dtype mismatch with query");
+     TORCH_CHECK(key->is_contiguous(), "key must be contiguous");
+   }
+ 
+   int query_hidden_size_calculated;
+   if (query.dim() == 2) {
+     query_hidden_size_calculated = (int)query.size(1);
+   } else {
+     query_hidden_size_calculated = (int)query.size(1) * (int)query.size(2);
+     TORCH_CHECK(query.size(2) == head_size, "query head_size mismatch in 3D tensor");
+   }
+   TORCH_CHECK(query_hidden_size_calculated % head_size == 0, "query_hidden_size not divisible by head_size");
+   int num_heads = query_hidden_size_calculated / (int)head_size;
+ 
+   int key_hidden_size_calculated = 0;
+   int num_kv_heads = num_heads;
+   if (key.has_value()) {
+     TORCH_CHECK((int)key->size(0) == num_tokens, "key num_tokens mismatch");
+     if (key->dim() == 2) {
+       key_hidden_size_calculated = (int)key->size(1);
+     } else {
+       key_hidden_size_calculated = (int)key->size(1) * (int)key->size(2);
+       TORCH_CHECK((int)key->size(2) == head_size, "key head_size mismatch in 3D tensor");
+     }
+     TORCH_CHECK(key_hidden_size_calculated % head_size == 0, "key_hidden_size not divisible by head_size");
+     num_kv_heads = key_hidden_size_calculated / (int)head_size;
+   }
+   TORCH_CHECK(num_heads % num_kv_heads == 0, "num_heads must be divisible by num_kv_heads");
+ 
+   // NeoX interleaved: if cos/sin are full-dim (head_size), downsample once.
+   if (interleaved && cos.size(1) == head_size) {
+     TORCH_CHECK(head_size % 2 == 0, "interleaved layout requires even head_size");
+     const int64_t half = head_size / 2;
+     std::vector<int64_t> new_shape = {cos.size(0), half, 2};
+     cos = cos.view(new_shape).select(2, 0).contiguous();
+     sin = sin.view(new_shape).select(2, 1).contiguous();
+   }
+ 
+   const int rot_dim_from_cache = (int)cos.size(1);
+ 
+   const int64_t query_token_stride = query_hidden_size_calculated;
+   const int64_t key_token_stride = key.has_value() ? key_hidden_size_calculated : 0;
+ 
+   int64_t head_stride_query;
+   if (query.dim() == 3 && query.size(1) == num_heads && query.size(2) == head_size) {
+     head_stride_query = query.stride(1);
+   } else {
+     head_stride_query = head_size;
+   }
+ 
+   int64_t head_stride_key = head_size;
+   if (key.has_value()) {
+     if (key->dim() == 3 && key->size(1) == num_kv_heads && key->size(2) == head_size) {
+       head_stride_key = key->stride(1);
+     } else {
+       head_stride_key = head_size;
+     }
+   }
+ 
+   const int embed_dim_for_rotation = interleaved ? rot_dim_from_cache : (rot_dim_from_cache / 2);
+   TORCH_CHECK(embed_dim_for_rotation > 0, "embed_dim_for_rotation must be > 0");
+ 
+   const int max_pairs_to_rotate_per_token =
+       std::max(num_heads * embed_dim_for_rotation, num_kv_heads * embed_dim_for_rotation);
+ 
+   const int threads_per_block = std::min<int>(256, std::max(128, embed_dim_for_rotation));
+   const int blocks_per_token = (max_pairs_to_rotate_per_token + threads_per_block - 1) / threads_per_block;
+ 
+   const bool use_grid_2d = (num_tokens <= 4) && (blocks_per_token > 1);
+   dim3 block(threads_per_block);
+   dim3 grid_2d((int)num_tokens, std::max(1, blocks_per_token));
+   dim3 grid_1d((int)num_tokens);
+ 
+   const at::cuda::OptionalCUDAGuard device_guard(device_of(query));
+   const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+ 
+   AT_DISPATCH_FLOATING_TYPES_AND2(
+       at::ScalarType::Half, at::ScalarType::BFloat16, query.scalar_type(), "rotary_embedding_cos_sin", [&] {
+         using torch_scalar_t = scalar_t;
+         using cuda_scalar_t = typename std::conditional<
+             std::is_same<torch_scalar_t, at::Half>::value,
+             nv_half,
+             typename std::conditional<std::is_same<torch_scalar_t, at::BFloat16>::value, nv_bfloat16, torch_scalar_t>::
+                 type>::type;
+         
+         // We need 2 arrays (cos, sin) of size 'rot_dim_from_cache', each element is 'sizeof(torch_scalar_t)'
+         size_t smem_size = rot_dim_from_cache * sizeof(torch_scalar_t) * 2;
+ 
+         if (interleaved) {
+           if (use_grid_2d) {
+             rotary_embedding_kernel_2d<cuda_scalar_t, true><<<grid_2d, block, smem_size, stream>>>(
+                 reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
+                 reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
+                 reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
+                 key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>()) : nullptr,
+                 rot_dim_from_cache,
+                 embed_dim_for_rotation,
+                 query_token_stride,
+                 key_token_stride,
+                 head_stride_query,
+                 head_stride_key,
+                 num_heads,
+                 num_kv_heads,
+                 (int)head_size,
+                 blocks_per_token);
+           } else {
+             rotary_embedding_kernel_1d<cuda_scalar_t, true><<<grid_1d, block, smem_size, stream>>>(
+                 reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
+                 reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
+                 reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
+                 key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>()) : nullptr,
+                 rot_dim_from_cache,
+                 embed_dim_for_rotation,
+                 query_token_stride,
+                 key_token_stride,
+                 head_stride_query,
+                 head_stride_key,
+                 num_heads,
+                 num_kv_heads,
+                 (int)head_size);
+           }
+         } else {
+           if (use_grid_2d) {
+             rotary_embedding_kernel_2d<cuda_scalar_t, false><<<grid_2d, block, smem_size, stream>>>(
+                 reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
+                 reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
+                 reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
+                 key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>()) : nullptr,
+                 rot_dim_from_cache,
+                 embed_dim_for_rotation,
+                 query_token_stride,
+                 key_token_stride,
+                 head_stride_query,
+                 head_stride_key,
+                 num_heads,
+                 num_kv_heads,
+                 (int)head_size,
+                 blocks_per_token);
+           } else {
+             rotary_embedding_kernel_1d<cuda_scalar_t, false><<<grid_1d, block, smem_size, stream>>>(
+                 reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
+                 reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
+                 reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
+                 key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>()) : nullptr,
+                 rot_dim_from_cache,
+                 embed_dim_for_rotation,
+                 query_token_stride,
+                 key_token_stride,
+                 head_stride_query,
+                 head_stride_key,
+                 num_heads,
+                 num_kv_heads,
+                 (int)head_size);
+           }
+         }
+       });
+ 
+   C10_CUDA_KERNEL_LAUNCH_CHECK();
+ }
\ No newline at end of file

From 00cf6057da06a180f6c579084815dbbf3079c1ad Mon Sep 17 00:00:00 2001
From: Ther-LF <2639852836@qq.com>
Date: Tue, 16 Dec 2025 16:26:05 +0800
Subject: [PATCH 07/37] replace __ldg with SGLANG_LDG

---
 .gitignore                                    |   2 +-
 .../runtime/models/dits/qwen_image.py         |  39 +-
 .../benchmark/bench_mm_rotary_embedding.py    |  45 +-
 sgl-kernel/csrc/common_extension.cc           |   9 +-
 .../sgl_diffusion/rope/rotary_embedding.cu    | 754 +++++++++---------
 sgl-kernel/python/sgl_kernel/__init__.py      |   2 +-
 .../python/sgl_kernel/rotary_embedding.py     |  12 +-
 sgl-kernel/tests/test_mm_rotary_embedding.py  |  30 +-
 8 files changed, 476 insertions(+), 417 deletions(-)
 mode change 100644 => 100755 sgl-kernel/benchmark/bench_mm_rotary_embedding.py
 mode change 100644 => 100755 sgl-kernel/tests/test_mm_rotary_embedding.py

diff --git a/.gitignore b/.gitignore
index 3039bf31d9bb..e1813ae42614 100644
--- a/.gitignore
+++ b/.gitignore
@@ -252,4 +252,4 @@ outputs/
 
 # Eval Cache
 .longbench_cache/
-logs/
\ No newline at end of file
+logs/
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 94b32b36a5a9..a950b6d94b18 100644
--- a/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py
+++ b/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py
@@ -23,7 +23,9 @@
 )
 
 try:
-    from sgl_kernel.rotary_embedding import rotary_embedding_cos_sin as sgl_rotary_embedding  # type: ignore
+    from sgl_kernel.rotary_embedding import (
+        rotary_embedding_cos_sin as sgl_rotary_embedding,  # type: ignore
+    )
 except Exception:
     sgl_rotary_embedding = None
 from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT
@@ -403,23 +405,42 @@ def forward(
             (img_cos, img_sin), (txt_cos, txt_sin) = image_rotary_emb
 
             if sgl_rotary_embedding is not None:
-                def _apply_sgl_rope(q: torch.Tensor, k: Optional[torch.Tensor], cos, sin):
+
+                def _apply_sgl_rope(
+                    q: torch.Tensor, k: Optional[torch.Tensor], cos, sin
+                ):
                     # sgl_kernel expects contiguous [num_tokens, num_heads, head_size]
                     q_shape = q.shape
                     q_view = q.contiguous().view(-1, q_shape[-2], q_shape[-1])
-                    k_view = k.contiguous().view(-1, k.shape[-2], k.shape[-1]) if k is not None else None
+                    k_view = (
+                        k.contiguous().view(-1, k.shape[-2], k.shape[-1])
+                        if k is not None
+                        else None
+                    )
                     sgl_rotary_embedding(cos, sin, q_view, k_view, q_shape[-1], True)
                     q_out = q_view.view(q_shape)
                     k_out = k_view.view(k.shape) if k_view is not None else None
                     return q_out, k_out
 
-                img_query, img_key = _apply_sgl_rope(img_query, img_key, img_cos, img_sin)
-                txt_query, txt_key = _apply_sgl_rope(txt_query, txt_key, txt_cos, txt_sin)
+                img_query, img_key = _apply_sgl_rope(
+                    img_query, img_key, img_cos, img_sin
+                )
+                txt_query, txt_key = _apply_sgl_rope(
+                    txt_query, txt_key, txt_cos, txt_sin
+                )
             else:
-                img_query = apply_rotary_embedding(img_query, img_cos, img_sin, interleaved=True)
-                img_key = apply_rotary_embedding(img_key, img_cos, img_sin, interleaved=True)
-                txt_query = apply_rotary_embedding(txt_query, txt_cos, txt_sin, interleaved=True)
-                txt_key = apply_rotary_embedding(txt_key, txt_cos, txt_sin, interleaved=True)
+                img_query = apply_rotary_embedding(
+                    img_query, img_cos, img_sin, interleaved=True
+                )
+                img_key = apply_rotary_embedding(
+                    img_key, img_cos, img_sin, interleaved=True
+                )
+                txt_query = apply_rotary_embedding(
+                    txt_query, txt_cos, txt_sin, interleaved=True
+                )
+                txt_key = apply_rotary_embedding(
+                    txt_key, txt_cos, txt_sin, interleaved=True
+                )
 
         # Concatenate for joint attention
         # Order: [text, image]
diff --git a/sgl-kernel/benchmark/bench_mm_rotary_embedding.py b/sgl-kernel/benchmark/bench_mm_rotary_embedding.py
old mode 100644
new mode 100755
index 8c4a2ec30a68..0cc2dc156557
--- a/sgl-kernel/benchmark/bench_mm_rotary_embedding.py
+++ b/sgl-kernel/benchmark/bench_mm_rotary_embedding.py
@@ -6,10 +6,10 @@
 import numpy as np
 import torch
 import triton
-
 from sgl_kernel.rotary_embedding import rotary_embedding_cos_sin as sgl_rotary_cos_sin
 from sgl_kernel.testing.rotary_embedding import RotaryEmbedding as NativeRotaryEmbedding
 
+
 def compute_cos_sin_cache(
     max_seq_len: int,
     rotary_dim: int,
@@ -28,7 +28,10 @@ def compute_cos_sin_cache(
 
 def benchmark_mm_rotary_embedding() -> None:
     try:
-        from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding as vLLMRotaryEmbedding
+        from vllm.model_executor.layers.rotary_embedding import (
+            RotaryEmbedding as vLLMRotaryEmbedding,
+        )
+
         HAS_VLLM = True
     except ImportError:
         vLLMRotaryEmbedding = None
@@ -57,9 +60,9 @@ def benchmark_mm_rotary_embedding() -> None:
         (1, 32, 8, 64),
         (1, 32, 8, 256),
         (32, 32, 8, 64),
-        (1, 32, 1, 128), 
-        (32, 8, 8, 128), 
-        (1, 32, 8, 80),  
+        (1, 32, 1, 128),
+        (32, 8, 8, 128),
+        (1, 32, 8, 80),
     ]
 
     for batch_size, num_heads, num_kv_heads, head_size in configs:
@@ -90,10 +93,14 @@ def benchmark_mm_rotary_embedding() -> None:
                 dtype=dtype,
             ).to(device)
 
-            print(f"\nConfig: batch_size={batch_size}, heads={num_heads}/{num_kv_heads}, head_size={head_size}, dtype={dtype}")
+            print(
+                f"\nConfig: batch_size={batch_size}, heads={num_heads}/{num_kv_heads}, head_size={head_size}, dtype={dtype}"
+            )
             print("-" * 100)
         except Exception as e:
-            print(f"\nSkipping config (batch_size={batch_size}, heads={num_heads}/{num_kv_heads}, head_size={head_size}): {e}")
+            print(
+                f"\nSkipping config (batch_size={batch_size}, heads={num_heads}/{num_kv_heads}, head_size={head_size}): {e}"
+            )
             continue
 
         header = f"{'seq_len':>8}"
@@ -113,9 +120,15 @@ def benchmark_mm_rotary_embedding() -> None:
         for seq_len in seq_lens:
             try:
                 num_tokens = batch_size * seq_len
-                query = torch.randn(num_tokens, num_heads * head_size, dtype=dtype, device=device)  
-                key = torch.randn(num_tokens, num_kv_heads * head_size, dtype=dtype, device=device)
-                positions = torch.arange(seq_len, device=device, dtype=torch.int64).repeat(batch_size)
+                query = torch.randn(
+                    num_tokens, num_heads * head_size, dtype=dtype, device=device
+                )
+                key = torch.randn(
+                    num_tokens, num_kv_heads * head_size, dtype=dtype, device=device
+                )
+                positions = torch.arange(
+                    seq_len, device=device, dtype=torch.int64
+                ).repeat(batch_size)
                 cos = cos_cache[positions]
                 sin = sin_cache[positions]
 
@@ -144,7 +157,9 @@ def fn_sgl_cos_sin() -> None:
                 k = key.clone()
                 sgl_rotary_cos_sin(cos, sin, q, k, head_size, True)
 
-            ms, _, _ = triton.testing.do_bench(fn_sgl_cos_sin, quantiles=[0.5, 0.2, 0.8])
+            ms, _, _ = triton.testing.do_bench(
+                fn_sgl_cos_sin, quantiles=[0.5, 0.2, 0.8]
+            )
             sgl_cos_sin_time = 1000 * ms
             row_str += f" | {sgl_cos_sin_time:16.4f}"
 
@@ -237,7 +252,9 @@ def fn_flash_attn() -> None:
         if native_vals:
             print(f"  Native: {np.mean(native_vals):.4f} ms")
 
-        sgl_cos_sin_vals = [r["sgl_cos_sin"] for r in results if r["sgl_cos_sin"] is not None]
+        sgl_cos_sin_vals = [
+            r["sgl_cos_sin"] for r in results if r["sgl_cos_sin"] is not None
+        ]
         if sgl_cos_sin_vals:
             print(f"  SGLang (cos/sin API): {np.mean(sgl_cos_sin_vals):.4f} ms")
 
@@ -252,7 +269,7 @@ def fn_flash_attn() -> None:
             fa_vals = [r["flash_attn"] for r in results if r["flash_attn"] is not None]
             if fa_vals:
                 print(f"  flash_attn: {np.mean(fa_vals):.4f} ms")
-   
+
 
 if __name__ == "__main__":
-    benchmark_mm_rotary_embedding()
\ No newline at end of file
+    benchmark_mm_rotary_embedding()
diff --git a/sgl-kernel/csrc/common_extension.cc b/sgl-kernel/csrc/common_extension.cc
index f82d15073cd9..64f821a49b70 100644
--- a/sgl-kernel/csrc/common_extension.cc
+++ b/sgl-kernel/csrc/common_extension.cc
@@ -91,15 +91,18 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) {
   m.impl("apply_rope_pos_ids_cos_sin_cache", torch::kCUDA, &apply_rope_pos_ids_cos_sin_cache);
 
   m.def(
-      "rotary_embedding(Tensor positions, Tensor! query,Tensor!? key, int head_size, Tensor cos_sin_cache, bool is_neox) -> ()");
+      "rotary_embedding(Tensor positions, Tensor! query,Tensor!? key, int head_size, Tensor cos_sin_cache, bool "
+      "is_neox) -> ()");
   m.impl(
       "rotary_embedding",
       torch::kCUDA,
-      static_cast<void (*)(torch::Tensor&, torch::Tensor&, std::optional<torch::Tensor>, int64_t, torch::Tensor&, bool)>(
+      static_cast<void (*)(
+          torch::Tensor&, torch::Tensor&, std::optional<torch::Tensor>, int64_t, torch::Tensor&, bool)>(
           &rotary_embedding));
 
   m.def(
-      "rotary_embedding_cos_sin(Tensor cos, Tensor sin, Tensor(a!) query, Tensor? key, int head_size, bool interleaved) -> ()");
+      "rotary_embedding_cos_sin(Tensor cos, Tensor sin, Tensor(a!) query, Tensor? key, int head_size, bool "
+      "interleaved) -> ()");
   m.impl(
       "rotary_embedding_cos_sin",
       torch::kCUDA,
diff --git a/sgl-kernel/csrc/sgl_diffusion/rope/rotary_embedding.cu b/sgl-kernel/csrc/sgl_diffusion/rope/rotary_embedding.cu
index e7f69ac76e99..eb12ba3eaa6a 100644
--- a/sgl-kernel/csrc/sgl_diffusion/rope/rotary_embedding.cu
+++ b/sgl-kernel/csrc/sgl_diffusion/rope/rotary_embedding.cu
@@ -15,105 +15,104 @@
  * limitations under the License.
  */
 
- #include <ATen/cuda/CUDAContext.h>
- #include <c10/cuda/CUDAException.h>
- #include <c10/cuda/CUDAGuard.h>
- #include <c10/cuda/CUDAStream.h>
- #include <torch/all.h>
- 
- #include <cmath>
- #include <vector>
- 
- #include "utils.h"
- 
- template <typename scalar_t, bool interleaved>
- inline __device__ void apply_token_rotary_embedding(
-     scalar_t* __restrict__ arr,            // [head_size]
-     const scalar_t* __restrict__ cos_ptr,  // [rot_dim]
-     const scalar_t* __restrict__ sin_ptr,  // [rot_dim]
-     int rot_offset,
-     int embed_dim) {  // for non-interleaved: half dim
-   if constexpr (interleaved) {
-     // NeoX-style: interleaved layout [x0, y0, x1, y1, ...].
-     // cos/sin: [..., rotary_dim/2], one entry per (x, y) pair.
-     const int x_index = 2 * rot_offset;
-     const int y_index = x_index + 1;
- 
-     // Directly access SMEM
-     const float cos_val = static_cast<float>(cos_ptr[rot_offset]);
-     const float sin_val = static_cast<float>(sin_ptr[rot_offset]);
- 
-     const float x = static_cast<float>(arr[x_index]);
-     const float y = static_cast<float>(arr[y_index]);
-     arr[x_index] = static_cast<scalar_t>(x * cos_val - y * sin_val);
-     arr[y_index] = static_cast<scalar_t>(y * cos_val + x * sin_val);
-   } else {
-     // GPT-J / LLaMA style: layout [x0, x1, ..., y0, y1, ...]
-     // cos/sin: [..., rotary_dim], one entry per (x, y) pair.
-     const int x_index = rot_offset;
-     const int y_index = rot_offset + embed_dim;
- 
-     // Directly access SMEM
-     const float cos_val_x = static_cast<float>(cos_ptr[rot_offset]);
-     const float sin_val_x = static_cast<float>(sin_ptr[rot_offset]);
-     const float cos_val_y = static_cast<float>(cos_ptr[rot_offset + embed_dim]);
-     const float sin_val_y = static_cast<float>(sin_ptr[rot_offset + embed_dim]);
- 
-     const float x = static_cast<float>(arr[x_index]);
-     const float y = static_cast<float>(arr[y_index]);
-     arr[x_index] = static_cast<scalar_t>(x * cos_val_x - y * sin_val_x);
-     arr[y_index] = static_cast<scalar_t>(y * cos_val_y + x * sin_val_y);
-   }
- }
- 
- template <>
- inline __device__ void apply_token_rotary_embedding<float, true>(
-     float* __restrict__ arr,
-     const float* __restrict__ cos_ptr,
-     const float* __restrict__ sin_ptr,
-     int rot_offset,
-     int /*embed_dim*/) {
-   float2 xy = *reinterpret_cast<const float2*>(arr + rot_offset * 2);
-   
-   // Directly access SMEM
-   const float cos_val = static_cast<float>(cos_ptr[rot_offset]);
-   const float sin_val = static_cast<float>(sin_ptr[rot_offset]);
- 
-   float2 out;
-   out.x = xy.x * cos_val - xy.y * sin_val;
-   out.y = xy.y * cos_val + xy.x * sin_val;
- 
-   *reinterpret_cast<float2*>(arr + rot_offset * 2) = out;
- }
- 
- // 2D grid kernel: parallel over tokens (grid.x) and pair-tiles (grid.y)
- template <typename scalar_t, bool interleaved>
- __global__ void rotary_embedding_kernel_2d(
-     const scalar_t* __restrict__ cos_data,  // [num_tokens, rot_dim_arg]
-     const scalar_t* __restrict__ sin_data,  // [num_tokens, rot_dim_arg]
-     scalar_t* __restrict__ query_total,
-     scalar_t* __restrict__ key_total,
-     const int rot_dim_arg,
-     const int embed_dim_for_rotation,
-     const int64_t query_token_stride,
-     const int64_t key_token_stride,
-     const int64_t head_stride_query,
-     const int64_t head_stride_key,
-     const int num_heads,
-     const int num_kv_heads,
-     const int head_size,
-     const int blocks_per_token) {
-   
-   extern __shared__ char smem_[];
-   scalar_t* s_cos = reinterpret_cast<scalar_t*>(smem_);
-   scalar_t* s_sin = s_cos + rot_dim_arg;
- 
-   const int token_idx = blockIdx.x;
-   if (token_idx >= gridDim.x) {
-     return;
-   }
- 
-   // Pointers to Global Memory for the current token
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAException.h>
+#include <c10/cuda/CUDAGuard.h>
+#include <c10/cuda/CUDAStream.h>
+#include <torch/all.h>
+
+#include <cmath>
+#include <vector>
+
+#include "utils.h"
+
+template <typename scalar_t, bool interleaved>
+inline __device__ void apply_token_rotary_embedding(
+    scalar_t* __restrict__ arr,            // [head_size]
+    const scalar_t* __restrict__ cos_ptr,  // [rot_dim]
+    const scalar_t* __restrict__ sin_ptr,  // [rot_dim]
+    int rot_offset,
+    int embed_dim) {  // for non-interleaved: half dim
+  if constexpr (interleaved) {
+    // NeoX-style: interleaved layout [x0, y0, x1, y1, ...].
+    // cos/sin: [..., rotary_dim/2], one entry per (x, y) pair.
+    const int x_index = 2 * rot_offset;
+    const int y_index = x_index + 1;
+
+    // Directly access SMEM
+    const float cos_val = static_cast<float>(cos_ptr[rot_offset]);
+    const float sin_val = static_cast<float>(sin_ptr[rot_offset]);
+
+    const float x = static_cast<float>(arr[x_index]);
+    const float y = static_cast<float>(arr[y_index]);
+    arr[x_index] = static_cast<scalar_t>(x * cos_val - y * sin_val);
+    arr[y_index] = static_cast<scalar_t>(y * cos_val + x * sin_val);
+  } else {
+    // GPT-J / LLaMA style: layout [x0, x1, ..., y0, y1, ...]
+    // cos/sin: [..., rotary_dim], one entry per (x, y) pair.
+    const int x_index = rot_offset;
+    const int y_index = rot_offset + embed_dim;
+
+    // Directly access SMEM
+    const float cos_val_x = static_cast<float>(cos_ptr[rot_offset]);
+    const float sin_val_x = static_cast<float>(sin_ptr[rot_offset]);
+    const float cos_val_y = static_cast<float>(cos_ptr[rot_offset + embed_dim]);
+    const float sin_val_y = static_cast<float>(sin_ptr[rot_offset + embed_dim]);
+
+    const float x = static_cast<float>(arr[x_index]);
+    const float y = static_cast<float>(arr[y_index]);
+    arr[x_index] = static_cast<scalar_t>(x * cos_val_x - y * sin_val_x);
+    arr[y_index] = static_cast<scalar_t>(y * cos_val_y + x * sin_val_y);
+  }
+}
+
+template <>
+inline __device__ void apply_token_rotary_embedding<float, true>(
+    float* __restrict__ arr,
+    const float* __restrict__ cos_ptr,
+    const float* __restrict__ sin_ptr,
+    int rot_offset,
+    int /*embed_dim*/) {
+  float2 xy = *reinterpret_cast<const float2*>(arr + rot_offset * 2);
+
+  // Directly access SMEM
+  const float cos_val = static_cast<float>(cos_ptr[rot_offset]);
+  const float sin_val = static_cast<float>(sin_ptr[rot_offset]);
+
+  float2 out;
+  out.x = xy.x * cos_val - xy.y * sin_val;
+  out.y = xy.y * cos_val + xy.x * sin_val;
+
+  *reinterpret_cast<float2*>(arr + rot_offset * 2) = out;
+}
+
+// 2D grid kernel: parallel over tokens (grid.x) and pair-tiles (grid.y)
+template <typename scalar_t, bool interleaved>
+__global__ void rotary_embedding_kernel_2d(
+    const scalar_t* __restrict__ cos_data,  // [num_tokens, rot_dim_arg]
+    const scalar_t* __restrict__ sin_data,  // [num_tokens, rot_dim_arg]
+    scalar_t* __restrict__ query_total,
+    scalar_t* __restrict__ key_total,
+    const int rot_dim_arg,
+    const int embed_dim_for_rotation,
+    const int64_t query_token_stride,
+    const int64_t key_token_stride,
+    const int64_t head_stride_query,
+    const int64_t head_stride_key,
+    const int num_heads,
+    const int num_kv_heads,
+    const int head_size,
+    const int blocks_per_token) {
+  extern __shared__ char smem_[];
+  scalar_t* s_cos = reinterpret_cast<scalar_t*>(smem_);
+  scalar_t* s_sin = s_cos + rot_dim_arg;
+
+  const int token_idx = blockIdx.x;
+  if (token_idx >= gridDim.x) {
+    return;
+  }
+
+  // Pointers to Global Memory for the current token
   const scalar_t* current_token_cos_ptr = cos_data + token_idx * rot_dim_arg;
   const scalar_t* current_token_sin_ptr = sin_data + token_idx * rot_dim_arg;
 
@@ -133,8 +132,8 @@
     vec_t* s_sin_vec = reinterpret_cast<vec_t*>(s_sin);
 
     for (int i = threadIdx.x; i < rot_dim_arg / vec_size; i += blockDim.x) {
-      s_cos_vec[i] = __ldg(cos_vec_ptr + i);
-      s_sin_vec[i] = __ldg(sin_vec_ptr + i);
+      s_cos_vec[i] = SGLANG_LDG(cos_vec_ptr + i);
+      s_sin_vec[i] = SGLANG_LDG(sin_vec_ptr + i);
     }
   } else {
     for (int i = threadIdx.x; i < rot_dim_arg; i += blockDim.x) {
@@ -142,64 +141,63 @@
       s_sin[i] = SGLANG_LDG(current_token_sin_ptr + i);
     }
   }
-   // Essential synchronization
-   __syncthreads();
- 
-   scalar_t* query_for_token = query_total + token_idx * (int)query_token_stride;
-   scalar_t* key_for_token = (key_total != nullptr) ? (key_total + token_idx * (int)key_token_stride) : nullptr;
- 
-   const int local_block_idx = blockIdx.y; 
-   const int pair_stride = blockDim.x * blocks_per_token;
-   const int thread_pair_offset = local_block_idx * blockDim.x + threadIdx.x;
- 
-   const int nq_pairs = num_heads * embed_dim_for_rotation;
-   for (int i = thread_pair_offset; i < nq_pairs; i += pair_stride) {
-     const int head_idx = i / embed_dim_for_rotation;
-     const int rot_offset = i % embed_dim_for_rotation;
- 
-     scalar_t* query_for_token_head = query_for_token + head_idx * (int)head_stride_query;
-     apply_token_rotary_embedding<scalar_t, interleaved>(
-         query_for_token_head, s_cos, s_sin, rot_offset, embed_dim_for_rotation);
-   }
- 
-   if (key_for_token != nullptr) {
-     const int nk_pairs = num_kv_heads * embed_dim_for_rotation;
-     for (int i = thread_pair_offset; i < nk_pairs; i += pair_stride) {
-       const int head_idx = i / embed_dim_for_rotation;
-       const int rot_offset = i % embed_dim_for_rotation;
- 
-       scalar_t* key_for_token_head = key_for_token + head_idx * (int)head_stride_key;
-       // Pass Shared Memory pointers
-       apply_token_rotary_embedding<scalar_t, interleaved>(
-           key_for_token_head, s_cos, s_sin, rot_offset, embed_dim_for_rotation);
-     }
-   }
- }
- 
- // 1D grid kernel: each block handles one token
- template <typename scalar_t, bool interleaved>
- __global__ void rotary_embedding_kernel_1d(
-     const scalar_t* __restrict__ cos_data,  // [num_tokens, rot_dim_arg]
-     const scalar_t* __restrict__ sin_data,  // [num_tokens, rot_dim_arg]
-     scalar_t* __restrict__ query_total,
-     scalar_t* __restrict__ key_total,
-     const int rot_dim_arg,
-     const int embed_dim_for_rotation,
-     const int64_t query_token_stride,
-     const int64_t key_token_stride,
-     const int64_t head_stride_query,
-     const int64_t head_stride_key,
-     const int num_heads,
-     const int num_kv_heads,
-     const int head_size) {
-   
-   extern __shared__ char smem_[];
-   scalar_t* s_cos = reinterpret_cast<scalar_t*>(smem_);
-   scalar_t* s_sin = s_cos + rot_dim_arg;
- 
-   const int token_idx = blockIdx.x;
-   if (token_idx >= gridDim.x) return;
- 
+  // Essential synchronization
+  __syncthreads();
+
+  scalar_t* query_for_token = query_total + token_idx * (int)query_token_stride;
+  scalar_t* key_for_token = (key_total != nullptr) ? (key_total + token_idx * (int)key_token_stride) : nullptr;
+
+  const int local_block_idx = blockIdx.y;
+  const int pair_stride = blockDim.x * blocks_per_token;
+  const int thread_pair_offset = local_block_idx * blockDim.x + threadIdx.x;
+
+  const int nq_pairs = num_heads * embed_dim_for_rotation;
+  for (int i = thread_pair_offset; i < nq_pairs; i += pair_stride) {
+    const int head_idx = i / embed_dim_for_rotation;
+    const int rot_offset = i % embed_dim_for_rotation;
+
+    scalar_t* query_for_token_head = query_for_token + head_idx * (int)head_stride_query;
+    apply_token_rotary_embedding<scalar_t, interleaved>(
+        query_for_token_head, s_cos, s_sin, rot_offset, embed_dim_for_rotation);
+  }
+
+  if (key_for_token != nullptr) {
+    const int nk_pairs = num_kv_heads * embed_dim_for_rotation;
+    for (int i = thread_pair_offset; i < nk_pairs; i += pair_stride) {
+      const int head_idx = i / embed_dim_for_rotation;
+      const int rot_offset = i % embed_dim_for_rotation;
+
+      scalar_t* key_for_token_head = key_for_token + head_idx * (int)head_stride_key;
+      // Pass Shared Memory pointers
+      apply_token_rotary_embedding<scalar_t, interleaved>(
+          key_for_token_head, s_cos, s_sin, rot_offset, embed_dim_for_rotation);
+    }
+  }
+}
+
+// 1D grid kernel: each block handles one token
+template <typename scalar_t, bool interleaved>
+__global__ void rotary_embedding_kernel_1d(
+    const scalar_t* __restrict__ cos_data,  // [num_tokens, rot_dim_arg]
+    const scalar_t* __restrict__ sin_data,  // [num_tokens, rot_dim_arg]
+    scalar_t* __restrict__ query_total,
+    scalar_t* __restrict__ key_total,
+    const int rot_dim_arg,
+    const int embed_dim_for_rotation,
+    const int64_t query_token_stride,
+    const int64_t key_token_stride,
+    const int64_t head_stride_query,
+    const int64_t head_stride_key,
+    const int num_heads,
+    const int num_kv_heads,
+    const int head_size) {
+  extern __shared__ char smem_[];
+  scalar_t* s_cos = reinterpret_cast<scalar_t*>(smem_);
+  scalar_t* s_sin = s_cos + rot_dim_arg;
+
+  const int token_idx = blockIdx.x;
+  if (token_idx >= gridDim.x) return;
+
   const scalar_t* current_token_cos_ptr = cos_data + token_idx * rot_dim_arg;
   const scalar_t* current_token_sin_ptr = sin_data + token_idx * rot_dim_arg;
 
@@ -219,8 +217,8 @@
     vec_t* s_sin_vec = reinterpret_cast<vec_t*>(s_sin);
 
     for (int i = threadIdx.x; i < rot_dim_arg / vec_size; i += blockDim.x) {
-      s_cos_vec[i] = __ldg(cos_vec_ptr + i);
-      s_sin_vec[i] = __ldg(sin_vec_ptr + i);
+      s_cos_vec[i] = SGLANG_LDG(cos_vec_ptr + i);
+      s_sin_vec[i] = SGLANG_LDG(sin_vec_ptr + i);
     }
   } else {
     for (int i = threadIdx.x; i < rot_dim_arg; i += blockDim.x) {
@@ -228,220 +226,220 @@
       s_sin[i] = SGLANG_LDG(current_token_sin_ptr + i);
     }
   }
-   __syncthreads();
- 
-   scalar_t* query_for_token = query_total + token_idx * (int)query_token_stride;
-   scalar_t* key_for_token = (key_total != nullptr) ? (key_total + token_idx * (int)key_token_stride) : nullptr;
- 
-   const int nq_pairs = num_heads * embed_dim_for_rotation;
-   for (int i = threadIdx.x; i < nq_pairs; i += blockDim.x) {
-     const int head_idx = i / embed_dim_for_rotation;
-     const int rot_offset = i % embed_dim_for_rotation;
-     scalar_t* query_for_token_head = query_for_token + head_idx * (int)head_stride_query;
-      
-     apply_token_rotary_embedding<scalar_t, interleaved>(
-         query_for_token_head, s_cos, s_sin, rot_offset, embed_dim_for_rotation);
-   }
- 
-   if (key_for_token != nullptr) {
-     const int nk_pairs = num_kv_heads * embed_dim_for_rotation;
-     for (int i = threadIdx.x; i < nk_pairs; i += blockDim.x) {
-       const int head_idx = i / embed_dim_for_rotation;
-       const int rot_offset = i % embed_dim_for_rotation;
-       
-       scalar_t* key_for_token_head = key_for_token + head_idx * (int)head_stride_key;
-       apply_token_rotary_embedding<scalar_t, interleaved>(
-           key_for_token_head, s_cos, s_sin, rot_offset, embed_dim_for_rotation);
-     }
-   }
- }
- 
- void rotary_embedding_cos_sin(
-     at::Tensor& cos,
-     at::Tensor& sin,
-     at::Tensor& query,
-     const std::optional<at::Tensor>& key,
-     int64_t head_size,
-     bool interleaved) {
-   TORCH_CHECK(
-       query.dim() == 2 || query.dim() == 3,
-       "query must be in shape [num_tokens, hidden_size] or [num_tokens, num_heads, head_size]");
-   if (key.has_value()) {
-     TORCH_CHECK(
-         key->dim() == 2 || key->dim() == 3,
-         "key must be in shape [num_tokens, hidden_size] or [num_tokens, num_kv_heads, head_size]");
-   }
- 
-   const int64_t num_tokens = query.size(0);
- 
-   TORCH_CHECK(cos.dim() == 2, "cos must be in shape [num_tokens, D_cos]");
-   TORCH_CHECK(sin.dim() == 2, "sin must be in shape [num_tokens, D_sin]");
-   TORCH_CHECK(cos.size(0) == num_tokens, "cos num_tokens mismatch with query");
-   TORCH_CHECK(sin.size(0) == num_tokens, "sin num_tokens mismatch with query");
-   TORCH_CHECK(cos.size(1) == sin.size(1), "cos and sin D_cos/D_sin mismatch");
- 
-   TORCH_CHECK(cos.scalar_type() == query.scalar_type(), "cos dtype mismatch with query");
-   TORCH_CHECK(sin.scalar_type() == query.scalar_type(), "sin dtype mismatch with query");
-   TORCH_CHECK(cos.is_cuda() && sin.is_cuda() && query.is_cuda(), "cos/sin/query must be CUDA tensors");
-   TORCH_CHECK(query.is_contiguous(), "query must be contiguous");
-   if (key.has_value()) {
-     TORCH_CHECK(key->is_cuda(), "key must be CUDA tensor if provided");
-     TORCH_CHECK(key->scalar_type() == query.scalar_type(), "key dtype mismatch with query");
-     TORCH_CHECK(key->is_contiguous(), "key must be contiguous");
-   }
- 
-   int query_hidden_size_calculated;
-   if (query.dim() == 2) {
-     query_hidden_size_calculated = (int)query.size(1);
-   } else {
-     query_hidden_size_calculated = (int)query.size(1) * (int)query.size(2);
-     TORCH_CHECK(query.size(2) == head_size, "query head_size mismatch in 3D tensor");
-   }
-   TORCH_CHECK(query_hidden_size_calculated % head_size == 0, "query_hidden_size not divisible by head_size");
-   int num_heads = query_hidden_size_calculated / (int)head_size;
- 
-   int key_hidden_size_calculated = 0;
-   int num_kv_heads = num_heads;
-   if (key.has_value()) {
-     TORCH_CHECK((int)key->size(0) == num_tokens, "key num_tokens mismatch");
-     if (key->dim() == 2) {
-       key_hidden_size_calculated = (int)key->size(1);
-     } else {
-       key_hidden_size_calculated = (int)key->size(1) * (int)key->size(2);
-       TORCH_CHECK((int)key->size(2) == head_size, "key head_size mismatch in 3D tensor");
-     }
-     TORCH_CHECK(key_hidden_size_calculated % head_size == 0, "key_hidden_size not divisible by head_size");
-     num_kv_heads = key_hidden_size_calculated / (int)head_size;
-   }
-   TORCH_CHECK(num_heads % num_kv_heads == 0, "num_heads must be divisible by num_kv_heads");
- 
-   // NeoX interleaved: if cos/sin are full-dim (head_size), downsample once.
-   if (interleaved && cos.size(1) == head_size) {
-     TORCH_CHECK(head_size % 2 == 0, "interleaved layout requires even head_size");
-     const int64_t half = head_size / 2;
-     std::vector<int64_t> new_shape = {cos.size(0), half, 2};
-     cos = cos.view(new_shape).select(2, 0).contiguous();
-     sin = sin.view(new_shape).select(2, 1).contiguous();
-   }
- 
-   const int rot_dim_from_cache = (int)cos.size(1);
- 
-   const int64_t query_token_stride = query_hidden_size_calculated;
-   const int64_t key_token_stride = key.has_value() ? key_hidden_size_calculated : 0;
- 
-   int64_t head_stride_query;
-   if (query.dim() == 3 && query.size(1) == num_heads && query.size(2) == head_size) {
-     head_stride_query = query.stride(1);
-   } else {
-     head_stride_query = head_size;
-   }
- 
-   int64_t head_stride_key = head_size;
-   if (key.has_value()) {
-     if (key->dim() == 3 && key->size(1) == num_kv_heads && key->size(2) == head_size) {
-       head_stride_key = key->stride(1);
-     } else {
-       head_stride_key = head_size;
-     }
-   }
- 
-   const int embed_dim_for_rotation = interleaved ? rot_dim_from_cache : (rot_dim_from_cache / 2);
-   TORCH_CHECK(embed_dim_for_rotation > 0, "embed_dim_for_rotation must be > 0");
- 
-   const int max_pairs_to_rotate_per_token =
-       std::max(num_heads * embed_dim_for_rotation, num_kv_heads * embed_dim_for_rotation);
- 
-   const int threads_per_block = std::min<int>(256, std::max(128, embed_dim_for_rotation));
-   const int blocks_per_token = (max_pairs_to_rotate_per_token + threads_per_block - 1) / threads_per_block;
- 
-   const bool use_grid_2d = (num_tokens <= 4) && (blocks_per_token > 1);
-   dim3 block(threads_per_block);
-   dim3 grid_2d((int)num_tokens, std::max(1, blocks_per_token));
-   dim3 grid_1d((int)num_tokens);
- 
-   const at::cuda::OptionalCUDAGuard device_guard(device_of(query));
-   const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
- 
-   AT_DISPATCH_FLOATING_TYPES_AND2(
-       at::ScalarType::Half, at::ScalarType::BFloat16, query.scalar_type(), "rotary_embedding_cos_sin", [&] {
-         using torch_scalar_t = scalar_t;
-         using cuda_scalar_t = typename std::conditional<
-             std::is_same<torch_scalar_t, at::Half>::value,
-             nv_half,
-             typename std::conditional<std::is_same<torch_scalar_t, at::BFloat16>::value, nv_bfloat16, torch_scalar_t>::
-                 type>::type;
-         
-         // We need 2 arrays (cos, sin) of size 'rot_dim_from_cache', each element is 'sizeof(torch_scalar_t)'
-         size_t smem_size = rot_dim_from_cache * sizeof(torch_scalar_t) * 2;
- 
-         if (interleaved) {
-           if (use_grid_2d) {
-             rotary_embedding_kernel_2d<cuda_scalar_t, true><<<grid_2d, block, smem_size, stream>>>(
-                 reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
-                 reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
-                 reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
-                 key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>()) : nullptr,
-                 rot_dim_from_cache,
-                 embed_dim_for_rotation,
-                 query_token_stride,
-                 key_token_stride,
-                 head_stride_query,
-                 head_stride_key,
-                 num_heads,
-                 num_kv_heads,
-                 (int)head_size,
-                 blocks_per_token);
-           } else {
-             rotary_embedding_kernel_1d<cuda_scalar_t, true><<<grid_1d, block, smem_size, stream>>>(
-                 reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
-                 reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
-                 reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
-                 key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>()) : nullptr,
-                 rot_dim_from_cache,
-                 embed_dim_for_rotation,
-                 query_token_stride,
-                 key_token_stride,
-                 head_stride_query,
-                 head_stride_key,
-                 num_heads,
-                 num_kv_heads,
-                 (int)head_size);
-           }
-         } else {
-           if (use_grid_2d) {
-             rotary_embedding_kernel_2d<cuda_scalar_t, false><<<grid_2d, block, smem_size, stream>>>(
-                 reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
-                 reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
-                 reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
-                 key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>()) : nullptr,
-                 rot_dim_from_cache,
-                 embed_dim_for_rotation,
-                 query_token_stride,
-                 key_token_stride,
-                 head_stride_query,
-                 head_stride_key,
-                 num_heads,
-                 num_kv_heads,
-                 (int)head_size,
-                 blocks_per_token);
-           } else {
-             rotary_embedding_kernel_1d<cuda_scalar_t, false><<<grid_1d, block, smem_size, stream>>>(
-                 reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
-                 reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
-                 reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
-                 key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>()) : nullptr,
-                 rot_dim_from_cache,
-                 embed_dim_for_rotation,
-                 query_token_stride,
-                 key_token_stride,
-                 head_stride_query,
-                 head_stride_key,
-                 num_heads,
-                 num_kv_heads,
-                 (int)head_size);
-           }
-         }
-       });
- 
-   C10_CUDA_KERNEL_LAUNCH_CHECK();
- }
\ No newline at end of file
+  __syncthreads();
+
+  scalar_t* query_for_token = query_total + token_idx * (int)query_token_stride;
+  scalar_t* key_for_token = (key_total != nullptr) ? (key_total + token_idx * (int)key_token_stride) : nullptr;
+
+  const int nq_pairs = num_heads * embed_dim_for_rotation;
+  for (int i = threadIdx.x; i < nq_pairs; i += blockDim.x) {
+    const int head_idx = i / embed_dim_for_rotation;
+    const int rot_offset = i % embed_dim_for_rotation;
+    scalar_t* query_for_token_head = query_for_token + head_idx * (int)head_stride_query;
+
+    apply_token_rotary_embedding<scalar_t, interleaved>(
+        query_for_token_head, s_cos, s_sin, rot_offset, embed_dim_for_rotation);
+  }
+
+  if (key_for_token != nullptr) {
+    const int nk_pairs = num_kv_heads * embed_dim_for_rotation;
+    for (int i = threadIdx.x; i < nk_pairs; i += blockDim.x) {
+      const int head_idx = i / embed_dim_for_rotation;
+      const int rot_offset = i % embed_dim_for_rotation;
+
+      scalar_t* key_for_token_head = key_for_token + head_idx * (int)head_stride_key;
+      apply_token_rotary_embedding<scalar_t, interleaved>(
+          key_for_token_head, s_cos, s_sin, rot_offset, embed_dim_for_rotation);
+    }
+  }
+}
+
+void rotary_embedding_cos_sin(
+    at::Tensor& cos,
+    at::Tensor& sin,
+    at::Tensor& query,
+    const std::optional<at::Tensor>& key,
+    int64_t head_size,
+    bool interleaved) {
+  TORCH_CHECK(
+      query.dim() == 2 || query.dim() == 3,
+      "query must be in shape [num_tokens, hidden_size] or [num_tokens, num_heads, head_size]");
+  if (key.has_value()) {
+    TORCH_CHECK(
+        key->dim() == 2 || key->dim() == 3,
+        "key must be in shape [num_tokens, hidden_size] or [num_tokens, num_kv_heads, head_size]");
+  }
+
+  const int64_t num_tokens = query.size(0);
+
+  TORCH_CHECK(cos.dim() == 2, "cos must be in shape [num_tokens, D_cos]");
+  TORCH_CHECK(sin.dim() == 2, "sin must be in shape [num_tokens, D_sin]");
+  TORCH_CHECK(cos.size(0) == num_tokens, "cos num_tokens mismatch with query");
+  TORCH_CHECK(sin.size(0) == num_tokens, "sin num_tokens mismatch with query");
+  TORCH_CHECK(cos.size(1) == sin.size(1), "cos and sin D_cos/D_sin mismatch");
+
+  TORCH_CHECK(cos.scalar_type() == query.scalar_type(), "cos dtype mismatch with query");
+  TORCH_CHECK(sin.scalar_type() == query.scalar_type(), "sin dtype mismatch with query");
+  TORCH_CHECK(cos.is_cuda() && sin.is_cuda() && query.is_cuda(), "cos/sin/query must be CUDA tensors");
+  TORCH_CHECK(query.is_contiguous(), "query must be contiguous");
+  if (key.has_value()) {
+    TORCH_CHECK(key->is_cuda(), "key must be CUDA tensor if provided");
+    TORCH_CHECK(key->scalar_type() == query.scalar_type(), "key dtype mismatch with query");
+    TORCH_CHECK(key->is_contiguous(), "key must be contiguous");
+  }
+
+  int query_hidden_size_calculated;
+  if (query.dim() == 2) {
+    query_hidden_size_calculated = (int)query.size(1);
+  } else {
+    query_hidden_size_calculated = (int)query.size(1) * (int)query.size(2);
+    TORCH_CHECK(query.size(2) == head_size, "query head_size mismatch in 3D tensor");
+  }
+  TORCH_CHECK(query_hidden_size_calculated % head_size == 0, "query_hidden_size not divisible by head_size");
+  int num_heads = query_hidden_size_calculated / (int)head_size;
+
+  int key_hidden_size_calculated = 0;
+  int num_kv_heads = num_heads;
+  if (key.has_value()) {
+    TORCH_CHECK((int)key->size(0) == num_tokens, "key num_tokens mismatch");
+    if (key->dim() == 2) {
+      key_hidden_size_calculated = (int)key->size(1);
+    } else {
+      key_hidden_size_calculated = (int)key->size(1) * (int)key->size(2);
+      TORCH_CHECK((int)key->size(2) == head_size, "key head_size mismatch in 3D tensor");
+    }
+    TORCH_CHECK(key_hidden_size_calculated % head_size == 0, "key_hidden_size not divisible by head_size");
+    num_kv_heads = key_hidden_size_calculated / (int)head_size;
+  }
+  TORCH_CHECK(num_heads % num_kv_heads == 0, "num_heads must be divisible by num_kv_heads");
+
+  // NeoX interleaved: if cos/sin are full-dim (head_size), downsample once.
+  if (interleaved && cos.size(1) == head_size) {
+    TORCH_CHECK(head_size % 2 == 0, "interleaved layout requires even head_size");
+    const int64_t half = head_size / 2;
+    std::vector<int64_t> new_shape = {cos.size(0), half, 2};
+    cos = cos.view(new_shape).select(2, 0).contiguous();
+    sin = sin.view(new_shape).select(2, 1).contiguous();
+  }
+
+  const int rot_dim_from_cache = (int)cos.size(1);
+
+  const int64_t query_token_stride = query_hidden_size_calculated;
+  const int64_t key_token_stride = key.has_value() ? key_hidden_size_calculated : 0;
+
+  int64_t head_stride_query;
+  if (query.dim() == 3 && query.size(1) == num_heads && query.size(2) == head_size) {
+    head_stride_query = query.stride(1);
+  } else {
+    head_stride_query = head_size;
+  }
+
+  int64_t head_stride_key = head_size;
+  if (key.has_value()) {
+    if (key->dim() == 3 && key->size(1) == num_kv_heads && key->size(2) == head_size) {
+      head_stride_key = key->stride(1);
+    } else {
+      head_stride_key = head_size;
+    }
+  }
+
+  const int embed_dim_for_rotation = interleaved ? rot_dim_from_cache : (rot_dim_from_cache / 2);
+  TORCH_CHECK(embed_dim_for_rotation > 0, "embed_dim_for_rotation must be > 0");
+
+  const int max_pairs_to_rotate_per_token =
+      std::max(num_heads * embed_dim_for_rotation, num_kv_heads * embed_dim_for_rotation);
+
+  const int threads_per_block = std::min<int>(256, std::max(128, embed_dim_for_rotation));
+  const int blocks_per_token = (max_pairs_to_rotate_per_token + threads_per_block - 1) / threads_per_block;
+
+  const bool use_grid_2d = (num_tokens <= 4) && (blocks_per_token > 1);
+  dim3 block(threads_per_block);
+  dim3 grid_2d((int)num_tokens, std::max(1, blocks_per_token));
+  dim3 grid_1d((int)num_tokens);
+
+  const at::cuda::OptionalCUDAGuard device_guard(device_of(query));
+  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+
+  AT_DISPATCH_FLOATING_TYPES_AND2(
+      at::ScalarType::Half, at::ScalarType::BFloat16, query.scalar_type(), "rotary_embedding_cos_sin", [&] {
+        using torch_scalar_t = scalar_t;
+        using cuda_scalar_t = typename std::conditional<
+            std::is_same<torch_scalar_t, at::Half>::value,
+            nv_half,
+            typename std::conditional<std::is_same<torch_scalar_t, at::BFloat16>::value, nv_bfloat16, torch_scalar_t>::
+                type>::type;
+
+        // We need 2 arrays (cos, sin) of size 'rot_dim_from_cache', each element is 'sizeof(torch_scalar_t)'
+        size_t smem_size = rot_dim_from_cache * sizeof(torch_scalar_t) * 2;
+
+        if (interleaved) {
+          if (use_grid_2d) {
+            rotary_embedding_kernel_2d<cuda_scalar_t, true><<<grid_2d, block, smem_size, stream>>>(
+                reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
+                reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
+                reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
+                key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>()) : nullptr,
+                rot_dim_from_cache,
+                embed_dim_for_rotation,
+                query_token_stride,
+                key_token_stride,
+                head_stride_query,
+                head_stride_key,
+                num_heads,
+                num_kv_heads,
+                (int)head_size,
+                blocks_per_token);
+          } else {
+            rotary_embedding_kernel_1d<cuda_scalar_t, true><<<grid_1d, block, smem_size, stream>>>(
+                reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
+                reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
+                reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
+                key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>()) : nullptr,
+                rot_dim_from_cache,
+                embed_dim_for_rotation,
+                query_token_stride,
+                key_token_stride,
+                head_stride_query,
+                head_stride_key,
+                num_heads,
+                num_kv_heads,
+                (int)head_size);
+          }
+        } else {
+          if (use_grid_2d) {
+            rotary_embedding_kernel_2d<cuda_scalar_t, false><<<grid_2d, block, smem_size, stream>>>(
+                reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
+                reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
+                reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
+                key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>()) : nullptr,
+                rot_dim_from_cache,
+                embed_dim_for_rotation,
+                query_token_stride,
+                key_token_stride,
+                head_stride_query,
+                head_stride_key,
+                num_heads,
+                num_kv_heads,
+                (int)head_size,
+                blocks_per_token);
+          } else {
+            rotary_embedding_kernel_1d<cuda_scalar_t, false><<<grid_1d, block, smem_size, stream>>>(
+                reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
+                reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
+                reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
+                key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>()) : nullptr,
+                rot_dim_from_cache,
+                embed_dim_for_rotation,
+                query_token_stride,
+                key_token_stride,
+                head_stride_query,
+                head_stride_key,
+                num_heads,
+                num_kv_heads,
+                (int)head_size);
+          }
+        }
+      });
+
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+}
diff --git a/sgl-kernel/python/sgl_kernel/__init__.py b/sgl-kernel/python/sgl_kernel/__init__.py
index 191360421816..c350dfbf468f 100644
--- a/sgl-kernel/python/sgl_kernel/__init__.py
+++ b/sgl-kernel/python/sgl_kernel/__init__.py
@@ -33,7 +33,6 @@
     rotary_embedding,
     silu_and_mul,
 )
-from sgl_kernel.rotary_embedding import rotary_embedding_cos_sin
 from sgl_kernel.expert_specialization import (
     es_fp8_blockwise_scaled_grouped_mm,
     es_sm100_mxfp8_blockscaled_grouped_mm,
@@ -108,6 +107,7 @@
     ggml_mul_mat_a8,
     ggml_mul_mat_vec_a8,
 )
+from sgl_kernel.rotary_embedding import rotary_embedding_cos_sin
 from sgl_kernel.sampling import (
     min_p_sampling_from_probs,
     top_k_mask_logits,
diff --git a/sgl-kernel/python/sgl_kernel/rotary_embedding.py b/sgl-kernel/python/sgl_kernel/rotary_embedding.py
index 0c187e2ae4c1..ed890b759096 100644
--- a/sgl-kernel/python/sgl_kernel/rotary_embedding.py
+++ b/sgl-kernel/python/sgl_kernel/rotary_embedding.py
@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 from typing import Literal, Optional
+
 import torch
 
 _HAS_COS_SIN = hasattr(torch.ops.sgl_kernel, "rotary_embedding_cos_sin")
@@ -74,14 +75,18 @@ def apply_rotary_embedding(
             )
             return
 
-        raise RuntimeError("No cos/sin rotary embedding kernel is available in torch.ops.sgl_kernel")
+        raise RuntimeError(
+            "No cos/sin rotary embedding kernel is available in torch.ops.sgl_kernel"
+        )
 
     if mode == "positions":
         if positions is None or cos_sin_cache is None:
             raise ValueError("mode='positions' requires positions and cos_sin_cache")
 
         if not _HAS_GENERIC:
-            raise RuntimeError("No positions rotary embedding kernel is available in torch.ops.sgl_kernel")
+            raise RuntimeError(
+                "No positions rotary embedding kernel is available in torch.ops.sgl_kernel"
+            )
 
         # Prefer keyword call if supported.
         try:
@@ -133,6 +138,7 @@ def rotary_embedding_cos_sin(
         is_neox=is_neox,
     )
 
+
 def rotary_embedding(
     positions: torch.Tensor,
     query: torch.Tensor,
@@ -152,4 +158,4 @@ def rotary_embedding(
         head_size=head_size,
         interleaved=interleaved,
         is_neox=is_neox,
-    )
\ No newline at end of file
+    )
diff --git a/sgl-kernel/tests/test_mm_rotary_embedding.py b/sgl-kernel/tests/test_mm_rotary_embedding.py
old mode 100644
new mode 100755
index 2a7ee8155361..4fc454fa9b44
--- a/sgl-kernel/tests/test_mm_rotary_embedding.py
+++ b/sgl-kernel/tests/test_mm_rotary_embedding.py
@@ -4,19 +4,23 @@
 import os
 import sys
 from dataclasses import dataclass
-from typing import Tuple, List
+from typing import List, Tuple
 
 import torch
 
 try:
     import pytest
+
     HAS_PYTEST = True
 except ImportError:
     HAS_PYTEST = False
 
 try:
     # Cos/sin-based rotary kernel from sgl_diffusion
-    from sgl_kernel.rotary_embedding import rotary_embedding_cos_sin as rotary_emb_module
+    from sgl_kernel.rotary_embedding import (
+        rotary_embedding_cos_sin as rotary_emb_module,
+    )
+
     HAS_ROTARY_EMBEDDING = True
 except ImportError:
     rotary_emb_module = None
@@ -42,7 +46,9 @@ def compute_cos_sin_cache(
     """
     Compute separate cos and sin caches.
     """
-    inv_freq = 1.0 / (base ** (torch.arange(0, rotary_dim, 2, dtype=torch.float32) / rotary_dim))
+    inv_freq = 1.0 / (
+        base ** (torch.arange(0, rotary_dim, 2, dtype=torch.float32) / rotary_dim)
+    )
     t = torch.arange(max_seq_len, dtype=torch.float32)
     freqs = torch.einsum("i,j->ij", t, inv_freq)
     cos = freqs.cos().to(dtype)
@@ -58,7 +64,7 @@ def reference_rotary_neox(
     head_size: int,
 ) -> Tuple[torch.Tensor, torch.Tensor]:
     """
-    Reference implementation of NeoX-style rotary embedding, 
+    Reference implementation of NeoX-style rotary embedding,
     using explicit cos/sin caches and Q/K in flattened (num_tokens, heads * dim) layout.
     """
     num_tokens = query.size(0)
@@ -124,7 +130,11 @@ def _check_rotary_correctness(
     device: str = "cuda",
 ) -> RotaryTestResult:
     if not HAS_ROTARY_EMBEDDING:
-        return RotaryTestResult("basic (skipped: rotary_embedding not built)", True, details="rotary_embedding extension not found")
+        return RotaryTestResult(
+            "basic (skipped: rotary_embedding not built)",
+            True,
+            details="rotary_embedding extension not found",
+        )
 
     rotary_dim = head_size
     max_seq_len = 8192
@@ -139,7 +149,9 @@ def _check_rotary_correctness(
     cos = cos_cache[positions]
     sin = sin_cache[positions]
 
-    q_ref, k_ref = reference_rotary_neox(query.clone(), key.clone(), cos, sin, head_size)
+    q_ref, k_ref = reference_rotary_neox(
+        query.clone(), key.clone(), cos, sin, head_size
+    )
 
     q_out = query.clone()
     k_out = key.clone()
@@ -168,7 +180,7 @@ def test_rotary_embedding_correctness(
 ) -> None:
     if not HAS_ROTARY_EMBEDDING:
         pytest.skip("sgl_kernel.rotary_embedding not available")
-    
+
     result = _check_rotary_correctness(
         batch_size=batch_size,
         seq_len=seq_len,
@@ -177,7 +189,9 @@ def test_rotary_embedding_correctness(
         head_size=head_size,
         dtype=dtype,
     )
-    assert result.passed, f"{result.name} failed: Q={result.q_diff:.2e}, K={result.k_diff:.2e}"
+    assert (
+        result.passed
+    ), f"{result.name} failed: Q={result.q_diff:.2e}, K={result.k_diff:.2e}"
 
 
 if __name__ == "__main__":

From d23edf281ff985e0b5690d620003ebbc55b33659 Mon Sep 17 00:00:00 2001
From: Ther-LF <2639852836@qq.com>
Date: Tue, 16 Dec 2025 21:42:11 +0800
Subject: [PATCH 08/37] perf:vec load and compute;modify threads_per_block

---
 .../sgl_diffusion/rope/rotary_embedding.cu    | 452 +++++++++++++-----
 1 file changed, 332 insertions(+), 120 deletions(-)

diff --git a/sgl-kernel/csrc/sgl_diffusion/rope/rotary_embedding.cu b/sgl-kernel/csrc/sgl_diffusion/rope/rotary_embedding.cu
index eb12ba3eaa6a..057e6cba0651 100644
--- a/sgl-kernel/csrc/sgl_diffusion/rope/rotary_embedding.cu
+++ b/sgl-kernel/csrc/sgl_diffusion/rope/rotary_embedding.cu
@@ -66,6 +66,88 @@ inline __device__ void apply_token_rotary_embedding(
   }
 }
 
+template <typename scalar_t, bool interleaved>
+inline __device__ void apply_token_rotary_embedding_vec(
+    scalar_t* __restrict__ arr,            // [head_size]
+    const scalar_t* __restrict__ cos_ptr,  // [rot_dim]
+    const scalar_t* __restrict__ sin_ptr,  // [rot_dim]
+    int rot_offset,
+    int embed_dim) {
+  
+  using vec_t = float4;
+  constexpr int kVecBytes = sizeof(vec_t);
+  constexpr int kScalarBytes = sizeof(scalar_t);
+  constexpr int kElePerVec = kVecBytes / kScalarBytes;
+  
+  // Union for type punning to avoid strict aliasing issues with reinterpret_cast
+  union VecUnion {
+    vec_t vec;
+    scalar_t elems[kElePerVec];
+  };
+
+  if constexpr (interleaved) {
+    // Interleaved: arr has [x0, y0, x1, y1...]
+    // A single vector load contains 'kElePerVec' elements, which is 'kElePerVec / 2' pairs.
+    // rot_offset is the index of the PAIR.
+    // Address in arr is rot_offset * 2.
+    
+    VecUnion data;
+    data.vec = *reinterpret_cast<const vec_t*>(arr + rot_offset * 2);
+    
+    #pragma unroll
+    for (int i = 0; i < kElePerVec; i += 2) {
+      // data.elems[i] is x, data.elems[i+1] is y
+      // They correspond to pair index: rot_offset + (i / 2)
+      int curr_rot_offset = rot_offset + i / 2;
+      
+      float cos_val = static_cast<float>(cos_ptr[curr_rot_offset]);
+      float sin_val = static_cast<float>(sin_ptr[curr_rot_offset]);
+      
+      float x = static_cast<float>(data.elems[i]);
+      float y = static_cast<float>(data.elems[i+1]);
+      
+      data.elems[i] = static_cast<scalar_t>(x * cos_val - y * sin_val);
+      data.elems[i+1] = static_cast<scalar_t>(y * cos_val + x * sin_val);
+    }
+    
+    *reinterpret_cast<vec_t*>(arr + rot_offset * 2) = data.vec;
+
+  } else {
+    // Non-interleaved: X and Y are separated by embed_dim.
+    // We process 'kElePerVec' PAIRS at once.
+    // Load X vector and Y vector.
+    
+    VecUnion data_x, data_y;
+    data_x.vec = *reinterpret_cast<const vec_t*>(arr + rot_offset);
+    data_y.vec = *reinterpret_cast<const vec_t*>(arr + rot_offset + embed_dim);
+    
+    #pragma unroll
+    for (int i = 0; i < kElePerVec; ++i) {
+      int curr_rot_offset = rot_offset + i;
+      
+      // In non-interleaved, we might need different cos/sin for X and Y depending on implementation,
+      // but standard RoPE uses the same angle for the pair.
+      // Based on original scalar code:
+      // cos_val_x = cos_ptr[rot_offset]
+      // cos_val_y = cos_ptr[rot_offset + embed_dim]
+      
+      float cos_val_x = static_cast<float>(cos_ptr[curr_rot_offset]);
+      float sin_val_x = static_cast<float>(sin_ptr[curr_rot_offset]);
+      float cos_val_y = static_cast<float>(cos_ptr[curr_rot_offset + embed_dim]);
+      float sin_val_y = static_cast<float>(sin_ptr[curr_rot_offset + embed_dim]);
+      
+      float x = static_cast<float>(data_x.elems[i]);
+      float y = static_cast<float>(data_y.elems[i]);
+      
+      data_x.elems[i] = static_cast<scalar_t>(x * cos_val_x - y * sin_val_x);
+      data_y.elems[i] = static_cast<scalar_t>(y * cos_val_y + x * sin_val_y);
+    }
+    
+    *reinterpret_cast<vec_t*>(arr + rot_offset) = data_x.vec;
+    *reinterpret_cast<vec_t*>(arr + rot_offset + embed_dim) = data_y.vec;
+  }
+}
+
 template <>
 inline __device__ void apply_token_rotary_embedding<float, true>(
     float* __restrict__ arr,
@@ -75,7 +157,6 @@ inline __device__ void apply_token_rotary_embedding<float, true>(
     int /*embed_dim*/) {
   float2 xy = *reinterpret_cast<const float2*>(arr + rot_offset * 2);
 
-  // Directly access SMEM
   const float cos_val = static_cast<float>(cos_ptr[rot_offset]);
   const float sin_val = static_cast<float>(sin_ptr[rot_offset]);
 
@@ -87,7 +168,7 @@ inline __device__ void apply_token_rotary_embedding<float, true>(
 }
 
 // 2D grid kernel: parallel over tokens (grid.x) and pair-tiles (grid.y)
-template <typename scalar_t, bool interleaved>
+template <typename scalar_t, bool interleaved, bool vectorized>
 __global__ void rotary_embedding_kernel_2d(
     const scalar_t* __restrict__ cos_data,  // [num_tokens, rot_dim_arg]
     const scalar_t* __restrict__ sin_data,  // [num_tokens, rot_dim_arg]
@@ -120,11 +201,9 @@ __global__ void rotary_embedding_kernel_2d(
   const int scalar_size = sizeof(scalar_t);
   const int vec_size = kVecBytes / scalar_size;
 
-  bool is_aligned = (rot_dim_arg % vec_size == 0) &&
-                    (reinterpret_cast<uintptr_t>(current_token_cos_ptr) % kVecBytes == 0) &&
-                    (reinterpret_cast<uintptr_t>(current_token_sin_ptr) % kVecBytes == 0);
-
-  if (is_aligned) {
+  // Load Cos/Sin to SMEM
+  // Always use vectorized load if available, but fallback to scalar loop if dim is not multiple
+  if constexpr (vectorized) {
     using vec_t = float4;
     const vec_t* cos_vec_ptr = reinterpret_cast<const vec_t*>(current_token_cos_ptr);
     const vec_t* sin_vec_ptr = reinterpret_cast<const vec_t*>(current_token_sin_ptr);
@@ -148,35 +227,67 @@ __global__ void rotary_embedding_kernel_2d(
   scalar_t* key_for_token = (key_total != nullptr) ? (key_total + token_idx * (int)key_token_stride) : nullptr;
 
   const int local_block_idx = blockIdx.y;
-  const int pair_stride = blockDim.x * blocks_per_token;
-  const int thread_pair_offset = local_block_idx * blockDim.x + threadIdx.x;
-
-  const int nq_pairs = num_heads * embed_dim_for_rotation;
-  for (int i = thread_pair_offset; i < nq_pairs; i += pair_stride) {
-    const int head_idx = i / embed_dim_for_rotation;
-    const int rot_offset = i % embed_dim_for_rotation;
-
-    scalar_t* query_for_token_head = query_for_token + head_idx * (int)head_stride_query;
-    apply_token_rotary_embedding<scalar_t, interleaved>(
-        query_for_token_head, s_cos, s_sin, rot_offset, embed_dim_for_rotation);
-  }
-
-  if (key_for_token != nullptr) {
-    const int nk_pairs = num_kv_heads * embed_dim_for_rotation;
-    for (int i = thread_pair_offset; i < nk_pairs; i += pair_stride) {
+  
+  if constexpr (vectorized) {
+    using vec_t = float4;
+    constexpr int kElePerVec = kVecBytes / sizeof(scalar_t);
+    constexpr int pairs_per_step = interleaved ? (kElePerVec / 2) : kElePerVec;
+    
+    const int pair_stride = blockDim.x * blocks_per_token * pairs_per_step;
+    const int thread_pair_offset = (local_block_idx * blockDim.x + threadIdx.x) * pairs_per_step;
+    const int nq_pairs = num_heads * embed_dim_for_rotation;
+    
+    for (int i = thread_pair_offset; i < nq_pairs; i += pair_stride) {
+      const int head_idx = i / embed_dim_for_rotation;
+      const int rot_offset = i % embed_dim_for_rotation;
+      
+      scalar_t* query_for_token_head = query_for_token + head_idx * (int)head_stride_query;
+      apply_token_rotary_embedding_vec<scalar_t, interleaved>(
+          query_for_token_head, s_cos, s_sin, rot_offset, embed_dim_for_rotation);
+    }
+    
+    if (key_for_token != nullptr) {
+      const int nk_pairs = num_kv_heads * embed_dim_for_rotation;
+      for (int i = thread_pair_offset; i < nk_pairs; i += pair_stride) {
+        const int head_idx = i / embed_dim_for_rotation;
+        const int rot_offset = i % embed_dim_for_rotation;
+        
+        scalar_t* key_for_token_head = key_for_token + head_idx * (int)head_stride_key;
+        apply_token_rotary_embedding_vec<scalar_t, interleaved>(
+            key_for_token_head, s_cos, s_sin, rot_offset, embed_dim_for_rotation);
+      }
+    }
+  } else {
+    // Fallback to scalar implementation
+    const int pair_stride = blockDim.x * blocks_per_token;
+    const int thread_pair_offset = local_block_idx * blockDim.x + threadIdx.x;
+    
+    const int nq_pairs = num_heads * embed_dim_for_rotation;
+    for (int i = thread_pair_offset; i < nq_pairs; i += pair_stride) {
       const int head_idx = i / embed_dim_for_rotation;
       const int rot_offset = i % embed_dim_for_rotation;
 
-      scalar_t* key_for_token_head = key_for_token + head_idx * (int)head_stride_key;
-      // Pass Shared Memory pointers
+      scalar_t* query_for_token_head = query_for_token + head_idx * (int)head_stride_query;
       apply_token_rotary_embedding<scalar_t, interleaved>(
-          key_for_token_head, s_cos, s_sin, rot_offset, embed_dim_for_rotation);
+          query_for_token_head, s_cos, s_sin, rot_offset, embed_dim_for_rotation);
+    }
+
+    if (key_for_token != nullptr) {
+      const int nk_pairs = num_kv_heads * embed_dim_for_rotation;
+      for (int i = thread_pair_offset; i < nk_pairs; i += pair_stride) {
+        const int head_idx = i / embed_dim_for_rotation;
+        const int rot_offset = i % embed_dim_for_rotation;
+
+        scalar_t* key_for_token_head = key_for_token + head_idx * (int)head_stride_key;
+        apply_token_rotary_embedding<scalar_t, interleaved>(
+            key_for_token_head, s_cos, s_sin, rot_offset, embed_dim_for_rotation);
+      }
     }
   }
 }
 
 // 1D grid kernel: each block handles one token
-template <typename scalar_t, bool interleaved>
+template <typename scalar_t, bool interleaved, bool vectorized>
 __global__ void rotary_embedding_kernel_1d(
     const scalar_t* __restrict__ cos_data,  // [num_tokens, rot_dim_arg]
     const scalar_t* __restrict__ sin_data,  // [num_tokens, rot_dim_arg]
@@ -205,11 +316,7 @@ __global__ void rotary_embedding_kernel_1d(
   const int scalar_size = sizeof(scalar_t);
   const int vec_size = kVecBytes / scalar_size;
 
-  bool is_aligned = (rot_dim_arg % vec_size == 0) &&
-                    (reinterpret_cast<uintptr_t>(current_token_cos_ptr) % kVecBytes == 0) &&
-                    (reinterpret_cast<uintptr_t>(current_token_sin_ptr) % kVecBytes == 0);
-
-  if (is_aligned) {
+  if constexpr (vectorized) {
     using vec_t = float4;
     const vec_t* cos_vec_ptr = reinterpret_cast<const vec_t*>(current_token_cos_ptr);
     const vec_t* sin_vec_ptr = reinterpret_cast<const vec_t*>(current_token_sin_ptr);
@@ -231,25 +338,54 @@ __global__ void rotary_embedding_kernel_1d(
   scalar_t* query_for_token = query_total + token_idx * (int)query_token_stride;
   scalar_t* key_for_token = (key_total != nullptr) ? (key_total + token_idx * (int)key_token_stride) : nullptr;
 
-  const int nq_pairs = num_heads * embed_dim_for_rotation;
-  for (int i = threadIdx.x; i < nq_pairs; i += blockDim.x) {
-    const int head_idx = i / embed_dim_for_rotation;
-    const int rot_offset = i % embed_dim_for_rotation;
-    scalar_t* query_for_token_head = query_for_token + head_idx * (int)head_stride_query;
-
-    apply_token_rotary_embedding<scalar_t, interleaved>(
-        query_for_token_head, s_cos, s_sin, rot_offset, embed_dim_for_rotation);
-  }
-
-  if (key_for_token != nullptr) {
-    const int nk_pairs = num_kv_heads * embed_dim_for_rotation;
-    for (int i = threadIdx.x; i < nk_pairs; i += blockDim.x) {
+  if constexpr (vectorized) {
+    using vec_t = float4;
+    constexpr int kElePerVec = kVecBytes / sizeof(scalar_t);
+    constexpr int pairs_per_step = interleaved ? (kElePerVec / 2) : kElePerVec;
+    
+    const int nq_pairs = num_heads * embed_dim_for_rotation;
+    for (int i = threadIdx.x * pairs_per_step; i < nq_pairs; i += blockDim.x * pairs_per_step) {
       const int head_idx = i / embed_dim_for_rotation;
       const int rot_offset = i % embed_dim_for_rotation;
-
-      scalar_t* key_for_token_head = key_for_token + head_idx * (int)head_stride_key;
+      scalar_t* query_for_token_head = query_for_token + head_idx * (int)head_stride_query;
+      
+      apply_token_rotary_embedding_vec<scalar_t, interleaved>(
+          query_for_token_head, s_cos, s_sin, rot_offset, embed_dim_for_rotation);
+    }
+    
+    if (key_for_token != nullptr) {
+      const int nk_pairs = num_kv_heads * embed_dim_for_rotation;
+      for (int i = threadIdx.x * pairs_per_step; i < nk_pairs; i += blockDim.x * pairs_per_step) {
+        const int head_idx = i / embed_dim_for_rotation;
+        const int rot_offset = i % embed_dim_for_rotation;
+        
+        scalar_t* key_for_token_head = key_for_token + head_idx * (int)head_stride_key;
+        apply_token_rotary_embedding_vec<scalar_t, interleaved>(
+            key_for_token_head, s_cos, s_sin, rot_offset, embed_dim_for_rotation);
+      }
+    }
+  } else {
+    // Fallback scalar
+    const int nq_pairs = num_heads * embed_dim_for_rotation;
+    for (int i = threadIdx.x; i < nq_pairs; i += blockDim.x) {
+      const int head_idx = i / embed_dim_for_rotation;
+      const int rot_offset = i % embed_dim_for_rotation;
+      scalar_t* query_for_token_head = query_for_token + head_idx * (int)head_stride_query;
+       
       apply_token_rotary_embedding<scalar_t, interleaved>(
-          key_for_token_head, s_cos, s_sin, rot_offset, embed_dim_for_rotation);
+          query_for_token_head, s_cos, s_sin, rot_offset, embed_dim_for_rotation);
+    }
+
+    if (key_for_token != nullptr) {
+      const int nk_pairs = num_kv_heads * embed_dim_for_rotation;
+      for (int i = threadIdx.x; i < nk_pairs; i += blockDim.x) {
+        const int head_idx = i / embed_dim_for_rotation;
+        const int rot_offset = i % embed_dim_for_rotation;
+        
+        scalar_t* key_for_token_head = key_for_token + head_idx * (int)head_stride_key;
+        apply_token_rotary_embedding<scalar_t, interleaved>(
+            key_for_token_head, s_cos, s_sin, rot_offset, embed_dim_for_rotation);
+      }
     }
   }
 }
@@ -349,14 +485,6 @@ void rotary_embedding_cos_sin(
   const int max_pairs_to_rotate_per_token =
       std::max(num_heads * embed_dim_for_rotation, num_kv_heads * embed_dim_for_rotation);
 
-  const int threads_per_block = std::min<int>(256, std::max(128, embed_dim_for_rotation));
-  const int blocks_per_token = (max_pairs_to_rotate_per_token + threads_per_block - 1) / threads_per_block;
-
-  const bool use_grid_2d = (num_tokens <= 4) && (blocks_per_token > 1);
-  dim3 block(threads_per_block);
-  dim3 grid_2d((int)num_tokens, std::max(1, blocks_per_token));
-  dim3 grid_1d((int)num_tokens);
-
   const at::cuda::OptionalCUDAGuard device_guard(device_of(query));
   const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
 
@@ -369,76 +497,160 @@ void rotary_embedding_cos_sin(
             typename std::conditional<std::is_same<torch_scalar_t, at::BFloat16>::value, nv_bfloat16, torch_scalar_t>::
                 type>::type;
 
+        // Constants for vectorization
+        constexpr int kVecBytes = 16;
+        constexpr int kElePerVec = kVecBytes / sizeof(torch_scalar_t);
+        
+        // Determine how many pairs one thread handles in one vector step
+        // Interleaved: 1 vector load (16B) contains 'kElePerVec' elements -> 'kElePerVec / 2' pairs.
+        // Non-interleaved: 2 vector loads (32B) contains 'kElePerVec' pairs (X vec + Y vec).
+        const int pairs_per_step = interleaved ? (kElePerVec / 2) : kElePerVec;
+
+        // Check if we can guarantee vectorization for ALL tokens
+        // We need Base Pointers, Strides, and Dimension to be aligned.
+        bool can_vectorize_all = true;
+        
+        // 1. Check Dimensions
+        if (embed_dim_for_rotation % pairs_per_step != 0) can_vectorize_all = false;
+        
+        // 2. Check Base Pointers
+        if (reinterpret_cast<uintptr_t>(query.data_ptr()) % kVecBytes != 0) can_vectorize_all = false;
+        if (reinterpret_cast<uintptr_t>(cos.data_ptr()) % kVecBytes != 0) can_vectorize_all = false;
+        if (reinterpret_cast<uintptr_t>(sin.data_ptr()) % kVecBytes != 0) can_vectorize_all = false;
+        if (key.has_value()) {
+          if (reinterpret_cast<uintptr_t>(key->data_ptr()) % kVecBytes != 0) can_vectorize_all = false;
+        }
+
+        // 3. Check Strides
+        // We need the stride between tokens to be a multiple of vector size 
+        // to ensure that if token 0 is aligned, token 1 is also aligned.
+        if (query_token_stride * sizeof(torch_scalar_t) % kVecBytes != 0) can_vectorize_all = false;
+        if (head_stride_query * sizeof(torch_scalar_t) % kVecBytes != 0) can_vectorize_all = false;
+        
+        if (key.has_value()) {
+           if (key_token_stride * sizeof(torch_scalar_t) % kVecBytes != 0) can_vectorize_all = false;
+           if (head_stride_key * sizeof(torch_scalar_t) % kVecBytes != 0) can_vectorize_all = false;
+        }
+
+        // Determine launch configuration
+        // If we can vectorize all, each thread handles 'pairs_per_step' pairs.
+        // Otherwise, fallback to conservative estimate (1 pair per thread) to ensure enough blocks.
+        const int launch_pairs_per_thread = can_vectorize_all ? pairs_per_step : 1;
+
+        const int total_threads_needed = (max_pairs_to_rotate_per_token + launch_pairs_per_thread - 1) / launch_pairs_per_thread;
+        
+        // Case 1: 2D Grid (Split one token across multiple blocks)
+        // Keep block size moderate (128-256) aligned with head size
+        const int threads_per_block_2d = std::min<int>(256, std::max(128, embed_dim_for_rotation));
+        const int blocks_per_token_2d = (total_threads_needed + threads_per_block_2d - 1) / threads_per_block_2d;
+        
+        // Decide grid strategy
+        const bool use_grid_2d = (num_tokens <= 4) && (blocks_per_token_2d > 1);
+
+        // Case 2: 1D Grid (One block per token)
+        // Maximize threads per block to cover all heads in one block if possible
+        // Cap at 512 threads to balance occupancy and register usage
+        const int threads_per_block_1d = std::min<int>(512, std::max(128, (total_threads_needed + 31) / 32 * 32));
+
+        // Final launch config
+        const int threads_per_block = use_grid_2d ? threads_per_block_2d : threads_per_block_1d;
+        const int blocks_per_token = use_grid_2d ? blocks_per_token_2d : 1;
+
+        dim3 block(threads_per_block);
+        dim3 grid_2d((int)num_tokens, std::max(1, blocks_per_token));
+        dim3 grid_1d((int)num_tokens);
+
         // We need 2 arrays (cos, sin) of size 'rot_dim_from_cache', each element is 'sizeof(torch_scalar_t)'
         size_t smem_size = rot_dim_from_cache * sizeof(torch_scalar_t) * 2;
 
-        if (interleaved) {
-          if (use_grid_2d) {
-            rotary_embedding_kernel_2d<cuda_scalar_t, true><<<grid_2d, block, smem_size, stream>>>(
-                reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
-                reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
-                reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
-                key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>()) : nullptr,
-                rot_dim_from_cache,
-                embed_dim_for_rotation,
-                query_token_stride,
-                key_token_stride,
-                head_stride_query,
-                head_stride_key,
-                num_heads,
-                num_kv_heads,
-                (int)head_size,
-                blocks_per_token);
-          } else {
-            rotary_embedding_kernel_1d<cuda_scalar_t, true><<<grid_1d, block, smem_size, stream>>>(
-                reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
-                reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
-                reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
-                key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>()) : nullptr,
-                rot_dim_from_cache,
-                embed_dim_for_rotation,
-                query_token_stride,
-                key_token_stride,
-                head_stride_query,
-                head_stride_key,
-                num_heads,
-                num_kv_heads,
-                (int)head_size);
+        auto launch_kernel = [&](bool vectorized) {
+          if (interleaved) {
+            if (use_grid_2d) {
+              if (vectorized) {
+                rotary_embedding_kernel_2d<cuda_scalar_t, true, true><<<grid_2d, block, smem_size, stream>>>(
+                  reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
+                  reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
+                  reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
+                  key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>()) : nullptr,
+                  rot_dim_from_cache, embed_dim_for_rotation, query_token_stride, key_token_stride,
+                  head_stride_query, head_stride_key, num_heads, num_kv_heads, (int)head_size, blocks_per_token);
+              } else {
+                rotary_embedding_kernel_2d<cuda_scalar_t, true, false><<<grid_2d, block, smem_size, stream>>>(
+                  reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
+                  reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
+                  reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
+                  key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>()) : nullptr,
+                  rot_dim_from_cache, embed_dim_for_rotation, query_token_stride, key_token_stride,
+                  head_stride_query, head_stride_key, num_heads, num_kv_heads, (int)head_size, blocks_per_token);
+              }
+            } else {
+              if (vectorized) {
+                rotary_embedding_kernel_1d<cuda_scalar_t, true, true><<<grid_1d, block, smem_size, stream>>>(
+                  reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
+                  reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
+                  reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
+                  key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>()) : nullptr,
+                  rot_dim_from_cache, embed_dim_for_rotation, query_token_stride, key_token_stride,
+                  head_stride_query, head_stride_key, num_heads, num_kv_heads, (int)head_size);
+              } else {
+                rotary_embedding_kernel_1d<cuda_scalar_t, true, false><<<grid_1d, block, smem_size, stream>>>(
+                  reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
+                  reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
+                  reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
+                  key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>()) : nullptr,
+                  rot_dim_from_cache, embed_dim_for_rotation, query_token_stride, key_token_stride,
+                  head_stride_query, head_stride_key, num_heads, num_kv_heads, (int)head_size);
+              }
+            }
+          } else { // non-interleaved
+            if (use_grid_2d) {
+              if (vectorized) {
+                rotary_embedding_kernel_2d<cuda_scalar_t, false, true><<<grid_2d, block, smem_size, stream>>>(
+                  reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
+                  reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
+                  reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
+                  key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>()) : nullptr,
+                  rot_dim_from_cache, embed_dim_for_rotation, query_token_stride, key_token_stride,
+                  head_stride_query, head_stride_key, num_heads, num_kv_heads, (int)head_size, blocks_per_token);
+              } else {
+                rotary_embedding_kernel_2d<cuda_scalar_t, false, false><<<grid_2d, block, smem_size, stream>>>(
+                  reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
+                  reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
+                  reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
+                  key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>()) : nullptr,
+                  rot_dim_from_cache, embed_dim_for_rotation, query_token_stride, key_token_stride,
+                  head_stride_query, head_stride_key, num_heads, num_kv_heads, (int)head_size, blocks_per_token);
+              }
+            } else {
+              if (vectorized) {
+                rotary_embedding_kernel_1d<cuda_scalar_t, false, true><<<grid_1d, block, smem_size, stream>>>(
+                  reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
+                  reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
+                  reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
+                  key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>()) : nullptr,
+                  rot_dim_from_cache, embed_dim_for_rotation, query_token_stride, key_token_stride,
+                  head_stride_query, head_stride_key, num_heads, num_kv_heads, (int)head_size);
+              } else {
+                rotary_embedding_kernel_1d<cuda_scalar_t, false, false><<<grid_1d, block, smem_size, stream>>>(
+                  reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
+                  reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
+                  reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
+                  key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>()) : nullptr,
+                  rot_dim_from_cache, embed_dim_for_rotation, query_token_stride, key_token_stride,
+                  head_stride_query, head_stride_key, num_heads, num_kv_heads, (int)head_size);
+              }
+            }
           }
+          
+        };
+
+        // Close the log file after all launches
+        if (can_vectorize_all) {
+          launch_kernel(true);
         } else {
-          if (use_grid_2d) {
-            rotary_embedding_kernel_2d<cuda_scalar_t, false><<<grid_2d, block, smem_size, stream>>>(
-                reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
-                reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
-                reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
-                key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>()) : nullptr,
-                rot_dim_from_cache,
-                embed_dim_for_rotation,
-                query_token_stride,
-                key_token_stride,
-                head_stride_query,
-                head_stride_key,
-                num_heads,
-                num_kv_heads,
-                (int)head_size,
-                blocks_per_token);
-          } else {
-            rotary_embedding_kernel_1d<cuda_scalar_t, false><<<grid_1d, block, smem_size, stream>>>(
-                reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
-                reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
-                reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
-                key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>()) : nullptr,
-                rot_dim_from_cache,
-                embed_dim_for_rotation,
-                query_token_stride,
-                key_token_stride,
-                head_stride_query,
-                head_stride_key,
-                num_heads,
-                num_kv_heads,
-                (int)head_size);
-          }
+          launch_kernel(false);
         }
+
       });
 
   C10_CUDA_KERNEL_LAUNCH_CHECK();

From 4a31c1befb33929a914d3c802f25942de3b6f5ec Mon Sep 17 00:00:00 2001
From: Ther-LF <2639852836@qq.com>
Date: Tue, 16 Dec 2025 21:42:53 +0800
Subject: [PATCH 09/37] fix lint

---
 .../sgl_diffusion/rope/rotary_embedding.cu    | 244 +++++++++++-------
 1 file changed, 151 insertions(+), 93 deletions(-)

diff --git a/sgl-kernel/csrc/sgl_diffusion/rope/rotary_embedding.cu b/sgl-kernel/csrc/sgl_diffusion/rope/rotary_embedding.cu
index 057e6cba0651..f0185bcf3992 100644
--- a/sgl-kernel/csrc/sgl_diffusion/rope/rotary_embedding.cu
+++ b/sgl-kernel/csrc/sgl_diffusion/rope/rotary_embedding.cu
@@ -73,12 +73,11 @@ inline __device__ void apply_token_rotary_embedding_vec(
     const scalar_t* __restrict__ sin_ptr,  // [rot_dim]
     int rot_offset,
     int embed_dim) {
-  
   using vec_t = float4;
   constexpr int kVecBytes = sizeof(vec_t);
   constexpr int kScalarBytes = sizeof(scalar_t);
   constexpr int kElePerVec = kVecBytes / kScalarBytes;
-  
+
   // Union for type punning to avoid strict aliasing issues with reinterpret_cast
   union VecUnion {
     vec_t vec;
@@ -90,59 +89,59 @@ inline __device__ void apply_token_rotary_embedding_vec(
     // A single vector load contains 'kElePerVec' elements, which is 'kElePerVec / 2' pairs.
     // rot_offset is the index of the PAIR.
     // Address in arr is rot_offset * 2.
-    
+
     VecUnion data;
     data.vec = *reinterpret_cast<const vec_t*>(arr + rot_offset * 2);
-    
-    #pragma unroll
+
+#pragma unroll
     for (int i = 0; i < kElePerVec; i += 2) {
       // data.elems[i] is x, data.elems[i+1] is y
       // They correspond to pair index: rot_offset + (i / 2)
       int curr_rot_offset = rot_offset + i / 2;
-      
+
       float cos_val = static_cast<float>(cos_ptr[curr_rot_offset]);
       float sin_val = static_cast<float>(sin_ptr[curr_rot_offset]);
-      
+
       float x = static_cast<float>(data.elems[i]);
-      float y = static_cast<float>(data.elems[i+1]);
-      
+      float y = static_cast<float>(data.elems[i + 1]);
+
       data.elems[i] = static_cast<scalar_t>(x * cos_val - y * sin_val);
-      data.elems[i+1] = static_cast<scalar_t>(y * cos_val + x * sin_val);
+      data.elems[i + 1] = static_cast<scalar_t>(y * cos_val + x * sin_val);
     }
-    
+
     *reinterpret_cast<vec_t*>(arr + rot_offset * 2) = data.vec;
 
   } else {
     // Non-interleaved: X and Y are separated by embed_dim.
     // We process 'kElePerVec' PAIRS at once.
     // Load X vector and Y vector.
-    
+
     VecUnion data_x, data_y;
     data_x.vec = *reinterpret_cast<const vec_t*>(arr + rot_offset);
     data_y.vec = *reinterpret_cast<const vec_t*>(arr + rot_offset + embed_dim);
-    
-    #pragma unroll
+
+#pragma unroll
     for (int i = 0; i < kElePerVec; ++i) {
       int curr_rot_offset = rot_offset + i;
-      
+
       // In non-interleaved, we might need different cos/sin for X and Y depending on implementation,
       // but standard RoPE uses the same angle for the pair.
       // Based on original scalar code:
       // cos_val_x = cos_ptr[rot_offset]
       // cos_val_y = cos_ptr[rot_offset + embed_dim]
-      
+
       float cos_val_x = static_cast<float>(cos_ptr[curr_rot_offset]);
       float sin_val_x = static_cast<float>(sin_ptr[curr_rot_offset]);
       float cos_val_y = static_cast<float>(cos_ptr[curr_rot_offset + embed_dim]);
       float sin_val_y = static_cast<float>(sin_ptr[curr_rot_offset + embed_dim]);
-      
+
       float x = static_cast<float>(data_x.elems[i]);
       float y = static_cast<float>(data_y.elems[i]);
-      
+
       data_x.elems[i] = static_cast<scalar_t>(x * cos_val_x - y * sin_val_x);
       data_y.elems[i] = static_cast<scalar_t>(y * cos_val_y + x * sin_val_y);
     }
-    
+
     *reinterpret_cast<vec_t*>(arr + rot_offset) = data_x.vec;
     *reinterpret_cast<vec_t*>(arr + rot_offset + embed_dim) = data_y.vec;
   }
@@ -227,31 +226,31 @@ __global__ void rotary_embedding_kernel_2d(
   scalar_t* key_for_token = (key_total != nullptr) ? (key_total + token_idx * (int)key_token_stride) : nullptr;
 
   const int local_block_idx = blockIdx.y;
-  
+
   if constexpr (vectorized) {
     using vec_t = float4;
     constexpr int kElePerVec = kVecBytes / sizeof(scalar_t);
     constexpr int pairs_per_step = interleaved ? (kElePerVec / 2) : kElePerVec;
-    
+
     const int pair_stride = blockDim.x * blocks_per_token * pairs_per_step;
     const int thread_pair_offset = (local_block_idx * blockDim.x + threadIdx.x) * pairs_per_step;
     const int nq_pairs = num_heads * embed_dim_for_rotation;
-    
+
     for (int i = thread_pair_offset; i < nq_pairs; i += pair_stride) {
       const int head_idx = i / embed_dim_for_rotation;
       const int rot_offset = i % embed_dim_for_rotation;
-      
+
       scalar_t* query_for_token_head = query_for_token + head_idx * (int)head_stride_query;
       apply_token_rotary_embedding_vec<scalar_t, interleaved>(
           query_for_token_head, s_cos, s_sin, rot_offset, embed_dim_for_rotation);
     }
-    
+
     if (key_for_token != nullptr) {
       const int nk_pairs = num_kv_heads * embed_dim_for_rotation;
       for (int i = thread_pair_offset; i < nk_pairs; i += pair_stride) {
         const int head_idx = i / embed_dim_for_rotation;
         const int rot_offset = i % embed_dim_for_rotation;
-        
+
         scalar_t* key_for_token_head = key_for_token + head_idx * (int)head_stride_key;
         apply_token_rotary_embedding_vec<scalar_t, interleaved>(
             key_for_token_head, s_cos, s_sin, rot_offset, embed_dim_for_rotation);
@@ -261,7 +260,7 @@ __global__ void rotary_embedding_kernel_2d(
     // Fallback to scalar implementation
     const int pair_stride = blockDim.x * blocks_per_token;
     const int thread_pair_offset = local_block_idx * blockDim.x + threadIdx.x;
-    
+
     const int nq_pairs = num_heads * embed_dim_for_rotation;
     for (int i = thread_pair_offset; i < nq_pairs; i += pair_stride) {
       const int head_idx = i / embed_dim_for_rotation;
@@ -342,23 +341,23 @@ __global__ void rotary_embedding_kernel_1d(
     using vec_t = float4;
     constexpr int kElePerVec = kVecBytes / sizeof(scalar_t);
     constexpr int pairs_per_step = interleaved ? (kElePerVec / 2) : kElePerVec;
-    
+
     const int nq_pairs = num_heads * embed_dim_for_rotation;
     for (int i = threadIdx.x * pairs_per_step; i < nq_pairs; i += blockDim.x * pairs_per_step) {
       const int head_idx = i / embed_dim_for_rotation;
       const int rot_offset = i % embed_dim_for_rotation;
       scalar_t* query_for_token_head = query_for_token + head_idx * (int)head_stride_query;
-      
+
       apply_token_rotary_embedding_vec<scalar_t, interleaved>(
           query_for_token_head, s_cos, s_sin, rot_offset, embed_dim_for_rotation);
     }
-    
+
     if (key_for_token != nullptr) {
       const int nk_pairs = num_kv_heads * embed_dim_for_rotation;
       for (int i = threadIdx.x * pairs_per_step; i < nk_pairs; i += blockDim.x * pairs_per_step) {
         const int head_idx = i / embed_dim_for_rotation;
         const int rot_offset = i % embed_dim_for_rotation;
-        
+
         scalar_t* key_for_token_head = key_for_token + head_idx * (int)head_stride_key;
         apply_token_rotary_embedding_vec<scalar_t, interleaved>(
             key_for_token_head, s_cos, s_sin, rot_offset, embed_dim_for_rotation);
@@ -371,7 +370,7 @@ __global__ void rotary_embedding_kernel_1d(
       const int head_idx = i / embed_dim_for_rotation;
       const int rot_offset = i % embed_dim_for_rotation;
       scalar_t* query_for_token_head = query_for_token + head_idx * (int)head_stride_query;
-       
+
       apply_token_rotary_embedding<scalar_t, interleaved>(
           query_for_token_head, s_cos, s_sin, rot_offset, embed_dim_for_rotation);
     }
@@ -381,7 +380,7 @@ __global__ void rotary_embedding_kernel_1d(
       for (int i = threadIdx.x; i < nk_pairs; i += blockDim.x) {
         const int head_idx = i / embed_dim_for_rotation;
         const int rot_offset = i % embed_dim_for_rotation;
-        
+
         scalar_t* key_for_token_head = key_for_token + head_idx * (int)head_stride_key;
         apply_token_rotary_embedding<scalar_t, interleaved>(
             key_for_token_head, s_cos, s_sin, rot_offset, embed_dim_for_rotation);
@@ -500,7 +499,7 @@ void rotary_embedding_cos_sin(
         // Constants for vectorization
         constexpr int kVecBytes = 16;
         constexpr int kElePerVec = kVecBytes / sizeof(torch_scalar_t);
-        
+
         // Determine how many pairs one thread handles in one vector step
         // Interleaved: 1 vector load (16B) contains 'kElePerVec' elements -> 'kElePerVec / 2' pairs.
         // Non-interleaved: 2 vector loads (32B) contains 'kElePerVec' pairs (X vec + Y vec).
@@ -509,10 +508,10 @@ void rotary_embedding_cos_sin(
         // Check if we can guarantee vectorization for ALL tokens
         // We need Base Pointers, Strides, and Dimension to be aligned.
         bool can_vectorize_all = true;
-        
+
         // 1. Check Dimensions
         if (embed_dim_for_rotation % pairs_per_step != 0) can_vectorize_all = false;
-        
+
         // 2. Check Base Pointers
         if (reinterpret_cast<uintptr_t>(query.data_ptr()) % kVecBytes != 0) can_vectorize_all = false;
         if (reinterpret_cast<uintptr_t>(cos.data_ptr()) % kVecBytes != 0) can_vectorize_all = false;
@@ -522,14 +521,14 @@ void rotary_embedding_cos_sin(
         }
 
         // 3. Check Strides
-        // We need the stride between tokens to be a multiple of vector size 
+        // We need the stride between tokens to be a multiple of vector size
         // to ensure that if token 0 is aligned, token 1 is also aligned.
         if (query_token_stride * sizeof(torch_scalar_t) % kVecBytes != 0) can_vectorize_all = false;
         if (head_stride_query * sizeof(torch_scalar_t) % kVecBytes != 0) can_vectorize_all = false;
-        
+
         if (key.has_value()) {
-           if (key_token_stride * sizeof(torch_scalar_t) % kVecBytes != 0) can_vectorize_all = false;
-           if (head_stride_key * sizeof(torch_scalar_t) % kVecBytes != 0) can_vectorize_all = false;
+          if (key_token_stride * sizeof(torch_scalar_t) % kVecBytes != 0) can_vectorize_all = false;
+          if (head_stride_key * sizeof(torch_scalar_t) % kVecBytes != 0) can_vectorize_all = false;
         }
 
         // Determine launch configuration
@@ -537,13 +536,14 @@ void rotary_embedding_cos_sin(
         // Otherwise, fallback to conservative estimate (1 pair per thread) to ensure enough blocks.
         const int launch_pairs_per_thread = can_vectorize_all ? pairs_per_step : 1;
 
-        const int total_threads_needed = (max_pairs_to_rotate_per_token + launch_pairs_per_thread - 1) / launch_pairs_per_thread;
-        
+        const int total_threads_needed =
+            (max_pairs_to_rotate_per_token + launch_pairs_per_thread - 1) / launch_pairs_per_thread;
+
         // Case 1: 2D Grid (Split one token across multiple blocks)
         // Keep block size moderate (128-256) aligned with head size
         const int threads_per_block_2d = std::min<int>(256, std::max(128, embed_dim_for_rotation));
         const int blocks_per_token_2d = (total_threads_needed + threads_per_block_2d - 1) / threads_per_block_2d;
-        
+
         // Decide grid strategy
         const bool use_grid_2d = (num_tokens <= 4) && (blocks_per_token_2d > 1);
 
@@ -568,80 +568,139 @@ void rotary_embedding_cos_sin(
             if (use_grid_2d) {
               if (vectorized) {
                 rotary_embedding_kernel_2d<cuda_scalar_t, true, true><<<grid_2d, block, smem_size, stream>>>(
-                  reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
-                  reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
-                  reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
-                  key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>()) : nullptr,
-                  rot_dim_from_cache, embed_dim_for_rotation, query_token_stride, key_token_stride,
-                  head_stride_query, head_stride_key, num_heads, num_kv_heads, (int)head_size, blocks_per_token);
+                    reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
+                    reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
+                    reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
+                    key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>()) : nullptr,
+                    rot_dim_from_cache,
+                    embed_dim_for_rotation,
+                    query_token_stride,
+                    key_token_stride,
+                    head_stride_query,
+                    head_stride_key,
+                    num_heads,
+                    num_kv_heads,
+                    (int)head_size,
+                    blocks_per_token);
               } else {
                 rotary_embedding_kernel_2d<cuda_scalar_t, true, false><<<grid_2d, block, smem_size, stream>>>(
-                  reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
-                  reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
-                  reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
-                  key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>()) : nullptr,
-                  rot_dim_from_cache, embed_dim_for_rotation, query_token_stride, key_token_stride,
-                  head_stride_query, head_stride_key, num_heads, num_kv_heads, (int)head_size, blocks_per_token);
+                    reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
+                    reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
+                    reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
+                    key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>()) : nullptr,
+                    rot_dim_from_cache,
+                    embed_dim_for_rotation,
+                    query_token_stride,
+                    key_token_stride,
+                    head_stride_query,
+                    head_stride_key,
+                    num_heads,
+                    num_kv_heads,
+                    (int)head_size,
+                    blocks_per_token);
               }
             } else {
               if (vectorized) {
                 rotary_embedding_kernel_1d<cuda_scalar_t, true, true><<<grid_1d, block, smem_size, stream>>>(
-                  reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
-                  reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
-                  reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
-                  key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>()) : nullptr,
-                  rot_dim_from_cache, embed_dim_for_rotation, query_token_stride, key_token_stride,
-                  head_stride_query, head_stride_key, num_heads, num_kv_heads, (int)head_size);
+                    reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
+                    reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
+                    reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
+                    key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>()) : nullptr,
+                    rot_dim_from_cache,
+                    embed_dim_for_rotation,
+                    query_token_stride,
+                    key_token_stride,
+                    head_stride_query,
+                    head_stride_key,
+                    num_heads,
+                    num_kv_heads,
+                    (int)head_size);
               } else {
                 rotary_embedding_kernel_1d<cuda_scalar_t, true, false><<<grid_1d, block, smem_size, stream>>>(
-                  reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
-                  reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
-                  reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
-                  key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>()) : nullptr,
-                  rot_dim_from_cache, embed_dim_for_rotation, query_token_stride, key_token_stride,
-                  head_stride_query, head_stride_key, num_heads, num_kv_heads, (int)head_size);
+                    reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
+                    reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
+                    reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
+                    key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>()) : nullptr,
+                    rot_dim_from_cache,
+                    embed_dim_for_rotation,
+                    query_token_stride,
+                    key_token_stride,
+                    head_stride_query,
+                    head_stride_key,
+                    num_heads,
+                    num_kv_heads,
+                    (int)head_size);
               }
             }
-          } else { // non-interleaved
+          } else {  // non-interleaved
             if (use_grid_2d) {
               if (vectorized) {
                 rotary_embedding_kernel_2d<cuda_scalar_t, false, true><<<grid_2d, block, smem_size, stream>>>(
-                  reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
-                  reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
-                  reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
-                  key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>()) : nullptr,
-                  rot_dim_from_cache, embed_dim_for_rotation, query_token_stride, key_token_stride,
-                  head_stride_query, head_stride_key, num_heads, num_kv_heads, (int)head_size, blocks_per_token);
+                    reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
+                    reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
+                    reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
+                    key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>()) : nullptr,
+                    rot_dim_from_cache,
+                    embed_dim_for_rotation,
+                    query_token_stride,
+                    key_token_stride,
+                    head_stride_query,
+                    head_stride_key,
+                    num_heads,
+                    num_kv_heads,
+                    (int)head_size,
+                    blocks_per_token);
               } else {
                 rotary_embedding_kernel_2d<cuda_scalar_t, false, false><<<grid_2d, block, smem_size, stream>>>(
-                  reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
-                  reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
-                  reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
-                  key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>()) : nullptr,
-                  rot_dim_from_cache, embed_dim_for_rotation, query_token_stride, key_token_stride,
-                  head_stride_query, head_stride_key, num_heads, num_kv_heads, (int)head_size, blocks_per_token);
+                    reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
+                    reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
+                    reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
+                    key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>()) : nullptr,
+                    rot_dim_from_cache,
+                    embed_dim_for_rotation,
+                    query_token_stride,
+                    key_token_stride,
+                    head_stride_query,
+                    head_stride_key,
+                    num_heads,
+                    num_kv_heads,
+                    (int)head_size,
+                    blocks_per_token);
               }
             } else {
               if (vectorized) {
                 rotary_embedding_kernel_1d<cuda_scalar_t, false, true><<<grid_1d, block, smem_size, stream>>>(
-                  reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
-                  reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
-                  reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
-                  key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>()) : nullptr,
-                  rot_dim_from_cache, embed_dim_for_rotation, query_token_stride, key_token_stride,
-                  head_stride_query, head_stride_key, num_heads, num_kv_heads, (int)head_size);
+                    reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
+                    reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
+                    reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
+                    key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>()) : nullptr,
+                    rot_dim_from_cache,
+                    embed_dim_for_rotation,
+                    query_token_stride,
+                    key_token_stride,
+                    head_stride_query,
+                    head_stride_key,
+                    num_heads,
+                    num_kv_heads,
+                    (int)head_size);
               } else {
                 rotary_embedding_kernel_1d<cuda_scalar_t, false, false><<<grid_1d, block, smem_size, stream>>>(
-                  reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
-                  reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
-                  reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
-                  key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>()) : nullptr,
-                  rot_dim_from_cache, embed_dim_for_rotation, query_token_stride, key_token_stride,
-                  head_stride_query, head_stride_key, num_heads, num_kv_heads, (int)head_size);
+                    reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
+                    reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
+                    reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
+                    key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>()) : nullptr,
+                    rot_dim_from_cache,
+                    embed_dim_for_rotation,
+                    query_token_stride,
+                    key_token_stride,
+                    head_stride_query,
+                    head_stride_key,
+                    num_heads,
+                    num_kv_heads,
+                    (int)head_size);
               }
             }
           }
-          
         };
 
         // Close the log file after all launches
@@ -650,7 +709,6 @@ void rotary_embedding_cos_sin(
         } else {
           launch_kernel(false);
         }
-
       });
 
   C10_CUDA_KERNEL_LAUNCH_CHECK();

From 0735ba3e6d576d5cee14b4881bbee0aeb8c455f6 Mon Sep 17 00:00:00 2001
From: Ther-LF <2639852836@qq.com>
Date: Tue, 16 Dec 2025 23:04:31 +0800
Subject: [PATCH 10/37] perf:replace int div/mod by bit-wise op

---
 .../sgl_diffusion/rope/rotary_embedding.cu    | 319 ++++++++++--------
 1 file changed, 173 insertions(+), 146 deletions(-)

diff --git a/sgl-kernel/csrc/sgl_diffusion/rope/rotary_embedding.cu b/sgl-kernel/csrc/sgl_diffusion/rope/rotary_embedding.cu
index f0185bcf3992..84b37e3da852 100644
--- a/sgl-kernel/csrc/sgl_diffusion/rope/rotary_embedding.cu
+++ b/sgl-kernel/csrc/sgl_diffusion/rope/rotary_embedding.cu
@@ -167,21 +167,21 @@ inline __device__ void apply_token_rotary_embedding<float, true>(
 }
 
 // 2D grid kernel: parallel over tokens (grid.x) and pair-tiles (grid.y)
-template <typename scalar_t, bool interleaved, bool vectorized>
+template <typename scalar_t, bool interleaved, bool vectorized, int ROT_EMBED_DIM>
 __global__ void rotary_embedding_kernel_2d(
     const scalar_t* __restrict__ cos_data,  // [num_tokens, rot_dim_arg]
     const scalar_t* __restrict__ sin_data,  // [num_tokens, rot_dim_arg]
     scalar_t* __restrict__ query_total,
     scalar_t* __restrict__ key_total,
     const int rot_dim_arg,
-    const int embed_dim_for_rotation,
+    const int embed_dim_for_rotation_arg,
     const int64_t query_token_stride,
     const int64_t key_token_stride,
     const int64_t head_stride_query,
     const int64_t head_stride_key,
     const int num_heads,
     const int num_kv_heads,
-    const int head_size,
+    const int head_size_arg,
     const int blocks_per_token) {
   extern __shared__ char smem_[];
   scalar_t* s_cos = reinterpret_cast<scalar_t*>(smem_);
@@ -227,6 +227,9 @@ __global__ void rotary_embedding_kernel_2d(
 
   const int local_block_idx = blockIdx.y;
 
+  // Use compile-time constant if ROT_EMBED_DIM > 0, otherwise fallback to runtime argument
+  const int embed_dim_for_rotation = (ROT_EMBED_DIM > 0) ? ROT_EMBED_DIM : embed_dim_for_rotation_arg;
+
   if constexpr (vectorized) {
     using vec_t = float4;
     constexpr int kElePerVec = kVecBytes / sizeof(scalar_t);
@@ -286,21 +289,21 @@ __global__ void rotary_embedding_kernel_2d(
 }
 
 // 1D grid kernel: each block handles one token
-template <typename scalar_t, bool interleaved, bool vectorized>
+template <typename scalar_t, bool interleaved, bool vectorized, int ROT_EMBED_DIM>
 __global__ void rotary_embedding_kernel_1d(
     const scalar_t* __restrict__ cos_data,  // [num_tokens, rot_dim_arg]
     const scalar_t* __restrict__ sin_data,  // [num_tokens, rot_dim_arg]
     scalar_t* __restrict__ query_total,
     scalar_t* __restrict__ key_total,
     const int rot_dim_arg,
-    const int embed_dim_for_rotation,
+    const int embed_dim_for_rotation_arg,
     const int64_t query_token_stride,
     const int64_t key_token_stride,
     const int64_t head_stride_query,
     const int64_t head_stride_key,
     const int num_heads,
     const int num_kv_heads,
-    const int head_size) {
+    const int head_size_arg) {
   extern __shared__ char smem_[];
   scalar_t* s_cos = reinterpret_cast<scalar_t*>(smem_);
   scalar_t* s_sin = s_cos + rot_dim_arg;
@@ -337,6 +340,9 @@ __global__ void rotary_embedding_kernel_1d(
   scalar_t* query_for_token = query_total + token_idx * (int)query_token_stride;
   scalar_t* key_for_token = (key_total != nullptr) ? (key_total + token_idx * (int)key_token_stride) : nullptr;
 
+  // Use compile-time constant if ROT_EMBED_DIM > 0, otherwise fallback to runtime argument
+  const int embed_dim_for_rotation = (ROT_EMBED_DIM > 0) ? ROT_EMBED_DIM : embed_dim_for_rotation_arg;
+
   if constexpr (vectorized) {
     using vec_t = float4;
     constexpr int kElePerVec = kVecBytes / sizeof(scalar_t);
@@ -389,6 +395,25 @@ __global__ void rotary_embedding_kernel_1d(
   }
 }
 
+// Define dispatch macro for Rotation Dimension (Pair Count)
+// Note: This is usually HEAD_SIZE / 2 for LLaMA-like models.
+// Case 32 -> head_size 64
+// Case 64 -> head_size 128
+// Case 128 -> head_size 256
+#define DISPATCH_ROT_DIM(embed_dim, ...) \
+    switch (embed_dim) { \
+        case 16: { constexpr int ROT_EMBED_DIM = 16; __VA_ARGS__; break; } \
+        case 32: { constexpr int ROT_EMBED_DIM = 32; __VA_ARGS__; break; } \
+        case 64: { constexpr int ROT_EMBED_DIM = 64; __VA_ARGS__; break; } \
+        case 96: { constexpr int ROT_EMBED_DIM = 96; __VA_ARGS__; break; } \
+        case 128: { constexpr int ROT_EMBED_DIM = 128; __VA_ARGS__; break; } \
+        case 256: { constexpr int ROT_EMBED_DIM = 256; __VA_ARGS__; break; } \
+        case 512: { constexpr int ROT_EMBED_DIM = 512; __VA_ARGS__; break; } \
+        case 1024: { constexpr int ROT_EMBED_DIM = 1024; __VA_ARGS__; break; } \
+        default: { constexpr int ROT_EMBED_DIM = 0; __VA_ARGS__; break; } \
+    }
+
+
 void rotary_embedding_cos_sin(
     at::Tensor& cos,
     at::Tensor& sin,
@@ -563,152 +588,154 @@ void rotary_embedding_cos_sin(
         // We need 2 arrays (cos, sin) of size 'rot_dim_from_cache', each element is 'sizeof(torch_scalar_t)'
         size_t smem_size = rot_dim_from_cache * sizeof(torch_scalar_t) * 2;
 
-        auto launch_kernel = [&](bool vectorized) {
-          if (interleaved) {
-            if (use_grid_2d) {
-              if (vectorized) {
-                rotary_embedding_kernel_2d<cuda_scalar_t, true, true><<<grid_2d, block, smem_size, stream>>>(
-                    reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
-                    reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
-                    reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
-                    key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>()) : nullptr,
-                    rot_dim_from_cache,
-                    embed_dim_for_rotation,
-                    query_token_stride,
-                    key_token_stride,
-                    head_stride_query,
-                    head_stride_key,
-                    num_heads,
-                    num_kv_heads,
-                    (int)head_size,
-                    blocks_per_token);
+        DISPATCH_ROT_DIM(embed_dim_for_rotation,
+          auto launch_kernel = [&](bool vectorized) {
+            if (interleaved) {
+              if (use_grid_2d) {
+                if (vectorized) {
+                  rotary_embedding_kernel_2d<cuda_scalar_t, true, true, ROT_EMBED_DIM><<<grid_2d, block, smem_size, stream>>>(
+                      reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
+                      reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
+                      reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
+                      key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>()) : nullptr,
+                      rot_dim_from_cache,
+                      embed_dim_for_rotation,
+                      query_token_stride,
+                      key_token_stride,
+                      head_stride_query,
+                      head_stride_key,
+                      num_heads,
+                      num_kv_heads,
+                      (int)head_size,
+                      blocks_per_token);
+                } else {
+                  rotary_embedding_kernel_2d<cuda_scalar_t, true, false, ROT_EMBED_DIM><<<grid_2d, block, smem_size, stream>>>(
+                      reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
+                      reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
+                      reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
+                      key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>()) : nullptr,
+                      rot_dim_from_cache,
+                      embed_dim_for_rotation,
+                      query_token_stride,
+                      key_token_stride,
+                      head_stride_query,
+                      head_stride_key,
+                      num_heads,
+                      num_kv_heads,
+                      (int)head_size,
+                      blocks_per_token);
+                }
               } else {
-                rotary_embedding_kernel_2d<cuda_scalar_t, true, false><<<grid_2d, block, smem_size, stream>>>(
-                    reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
-                    reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
-                    reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
-                    key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>()) : nullptr,
-                    rot_dim_from_cache,
-                    embed_dim_for_rotation,
-                    query_token_stride,
-                    key_token_stride,
-                    head_stride_query,
-                    head_stride_key,
-                    num_heads,
-                    num_kv_heads,
-                    (int)head_size,
-                    blocks_per_token);
+                if (vectorized) {
+                  rotary_embedding_kernel_1d<cuda_scalar_t, true, true, ROT_EMBED_DIM><<<grid_1d, block, smem_size, stream>>>(
+                      reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
+                      reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
+                      reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
+                      key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>()) : nullptr,
+                      rot_dim_from_cache,
+                      embed_dim_for_rotation,
+                      query_token_stride,
+                      key_token_stride,
+                      head_stride_query,
+                      head_stride_key,
+                      num_heads,
+                      num_kv_heads,
+                      (int)head_size);
+                } else {
+                  rotary_embedding_kernel_1d<cuda_scalar_t, true, false, ROT_EMBED_DIM><<<grid_1d, block, smem_size, stream>>>(
+                      reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
+                      reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
+                      reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
+                      key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>()) : nullptr,
+                      rot_dim_from_cache,
+                      embed_dim_for_rotation,
+                      query_token_stride,
+                      key_token_stride,
+                      head_stride_query,
+                      head_stride_key,
+                      num_heads,
+                      num_kv_heads,
+                      (int)head_size);
+                }
               }
-            } else {
-              if (vectorized) {
-                rotary_embedding_kernel_1d<cuda_scalar_t, true, true><<<grid_1d, block, smem_size, stream>>>(
-                    reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
-                    reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
-                    reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
-                    key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>()) : nullptr,
-                    rot_dim_from_cache,
-                    embed_dim_for_rotation,
-                    query_token_stride,
-                    key_token_stride,
-                    head_stride_query,
-                    head_stride_key,
-                    num_heads,
-                    num_kv_heads,
-                    (int)head_size);
+            } else {  // non-interleaved
+              if (use_grid_2d) {
+                if (vectorized) {
+                  rotary_embedding_kernel_2d<cuda_scalar_t, false, true, ROT_EMBED_DIM><<<grid_2d, block, smem_size, stream>>>(
+                      reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
+                      reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
+                      reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
+                      key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>()) : nullptr,
+                      rot_dim_from_cache,
+                      embed_dim_for_rotation,
+                      query_token_stride,
+                      key_token_stride,
+                      head_stride_query,
+                      head_stride_key,
+                      num_heads,
+                      num_kv_heads,
+                      (int)head_size,
+                      blocks_per_token);
+                } else {
+                  rotary_embedding_kernel_2d<cuda_scalar_t, false, false, ROT_EMBED_DIM><<<grid_2d, block, smem_size, stream>>>(
+                      reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
+                      reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
+                      reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
+                      key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>()) : nullptr,
+                      rot_dim_from_cache,
+                      embed_dim_for_rotation,
+                      query_token_stride,
+                      key_token_stride,
+                      head_stride_query,
+                      head_stride_key,
+                      num_heads,
+                      num_kv_heads,
+                      (int)head_size,
+                      blocks_per_token);
+                }
               } else {
-                rotary_embedding_kernel_1d<cuda_scalar_t, true, false><<<grid_1d, block, smem_size, stream>>>(
-                    reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
-                    reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
-                    reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
-                    key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>()) : nullptr,
-                    rot_dim_from_cache,
-                    embed_dim_for_rotation,
-                    query_token_stride,
-                    key_token_stride,
-                    head_stride_query,
-                    head_stride_key,
-                    num_heads,
-                    num_kv_heads,
-                    (int)head_size);
+                if (vectorized) {
+                  rotary_embedding_kernel_1d<cuda_scalar_t, false, true, ROT_EMBED_DIM><<<grid_1d, block, smem_size, stream>>>(
+                      reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
+                      reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
+                      reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
+                      key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>()) : nullptr,
+                      rot_dim_from_cache,
+                      embed_dim_for_rotation,
+                      query_token_stride,
+                      key_token_stride,
+                      head_stride_query,
+                      head_stride_key,
+                      num_heads,
+                      num_kv_heads,
+                      (int)head_size);
+                } else {
+                  rotary_embedding_kernel_1d<cuda_scalar_t, false, false, ROT_EMBED_DIM><<<grid_1d, block, smem_size, stream>>>(
+                      reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
+                      reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
+                      reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
+                      key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>()) : nullptr,
+                      rot_dim_from_cache,
+                      embed_dim_for_rotation,
+                      query_token_stride,
+                      key_token_stride,
+                      head_stride_query,
+                      head_stride_key,
+                      num_heads,
+                      num_kv_heads,
+                      (int)head_size);
+                }
               }
             }
-          } else {  // non-interleaved
-            if (use_grid_2d) {
-              if (vectorized) {
-                rotary_embedding_kernel_2d<cuda_scalar_t, false, true><<<grid_2d, block, smem_size, stream>>>(
-                    reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
-                    reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
-                    reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
-                    key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>()) : nullptr,
-                    rot_dim_from_cache,
-                    embed_dim_for_rotation,
-                    query_token_stride,
-                    key_token_stride,
-                    head_stride_query,
-                    head_stride_key,
-                    num_heads,
-                    num_kv_heads,
-                    (int)head_size,
-                    blocks_per_token);
-              } else {
-                rotary_embedding_kernel_2d<cuda_scalar_t, false, false><<<grid_2d, block, smem_size, stream>>>(
-                    reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
-                    reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
-                    reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
-                    key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>()) : nullptr,
-                    rot_dim_from_cache,
-                    embed_dim_for_rotation,
-                    query_token_stride,
-                    key_token_stride,
-                    head_stride_query,
-                    head_stride_key,
-                    num_heads,
-                    num_kv_heads,
-                    (int)head_size,
-                    blocks_per_token);
-              }
-            } else {
-              if (vectorized) {
-                rotary_embedding_kernel_1d<cuda_scalar_t, false, true><<<grid_1d, block, smem_size, stream>>>(
-                    reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
-                    reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
-                    reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
-                    key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>()) : nullptr,
-                    rot_dim_from_cache,
-                    embed_dim_for_rotation,
-                    query_token_stride,
-                    key_token_stride,
-                    head_stride_query,
-                    head_stride_key,
-                    num_heads,
-                    num_kv_heads,
-                    (int)head_size);
-              } else {
-                rotary_embedding_kernel_1d<cuda_scalar_t, false, false><<<grid_1d, block, smem_size, stream>>>(
-                    reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
-                    reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
-                    reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
-                    key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>()) : nullptr,
-                    rot_dim_from_cache,
-                    embed_dim_for_rotation,
-                    query_token_stride,
-                    key_token_stride,
-                    head_stride_query,
-                    head_stride_key,
-                    num_heads,
-                    num_kv_heads,
-                    (int)head_size);
-              }
-            }
-          }
-        };
+          };
 
-        // Close the log file after all launches
-        if (can_vectorize_all) {
-          launch_kernel(true);
-        } else {
-          launch_kernel(false);
-        }
+          // Close the log file after all launches
+          if (can_vectorize_all) {
+            launch_kernel(true);
+          } else {
+            launch_kernel(false);
+          }
+        );
       });
 
   C10_CUDA_KERNEL_LAUNCH_CHECK();

From 0466c59c3880c61962c0b52748882f32dfceaede Mon Sep 17 00:00:00 2001
From: Ther-LF <2639852836@qq.com>
Date: Tue, 16 Dec 2025 23:04:54 +0800
Subject: [PATCH 11/37] perf:replace int div/mod by bit-wise op

---
 .../sgl_diffusion/rope/rotary_embedding.cu    | 370 ++++++++++--------
 1 file changed, 209 insertions(+), 161 deletions(-)

diff --git a/sgl-kernel/csrc/sgl_diffusion/rope/rotary_embedding.cu b/sgl-kernel/csrc/sgl_diffusion/rope/rotary_embedding.cu
index 84b37e3da852..6fc83a5a7d58 100644
--- a/sgl-kernel/csrc/sgl_diffusion/rope/rotary_embedding.cu
+++ b/sgl-kernel/csrc/sgl_diffusion/rope/rotary_embedding.cu
@@ -400,19 +400,54 @@ __global__ void rotary_embedding_kernel_1d(
 // Case 32 -> head_size 64
 // Case 64 -> head_size 128
 // Case 128 -> head_size 256
-#define DISPATCH_ROT_DIM(embed_dim, ...) \
-    switch (embed_dim) { \
-        case 16: { constexpr int ROT_EMBED_DIM = 16; __VA_ARGS__; break; } \
-        case 32: { constexpr int ROT_EMBED_DIM = 32; __VA_ARGS__; break; } \
-        case 64: { constexpr int ROT_EMBED_DIM = 64; __VA_ARGS__; break; } \
-        case 96: { constexpr int ROT_EMBED_DIM = 96; __VA_ARGS__; break; } \
-        case 128: { constexpr int ROT_EMBED_DIM = 128; __VA_ARGS__; break; } \
-        case 256: { constexpr int ROT_EMBED_DIM = 256; __VA_ARGS__; break; } \
-        case 512: { constexpr int ROT_EMBED_DIM = 512; __VA_ARGS__; break; } \
-        case 1024: { constexpr int ROT_EMBED_DIM = 1024; __VA_ARGS__; break; } \
-        default: { constexpr int ROT_EMBED_DIM = 0; __VA_ARGS__; break; } \
-    }
-
+#define DISPATCH_ROT_DIM(embed_dim, ...)  \
+  switch (embed_dim) {                    \
+    case 16: {                            \
+      constexpr int ROT_EMBED_DIM = 16;   \
+      __VA_ARGS__;                        \
+      break;                              \
+    }                                     \
+    case 32: {                            \
+      constexpr int ROT_EMBED_DIM = 32;   \
+      __VA_ARGS__;                        \
+      break;                              \
+    }                                     \
+    case 64: {                            \
+      constexpr int ROT_EMBED_DIM = 64;   \
+      __VA_ARGS__;                        \
+      break;                              \
+    }                                     \
+    case 96: {                            \
+      constexpr int ROT_EMBED_DIM = 96;   \
+      __VA_ARGS__;                        \
+      break;                              \
+    }                                     \
+    case 128: {                           \
+      constexpr int ROT_EMBED_DIM = 128;  \
+      __VA_ARGS__;                        \
+      break;                              \
+    }                                     \
+    case 256: {                           \
+      constexpr int ROT_EMBED_DIM = 256;  \
+      __VA_ARGS__;                        \
+      break;                              \
+    }                                     \
+    case 512: {                           \
+      constexpr int ROT_EMBED_DIM = 512;  \
+      __VA_ARGS__;                        \
+      break;                              \
+    }                                     \
+    case 1024: {                          \
+      constexpr int ROT_EMBED_DIM = 1024; \
+      __VA_ARGS__;                        \
+      break;                              \
+    }                                     \
+    default: {                            \
+      constexpr int ROT_EMBED_DIM = 0;    \
+      __VA_ARGS__;                        \
+      break;                              \
+    }                                     \
+  }
 
 void rotary_embedding_cos_sin(
     at::Tensor& cos,
@@ -588,154 +623,167 @@ void rotary_embedding_cos_sin(
         // We need 2 arrays (cos, sin) of size 'rot_dim_from_cache', each element is 'sizeof(torch_scalar_t)'
         size_t smem_size = rot_dim_from_cache * sizeof(torch_scalar_t) * 2;
 
-        DISPATCH_ROT_DIM(embed_dim_for_rotation,
-          auto launch_kernel = [&](bool vectorized) {
-            if (interleaved) {
-              if (use_grid_2d) {
-                if (vectorized) {
-                  rotary_embedding_kernel_2d<cuda_scalar_t, true, true, ROT_EMBED_DIM><<<grid_2d, block, smem_size, stream>>>(
-                      reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
-                      reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
-                      reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
-                      key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>()) : nullptr,
-                      rot_dim_from_cache,
-                      embed_dim_for_rotation,
-                      query_token_stride,
-                      key_token_stride,
-                      head_stride_query,
-                      head_stride_key,
-                      num_heads,
-                      num_kv_heads,
-                      (int)head_size,
-                      blocks_per_token);
-                } else {
-                  rotary_embedding_kernel_2d<cuda_scalar_t, true, false, ROT_EMBED_DIM><<<grid_2d, block, smem_size, stream>>>(
-                      reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
-                      reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
-                      reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
-                      key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>()) : nullptr,
-                      rot_dim_from_cache,
-                      embed_dim_for_rotation,
-                      query_token_stride,
-                      key_token_stride,
-                      head_stride_query,
-                      head_stride_key,
-                      num_heads,
-                      num_kv_heads,
-                      (int)head_size,
-                      blocks_per_token);
-                }
-              } else {
-                if (vectorized) {
-                  rotary_embedding_kernel_1d<cuda_scalar_t, true, true, ROT_EMBED_DIM><<<grid_1d, block, smem_size, stream>>>(
-                      reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
-                      reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
-                      reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
-                      key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>()) : nullptr,
-                      rot_dim_from_cache,
-                      embed_dim_for_rotation,
-                      query_token_stride,
-                      key_token_stride,
-                      head_stride_query,
-                      head_stride_key,
-                      num_heads,
-                      num_kv_heads,
-                      (int)head_size);
-                } else {
-                  rotary_embedding_kernel_1d<cuda_scalar_t, true, false, ROT_EMBED_DIM><<<grid_1d, block, smem_size, stream>>>(
-                      reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
-                      reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
-                      reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
-                      key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>()) : nullptr,
-                      rot_dim_from_cache,
-                      embed_dim_for_rotation,
-                      query_token_stride,
-                      key_token_stride,
-                      head_stride_query,
-                      head_stride_key,
-                      num_heads,
-                      num_kv_heads,
-                      (int)head_size);
-                }
-              }
-            } else {  // non-interleaved
-              if (use_grid_2d) {
-                if (vectorized) {
-                  rotary_embedding_kernel_2d<cuda_scalar_t, false, true, ROT_EMBED_DIM><<<grid_2d, block, smem_size, stream>>>(
-                      reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
-                      reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
-                      reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
-                      key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>()) : nullptr,
-                      rot_dim_from_cache,
-                      embed_dim_for_rotation,
-                      query_token_stride,
-                      key_token_stride,
-                      head_stride_query,
-                      head_stride_key,
-                      num_heads,
-                      num_kv_heads,
-                      (int)head_size,
-                      blocks_per_token);
-                } else {
-                  rotary_embedding_kernel_2d<cuda_scalar_t, false, false, ROT_EMBED_DIM><<<grid_2d, block, smem_size, stream>>>(
-                      reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
-                      reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
-                      reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
-                      key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>()) : nullptr,
-                      rot_dim_from_cache,
-                      embed_dim_for_rotation,
-                      query_token_stride,
-                      key_token_stride,
-                      head_stride_query,
-                      head_stride_key,
-                      num_heads,
-                      num_kv_heads,
-                      (int)head_size,
-                      blocks_per_token);
-                }
-              } else {
-                if (vectorized) {
-                  rotary_embedding_kernel_1d<cuda_scalar_t, false, true, ROT_EMBED_DIM><<<grid_1d, block, smem_size, stream>>>(
-                      reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
-                      reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
-                      reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
-                      key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>()) : nullptr,
-                      rot_dim_from_cache,
-                      embed_dim_for_rotation,
-                      query_token_stride,
-                      key_token_stride,
-                      head_stride_query,
-                      head_stride_key,
-                      num_heads,
-                      num_kv_heads,
-                      (int)head_size);
-                } else {
-                  rotary_embedding_kernel_1d<cuda_scalar_t, false, false, ROT_EMBED_DIM><<<grid_1d, block, smem_size, stream>>>(
-                      reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
-                      reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
-                      reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
-                      key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>()) : nullptr,
-                      rot_dim_from_cache,
-                      embed_dim_for_rotation,
-                      query_token_stride,
-                      key_token_stride,
-                      head_stride_query,
-                      head_stride_key,
-                      num_heads,
-                      num_kv_heads,
-                      (int)head_size);
-                }
-              }
-            }
-          };
-
-          // Close the log file after all launches
-          if (can_vectorize_all) {
-            launch_kernel(true);
-          } else {
-            launch_kernel(false);
-          }
-        );
+        DISPATCH_ROT_DIM(
+            embed_dim_for_rotation,
+            auto launch_kernel =
+                [&](bool vectorized) {
+                  if (interleaved) {
+                    if (use_grid_2d) {
+                      if (vectorized) {
+                        rotary_embedding_kernel_2d<cuda_scalar_t, true, true, ROT_EMBED_DIM>
+                            <<<grid_2d, block, smem_size, stream>>>(
+                                reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
+                                reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
+                                reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
+                                key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>())
+                                                : nullptr,
+                                rot_dim_from_cache,
+                                embed_dim_for_rotation,
+                                query_token_stride,
+                                key_token_stride,
+                                head_stride_query,
+                                head_stride_key,
+                                num_heads,
+                                num_kv_heads,
+                                (int)head_size,
+                                blocks_per_token);
+                      } else {
+                        rotary_embedding_kernel_2d<cuda_scalar_t, true, false, ROT_EMBED_DIM>
+                            <<<grid_2d, block, smem_size, stream>>>(
+                                reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
+                                reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
+                                reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
+                                key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>())
+                                                : nullptr,
+                                rot_dim_from_cache,
+                                embed_dim_for_rotation,
+                                query_token_stride,
+                                key_token_stride,
+                                head_stride_query,
+                                head_stride_key,
+                                num_heads,
+                                num_kv_heads,
+                                (int)head_size,
+                                blocks_per_token);
+                      }
+                    } else {
+                      if (vectorized) {
+                        rotary_embedding_kernel_1d<cuda_scalar_t, true, true, ROT_EMBED_DIM>
+                            <<<grid_1d, block, smem_size, stream>>>(
+                                reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
+                                reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
+                                reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
+                                key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>())
+                                                : nullptr,
+                                rot_dim_from_cache,
+                                embed_dim_for_rotation,
+                                query_token_stride,
+                                key_token_stride,
+                                head_stride_query,
+                                head_stride_key,
+                                num_heads,
+                                num_kv_heads,
+                                (int)head_size);
+                      } else {
+                        rotary_embedding_kernel_1d<cuda_scalar_t, true, false, ROT_EMBED_DIM>
+                            <<<grid_1d, block, smem_size, stream>>>(
+                                reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
+                                reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
+                                reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
+                                key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>())
+                                                : nullptr,
+                                rot_dim_from_cache,
+                                embed_dim_for_rotation,
+                                query_token_stride,
+                                key_token_stride,
+                                head_stride_query,
+                                head_stride_key,
+                                num_heads,
+                                num_kv_heads,
+                                (int)head_size);
+                      }
+                    }
+                  } else {  // non-interleaved
+                    if (use_grid_2d) {
+                      if (vectorized) {
+                        rotary_embedding_kernel_2d<cuda_scalar_t, false, true, ROT_EMBED_DIM>
+                            <<<grid_2d, block, smem_size, stream>>>(
+                                reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
+                                reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
+                                reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
+                                key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>())
+                                                : nullptr,
+                                rot_dim_from_cache,
+                                embed_dim_for_rotation,
+                                query_token_stride,
+                                key_token_stride,
+                                head_stride_query,
+                                head_stride_key,
+                                num_heads,
+                                num_kv_heads,
+                                (int)head_size,
+                                blocks_per_token);
+                      } else {
+                        rotary_embedding_kernel_2d<cuda_scalar_t, false, false, ROT_EMBED_DIM>
+                            <<<grid_2d, block, smem_size, stream>>>(
+                                reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
+                                reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
+                                reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
+                                key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>())
+                                                : nullptr,
+                                rot_dim_from_cache,
+                                embed_dim_for_rotation,
+                                query_token_stride,
+                                key_token_stride,
+                                head_stride_query,
+                                head_stride_key,
+                                num_heads,
+                                num_kv_heads,
+                                (int)head_size,
+                                blocks_per_token);
+                      }
+                    } else {
+                      if (vectorized) {
+                        rotary_embedding_kernel_1d<cuda_scalar_t, false, true, ROT_EMBED_DIM>
+                            <<<grid_1d, block, smem_size, stream>>>(
+                                reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
+                                reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
+                                reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
+                                key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>())
+                                                : nullptr,
+                                rot_dim_from_cache,
+                                embed_dim_for_rotation,
+                                query_token_stride,
+                                key_token_stride,
+                                head_stride_query,
+                                head_stride_key,
+                                num_heads,
+                                num_kv_heads,
+                                (int)head_size);
+                      } else {
+                        rotary_embedding_kernel_1d<cuda_scalar_t, false, false, ROT_EMBED_DIM>
+                            <<<grid_1d, block, smem_size, stream>>>(
+                                reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
+                                reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
+                                reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
+                                key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>())
+                                                : nullptr,
+                                rot_dim_from_cache,
+                                embed_dim_for_rotation,
+                                query_token_stride,
+                                key_token_stride,
+                                head_stride_query,
+                                head_stride_key,
+                                num_heads,
+                                num_kv_heads,
+                                (int)head_size);
+                      }
+                    }
+                  }
+                };
+
+            // Close the log file after all launches
+            if (can_vectorize_all) { launch_kernel(true); } else { launch_kernel(false); });
       });
 
   C10_CUDA_KERNEL_LAUNCH_CHECK();

From aa02f4db78983b8f3453ea3caaec709ca8179c2e Mon Sep 17 00:00:00 2001
From: Ther-LF <2639852836@qq.com>
Date: Tue, 16 Dec 2025 23:50:59 +0800
Subject: [PATCH 12/37] perf:remove smem add vec load for sin/cos

---
 .../sgl_diffusion/rope/rotary_embedding.cu    | 128 ++++++++----------
 1 file changed, 53 insertions(+), 75 deletions(-)

diff --git a/sgl-kernel/csrc/sgl_diffusion/rope/rotary_embedding.cu b/sgl-kernel/csrc/sgl_diffusion/rope/rotary_embedding.cu
index 6fc83a5a7d58..dceeee4a3a13 100644
--- a/sgl-kernel/csrc/sgl_diffusion/rope/rotary_embedding.cu
+++ b/sgl-kernel/csrc/sgl_diffusion/rope/rotary_embedding.cu
@@ -39,7 +39,7 @@ inline __device__ void apply_token_rotary_embedding(
     const int x_index = 2 * rot_offset;
     const int y_index = x_index + 1;
 
-    // Directly access SMEM
+    // Directly access Global Memory
     const float cos_val = static_cast<float>(cos_ptr[rot_offset]);
     const float sin_val = static_cast<float>(sin_ptr[rot_offset]);
 
@@ -53,7 +53,7 @@ inline __device__ void apply_token_rotary_embedding(
     const int x_index = rot_offset;
     const int y_index = rot_offset + embed_dim;
 
-    // Directly access SMEM
+    // Directly access Global Memory
     const float cos_val_x = static_cast<float>(cos_ptr[rot_offset]);
     const float sin_val_x = static_cast<float>(sin_ptr[rot_offset]);
     const float cos_val_y = static_cast<float>(cos_ptr[rot_offset + embed_dim]);
@@ -93,14 +93,26 @@ inline __device__ void apply_token_rotary_embedding_vec(
     VecUnion data;
     data.vec = *reinterpret_cast<const vec_t*>(arr + rot_offset * 2);
 
+    // Vectorized load for cos/sin: always 8 bytes (half of vector size, float2)
+    // float2 covers (kElePerVec / 2) scalar_t elements (pairs)
+    float2 cos_vec = *reinterpret_cast<const float2*>(cos_ptr + rot_offset);
+    float2 sin_vec = *reinterpret_cast<const float2*>(sin_ptr + rot_offset);
+
+    union CosSinVec {
+      float2 vec;
+      scalar_t elems[kElePerVec / 2];
+    };
+    CosSinVec cos_u, sin_u;
+    cos_u.vec = cos_vec;
+    sin_u.vec = sin_vec;
+
 #pragma unroll
     for (int i = 0; i < kElePerVec; i += 2) {
       // data.elems[i] is x, data.elems[i+1] is y
       // They correspond to pair index: rot_offset + (i / 2)
-      int curr_rot_offset = rot_offset + i / 2;
-
-      float cos_val = static_cast<float>(cos_ptr[curr_rot_offset]);
-      float sin_val = static_cast<float>(sin_ptr[curr_rot_offset]);
+      int idx = i / 2;
+      float cos_val = static_cast<float>(cos_u.elems[idx]);
+      float sin_val = static_cast<float>(sin_u.elems[idx]);
 
       float x = static_cast<float>(data.elems[i]);
       float y = static_cast<float>(data.elems[i + 1]);
@@ -120,20 +132,29 @@ inline __device__ void apply_token_rotary_embedding_vec(
     data_x.vec = *reinterpret_cast<const vec_t*>(arr + rot_offset);
     data_y.vec = *reinterpret_cast<const vec_t*>(arr + rot_offset + embed_dim);
 
+    // Vectorized load for cos/sin: always 16 bytes (full vector size, float4)
+    // float4 covers kElePerVec scalar_t elements
+    float4 cos_vec_x = *reinterpret_cast<const float4*>(cos_ptr + rot_offset);
+    float4 sin_vec_x = *reinterpret_cast<const float4*>(sin_ptr + rot_offset);
+    float4 cos_vec_y = *reinterpret_cast<const float4*>(cos_ptr + rot_offset + embed_dim);
+    float4 sin_vec_y = *reinterpret_cast<const float4*>(sin_ptr + rot_offset + embed_dim);
+
+    union CosSinVec {
+      float4 vec;
+      scalar_t elems[kElePerVec];
+    };
+    CosSinVec cos_u_x, sin_u_x, cos_u_y, sin_u_y;
+    cos_u_x.vec = cos_vec_x;
+    sin_u_x.vec = sin_vec_x;
+    cos_u_y.vec = cos_vec_y;
+    sin_u_y.vec = sin_vec_y;
+
 #pragma unroll
     for (int i = 0; i < kElePerVec; ++i) {
-      int curr_rot_offset = rot_offset + i;
-
-      // In non-interleaved, we might need different cos/sin for X and Y depending on implementation,
-      // but standard RoPE uses the same angle for the pair.
-      // Based on original scalar code:
-      // cos_val_x = cos_ptr[rot_offset]
-      // cos_val_y = cos_ptr[rot_offset + embed_dim]
-
-      float cos_val_x = static_cast<float>(cos_ptr[curr_rot_offset]);
-      float sin_val_x = static_cast<float>(sin_ptr[curr_rot_offset]);
-      float cos_val_y = static_cast<float>(cos_ptr[curr_rot_offset + embed_dim]);
-      float sin_val_y = static_cast<float>(sin_ptr[curr_rot_offset + embed_dim]);
+      float cos_val_x = static_cast<float>(cos_u_x.elems[i]);
+      float sin_val_x = static_cast<float>(sin_u_x.elems[i]);
+      float cos_val_y = static_cast<float>(cos_u_y.elems[i]);
+      float sin_val_y = static_cast<float>(sin_u_y.elems[i]);
 
       float x = static_cast<float>(data_x.elems[i]);
       float y = static_cast<float>(data_y.elems[i]);
@@ -183,9 +204,7 @@ __global__ void rotary_embedding_kernel_2d(
     const int num_kv_heads,
     const int head_size_arg,
     const int blocks_per_token) {
-  extern __shared__ char smem_[];
-  scalar_t* s_cos = reinterpret_cast<scalar_t*>(smem_);
-  scalar_t* s_sin = s_cos + rot_dim_arg;
+  // Removed shared memory allocation and loading logic
 
   const int token_idx = blockIdx.x;
   if (token_idx >= gridDim.x) {
@@ -198,29 +217,8 @@ __global__ void rotary_embedding_kernel_2d(
 
   constexpr int kVecBytes = 16;
   const int scalar_size = sizeof(scalar_t);
-  const int vec_size = kVecBytes / scalar_size;
 
-  // Load Cos/Sin to SMEM
-  // Always use vectorized load if available, but fallback to scalar loop if dim is not multiple
-  if constexpr (vectorized) {
-    using vec_t = float4;
-    const vec_t* cos_vec_ptr = reinterpret_cast<const vec_t*>(current_token_cos_ptr);
-    const vec_t* sin_vec_ptr = reinterpret_cast<const vec_t*>(current_token_sin_ptr);
-    vec_t* s_cos_vec = reinterpret_cast<vec_t*>(s_cos);
-    vec_t* s_sin_vec = reinterpret_cast<vec_t*>(s_sin);
-
-    for (int i = threadIdx.x; i < rot_dim_arg / vec_size; i += blockDim.x) {
-      s_cos_vec[i] = SGLANG_LDG(cos_vec_ptr + i);
-      s_sin_vec[i] = SGLANG_LDG(sin_vec_ptr + i);
-    }
-  } else {
-    for (int i = threadIdx.x; i < rot_dim_arg; i += blockDim.x) {
-      s_cos[i] = SGLANG_LDG(current_token_cos_ptr + i);
-      s_sin[i] = SGLANG_LDG(current_token_sin_ptr + i);
-    }
-  }
-  // Essential synchronization
-  __syncthreads();
+  // Removed SMEM loading and __syncthreads()
 
   scalar_t* query_for_token = query_total + token_idx * (int)query_token_stride;
   scalar_t* key_for_token = (key_total != nullptr) ? (key_total + token_idx * (int)key_token_stride) : nullptr;
@@ -245,7 +243,7 @@ __global__ void rotary_embedding_kernel_2d(
 
       scalar_t* query_for_token_head = query_for_token + head_idx * (int)head_stride_query;
       apply_token_rotary_embedding_vec<scalar_t, interleaved>(
-          query_for_token_head, s_cos, s_sin, rot_offset, embed_dim_for_rotation);
+          query_for_token_head, current_token_cos_ptr, current_token_sin_ptr, rot_offset, embed_dim_for_rotation);
     }
 
     if (key_for_token != nullptr) {
@@ -256,7 +254,7 @@ __global__ void rotary_embedding_kernel_2d(
 
         scalar_t* key_for_token_head = key_for_token + head_idx * (int)head_stride_key;
         apply_token_rotary_embedding_vec<scalar_t, interleaved>(
-            key_for_token_head, s_cos, s_sin, rot_offset, embed_dim_for_rotation);
+            key_for_token_head, current_token_cos_ptr, current_token_sin_ptr, rot_offset, embed_dim_for_rotation);
       }
     }
   } else {
@@ -271,7 +269,7 @@ __global__ void rotary_embedding_kernel_2d(
 
       scalar_t* query_for_token_head = query_for_token + head_idx * (int)head_stride_query;
       apply_token_rotary_embedding<scalar_t, interleaved>(
-          query_for_token_head, s_cos, s_sin, rot_offset, embed_dim_for_rotation);
+          query_for_token_head, current_token_cos_ptr, current_token_sin_ptr, rot_offset, embed_dim_for_rotation);
     }
 
     if (key_for_token != nullptr) {
@@ -282,7 +280,7 @@ __global__ void rotary_embedding_kernel_2d(
 
         scalar_t* key_for_token_head = key_for_token + head_idx * (int)head_stride_key;
         apply_token_rotary_embedding<scalar_t, interleaved>(
-            key_for_token_head, s_cos, s_sin, rot_offset, embed_dim_for_rotation);
+            key_for_token_head, current_token_cos_ptr, current_token_sin_ptr, rot_offset, embed_dim_for_rotation);
       }
     }
   }
@@ -304,9 +302,7 @@ __global__ void rotary_embedding_kernel_1d(
     const int num_heads,
     const int num_kv_heads,
     const int head_size_arg) {
-  extern __shared__ char smem_[];
-  scalar_t* s_cos = reinterpret_cast<scalar_t*>(smem_);
-  scalar_t* s_sin = s_cos + rot_dim_arg;
+  // Removed shared memory allocation and loading logic
 
   const int token_idx = blockIdx.x;
   if (token_idx >= gridDim.x) return;
@@ -316,26 +312,8 @@ __global__ void rotary_embedding_kernel_1d(
 
   constexpr int kVecBytes = 16;
   const int scalar_size = sizeof(scalar_t);
-  const int vec_size = kVecBytes / scalar_size;
 
-  if constexpr (vectorized) {
-    using vec_t = float4;
-    const vec_t* cos_vec_ptr = reinterpret_cast<const vec_t*>(current_token_cos_ptr);
-    const vec_t* sin_vec_ptr = reinterpret_cast<const vec_t*>(current_token_sin_ptr);
-    vec_t* s_cos_vec = reinterpret_cast<vec_t*>(s_cos);
-    vec_t* s_sin_vec = reinterpret_cast<vec_t*>(s_sin);
-
-    for (int i = threadIdx.x; i < rot_dim_arg / vec_size; i += blockDim.x) {
-      s_cos_vec[i] = SGLANG_LDG(cos_vec_ptr + i);
-      s_sin_vec[i] = SGLANG_LDG(sin_vec_ptr + i);
-    }
-  } else {
-    for (int i = threadIdx.x; i < rot_dim_arg; i += blockDim.x) {
-      s_cos[i] = SGLANG_LDG(current_token_cos_ptr + i);
-      s_sin[i] = SGLANG_LDG(current_token_sin_ptr + i);
-    }
-  }
-  __syncthreads();
+  // Removed SMEM loading and __syncthreads()
 
   scalar_t* query_for_token = query_total + token_idx * (int)query_token_stride;
   scalar_t* key_for_token = (key_total != nullptr) ? (key_total + token_idx * (int)key_token_stride) : nullptr;
@@ -355,7 +333,7 @@ __global__ void rotary_embedding_kernel_1d(
       scalar_t* query_for_token_head = query_for_token + head_idx * (int)head_stride_query;
 
       apply_token_rotary_embedding_vec<scalar_t, interleaved>(
-          query_for_token_head, s_cos, s_sin, rot_offset, embed_dim_for_rotation);
+          query_for_token_head, current_token_cos_ptr, current_token_sin_ptr, rot_offset, embed_dim_for_rotation);
     }
 
     if (key_for_token != nullptr) {
@@ -366,7 +344,7 @@ __global__ void rotary_embedding_kernel_1d(
 
         scalar_t* key_for_token_head = key_for_token + head_idx * (int)head_stride_key;
         apply_token_rotary_embedding_vec<scalar_t, interleaved>(
-            key_for_token_head, s_cos, s_sin, rot_offset, embed_dim_for_rotation);
+            key_for_token_head, current_token_cos_ptr, current_token_sin_ptr, rot_offset, embed_dim_for_rotation);
       }
     }
   } else {
@@ -378,7 +356,7 @@ __global__ void rotary_embedding_kernel_1d(
       scalar_t* query_for_token_head = query_for_token + head_idx * (int)head_stride_query;
 
       apply_token_rotary_embedding<scalar_t, interleaved>(
-          query_for_token_head, s_cos, s_sin, rot_offset, embed_dim_for_rotation);
+          query_for_token_head, current_token_cos_ptr, current_token_sin_ptr, rot_offset, embed_dim_for_rotation);
     }
 
     if (key_for_token != nullptr) {
@@ -389,7 +367,7 @@ __global__ void rotary_embedding_kernel_1d(
 
         scalar_t* key_for_token_head = key_for_token + head_idx * (int)head_stride_key;
         apply_token_rotary_embedding<scalar_t, interleaved>(
-            key_for_token_head, s_cos, s_sin, rot_offset, embed_dim_for_rotation);
+            key_for_token_head, current_token_cos_ptr, current_token_sin_ptr, rot_offset, embed_dim_for_rotation);
       }
     }
   }
@@ -620,8 +598,8 @@ void rotary_embedding_cos_sin(
         dim3 grid_2d((int)num_tokens, std::max(1, blocks_per_token));
         dim3 grid_1d((int)num_tokens);
 
-        // We need 2 arrays (cos, sin) of size 'rot_dim_from_cache', each element is 'sizeof(torch_scalar_t)'
-        size_t smem_size = rot_dim_from_cache * sizeof(torch_scalar_t) * 2;
+        // No shared memory needed
+        size_t smem_size = 0;
 
         DISPATCH_ROT_DIM(
             embed_dim_for_rotation,

From ead2093f7b419814bbd4cda6b5999e145bfd0e48 Mon Sep 17 00:00:00 2001
From: Ther-LF <2639852836@qq.com>
Date: Wed, 17 Dec 2025 17:38:14 +0800
Subject: [PATCH 13/37] perf:pipeline

---
 .../benchmark/bench_mm_rotary_embedding.py    |   1 +
 .../benchmark/profile_rotary_embedding.py     | 103 ++++++
 .../sgl_diffusion/rope/rotary_embedding.cu    | 300 ++++++++++++++++--
 3 files changed, 376 insertions(+), 28 deletions(-)
 create mode 100644 sgl-kernel/benchmark/profile_rotary_embedding.py

diff --git a/sgl-kernel/benchmark/bench_mm_rotary_embedding.py b/sgl-kernel/benchmark/bench_mm_rotary_embedding.py
index 0cc2dc156557..180b2e800b67 100755
--- a/sgl-kernel/benchmark/bench_mm_rotary_embedding.py
+++ b/sgl-kernel/benchmark/bench_mm_rotary_embedding.py
@@ -63,6 +63,7 @@ def benchmark_mm_rotary_embedding() -> None:
         (1, 32, 1, 128),
         (32, 8, 8, 128),
         (1, 32, 8, 80),
+        (1, 64, 8, 256),
     ]
 
     for batch_size, num_heads, num_kv_heads, head_size in configs:
diff --git a/sgl-kernel/benchmark/profile_rotary_embedding.py b/sgl-kernel/benchmark/profile_rotary_embedding.py
new file mode 100644
index 000000000000..a8d981e655fb
--- /dev/null
+++ b/sgl-kernel/benchmark/profile_rotary_embedding.py
@@ -0,0 +1,103 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+
+import sys
+
+import torch
+from sgl_kernel.rotary_embedding import rotary_embedding_cos_sin
+
+
+def compute_cos_sin_cache(
+    max_seq_len: int,
+    rotary_dim: int,
+    base: float = 10000.0,
+    dtype: torch.dtype = torch.float32,
+):
+    inv_freq = 1.0 / (
+        base ** (torch.arange(0, rotary_dim, 2, dtype=torch.float32) / rotary_dim)
+    )
+    t = torch.arange(max_seq_len, dtype=torch.float32)
+    freqs = torch.einsum("i,j->ij", t, inv_freq)
+    cos = freqs.cos().to(dtype)
+    sin = freqs.sin().to(dtype)
+    return cos, sin
+
+
+def profile_rotary_embedding():
+    # 配置参数：选取 Benchmark 中出现过的最大维度组合
+    # Batch Size: 32
+    # Seq Len: 8192 (Benchmark 列表中的最大值)
+    # Num Heads: 32 (常规配置，保证计算量)
+    # Num KV Heads: 32 (与 Heads 相同，最大化 Key 的负载)
+    # Head Size: 256 (Benchmark Configs 中的最大值)
+    batch_size = 1
+    seq_len = 8192
+    num_heads = 32
+    num_kv_heads = 8
+    head_size = 80
+    rotary_dim = head_size
+    dtype = torch.bfloat16
+    device = "cuda"
+
+    print(
+        f"Profiling Config: Batch={batch_size}, SeqLen={seq_len}, "
+        f"Heads={num_heads}, KV_Heads={num_kv_heads}, HeadSize={head_size}, Dtype={dtype}"
+    )
+
+    # 1. 准备 Cos/Sin Cache
+    try:
+        cos_cache, sin_cache = compute_cos_sin_cache(seq_len, rotary_dim, dtype=dtype)
+        cos_cache = cos_cache.to(device)
+        sin_cache = sin_cache.to(device)
+    except Exception as e:
+        print(f"Error creating cache: {e}")
+        return
+
+    # 2. 准备输入数据
+    # 计算总 token 数
+    num_tokens = batch_size * seq_len
+    print(f"Total tokens: {num_tokens}")
+
+    try:
+        # 生成 Query 和 Key
+        query = torch.randn(
+            num_tokens, num_heads * head_size, dtype=dtype, device=device
+        )
+        key = torch.randn(
+            num_tokens, num_kv_heads * head_size, dtype=dtype, device=device
+        )
+
+        # 生成对应的 Cos/Sin (Expand per token)
+        positions = torch.arange(seq_len, device=device, dtype=torch.int64).repeat(
+            batch_size
+        )
+
+        cos = cos_cache[positions]
+        sin = sin_cache[positions]
+    except torch.cuda.OutOfMemoryError:
+        print("OOM: The configuration is too large for the current GPU memory.")
+        return
+
+    # 3. Warmup
+    print("Warming up...")
+    for _ in range(5):
+        rotary_embedding_cos_sin(cos, sin, query, key, head_size, True)
+    torch.cuda.synchronize()
+
+    # 4. Profile Run (只运行一次，便于 NCU 捕获)
+    print("Running for profiling...")
+
+    # 开启 Profiler (如果使用 ncu --launch-count 1 ... python script.py，这行其实是可选的，
+    # 但显式开启有助于 ncu --profile-from-start off 模式)
+    torch.cuda.profiler.start()
+
+    rotary_embedding_cos_sin(cos, sin, query, key, head_size, True)
+
+    torch.cuda.profiler.stop()
+
+    torch.cuda.synchronize()
+    print("Done.")
+
+
+if __name__ == "__main__":
+    profile_rotary_embedding()
diff --git a/sgl-kernel/csrc/sgl_diffusion/rope/rotary_embedding.cu b/sgl-kernel/csrc/sgl_diffusion/rope/rotary_embedding.cu
index dceeee4a3a13..ea817567b348 100644
--- a/sgl-kernel/csrc/sgl_diffusion/rope/rotary_embedding.cu
+++ b/sgl-kernel/csrc/sgl_diffusion/rope/rotary_embedding.cu
@@ -26,6 +26,123 @@
 
 #include "utils.h"
 
+// -------------------------------------------------------------------------
+// Helper structures and functions for manual pipeline interleaving
+// -------------------------------------------------------------------------
+
+template <typename scalar_t, bool interleaved>
+struct RotaryVecData;
+
+template <typename scalar_t>
+struct RotaryVecData<scalar_t, true> {
+  float4 vec;
+  float2 cos_vec;
+  float2 sin_vec;
+};
+
+template <typename scalar_t>
+struct RotaryVecData<scalar_t, false> {
+  float4 vec_x;
+  float4 vec_y;
+  float4 cos_x;
+  float4 sin_x;
+  float4 cos_y;
+  float4 sin_y;
+};
+
+template <typename scalar_t, bool interleaved>
+inline __device__ RotaryVecData<scalar_t, interleaved> load_rotary_vec(
+    const scalar_t* __restrict__ arr,
+    const scalar_t* __restrict__ cos_ptr,
+    const scalar_t* __restrict__ sin_ptr,
+    int rot_offset,
+    int embed_dim) {
+  RotaryVecData<scalar_t, interleaved> data;
+  if constexpr (interleaved) {
+    data.vec = *reinterpret_cast<const float4*>(arr + rot_offset * 2);
+    data.cos_vec = *reinterpret_cast<const float2*>(cos_ptr + rot_offset);
+    data.sin_vec = *reinterpret_cast<const float2*>(sin_ptr + rot_offset);
+  } else {
+    data.vec_x = *reinterpret_cast<const float4*>(arr + rot_offset);
+    data.vec_y = *reinterpret_cast<const float4*>(arr + rot_offset + embed_dim);
+    data.cos_x = *reinterpret_cast<const float4*>(cos_ptr + rot_offset);
+    data.sin_x = *reinterpret_cast<const float4*>(sin_ptr + rot_offset);
+    data.cos_y = *reinterpret_cast<const float4*>(cos_ptr + rot_offset + embed_dim);
+    data.sin_y = *reinterpret_cast<const float4*>(sin_ptr + rot_offset + embed_dim);
+  }
+  return data;
+}
+
+template <typename scalar_t, bool interleaved>
+inline __device__ void compute_store_rotary_vec(
+    scalar_t* __restrict__ arr, const RotaryVecData<scalar_t, interleaved>& data, int rot_offset, int embed_dim) {
+  constexpr int kVecBytes = 16;
+  constexpr int kElePerVec = kVecBytes / sizeof(scalar_t);
+
+  if constexpr (interleaved) {
+    union VecUnion {
+      float4 vec;
+      scalar_t elems[kElePerVec];
+    };
+    union CosSinVec {
+      float2 vec;
+      scalar_t elems[kElePerVec / 2];
+    };
+
+    VecUnion v;
+    v.vec = data.vec;
+    CosSinVec c, s;
+    c.vec = data.cos_vec;
+    s.vec = data.sin_vec;
+
+#pragma unroll
+    for (int i = 0; i < kElePerVec; i += 2) {
+      int idx = i / 2;
+      float cos_val = static_cast<float>(c.elems[idx]);
+      float sin_val = static_cast<float>(s.elems[idx]);
+      float x = static_cast<float>(v.elems[i]);
+      float y = static_cast<float>(v.elems[i + 1]);
+      v.elems[i] = static_cast<scalar_t>(x * cos_val - y * sin_val);
+      v.elems[i + 1] = static_cast<scalar_t>(y * cos_val + x * sin_val);
+    }
+    *reinterpret_cast<float4*>(arr + rot_offset * 2) = v.vec;
+
+  } else {
+    union VecUnion {
+      float4 vec;
+      scalar_t elems[kElePerVec];
+    };
+    union CosSinVec {
+      float4 vec;
+      scalar_t elems[kElePerVec];
+    };
+
+    VecUnion vx, vy;
+    vx.vec = data.vec_x;
+    vy.vec = data.vec_y;
+    CosSinVec cx, sx, cy, sy;
+    cx.vec = data.cos_x;
+    sx.vec = data.sin_x;
+    cy.vec = data.cos_y;
+    sy.vec = data.sin_y;
+
+#pragma unroll
+    for (int i = 0; i < kElePerVec; ++i) {
+      float cos_val_x = static_cast<float>(cx.elems[i]);
+      float sin_val_x = static_cast<float>(sx.elems[i]);
+      float cos_val_y = static_cast<float>(cy.elems[i]);
+      float sin_val_y = static_cast<float>(sy.elems[i]);
+      float x = static_cast<float>(vx.elems[i]);
+      float y = static_cast<float>(vy.elems[i]);
+
+      vx.elems[i] = static_cast<scalar_t>(x * cos_val_x - y * sin_val_x);
+      vy.elems[i] = static_cast<scalar_t>(y * cos_val_y + x * sin_val_y);
+    }
+    *reinterpret_cast<float4*>(arr + rot_offset) = vx.vec;
+    *reinterpret_cast<float4*>(arr + rot_offset + embed_dim) = vy.vec;
+  }
+}
+
 template <typename scalar_t, bool interleaved>
 inline __device__ void apply_token_rotary_embedding(
     scalar_t* __restrict__ arr,            // [head_size]
@@ -234,27 +351,82 @@ __global__ void rotary_embedding_kernel_2d(
     constexpr int pairs_per_step = interleaved ? (kElePerVec / 2) : kElePerVec;
 
     const int pair_stride = blockDim.x * blocks_per_token * pairs_per_step;
-    const int thread_pair_offset = (local_block_idx * blockDim.x + threadIdx.x) * pairs_per_step;
+    int i = (local_block_idx * blockDim.x + threadIdx.x) * pairs_per_step;
     const int nq_pairs = num_heads * embed_dim_for_rotation;
 
-    for (int i = thread_pair_offset; i < nq_pairs; i += pair_stride) {
-      const int head_idx = i / embed_dim_for_rotation;
-      const int rot_offset = i % embed_dim_for_rotation;
+    // PROLOGUE: Preload 0-th batch
+    RotaryVecData<scalar_t, interleaved> curr_data;
+    if (i < nq_pairs) {
+      int head_idx = i / embed_dim_for_rotation;
+      int rot_offset = i % embed_dim_for_rotation;
+      scalar_t* ptr = query_for_token + head_idx * (int)head_stride_query;
+      curr_data = load_rotary_vec<scalar_t, interleaved>(
+          ptr, current_token_cos_ptr, current_token_sin_ptr, rot_offset, embed_dim_for_rotation);
+    }
 
-      scalar_t* query_for_token_head = query_for_token + head_idx * (int)head_stride_query;
-      apply_token_rotary_embedding_vec<scalar_t, interleaved>(
-          query_for_token_head, current_token_cos_ptr, current_token_sin_ptr, rot_offset, embed_dim_for_rotation);
+    int next_i = i + pair_stride;
+
+    // MAIN LOOP
+    for (; i < nq_pairs; i += pair_stride, next_i += pair_stride) {
+      // 1. PREFETCH NEXT
+      RotaryVecData<scalar_t, interleaved> next_data;
+      bool active_next = (next_i < nq_pairs);
+      if (active_next) {
+        int head_idx_next = next_i / embed_dim_for_rotation;
+        int rot_offset_next = next_i % embed_dim_for_rotation;
+        scalar_t* ptr_next = query_for_token + head_idx_next * (int)head_stride_query;
+        next_data = load_rotary_vec<scalar_t, interleaved>(
+            ptr_next, current_token_cos_ptr, current_token_sin_ptr, rot_offset_next, embed_dim_for_rotation);
+      }
+
+      // 2. COMPUTE CURRENT
+      // The compute uses curr_data which corresponds to index 'i'
+      int head_idx = i / embed_dim_for_rotation;
+      int rot_offset = i % embed_dim_for_rotation;
+      scalar_t* ptr = query_for_token + head_idx * (int)head_stride_query;
+      compute_store_rotary_vec<scalar_t, interleaved>(ptr, curr_data, rot_offset, embed_dim_for_rotation);
+
+      // 3. SHIFT
+      curr_data = next_data;
     }
 
     if (key_for_token != nullptr) {
       const int nk_pairs = num_kv_heads * embed_dim_for_rotation;
-      for (int i = thread_pair_offset; i < nk_pairs; i += pair_stride) {
-        const int head_idx = i / embed_dim_for_rotation;
-        const int rot_offset = i % embed_dim_for_rotation;
+      int k_i = (local_block_idx * blockDim.x + threadIdx.x) * pairs_per_step;
+
+      // PROLOGUE (Key)
+      RotaryVecData<scalar_t, interleaved> curr_data_k;
+      if (k_i < nk_pairs) {
+        int head_idx = k_i / embed_dim_for_rotation;
+        int rot_offset = k_i % embed_dim_for_rotation;
+        scalar_t* ptr = key_for_token + head_idx * (int)head_stride_key;
+        curr_data_k = load_rotary_vec<scalar_t, interleaved>(
+            ptr, current_token_cos_ptr, current_token_sin_ptr, rot_offset, embed_dim_for_rotation);
+      }
 
-        scalar_t* key_for_token_head = key_for_token + head_idx * (int)head_stride_key;
-        apply_token_rotary_embedding_vec<scalar_t, interleaved>(
-            key_for_token_head, current_token_cos_ptr, current_token_sin_ptr, rot_offset, embed_dim_for_rotation);
+      int next_k_i = k_i + pair_stride;
+
+      // MAIN LOOP (Key)
+      for (; k_i < nk_pairs; k_i += pair_stride, next_k_i += pair_stride) {
+        // 1. PREFETCH NEXT
+        RotaryVecData<scalar_t, interleaved> next_data_k;
+        bool active_next = (next_k_i < nk_pairs);
+        if (active_next) {
+          int head_idx_next = next_k_i / embed_dim_for_rotation;
+          int rot_offset_next = next_k_i % embed_dim_for_rotation;
+          scalar_t* ptr_next = key_for_token + head_idx_next * (int)head_stride_key;
+          next_data_k = load_rotary_vec<scalar_t, interleaved>(
+              ptr_next, current_token_cos_ptr, current_token_sin_ptr, rot_offset_next, embed_dim_for_rotation);
+        }
+
+        // 2. COMPUTE CURRENT
+        int head_idx = k_i / embed_dim_for_rotation;
+        int rot_offset = k_i % embed_dim_for_rotation;
+        scalar_t* ptr = key_for_token + head_idx * (int)head_stride_key;
+        compute_store_rotary_vec<scalar_t, interleaved>(ptr, curr_data_k, rot_offset, embed_dim_for_rotation);
+
+        // 3. SHIFT
+        curr_data_k = next_data_k;
       }
     }
   } else {
@@ -288,7 +460,7 @@ __global__ void rotary_embedding_kernel_2d(
 
 // 1D grid kernel: each block handles one token
 template <typename scalar_t, bool interleaved, bool vectorized, int ROT_EMBED_DIM>
-__global__ void rotary_embedding_kernel_1d(
+__launch_bounds__(512) __global__ void rotary_embedding_kernel_1d(
     const scalar_t* __restrict__ cos_data,  // [num_tokens, rot_dim_arg]
     const scalar_t* __restrict__ sin_data,  // [num_tokens, rot_dim_arg]
     scalar_t* __restrict__ query_total,
@@ -322,29 +494,101 @@ __global__ void rotary_embedding_kernel_1d(
   const int embed_dim_for_rotation = (ROT_EMBED_DIM > 0) ? ROT_EMBED_DIM : embed_dim_for_rotation_arg;
 
   if constexpr (vectorized) {
-    using vec_t = float4;
+    constexpr int kVecBytes = 16;
     constexpr int kElePerVec = kVecBytes / sizeof(scalar_t);
+    // Interleaved: 1 vec = kElePerVec elements = kElePerVec/2 pairs
+    // Non-Interleaved: 1 vec (X) + 1 vec (Y) = kElePerVec elements = kElePerVec pairs
     constexpr int pairs_per_step = interleaved ? (kElePerVec / 2) : kElePerVec;
-
     const int nq_pairs = num_heads * embed_dim_for_rotation;
-    for (int i = threadIdx.x * pairs_per_step; i < nq_pairs; i += blockDim.x * pairs_per_step) {
-      const int head_idx = i / embed_dim_for_rotation;
-      const int rot_offset = i % embed_dim_for_rotation;
-      scalar_t* query_for_token_head = query_for_token + head_idx * (int)head_stride_query;
 
-      apply_token_rotary_embedding_vec<scalar_t, interleaved>(
-          query_for_token_head, current_token_cos_ptr, current_token_sin_ptr, rot_offset, embed_dim_for_rotation);
+    // Original Stride
+    const int stride = blockDim.x * pairs_per_step;
+
+    int i = threadIdx.x * pairs_per_step;
+
+    // ==========================================
+    // PROLOGUE: Preload first batch
+    // ==========================================
+    RotaryVecData<scalar_t, interleaved> curr_data;
+
+    if (i < nq_pairs) {
+      int head_idx = i / embed_dim_for_rotation;
+      int rot_offset = i % embed_dim_for_rotation;
+      scalar_t* ptr = query_for_token + head_idx * (int)head_stride_query;
+      curr_data = load_rotary_vec<scalar_t, interleaved>(
+          ptr, current_token_cos_ptr, current_token_sin_ptr, rot_offset, embed_dim_for_rotation);
     }
 
+    int next_i = i + stride;
+
+    // ==========================================
+    // MAIN LOOP
+    // ==========================================
+    for (; i < nq_pairs; i += stride, next_i += stride) {
+      // --- 1. PREFETCH NEXT BATCH ---
+      RotaryVecData<scalar_t, interleaved> next_data;
+
+      if (next_i < nq_pairs) {
+        int head_idx = next_i / embed_dim_for_rotation;
+        int rot_offset = next_i % embed_dim_for_rotation;
+        scalar_t* ptr = query_for_token + head_idx * (int)head_stride_query;
+        next_data = load_rotary_vec<scalar_t, interleaved>(
+            ptr, current_token_cos_ptr, current_token_sin_ptr, rot_offset, embed_dim_for_rotation);
+      }
+
+      // --- 2. COMPUTE CURRENT BATCH ---
+      if (i < nq_pairs) {
+        int head_idx = i / embed_dim_for_rotation;
+        int rot_offset = i % embed_dim_for_rotation;
+        scalar_t* ptr = query_for_token + head_idx * (int)head_stride_query;
+        compute_store_rotary_vec<scalar_t, interleaved>(ptr, curr_data, rot_offset, embed_dim_for_rotation);
+      }
+
+      // --- 3. SHIFT ---
+      curr_data = next_data;
+    }
+
+    // Key branch
     if (key_for_token != nullptr) {
       const int nk_pairs = num_kv_heads * embed_dim_for_rotation;
-      for (int i = threadIdx.x * pairs_per_step; i < nk_pairs; i += blockDim.x * pairs_per_step) {
-        const int head_idx = i / embed_dim_for_rotation;
-        const int rot_offset = i % embed_dim_for_rotation;
+      int k_i = threadIdx.x * pairs_per_step;
 
-        scalar_t* key_for_token_head = key_for_token + head_idx * (int)head_stride_key;
-        apply_token_rotary_embedding_vec<scalar_t, interleaved>(
-            key_for_token_head, current_token_cos_ptr, current_token_sin_ptr, rot_offset, embed_dim_for_rotation);
+      // PROLOGUE (Key): Preload first batch
+      RotaryVecData<scalar_t, interleaved> curr_data_k;
+
+      if (k_i < nk_pairs) {
+        int head_idx = k_i / embed_dim_for_rotation;
+        int rot_offset = k_i % embed_dim_for_rotation;
+        scalar_t* ptr = key_for_token + head_idx * (int)head_stride_key;
+        curr_data_k = load_rotary_vec<scalar_t, interleaved>(
+            ptr, current_token_cos_ptr, current_token_sin_ptr, rot_offset, embed_dim_for_rotation);
+      }
+
+      int next_k_i = k_i + stride;
+
+      // MAIN LOOP (Key)
+      for (; k_i < nk_pairs; k_i += stride, next_k_i += stride) {
+        // --- 1. PREFETCH NEXT BATCH ---
+        RotaryVecData<scalar_t, interleaved> next_data_k;
+
+        if (next_k_i < nk_pairs) {
+          int head_idx = next_k_i / embed_dim_for_rotation;
+          int rot_offset = next_k_i % embed_dim_for_rotation;
+          scalar_t* ptr = key_for_token + head_idx * (int)head_stride_key;
+          next_data_k = load_rotary_vec<scalar_t, interleaved>(
+              ptr, current_token_cos_ptr, current_token_sin_ptr, rot_offset, embed_dim_for_rotation);
+        }
+
+        // --- 2. COMPUTE CURRENT BATCH ---
+        if (k_i < nk_pairs) {
+          int head_idx = k_i / embed_dim_for_rotation;
+          int rot_offset = k_i % embed_dim_for_rotation;
+          scalar_t* ptr = key_for_token + head_idx * (int)head_stride_key;
+          compute_store_rotary_vec<scalar_t, interleaved>(ptr, curr_data_k, rot_offset, embed_dim_for_rotation);
+        }
+
+        // --- 3. SHIFT ---
+        curr_data_k = next_data_k;
       }
     }
   } else {

From 9bbc0a4d3a5a06bb0a465beab29c3245fd1fb0cd Mon Sep 17 00:00:00 2001
From: RubiaCx <1084281732@qq.com>
Date: Wed, 17 Dec 2025 04:45:19 -0800
Subject: [PATCH 14/37] Bench: add qwen-image shape

---
 .../benchmark/bench_mm_rotary_embedding.py    |    5 +-
 .../sgl_diffusion/rope/rotary_embedding.cu    | 1989 ++++++++---------
 2 files changed, 998 insertions(+), 996 deletions(-)

diff --git a/sgl-kernel/benchmark/bench_mm_rotary_embedding.py b/sgl-kernel/benchmark/bench_mm_rotary_embedding.py
index 180b2e800b67..2692ed62909e 100755
--- a/sgl-kernel/benchmark/bench_mm_rotary_embedding.py
+++ b/sgl-kernel/benchmark/bench_mm_rotary_embedding.py
@@ -51,7 +51,8 @@ def benchmark_mm_rotary_embedding() -> None:
     dtype = torch.bfloat16
     max_seq_len = 65536
 
-    seq_lens = [1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192]
+    # Include Qwen-image observed shapes (e.g., 3015 image tokens, small text token counts)
+    seq_lens = [1, 2, 4, 6, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 3015, 4096, 8192]
 
     # Test configurations: (batch_size, num_heads, num_kv_heads, head_size)
     configs = [
@@ -64,6 +65,8 @@ def benchmark_mm_rotary_embedding() -> None:
         (32, 8, 8, 128),
         (1, 32, 8, 80),
         (1, 64, 8, 256),
+        # Qwen-image observed: q/k view is [num_tokens, 24, 128] and cos/sin is [num_tokens, 64], interleaved=True
+        (1, 24, 24, 128),
     ]
 
     for batch_size, num_heads, num_kv_heads, head_size in configs:
diff --git a/sgl-kernel/csrc/sgl_diffusion/rope/rotary_embedding.cu b/sgl-kernel/csrc/sgl_diffusion/rope/rotary_embedding.cu
index ea817567b348..ad8885c46d9c 100644
--- a/sgl-kernel/csrc/sgl_diffusion/rope/rotary_embedding.cu
+++ b/sgl-kernel/csrc/sgl_diffusion/rope/rotary_embedding.cu
@@ -15,998 +15,997 @@
  * limitations under the License.
  */
 
-#include <ATen/cuda/CUDAContext.h>
-#include <c10/cuda/CUDAException.h>
-#include <c10/cuda/CUDAGuard.h>
-#include <c10/cuda/CUDAStream.h>
-#include <torch/all.h>
-
-#include <cmath>
-#include <vector>
-
-#include "utils.h"
-
-// -------------------------------------------------------------------------
-// Helper structures and functions for manual pipeline interleaving
-// -------------------------------------------------------------------------
-
-template <typename scalar_t, bool interleaved>
-struct RotaryVecData;
-
-template <typename scalar_t>
-struct RotaryVecData<scalar_t, true> {
-  float4 vec;
-  float2 cos_vec;
-  float2 sin_vec;
-};
-
-template <typename scalar_t>
-struct RotaryVecData<scalar_t, false> {
-  float4 vec_x;
-  float4 vec_y;
-  float4 cos_x;
-  float4 sin_x;
-  float4 cos_y;
-  float4 sin_y;
-};
-
-template <typename scalar_t, bool interleaved>
-inline __device__ RotaryVecData<scalar_t, interleaved> load_rotary_vec(
-    const scalar_t* __restrict__ arr,
-    const scalar_t* __restrict__ cos_ptr,
-    const scalar_t* __restrict__ sin_ptr,
-    int rot_offset,
-    int embed_dim) {
-  RotaryVecData<scalar_t, interleaved> data;
-  if constexpr (interleaved) {
-    data.vec = *reinterpret_cast<const float4*>(arr + rot_offset * 2);
-    data.cos_vec = *reinterpret_cast<const float2*>(cos_ptr + rot_offset);
-    data.sin_vec = *reinterpret_cast<const float2*>(sin_ptr + rot_offset);
-  } else {
-    data.vec_x = *reinterpret_cast<const float4*>(arr + rot_offset);
-    data.vec_y = *reinterpret_cast<const float4*>(arr + rot_offset + embed_dim);
-    data.cos_x = *reinterpret_cast<const float4*>(cos_ptr + rot_offset);
-    data.sin_x = *reinterpret_cast<const float4*>(sin_ptr + rot_offset);
-    data.cos_y = *reinterpret_cast<const float4*>(cos_ptr + rot_offset + embed_dim);
-    data.sin_y = *reinterpret_cast<const float4*>(sin_ptr + rot_offset + embed_dim);
-  }
-  return data;
-}
-
-template <typename scalar_t, bool interleaved>
-inline __device__ void compute_store_rotary_vec(
-    scalar_t* __restrict__ arr, const RotaryVecData<scalar_t, interleaved>& data, int rot_offset, int embed_dim) {
-  constexpr int kVecBytes = 16;
-  constexpr int kElePerVec = kVecBytes / sizeof(scalar_t);
-
-  if constexpr (interleaved) {
-    union VecUnion {
-      float4 vec;
-      scalar_t elems[kElePerVec];
-    };
-    union CosSinVec {
-      float2 vec;
-      scalar_t elems[kElePerVec / 2];
-    };
-
-    VecUnion v;
-    v.vec = data.vec;
-    CosSinVec c, s;
-    c.vec = data.cos_vec;
-    s.vec = data.sin_vec;
-
-#pragma unroll
-    for (int i = 0; i < kElePerVec; i += 2) {
-      int idx = i / 2;
-      float cos_val = static_cast<float>(c.elems[idx]);
-      float sin_val = static_cast<float>(s.elems[idx]);
-      float x = static_cast<float>(v.elems[i]);
-      float y = static_cast<float>(v.elems[i + 1]);
-      v.elems[i] = static_cast<scalar_t>(x * cos_val - y * sin_val);
-      v.elems[i + 1] = static_cast<scalar_t>(y * cos_val + x * sin_val);
-    }
-    *reinterpret_cast<float4*>(arr + rot_offset * 2) = v.vec;
-
-  } else {
-    union VecUnion {
-      float4 vec;
-      scalar_t elems[kElePerVec];
-    };
-    union CosSinVec {
-      float4 vec;
-      scalar_t elems[kElePerVec];
-    };
-
-    VecUnion vx, vy;
-    vx.vec = data.vec_x;
-    vy.vec = data.vec_y;
-    CosSinVec cx, sx, cy, sy;
-    cx.vec = data.cos_x;
-    sx.vec = data.sin_x;
-    cy.vec = data.cos_y;
-    sy.vec = data.sin_y;
-
-#pragma unroll
-    for (int i = 0; i < kElePerVec; ++i) {
-      float cos_val_x = static_cast<float>(cx.elems[i]);
-      float sin_val_x = static_cast<float>(sx.elems[i]);
-      float cos_val_y = static_cast<float>(cy.elems[i]);
-      float sin_val_y = static_cast<float>(sy.elems[i]);
-      float x = static_cast<float>(vx.elems[i]);
-      float y = static_cast<float>(vy.elems[i]);
-
-      vx.elems[i] = static_cast<scalar_t>(x * cos_val_x - y * sin_val_x);
-      vy.elems[i] = static_cast<scalar_t>(y * cos_val_y + x * sin_val_y);
-    }
-    *reinterpret_cast<float4*>(arr + rot_offset) = vx.vec;
-    *reinterpret_cast<float4*>(arr + rot_offset + embed_dim) = vy.vec;
-  }
-}
-
-template <typename scalar_t, bool interleaved>
-inline __device__ void apply_token_rotary_embedding(
-    scalar_t* __restrict__ arr,            // [head_size]
-    const scalar_t* __restrict__ cos_ptr,  // [rot_dim]
-    const scalar_t* __restrict__ sin_ptr,  // [rot_dim]
-    int rot_offset,
-    int embed_dim) {  // for non-interleaved: half dim
-  if constexpr (interleaved) {
-    // NeoX-style: interleaved layout [x0, y0, x1, y1, ...].
-    // cos/sin: [..., rotary_dim/2], one entry per (x, y) pair.
-    const int x_index = 2 * rot_offset;
-    const int y_index = x_index + 1;
-
-    // Directly access Global Memory
-    const float cos_val = static_cast<float>(cos_ptr[rot_offset]);
-    const float sin_val = static_cast<float>(sin_ptr[rot_offset]);
-
-    const float x = static_cast<float>(arr[x_index]);
-    const float y = static_cast<float>(arr[y_index]);
-    arr[x_index] = static_cast<scalar_t>(x * cos_val - y * sin_val);
-    arr[y_index] = static_cast<scalar_t>(y * cos_val + x * sin_val);
-  } else {
-    // GPT-J / LLaMA style: layout [x0, x1, ..., y0, y1, ...]
-    // cos/sin: [..., rotary_dim], one entry per (x, y) pair.
-    const int x_index = rot_offset;
-    const int y_index = rot_offset + embed_dim;
-
-    // Directly access Global Memory
-    const float cos_val_x = static_cast<float>(cos_ptr[rot_offset]);
-    const float sin_val_x = static_cast<float>(sin_ptr[rot_offset]);
-    const float cos_val_y = static_cast<float>(cos_ptr[rot_offset + embed_dim]);
-    const float sin_val_y = static_cast<float>(sin_ptr[rot_offset + embed_dim]);
-
-    const float x = static_cast<float>(arr[x_index]);
-    const float y = static_cast<float>(arr[y_index]);
-    arr[x_index] = static_cast<scalar_t>(x * cos_val_x - y * sin_val_x);
-    arr[y_index] = static_cast<scalar_t>(y * cos_val_y + x * sin_val_y);
-  }
-}
-
-template <typename scalar_t, bool interleaved>
-inline __device__ void apply_token_rotary_embedding_vec(
-    scalar_t* __restrict__ arr,            // [head_size]
-    const scalar_t* __restrict__ cos_ptr,  // [rot_dim]
-    const scalar_t* __restrict__ sin_ptr,  // [rot_dim]
-    int rot_offset,
-    int embed_dim) {
-  using vec_t = float4;
-  constexpr int kVecBytes = sizeof(vec_t);
-  constexpr int kScalarBytes = sizeof(scalar_t);
-  constexpr int kElePerVec = kVecBytes / kScalarBytes;
-
-  // Union for type punning to avoid strict aliasing issues with reinterpret_cast
-  union VecUnion {
-    vec_t vec;
-    scalar_t elems[kElePerVec];
-  };
-
-  if constexpr (interleaved) {
-    // Interleaved: arr has [x0, y0, x1, y1...]
-    // A single vector load contains 'kElePerVec' elements, which is 'kElePerVec / 2' pairs.
-    // rot_offset is the index of the PAIR.
-    // Address in arr is rot_offset * 2.
-
-    VecUnion data;
-    data.vec = *reinterpret_cast<const vec_t*>(arr + rot_offset * 2);
-
-    // Vectorized load for cos/sin: always 8 bytes (half of vector size, float2)
-    // float2 covers (kElePerVec / 2) scalar_t elements (pairs)
-    float2 cos_vec = *reinterpret_cast<const float2*>(cos_ptr + rot_offset);
-    float2 sin_vec = *reinterpret_cast<const float2*>(sin_ptr + rot_offset);
-
-    union CosSinVec {
-      float2 vec;
-      scalar_t elems[kElePerVec / 2];
-    };
-    CosSinVec cos_u, sin_u;
-    cos_u.vec = cos_vec;
-    sin_u.vec = sin_vec;
-
-#pragma unroll
-    for (int i = 0; i < kElePerVec; i += 2) {
-      // data.elems[i] is x, data.elems[i+1] is y
-      // They correspond to pair index: rot_offset + (i / 2)
-      int idx = i / 2;
-      float cos_val = static_cast<float>(cos_u.elems[idx]);
-      float sin_val = static_cast<float>(sin_u.elems[idx]);
-
-      float x = static_cast<float>(data.elems[i]);
-      float y = static_cast<float>(data.elems[i + 1]);
-
-      data.elems[i] = static_cast<scalar_t>(x * cos_val - y * sin_val);
-      data.elems[i + 1] = static_cast<scalar_t>(y * cos_val + x * sin_val);
-    }
-
-    *reinterpret_cast<vec_t*>(arr + rot_offset * 2) = data.vec;
-
-  } else {
-    // Non-interleaved: X and Y are separated by embed_dim.
-    // We process 'kElePerVec' PAIRS at once.
-    // Load X vector and Y vector.
-
-    VecUnion data_x, data_y;
-    data_x.vec = *reinterpret_cast<const vec_t*>(arr + rot_offset);
-    data_y.vec = *reinterpret_cast<const vec_t*>(arr + rot_offset + embed_dim);
-
-    // Vectorized load for cos/sin: always 16 bytes (full vector size, float4)
-    // float4 covers kElePerVec scalar_t elements
-    float4 cos_vec_x = *reinterpret_cast<const float4*>(cos_ptr + rot_offset);
-    float4 sin_vec_x = *reinterpret_cast<const float4*>(sin_ptr + rot_offset);
-    float4 cos_vec_y = *reinterpret_cast<const float4*>(cos_ptr + rot_offset + embed_dim);
-    float4 sin_vec_y = *reinterpret_cast<const float4*>(sin_ptr + rot_offset + embed_dim);
-
-    union CosSinVec {
-      float4 vec;
-      scalar_t elems[kElePerVec];
-    };
-    CosSinVec cos_u_x, sin_u_x, cos_u_y, sin_u_y;
-    cos_u_x.vec = cos_vec_x;
-    sin_u_x.vec = sin_vec_x;
-    cos_u_y.vec = cos_vec_y;
-    sin_u_y.vec = sin_vec_y;
-
-#pragma unroll
-    for (int i = 0; i < kElePerVec; ++i) {
-      float cos_val_x = static_cast<float>(cos_u_x.elems[i]);
-      float sin_val_x = static_cast<float>(sin_u_x.elems[i]);
-      float cos_val_y = static_cast<float>(cos_u_y.elems[i]);
-      float sin_val_y = static_cast<float>(sin_u_y.elems[i]);
-
-      float x = static_cast<float>(data_x.elems[i]);
-      float y = static_cast<float>(data_y.elems[i]);
-
-      data_x.elems[i] = static_cast<scalar_t>(x * cos_val_x - y * sin_val_x);
-      data_y.elems[i] = static_cast<scalar_t>(y * cos_val_y + x * sin_val_y);
-    }
-
-    *reinterpret_cast<vec_t*>(arr + rot_offset) = data_x.vec;
-    *reinterpret_cast<vec_t*>(arr + rot_offset + embed_dim) = data_y.vec;
-  }
-}
-
-template <>
-inline __device__ void apply_token_rotary_embedding<float, true>(
-    float* __restrict__ arr,
-    const float* __restrict__ cos_ptr,
-    const float* __restrict__ sin_ptr,
-    int rot_offset,
-    int /*embed_dim*/) {
-  float2 xy = *reinterpret_cast<const float2*>(arr + rot_offset * 2);
-
-  const float cos_val = static_cast<float>(cos_ptr[rot_offset]);
-  const float sin_val = static_cast<float>(sin_ptr[rot_offset]);
-
-  float2 out;
-  out.x = xy.x * cos_val - xy.y * sin_val;
-  out.y = xy.y * cos_val + xy.x * sin_val;
-
-  *reinterpret_cast<float2*>(arr + rot_offset * 2) = out;
-}
-
-// 2D grid kernel: parallel over tokens (grid.x) and pair-tiles (grid.y)
-template <typename scalar_t, bool interleaved, bool vectorized, int ROT_EMBED_DIM>
-__global__ void rotary_embedding_kernel_2d(
-    const scalar_t* __restrict__ cos_data,  // [num_tokens, rot_dim_arg]
-    const scalar_t* __restrict__ sin_data,  // [num_tokens, rot_dim_arg]
-    scalar_t* __restrict__ query_total,
-    scalar_t* __restrict__ key_total,
-    const int rot_dim_arg,
-    const int embed_dim_for_rotation_arg,
-    const int64_t query_token_stride,
-    const int64_t key_token_stride,
-    const int64_t head_stride_query,
-    const int64_t head_stride_key,
-    const int num_heads,
-    const int num_kv_heads,
-    const int head_size_arg,
-    const int blocks_per_token) {
-  // Removed shared memory allocation and loading logic
-
-  const int token_idx = blockIdx.x;
-  if (token_idx >= gridDim.x) {
-    return;
-  }
-
-  // Pointers to Global Memory for the current token
-  const scalar_t* current_token_cos_ptr = cos_data + token_idx * rot_dim_arg;
-  const scalar_t* current_token_sin_ptr = sin_data + token_idx * rot_dim_arg;
-
-  constexpr int kVecBytes = 16;
-  const int scalar_size = sizeof(scalar_t);
-
-  // Removed SMEM loading and __syncthreads()
-
-  scalar_t* query_for_token = query_total + token_idx * (int)query_token_stride;
-  scalar_t* key_for_token = (key_total != nullptr) ? (key_total + token_idx * (int)key_token_stride) : nullptr;
-
-  const int local_block_idx = blockIdx.y;
-
-  // Use compile-time constant if ROT_EMBED_DIM > 0, otherwise fallback to runtime argument
-  const int embed_dim_for_rotation = (ROT_EMBED_DIM > 0) ? ROT_EMBED_DIM : embed_dim_for_rotation_arg;
-
-  if constexpr (vectorized) {
-    using vec_t = float4;
-    constexpr int kElePerVec = kVecBytes / sizeof(scalar_t);
-    constexpr int pairs_per_step = interleaved ? (kElePerVec / 2) : kElePerVec;
-
-    const int pair_stride = blockDim.x * blocks_per_token * pairs_per_step;
-    int i = (local_block_idx * blockDim.x + threadIdx.x) * pairs_per_step;
-    const int nq_pairs = num_heads * embed_dim_for_rotation;
-
-    // PROLOGUE: Preload 0-th batch
-    RotaryVecData<scalar_t, interleaved> curr_data;
-    if (i < nq_pairs) {
-      int head_idx = i / embed_dim_for_rotation;
-      int rot_offset = i % embed_dim_for_rotation;
-      scalar_t* ptr = query_for_token + head_idx * (int)head_stride_query;
-      curr_data = load_rotary_vec<scalar_t, interleaved>(
-          ptr, current_token_cos_ptr, current_token_sin_ptr, rot_offset, embed_dim_for_rotation);
-    }
-
-    int next_i = i + pair_stride;
-
-    // MAIN LOOP
-    for (; i < nq_pairs; i += pair_stride, next_i += pair_stride) {
-      // 1. PREFETCH NEXT
-      RotaryVecData<scalar_t, interleaved> next_data;
-      bool active_next = (next_i < nq_pairs);
-      if (active_next) {
-        int head_idx_next = next_i / embed_dim_for_rotation;
-        int rot_offset_next = next_i % embed_dim_for_rotation;
-        scalar_t* ptr_next = query_for_token + head_idx_next * (int)head_stride_query;
-        next_data = load_rotary_vec<scalar_t, interleaved>(
-            ptr_next, current_token_cos_ptr, current_token_sin_ptr, rot_offset_next, embed_dim_for_rotation);
-      }
-
-      // 2. COMPUTE CURRENT
-      // The compute uses curr_data which corresponds to index 'i'
-      int head_idx = i / embed_dim_for_rotation;
-      int rot_offset = i % embed_dim_for_rotation;
-      scalar_t* ptr = query_for_token + head_idx * (int)head_stride_query;
-      compute_store_rotary_vec<scalar_t, interleaved>(ptr, curr_data, rot_offset, embed_dim_for_rotation);
-
-      // 3. SHIFT
-      curr_data = next_data;
-    }
-
-    if (key_for_token != nullptr) {
-      const int nk_pairs = num_kv_heads * embed_dim_for_rotation;
-      int k_i = (local_block_idx * blockDim.x + threadIdx.x) * pairs_per_step;
-
-      // PROLOGUE (Key)
-      RotaryVecData<scalar_t, interleaved> curr_data_k;
-      if (k_i < nk_pairs) {
-        int head_idx = k_i / embed_dim_for_rotation;
-        int rot_offset = k_i % embed_dim_for_rotation;
-        scalar_t* ptr = key_for_token + head_idx * (int)head_stride_key;
-        curr_data_k = load_rotary_vec<scalar_t, interleaved>(
-            ptr, current_token_cos_ptr, current_token_sin_ptr, rot_offset, embed_dim_for_rotation);
-      }
-
-      int next_k_i = k_i + pair_stride;
-
-      // MAIN LOOP (Key)
-      for (; k_i < nk_pairs; k_i += pair_stride, next_k_i += pair_stride) {
-        // 1. PREFETCH NEXT
-        RotaryVecData<scalar_t, interleaved> next_data_k;
-        bool active_next = (next_k_i < nk_pairs);
-        if (active_next) {
-          int head_idx_next = next_k_i / embed_dim_for_rotation;
-          int rot_offset_next = next_k_i % embed_dim_for_rotation;
-          scalar_t* ptr_next = key_for_token + head_idx_next * (int)head_stride_key;
-          next_data_k = load_rotary_vec<scalar_t, interleaved>(
-              ptr_next, current_token_cos_ptr, current_token_sin_ptr, rot_offset_next, embed_dim_for_rotation);
-        }
-
-        // 2. COMPUTE CURRENT
-        int head_idx = k_i / embed_dim_for_rotation;
-        int rot_offset = k_i % embed_dim_for_rotation;
-        scalar_t* ptr = key_for_token + head_idx * (int)head_stride_key;
-        compute_store_rotary_vec<scalar_t, interleaved>(ptr, curr_data_k, rot_offset, embed_dim_for_rotation);
-
-        // 3. SHIFT
-        curr_data_k = next_data_k;
-      }
-    }
-  } else {
-    // Fallback to scalar implementation
-    const int pair_stride = blockDim.x * blocks_per_token;
-    const int thread_pair_offset = local_block_idx * blockDim.x + threadIdx.x;
-
-    const int nq_pairs = num_heads * embed_dim_for_rotation;
-    for (int i = thread_pair_offset; i < nq_pairs; i += pair_stride) {
-      const int head_idx = i / embed_dim_for_rotation;
-      const int rot_offset = i % embed_dim_for_rotation;
-
-      scalar_t* query_for_token_head = query_for_token + head_idx * (int)head_stride_query;
-      apply_token_rotary_embedding<scalar_t, interleaved>(
-          query_for_token_head, current_token_cos_ptr, current_token_sin_ptr, rot_offset, embed_dim_for_rotation);
-    }
-
-    if (key_for_token != nullptr) {
-      const int nk_pairs = num_kv_heads * embed_dim_for_rotation;
-      for (int i = thread_pair_offset; i < nk_pairs; i += pair_stride) {
-        const int head_idx = i / embed_dim_for_rotation;
-        const int rot_offset = i % embed_dim_for_rotation;
-
-        scalar_t* key_for_token_head = key_for_token + head_idx * (int)head_stride_key;
-        apply_token_rotary_embedding<scalar_t, interleaved>(
-            key_for_token_head, current_token_cos_ptr, current_token_sin_ptr, rot_offset, embed_dim_for_rotation);
-      }
-    }
-  }
-}
-
-// 1D grid kernel: each block handles one token
-template <typename scalar_t, bool interleaved, bool vectorized, int ROT_EMBED_DIM>
-__launch_bounds__(512) __global__ void rotary_embedding_kernel_1d(
-    const scalar_t* __restrict__ cos_data,  // [num_tokens, rot_dim_arg]
-    const scalar_t* __restrict__ sin_data,  // [num_tokens, rot_dim_arg]
-    scalar_t* __restrict__ query_total,
-    scalar_t* __restrict__ key_total,
-    const int rot_dim_arg,
-    const int embed_dim_for_rotation_arg,
-    const int64_t query_token_stride,
-    const int64_t key_token_stride,
-    const int64_t head_stride_query,
-    const int64_t head_stride_key,
-    const int num_heads,
-    const int num_kv_heads,
-    const int head_size_arg) {
-  // Removed shared memory allocation and loading logic
-
-  const int token_idx = blockIdx.x;
-  if (token_idx >= gridDim.x) return;
-
-  const scalar_t* current_token_cos_ptr = cos_data + token_idx * rot_dim_arg;
-  const scalar_t* current_token_sin_ptr = sin_data + token_idx * rot_dim_arg;
-
-  constexpr int kVecBytes = 16;
-  const int scalar_size = sizeof(scalar_t);
-
-  // Removed SMEM loading and __syncthreads()
-
-  scalar_t* query_for_token = query_total + token_idx * (int)query_token_stride;
-  scalar_t* key_for_token = (key_total != nullptr) ? (key_total + token_idx * (int)key_token_stride) : nullptr;
-
-  // Use compile-time constant if ROT_EMBED_DIM > 0, otherwise fallback to runtime argument
-  const int embed_dim_for_rotation = (ROT_EMBED_DIM > 0) ? ROT_EMBED_DIM : embed_dim_for_rotation_arg;
-
-  if constexpr (vectorized) {
-    constexpr int kVecBytes = 16;
-    constexpr int kElePerVec = kVecBytes / sizeof(scalar_t);
-    // Interleaved: 1 vec = kElePerVec elements = kElePerVec/2 pairs
-    // Non-Interleaved: 1 vec (X) + 1 vec (Y) = kElePerVec elements = kElePerVec pairs
-    constexpr int pairs_per_step = interleaved ? (kElePerVec / 2) : kElePerVec;
-    const int nq_pairs = num_heads * embed_dim_for_rotation;
-
-    // Original Stride
-    const int stride = blockDim.x * pairs_per_step;
-
-    int i = threadIdx.x * pairs_per_step;
-
-    // ==========================================
-    // PROLOGUE: Preload first batch
-    // ==========================================
-    RotaryVecData<scalar_t, interleaved> curr_data;
-
-    if (i < nq_pairs) {
-      int head_idx = i / embed_dim_for_rotation;
-      int rot_offset = i % embed_dim_for_rotation;
-      scalar_t* ptr = query_for_token + head_idx * (int)head_stride_query;
-      curr_data = load_rotary_vec<scalar_t, interleaved>(
-          ptr, current_token_cos_ptr, current_token_sin_ptr, rot_offset, embed_dim_for_rotation);
-    }
-
-    int next_i = i + stride;
-
-    // ==========================================
-    // MAIN LOOP
-    // ==========================================
-    for (; i < nq_pairs; i += stride, next_i += stride) {
-      // --- 1. PREFETCH NEXT BATCH ---
-      RotaryVecData<scalar_t, interleaved> next_data;
-
-      if (next_i < nq_pairs) {
-        int head_idx = next_i / embed_dim_for_rotation;
-        int rot_offset = next_i % embed_dim_for_rotation;
-        scalar_t* ptr = query_for_token + head_idx * (int)head_stride_query;
-        next_data = load_rotary_vec<scalar_t, interleaved>(
-            ptr, current_token_cos_ptr, current_token_sin_ptr, rot_offset, embed_dim_for_rotation);
-      }
-
-      // --- 2. COMPUTE CURRENT BATCH ---
-      if (i < nq_pairs) {
-        int head_idx = i / embed_dim_for_rotation;
-        int rot_offset = i % embed_dim_for_rotation;
-        scalar_t* ptr = query_for_token + head_idx * (int)head_stride_query;
-        compute_store_rotary_vec<scalar_t, interleaved>(ptr, curr_data, rot_offset, embed_dim_for_rotation);
-      }
-
-      // --- 3. SHIFT ---
-      curr_data = next_data;
-    }
-
-    // Key branch
-    if (key_for_token != nullptr) {
-      const int nk_pairs = num_kv_heads * embed_dim_for_rotation;
-      int k_i = threadIdx.x * pairs_per_step;
-
-      // PROLOGUE (Key): Preload first batch
-      RotaryVecData<scalar_t, interleaved> curr_data_k;
-
-      if (k_i < nk_pairs) {
-        int head_idx = k_i / embed_dim_for_rotation;
-        int rot_offset = k_i % embed_dim_for_rotation;
-        scalar_t* ptr = key_for_token + head_idx * (int)head_stride_key;
-        curr_data_k = load_rotary_vec<scalar_t, interleaved>(
-            ptr, current_token_cos_ptr, current_token_sin_ptr, rot_offset, embed_dim_for_rotation);
-      }
-
-      int next_k_i = k_i + stride;
-
-      // MAIN LOOP (Key)
-      for (; k_i < nk_pairs; k_i += stride, next_k_i += stride) {
-        // --- 1. PREFETCH NEXT BATCH ---
-        RotaryVecData<scalar_t, interleaved> next_data_k;
-
-        if (next_k_i < nk_pairs) {
-          int head_idx = next_k_i / embed_dim_for_rotation;
-          int rot_offset = next_k_i % embed_dim_for_rotation;
-          scalar_t* ptr = key_for_token + head_idx * (int)head_stride_key;
-          next_data_k = load_rotary_vec<scalar_t, interleaved>(
-              ptr, current_token_cos_ptr, current_token_sin_ptr, rot_offset, embed_dim_for_rotation);
-        }
-
-        // --- 2. COMPUTE CURRENT BATCH ---
-        if (k_i < nk_pairs) {
-          int head_idx = k_i / embed_dim_for_rotation;
-          int rot_offset = k_i % embed_dim_for_rotation;
-          scalar_t* ptr = key_for_token + head_idx * (int)head_stride_key;
-          compute_store_rotary_vec<scalar_t, interleaved>(ptr, curr_data_k, rot_offset, embed_dim_for_rotation);
-        }
-
-        // --- 3. SHIFT ---
-        curr_data_k = next_data_k;
-      }
-    }
-  } else {
-    // Fallback scalar
-    const int nq_pairs = num_heads * embed_dim_for_rotation;
-    for (int i = threadIdx.x; i < nq_pairs; i += blockDim.x) {
-      const int head_idx = i / embed_dim_for_rotation;
-      const int rot_offset = i % embed_dim_for_rotation;
-      scalar_t* query_for_token_head = query_for_token + head_idx * (int)head_stride_query;
-
-      apply_token_rotary_embedding<scalar_t, interleaved>(
-          query_for_token_head, current_token_cos_ptr, current_token_sin_ptr, rot_offset, embed_dim_for_rotation);
-    }
-
-    if (key_for_token != nullptr) {
-      const int nk_pairs = num_kv_heads * embed_dim_for_rotation;
-      for (int i = threadIdx.x; i < nk_pairs; i += blockDim.x) {
-        const int head_idx = i / embed_dim_for_rotation;
-        const int rot_offset = i % embed_dim_for_rotation;
-
-        scalar_t* key_for_token_head = key_for_token + head_idx * (int)head_stride_key;
-        apply_token_rotary_embedding<scalar_t, interleaved>(
-            key_for_token_head, current_token_cos_ptr, current_token_sin_ptr, rot_offset, embed_dim_for_rotation);
-      }
-    }
-  }
-}
-
-// Define dispatch macro for Rotation Dimension (Pair Count)
-// Note: This is usually HEAD_SIZE / 2 for LLaMA-like models.
-// Case 32 -> head_size 64
-// Case 64 -> head_size 128
-// Case 128 -> head_size 256
-#define DISPATCH_ROT_DIM(embed_dim, ...)  \
-  switch (embed_dim) {                    \
-    case 16: {                            \
-      constexpr int ROT_EMBED_DIM = 16;   \
-      __VA_ARGS__;                        \
-      break;                              \
-    }                                     \
-    case 32: {                            \
-      constexpr int ROT_EMBED_DIM = 32;   \
-      __VA_ARGS__;                        \
-      break;                              \
-    }                                     \
-    case 64: {                            \
-      constexpr int ROT_EMBED_DIM = 64;   \
-      __VA_ARGS__;                        \
-      break;                              \
-    }                                     \
-    case 96: {                            \
-      constexpr int ROT_EMBED_DIM = 96;   \
-      __VA_ARGS__;                        \
-      break;                              \
-    }                                     \
-    case 128: {                           \
-      constexpr int ROT_EMBED_DIM = 128;  \
-      __VA_ARGS__;                        \
-      break;                              \
-    }                                     \
-    case 256: {                           \
-      constexpr int ROT_EMBED_DIM = 256;  \
-      __VA_ARGS__;                        \
-      break;                              \
-    }                                     \
-    case 512: {                           \
-      constexpr int ROT_EMBED_DIM = 512;  \
-      __VA_ARGS__;                        \
-      break;                              \
-    }                                     \
-    case 1024: {                          \
-      constexpr int ROT_EMBED_DIM = 1024; \
-      __VA_ARGS__;                        \
-      break;                              \
-    }                                     \
-    default: {                            \
-      constexpr int ROT_EMBED_DIM = 0;    \
-      __VA_ARGS__;                        \
-      break;                              \
-    }                                     \
-  }
-
-void rotary_embedding_cos_sin(
-    at::Tensor& cos,
-    at::Tensor& sin,
-    at::Tensor& query,
-    const std::optional<at::Tensor>& key,
-    int64_t head_size,
-    bool interleaved) {
-  TORCH_CHECK(
-      query.dim() == 2 || query.dim() == 3,
-      "query must be in shape [num_tokens, hidden_size] or [num_tokens, num_heads, head_size]");
-  if (key.has_value()) {
-    TORCH_CHECK(
-        key->dim() == 2 || key->dim() == 3,
-        "key must be in shape [num_tokens, hidden_size] or [num_tokens, num_kv_heads, head_size]");
-  }
-
-  const int64_t num_tokens = query.size(0);
-
-  TORCH_CHECK(cos.dim() == 2, "cos must be in shape [num_tokens, D_cos]");
-  TORCH_CHECK(sin.dim() == 2, "sin must be in shape [num_tokens, D_sin]");
-  TORCH_CHECK(cos.size(0) == num_tokens, "cos num_tokens mismatch with query");
-  TORCH_CHECK(sin.size(0) == num_tokens, "sin num_tokens mismatch with query");
-  TORCH_CHECK(cos.size(1) == sin.size(1), "cos and sin D_cos/D_sin mismatch");
-
-  TORCH_CHECK(cos.scalar_type() == query.scalar_type(), "cos dtype mismatch with query");
-  TORCH_CHECK(sin.scalar_type() == query.scalar_type(), "sin dtype mismatch with query");
-  TORCH_CHECK(cos.is_cuda() && sin.is_cuda() && query.is_cuda(), "cos/sin/query must be CUDA tensors");
-  TORCH_CHECK(query.is_contiguous(), "query must be contiguous");
-  if (key.has_value()) {
-    TORCH_CHECK(key->is_cuda(), "key must be CUDA tensor if provided");
-    TORCH_CHECK(key->scalar_type() == query.scalar_type(), "key dtype mismatch with query");
-    TORCH_CHECK(key->is_contiguous(), "key must be contiguous");
-  }
-
-  int query_hidden_size_calculated;
-  if (query.dim() == 2) {
-    query_hidden_size_calculated = (int)query.size(1);
-  } else {
-    query_hidden_size_calculated = (int)query.size(1) * (int)query.size(2);
-    TORCH_CHECK(query.size(2) == head_size, "query head_size mismatch in 3D tensor");
-  }
-  TORCH_CHECK(query_hidden_size_calculated % head_size == 0, "query_hidden_size not divisible by head_size");
-  int num_heads = query_hidden_size_calculated / (int)head_size;
-
-  int key_hidden_size_calculated = 0;
-  int num_kv_heads = num_heads;
-  if (key.has_value()) {
-    TORCH_CHECK((int)key->size(0) == num_tokens, "key num_tokens mismatch");
-    if (key->dim() == 2) {
-      key_hidden_size_calculated = (int)key->size(1);
-    } else {
-      key_hidden_size_calculated = (int)key->size(1) * (int)key->size(2);
-      TORCH_CHECK((int)key->size(2) == head_size, "key head_size mismatch in 3D tensor");
-    }
-    TORCH_CHECK(key_hidden_size_calculated % head_size == 0, "key_hidden_size not divisible by head_size");
-    num_kv_heads = key_hidden_size_calculated / (int)head_size;
-  }
-  TORCH_CHECK(num_heads % num_kv_heads == 0, "num_heads must be divisible by num_kv_heads");
-
-  // NeoX interleaved: if cos/sin are full-dim (head_size), downsample once.
-  if (interleaved && cos.size(1) == head_size) {
-    TORCH_CHECK(head_size % 2 == 0, "interleaved layout requires even head_size");
-    const int64_t half = head_size / 2;
-    std::vector<int64_t> new_shape = {cos.size(0), half, 2};
-    cos = cos.view(new_shape).select(2, 0).contiguous();
-    sin = sin.view(new_shape).select(2, 1).contiguous();
-  }
-
-  const int rot_dim_from_cache = (int)cos.size(1);
-
-  const int64_t query_token_stride = query_hidden_size_calculated;
-  const int64_t key_token_stride = key.has_value() ? key_hidden_size_calculated : 0;
-
-  int64_t head_stride_query;
-  if (query.dim() == 3 && query.size(1) == num_heads && query.size(2) == head_size) {
-    head_stride_query = query.stride(1);
-  } else {
-    head_stride_query = head_size;
-  }
-
-  int64_t head_stride_key = head_size;
-  if (key.has_value()) {
-    if (key->dim() == 3 && key->size(1) == num_kv_heads && key->size(2) == head_size) {
-      head_stride_key = key->stride(1);
-    } else {
-      head_stride_key = head_size;
-    }
-  }
-
-  const int embed_dim_for_rotation = interleaved ? rot_dim_from_cache : (rot_dim_from_cache / 2);
-  TORCH_CHECK(embed_dim_for_rotation > 0, "embed_dim_for_rotation must be > 0");
-
-  const int max_pairs_to_rotate_per_token =
-      std::max(num_heads * embed_dim_for_rotation, num_kv_heads * embed_dim_for_rotation);
-
-  const at::cuda::OptionalCUDAGuard device_guard(device_of(query));
-  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
-
-  AT_DISPATCH_FLOATING_TYPES_AND2(
-      at::ScalarType::Half, at::ScalarType::BFloat16, query.scalar_type(), "rotary_embedding_cos_sin", [&] {
-        using torch_scalar_t = scalar_t;
-        using cuda_scalar_t = typename std::conditional<
-            std::is_same<torch_scalar_t, at::Half>::value,
-            nv_half,
-            typename std::conditional<std::is_same<torch_scalar_t, at::BFloat16>::value, nv_bfloat16, torch_scalar_t>::
-                type>::type;
-
-        // Constants for vectorization
-        constexpr int kVecBytes = 16;
-        constexpr int kElePerVec = kVecBytes / sizeof(torch_scalar_t);
-
-        // Determine how many pairs one thread handles in one vector step
-        // Interleaved: 1 vector load (16B) contains 'kElePerVec' elements -> 'kElePerVec / 2' pairs.
-        // Non-interleaved: 2 vector loads (32B) contains 'kElePerVec' pairs (X vec + Y vec).
-        const int pairs_per_step = interleaved ? (kElePerVec / 2) : kElePerVec;
-
-        // Check if we can guarantee vectorization for ALL tokens
-        // We need Base Pointers, Strides, and Dimension to be aligned.
-        bool can_vectorize_all = true;
-
-        // 1. Check Dimensions
-        if (embed_dim_for_rotation % pairs_per_step != 0) can_vectorize_all = false;
-
-        // 2. Check Base Pointers
-        if (reinterpret_cast<uintptr_t>(query.data_ptr()) % kVecBytes != 0) can_vectorize_all = false;
-        if (reinterpret_cast<uintptr_t>(cos.data_ptr()) % kVecBytes != 0) can_vectorize_all = false;
-        if (reinterpret_cast<uintptr_t>(sin.data_ptr()) % kVecBytes != 0) can_vectorize_all = false;
-        if (key.has_value()) {
-          if (reinterpret_cast<uintptr_t>(key->data_ptr()) % kVecBytes != 0) can_vectorize_all = false;
-        }
-
-        // 3. Check Strides
-        // We need the stride between tokens to be a multiple of vector size
-        // to ensure that if token 0 is aligned, token 1 is also aligned.
-        if (query_token_stride * sizeof(torch_scalar_t) % kVecBytes != 0) can_vectorize_all = false;
-        if (head_stride_query * sizeof(torch_scalar_t) % kVecBytes != 0) can_vectorize_all = false;
-
-        if (key.has_value()) {
-          if (key_token_stride * sizeof(torch_scalar_t) % kVecBytes != 0) can_vectorize_all = false;
-          if (head_stride_key * sizeof(torch_scalar_t) % kVecBytes != 0) can_vectorize_all = false;
-        }
-
-        // Determine launch configuration
-        // If we can vectorize all, each thread handles 'pairs_per_step' pairs.
-        // Otherwise, fallback to conservative estimate (1 pair per thread) to ensure enough blocks.
-        const int launch_pairs_per_thread = can_vectorize_all ? pairs_per_step : 1;
-
-        const int total_threads_needed =
-            (max_pairs_to_rotate_per_token + launch_pairs_per_thread - 1) / launch_pairs_per_thread;
-
-        // Case 1: 2D Grid (Split one token across multiple blocks)
-        // Keep block size moderate (128-256) aligned with head size
-        const int threads_per_block_2d = std::min<int>(256, std::max(128, embed_dim_for_rotation));
-        const int blocks_per_token_2d = (total_threads_needed + threads_per_block_2d - 1) / threads_per_block_2d;
-
-        // Decide grid strategy
-        const bool use_grid_2d = (num_tokens <= 4) && (blocks_per_token_2d > 1);
-
-        // Case 2: 1D Grid (One block per token)
-        // Maximize threads per block to cover all heads in one block if possible
-        // Cap at 512 threads to balance occupancy and register usage
-        const int threads_per_block_1d = std::min<int>(512, std::max(128, (total_threads_needed + 31) / 32 * 32));
-
-        // Final launch config
-        const int threads_per_block = use_grid_2d ? threads_per_block_2d : threads_per_block_1d;
-        const int blocks_per_token = use_grid_2d ? blocks_per_token_2d : 1;
-
-        dim3 block(threads_per_block);
-        dim3 grid_2d((int)num_tokens, std::max(1, blocks_per_token));
-        dim3 grid_1d((int)num_tokens);
-
-        // No shared memory needed
-        size_t smem_size = 0;
-
-        DISPATCH_ROT_DIM(
-            embed_dim_for_rotation,
-            auto launch_kernel =
-                [&](bool vectorized) {
-                  if (interleaved) {
-                    if (use_grid_2d) {
-                      if (vectorized) {
-                        rotary_embedding_kernel_2d<cuda_scalar_t, true, true, ROT_EMBED_DIM>
-                            <<<grid_2d, block, smem_size, stream>>>(
-                                reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
-                                reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
-                                reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
-                                key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>())
-                                                : nullptr,
-                                rot_dim_from_cache,
-                                embed_dim_for_rotation,
-                                query_token_stride,
-                                key_token_stride,
-                                head_stride_query,
-                                head_stride_key,
-                                num_heads,
-                                num_kv_heads,
-                                (int)head_size,
-                                blocks_per_token);
-                      } else {
-                        rotary_embedding_kernel_2d<cuda_scalar_t, true, false, ROT_EMBED_DIM>
-                            <<<grid_2d, block, smem_size, stream>>>(
-                                reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
-                                reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
-                                reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
-                                key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>())
-                                                : nullptr,
-                                rot_dim_from_cache,
-                                embed_dim_for_rotation,
-                                query_token_stride,
-                                key_token_stride,
-                                head_stride_query,
-                                head_stride_key,
-                                num_heads,
-                                num_kv_heads,
-                                (int)head_size,
-                                blocks_per_token);
-                      }
-                    } else {
-                      if (vectorized) {
-                        rotary_embedding_kernel_1d<cuda_scalar_t, true, true, ROT_EMBED_DIM>
-                            <<<grid_1d, block, smem_size, stream>>>(
-                                reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
-                                reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
-                                reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
-                                key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>())
-                                                : nullptr,
-                                rot_dim_from_cache,
-                                embed_dim_for_rotation,
-                                query_token_stride,
-                                key_token_stride,
-                                head_stride_query,
-                                head_stride_key,
-                                num_heads,
-                                num_kv_heads,
-                                (int)head_size);
-                      } else {
-                        rotary_embedding_kernel_1d<cuda_scalar_t, true, false, ROT_EMBED_DIM>
-                            <<<grid_1d, block, smem_size, stream>>>(
-                                reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
-                                reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
-                                reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
-                                key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>())
-                                                : nullptr,
-                                rot_dim_from_cache,
-                                embed_dim_for_rotation,
-                                query_token_stride,
-                                key_token_stride,
-                                head_stride_query,
-                                head_stride_key,
-                                num_heads,
-                                num_kv_heads,
-                                (int)head_size);
-                      }
-                    }
-                  } else {  // non-interleaved
-                    if (use_grid_2d) {
-                      if (vectorized) {
-                        rotary_embedding_kernel_2d<cuda_scalar_t, false, true, ROT_EMBED_DIM>
-                            <<<grid_2d, block, smem_size, stream>>>(
-                                reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
-                                reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
-                                reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
-                                key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>())
-                                                : nullptr,
-                                rot_dim_from_cache,
-                                embed_dim_for_rotation,
-                                query_token_stride,
-                                key_token_stride,
-                                head_stride_query,
-                                head_stride_key,
-                                num_heads,
-                                num_kv_heads,
-                                (int)head_size,
-                                blocks_per_token);
-                      } else {
-                        rotary_embedding_kernel_2d<cuda_scalar_t, false, false, ROT_EMBED_DIM>
-                            <<<grid_2d, block, smem_size, stream>>>(
-                                reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
-                                reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
-                                reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
-                                key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>())
-                                                : nullptr,
-                                rot_dim_from_cache,
-                                embed_dim_for_rotation,
-                                query_token_stride,
-                                key_token_stride,
-                                head_stride_query,
-                                head_stride_key,
-                                num_heads,
-                                num_kv_heads,
-                                (int)head_size,
-                                blocks_per_token);
-                      }
-                    } else {
-                      if (vectorized) {
-                        rotary_embedding_kernel_1d<cuda_scalar_t, false, true, ROT_EMBED_DIM>
-                            <<<grid_1d, block, smem_size, stream>>>(
-                                reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
-                                reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
-                                reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
-                                key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>())
-                                                : nullptr,
-                                rot_dim_from_cache,
-                                embed_dim_for_rotation,
-                                query_token_stride,
-                                key_token_stride,
-                                head_stride_query,
-                                head_stride_key,
-                                num_heads,
-                                num_kv_heads,
-                                (int)head_size);
-                      } else {
-                        rotary_embedding_kernel_1d<cuda_scalar_t, false, false, ROT_EMBED_DIM>
-                            <<<grid_1d, block, smem_size, stream>>>(
-                                reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
-                                reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
-                                reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
-                                key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>())
-                                                : nullptr,
-                                rot_dim_from_cache,
-                                embed_dim_for_rotation,
-                                query_token_stride,
-                                key_token_stride,
-                                head_stride_query,
-                                head_stride_key,
-                                num_heads,
-                                num_kv_heads,
-                                (int)head_size);
-                      }
-                    }
-                  }
-                };
-
-            // Close the log file after all launches
-            if (can_vectorize_all) { launch_kernel(true); } else { launch_kernel(false); });
-      });
-
-  C10_CUDA_KERNEL_LAUNCH_CHECK();
-}
+ #include <ATen/cuda/CUDAContext.h>
+ #include <c10/cuda/CUDAException.h>
+ #include <c10/cuda/CUDAGuard.h>
+ #include <c10/cuda/CUDAStream.h>
+ #include <torch/all.h>
+ 
+ #include <cmath>
+ #include <vector>
+ 
+ #include "utils.h"
+ 
+ // -------------------------------------------------------------------------
+ // Helper structures and functions for manual pipeline interleaving
+ // -------------------------------------------------------------------------
+ 
+ template <typename scalar_t, bool interleaved>
+ struct RotaryVecData;
+ 
+ template <typename scalar_t>
+ struct RotaryVecData<scalar_t, true> {
+   float4 vec;
+   float2 cos_vec;
+   float2 sin_vec;
+ };
+ 
+ template <typename scalar_t>
+ struct RotaryVecData<scalar_t, false> {
+   float4 vec_x;
+   float4 vec_y;
+   float4 cos_x;
+   float4 sin_x;
+   float4 cos_y;
+   float4 sin_y;
+ };
+ 
+ template <typename scalar_t, bool interleaved>
+ inline __device__ RotaryVecData<scalar_t, interleaved> load_rotary_vec(
+     const scalar_t* __restrict__ arr,
+     const scalar_t* __restrict__ cos_ptr,
+     const scalar_t* __restrict__ sin_ptr,
+     int rot_offset,
+     int embed_dim) {
+   RotaryVecData<scalar_t, interleaved> data;
+   if constexpr (interleaved) {
+     data.vec = *reinterpret_cast<const float4*>(arr + rot_offset * 2);
+     data.cos_vec = *reinterpret_cast<const float2*>(cos_ptr + rot_offset);
+     data.sin_vec = *reinterpret_cast<const float2*>(sin_ptr + rot_offset);
+   } else {
+     data.vec_x = *reinterpret_cast<const float4*>(arr + rot_offset);
+     data.vec_y = *reinterpret_cast<const float4*>(arr + rot_offset + embed_dim);
+     data.cos_x = *reinterpret_cast<const float4*>(cos_ptr + rot_offset);
+     data.sin_x = *reinterpret_cast<const float4*>(sin_ptr + rot_offset);
+     data.cos_y = *reinterpret_cast<const float4*>(cos_ptr + rot_offset + embed_dim);
+     data.sin_y = *reinterpret_cast<const float4*>(sin_ptr + rot_offset + embed_dim);
+   }
+   return data;
+ }
+ 
+ template <typename scalar_t, bool interleaved>
+ inline __device__ void compute_store_rotary_vec(
+     scalar_t* __restrict__ arr, const RotaryVecData<scalar_t, interleaved>& data, int rot_offset, int embed_dim) {
+   constexpr int kVecBytes = 16;
+   constexpr int kElePerVec = kVecBytes / sizeof(scalar_t);
+ 
+   if constexpr (interleaved) {
+     union VecUnion {
+       float4 vec;
+       scalar_t elems[kElePerVec];
+     };
+     union CosSinVec {
+       float2 vec;
+       scalar_t elems[kElePerVec / 2];
+     };
+ 
+     VecUnion v;
+     v.vec = data.vec;
+     CosSinVec c, s;
+     c.vec = data.cos_vec;
+     s.vec = data.sin_vec;
+ 
+ #pragma unroll
+     for (int i = 0; i < kElePerVec; i += 2) {
+       int idx = i / 2;
+       float cos_val = static_cast<float>(c.elems[idx]);
+       float sin_val = static_cast<float>(s.elems[idx]);
+       float x = static_cast<float>(v.elems[i]);
+       float y = static_cast<float>(v.elems[i + 1]);
+       v.elems[i] = static_cast<scalar_t>(x * cos_val - y * sin_val);
+       v.elems[i + 1] = static_cast<scalar_t>(y * cos_val + x * sin_val);
+     }
+     *reinterpret_cast<float4*>(arr + rot_offset * 2) = v.vec;
+ 
+   } else {
+     union VecUnion {
+       float4 vec;
+       scalar_t elems[kElePerVec];
+     };
+     union CosSinVec {
+       float4 vec;
+       scalar_t elems[kElePerVec];
+     };
+ 
+     VecUnion vx, vy;
+     vx.vec = data.vec_x;
+     vy.vec = data.vec_y;
+     CosSinVec cx, sx, cy, sy;
+     cx.vec = data.cos_x;
+     sx.vec = data.sin_x;
+     cy.vec = data.cos_y;
+     sy.vec = data.sin_y;
+ 
+ #pragma unroll
+     for (int i = 0; i < kElePerVec; ++i) {
+       float cos_val_x = static_cast<float>(cx.elems[i]);
+       float sin_val_x = static_cast<float>(sx.elems[i]);
+       float cos_val_y = static_cast<float>(cy.elems[i]);
+       float sin_val_y = static_cast<float>(sy.elems[i]);
+       float x = static_cast<float>(vx.elems[i]);
+       float y = static_cast<float>(vy.elems[i]);
+ 
+       vx.elems[i] = static_cast<scalar_t>(x * cos_val_x - y * sin_val_x);
+       vy.elems[i] = static_cast<scalar_t>(y * cos_val_y + x * sin_val_y);
+     }
+     *reinterpret_cast<float4*>(arr + rot_offset) = vx.vec;
+     *reinterpret_cast<float4*>(arr + rot_offset + embed_dim) = vy.vec;
+   }
+ }
+ 
+ template <typename scalar_t, bool interleaved>
+ inline __device__ void apply_token_rotary_embedding(
+     scalar_t* __restrict__ arr,            // [head_size]
+     const scalar_t* __restrict__ cos_ptr,  // [rot_dim]
+     const scalar_t* __restrict__ sin_ptr,  // [rot_dim]
+     int rot_offset,
+     int embed_dim) {  // for non-interleaved: half dim
+   if constexpr (interleaved) {
+     // NeoX-style: interleaved layout [x0, y0, x1, y1, ...].
+     // cos/sin: [..., rotary_dim/2], one entry per (x, y) pair.
+     const int x_index = 2 * rot_offset;
+     const int y_index = x_index + 1;
+ 
+     // Directly access Global Memory
+     const float cos_val = static_cast<float>(cos_ptr[rot_offset]);
+     const float sin_val = static_cast<float>(sin_ptr[rot_offset]);
+ 
+     const float x = static_cast<float>(arr[x_index]);
+     const float y = static_cast<float>(arr[y_index]);
+     arr[x_index] = static_cast<scalar_t>(x * cos_val - y * sin_val);
+     arr[y_index] = static_cast<scalar_t>(y * cos_val + x * sin_val);
+   } else {
+     // GPT-J / LLaMA style: layout [x0, x1, ..., y0, y1, ...]
+     // cos/sin: [..., rotary_dim], one entry per (x, y) pair.
+     const int x_index = rot_offset;
+     const int y_index = rot_offset + embed_dim;
+ 
+     // Directly access Global Memory
+     const float cos_val_x = static_cast<float>(cos_ptr[rot_offset]);
+     const float sin_val_x = static_cast<float>(sin_ptr[rot_offset]);
+     const float cos_val_y = static_cast<float>(cos_ptr[rot_offset + embed_dim]);
+     const float sin_val_y = static_cast<float>(sin_ptr[rot_offset + embed_dim]);
+ 
+     const float x = static_cast<float>(arr[x_index]);
+     const float y = static_cast<float>(arr[y_index]);
+     arr[x_index] = static_cast<scalar_t>(x * cos_val_x - y * sin_val_x);
+     arr[y_index] = static_cast<scalar_t>(y * cos_val_y + x * sin_val_y);
+   }
+ }
+ 
+ template <typename scalar_t, bool interleaved>
+ inline __device__ void apply_token_rotary_embedding_vec(
+     scalar_t* __restrict__ arr,            // [head_size]
+     const scalar_t* __restrict__ cos_ptr,  // [rot_dim]
+     const scalar_t* __restrict__ sin_ptr,  // [rot_dim]
+     int rot_offset,
+     int embed_dim) {
+   using vec_t = float4;
+   constexpr int kVecBytes = sizeof(vec_t);
+   constexpr int kScalarBytes = sizeof(scalar_t);
+   constexpr int kElePerVec = kVecBytes / kScalarBytes;
+ 
+   // Union for type punning to avoid strict aliasing issues with reinterpret_cast
+   union VecUnion {
+     vec_t vec;
+     scalar_t elems[kElePerVec];
+   };
+ 
+   if constexpr (interleaved) {
+     // Interleaved: arr has [x0, y0, x1, y1...]
+     // A single vector load contains 'kElePerVec' elements, which is 'kElePerVec / 2' pairs.
+     // rot_offset is the index of the PAIR.
+     // Address in arr is rot_offset * 2.
+ 
+     VecUnion data;
+     data.vec = *reinterpret_cast<const vec_t*>(arr + rot_offset * 2);
+ 
+     // Vectorized load for cos/sin: always 8 bytes (half of vector size, float2)
+     // float2 covers (kElePerVec / 2) scalar_t elements (pairs)
+     float2 cos_vec = *reinterpret_cast<const float2*>(cos_ptr + rot_offset);
+     float2 sin_vec = *reinterpret_cast<const float2*>(sin_ptr + rot_offset);
+ 
+     union CosSinVec {
+       float2 vec;
+       scalar_t elems[kElePerVec / 2];
+     };
+     CosSinVec cos_u, sin_u;
+     cos_u.vec = cos_vec;
+     sin_u.vec = sin_vec;
+ 
+ #pragma unroll
+     for (int i = 0; i < kElePerVec; i += 2) {
+       // data.elems[i] is x, data.elems[i+1] is y
+       // They correspond to pair index: rot_offset + (i / 2)
+       int idx = i / 2;
+       float cos_val = static_cast<float>(cos_u.elems[idx]);
+       float sin_val = static_cast<float>(sin_u.elems[idx]);
+ 
+       float x = static_cast<float>(data.elems[i]);
+       float y = static_cast<float>(data.elems[i + 1]);
+ 
+       data.elems[i] = static_cast<scalar_t>(x * cos_val - y * sin_val);
+       data.elems[i + 1] = static_cast<scalar_t>(y * cos_val + x * sin_val);
+     }
+ 
+     *reinterpret_cast<vec_t*>(arr + rot_offset * 2) = data.vec;
+ 
+   } else {
+     // Non-interleaved: X and Y are separated by embed_dim.
+     // We process 'kElePerVec' PAIRS at once.
+     // Load X vector and Y vector.
+ 
+     VecUnion data_x, data_y;
+     data_x.vec = *reinterpret_cast<const vec_t*>(arr + rot_offset);
+     data_y.vec = *reinterpret_cast<const vec_t*>(arr + rot_offset + embed_dim);
+ 
+     // Vectorized load for cos/sin: always 16 bytes (full vector size, float4)
+     // float4 covers kElePerVec scalar_t elements
+     float4 cos_vec_x = *reinterpret_cast<const float4*>(cos_ptr + rot_offset);
+     float4 sin_vec_x = *reinterpret_cast<const float4*>(sin_ptr + rot_offset);
+     float4 cos_vec_y = *reinterpret_cast<const float4*>(cos_ptr + rot_offset + embed_dim);
+     float4 sin_vec_y = *reinterpret_cast<const float4*>(sin_ptr + rot_offset + embed_dim);
+ 
+     union CosSinVec {
+       float4 vec;
+       scalar_t elems[kElePerVec];
+     };
+     CosSinVec cos_u_x, sin_u_x, cos_u_y, sin_u_y;
+     cos_u_x.vec = cos_vec_x;
+     sin_u_x.vec = sin_vec_x;
+     cos_u_y.vec = cos_vec_y;
+     sin_u_y.vec = sin_vec_y;
+ 
+ #pragma unroll
+     for (int i = 0; i < kElePerVec; ++i) {
+       float cos_val_x = static_cast<float>(cos_u_x.elems[i]);
+       float sin_val_x = static_cast<float>(sin_u_x.elems[i]);
+       float cos_val_y = static_cast<float>(cos_u_y.elems[i]);
+       float sin_val_y = static_cast<float>(sin_u_y.elems[i]);
+ 
+       float x = static_cast<float>(data_x.elems[i]);
+       float y = static_cast<float>(data_y.elems[i]);
+ 
+       data_x.elems[i] = static_cast<scalar_t>(x * cos_val_x - y * sin_val_x);
+       data_y.elems[i] = static_cast<scalar_t>(y * cos_val_y + x * sin_val_y);
+     }
+ 
+     *reinterpret_cast<vec_t*>(arr + rot_offset) = data_x.vec;
+     *reinterpret_cast<vec_t*>(arr + rot_offset + embed_dim) = data_y.vec;
+   }
+ }
+ 
+ template <>
+ inline __device__ void apply_token_rotary_embedding<float, true>(
+     float* __restrict__ arr,
+     const float* __restrict__ cos_ptr,
+     const float* __restrict__ sin_ptr,
+     int rot_offset,
+     int /*embed_dim*/) {
+   float2 xy = *reinterpret_cast<const float2*>(arr + rot_offset * 2);
+ 
+   const float cos_val = static_cast<float>(cos_ptr[rot_offset]);
+   const float sin_val = static_cast<float>(sin_ptr[rot_offset]);
+ 
+   float2 out;
+   out.x = xy.x * cos_val - xy.y * sin_val;
+   out.y = xy.y * cos_val + xy.x * sin_val;
+ 
+   *reinterpret_cast<float2*>(arr + rot_offset * 2) = out;
+ }
+ 
+ // 2D grid kernel: parallel over tokens (grid.x) and pair-tiles (grid.y)
+ template <typename scalar_t, bool interleaved, bool vectorized, int ROT_EMBED_DIM>
+ __global__ void rotary_embedding_kernel_2d(
+     const scalar_t* __restrict__ cos_data,  // [num_tokens, rot_dim_arg]
+     const scalar_t* __restrict__ sin_data,  // [num_tokens, rot_dim_arg]
+     scalar_t* __restrict__ query_total,
+     scalar_t* __restrict__ key_total,
+     const int rot_dim_arg,
+     const int embed_dim_for_rotation_arg,
+     const int64_t query_token_stride,
+     const int64_t key_token_stride,
+     const int64_t head_stride_query,
+     const int64_t head_stride_key,
+     const int num_heads,
+     const int num_kv_heads,
+     const int head_size_arg,
+     const int blocks_per_token) {
+   // Removed shared memory allocation and loading logic
+ 
+   const int token_idx = blockIdx.x;
+   if (token_idx >= gridDim.x) {
+     return;
+   }
+ 
+   // Pointers to Global Memory for the current token
+   const scalar_t* current_token_cos_ptr = cos_data + token_idx * rot_dim_arg;
+   const scalar_t* current_token_sin_ptr = sin_data + token_idx * rot_dim_arg;
+ 
+   constexpr int kVecBytes = 16;
+   const int scalar_size = sizeof(scalar_t);
+ 
+   // Removed SMEM loading and __syncthreads()
+ 
+   scalar_t* query_for_token = query_total + token_idx * (int)query_token_stride;
+   scalar_t* key_for_token = (key_total != nullptr) ? (key_total + token_idx * (int)key_token_stride) : nullptr;
+ 
+   const int local_block_idx = blockIdx.y;
+ 
+   // Use compile-time constant if ROT_EMBED_DIM > 0, otherwise fallback to runtime argument
+   const int embed_dim_for_rotation = (ROT_EMBED_DIM > 0) ? ROT_EMBED_DIM : embed_dim_for_rotation_arg;
+ 
+   if constexpr (vectorized) {
+     using vec_t = float4;
+     constexpr int kElePerVec = kVecBytes / sizeof(scalar_t);
+     constexpr int pairs_per_step = interleaved ? (kElePerVec / 2) : kElePerVec;
+ 
+     const int pair_stride = blockDim.x * blocks_per_token * pairs_per_step;
+     int i = (local_block_idx * blockDim.x + threadIdx.x) * pairs_per_step;
+     const int nq_pairs = num_heads * embed_dim_for_rotation;
+ 
+     // PROLOGUE: Preload 0-th batch
+     RotaryVecData<scalar_t, interleaved> curr_data;
+     if (i < nq_pairs) {
+       int head_idx = i / embed_dim_for_rotation;
+       int rot_offset = i % embed_dim_for_rotation;
+       scalar_t* ptr = query_for_token + head_idx * (int)head_stride_query;
+       curr_data = load_rotary_vec<scalar_t, interleaved>(
+           ptr, current_token_cos_ptr, current_token_sin_ptr, rot_offset, embed_dim_for_rotation);
+     }
+ 
+     int next_i = i + pair_stride;
+ 
+     // MAIN LOOP
+     for (; i < nq_pairs; i += pair_stride, next_i += pair_stride) {
+       // 1. PREFETCH NEXT
+       RotaryVecData<scalar_t, interleaved> next_data;
+       bool active_next = (next_i < nq_pairs);
+       if (active_next) {
+         int head_idx_next = next_i / embed_dim_for_rotation;
+         int rot_offset_next = next_i % embed_dim_for_rotation;
+         scalar_t* ptr_next = query_for_token + head_idx_next * (int)head_stride_query;
+         next_data = load_rotary_vec<scalar_t, interleaved>(
+             ptr_next, current_token_cos_ptr, current_token_sin_ptr, rot_offset_next, embed_dim_for_rotation);
+       }
+ 
+       // 2. COMPUTE CURRENT
+       // The compute uses curr_data which corresponds to index 'i'
+       int head_idx = i / embed_dim_for_rotation;
+       int rot_offset = i % embed_dim_for_rotation;
+       scalar_t* ptr = query_for_token + head_idx * (int)head_stride_query;
+       compute_store_rotary_vec<scalar_t, interleaved>(ptr, curr_data, rot_offset, embed_dim_for_rotation);
+ 
+       // 3. SHIFT
+       curr_data = next_data;
+     }
+ 
+     if (key_for_token != nullptr) {
+       const int nk_pairs = num_kv_heads * embed_dim_for_rotation;
+       int k_i = (local_block_idx * blockDim.x + threadIdx.x) * pairs_per_step;
+ 
+       // PROLOGUE (Key)
+       RotaryVecData<scalar_t, interleaved> curr_data_k;
+       if (k_i < nk_pairs) {
+         int head_idx = k_i / embed_dim_for_rotation;
+         int rot_offset = k_i % embed_dim_for_rotation;
+         scalar_t* ptr = key_for_token + head_idx * (int)head_stride_key;
+         curr_data_k = load_rotary_vec<scalar_t, interleaved>(
+             ptr, current_token_cos_ptr, current_token_sin_ptr, rot_offset, embed_dim_for_rotation);
+       }
+ 
+       int next_k_i = k_i + pair_stride;
+ 
+       // MAIN LOOP (Key)
+       for (; k_i < nk_pairs; k_i += pair_stride, next_k_i += pair_stride) {
+         // 1. PREFETCH NEXT
+         RotaryVecData<scalar_t, interleaved> next_data_k;
+         bool active_next = (next_k_i < nk_pairs);
+         if (active_next) {
+           int head_idx_next = next_k_i / embed_dim_for_rotation;
+           int rot_offset_next = next_k_i % embed_dim_for_rotation;
+           scalar_t* ptr_next = key_for_token + head_idx_next * (int)head_stride_key;
+           next_data_k = load_rotary_vec<scalar_t, interleaved>(
+               ptr_next, current_token_cos_ptr, current_token_sin_ptr, rot_offset_next, embed_dim_for_rotation);
+         }
+ 
+         // 2. COMPUTE CURRENT
+         int head_idx = k_i / embed_dim_for_rotation;
+         int rot_offset = k_i % embed_dim_for_rotation;
+         scalar_t* ptr = key_for_token + head_idx * (int)head_stride_key;
+         compute_store_rotary_vec<scalar_t, interleaved>(ptr, curr_data_k, rot_offset, embed_dim_for_rotation);
+ 
+         // 3. SHIFT
+         curr_data_k = next_data_k;
+       }
+     }
+   } else {
+     // Fallback to scalar implementation
+     const int pair_stride = blockDim.x * blocks_per_token;
+     const int thread_pair_offset = local_block_idx * blockDim.x + threadIdx.x;
+ 
+     const int nq_pairs = num_heads * embed_dim_for_rotation;
+     for (int i = thread_pair_offset; i < nq_pairs; i += pair_stride) {
+       const int head_idx = i / embed_dim_for_rotation;
+       const int rot_offset = i % embed_dim_for_rotation;
+ 
+       scalar_t* query_for_token_head = query_for_token + head_idx * (int)head_stride_query;
+       apply_token_rotary_embedding<scalar_t, interleaved>(
+           query_for_token_head, current_token_cos_ptr, current_token_sin_ptr, rot_offset, embed_dim_for_rotation);
+     }
+ 
+     if (key_for_token != nullptr) {
+       const int nk_pairs = num_kv_heads * embed_dim_for_rotation;
+       for (int i = thread_pair_offset; i < nk_pairs; i += pair_stride) {
+         const int head_idx = i / embed_dim_for_rotation;
+         const int rot_offset = i % embed_dim_for_rotation;
+ 
+         scalar_t* key_for_token_head = key_for_token + head_idx * (int)head_stride_key;
+         apply_token_rotary_embedding<scalar_t, interleaved>(
+             key_for_token_head, current_token_cos_ptr, current_token_sin_ptr, rot_offset, embed_dim_for_rotation);
+       }
+     }
+   }
+ }
+ 
+ // 1D grid kernel: each block handles one token
+ template <typename scalar_t, bool interleaved, bool vectorized, int ROT_EMBED_DIM>
+ __launch_bounds__(512) __global__ void rotary_embedding_kernel_1d(
+     const scalar_t* __restrict__ cos_data,  // [num_tokens, rot_dim_arg]
+     const scalar_t* __restrict__ sin_data,  // [num_tokens, rot_dim_arg]
+     scalar_t* __restrict__ query_total,
+     scalar_t* __restrict__ key_total,
+     const int rot_dim_arg,
+     const int embed_dim_for_rotation_arg,
+     const int64_t query_token_stride,
+     const int64_t key_token_stride,
+     const int64_t head_stride_query,
+     const int64_t head_stride_key,
+     const int num_heads,
+     const int num_kv_heads,
+     const int head_size_arg) {
+   // Removed shared memory allocation and loading logic
+ 
+   const int token_idx = blockIdx.x;
+   if (token_idx >= gridDim.x) return;
+ 
+   const scalar_t* current_token_cos_ptr = cos_data + token_idx * rot_dim_arg;
+   const scalar_t* current_token_sin_ptr = sin_data + token_idx * rot_dim_arg;
+ 
+   constexpr int kVecBytes = 16;
+   const int scalar_size = sizeof(scalar_t);
+ 
+   // Removed SMEM loading and __syncthreads()
+ 
+   scalar_t* query_for_token = query_total + token_idx * (int)query_token_stride;
+   scalar_t* key_for_token = (key_total != nullptr) ? (key_total + token_idx * (int)key_token_stride) : nullptr;
+ 
+   // Use compile-time constant if ROT_EMBED_DIM > 0, otherwise fallback to runtime argument
+   const int embed_dim_for_rotation = (ROT_EMBED_DIM > 0) ? ROT_EMBED_DIM : embed_dim_for_rotation_arg;
+ 
+   if constexpr (vectorized) {
+     constexpr int kVecBytes = 16;
+     constexpr int kElePerVec = kVecBytes / sizeof(scalar_t);
+     // Interleaved: 1 vec = kElePerVec elements = kElePerVec/2 pairs
+     // Non-Interleaved: 1 vec (X) + 1 vec (Y) = kElePerVec elements = kElePerVec pairs
+     constexpr int pairs_per_step = interleaved ? (kElePerVec / 2) : kElePerVec;
+     const int nq_pairs = num_heads * embed_dim_for_rotation;
+ 
+     // Original Stride
+     const int stride = blockDim.x * pairs_per_step;
+ 
+     int i = threadIdx.x * pairs_per_step;
+ 
+     // ==========================================
+     // PROLOGUE: Preload first batch
+     // ==========================================
+     RotaryVecData<scalar_t, interleaved> curr_data;
+ 
+     if (i < nq_pairs) {
+       int head_idx = i / embed_dim_for_rotation;
+       int rot_offset = i % embed_dim_for_rotation;
+       scalar_t* ptr = query_for_token + head_idx * (int)head_stride_query;
+       curr_data = load_rotary_vec<scalar_t, interleaved>(
+           ptr, current_token_cos_ptr, current_token_sin_ptr, rot_offset, embed_dim_for_rotation);
+     }
+ 
+     int next_i = i + stride;
+ 
+     // ==========================================
+     // MAIN LOOP
+     // ==========================================
+     for (; i < nq_pairs; i += stride, next_i += stride) {
+       // --- 1. PREFETCH NEXT BATCH ---
+       RotaryVecData<scalar_t, interleaved> next_data;
+ 
+       if (next_i < nq_pairs) {
+         int head_idx = next_i / embed_dim_for_rotation;
+         int rot_offset = next_i % embed_dim_for_rotation;
+         scalar_t* ptr = query_for_token + head_idx * (int)head_stride_query;
+         next_data = load_rotary_vec<scalar_t, interleaved>(
+             ptr, current_token_cos_ptr, current_token_sin_ptr, rot_offset, embed_dim_for_rotation);
+       }
+ 
+       // --- 2. COMPUTE CURRENT BATCH ---
+       if (i < nq_pairs) {
+         int head_idx = i / embed_dim_for_rotation;
+         int rot_offset = i % embed_dim_for_rotation;
+         scalar_t* ptr = query_for_token + head_idx * (int)head_stride_query;
+         compute_store_rotary_vec<scalar_t, interleaved>(ptr, curr_data, rot_offset, embed_dim_for_rotation);
+       }
+ 
+       // --- 3. SHIFT ---
+       curr_data = next_data;
+     }
+ 
+     // Key branch
+     if (key_for_token != nullptr) {
+       const int nk_pairs = num_kv_heads * embed_dim_for_rotation;
+       int k_i = threadIdx.x * pairs_per_step;
+ 
+       // PROLOGUE (Key): Preload first batch
+       RotaryVecData<scalar_t, interleaved> curr_data_k;
+ 
+       if (k_i < nk_pairs) {
+         int head_idx = k_i / embed_dim_for_rotation;
+         int rot_offset = k_i % embed_dim_for_rotation;
+         scalar_t* ptr = key_for_token + head_idx * (int)head_stride_key;
+         curr_data_k = load_rotary_vec<scalar_t, interleaved>(
+             ptr, current_token_cos_ptr, current_token_sin_ptr, rot_offset, embed_dim_for_rotation);
+       }
+ 
+       int next_k_i = k_i + stride;
+ 
+       // MAIN LOOP (Key)
+       for (; k_i < nk_pairs; k_i += stride, next_k_i += stride) {
+         // --- 1. PREFETCH NEXT BATCH ---
+         RotaryVecData<scalar_t, interleaved> next_data_k;
+ 
+         if (next_k_i < nk_pairs) {
+           int head_idx = next_k_i / embed_dim_for_rotation;
+           int rot_offset = next_k_i % embed_dim_for_rotation;
+           scalar_t* ptr = key_for_token + head_idx * (int)head_stride_key;
+           next_data_k = load_rotary_vec<scalar_t, interleaved>(
+               ptr, current_token_cos_ptr, current_token_sin_ptr, rot_offset, embed_dim_for_rotation);
+         }
+ 
+         // --- 2. COMPUTE CURRENT BATCH ---
+         if (k_i < nk_pairs) {
+           int head_idx = k_i / embed_dim_for_rotation;
+           int rot_offset = k_i % embed_dim_for_rotation;
+           scalar_t* ptr = key_for_token + head_idx * (int)head_stride_key;
+           compute_store_rotary_vec<scalar_t, interleaved>(ptr, curr_data_k, rot_offset, embed_dim_for_rotation);
+         }
+ 
+         // --- 3. SHIFT ---
+         curr_data_k = next_data_k;
+       }
+     }
+   } else {
+     // Fallback scalar
+     const int nq_pairs = num_heads * embed_dim_for_rotation;
+     for (int i = threadIdx.x; i < nq_pairs; i += blockDim.x) {
+       const int head_idx = i / embed_dim_for_rotation;
+       const int rot_offset = i % embed_dim_for_rotation;
+       scalar_t* query_for_token_head = query_for_token + head_idx * (int)head_stride_query;
+ 
+       apply_token_rotary_embedding<scalar_t, interleaved>(
+           query_for_token_head, current_token_cos_ptr, current_token_sin_ptr, rot_offset, embed_dim_for_rotation);
+     }
+ 
+     if (key_for_token != nullptr) {
+       const int nk_pairs = num_kv_heads * embed_dim_for_rotation;
+       for (int i = threadIdx.x; i < nk_pairs; i += blockDim.x) {
+         const int head_idx = i / embed_dim_for_rotation;
+         const int rot_offset = i % embed_dim_for_rotation;
+ 
+         scalar_t* key_for_token_head = key_for_token + head_idx * (int)head_stride_key;
+         apply_token_rotary_embedding<scalar_t, interleaved>(
+             key_for_token_head, current_token_cos_ptr, current_token_sin_ptr, rot_offset, embed_dim_for_rotation);
+       }
+     }
+   }
+ }
+ 
+ // Define dispatch macro for Rotation Dimension (Pair Count)
+ // Note: This is usually HEAD_SIZE / 2 for LLaMA-like models.
+ // Case 32 -> head_size 64
+ // Case 64 -> head_size 128
+ // Case 128 -> head_size 256
+ #define DISPATCH_ROT_DIM(embed_dim, ...)  \
+   switch (embed_dim) {                    \
+     case 16: {                            \
+       constexpr int ROT_EMBED_DIM = 16;   \
+       __VA_ARGS__;                        \
+       break;                              \
+     }                                     \
+     case 32: {                            \
+       constexpr int ROT_EMBED_DIM = 32;   \
+       __VA_ARGS__;                        \
+       break;                              \
+     }                                     \
+     case 64: {                            \
+       constexpr int ROT_EMBED_DIM = 64;   \
+       __VA_ARGS__;                        \
+       break;                              \
+     }                                     \
+     case 96: {                            \
+       constexpr int ROT_EMBED_DIM = 96;   \
+       __VA_ARGS__;                        \
+       break;                              \
+     }                                     \
+     case 128: {                           \
+       constexpr int ROT_EMBED_DIM = 128;  \
+       __VA_ARGS__;                        \
+       break;                              \
+     }                                     \
+     case 256: {                           \
+       constexpr int ROT_EMBED_DIM = 256;  \
+       __VA_ARGS__;                        \
+       break;                              \
+     }                                     \
+     case 512: {                           \
+       constexpr int ROT_EMBED_DIM = 512;  \
+       __VA_ARGS__;                        \
+       break;                              \
+     }                                     \
+     case 1024: {                          \
+       constexpr int ROT_EMBED_DIM = 1024; \
+       __VA_ARGS__;                        \
+       break;                              \
+     }                                     \
+     default: {                            \
+       constexpr int ROT_EMBED_DIM = 0;    \
+       __VA_ARGS__;                        \
+       break;                              \
+     }                                     \
+   }
+ 
+ void rotary_embedding_cos_sin(
+     at::Tensor& cos,
+     at::Tensor& sin,
+     at::Tensor& query,
+     const std::optional<at::Tensor>& key,
+     int64_t head_size,
+     bool interleaved) {
+   TORCH_CHECK(
+       query.dim() == 2 || query.dim() == 3,
+       "query must be in shape [num_tokens, hidden_size] or [num_tokens, num_heads, head_size]");
+   if (key.has_value()) {
+     TORCH_CHECK(
+         key->dim() == 2 || key->dim() == 3,
+         "key must be in shape [num_tokens, hidden_size] or [num_tokens, num_kv_heads, head_size]");
+   }
+ 
+   const int64_t num_tokens = query.size(0);
+ 
+   TORCH_CHECK(cos.dim() == 2, "cos must be in shape [num_tokens, D_cos]");
+   TORCH_CHECK(sin.dim() == 2, "sin must be in shape [num_tokens, D_sin]");
+   TORCH_CHECK(cos.size(0) == num_tokens, "cos num_tokens mismatch with query");
+   TORCH_CHECK(sin.size(0) == num_tokens, "sin num_tokens mismatch with query");
+   TORCH_CHECK(cos.size(1) == sin.size(1), "cos and sin D_cos/D_sin mismatch");
+ 
+   TORCH_CHECK(cos.scalar_type() == query.scalar_type(), "cos dtype mismatch with query");
+   TORCH_CHECK(sin.scalar_type() == query.scalar_type(), "sin dtype mismatch with query");
+   TORCH_CHECK(cos.is_cuda() && sin.is_cuda() && query.is_cuda(), "cos/sin/query must be CUDA tensors");
+   TORCH_CHECK(query.is_contiguous(), "query must be contiguous");
+   if (key.has_value()) {
+     TORCH_CHECK(key->is_cuda(), "key must be CUDA tensor if provided");
+     TORCH_CHECK(key->scalar_type() == query.scalar_type(), "key dtype mismatch with query");
+     TORCH_CHECK(key->is_contiguous(), "key must be contiguous");
+   }
+ 
+   int query_hidden_size_calculated;
+   if (query.dim() == 2) {
+     query_hidden_size_calculated = (int)query.size(1);
+   } else {
+     query_hidden_size_calculated = (int)query.size(1) * (int)query.size(2);
+     TORCH_CHECK(query.size(2) == head_size, "query head_size mismatch in 3D tensor");
+   }
+   TORCH_CHECK(query_hidden_size_calculated % head_size == 0, "query_hidden_size not divisible by head_size");
+   int num_heads = query_hidden_size_calculated / (int)head_size;
+ 
+   int key_hidden_size_calculated = 0;
+   int num_kv_heads = num_heads;
+   if (key.has_value()) {
+     TORCH_CHECK((int)key->size(0) == num_tokens, "key num_tokens mismatch");
+     if (key->dim() == 2) {
+       key_hidden_size_calculated = (int)key->size(1);
+     } else {
+       key_hidden_size_calculated = (int)key->size(1) * (int)key->size(2);
+       TORCH_CHECK((int)key->size(2) == head_size, "key head_size mismatch in 3D tensor");
+     }
+     TORCH_CHECK(key_hidden_size_calculated % head_size == 0, "key_hidden_size not divisible by head_size");
+     num_kv_heads = key_hidden_size_calculated / (int)head_size;
+   }
+   TORCH_CHECK(num_heads % num_kv_heads == 0, "num_heads must be divisible by num_kv_heads");
+ 
+   // NeoX interleaved: if cos/sin are full-dim (head_size), downsample once.
+   if (interleaved && cos.size(1) == head_size) {
+     TORCH_CHECK(head_size % 2 == 0, "interleaved layout requires even head_size");
+     const int64_t half = head_size / 2;
+     std::vector<int64_t> new_shape = {cos.size(0), half, 2};
+     cos = cos.view(new_shape).select(2, 0).contiguous();
+     sin = sin.view(new_shape).select(2, 1).contiguous();
+   }
+ 
+   const int rot_dim_from_cache = (int)cos.size(1);
+   const int64_t query_token_stride = query_hidden_size_calculated;
+   const int64_t key_token_stride = key.has_value() ? key_hidden_size_calculated : 0;
+ 
+   int64_t head_stride_query;
+   if (query.dim() == 3 && query.size(1) == num_heads && query.size(2) == head_size) {
+     head_stride_query = query.stride(1);
+   } else {
+     head_stride_query = head_size;
+   }
+ 
+   int64_t head_stride_key = head_size;
+   if (key.has_value()) {
+     if (key->dim() == 3 && key->size(1) == num_kv_heads && key->size(2) == head_size) {
+       head_stride_key = key->stride(1);
+     } else {
+       head_stride_key = head_size;
+     }
+   }
+ 
+   const int embed_dim_for_rotation = interleaved ? rot_dim_from_cache : (rot_dim_from_cache / 2);
+   TORCH_CHECK(embed_dim_for_rotation > 0, "embed_dim_for_rotation must be > 0");
+ 
+   const int max_pairs_to_rotate_per_token =
+       std::max(num_heads * embed_dim_for_rotation, num_kv_heads * embed_dim_for_rotation);
+ 
+   const at::cuda::OptionalCUDAGuard device_guard(device_of(query));
+   const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+ 
+   AT_DISPATCH_FLOATING_TYPES_AND2(
+       at::ScalarType::Half, at::ScalarType::BFloat16, query.scalar_type(), "rotary_embedding_cos_sin", [&] {
+         using torch_scalar_t = scalar_t;
+         using cuda_scalar_t = typename std::conditional<
+             std::is_same<torch_scalar_t, at::Half>::value,
+             nv_half,
+             typename std::conditional<std::is_same<torch_scalar_t, at::BFloat16>::value, nv_bfloat16, torch_scalar_t>::
+                 type>::type;
+ 
+         // Constants for vectorization
+         constexpr int kVecBytes = 16;
+         constexpr int kElePerVec = kVecBytes / sizeof(torch_scalar_t);
+ 
+         // Determine how many pairs one thread handles in one vector step
+         // Interleaved: 1 vector load (16B) contains 'kElePerVec' elements -> 'kElePerVec / 2' pairs.
+         // Non-interleaved: 2 vector loads (32B) contains 'kElePerVec' pairs (X vec + Y vec).
+         const int pairs_per_step = interleaved ? (kElePerVec / 2) : kElePerVec;
+ 
+         // Check if we can guarantee vectorization for ALL tokens
+         // We need Base Pointers, Strides, and Dimension to be aligned.
+         bool can_vectorize_all = true;
+ 
+         // 1. Check Dimensions
+         if (embed_dim_for_rotation % pairs_per_step != 0) can_vectorize_all = false;
+ 
+         // 2. Check Base Pointers
+         if (reinterpret_cast<uintptr_t>(query.data_ptr()) % kVecBytes != 0) can_vectorize_all = false;
+         if (reinterpret_cast<uintptr_t>(cos.data_ptr()) % kVecBytes != 0) can_vectorize_all = false;
+         if (reinterpret_cast<uintptr_t>(sin.data_ptr()) % kVecBytes != 0) can_vectorize_all = false;
+         if (key.has_value()) {
+           if (reinterpret_cast<uintptr_t>(key->data_ptr()) % kVecBytes != 0) can_vectorize_all = false;
+         }
+ 
+         // 3. Check Strides
+         // We need the stride between tokens to be a multiple of vector size
+         // to ensure that if token 0 is aligned, token 1 is also aligned.
+         if (query_token_stride * sizeof(torch_scalar_t) % kVecBytes != 0) can_vectorize_all = false;
+         if (head_stride_query * sizeof(torch_scalar_t) % kVecBytes != 0) can_vectorize_all = false;
+ 
+         if (key.has_value()) {
+           if (key_token_stride * sizeof(torch_scalar_t) % kVecBytes != 0) can_vectorize_all = false;
+           if (head_stride_key * sizeof(torch_scalar_t) % kVecBytes != 0) can_vectorize_all = false;
+         }
+ 
+         // Determine launch configuration
+         // If we can vectorize all, each thread handles 'pairs_per_step' pairs.
+         // Otherwise, fallback to conservative estimate (1 pair per thread) to ensure enough blocks.
+         const int launch_pairs_per_thread = can_vectorize_all ? pairs_per_step : 1;
+ 
+         const int total_threads_needed =
+             (max_pairs_to_rotate_per_token + launch_pairs_per_thread - 1) / launch_pairs_per_thread;
+ 
+         // Case 1: 2D Grid (Split one token across multiple blocks)
+         // Keep block size moderate (128-256) aligned with head size
+         const int threads_per_block_2d = std::min<int>(256, std::max(128, embed_dim_for_rotation));
+         const int blocks_per_token_2d = (total_threads_needed + threads_per_block_2d - 1) / threads_per_block_2d;
+ 
+         // Decide grid strategy
+         const bool use_grid_2d = (num_tokens <= 4) && (blocks_per_token_2d > 1);
+ 
+         // Case 2: 1D Grid (One block per token)
+         // Maximize threads per block to cover all heads in one block if possible
+         // Cap at 512 threads to balance occupancy and register usage
+         const int threads_per_block_1d = std::min<int>(512, std::max(128, (total_threads_needed + 31) / 32 * 32));
+ 
+         // Final launch config
+         const int threads_per_block = use_grid_2d ? threads_per_block_2d : threads_per_block_1d;
+         const int blocks_per_token = use_grid_2d ? blocks_per_token_2d : 1;
+ 
+         dim3 block(threads_per_block);
+         dim3 grid_2d((int)num_tokens, std::max(1, blocks_per_token));
+         dim3 grid_1d((int)num_tokens);
+ 
+         // No shared memory needed
+         size_t smem_size = 0;
+ 
+         DISPATCH_ROT_DIM(
+             embed_dim_for_rotation,
+             auto launch_kernel =
+                 [&](bool vectorized) {
+                   if (interleaved) {
+                     if (use_grid_2d) {
+                       if (vectorized) {
+                         rotary_embedding_kernel_2d<cuda_scalar_t, true, true, ROT_EMBED_DIM>
+                             <<<grid_2d, block, smem_size, stream>>>(
+                                 reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
+                                 reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
+                                 reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
+                                 key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>())
+                                                 : nullptr,
+                                 rot_dim_from_cache,
+                                 embed_dim_for_rotation,
+                                 query_token_stride,
+                                 key_token_stride,
+                                 head_stride_query,
+                                 head_stride_key,
+                                 num_heads,
+                                 num_kv_heads,
+                                 (int)head_size,
+                                 blocks_per_token);
+                       } else {
+                         rotary_embedding_kernel_2d<cuda_scalar_t, true, false, ROT_EMBED_DIM>
+                             <<<grid_2d, block, smem_size, stream>>>(
+                                 reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
+                                 reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
+                                 reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
+                                 key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>())
+                                                 : nullptr,
+                                 rot_dim_from_cache,
+                                 embed_dim_for_rotation,
+                                 query_token_stride,
+                                 key_token_stride,
+                                 head_stride_query,
+                                 head_stride_key,
+                                 num_heads,
+                                 num_kv_heads,
+                                 (int)head_size,
+                                 blocks_per_token);
+                       }
+                     } else {
+                       if (vectorized) {
+                         rotary_embedding_kernel_1d<cuda_scalar_t, true, true, ROT_EMBED_DIM>
+                             <<<grid_1d, block, smem_size, stream>>>(
+                                 reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
+                                 reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
+                                 reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
+                                 key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>())
+                                                 : nullptr,
+                                 rot_dim_from_cache,
+                                 embed_dim_for_rotation,
+                                 query_token_stride,
+                                 key_token_stride,
+                                 head_stride_query,
+                                 head_stride_key,
+                                 num_heads,
+                                 num_kv_heads,
+                                 (int)head_size);
+                       } else {
+                         rotary_embedding_kernel_1d<cuda_scalar_t, true, false, ROT_EMBED_DIM>
+                             <<<grid_1d, block, smem_size, stream>>>(
+                                 reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
+                                 reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
+                                 reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
+                                 key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>())
+                                                 : nullptr,
+                                 rot_dim_from_cache,
+                                 embed_dim_for_rotation,
+                                 query_token_stride,
+                                 key_token_stride,
+                                 head_stride_query,
+                                 head_stride_key,
+                                 num_heads,
+                                 num_kv_heads,
+                                 (int)head_size);
+                       }
+                     }
+                   } else {  // non-interleaved
+                     if (use_grid_2d) {
+                       if (vectorized) {
+                         rotary_embedding_kernel_2d<cuda_scalar_t, false, true, ROT_EMBED_DIM>
+                             <<<grid_2d, block, smem_size, stream>>>(
+                                 reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
+                                 reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
+                                 reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
+                                 key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>())
+                                                 : nullptr,
+                                 rot_dim_from_cache,
+                                 embed_dim_for_rotation,
+                                 query_token_stride,
+                                 key_token_stride,
+                                 head_stride_query,
+                                 head_stride_key,
+                                 num_heads,
+                                 num_kv_heads,
+                                 (int)head_size,
+                                 blocks_per_token);
+                       } else {
+                         rotary_embedding_kernel_2d<cuda_scalar_t, false, false, ROT_EMBED_DIM>
+                             <<<grid_2d, block, smem_size, stream>>>(
+                                 reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
+                                 reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
+                                 reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
+                                 key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>())
+                                                 : nullptr,
+                                 rot_dim_from_cache,
+                                 embed_dim_for_rotation,
+                                 query_token_stride,
+                                 key_token_stride,
+                                 head_stride_query,
+                                 head_stride_key,
+                                 num_heads,
+                                 num_kv_heads,
+                                 (int)head_size,
+                                 blocks_per_token);
+                       }
+                     } else {
+                       if (vectorized) {
+                         rotary_embedding_kernel_1d<cuda_scalar_t, false, true, ROT_EMBED_DIM>
+                             <<<grid_1d, block, smem_size, stream>>>(
+                                 reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
+                                 reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
+                                 reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
+                                 key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>())
+                                                 : nullptr,
+                                 rot_dim_from_cache,
+                                 embed_dim_for_rotation,
+                                 query_token_stride,
+                                 key_token_stride,
+                                 head_stride_query,
+                                 head_stride_key,
+                                 num_heads,
+                                 num_kv_heads,
+                                 (int)head_size);
+                       } else {
+                         rotary_embedding_kernel_1d<cuda_scalar_t, false, false, ROT_EMBED_DIM>
+                             <<<grid_1d, block, smem_size, stream>>>(
+                                 reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
+                                 reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
+                                 reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
+                                 key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>())
+                                                 : nullptr,
+                                 rot_dim_from_cache,
+                                 embed_dim_for_rotation,
+                                 query_token_stride,
+                                 key_token_stride,
+                                 head_stride_query,
+                                 head_stride_key,
+                                 num_heads,
+                                 num_kv_heads,
+                                 (int)head_size);
+                       }
+                     }
+                   }
+                 };
+ 
+             // Close the log file after all launches
+             if (can_vectorize_all) { launch_kernel(true); } else { launch_kernel(false); });
+       });
+ 
+   C10_CUDA_KERNEL_LAUNCH_CHECK();
+ }
\ No newline at end of file

From 7660b8384c17e2668e17394393b5ac9e62818bc9 Mon Sep 17 00:00:00 2001
From: RubiaCx <1084281732@qq.com>
Date: Wed, 17 Dec 2025 05:47:52 -0800
Subject: [PATCH 15/37] Bench

---
 .../benchmark/bench_mm_rotary_embedding.py    | 680 +++++++++++++++---
 1 file changed, 582 insertions(+), 98 deletions(-)

diff --git a/sgl-kernel/benchmark/bench_mm_rotary_embedding.py b/sgl-kernel/benchmark/bench_mm_rotary_embedding.py
index 2692ed62909e..7962f68a524e 100755
--- a/sgl-kernel/benchmark/bench_mm_rotary_embedding.py
+++ b/sgl-kernel/benchmark/bench_mm_rotary_embedding.py
@@ -1,7 +1,11 @@
 #!/usr/bin/env python
 # -*- coding: utf-8 -*-
 
-from typing import List, Tuple
+from __future__ import annotations
+
+import argparse
+from dataclasses import dataclass
+from typing import List, Optional, Tuple
 
 import numpy as np
 import torch
@@ -10,7 +14,7 @@
 from sgl_kernel.testing.rotary_embedding import RotaryEmbedding as NativeRotaryEmbedding
 
 
-def compute_cos_sin_cache(
+def compute_cos_sin_cache_half(
     max_seq_len: int,
     rotary_dim: int,
     base: float = 10000.0,
@@ -26,7 +30,207 @@ def compute_cos_sin_cache(
     return cos, sin
 
 
-def benchmark_mm_rotary_embedding() -> None:
+def _expand_neox_full_repeat(half: torch.Tensor) -> torch.Tensor:
+    """Expand [T, D] -> [T, 2D] by repeating each element twice: [a0,a0,a1,a1,...].
+
+    This matches some NeoX caches expanded to head_size for interleaved layout.
+    The CUDA kernel has a fast path that downsamples such 'full' caches to half.
+    """
+    # half: [T, D]
+    out = torch.empty((half.shape[0], half.shape[1] * 2), dtype=half.dtype, device=half.device)
+    out[:, 0::2] = half
+    out[:, 1::2] = half
+    return out
+
+
+def _expand_llama_full_cat(half: torch.Tensor) -> torch.Tensor:
+    """Expand [T, D] -> [T, 2D] by concatenation: [a0,a1,..., a0,a1,...]."""
+    return torch.cat([half, half], dim=-1)
+
+
+def _make_misaligned_contiguous_like(t: torch.Tensor, offset_elems: int = 1) -> torch.Tensor:
+    """Create a contiguous tensor with same shape/dtype/device but misaligned base pointer.
+
+    Achieved via a 1D buffer and a view with a storage offset.
+    """
+    numel = t.numel()
+    buf = torch.empty((numel + offset_elems,), device=t.device, dtype=t.dtype)
+    out = buf[offset_elems : offset_elems + numel].view_as(t)
+    out.copy_(t)
+    assert out.is_contiguous()
+    return out
+
+
+def _make_tensor(shape: Tuple[int, ...], *, dtype: torch.dtype, device: str, misalign: bool) -> torch.Tensor:
+    numel = int(np.prod(shape))
+    if not misalign:
+        return torch.randn(*shape, dtype=dtype, device=device)
+    # Misaligned but contiguous
+    buf = torch.randn((numel + 1,), dtype=dtype, device=device)
+    return buf[1 : 1 + numel].view(shape)
+
+
+def _torch_naive_rope_inplace(
+    *,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    q: torch.Tensor,
+    k: Optional[torch.Tensor],
+    num_heads: int,
+    num_kv_heads: int,
+    head_size: int,
+    interleaved: bool,
+) -> None:
+    """Naive RoPE on GPU using PyTorch tensor ops (in-place on q/k).
+
+    Supports:
+      - interleaved=True  (NeoX-style [x0,y0,x1,y1,...])
+      - interleaved=False (LLaMA/GPT-J-style [x...][y...])
+      - q/k in 2D ([T, H*D]) or 3D ([T, H, D]) layouts
+      - k can be None (Q-only)
+    """
+
+    def _as_3d(x: torch.Tensor, h: int) -> torch.Tensor:
+        if x.dim() == 3:
+            return x
+        # 2D flattened
+        return x.view(x.shape[0], h, head_size)
+
+    # Normalize cos/sin for interleaved fast-path compat:
+    # - interleaved expects cos/sin: [T, embed_dim] where embed_dim = rotary_dim/2
+    # - if provided as full head_size (repeat format), downsample like the CUDA kernel does
+    if interleaved and cos.shape[1] == head_size:
+        half = head_size // 2
+        cos = cos.view(cos.shape[0], half, 2).select(2, 0).contiguous()
+        sin = sin.view(sin.shape[0], half, 2).select(2, 1).contiguous()
+
+    if interleaved:
+        embed_dim = int(cos.shape[1])
+        rot_dim = embed_dim * 2
+    else:
+        embed_dim = int(cos.shape[1]) // 2
+        rot_dim = embed_dim * 2
+
+    cos_b = cos[:, None, :embed_dim]
+    sin_b = sin[:, None, :embed_dim]
+
+    def _apply(x: torch.Tensor, h: int) -> None:
+        x3 = _as_3d(x, h)
+        # Only rotate the first rot_dim elements; leave tail intact.
+        xr = x3[..., :rot_dim]
+        if interleaved:
+            # [T,H,embed_dim,2]
+            xr2 = xr.view(xr.shape[0], xr.shape[1], embed_dim, 2)
+            x0 = xr2[..., 0]
+            x1 = xr2[..., 1]
+            out0 = x0 * cos_b - x1 * sin_b
+            out1 = x1 * cos_b + x0 * sin_b
+            xr2[..., 0].copy_(out0)
+            xr2[..., 1].copy_(out1)
+        else:
+            x0 = xr[..., :embed_dim]
+            x1 = xr[..., embed_dim:rot_dim]
+            cos_x = cos[:, None, :embed_dim]
+            sin_x = sin[:, None, :embed_dim]
+            cos_y = cos[:, None, embed_dim:rot_dim]
+            sin_y = sin[:, None, embed_dim:rot_dim]
+            out0 = x0 * cos_x - x1 * sin_x
+            out1 = x1 * cos_y + x0 * sin_y
+            xr[..., :embed_dim].copy_(out0)
+            xr[..., embed_dim:rot_dim].copy_(out1)
+
+    _apply(q, num_heads)
+    if k is not None:
+        _apply(k, num_kv_heads)
+
+
+def _predict_can_vectorize_all(
+    *,
+    dtype: torch.dtype,
+    interleaved: bool,
+    embed_dim_for_rotation: int,
+    query: torch.Tensor,
+    key: Optional[torch.Tensor],
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    query_token_stride_elems: int,
+    key_token_stride_elems: int,
+    head_stride_query_elems: int,
+    head_stride_key_elems: int,
+) -> bool:
+    kVecBytes = 16
+    elem_bytes = torch.tensor([], dtype=dtype).element_size()
+    kElePerVec = kVecBytes // elem_bytes
+    pairs_per_step = (kElePerVec // 2) if interleaved else kElePerVec
+
+    can = True
+    if embed_dim_for_rotation % pairs_per_step != 0:
+        can = False
+    if (query.data_ptr() % kVecBytes) != 0:
+        can = False
+    if (cos.data_ptr() % kVecBytes) != 0:
+        can = False
+    if (sin.data_ptr() % kVecBytes) != 0:
+        can = False
+    if key is not None and (key.data_ptr() % kVecBytes) != 0:
+        can = False
+
+    if (query_token_stride_elems * elem_bytes) % kVecBytes != 0:
+        can = False
+    if (head_stride_query_elems * elem_bytes) % kVecBytes != 0:
+        can = False
+    if key is not None:
+        if (key_token_stride_elems * elem_bytes) % kVecBytes != 0:
+            can = False
+        if (head_stride_key_elems * elem_bytes) % kVecBytes != 0:
+            can = False
+    return can
+
+
+def _predict_use_grid_2d(
+    *,
+    num_tokens: int,
+    num_heads: int,
+    num_kv_heads: int,
+    embed_dim_for_rotation: int,
+    can_vectorize_all: bool,
+    dtype: torch.dtype,
+    interleaved: bool,
+) -> Tuple[bool, int]:
+    # Mirror the C++ launch config logic in rotary_embedding.cu (approx).
+    kVecBytes = 16
+    elem_bytes = torch.tensor([], dtype=dtype).element_size()
+    kElePerVec = kVecBytes // elem_bytes
+    pairs_per_step = (kElePerVec // 2) if interleaved else kElePerVec
+    launch_pairs_per_thread = pairs_per_step if can_vectorize_all else 1
+
+    max_pairs = max(num_heads * embed_dim_for_rotation, num_kv_heads * embed_dim_for_rotation)
+    total_threads_needed = (max_pairs + launch_pairs_per_thread - 1) // launch_pairs_per_thread
+    threads_per_block_2d = min(256, max(128, embed_dim_for_rotation))
+    blocks_per_token_2d = (total_threads_needed + threads_per_block_2d - 1) // threads_per_block_2d
+    use_grid_2d = (num_tokens <= 4) and (blocks_per_token_2d > 1)
+    return use_grid_2d, blocks_per_token_2d
+
+
+@dataclass(frozen=True)
+class BenchCase:
+    name: str
+    batch_size: int
+    num_heads: int
+    num_kv_heads: int
+    head_size: int
+    rotary_dim: int
+    interleaved: bool
+    dtype: torch.dtype
+    layout: str  # "2d" or "3d"
+    key_mode: str  # "with_k" or "no_k"
+    cache_format: str  # "half" | "neox_full_repeat" | "llama_full_cat"
+    misalign: str  # "none" | "q" | "cos"
+    seq_lens: List[int]
+    bench_native: bool = True
+
+
+def benchmark_mm_rotary_embedding(args: argparse.Namespace) -> None:
     try:
         from vllm.model_executor.layers.rotary_embedding import (
             RotaryEmbedding as vLLMRotaryEmbedding,
@@ -48,28 +252,179 @@ def benchmark_mm_rotary_embedding() -> None:
         print("flash_attn not available")
 
     device = "cuda"
-    dtype = torch.bfloat16
     max_seq_len = 65536
 
-    # Include Qwen-image observed shapes (e.g., 3015 image tokens, small text token counts)
-    seq_lens = [1, 2, 4, 6, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 3015, 4096, 8192]
-
-    # Test configurations: (batch_size, num_heads, num_kv_heads, head_size)
-    configs = [
-        (1, 32, 8, 128),
-        (1, 64, 8, 128),
-        (1, 32, 8, 64),
-        (1, 32, 8, 256),
-        (32, 32, 8, 64),
-        (1, 32, 1, 128),
-        (32, 8, 8, 128),
-        (1, 32, 8, 80),
-        (1, 64, 8, 256),
-        # Qwen-image observed: q/k view is [num_tokens, 24, 128] and cos/sin is [num_tokens, 64], interleaved=True
-        (1, 24, 24, 128),
+    print(
+        "Legend:\n"
+        "  - vec: vec16 => predicted to take the 16B vectorized path; scalar => predicted scalar/fallback path\n"
+        "  - grid: 1D => one block per token; 2D => split one token across multiple blocks (only when num_tokens<=4)\n"
+        "  - bpt2d: estimated blocks-per-token if 2D were used (useful to judge 'small grid but many blocks per token')\n"
+    )
+
+    cases: List[BenchCase] = [
+        # Baseline (vectorized, interleaved)
+        BenchCase(
+            name="base_bf16_int2d_k_halfcache",
+            batch_size=1,
+            num_heads=32,
+            num_kv_heads=8,
+            head_size=128,
+            rotary_dim=128,
+            interleaved=True,
+            dtype=torch.bfloat16,
+            layout="2d",
+            key_mode="with_k",
+            cache_format="half",
+            misalign="none",
+            seq_lens=[1, 2, 4, 32, 256, 3015, 8192],
+            bench_native=True,
+        ),
+        # Non-interleaved (GPT-J/LLaMA-style)
+        BenchCase(
+            name="base_bf16_nonint2d_k_llamacache",
+            batch_size=1,
+            num_heads=32,
+            num_kv_heads=8,
+            head_size=128,
+            rotary_dim=128,
+            interleaved=False,
+            dtype=torch.bfloat16,
+            layout="2d",
+            key_mode="with_k",
+            cache_format="llama_full_cat",
+            misalign="none",
+            seq_lens=[1, 4, 256, 8192],
+            bench_native=False,
+        ),
+        # Force non-vectorized path via misaligned Q pointer
+        BenchCase(
+            name="base_bf16_int2d_k_halfcache_misalignQ",
+            batch_size=1,
+            num_heads=32,
+            num_kv_heads=8,
+            head_size=128,
+            rotary_dim=128,
+            interleaved=True,
+            dtype=torch.bfloat16,
+            layout="2d",
+            key_mode="with_k",
+            cache_format="half",
+            misalign="q",
+            seq_lens=[1, 4, 256, 8192],
+            bench_native=False,
+        ),
+        # Key is optional (Q-only)
+        BenchCase(
+            name="base_bf16_int2d_noK_halfcache",
+            batch_size=1,
+            num_heads=32,
+            num_kv_heads=8,
+            head_size=128,
+            rotary_dim=128,
+            interleaved=True,
+            dtype=torch.bfloat16,
+            layout="2d",
+            key_mode="no_k",
+            cache_format="half",
+            misalign="none",
+            seq_lens=[1, 4, 256, 8192],
+            bench_native=False,
+        ),
+        # 3D layout path ([T, H, D])
+        BenchCase(
+            name="base_bf16_int3d_k_halfcache",
+            batch_size=1,
+            num_heads=32,
+            num_kv_heads=8,
+            head_size=128,
+            rotary_dim=128,
+            interleaved=True,
+            dtype=torch.bfloat16,
+            layout="3d",
+            key_mode="with_k",
+            cache_format="half",
+            misalign="none",
+            seq_lens=[1, 4, 256, 8192],
+            bench_native=False,
+        ),
+        # fp16 dtype
+        BenchCase(
+            name="base_fp16_int2d_k_halfcache",
+            batch_size=1,
+            num_heads=32,
+            num_kv_heads=8,
+            head_size=128,
+            rotary_dim=128,
+            interleaved=True,
+            dtype=torch.float16,
+            layout="2d",
+            key_mode="with_k",
+            cache_format="half",
+            misalign="none",
+            seq_lens=[1, 4, 256, 8192],
+            bench_native=False,
+        ),
+        # Qwen-image observed config
+        BenchCase(
+            name="qwen_image_bf16_int2d_k_halfcache",
+            batch_size=1,
+            num_heads=24,
+            num_kv_heads=24,
+            head_size=128,
+            rotary_dim=128,
+            interleaved=True,
+            dtype=torch.bfloat16,
+            layout="2d",
+            key_mode="with_k",
+            cache_format="half",
+            misalign="none",
+            seq_lens=[1, 4, 3015],
+            bench_native=False,
+        ),
+        # Grid-2D stress (num_tokens<=4 and blocks_per_token>1)
+        BenchCase(
+            name="grid2d_stress_bf16_int2d_k_neoxfullrepeat",
+            batch_size=1,
+            num_heads=128,
+            num_kv_heads=128,
+            head_size=256,
+            rotary_dim=256,
+            interleaved=True,
+            dtype=torch.bfloat16,
+            layout="2d",
+            key_mode="with_k",
+            cache_format="neox_full_repeat",
+            misalign="none",
+            seq_lens=[1, 2, 4],
+            bench_native=False,
+        ),
+        # Large batch representative
+        BenchCase(
+            name="batch32_bf16_int2d_k_halfcache",
+            batch_size=32,
+            num_heads=32,
+            num_kv_heads=8,
+            head_size=64,
+            rotary_dim=64,
+            interleaved=True,
+            dtype=torch.bfloat16,
+            layout="2d",
+            key_mode="with_k",
+            cache_format="half",
+            misalign="none",
+            seq_lens=[1, 4, 128, 1024],
+            bench_native=False,
+        ),
     ]
 
-    for batch_size, num_heads, num_kv_heads, head_size in configs:
+    if args.case is not None:
+        allow = set(args.case)
+        cases = [c for c in cases if c.name in allow]
+        missing = allow.difference({c.name for c in cases})
+        if missing:
+            print(f"WARNING: unknown --case entries ignored: {sorted(missing)}")
+
+    for case in cases:
         try:
             torch.cuda.synchronize()
         except torch.AcceleratorError:
@@ -77,67 +432,95 @@ def benchmark_mm_rotary_embedding() -> None:
             pass
 
         try:
-            rotary_dim = head_size
-            cos_cache, sin_cache = compute_cos_sin_cache(
-                max_seq_len, rotary_dim, dtype=dtype
-            )
-            cos_cache = cos_cache.to(device)
-            sin_cache = sin_cache.to(device)
-
-            cos_sin_cache = torch.cat([cos_cache, sin_cache], dim=-1).to(
-                device=device, dtype=dtype
+            cos_half, sin_half = compute_cos_sin_cache_half(
+                max_seq_len, case.rotary_dim, dtype=torch.float32
             )
+            cos_half = cos_half.to(device=device, dtype=case.dtype)
+            sin_half = sin_half.to(device=device, dtype=case.dtype)
+
+            if case.cache_format == "half":
+                cos_cache = cos_half
+                sin_cache = sin_half
+            elif case.cache_format == "neox_full_repeat":
+                cos_cache = _expand_neox_full_repeat(cos_half)
+                sin_cache = _expand_neox_full_repeat(sin_half)
+            elif case.cache_format == "llama_full_cat":
+                cos_cache = _expand_llama_full_cat(cos_half)
+                sin_cache = _expand_llama_full_cat(sin_half)
+            else:
+                raise ValueError(f"Unknown cache_format={case.cache_format}")
+
+            cos_sin_cache = torch.cat([cos_cache, sin_cache], dim=-1)
 
             native_rope = NativeRotaryEmbedding(
-                head_size=head_size,
-                rotary_dim=rotary_dim,
+                head_size=case.head_size,
+                rotary_dim=case.rotary_dim,
                 max_position_embeddings=max_seq_len,
                 base=10000,
-                is_neox_style=True,
-                dtype=dtype,
+                is_neox_style=case.interleaved,
+                dtype=case.dtype,
             ).to(device)
 
             print(
-                f"\nConfig: batch_size={batch_size}, heads={num_heads}/{num_kv_heads}, head_size={head_size}, dtype={dtype}"
+                f"\nCase: {case.name} | "
+                f"bs={case.batch_size} heads={case.num_heads}/{case.num_kv_heads} "
+                f"hs={case.head_size} rd={case.rotary_dim} "
+                f"dtype={case.dtype} interleaved={case.interleaved} layout={case.layout} "
+                f"key={case.key_mode} cache={case.cache_format} misalign={case.misalign}"
             )
             print("-" * 100)
         except Exception as e:
             print(
-                f"\nSkipping config (batch_size={batch_size}, heads={num_heads}/{num_kv_heads}, head_size={head_size}): {e}"
+                f"\nSkipping case {case.name}: {e}"
             )
             continue
 
         header = f"{'seq_len':>8}"
-        header += f" | {'native (ms)':>12}"
+        if case.bench_native:
+            header += f" | {'native (ms)':>12}"
+        header += f" | {'naive (ms)':>11}"
         header += f" | {'sgl_cos_sin (ms)':>16}"
         header += f" | {'sgl_pos (ms)':>12}"
+        header += f" | {'vec':>6} | {'grid':>4} | {'bpt2d':>5}"
         if HAS_VLLM:
             header += f" | {'vLLM (ms)':>10}"
         if HAS_FLASH_ATTN:
             header += f" | {'flash_attn (ms)':>14}"
-        header += f" | {'speedup':>9}"
+        header += f" | {'pos/cos':>9}"
         print(header)
         print("-" * 100)
 
         results = []
 
-        for seq_len in seq_lens:
+        for seq_len in case.seq_lens:
             try:
-                num_tokens = batch_size * seq_len
-                query = torch.randn(
-                    num_tokens, num_heads * head_size, dtype=dtype, device=device
-                )
-                key = torch.randn(
-                    num_tokens, num_kv_heads * head_size, dtype=dtype, device=device
-                )
+                num_tokens = case.batch_size * seq_len
+                if case.layout == "2d":
+                    query_shape = (num_tokens, case.num_heads * case.head_size)
+                    key_shape = (num_tokens, case.num_kv_heads * case.head_size)
+                elif case.layout == "3d":
+                    query_shape = (num_tokens, case.num_heads, case.head_size)
+                    key_shape = (num_tokens, case.num_kv_heads, case.head_size)
+                else:
+                    raise ValueError(f"Unknown layout={case.layout}")
+
+                query = _make_tensor(query_shape, dtype=case.dtype, device=device, misalign=(case.misalign == "q"))
+                key = _make_tensor(key_shape, dtype=case.dtype, device=device, misalign=False)
+                if case.key_mode == "no_k":
+                    key = None
                 positions = torch.arange(
                     seq_len, device=device, dtype=torch.int64
-                ).repeat(batch_size)
+                ).repeat(case.batch_size)
                 cos = cos_cache[positions]
                 sin = sin_cache[positions]
 
+                if case.misalign == "cos":
+                    cos = _make_misaligned_contiguous_like(cos)
+                    sin = _make_misaligned_contiguous_like(sin)
+
                 row_str = f"{seq_len:8d}"
                 native_time = None
+                naive_time = None
                 sgl_cos_sin_time = None
                 sgl_pos_time = None
                 vllm_time = None
@@ -149,17 +532,43 @@ def benchmark_mm_rotary_embedding() -> None:
 
             def fn_native() -> None:
                 q = query.clone()
+                if key is None:
+                    # Native reference expects both q and k; skip in this mode.
+                    raise RuntimeError("native requires key")
                 k = key.clone()
                 native_rope.forward_native(positions, q, k)
 
-            ms, _, _ = triton.testing.do_bench(fn_native, quantiles=[0.5, 0.2, 0.8])
-            native_time = 1000 * ms
-            row_str += f" | {native_time:12.4f}"
+            if case.bench_native and key is not None:
+                ms, _, _ = triton.testing.do_bench(fn_native, quantiles=[0.5, 0.2, 0.8])
+                native_time = 1000 * ms
+                row_str += f" | {native_time:12.4f}"
+            elif case.bench_native:
+                row_str += f" | {'N/A':>12}"
+
+            # Torch naive baseline (works for key=None / interleaved=False / 2D/3D layouts).
+            @torch.no_grad()
+            def fn_naive() -> None:
+                qn = query.clone()
+                kn = None if key is None else key.clone()
+                _torch_naive_rope_inplace(
+                    cos=cos,
+                    sin=sin,
+                    q=qn,
+                    k=kn,
+                    num_heads=case.num_heads,
+                    num_kv_heads=case.num_kv_heads,
+                    head_size=case.head_size,
+                    interleaved=case.interleaved,
+                )
+
+            ms, _, _ = triton.testing.do_bench(fn_naive, quantiles=[0.5, 0.2, 0.8])
+            naive_time = 1000 * ms
+            row_str += f" | {naive_time:11.4f}"
 
             def fn_sgl_cos_sin() -> None:
                 q = query.clone()
-                k = key.clone()
-                sgl_rotary_cos_sin(cos, sin, q, k, head_size, True)
+                k = None if key is None else key.clone()
+                sgl_rotary_cos_sin(cos, sin, q, k, case.head_size, case.interleaved)
 
             ms, _, _ = triton.testing.do_bench(
                 fn_sgl_cos_sin, quantiles=[0.5, 0.2, 0.8]
@@ -169,49 +578,115 @@ def fn_sgl_cos_sin() -> None:
 
             def fn_sgl_pos() -> None:
                 q = query.clone()
-                k = key.clone()
+                k = None if key is None else key.clone()
                 torch.ops.sgl_kernel.rotary_embedding(
                     positions,
                     q,
                     k,
-                    head_size,
+                    case.head_size,
                     cos_sin_cache,
-                    True,
+                    case.interleaved,
                 )
 
-            ms, _, _ = triton.testing.do_bench(fn_sgl_pos, quantiles=[0.5, 0.2, 0.8])
-            sgl_pos_time = 1000 * ms
-            row_str += f" | {sgl_pos_time:12.4f}"
+            try:
+                ms, _, _ = triton.testing.do_bench(fn_sgl_pos, quantiles=[0.5, 0.2, 0.8])
+                sgl_pos_time = 1000 * ms
+                row_str += f" | {sgl_pos_time:12.4f}"
+            except Exception:
+                row_str += f" | {'ERROR':>12}"
+                sgl_pos_time = None
+
+            # Flags: predicted vectorization + whether C++ would pick grid2d (approx)
+            # Compute embed_dim_for_rotation consistent with C++ checks:
+            rot_dim_from_cache = int(cos.shape[1])
+            embed_dim_for_rotation = rot_dim_from_cache if case.interleaved else (rot_dim_from_cache // 2)
+
+            if query.dim() == 2:
+                query_hidden = int(query.size(1))
+                query_token_stride = query_hidden
+                head_stride_query = case.head_size
+            else:
+                query_hidden = int(query.size(1) * query.size(2))
+                query_token_stride = query_hidden
+                head_stride_query = int(query.stride(1))
+
+            if key is None:
+                key_token_stride = 0
+                head_stride_key = case.head_size
+            else:
+                if key.dim() == 2:
+                    key_hidden = int(key.size(1))
+                    key_token_stride = key_hidden
+                    head_stride_key = case.head_size
+                else:
+                    key_hidden = int(key.size(1) * key.size(2))
+                    key_token_stride = key_hidden
+                    head_stride_key = int(key.stride(1))
+
+            can_vec = _predict_can_vectorize_all(
+                dtype=case.dtype,
+                interleaved=case.interleaved,
+                embed_dim_for_rotation=embed_dim_for_rotation,
+                query=query,
+                key=key,
+                cos=cos,
+                sin=sin,
+                query_token_stride_elems=query_token_stride,
+                key_token_stride_elems=key_token_stride,
+                head_stride_query_elems=head_stride_query,
+                head_stride_key_elems=head_stride_key,
+            )
+
+            use_grid2d, bpt2d = _predict_use_grid_2d(
+                num_tokens=num_tokens,
+                num_heads=case.num_heads,
+                num_kv_heads=(0 if key is None else case.num_kv_heads),
+                embed_dim_for_rotation=embed_dim_for_rotation,
+                can_vectorize_all=can_vec,
+                dtype=case.dtype,
+                interleaved=case.interleaved,
+            )
+
+            vec_str = "vec16" if can_vec else "scalar"
+            grid_str = "2D" if use_grid2d else "1D"
+            bpt2d_str = str(bpt2d)
+            row_str += f" | {vec_str:>6} | {grid_str:>4} | {bpt2d_str:>5}"
 
             if HAS_VLLM:
                 vllm_rope = vLLMRotaryEmbedding(
-                    head_size=head_size,
-                    rotary_dim=rotary_dim,
+                    head_size=case.head_size,
+                    rotary_dim=case.rotary_dim,
                     max_position_embeddings=max_seq_len,
                     base=10000,
-                    is_neox_style=True,
-                    dtype=dtype,
+                    is_neox_style=case.interleaved,
+                    dtype=case.dtype,
                 ).cuda()
 
                 def fn_vllm() -> None:
                     q = query.clone()
+                    if key is None:
+                        raise RuntimeError("vLLM requires key")
                     k = key.clone()
                     vllm_rope.forward_cuda(positions, q, k)
 
-                ms, _, _ = triton.testing.do_bench(fn_vllm, quantiles=[0.5, 0.2, 0.8])
-                vllm_time = 1000 * ms
-                row_str += f" | {vllm_time:10.4f}"
+                try:
+                    ms, _, _ = triton.testing.do_bench(fn_vllm, quantiles=[0.5, 0.2, 0.8])
+                    vllm_time = 1000 * ms
+                    row_str += f" | {vllm_time:10.4f}"
+                except Exception:
+                    row_str += f" | {'ERROR':>10}"
+                    vllm_time = None
 
             if HAS_FLASH_ATTN:
                 try:
-                    flash_rotary = FlashRotaryEmbedding(rotary_dim, device=device)
+                    flash_rotary = FlashRotaryEmbedding(case.rotary_dim, device=device)
                     qkv = torch.randn(
-                        batch_size,
+                        case.batch_size,
                         seq_len,
                         3,
-                        num_heads,
-                        head_size,
-                        dtype=dtype,
+                        case.num_heads,
+                        case.head_size,
+                        dtype=case.dtype,
                         device=device,
                     )
 
@@ -228,21 +703,18 @@ def fn_flash_attn() -> None:
                     row_str += f" | {'ERROR':>14}"
                     fa_time = None
 
-            if sgl_cos_sin_time is not None:
-                if HAS_VLLM and vllm_time is not None:
-                    speedup = vllm_time / sgl_cos_sin_time
-                    row_str += f" | {speedup:9.2f}x"
-                elif (not HAS_VLLM) and native_time is not None:
-                    speedup = native_time / sgl_cos_sin_time
-                    row_str += f" | {speedup:9.2f}x"
-                else:
-                    row_str += f" | {'N/A':>9}"
+            if (sgl_pos_time is not None) and (sgl_cos_sin_time is not None) and (sgl_cos_sin_time > 0):
+                speedup = sgl_pos_time / sgl_cos_sin_time
+                row_str += f" | {speedup:9.2f}x"
+            else:
+                row_str += f" | {'N/A':>9}"
 
             print(row_str)
             results.append(
                 {
                     "seq_len": seq_len,
                     "native": native_time,
+                    "naive": naive_time,
                     "sgl_cos_sin": sgl_cos_sin_time,
                     "sgl_pos": sgl_pos_time,
                     "vllm": vllm_time,
@@ -250,30 +722,42 @@ def fn_flash_attn() -> None:
                 }
             )
 
-        print("-" * 100)
-        print("\nAverage time:")
+        # Print averages as the last row of the table (instead of a multi-line block).
         native_vals = [r["native"] for r in results if r["native"] is not None]
-        if native_vals:
-            print(f"  Native: {np.mean(native_vals):.4f} ms")
+        naive_vals = [r["naive"] for r in results if r["naive"] is not None]
+        sgl_cos_sin_vals = [r["sgl_cos_sin"] for r in results if r["sgl_cos_sin"] is not None]
+        sgl_pos_vals = [r["sgl_pos"] for r in results if r["sgl_pos"] is not None]
+        vllm_vals = [r["vllm"] for r in results if r["vllm"] is not None]
+        fa_vals = [r["flash_attn"] for r in results if r["flash_attn"] is not None]
 
-        sgl_cos_sin_vals = [
-            r["sgl_cos_sin"] for r in results if r["sgl_cos_sin"] is not None
-        ]
-        if sgl_cos_sin_vals:
-            print(f"  SGLang (cos/sin API): {np.mean(sgl_cos_sin_vals):.4f} ms")
+        avg_row = f"{'AVG':>8}"
+        if case.bench_native:
+            avg_row += f" | {np.mean(native_vals):12.4f}" if native_vals else f" | {'N/A':>12}"
+        avg_row += f" | {np.mean(naive_vals):11.4f}" if naive_vals else f" | {'N/A':>11}"
+        avg_row += f" | {np.mean(sgl_cos_sin_vals):16.4f}" if sgl_cos_sin_vals else f" | {'N/A':>16}"
+        avg_row += f" | {np.mean(sgl_pos_vals):12.4f}" if sgl_pos_vals else f" | {'N/A':>12}"
+        avg_row += f" | {'-':>6} | {'-':>4} | {'-':>5}"
 
-        sgl_pos_vals = [r["sgl_pos"] for r in results if r["sgl_pos"] is not None]
-        if sgl_pos_vals:
-            print(f"  SGLang (positions API): {np.mean(sgl_pos_vals):.4f} ms")
         if HAS_VLLM:
-            vllm_vals = [r["vllm"] for r in results if r["vllm"] is not None]
-            if vllm_vals:
-                print(f"  vLLM: {np.mean(vllm_vals):.4f} ms")
+            avg_row += f" | {np.mean(vllm_vals):10.4f}" if vllm_vals else f" | {'N/A':>10}"
         if HAS_FLASH_ATTN:
-            fa_vals = [r["flash_attn"] for r in results if r["flash_attn"] is not None]
-            if fa_vals:
-                print(f"  flash_attn: {np.mean(fa_vals):.4f} ms")
+            avg_row += f" | {np.mean(fa_vals):14.4f}" if fa_vals else f" | {'N/A':>14}"
+        if sgl_pos_vals and sgl_cos_sin_vals and (np.mean(sgl_cos_sin_vals) > 0):
+            avg_speedup = float(np.mean(sgl_pos_vals) / np.mean(sgl_cos_sin_vals))
+            avg_row += f" | {avg_speedup:9.2f}x"
+        else:
+            avg_row += f" | {'N/A':>9}"
+
+        print("-" * 100)
+        print(avg_row)
 
 
 if __name__ == "__main__":
-    benchmark_mm_rotary_embedding()
+    parser = argparse.ArgumentParser(description="Benchmark rotary embedding ")
+    parser.add_argument(
+        "--case",
+        action="append",
+        default=None,
+        help="Run only the named case. If omitted, runs all cases.",
+    )
+    benchmark_mm_rotary_embedding(parser.parse_args())

From a62f591c53b4c06c92846b2bce27b24677313c6f Mon Sep 17 00:00:00 2001
From: TherLF <54900723+Ther-LF@users.noreply.github.com>
Date: Thu, 18 Dec 2025 12:25:50 +0800
Subject: [PATCH 16/37] Delete sgl-kernel/benchmark/profile_rotary_embedding.py

---
 .../benchmark/profile_rotary_embedding.py     | 103 ------------------
 1 file changed, 103 deletions(-)
 delete mode 100644 sgl-kernel/benchmark/profile_rotary_embedding.py

diff --git a/sgl-kernel/benchmark/profile_rotary_embedding.py b/sgl-kernel/benchmark/profile_rotary_embedding.py
deleted file mode 100644
index a8d981e655fb..000000000000
--- a/sgl-kernel/benchmark/profile_rotary_embedding.py
+++ /dev/null
@@ -1,103 +0,0 @@
-#!/usr/bin/env python
-# -*- coding: utf-8 -*-
-
-import sys
-
-import torch
-from sgl_kernel.rotary_embedding import rotary_embedding_cos_sin
-
-
-def compute_cos_sin_cache(
-    max_seq_len: int,
-    rotary_dim: int,
-    base: float = 10000.0,
-    dtype: torch.dtype = torch.float32,
-):
-    inv_freq = 1.0 / (
-        base ** (torch.arange(0, rotary_dim, 2, dtype=torch.float32) / rotary_dim)
-    )
-    t = torch.arange(max_seq_len, dtype=torch.float32)
-    freqs = torch.einsum("i,j->ij", t, inv_freq)
-    cos = freqs.cos().to(dtype)
-    sin = freqs.sin().to(dtype)
-    return cos, sin
-
-
-def profile_rotary_embedding():
-    # 配置参数：选取 Benchmark 中出现过的最大维度组合
-    # Batch Size: 32
-    # Seq Len: 8192 (Benchmark 列表中的最大值)
-    # Num Heads: 32 (常规配置，保证计算量)
-    # Num KV Heads: 32 (与 Heads 相同，最大化 Key 的负载)
-    # Head Size: 256 (Benchmark Configs 中的最大值)
-    batch_size = 1
-    seq_len = 8192
-    num_heads = 32
-    num_kv_heads = 8
-    head_size = 80
-    rotary_dim = head_size
-    dtype = torch.bfloat16
-    device = "cuda"
-
-    print(
-        f"Profiling Config: Batch={batch_size}, SeqLen={seq_len}, "
-        f"Heads={num_heads}, KV_Heads={num_kv_heads}, HeadSize={head_size}, Dtype={dtype}"
-    )
-
-    # 1. 准备 Cos/Sin Cache
-    try:
-        cos_cache, sin_cache = compute_cos_sin_cache(seq_len, rotary_dim, dtype=dtype)
-        cos_cache = cos_cache.to(device)
-        sin_cache = sin_cache.to(device)
-    except Exception as e:
-        print(f"Error creating cache: {e}")
-        return
-
-    # 2. 准备输入数据
-    # 计算总 token 数
-    num_tokens = batch_size * seq_len
-    print(f"Total tokens: {num_tokens}")
-
-    try:
-        # 生成 Query 和 Key
-        query = torch.randn(
-            num_tokens, num_heads * head_size, dtype=dtype, device=device
-        )
-        key = torch.randn(
-            num_tokens, num_kv_heads * head_size, dtype=dtype, device=device
-        )
-
-        # 生成对应的 Cos/Sin (Expand per token)
-        positions = torch.arange(seq_len, device=device, dtype=torch.int64).repeat(
-            batch_size
-        )
-
-        cos = cos_cache[positions]
-        sin = sin_cache[positions]
-    except torch.cuda.OutOfMemoryError:
-        print("OOM: The configuration is too large for the current GPU memory.")
-        return
-
-    # 3. Warmup
-    print("Warming up...")
-    for _ in range(5):
-        rotary_embedding_cos_sin(cos, sin, query, key, head_size, True)
-    torch.cuda.synchronize()
-
-    # 4. Profile Run (只运行一次，便于 NCU 捕获)
-    print("Running for profiling...")
-
-    # 开启 Profiler (如果使用 ncu --launch-count 1 ... python script.py，这行其实是可选的，
-    # 但显式开启有助于 ncu --profile-from-start off 模式)
-    torch.cuda.profiler.start()
-
-    rotary_embedding_cos_sin(cos, sin, query, key, head_size, True)
-
-    torch.cuda.profiler.stop()
-
-    torch.cuda.synchronize()
-    print("Done.")
-
-
-if __name__ == "__main__":
-    profile_rotary_embedding()

From 2610ae6002f32c3da811e65fff9f582d6e40f9a0 Mon Sep 17 00:00:00 2001
From: RubiaCx <1084281732@qq.com>
Date: Thu, 18 Dec 2025 05:10:59 -0800
Subject: [PATCH 17/37] [sgl-kernel]Optimize RoPE cos/sin kernel and clean up
 comments

---
 .../benchmark/bench_mm_rotary_embedding.py    |  147 +-
 .../sgl_diffusion/rope/rotary_embedding.cu    | 2032 +++++++++--------
 sgl-kernel/include/utils.h                    |   58 +-
 3 files changed, 1157 insertions(+), 1080 deletions(-)

diff --git a/sgl-kernel/benchmark/bench_mm_rotary_embedding.py b/sgl-kernel/benchmark/bench_mm_rotary_embedding.py
index 7962f68a524e..8927e7e98abf 100755
--- a/sgl-kernel/benchmark/bench_mm_rotary_embedding.py
+++ b/sgl-kernel/benchmark/bench_mm_rotary_embedding.py
@@ -31,28 +31,18 @@ def compute_cos_sin_cache_half(
 
 
 def _expand_neox_full_repeat(half: torch.Tensor) -> torch.Tensor:
-    """Expand [T, D] -> [T, 2D] by repeating each element twice: [a0,a0,a1,a1,...].
-
-    This matches some NeoX caches expanded to head_size for interleaved layout.
-    The CUDA kernel has a fast path that downsamples such 'full' caches to half.
-    """
-    # half: [T, D]
-    out = torch.empty((half.shape[0], half.shape[1] * 2), dtype=half.dtype, device=half.device)
+    out = torch.empty(
+        (half.shape[0], half.shape[1] * 2), dtype=half.dtype, device=half.device
+    )
     out[:, 0::2] = half
     out[:, 1::2] = half
     return out
 
-
 def _expand_llama_full_cat(half: torch.Tensor) -> torch.Tensor:
-    """Expand [T, D] -> [T, 2D] by concatenation: [a0,a1,..., a0,a1,...]."""
     return torch.cat([half, half], dim=-1)
 
-
 def _make_misaligned_contiguous_like(t: torch.Tensor, offset_elems: int = 1) -> torch.Tensor:
-    """Create a contiguous tensor with same shape/dtype/device but misaligned base pointer.
-
-    Achieved via a 1D buffer and a view with a storage offset.
-    """
+    # Create a contiguous tensor with same shape/dtype/device but misaligned base pointer.
     numel = t.numel()
     buf = torch.empty((numel + offset_elems,), device=t.device, dtype=t.dtype)
     out = buf[offset_elems : offset_elems + numel].view_as(t)
@@ -65,7 +55,6 @@ def _make_tensor(shape: Tuple[int, ...], *, dtype: torch.dtype, device: str, mis
     numel = int(np.prod(shape))
     if not misalign:
         return torch.randn(*shape, dtype=dtype, device=device)
-    # Misaligned but contiguous
     buf = torch.randn((numel + 1,), dtype=dtype, device=device)
     return buf[1 : 1 + numel].view(shape)
 
@@ -81,24 +70,13 @@ def _torch_naive_rope_inplace(
     head_size: int,
     interleaved: bool,
 ) -> None:
-    """Naive RoPE on GPU using PyTorch tensor ops (in-place on q/k).
-
-    Supports:
-      - interleaved=True  (NeoX-style [x0,y0,x1,y1,...])
-      - interleaved=False (LLaMA/GPT-J-style [x...][y...])
-      - q/k in 2D ([T, H*D]) or 3D ([T, H, D]) layouts
-      - k can be None (Q-only)
-    """
+    """Naive RoPE on GPU using PyTorch tensor ops (in-place on q/k)."""
 
     def _as_3d(x: torch.Tensor, h: int) -> torch.Tensor:
         if x.dim() == 3:
             return x
-        # 2D flattened
         return x.view(x.shape[0], h, head_size)
 
-    # Normalize cos/sin for interleaved fast-path compat:
-    # - interleaved expects cos/sin: [T, embed_dim] where embed_dim = rotary_dim/2
-    # - if provided as full head_size (repeat format), downsample like the CUDA kernel does
     if interleaved and cos.shape[1] == head_size:
         half = head_size // 2
         cos = cos.view(cos.shape[0], half, 2).select(2, 0).contiguous()
@@ -116,10 +94,8 @@ def _as_3d(x: torch.Tensor, h: int) -> torch.Tensor:
 
     def _apply(x: torch.Tensor, h: int) -> None:
         x3 = _as_3d(x, h)
-        # Only rotate the first rot_dim elements; leave tail intact.
         xr = x3[..., :rot_dim]
         if interleaved:
-            # [T,H,embed_dim,2]
             xr2 = xr.view(xr.shape[0], xr.shape[1], embed_dim, 2)
             x0 = xr2[..., 0]
             x1 = xr2[..., 1]
@@ -144,7 +120,7 @@ def _apply(x: torch.Tensor, h: int) -> None:
         _apply(k, num_kv_heads)
 
 
-def _predict_can_vectorize_all(
+def _predict_vec_flags(
     *,
     dtype: torch.dtype,
     interleaved: bool,
@@ -157,34 +133,49 @@ def _predict_can_vectorize_all(
     key_token_stride_elems: int,
     head_stride_query_elems: int,
     head_stride_key_elems: int,
-) -> bool:
+) -> Tuple[bool, bool]:
+    """Predict flags for the *optimized* kernel dispatch:
+    - use_vec_compute: vectorized compute loop (pairs_per_step) is allowed
+    - qk_aligned16: Q/K supports 16B vec load/store (otherwise use unaligned Q/K path)
+    """
     kVecBytes = 16
     elem_bytes = torch.tensor([], dtype=dtype).element_size()
     kElePerVec = kVecBytes // elem_bytes
     pairs_per_step = (kElePerVec // 2) if interleaved else kElePerVec
 
-    can = True
+    # A) vectorized compute viability (cos/sin alignment + embed_dim multiple)
+    use_vec_compute = True
     if embed_dim_for_rotation % pairs_per_step != 0:
-        can = False
-    if (query.data_ptr() % kVecBytes) != 0:
-        can = False
+        use_vec_compute = False
+
+    # cos/sin alignment & row stride for vector loads (float2/float4 reinterprets)
     if (cos.data_ptr() % kVecBytes) != 0:
-        can = False
+        use_vec_compute = False
     if (sin.data_ptr() % kVecBytes) != 0:
-        can = False
-    if key is not None and (key.data_ptr() % kVecBytes) != 0:
-        can = False
-
+        use_vec_compute = False
+    if ((cos.shape[1] * elem_bytes) % kVecBytes) != 0:
+        use_vec_compute = False
+    if ((sin.shape[1] * elem_bytes) % kVecBytes) != 0:
+        use_vec_compute = False
+
+    # B) Q/K 16B alignment (only affects aligned_qk template)
+    qk_aligned16 = True
+    if (query.data_ptr() % kVecBytes) != 0:
+        qk_aligned16 = False
     if (query_token_stride_elems * elem_bytes) % kVecBytes != 0:
-        can = False
+        qk_aligned16 = False
     if (head_stride_query_elems * elem_bytes) % kVecBytes != 0:
-        can = False
+        qk_aligned16 = False
+
     if key is not None:
+        if (key.data_ptr() % kVecBytes) != 0:
+            qk_aligned16 = False
         if (key_token_stride_elems * elem_bytes) % kVecBytes != 0:
-            can = False
+            qk_aligned16 = False
         if (head_stride_key_elems * elem_bytes) % kVecBytes != 0:
-            can = False
-    return can
+            qk_aligned16 = False
+
+    return use_vec_compute, qk_aligned16
 
 
 def _predict_use_grid_2d(
@@ -193,20 +184,23 @@ def _predict_use_grid_2d(
     num_heads: int,
     num_kv_heads: int,
     embed_dim_for_rotation: int,
-    can_vectorize_all: bool,
+    use_vec_compute: bool,
     dtype: torch.dtype,
     interleaved: bool,
 ) -> Tuple[bool, int]:
-    # Mirror the C++ launch config logic in rotary_embedding.cu (approx).
     kVecBytes = 16
     elem_bytes = torch.tensor([], dtype=dtype).element_size()
     kElePerVec = kVecBytes // elem_bytes
     pairs_per_step = (kElePerVec // 2) if interleaved else kElePerVec
-    launch_pairs_per_thread = pairs_per_step if can_vectorize_all else 1
+    launch_pairs_per_thread = pairs_per_step if use_vec_compute else 1
 
     max_pairs = max(num_heads * embed_dim_for_rotation, num_kv_heads * embed_dim_for_rotation)
     total_threads_needed = (max_pairs + launch_pairs_per_thread - 1) // launch_pairs_per_thread
-    threads_per_block_2d = min(256, max(128, embed_dim_for_rotation))
+
+    def _round_up32(x: int) -> int:
+        return ((x + 31) // 32) * 32
+
+    threads_per_block_2d = min(512, max(128, _round_up32(min(total_threads_needed, 512))))
     blocks_per_token_2d = (total_threads_needed + threads_per_block_2d - 1) // threads_per_block_2d
     use_grid_2d = (num_tokens <= 4) and (blocks_per_token_2d > 1)
     return use_grid_2d, blocks_per_token_2d
@@ -256,13 +250,12 @@ def benchmark_mm_rotary_embedding(args: argparse.Namespace) -> None:
 
     print(
         "Legend:\n"
-        "  - vec: vec16 => predicted to take the 16B vectorized path; scalar => predicted scalar/fallback path\n"
+        "  - vec: vec16 => predicted 16B-aligned vector path; vec16u => vectorized compute but unaligned Q/K load/store; scalar => scalar/fallback\n"
         "  - grid: 1D => one block per token; 2D => split one token across multiple blocks (only when num_tokens<=4)\n"
         "  - bpt2d: estimated blocks-per-token if 2D were used (useful to judge 'small grid but many blocks per token')\n"
     )
 
     cases: List[BenchCase] = [
-        # Baseline (vectorized, interleaved)
         BenchCase(
             name="base_bf16_int2d_k_halfcache",
             batch_size=1,
@@ -279,7 +272,6 @@ def benchmark_mm_rotary_embedding(args: argparse.Namespace) -> None:
             seq_lens=[1, 2, 4, 32, 256, 3015, 8192],
             bench_native=True,
         ),
-        # Non-interleaved (GPT-J/LLaMA-style)
         BenchCase(
             name="base_bf16_nonint2d_k_llamacache",
             batch_size=1,
@@ -296,7 +288,6 @@ def benchmark_mm_rotary_embedding(args: argparse.Namespace) -> None:
             seq_lens=[1, 4, 256, 8192],
             bench_native=False,
         ),
-        # Force non-vectorized path via misaligned Q pointer
         BenchCase(
             name="base_bf16_int2d_k_halfcache_misalignQ",
             batch_size=1,
@@ -313,7 +304,6 @@ def benchmark_mm_rotary_embedding(args: argparse.Namespace) -> None:
             seq_lens=[1, 4, 256, 8192],
             bench_native=False,
         ),
-        # Key is optional (Q-only)
         BenchCase(
             name="base_bf16_int2d_noK_halfcache",
             batch_size=1,
@@ -330,7 +320,6 @@ def benchmark_mm_rotary_embedding(args: argparse.Namespace) -> None:
             seq_lens=[1, 4, 256, 8192],
             bench_native=False,
         ),
-        # 3D layout path ([T, H, D])
         BenchCase(
             name="base_bf16_int3d_k_halfcache",
             batch_size=1,
@@ -347,7 +336,6 @@ def benchmark_mm_rotary_embedding(args: argparse.Namespace) -> None:
             seq_lens=[1, 4, 256, 8192],
             bench_native=False,
         ),
-        # fp16 dtype
         BenchCase(
             name="base_fp16_int2d_k_halfcache",
             batch_size=1,
@@ -364,7 +352,6 @@ def benchmark_mm_rotary_embedding(args: argparse.Namespace) -> None:
             seq_lens=[1, 4, 256, 8192],
             bench_native=False,
         ),
-        # Qwen-image observed config
         BenchCase(
             name="qwen_image_bf16_int2d_k_halfcache",
             batch_size=1,
@@ -381,7 +368,6 @@ def benchmark_mm_rotary_embedding(args: argparse.Namespace) -> None:
             seq_lens=[1, 4, 3015],
             bench_native=False,
         ),
-        # Grid-2D stress (num_tokens<=4 and blocks_per_token>1)
         BenchCase(
             name="grid2d_stress_bf16_int2d_k_neoxfullrepeat",
             batch_size=1,
@@ -398,7 +384,6 @@ def benchmark_mm_rotary_embedding(args: argparse.Namespace) -> None:
             seq_lens=[1, 2, 4],
             bench_native=False,
         ),
-        # Large batch representative
         BenchCase(
             name="batch32_bf16_int2d_k_halfcache",
             batch_size=32,
@@ -470,9 +455,7 @@ def benchmark_mm_rotary_embedding(args: argparse.Namespace) -> None:
             )
             print("-" * 100)
         except Exception as e:
-            print(
-                f"\nSkipping case {case.name}: {e}"
-            )
+            print(f"\nSkipping case {case.name}: {e}")
             continue
 
         header = f"{'seq_len':>8}"
@@ -504,13 +487,14 @@ def benchmark_mm_rotary_embedding(args: argparse.Namespace) -> None:
                 else:
                     raise ValueError(f"Unknown layout={case.layout}")
 
-                query = _make_tensor(query_shape, dtype=case.dtype, device=device, misalign=(case.misalign == "q"))
+                query = _make_tensor(
+                    query_shape, dtype=case.dtype, device=device, misalign=(case.misalign == "q")
+                )
                 key = _make_tensor(key_shape, dtype=case.dtype, device=device, misalign=False)
                 if case.key_mode == "no_k":
                     key = None
-                positions = torch.arange(
-                    seq_len, device=device, dtype=torch.int64
-                ).repeat(case.batch_size)
+
+                positions = torch.arange(seq_len, device=device, dtype=torch.int64).repeat(case.batch_size)
                 cos = cos_cache[positions]
                 sin = sin_cache[positions]
 
@@ -525,7 +509,7 @@ def benchmark_mm_rotary_embedding(args: argparse.Namespace) -> None:
                 sgl_pos_time = None
                 vllm_time = None
                 fa_time = None
-            except (RuntimeError, torch.AcceleratorError) as e:
+            except (RuntimeError, torch.AcceleratorError):
                 print(f"{seq_len:8d} | SKIP (CUDA error from previous iteration)")
                 torch.cuda.synchronize()
                 continue
@@ -533,7 +517,6 @@ def benchmark_mm_rotary_embedding(args: argparse.Namespace) -> None:
             def fn_native() -> None:
                 q = query.clone()
                 if key is None:
-                    # Native reference expects both q and k; skip in this mode.
                     raise RuntimeError("native requires key")
                 k = key.clone()
                 native_rope.forward_native(positions, q, k)
@@ -545,7 +528,6 @@ def fn_native() -> None:
             elif case.bench_native:
                 row_str += f" | {'N/A':>12}"
 
-            # Torch naive baseline (works for key=None / interleaved=False / 2D/3D layouts).
             @torch.no_grad()
             def fn_naive() -> None:
                 qn = query.clone()
@@ -570,9 +552,7 @@ def fn_sgl_cos_sin() -> None:
                 k = None if key is None else key.clone()
                 sgl_rotary_cos_sin(cos, sin, q, k, case.head_size, case.interleaved)
 
-            ms, _, _ = triton.testing.do_bench(
-                fn_sgl_cos_sin, quantiles=[0.5, 0.2, 0.8]
-            )
+            ms, _, _ = triton.testing.do_bench(fn_sgl_cos_sin, quantiles=[0.5, 0.2, 0.8])
             sgl_cos_sin_time = 1000 * ms
             row_str += f" | {sgl_cos_sin_time:16.4f}"
 
@@ -596,8 +576,7 @@ def fn_sgl_pos() -> None:
                 row_str += f" | {'ERROR':>12}"
                 sgl_pos_time = None
 
-            # Flags: predicted vectorization + whether C++ would pick grid2d (approx)
-            # Compute embed_dim_for_rotation consistent with C++ checks:
+            # Prediction flags (updated for optimized kernel)
             rot_dim_from_cache = int(cos.shape[1])
             embed_dim_for_rotation = rot_dim_from_cache if case.interleaved else (rot_dim_from_cache // 2)
 
@@ -623,7 +602,7 @@ def fn_sgl_pos() -> None:
                     key_token_stride = key_hidden
                     head_stride_key = int(key.stride(1))
 
-            can_vec = _predict_can_vectorize_all(
+            use_vec_compute, qk_aligned16 = _predict_vec_flags(
                 dtype=case.dtype,
                 interleaved=case.interleaved,
                 embed_dim_for_rotation=embed_dim_for_rotation,
@@ -642,15 +621,18 @@ def fn_sgl_pos() -> None:
                 num_heads=case.num_heads,
                 num_kv_heads=(0 if key is None else case.num_kv_heads),
                 embed_dim_for_rotation=embed_dim_for_rotation,
-                can_vectorize_all=can_vec,
+                use_vec_compute=use_vec_compute,
                 dtype=case.dtype,
                 interleaved=case.interleaved,
             )
 
-            vec_str = "vec16" if can_vec else "scalar"
+            if use_vec_compute:
+                vec_str = "vec16" if qk_aligned16 else "vec16u"
+            else:
+                vec_str = "scalar"
+
             grid_str = "2D" if use_grid2d else "1D"
-            bpt2d_str = str(bpt2d)
-            row_str += f" | {vec_str:>6} | {grid_str:>4} | {bpt2d_str:>5}"
+            row_str += f" | {vec_str:>6} | {grid_str:>4} | {str(bpt2d):>5}"
 
             if HAS_VLLM:
                 vllm_rope = vLLMRotaryEmbedding(
@@ -679,6 +661,8 @@ def fn_vllm() -> None:
 
             if HAS_FLASH_ATTN:
                 try:
+                    from flash_attn.layers.rotary import RotaryEmbedding as FlashRotaryEmbedding  # noqa
+
                     flash_rotary = FlashRotaryEmbedding(case.rotary_dim, device=device)
                     qkv = torch.randn(
                         case.batch_size,
@@ -694,9 +678,7 @@ def fn_flash_attn() -> None:
                         qkv_fa = qkv.clone()
                         flash_rotary(qkv_fa, seqlen_offset=0)
 
-                    ms, _, _ = triton.testing.do_bench(
-                        fn_flash_attn, quantiles=[0.5, 0.2, 0.8]
-                    )
+                    ms, _, _ = triton.testing.do_bench(fn_flash_attn, quantiles=[0.5, 0.2, 0.8])
                     fa_time = 1000 * ms
                     row_str += f" | {fa_time:14.4f}"
                 except Exception:
@@ -722,7 +704,6 @@ def fn_flash_attn() -> None:
                 }
             )
 
-        # Print averages as the last row of the table (instead of a multi-line block).
         native_vals = [r["native"] for r in results if r["native"] is not None]
         naive_vals = [r["naive"] for r in results if r["naive"] is not None]
         sgl_cos_sin_vals = [r["sgl_cos_sin"] for r in results if r["sgl_cos_sin"] is not None]
diff --git a/sgl-kernel/csrc/sgl_diffusion/rope/rotary_embedding.cu b/sgl-kernel/csrc/sgl_diffusion/rope/rotary_embedding.cu
index ad8885c46d9c..7b8c57d964c8 100644
--- a/sgl-kernel/csrc/sgl_diffusion/rope/rotary_embedding.cu
+++ b/sgl-kernel/csrc/sgl_diffusion/rope/rotary_embedding.cu
@@ -15,997 +15,1041 @@
  * limitations under the License.
  */
 
- #include <ATen/cuda/CUDAContext.h>
- #include <c10/cuda/CUDAException.h>
- #include <c10/cuda/CUDAGuard.h>
- #include <c10/cuda/CUDAStream.h>
- #include <torch/all.h>
- 
- #include <cmath>
- #include <vector>
- 
- #include "utils.h"
- 
- // -------------------------------------------------------------------------
- // Helper structures and functions for manual pipeline interleaving
- // -------------------------------------------------------------------------
- 
- template <typename scalar_t, bool interleaved>
- struct RotaryVecData;
- 
- template <typename scalar_t>
- struct RotaryVecData<scalar_t, true> {
-   float4 vec;
-   float2 cos_vec;
-   float2 sin_vec;
- };
- 
- template <typename scalar_t>
- struct RotaryVecData<scalar_t, false> {
-   float4 vec_x;
-   float4 vec_y;
-   float4 cos_x;
-   float4 sin_x;
-   float4 cos_y;
-   float4 sin_y;
- };
- 
- template <typename scalar_t, bool interleaved>
- inline __device__ RotaryVecData<scalar_t, interleaved> load_rotary_vec(
-     const scalar_t* __restrict__ arr,
-     const scalar_t* __restrict__ cos_ptr,
-     const scalar_t* __restrict__ sin_ptr,
-     int rot_offset,
-     int embed_dim) {
-   RotaryVecData<scalar_t, interleaved> data;
-   if constexpr (interleaved) {
-     data.vec = *reinterpret_cast<const float4*>(arr + rot_offset * 2);
-     data.cos_vec = *reinterpret_cast<const float2*>(cos_ptr + rot_offset);
-     data.sin_vec = *reinterpret_cast<const float2*>(sin_ptr + rot_offset);
-   } else {
-     data.vec_x = *reinterpret_cast<const float4*>(arr + rot_offset);
-     data.vec_y = *reinterpret_cast<const float4*>(arr + rot_offset + embed_dim);
-     data.cos_x = *reinterpret_cast<const float4*>(cos_ptr + rot_offset);
-     data.sin_x = *reinterpret_cast<const float4*>(sin_ptr + rot_offset);
-     data.cos_y = *reinterpret_cast<const float4*>(cos_ptr + rot_offset + embed_dim);
-     data.sin_y = *reinterpret_cast<const float4*>(sin_ptr + rot_offset + embed_dim);
-   }
-   return data;
- }
- 
- template <typename scalar_t, bool interleaved>
- inline __device__ void compute_store_rotary_vec(
-     scalar_t* __restrict__ arr, const RotaryVecData<scalar_t, interleaved>& data, int rot_offset, int embed_dim) {
-   constexpr int kVecBytes = 16;
-   constexpr int kElePerVec = kVecBytes / sizeof(scalar_t);
- 
-   if constexpr (interleaved) {
-     union VecUnion {
-       float4 vec;
-       scalar_t elems[kElePerVec];
-     };
-     union CosSinVec {
-       float2 vec;
-       scalar_t elems[kElePerVec / 2];
-     };
- 
-     VecUnion v;
-     v.vec = data.vec;
-     CosSinVec c, s;
-     c.vec = data.cos_vec;
-     s.vec = data.sin_vec;
- 
- #pragma unroll
-     for (int i = 0; i < kElePerVec; i += 2) {
-       int idx = i / 2;
-       float cos_val = static_cast<float>(c.elems[idx]);
-       float sin_val = static_cast<float>(s.elems[idx]);
-       float x = static_cast<float>(v.elems[i]);
-       float y = static_cast<float>(v.elems[i + 1]);
-       v.elems[i] = static_cast<scalar_t>(x * cos_val - y * sin_val);
-       v.elems[i + 1] = static_cast<scalar_t>(y * cos_val + x * sin_val);
-     }
-     *reinterpret_cast<float4*>(arr + rot_offset * 2) = v.vec;
- 
-   } else {
-     union VecUnion {
-       float4 vec;
-       scalar_t elems[kElePerVec];
-     };
-     union CosSinVec {
-       float4 vec;
-       scalar_t elems[kElePerVec];
-     };
- 
-     VecUnion vx, vy;
-     vx.vec = data.vec_x;
-     vy.vec = data.vec_y;
-     CosSinVec cx, sx, cy, sy;
-     cx.vec = data.cos_x;
-     sx.vec = data.sin_x;
-     cy.vec = data.cos_y;
-     sy.vec = data.sin_y;
- 
- #pragma unroll
-     for (int i = 0; i < kElePerVec; ++i) {
-       float cos_val_x = static_cast<float>(cx.elems[i]);
-       float sin_val_x = static_cast<float>(sx.elems[i]);
-       float cos_val_y = static_cast<float>(cy.elems[i]);
-       float sin_val_y = static_cast<float>(sy.elems[i]);
-       float x = static_cast<float>(vx.elems[i]);
-       float y = static_cast<float>(vy.elems[i]);
- 
-       vx.elems[i] = static_cast<scalar_t>(x * cos_val_x - y * sin_val_x);
-       vy.elems[i] = static_cast<scalar_t>(y * cos_val_y + x * sin_val_y);
-     }
-     *reinterpret_cast<float4*>(arr + rot_offset) = vx.vec;
-     *reinterpret_cast<float4*>(arr + rot_offset + embed_dim) = vy.vec;
-   }
- }
- 
- template <typename scalar_t, bool interleaved>
- inline __device__ void apply_token_rotary_embedding(
-     scalar_t* __restrict__ arr,            // [head_size]
-     const scalar_t* __restrict__ cos_ptr,  // [rot_dim]
-     const scalar_t* __restrict__ sin_ptr,  // [rot_dim]
-     int rot_offset,
-     int embed_dim) {  // for non-interleaved: half dim
-   if constexpr (interleaved) {
-     // NeoX-style: interleaved layout [x0, y0, x1, y1, ...].
-     // cos/sin: [..., rotary_dim/2], one entry per (x, y) pair.
-     const int x_index = 2 * rot_offset;
-     const int y_index = x_index + 1;
- 
-     // Directly access Global Memory
-     const float cos_val = static_cast<float>(cos_ptr[rot_offset]);
-     const float sin_val = static_cast<float>(sin_ptr[rot_offset]);
- 
-     const float x = static_cast<float>(arr[x_index]);
-     const float y = static_cast<float>(arr[y_index]);
-     arr[x_index] = static_cast<scalar_t>(x * cos_val - y * sin_val);
-     arr[y_index] = static_cast<scalar_t>(y * cos_val + x * sin_val);
-   } else {
-     // GPT-J / LLaMA style: layout [x0, x1, ..., y0, y1, ...]
-     // cos/sin: [..., rotary_dim], one entry per (x, y) pair.
-     const int x_index = rot_offset;
-     const int y_index = rot_offset + embed_dim;
- 
-     // Directly access Global Memory
-     const float cos_val_x = static_cast<float>(cos_ptr[rot_offset]);
-     const float sin_val_x = static_cast<float>(sin_ptr[rot_offset]);
-     const float cos_val_y = static_cast<float>(cos_ptr[rot_offset + embed_dim]);
-     const float sin_val_y = static_cast<float>(sin_ptr[rot_offset + embed_dim]);
- 
-     const float x = static_cast<float>(arr[x_index]);
-     const float y = static_cast<float>(arr[y_index]);
-     arr[x_index] = static_cast<scalar_t>(x * cos_val_x - y * sin_val_x);
-     arr[y_index] = static_cast<scalar_t>(y * cos_val_y + x * sin_val_y);
-   }
- }
- 
- template <typename scalar_t, bool interleaved>
- inline __device__ void apply_token_rotary_embedding_vec(
-     scalar_t* __restrict__ arr,            // [head_size]
-     const scalar_t* __restrict__ cos_ptr,  // [rot_dim]
-     const scalar_t* __restrict__ sin_ptr,  // [rot_dim]
-     int rot_offset,
-     int embed_dim) {
-   using vec_t = float4;
-   constexpr int kVecBytes = sizeof(vec_t);
-   constexpr int kScalarBytes = sizeof(scalar_t);
-   constexpr int kElePerVec = kVecBytes / kScalarBytes;
- 
-   // Union for type punning to avoid strict aliasing issues with reinterpret_cast
-   union VecUnion {
-     vec_t vec;
-     scalar_t elems[kElePerVec];
-   };
- 
-   if constexpr (interleaved) {
-     // Interleaved: arr has [x0, y0, x1, y1...]
-     // A single vector load contains 'kElePerVec' elements, which is 'kElePerVec / 2' pairs.
-     // rot_offset is the index of the PAIR.
-     // Address in arr is rot_offset * 2.
- 
-     VecUnion data;
-     data.vec = *reinterpret_cast<const vec_t*>(arr + rot_offset * 2);
- 
-     // Vectorized load for cos/sin: always 8 bytes (half of vector size, float2)
-     // float2 covers (kElePerVec / 2) scalar_t elements (pairs)
-     float2 cos_vec = *reinterpret_cast<const float2*>(cos_ptr + rot_offset);
-     float2 sin_vec = *reinterpret_cast<const float2*>(sin_ptr + rot_offset);
- 
-     union CosSinVec {
-       float2 vec;
-       scalar_t elems[kElePerVec / 2];
-     };
-     CosSinVec cos_u, sin_u;
-     cos_u.vec = cos_vec;
-     sin_u.vec = sin_vec;
- 
- #pragma unroll
-     for (int i = 0; i < kElePerVec; i += 2) {
-       // data.elems[i] is x, data.elems[i+1] is y
-       // They correspond to pair index: rot_offset + (i / 2)
-       int idx = i / 2;
-       float cos_val = static_cast<float>(cos_u.elems[idx]);
-       float sin_val = static_cast<float>(sin_u.elems[idx]);
- 
-       float x = static_cast<float>(data.elems[i]);
-       float y = static_cast<float>(data.elems[i + 1]);
- 
-       data.elems[i] = static_cast<scalar_t>(x * cos_val - y * sin_val);
-       data.elems[i + 1] = static_cast<scalar_t>(y * cos_val + x * sin_val);
-     }
- 
-     *reinterpret_cast<vec_t*>(arr + rot_offset * 2) = data.vec;
- 
-   } else {
-     // Non-interleaved: X and Y are separated by embed_dim.
-     // We process 'kElePerVec' PAIRS at once.
-     // Load X vector and Y vector.
- 
-     VecUnion data_x, data_y;
-     data_x.vec = *reinterpret_cast<const vec_t*>(arr + rot_offset);
-     data_y.vec = *reinterpret_cast<const vec_t*>(arr + rot_offset + embed_dim);
- 
-     // Vectorized load for cos/sin: always 16 bytes (full vector size, float4)
-     // float4 covers kElePerVec scalar_t elements
-     float4 cos_vec_x = *reinterpret_cast<const float4*>(cos_ptr + rot_offset);
-     float4 sin_vec_x = *reinterpret_cast<const float4*>(sin_ptr + rot_offset);
-     float4 cos_vec_y = *reinterpret_cast<const float4*>(cos_ptr + rot_offset + embed_dim);
-     float4 sin_vec_y = *reinterpret_cast<const float4*>(sin_ptr + rot_offset + embed_dim);
- 
-     union CosSinVec {
-       float4 vec;
-       scalar_t elems[kElePerVec];
-     };
-     CosSinVec cos_u_x, sin_u_x, cos_u_y, sin_u_y;
-     cos_u_x.vec = cos_vec_x;
-     sin_u_x.vec = sin_vec_x;
-     cos_u_y.vec = cos_vec_y;
-     sin_u_y.vec = sin_vec_y;
- 
- #pragma unroll
-     for (int i = 0; i < kElePerVec; ++i) {
-       float cos_val_x = static_cast<float>(cos_u_x.elems[i]);
-       float sin_val_x = static_cast<float>(sin_u_x.elems[i]);
-       float cos_val_y = static_cast<float>(cos_u_y.elems[i]);
-       float sin_val_y = static_cast<float>(sin_u_y.elems[i]);
- 
-       float x = static_cast<float>(data_x.elems[i]);
-       float y = static_cast<float>(data_y.elems[i]);
- 
-       data_x.elems[i] = static_cast<scalar_t>(x * cos_val_x - y * sin_val_x);
-       data_y.elems[i] = static_cast<scalar_t>(y * cos_val_y + x * sin_val_y);
-     }
- 
-     *reinterpret_cast<vec_t*>(arr + rot_offset) = data_x.vec;
-     *reinterpret_cast<vec_t*>(arr + rot_offset + embed_dim) = data_y.vec;
-   }
- }
- 
- template <>
- inline __device__ void apply_token_rotary_embedding<float, true>(
-     float* __restrict__ arr,
-     const float* __restrict__ cos_ptr,
-     const float* __restrict__ sin_ptr,
-     int rot_offset,
-     int /*embed_dim*/) {
-   float2 xy = *reinterpret_cast<const float2*>(arr + rot_offset * 2);
- 
-   const float cos_val = static_cast<float>(cos_ptr[rot_offset]);
-   const float sin_val = static_cast<float>(sin_ptr[rot_offset]);
- 
-   float2 out;
-   out.x = xy.x * cos_val - xy.y * sin_val;
-   out.y = xy.y * cos_val + xy.x * sin_val;
- 
-   *reinterpret_cast<float2*>(arr + rot_offset * 2) = out;
- }
- 
- // 2D grid kernel: parallel over tokens (grid.x) and pair-tiles (grid.y)
- template <typename scalar_t, bool interleaved, bool vectorized, int ROT_EMBED_DIM>
- __global__ void rotary_embedding_kernel_2d(
-     const scalar_t* __restrict__ cos_data,  // [num_tokens, rot_dim_arg]
-     const scalar_t* __restrict__ sin_data,  // [num_tokens, rot_dim_arg]
-     scalar_t* __restrict__ query_total,
-     scalar_t* __restrict__ key_total,
-     const int rot_dim_arg,
-     const int embed_dim_for_rotation_arg,
-     const int64_t query_token_stride,
-     const int64_t key_token_stride,
-     const int64_t head_stride_query,
-     const int64_t head_stride_key,
-     const int num_heads,
-     const int num_kv_heads,
-     const int head_size_arg,
-     const int blocks_per_token) {
-   // Removed shared memory allocation and loading logic
- 
-   const int token_idx = blockIdx.x;
-   if (token_idx >= gridDim.x) {
-     return;
-   }
- 
-   // Pointers to Global Memory for the current token
-   const scalar_t* current_token_cos_ptr = cos_data + token_idx * rot_dim_arg;
-   const scalar_t* current_token_sin_ptr = sin_data + token_idx * rot_dim_arg;
- 
-   constexpr int kVecBytes = 16;
-   const int scalar_size = sizeof(scalar_t);
- 
-   // Removed SMEM loading and __syncthreads()
- 
-   scalar_t* query_for_token = query_total + token_idx * (int)query_token_stride;
-   scalar_t* key_for_token = (key_total != nullptr) ? (key_total + token_idx * (int)key_token_stride) : nullptr;
- 
-   const int local_block_idx = blockIdx.y;
- 
-   // Use compile-time constant if ROT_EMBED_DIM > 0, otherwise fallback to runtime argument
-   const int embed_dim_for_rotation = (ROT_EMBED_DIM > 0) ? ROT_EMBED_DIM : embed_dim_for_rotation_arg;
- 
-   if constexpr (vectorized) {
-     using vec_t = float4;
-     constexpr int kElePerVec = kVecBytes / sizeof(scalar_t);
-     constexpr int pairs_per_step = interleaved ? (kElePerVec / 2) : kElePerVec;
- 
-     const int pair_stride = blockDim.x * blocks_per_token * pairs_per_step;
-     int i = (local_block_idx * blockDim.x + threadIdx.x) * pairs_per_step;
-     const int nq_pairs = num_heads * embed_dim_for_rotation;
- 
-     // PROLOGUE: Preload 0-th batch
-     RotaryVecData<scalar_t, interleaved> curr_data;
-     if (i < nq_pairs) {
-       int head_idx = i / embed_dim_for_rotation;
-       int rot_offset = i % embed_dim_for_rotation;
-       scalar_t* ptr = query_for_token + head_idx * (int)head_stride_query;
-       curr_data = load_rotary_vec<scalar_t, interleaved>(
-           ptr, current_token_cos_ptr, current_token_sin_ptr, rot_offset, embed_dim_for_rotation);
-     }
- 
-     int next_i = i + pair_stride;
- 
-     // MAIN LOOP
-     for (; i < nq_pairs; i += pair_stride, next_i += pair_stride) {
-       // 1. PREFETCH NEXT
-       RotaryVecData<scalar_t, interleaved> next_data;
-       bool active_next = (next_i < nq_pairs);
-       if (active_next) {
-         int head_idx_next = next_i / embed_dim_for_rotation;
-         int rot_offset_next = next_i % embed_dim_for_rotation;
-         scalar_t* ptr_next = query_for_token + head_idx_next * (int)head_stride_query;
-         next_data = load_rotary_vec<scalar_t, interleaved>(
-             ptr_next, current_token_cos_ptr, current_token_sin_ptr, rot_offset_next, embed_dim_for_rotation);
-       }
- 
-       // 2. COMPUTE CURRENT
-       // The compute uses curr_data which corresponds to index 'i'
-       int head_idx = i / embed_dim_for_rotation;
-       int rot_offset = i % embed_dim_for_rotation;
-       scalar_t* ptr = query_for_token + head_idx * (int)head_stride_query;
-       compute_store_rotary_vec<scalar_t, interleaved>(ptr, curr_data, rot_offset, embed_dim_for_rotation);
- 
-       // 3. SHIFT
-       curr_data = next_data;
-     }
- 
-     if (key_for_token != nullptr) {
-       const int nk_pairs = num_kv_heads * embed_dim_for_rotation;
-       int k_i = (local_block_idx * blockDim.x + threadIdx.x) * pairs_per_step;
- 
-       // PROLOGUE (Key)
-       RotaryVecData<scalar_t, interleaved> curr_data_k;
-       if (k_i < nk_pairs) {
-         int head_idx = k_i / embed_dim_for_rotation;
-         int rot_offset = k_i % embed_dim_for_rotation;
-         scalar_t* ptr = key_for_token + head_idx * (int)head_stride_key;
-         curr_data_k = load_rotary_vec<scalar_t, interleaved>(
-             ptr, current_token_cos_ptr, current_token_sin_ptr, rot_offset, embed_dim_for_rotation);
-       }
- 
-       int next_k_i = k_i + pair_stride;
- 
-       // MAIN LOOP (Key)
-       for (; k_i < nk_pairs; k_i += pair_stride, next_k_i += pair_stride) {
-         // 1. PREFETCH NEXT
-         RotaryVecData<scalar_t, interleaved> next_data_k;
-         bool active_next = (next_k_i < nk_pairs);
-         if (active_next) {
-           int head_idx_next = next_k_i / embed_dim_for_rotation;
-           int rot_offset_next = next_k_i % embed_dim_for_rotation;
-           scalar_t* ptr_next = key_for_token + head_idx_next * (int)head_stride_key;
-           next_data_k = load_rotary_vec<scalar_t, interleaved>(
-               ptr_next, current_token_cos_ptr, current_token_sin_ptr, rot_offset_next, embed_dim_for_rotation);
-         }
- 
-         // 2. COMPUTE CURRENT
-         int head_idx = k_i / embed_dim_for_rotation;
-         int rot_offset = k_i % embed_dim_for_rotation;
-         scalar_t* ptr = key_for_token + head_idx * (int)head_stride_key;
-         compute_store_rotary_vec<scalar_t, interleaved>(ptr, curr_data_k, rot_offset, embed_dim_for_rotation);
- 
-         // 3. SHIFT
-         curr_data_k = next_data_k;
-       }
-     }
-   } else {
-     // Fallback to scalar implementation
-     const int pair_stride = blockDim.x * blocks_per_token;
-     const int thread_pair_offset = local_block_idx * blockDim.x + threadIdx.x;
- 
-     const int nq_pairs = num_heads * embed_dim_for_rotation;
-     for (int i = thread_pair_offset; i < nq_pairs; i += pair_stride) {
-       const int head_idx = i / embed_dim_for_rotation;
-       const int rot_offset = i % embed_dim_for_rotation;
- 
-       scalar_t* query_for_token_head = query_for_token + head_idx * (int)head_stride_query;
-       apply_token_rotary_embedding<scalar_t, interleaved>(
-           query_for_token_head, current_token_cos_ptr, current_token_sin_ptr, rot_offset, embed_dim_for_rotation);
-     }
- 
-     if (key_for_token != nullptr) {
-       const int nk_pairs = num_kv_heads * embed_dim_for_rotation;
-       for (int i = thread_pair_offset; i < nk_pairs; i += pair_stride) {
-         const int head_idx = i / embed_dim_for_rotation;
-         const int rot_offset = i % embed_dim_for_rotation;
- 
-         scalar_t* key_for_token_head = key_for_token + head_idx * (int)head_stride_key;
-         apply_token_rotary_embedding<scalar_t, interleaved>(
-             key_for_token_head, current_token_cos_ptr, current_token_sin_ptr, rot_offset, embed_dim_for_rotation);
-       }
-     }
-   }
- }
- 
- // 1D grid kernel: each block handles one token
- template <typename scalar_t, bool interleaved, bool vectorized, int ROT_EMBED_DIM>
- __launch_bounds__(512) __global__ void rotary_embedding_kernel_1d(
-     const scalar_t* __restrict__ cos_data,  // [num_tokens, rot_dim_arg]
-     const scalar_t* __restrict__ sin_data,  // [num_tokens, rot_dim_arg]
-     scalar_t* __restrict__ query_total,
-     scalar_t* __restrict__ key_total,
-     const int rot_dim_arg,
-     const int embed_dim_for_rotation_arg,
-     const int64_t query_token_stride,
-     const int64_t key_token_stride,
-     const int64_t head_stride_query,
-     const int64_t head_stride_key,
-     const int num_heads,
-     const int num_kv_heads,
-     const int head_size_arg) {
-   // Removed shared memory allocation and loading logic
- 
-   const int token_idx = blockIdx.x;
-   if (token_idx >= gridDim.x) return;
- 
-   const scalar_t* current_token_cos_ptr = cos_data + token_idx * rot_dim_arg;
-   const scalar_t* current_token_sin_ptr = sin_data + token_idx * rot_dim_arg;
- 
-   constexpr int kVecBytes = 16;
-   const int scalar_size = sizeof(scalar_t);
- 
-   // Removed SMEM loading and __syncthreads()
- 
-   scalar_t* query_for_token = query_total + token_idx * (int)query_token_stride;
-   scalar_t* key_for_token = (key_total != nullptr) ? (key_total + token_idx * (int)key_token_stride) : nullptr;
- 
-   // Use compile-time constant if ROT_EMBED_DIM > 0, otherwise fallback to runtime argument
-   const int embed_dim_for_rotation = (ROT_EMBED_DIM > 0) ? ROT_EMBED_DIM : embed_dim_for_rotation_arg;
- 
-   if constexpr (vectorized) {
-     constexpr int kVecBytes = 16;
-     constexpr int kElePerVec = kVecBytes / sizeof(scalar_t);
-     // Interleaved: 1 vec = kElePerVec elements = kElePerVec/2 pairs
-     // Non-Interleaved: 1 vec (X) + 1 vec (Y) = kElePerVec elements = kElePerVec pairs
-     constexpr int pairs_per_step = interleaved ? (kElePerVec / 2) : kElePerVec;
-     const int nq_pairs = num_heads * embed_dim_for_rotation;
- 
-     // Original Stride
-     const int stride = blockDim.x * pairs_per_step;
- 
-     int i = threadIdx.x * pairs_per_step;
- 
-     // ==========================================
-     // PROLOGUE: Preload first batch
-     // ==========================================
-     RotaryVecData<scalar_t, interleaved> curr_data;
- 
-     if (i < nq_pairs) {
-       int head_idx = i / embed_dim_for_rotation;
-       int rot_offset = i % embed_dim_for_rotation;
-       scalar_t* ptr = query_for_token + head_idx * (int)head_stride_query;
-       curr_data = load_rotary_vec<scalar_t, interleaved>(
-           ptr, current_token_cos_ptr, current_token_sin_ptr, rot_offset, embed_dim_for_rotation);
-     }
- 
-     int next_i = i + stride;
- 
-     // ==========================================
-     // MAIN LOOP
-     // ==========================================
-     for (; i < nq_pairs; i += stride, next_i += stride) {
-       // --- 1. PREFETCH NEXT BATCH ---
-       RotaryVecData<scalar_t, interleaved> next_data;
- 
-       if (next_i < nq_pairs) {
-         int head_idx = next_i / embed_dim_for_rotation;
-         int rot_offset = next_i % embed_dim_for_rotation;
-         scalar_t* ptr = query_for_token + head_idx * (int)head_stride_query;
-         next_data = load_rotary_vec<scalar_t, interleaved>(
-             ptr, current_token_cos_ptr, current_token_sin_ptr, rot_offset, embed_dim_for_rotation);
-       }
- 
-       // --- 2. COMPUTE CURRENT BATCH ---
-       if (i < nq_pairs) {
-         int head_idx = i / embed_dim_for_rotation;
-         int rot_offset = i % embed_dim_for_rotation;
-         scalar_t* ptr = query_for_token + head_idx * (int)head_stride_query;
-         compute_store_rotary_vec<scalar_t, interleaved>(ptr, curr_data, rot_offset, embed_dim_for_rotation);
-       }
- 
-       // --- 3. SHIFT ---
-       curr_data = next_data;
-     }
- 
-     // Key branch
-     if (key_for_token != nullptr) {
-       const int nk_pairs = num_kv_heads * embed_dim_for_rotation;
-       int k_i = threadIdx.x * pairs_per_step;
- 
-       // PROLOGUE (Key): Preload first batch
-       RotaryVecData<scalar_t, interleaved> curr_data_k;
- 
-       if (k_i < nk_pairs) {
-         int head_idx = k_i / embed_dim_for_rotation;
-         int rot_offset = k_i % embed_dim_for_rotation;
-         scalar_t* ptr = key_for_token + head_idx * (int)head_stride_key;
-         curr_data_k = load_rotary_vec<scalar_t, interleaved>(
-             ptr, current_token_cos_ptr, current_token_sin_ptr, rot_offset, embed_dim_for_rotation);
-       }
- 
-       int next_k_i = k_i + stride;
- 
-       // MAIN LOOP (Key)
-       for (; k_i < nk_pairs; k_i += stride, next_k_i += stride) {
-         // --- 1. PREFETCH NEXT BATCH ---
-         RotaryVecData<scalar_t, interleaved> next_data_k;
- 
-         if (next_k_i < nk_pairs) {
-           int head_idx = next_k_i / embed_dim_for_rotation;
-           int rot_offset = next_k_i % embed_dim_for_rotation;
-           scalar_t* ptr = key_for_token + head_idx * (int)head_stride_key;
-           next_data_k = load_rotary_vec<scalar_t, interleaved>(
-               ptr, current_token_cos_ptr, current_token_sin_ptr, rot_offset, embed_dim_for_rotation);
-         }
- 
-         // --- 2. COMPUTE CURRENT BATCH ---
-         if (k_i < nk_pairs) {
-           int head_idx = k_i / embed_dim_for_rotation;
-           int rot_offset = k_i % embed_dim_for_rotation;
-           scalar_t* ptr = key_for_token + head_idx * (int)head_stride_key;
-           compute_store_rotary_vec<scalar_t, interleaved>(ptr, curr_data_k, rot_offset, embed_dim_for_rotation);
-         }
- 
-         // --- 3. SHIFT ---
-         curr_data_k = next_data_k;
-       }
-     }
-   } else {
-     // Fallback scalar
-     const int nq_pairs = num_heads * embed_dim_for_rotation;
-     for (int i = threadIdx.x; i < nq_pairs; i += blockDim.x) {
-       const int head_idx = i / embed_dim_for_rotation;
-       const int rot_offset = i % embed_dim_for_rotation;
-       scalar_t* query_for_token_head = query_for_token + head_idx * (int)head_stride_query;
- 
-       apply_token_rotary_embedding<scalar_t, interleaved>(
-           query_for_token_head, current_token_cos_ptr, current_token_sin_ptr, rot_offset, embed_dim_for_rotation);
-     }
- 
-     if (key_for_token != nullptr) {
-       const int nk_pairs = num_kv_heads * embed_dim_for_rotation;
-       for (int i = threadIdx.x; i < nk_pairs; i += blockDim.x) {
-         const int head_idx = i / embed_dim_for_rotation;
-         const int rot_offset = i % embed_dim_for_rotation;
- 
-         scalar_t* key_for_token_head = key_for_token + head_idx * (int)head_stride_key;
-         apply_token_rotary_embedding<scalar_t, interleaved>(
-             key_for_token_head, current_token_cos_ptr, current_token_sin_ptr, rot_offset, embed_dim_for_rotation);
-       }
-     }
-   }
- }
- 
- // Define dispatch macro for Rotation Dimension (Pair Count)
- // Note: This is usually HEAD_SIZE / 2 for LLaMA-like models.
- // Case 32 -> head_size 64
- // Case 64 -> head_size 128
- // Case 128 -> head_size 256
- #define DISPATCH_ROT_DIM(embed_dim, ...)  \
-   switch (embed_dim) {                    \
-     case 16: {                            \
-       constexpr int ROT_EMBED_DIM = 16;   \
-       __VA_ARGS__;                        \
-       break;                              \
-     }                                     \
-     case 32: {                            \
-       constexpr int ROT_EMBED_DIM = 32;   \
-       __VA_ARGS__;                        \
-       break;                              \
-     }                                     \
-     case 64: {                            \
-       constexpr int ROT_EMBED_DIM = 64;   \
-       __VA_ARGS__;                        \
-       break;                              \
-     }                                     \
-     case 96: {                            \
-       constexpr int ROT_EMBED_DIM = 96;   \
-       __VA_ARGS__;                        \
-       break;                              \
-     }                                     \
-     case 128: {                           \
-       constexpr int ROT_EMBED_DIM = 128;  \
-       __VA_ARGS__;                        \
-       break;                              \
-     }                                     \
-     case 256: {                           \
-       constexpr int ROT_EMBED_DIM = 256;  \
-       __VA_ARGS__;                        \
-       break;                              \
-     }                                     \
-     case 512: {                           \
-       constexpr int ROT_EMBED_DIM = 512;  \
-       __VA_ARGS__;                        \
-       break;                              \
-     }                                     \
-     case 1024: {                          \
-       constexpr int ROT_EMBED_DIM = 1024; \
-       __VA_ARGS__;                        \
-       break;                              \
-     }                                     \
-     default: {                            \
-       constexpr int ROT_EMBED_DIM = 0;    \
-       __VA_ARGS__;                        \
-       break;                              \
-     }                                     \
-   }
- 
- void rotary_embedding_cos_sin(
-     at::Tensor& cos,
-     at::Tensor& sin,
-     at::Tensor& query,
-     const std::optional<at::Tensor>& key,
-     int64_t head_size,
-     bool interleaved) {
-   TORCH_CHECK(
-       query.dim() == 2 || query.dim() == 3,
-       "query must be in shape [num_tokens, hidden_size] or [num_tokens, num_heads, head_size]");
-   if (key.has_value()) {
-     TORCH_CHECK(
-         key->dim() == 2 || key->dim() == 3,
-         "key must be in shape [num_tokens, hidden_size] or [num_tokens, num_kv_heads, head_size]");
-   }
- 
-   const int64_t num_tokens = query.size(0);
- 
-   TORCH_CHECK(cos.dim() == 2, "cos must be in shape [num_tokens, D_cos]");
-   TORCH_CHECK(sin.dim() == 2, "sin must be in shape [num_tokens, D_sin]");
-   TORCH_CHECK(cos.size(0) == num_tokens, "cos num_tokens mismatch with query");
-   TORCH_CHECK(sin.size(0) == num_tokens, "sin num_tokens mismatch with query");
-   TORCH_CHECK(cos.size(1) == sin.size(1), "cos and sin D_cos/D_sin mismatch");
- 
-   TORCH_CHECK(cos.scalar_type() == query.scalar_type(), "cos dtype mismatch with query");
-   TORCH_CHECK(sin.scalar_type() == query.scalar_type(), "sin dtype mismatch with query");
-   TORCH_CHECK(cos.is_cuda() && sin.is_cuda() && query.is_cuda(), "cos/sin/query must be CUDA tensors");
-   TORCH_CHECK(query.is_contiguous(), "query must be contiguous");
-   if (key.has_value()) {
-     TORCH_CHECK(key->is_cuda(), "key must be CUDA tensor if provided");
-     TORCH_CHECK(key->scalar_type() == query.scalar_type(), "key dtype mismatch with query");
-     TORCH_CHECK(key->is_contiguous(), "key must be contiguous");
-   }
- 
-   int query_hidden_size_calculated;
-   if (query.dim() == 2) {
-     query_hidden_size_calculated = (int)query.size(1);
-   } else {
-     query_hidden_size_calculated = (int)query.size(1) * (int)query.size(2);
-     TORCH_CHECK(query.size(2) == head_size, "query head_size mismatch in 3D tensor");
-   }
-   TORCH_CHECK(query_hidden_size_calculated % head_size == 0, "query_hidden_size not divisible by head_size");
-   int num_heads = query_hidden_size_calculated / (int)head_size;
- 
-   int key_hidden_size_calculated = 0;
-   int num_kv_heads = num_heads;
-   if (key.has_value()) {
-     TORCH_CHECK((int)key->size(0) == num_tokens, "key num_tokens mismatch");
-     if (key->dim() == 2) {
-       key_hidden_size_calculated = (int)key->size(1);
-     } else {
-       key_hidden_size_calculated = (int)key->size(1) * (int)key->size(2);
-       TORCH_CHECK((int)key->size(2) == head_size, "key head_size mismatch in 3D tensor");
-     }
-     TORCH_CHECK(key_hidden_size_calculated % head_size == 0, "key_hidden_size not divisible by head_size");
-     num_kv_heads = key_hidden_size_calculated / (int)head_size;
-   }
-   TORCH_CHECK(num_heads % num_kv_heads == 0, "num_heads must be divisible by num_kv_heads");
- 
-   // NeoX interleaved: if cos/sin are full-dim (head_size), downsample once.
-   if (interleaved && cos.size(1) == head_size) {
-     TORCH_CHECK(head_size % 2 == 0, "interleaved layout requires even head_size");
-     const int64_t half = head_size / 2;
-     std::vector<int64_t> new_shape = {cos.size(0), half, 2};
-     cos = cos.view(new_shape).select(2, 0).contiguous();
-     sin = sin.view(new_shape).select(2, 1).contiguous();
-   }
- 
-   const int rot_dim_from_cache = (int)cos.size(1);
-   const int64_t query_token_stride = query_hidden_size_calculated;
-   const int64_t key_token_stride = key.has_value() ? key_hidden_size_calculated : 0;
- 
-   int64_t head_stride_query;
-   if (query.dim() == 3 && query.size(1) == num_heads && query.size(2) == head_size) {
-     head_stride_query = query.stride(1);
-   } else {
-     head_stride_query = head_size;
-   }
- 
-   int64_t head_stride_key = head_size;
-   if (key.has_value()) {
-     if (key->dim() == 3 && key->size(1) == num_kv_heads && key->size(2) == head_size) {
-       head_stride_key = key->stride(1);
-     } else {
-       head_stride_key = head_size;
-     }
-   }
- 
-   const int embed_dim_for_rotation = interleaved ? rot_dim_from_cache : (rot_dim_from_cache / 2);
-   TORCH_CHECK(embed_dim_for_rotation > 0, "embed_dim_for_rotation must be > 0");
- 
-   const int max_pairs_to_rotate_per_token =
-       std::max(num_heads * embed_dim_for_rotation, num_kv_heads * embed_dim_for_rotation);
- 
-   const at::cuda::OptionalCUDAGuard device_guard(device_of(query));
-   const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
- 
-   AT_DISPATCH_FLOATING_TYPES_AND2(
-       at::ScalarType::Half, at::ScalarType::BFloat16, query.scalar_type(), "rotary_embedding_cos_sin", [&] {
-         using torch_scalar_t = scalar_t;
-         using cuda_scalar_t = typename std::conditional<
-             std::is_same<torch_scalar_t, at::Half>::value,
-             nv_half,
-             typename std::conditional<std::is_same<torch_scalar_t, at::BFloat16>::value, nv_bfloat16, torch_scalar_t>::
-                 type>::type;
- 
-         // Constants for vectorization
-         constexpr int kVecBytes = 16;
-         constexpr int kElePerVec = kVecBytes / sizeof(torch_scalar_t);
- 
-         // Determine how many pairs one thread handles in one vector step
-         // Interleaved: 1 vector load (16B) contains 'kElePerVec' elements -> 'kElePerVec / 2' pairs.
-         // Non-interleaved: 2 vector loads (32B) contains 'kElePerVec' pairs (X vec + Y vec).
-         const int pairs_per_step = interleaved ? (kElePerVec / 2) : kElePerVec;
- 
-         // Check if we can guarantee vectorization for ALL tokens
-         // We need Base Pointers, Strides, and Dimension to be aligned.
-         bool can_vectorize_all = true;
- 
-         // 1. Check Dimensions
-         if (embed_dim_for_rotation % pairs_per_step != 0) can_vectorize_all = false;
- 
-         // 2. Check Base Pointers
-         if (reinterpret_cast<uintptr_t>(query.data_ptr()) % kVecBytes != 0) can_vectorize_all = false;
-         if (reinterpret_cast<uintptr_t>(cos.data_ptr()) % kVecBytes != 0) can_vectorize_all = false;
-         if (reinterpret_cast<uintptr_t>(sin.data_ptr()) % kVecBytes != 0) can_vectorize_all = false;
-         if (key.has_value()) {
-           if (reinterpret_cast<uintptr_t>(key->data_ptr()) % kVecBytes != 0) can_vectorize_all = false;
-         }
- 
-         // 3. Check Strides
-         // We need the stride between tokens to be a multiple of vector size
-         // to ensure that if token 0 is aligned, token 1 is also aligned.
-         if (query_token_stride * sizeof(torch_scalar_t) % kVecBytes != 0) can_vectorize_all = false;
-         if (head_stride_query * sizeof(torch_scalar_t) % kVecBytes != 0) can_vectorize_all = false;
- 
-         if (key.has_value()) {
-           if (key_token_stride * sizeof(torch_scalar_t) % kVecBytes != 0) can_vectorize_all = false;
-           if (head_stride_key * sizeof(torch_scalar_t) % kVecBytes != 0) can_vectorize_all = false;
-         }
- 
-         // Determine launch configuration
-         // If we can vectorize all, each thread handles 'pairs_per_step' pairs.
-         // Otherwise, fallback to conservative estimate (1 pair per thread) to ensure enough blocks.
-         const int launch_pairs_per_thread = can_vectorize_all ? pairs_per_step : 1;
- 
-         const int total_threads_needed =
-             (max_pairs_to_rotate_per_token + launch_pairs_per_thread - 1) / launch_pairs_per_thread;
- 
-         // Case 1: 2D Grid (Split one token across multiple blocks)
-         // Keep block size moderate (128-256) aligned with head size
-         const int threads_per_block_2d = std::min<int>(256, std::max(128, embed_dim_for_rotation));
-         const int blocks_per_token_2d = (total_threads_needed + threads_per_block_2d - 1) / threads_per_block_2d;
- 
-         // Decide grid strategy
-         const bool use_grid_2d = (num_tokens <= 4) && (blocks_per_token_2d > 1);
- 
-         // Case 2: 1D Grid (One block per token)
-         // Maximize threads per block to cover all heads in one block if possible
-         // Cap at 512 threads to balance occupancy and register usage
-         const int threads_per_block_1d = std::min<int>(512, std::max(128, (total_threads_needed + 31) / 32 * 32));
- 
-         // Final launch config
-         const int threads_per_block = use_grid_2d ? threads_per_block_2d : threads_per_block_1d;
-         const int blocks_per_token = use_grid_2d ? blocks_per_token_2d : 1;
- 
-         dim3 block(threads_per_block);
-         dim3 grid_2d((int)num_tokens, std::max(1, blocks_per_token));
-         dim3 grid_1d((int)num_tokens);
- 
-         // No shared memory needed
-         size_t smem_size = 0;
- 
-         DISPATCH_ROT_DIM(
-             embed_dim_for_rotation,
-             auto launch_kernel =
-                 [&](bool vectorized) {
-                   if (interleaved) {
-                     if (use_grid_2d) {
-                       if (vectorized) {
-                         rotary_embedding_kernel_2d<cuda_scalar_t, true, true, ROT_EMBED_DIM>
-                             <<<grid_2d, block, smem_size, stream>>>(
-                                 reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
-                                 reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
-                                 reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
-                                 key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>())
-                                                 : nullptr,
-                                 rot_dim_from_cache,
-                                 embed_dim_for_rotation,
-                                 query_token_stride,
-                                 key_token_stride,
-                                 head_stride_query,
-                                 head_stride_key,
-                                 num_heads,
-                                 num_kv_heads,
-                                 (int)head_size,
-                                 blocks_per_token);
-                       } else {
-                         rotary_embedding_kernel_2d<cuda_scalar_t, true, false, ROT_EMBED_DIM>
-                             <<<grid_2d, block, smem_size, stream>>>(
-                                 reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
-                                 reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
-                                 reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
-                                 key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>())
-                                                 : nullptr,
-                                 rot_dim_from_cache,
-                                 embed_dim_for_rotation,
-                                 query_token_stride,
-                                 key_token_stride,
-                                 head_stride_query,
-                                 head_stride_key,
-                                 num_heads,
-                                 num_kv_heads,
-                                 (int)head_size,
-                                 blocks_per_token);
-                       }
-                     } else {
-                       if (vectorized) {
-                         rotary_embedding_kernel_1d<cuda_scalar_t, true, true, ROT_EMBED_DIM>
-                             <<<grid_1d, block, smem_size, stream>>>(
-                                 reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
-                                 reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
-                                 reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
-                                 key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>())
-                                                 : nullptr,
-                                 rot_dim_from_cache,
-                                 embed_dim_for_rotation,
-                                 query_token_stride,
-                                 key_token_stride,
-                                 head_stride_query,
-                                 head_stride_key,
-                                 num_heads,
-                                 num_kv_heads,
-                                 (int)head_size);
-                       } else {
-                         rotary_embedding_kernel_1d<cuda_scalar_t, true, false, ROT_EMBED_DIM>
-                             <<<grid_1d, block, smem_size, stream>>>(
-                                 reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
-                                 reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
-                                 reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
-                                 key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>())
-                                                 : nullptr,
-                                 rot_dim_from_cache,
-                                 embed_dim_for_rotation,
-                                 query_token_stride,
-                                 key_token_stride,
-                                 head_stride_query,
-                                 head_stride_key,
-                                 num_heads,
-                                 num_kv_heads,
-                                 (int)head_size);
-                       }
-                     }
-                   } else {  // non-interleaved
-                     if (use_grid_2d) {
-                       if (vectorized) {
-                         rotary_embedding_kernel_2d<cuda_scalar_t, false, true, ROT_EMBED_DIM>
-                             <<<grid_2d, block, smem_size, stream>>>(
-                                 reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
-                                 reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
-                                 reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
-                                 key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>())
-                                                 : nullptr,
-                                 rot_dim_from_cache,
-                                 embed_dim_for_rotation,
-                                 query_token_stride,
-                                 key_token_stride,
-                                 head_stride_query,
-                                 head_stride_key,
-                                 num_heads,
-                                 num_kv_heads,
-                                 (int)head_size,
-                                 blocks_per_token);
-                       } else {
-                         rotary_embedding_kernel_2d<cuda_scalar_t, false, false, ROT_EMBED_DIM>
-                             <<<grid_2d, block, smem_size, stream>>>(
-                                 reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
-                                 reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
-                                 reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
-                                 key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>())
-                                                 : nullptr,
-                                 rot_dim_from_cache,
-                                 embed_dim_for_rotation,
-                                 query_token_stride,
-                                 key_token_stride,
-                                 head_stride_query,
-                                 head_stride_key,
-                                 num_heads,
-                                 num_kv_heads,
-                                 (int)head_size,
-                                 blocks_per_token);
-                       }
-                     } else {
-                       if (vectorized) {
-                         rotary_embedding_kernel_1d<cuda_scalar_t, false, true, ROT_EMBED_DIM>
-                             <<<grid_1d, block, smem_size, stream>>>(
-                                 reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
-                                 reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
-                                 reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
-                                 key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>())
-                                                 : nullptr,
-                                 rot_dim_from_cache,
-                                 embed_dim_for_rotation,
-                                 query_token_stride,
-                                 key_token_stride,
-                                 head_stride_query,
-                                 head_stride_key,
-                                 num_heads,
-                                 num_kv_heads,
-                                 (int)head_size);
-                       } else {
-                         rotary_embedding_kernel_1d<cuda_scalar_t, false, false, ROT_EMBED_DIM>
-                             <<<grid_1d, block, smem_size, stream>>>(
-                                 reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
-                                 reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
-                                 reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
-                                 key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>())
-                                                 : nullptr,
-                                 rot_dim_from_cache,
-                                 embed_dim_for_rotation,
-                                 query_token_stride,
-                                 key_token_stride,
-                                 head_stride_query,
-                                 head_stride_key,
-                                 num_heads,
-                                 num_kv_heads,
-                                 (int)head_size);
-                       }
-                     }
-                   }
-                 };
- 
-             // Close the log file after all launches
-             if (can_vectorize_all) { launch_kernel(true); } else { launch_kernel(false); });
-       });
- 
-   C10_CUDA_KERNEL_LAUNCH_CHECK();
- }
\ No newline at end of file
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAException.h>
+#include <c10/cuda/CUDAGuard.h>
+#include <c10/cuda/CUDAStream.h>
+#include <torch/all.h>
+
+#include <cmath>
+#include <vector>
+
+#include "utils.h"
+
+// Packed vector payload per step (16B compute tile)
+template <typename scalar_t, bool interleaved>
+struct RotaryVecData;
+
+template <typename scalar_t>
+struct RotaryVecData<scalar_t, true> {
+  float4 vec;
+  float2 cos_vec;
+  float2 sin_vec;
+};
+
+template <typename scalar_t>
+struct RotaryVecData<scalar_t, false> {
+  float4 vec_x;
+  float4 vec_y;
+  float4 cos_x;
+  float4 sin_x;
+  float4 cos_y;
+  float4 sin_y;
+};
+
+// Vector load (q/k aligned or scalar-gather) + always vector-load cos/sin.
+// rot_offset is "pair index":
+//   - interleaved: pair i maps to q[2*i], q[2*i+1]
+//   - non-interleaved: pair i maps to q[i] and q[i+embed_dim]
+// embed_dim is "half dimension" (d).
+template <typename scalar_t, bool interleaved, bool aligned_qk>
+__device__ __forceinline__ RotaryVecData<scalar_t, interleaved> load_rotary_vec(
+    const scalar_t* __restrict__ arr,
+    const scalar_t* __restrict__ cos_ptr,
+    const scalar_t* __restrict__ sin_ptr,
+    int rot_offset,
+    int embed_dim) {
+  RotaryVecData<scalar_t, interleaved> data;
+
+  constexpr int kVecBytes = 16;
+  constexpr int kElePerVec = kVecBytes / sizeof(scalar_t);
+
+  if constexpr (interleaved) {
+    const int base = rot_offset * 2;
+
+    if constexpr (aligned_qk) {
+      data.vec = *reinterpret_cast<const float4*>(arr + base);
+    } else {
+      union VecU {
+        float4 v;
+        scalar_t e[kElePerVec];
+      } tmp;
+#pragma unroll
+      for (int i = 0; i < kElePerVec; ++i)
+        tmp.e[i] = arr[base + i];
+      data.vec = tmp.v;
+    }
+
+    data.cos_vec = *reinterpret_cast<const float2*>(cos_ptr + rot_offset);
+    data.sin_vec = *reinterpret_cast<const float2*>(sin_ptr + rot_offset);
+
+  } else {
+    const int bx = rot_offset;
+    const int by = rot_offset + embed_dim;
+
+    if constexpr (aligned_qk) {
+      data.vec_x = *reinterpret_cast<const float4*>(arr + bx);
+      data.vec_y = *reinterpret_cast<const float4*>(arr + by);
+    } else {
+      union VecU {
+        float4 v;
+        scalar_t e[kElePerVec];
+      } tx, ty;
+#pragma unroll
+      for (int i = 0; i < kElePerVec; ++i) {
+        tx.e[i] = arr[bx + i];
+        ty.e[i] = arr[by + i];
+      }
+      data.vec_x = tx.v;
+      data.vec_y = ty.v;
+    }
+
+    data.cos_x = *reinterpret_cast<const float4*>(cos_ptr + bx);
+    data.sin_x = *reinterpret_cast<const float4*>(sin_ptr + bx);
+    data.cos_y = *reinterpret_cast<const float4*>(cos_ptr + by);
+    data.sin_y = *reinterpret_cast<const float4*>(sin_ptr + by);
+  }
+
+  return data;
+}
+
+// Compute RoPE on the loaded vector tile, then store back to q/k.
+template <typename scalar_t, bool interleaved, bool aligned_qk>
+__device__ __forceinline__ void compute_store_rotary_vec(
+    scalar_t* __restrict__ arr, 
+    const RotaryVecData<scalar_t, interleaved>& data, 
+    int rot_offset, 
+    int embed_dim) {
+  constexpr int kVecBytes = 16;
+  constexpr int kElePerVec = kVecBytes / sizeof(scalar_t);
+
+  if constexpr (interleaved) {
+    union VecU {
+      float4 v;
+      scalar_t e[kElePerVec];
+    } v;
+    v.v = data.vec;
+
+    union CSU {
+      float2 v;
+      scalar_t e[kElePerVec / 2];
+    } c, s;
+    c.v = data.cos_vec;
+    s.v = data.sin_vec;
+
+#pragma unroll
+    for (int i = 0; i < kElePerVec; i += 2) {
+      int idx = i / 2;
+      float cos_val = static_cast<float>(c.e[idx]);
+      float sin_val = static_cast<float>(s.e[idx]);
+      float x = static_cast<float>(v.e[i]);
+      float y = static_cast<float>(v.e[i + 1]);
+      v.e[i] = static_cast<scalar_t>(x * cos_val - y * sin_val);
+      v.e[i + 1] = static_cast<scalar_t>(y * cos_val + x * sin_val);
+    }
+
+    const int base = rot_offset * 2;
+    if constexpr (aligned_qk) {
+      *reinterpret_cast<float4*>(arr + base) = v.v;
+    } else {
+#pragma unroll
+      for (int i = 0; i < kElePerVec; ++i)
+        arr[base + i] = v.e[i];
+    }
+
+  } else {
+    union VecU {
+      float4 v;
+      scalar_t e[kElePerVec];
+    } vx, vy;
+    vx.v = data.vec_x;
+    vy.v = data.vec_y;
+
+    union CSU {
+      float4 v;
+      scalar_t e[kElePerVec];
+    } cx, sx, cy, sy;
+    cx.v = data.cos_x;
+    sx.v = data.sin_x;
+    cy.v = data.cos_y;
+    sy.v = data.sin_y;
+
+#pragma unroll
+    for (int i = 0; i < kElePerVec; ++i) {
+      float cos_x = static_cast<float>(cx.e[i]);
+      float sin_x = static_cast<float>(sx.e[i]);
+      float cos_y = static_cast<float>(cy.e[i]);
+      float sin_y = static_cast<float>(sy.e[i]);
+
+      float x = static_cast<float>(vx.e[i]);
+      float y = static_cast<float>(vy.e[i]);
+
+      vx.e[i] = static_cast<scalar_t>(x * cos_x - y * sin_x);
+      vy.e[i] = static_cast<scalar_t>(y * cos_y + x * sin_y);
+    }
+
+    const int bx = rot_offset;
+    const int by = rot_offset + embed_dim;
+
+    if constexpr (aligned_qk) {
+      *reinterpret_cast<float4*>(arr + bx) = vx.v;
+      *reinterpret_cast<float4*>(arr + by) = vy.v;
+    } else {
+#pragma unroll
+      for (int i = 0; i < kElePerVec; ++i) {
+        arr[bx + i] = vx.e[i];
+        arr[by + i] = vy.e[i];
+      }
+    }
+  }
+}
+
+// Scalar fallback: one (x,y) pair per thread-iteration.
+template <typename scalar_t, bool interleaved>
+inline __device__ void apply_token_rotary_embedding(
+    scalar_t* __restrict__ arr,            // [head_size]
+    const scalar_t* __restrict__ cos_ptr,  // [rot_dim]
+    const scalar_t* __restrict__ sin_ptr,  // [rot_dim]
+    int rot_offset,
+    int embed_dim) {  // for non-interleaved: half dim
+  if constexpr (interleaved) {  // NeoX-style: interleaved layout [x0, y0, x1, y1, ...].
+    const int x_index = 2 * rot_offset;
+    const int y_index = x_index + 1;
+
+    // Directly access Global Memory
+    const float cos_val = static_cast<float>(cos_ptr[rot_offset]);
+    const float sin_val = static_cast<float>(sin_ptr[rot_offset]);
+
+    const float x = static_cast<float>(arr[x_index]);
+    const float y = static_cast<float>(arr[y_index]);
+    arr[x_index] = static_cast<scalar_t>(x * cos_val - y * sin_val);
+    arr[y_index] = static_cast<scalar_t>(y * cos_val + x * sin_val);
+  } else { // GPT-J / LLaMA style: layout [x0, x1, ..., y0, y1, ...]
+    const int x_index = rot_offset;
+    const int y_index = rot_offset + embed_dim;
+
+    // Directly access Global Memory
+    const float cos_val_x = static_cast<float>(cos_ptr[rot_offset]);
+    const float sin_val_x = static_cast<float>(sin_ptr[rot_offset]);
+    const float cos_val_y = static_cast<float>(cos_ptr[rot_offset + embed_dim]);
+    const float sin_val_y = static_cast<float>(sin_ptr[rot_offset + embed_dim]);
+
+    const float x = static_cast<float>(arr[x_index]);
+    const float y = static_cast<float>(arr[y_index]);
+    arr[x_index] = static_cast<scalar_t>(x * cos_val_x - y * sin_val_x);
+    arr[y_index] = static_cast<scalar_t>(y * cos_val_y + x * sin_val_y);
+  }
+}
+
+template <typename scalar_t, bool interleaved>
+inline __device__ void apply_token_rotary_embedding_vec(
+    scalar_t* __restrict__ arr,            // [head_size]
+    const scalar_t* __restrict__ cos_ptr,  // [rot_dim]
+    const scalar_t* __restrict__ sin_ptr,  // [rot_dim]
+    int rot_offset,
+    int embed_dim) {
+  using vec_t = float4;
+  constexpr int kVecBytes = sizeof(vec_t);
+  constexpr int kScalarBytes = sizeof(scalar_t);
+  constexpr int kElePerVec = kVecBytes / kScalarBytes;
+
+  // Union for type punning to avoid strict aliasing issues with reinterpret_cast
+  union VecUnion {
+    vec_t vec;
+    scalar_t elems[kElePerVec];
+  };
+
+  if constexpr (interleaved) {
+    VecUnion data;
+    data.vec = *reinterpret_cast<const vec_t*>(arr + rot_offset * 2);
+
+    // Vectorized load for cos/sin: always 8 bytes (half of vector size, float2)
+    float2 cos_vec = *reinterpret_cast<const float2*>(cos_ptr + rot_offset);
+    float2 sin_vec = *reinterpret_cast<const float2*>(sin_ptr + rot_offset);
+
+    union CosSinVec {
+      float2 vec;
+      scalar_t elems[kElePerVec / 2];
+    };
+    CosSinVec cos_u, sin_u;
+    cos_u.vec = cos_vec;
+    sin_u.vec = sin_vec;
+
+#pragma unroll
+    for (int i = 0; i < kElePerVec; i += 2) {
+      int idx = i / 2;
+      float cos_val = static_cast<float>(cos_u.elems[idx]);
+      float sin_val = static_cast<float>(sin_u.elems[idx]);
+
+      float x = static_cast<float>(data.elems[i]);
+      float y = static_cast<float>(data.elems[i + 1]);
+
+      data.elems[i] = static_cast<scalar_t>(x * cos_val - y * sin_val);
+      data.elems[i + 1] = static_cast<scalar_t>(y * cos_val + x * sin_val);
+    }
+
+    *reinterpret_cast<vec_t*>(arr + rot_offset * 2) = data.vec;
+
+  } else { // Non-interleaved
+    VecUnion data_x, data_y;
+    data_x.vec = *reinterpret_cast<const vec_t*>(arr + rot_offset);
+    data_y.vec = *reinterpret_cast<const vec_t*>(arr + rot_offset + embed_dim);
+
+    // Vectorized load for cos/sin: always 16 bytes (full vector size, float4)
+    float4 cos_vec_x = *reinterpret_cast<const float4*>(cos_ptr + rot_offset);
+    float4 sin_vec_x = *reinterpret_cast<const float4*>(sin_ptr + rot_offset);
+    float4 cos_vec_y = *reinterpret_cast<const float4*>(cos_ptr + rot_offset + embed_dim);
+    float4 sin_vec_y = *reinterpret_cast<const float4*>(sin_ptr + rot_offset + embed_dim);
+
+    union CosSinVec {
+      float4 vec;
+      scalar_t elems[kElePerVec];
+    };
+    CosSinVec cos_u_x, sin_u_x, cos_u_y, sin_u_y;
+    cos_u_x.vec = cos_vec_x;
+    sin_u_x.vec = sin_vec_x;
+    cos_u_y.vec = cos_vec_y;
+    sin_u_y.vec = sin_vec_y;
+
+#pragma unroll
+    for (int i = 0; i < kElePerVec; ++i) {
+      float cos_val_x = static_cast<float>(cos_u_x.elems[i]);
+      float sin_val_x = static_cast<float>(sin_u_x.elems[i]);
+      float cos_val_y = static_cast<float>(cos_u_y.elems[i]);
+      float sin_val_y = static_cast<float>(sin_u_y.elems[i]);
+
+      float x = static_cast<float>(data_x.elems[i]);
+      float y = static_cast<float>(data_y.elems[i]);
+
+      data_x.elems[i] = static_cast<scalar_t>(x * cos_val_x - y * sin_val_x);
+      data_y.elems[i] = static_cast<scalar_t>(y * cos_val_y + x * sin_val_y);
+    }
+
+    *reinterpret_cast<vec_t*>(arr + rot_offset) = data_x.vec;
+    *reinterpret_cast<vec_t*>(arr + rot_offset + embed_dim) = data_y.vec;
+  }
+}
+
+template <>
+inline __device__ void apply_token_rotary_embedding<float, true>(
+    float* __restrict__ arr,
+    const float* __restrict__ cos_ptr,
+    const float* __restrict__ sin_ptr,
+    int rot_offset,
+    int /*embed_dim*/) {
+  float2 xy = *reinterpret_cast<const float2*>(arr + rot_offset * 2);
+
+  const float cos_val = static_cast<float>(cos_ptr[rot_offset]);
+  const float sin_val = static_cast<float>(sin_ptr[rot_offset]);
+
+  float2 out;
+  out.x = xy.x * cos_val - xy.y * sin_val;
+  out.y = xy.y * cos_val + xy.x * sin_val;
+
+  *reinterpret_cast<float2*>(arr + rot_offset * 2) = out;
+}
+
+// 2D grid kernel: blockIdx.x = token, blockIdx.y = tile within token
+// Each thread processes "pairs_per_step" pairs per iteration when vectorized.
+template <typename scalar_t, bool interleaved, bool vectorized, bool aligned_qk, int ROT_EMBED_DIM>
+__global__ void rotary_embedding_kernel_2d(
+    const scalar_t* __restrict__ cos_data,  // [num_tokens, rot_dim_arg]
+    const scalar_t* __restrict__ sin_data,  // [num_tokens, rot_dim_arg]
+    scalar_t* __restrict__ query_total,
+    scalar_t* __restrict__ key_total,
+    const int rot_dim_arg,
+    const int embed_dim_for_rotation_arg,
+    const int64_t query_token_stride,
+    const int64_t key_token_stride,
+    const int64_t head_stride_query,
+    const int64_t head_stride_key,
+    const int num_heads,
+    const int num_kv_heads,
+    const int head_size_arg,
+    const int blocks_per_token) {
+  const int token_idx = blockIdx.x;
+  if (token_idx >= gridDim.x) {
+    return;
+  }
+
+  // Pointers to Global Memory for the current token
+  const scalar_t* current_token_cos_ptr = cos_data + token_idx * rot_dim_arg;
+  const scalar_t* current_token_sin_ptr = sin_data + token_idx * rot_dim_arg;
+
+  constexpr int kVecBytes = 16;
+  const int scalar_size = sizeof(scalar_t);
+
+  scalar_t* query_for_token = query_total + token_idx * (int)query_token_stride;
+  scalar_t* key_for_token = (key_total != nullptr) ? (key_total + token_idx * (int)key_token_stride) : nullptr;
+
+  const int local_block_idx = blockIdx.y;
+
+  // Otherwise fallback to runtime argument
+  const int embed_dim_for_rotation = (ROT_EMBED_DIM > 0) ? ROT_EMBED_DIM : embed_dim_for_rotation_arg;
+
+  if constexpr (vectorized) {
+    using vec_t = float4;
+    constexpr int kElePerVec = kVecBytes / sizeof(scalar_t);
+    constexpr int pairs_per_step = interleaved ? (kElePerVec / 2) : kElePerVec;
+
+    const int pair_stride = blockDim.x * blocks_per_token * pairs_per_step;
+    int i = (local_block_idx * blockDim.x + threadIdx.x) * pairs_per_step;
+    const int nq_pairs = num_heads * embed_dim_for_rotation;
+
+    // PROLOGUE
+    RotaryVecData<scalar_t, interleaved> curr_data;
+    if (i < nq_pairs) {
+      int head_idx = i / embed_dim_for_rotation;
+      int rot_offset = i % embed_dim_for_rotation;
+      scalar_t* ptr = query_for_token + head_idx * (int)head_stride_query;
+      curr_data = load_rotary_vec<scalar_t, interleaved, aligned_qk>(
+          ptr, current_token_cos_ptr, current_token_sin_ptr, rot_offset, embed_dim_for_rotation);
+    }
+
+    int next_i = i + pair_stride;
+    // MAIN LOOP
+    for (; i < nq_pairs; i += pair_stride, next_i += pair_stride) {
+      // PREFETCH NEXT
+      RotaryVecData<scalar_t, interleaved> next_data;
+      bool active_next = (next_i < nq_pairs);
+
+      if (active_next) {
+        int head_idx_next = next_i / embed_dim_for_rotation;
+        int rot_offset_next = next_i % embed_dim_for_rotation;
+        scalar_t* ptr_next = query_for_token + head_idx_next * (int)head_stride_query;
+        next_data = load_rotary_vec<scalar_t, interleaved, aligned_qk>(
+            ptr_next, current_token_cos_ptr, current_token_sin_ptr, rot_offset_next, embed_dim_for_rotation);
+      }
+      // COMPUTE CURRENT
+      int head_idx = i / embed_dim_for_rotation;
+      int rot_offset = i % embed_dim_for_rotation;
+      scalar_t* ptr = query_for_token + head_idx * (int)head_stride_query;
+      compute_store_rotary_vec<scalar_t, interleaved, aligned_qk>(ptr, curr_data, rot_offset, embed_dim_for_rotation);
+      // SHIFT
+      if (active_next) curr_data = next_data;
+    }
+
+    if (key_for_token != nullptr) {
+      const int nk_pairs = num_kv_heads * embed_dim_for_rotation;
+      int k_i = (local_block_idx * blockDim.x + threadIdx.x) * pairs_per_step;
+
+      // PROLOGUE
+      RotaryVecData<scalar_t, interleaved> curr_data_k;
+      if (k_i < nk_pairs) {
+        int head_idx = k_i / embed_dim_for_rotation;
+        int rot_offset = k_i % embed_dim_for_rotation;
+        scalar_t* ptr = key_for_token + head_idx * (int)head_stride_key;
+        curr_data_k = load_rotary_vec<scalar_t, interleaved, aligned_qk>(
+            ptr, current_token_cos_ptr, current_token_sin_ptr, rot_offset, embed_dim_for_rotation);
+      }
+
+      int next_k_i = k_i + pair_stride;
+
+      // MAIN LOOP
+      for (; k_i < nk_pairs; k_i += pair_stride, next_k_i += pair_stride) {
+        // PREFETCH NEXT
+        RotaryVecData<scalar_t, interleaved> next_data_k;
+        bool active_next = (next_k_i < nk_pairs);
+        if (active_next) {
+          int head_idx_next = next_k_i / embed_dim_for_rotation;
+          int rot_offset_next = next_k_i % embed_dim_for_rotation;
+          scalar_t* ptr_next = key_for_token + head_idx_next * (int)head_stride_key;
+          next_data_k = load_rotary_vec<scalar_t, interleaved, aligned_qk>(
+              ptr_next, current_token_cos_ptr, current_token_sin_ptr, rot_offset_next, embed_dim_for_rotation);
+        }
+
+        // COMPUTE CURRENT
+        int head_idx = k_i / embed_dim_for_rotation;
+        int rot_offset = k_i % embed_dim_for_rotation;
+        scalar_t* ptr = key_for_token + head_idx * (int)head_stride_key;
+        compute_store_rotary_vec<scalar_t, interleaved, aligned_qk>(ptr, curr_data_k, rot_offset, embed_dim_for_rotation);
+
+        // SHIFT
+        curr_data_k = next_data_k;
+      }
+    }
+  } else {
+    // Fallback to scalar implementation
+    const int pair_stride = blockDim.x * blocks_per_token;
+    const int thread_pair_offset = local_block_idx * blockDim.x + threadIdx.x;
+
+    const int nq_pairs = num_heads * embed_dim_for_rotation;
+    for (int i = thread_pair_offset; i < nq_pairs; i += pair_stride) {
+      const int head_idx = i / embed_dim_for_rotation;
+      const int rot_offset = i % embed_dim_for_rotation;
+
+      scalar_t* query_for_token_head = query_for_token + head_idx * (int)head_stride_query;
+      apply_token_rotary_embedding<scalar_t, interleaved>(
+          query_for_token_head, current_token_cos_ptr, current_token_sin_ptr, rot_offset, embed_dim_for_rotation);
+    }
+
+    if (key_for_token != nullptr) {
+      const int nk_pairs = num_kv_heads * embed_dim_for_rotation;
+      for (int i = thread_pair_offset; i < nk_pairs; i += pair_stride) {
+        const int head_idx = i / embed_dim_for_rotation;
+        const int rot_offset = i % embed_dim_for_rotation;
+
+        scalar_t* key_for_token_head = key_for_token + head_idx * (int)head_stride_key;
+        apply_token_rotary_embedding<scalar_t, interleaved>(
+            key_for_token_head, current_token_cos_ptr, current_token_sin_ptr, rot_offset, embed_dim_for_rotation);
+      }
+    }
+  }
+}
+
+// 1D grid kernel: one block per token (blockIdx.x = token)
+template <typename scalar_t, bool interleaved, bool vectorized, bool aligned_qk, int ROT_EMBED_DIM>
+__launch_bounds__(512) __global__ void rotary_embedding_kernel_1d(
+    const scalar_t* __restrict__ cos_data,  // [num_tokens, rot_dim_arg]
+    const scalar_t* __restrict__ sin_data,  // [num_tokens, rot_dim_arg]
+    scalar_t* __restrict__ query_total,
+    scalar_t* __restrict__ key_total,
+    const int rot_dim_arg,
+    const int embed_dim_for_rotation_arg,
+    const int64_t query_token_stride,
+    const int64_t key_token_stride,
+    const int64_t head_stride_query,
+    const int64_t head_stride_key,
+    const int num_heads,
+    const int num_kv_heads,
+    const int head_size_arg) {
+  const int token_idx = blockIdx.x;
+  if (token_idx >= gridDim.x) return;
+
+  const scalar_t* current_token_cos_ptr = cos_data + token_idx * rot_dim_arg;
+  const scalar_t* current_token_sin_ptr = sin_data + token_idx * rot_dim_arg;
+
+  constexpr int kVecBytes = 16;
+  const int scalar_size = sizeof(scalar_t);
+
+  scalar_t* query_for_token = query_total + token_idx * (int)query_token_stride;
+  scalar_t* key_for_token = (key_total != nullptr) ? (key_total + token_idx * (int)key_token_stride) : nullptr;
+
+  // Otherwise fallback to runtime argument
+  const int embed_dim_for_rotation = (ROT_EMBED_DIM > 0) ? ROT_EMBED_DIM : embed_dim_for_rotation_arg;
+
+  if constexpr (vectorized) {
+    constexpr int kVecBytes = 16;
+    constexpr int kElePerVec = kVecBytes / sizeof(scalar_t);
+    // Interleaved: 1 vec = kElePerVec elements = kElePerVec/2 pairs
+    // Non-Interleaved: 1 vec (X) + 1 vec (Y) = kElePerVec elements = kElePerVec pairs
+    constexpr int pairs_per_step = interleaved ? (kElePerVec / 2) : kElePerVec;
+    const int nq_pairs = num_heads * embed_dim_for_rotation;
+
+    // Original Stride
+    const int stride = blockDim.x * pairs_per_step;
+
+    int i = threadIdx.x * pairs_per_step;
+
+    // PROLOGUE
+    RotaryVecData<scalar_t, interleaved> curr_data;
+
+    if (i < nq_pairs) {
+      int head_idx = i / embed_dim_for_rotation;
+      int rot_offset = i % embed_dim_for_rotation;
+      scalar_t* ptr = query_for_token + head_idx * (int)head_stride_query;
+      curr_data = load_rotary_vec<scalar_t, interleaved, aligned_qk>(
+          ptr, current_token_cos_ptr, current_token_sin_ptr, rot_offset, embed_dim_for_rotation);
+    }
+
+    int next_i = i + stride;
+
+    // MAIN LOOP
+    for (; i < nq_pairs; i += stride, next_i += stride) {
+      // PREFETCH NEXT
+      RotaryVecData<scalar_t, interleaved> next_data;
+
+      if (next_i < nq_pairs) {
+        int head_idx = next_i / embed_dim_for_rotation;
+        int rot_offset = next_i % embed_dim_for_rotation;
+        scalar_t* ptr = query_for_token + head_idx * (int)head_stride_query;
+        next_data = load_rotary_vec<scalar_t, interleaved, aligned_qk>(
+            ptr, current_token_cos_ptr, current_token_sin_ptr, rot_offset, embed_dim_for_rotation);
+      }
+
+      // COMPUTE CURRENT
+      if (i < nq_pairs) {
+        int head_idx = i / embed_dim_for_rotation;
+        int rot_offset = i % embed_dim_for_rotation;
+        scalar_t* ptr = query_for_token + head_idx * (int)head_stride_query;
+        compute_store_rotary_vec<scalar_t, interleaved, aligned_qk>(ptr, curr_data, rot_offset, embed_dim_for_rotation);
+      }
+
+      // SHIFT
+      curr_data = next_data;
+    }
+
+    if (key_for_token != nullptr) {
+      const int nk_pairs = num_kv_heads * embed_dim_for_rotation;
+      int k_i = threadIdx.x * pairs_per_step;
+
+      // PROLOGUE
+      RotaryVecData<scalar_t, interleaved> curr_data_k;
+
+      if (k_i < nk_pairs) {
+        int head_idx = k_i / embed_dim_for_rotation;
+        int rot_offset = k_i % embed_dim_for_rotation;
+        scalar_t* ptr = key_for_token + head_idx * (int)head_stride_key;
+        curr_data_k = load_rotary_vec<scalar_t, interleaved, aligned_qk>(
+            ptr, current_token_cos_ptr, current_token_sin_ptr, rot_offset, embed_dim_for_rotation);
+      }
+
+      int next_k_i = k_i + stride;
+
+      // MAIN LOOP
+      for (; k_i < nk_pairs; k_i += stride, next_k_i += stride) {
+        // PREFETCH NEXT
+        RotaryVecData<scalar_t, interleaved> next_data_k;
+
+        if (next_k_i < nk_pairs) {
+          int head_idx = next_k_i / embed_dim_for_rotation;
+          int rot_offset = next_k_i % embed_dim_for_rotation;
+          scalar_t* ptr = key_for_token + head_idx * (int)head_stride_key;
+          next_data_k = load_rotary_vec<scalar_t, interleaved, aligned_qk>(
+              ptr, current_token_cos_ptr, current_token_sin_ptr, rot_offset, embed_dim_for_rotation);
+        }
+
+        // COMPUTE CURRENT
+        if (k_i < nk_pairs) {
+          int head_idx = k_i / embed_dim_for_rotation;
+          int rot_offset = k_i % embed_dim_for_rotation;
+          scalar_t* ptr = key_for_token + head_idx * (int)head_stride_key;
+          compute_store_rotary_vec<scalar_t, interleaved, aligned_qk>(ptr, curr_data_k, rot_offset, embed_dim_for_rotation);
+        }
+
+        // SHIFT
+        curr_data_k = next_data_k;
+      }
+    }
+  } else {
+    // Fallback to scalar implementation
+    const int nq_pairs = num_heads * embed_dim_for_rotation;
+    for (int i = threadIdx.x; i < nq_pairs; i += blockDim.x) {
+      const int head_idx = i / embed_dim_for_rotation;
+      const int rot_offset = i % embed_dim_for_rotation;
+      scalar_t* query_for_token_head = query_for_token + head_idx * (int)head_stride_query;
+
+      apply_token_rotary_embedding<scalar_t, interleaved>(
+          query_for_token_head, current_token_cos_ptr, current_token_sin_ptr, rot_offset, embed_dim_for_rotation);
+    }
+
+    if (key_for_token != nullptr) {
+      const int nk_pairs = num_kv_heads * embed_dim_for_rotation;
+      for (int i = threadIdx.x; i < nk_pairs; i += blockDim.x) {
+        const int head_idx = i / embed_dim_for_rotation;
+        const int rot_offset = i % embed_dim_for_rotation;
+
+        scalar_t* key_for_token_head = key_for_token + head_idx * (int)head_stride_key;
+        apply_token_rotary_embedding<scalar_t, interleaved>(
+            key_for_token_head, current_token_cos_ptr, current_token_sin_ptr, rot_offset, embed_dim_for_rotation);
+      }
+    }
+  }
+}
+
+void rotary_embedding_cos_sin(
+    at::Tensor& cos,
+    at::Tensor& sin,
+    at::Tensor& query,
+    const std::optional<at::Tensor>& key,
+    int64_t head_size,
+    bool interleaved) {
+  TORCH_CHECK(
+      query.dim() == 2 || query.dim() == 3,
+      "query must be in shape [num_tokens, hidden_size] or [num_tokens, num_heads, head_size]");
+  if (key.has_value()) {
+    TORCH_CHECK(
+        key->dim() == 2 || key->dim() == 3,
+        "key must be in shape [num_tokens, hidden_size] or [num_tokens, num_kv_heads, head_size]");
+  }
+
+  const int64_t num_tokens = query.size(0);
+
+  TORCH_CHECK(cos.dim() == 2, "cos must be in shape [num_tokens, D_cos]");
+  TORCH_CHECK(sin.dim() == 2, "sin must be in shape [num_tokens, D_sin]");
+  TORCH_CHECK(cos.size(0) == num_tokens, "cos num_tokens mismatch with query");
+  TORCH_CHECK(sin.size(0) == num_tokens, "sin num_tokens mismatch with query");
+  TORCH_CHECK(cos.size(1) == sin.size(1), "cos and sin D_cos/D_sin mismatch");
+
+  TORCH_CHECK(cos.scalar_type() == query.scalar_type(), "cos dtype mismatch with query");
+  TORCH_CHECK(sin.scalar_type() == query.scalar_type(), "sin dtype mismatch with query");
+  TORCH_CHECK(cos.is_cuda() && sin.is_cuda() && query.is_cuda(), "cos/sin/query must be CUDA tensors");
+  TORCH_CHECK(query.is_contiguous(), "query must be contiguous");
+  if (key.has_value()) {
+    TORCH_CHECK(key->is_cuda(), "key must be CUDA tensor if provided");
+    TORCH_CHECK(key->scalar_type() == query.scalar_type(), "key dtype mismatch with query");
+    TORCH_CHECK(key->is_contiguous(), "key must be contiguous");
+  }
+
+  int query_hidden_size_calculated;
+  if (query.dim() == 2) {
+    query_hidden_size_calculated = (int)query.size(1);
+  } else {
+    query_hidden_size_calculated = (int)query.size(1) * (int)query.size(2);
+    TORCH_CHECK(query.size(2) == head_size, "query head_size mismatch in 3D tensor");
+  }
+  TORCH_CHECK(query_hidden_size_calculated % head_size == 0, "query_hidden_size not divisible by head_size");
+  int num_heads = query_hidden_size_calculated / (int)head_size;
+
+  int key_hidden_size_calculated = 0;
+  int num_kv_heads = num_heads;
+  if (key.has_value()) {
+    TORCH_CHECK((int)key->size(0) == num_tokens, "key num_tokens mismatch");
+    if (key->dim() == 2) {
+      key_hidden_size_calculated = (int)key->size(1);
+    } else {
+      key_hidden_size_calculated = (int)key->size(1) * (int)key->size(2);
+      TORCH_CHECK((int)key->size(2) == head_size, "key head_size mismatch in 3D tensor");
+    }
+    TORCH_CHECK(key_hidden_size_calculated % head_size == 0, "key_hidden_size not divisible by head_size");
+    num_kv_heads = key_hidden_size_calculated / (int)head_size;
+  }
+  TORCH_CHECK(num_heads % num_kv_heads == 0, "num_heads must be divisible by num_kv_heads");
+
+  // NeoX interleaved: if cos/sin are full-dim (head_size), downsample once.
+  if (interleaved && cos.size(1) == head_size) {
+    TORCH_CHECK(head_size % 2 == 0, "interleaved layout requires even head_size");
+    const int64_t half = head_size / 2;
+    std::vector<int64_t> new_shape = {cos.size(0), half, 2};
+    cos = cos.view(new_shape).select(2, 0).contiguous();
+    sin = sin.view(new_shape).select(2, 1).contiguous();
+  }
+
+  const int rot_dim_from_cache = (int)cos.size(1);
+  const int64_t query_token_stride = query_hidden_size_calculated;
+  const int64_t key_token_stride = key.has_value() ? key_hidden_size_calculated : 0;
+
+  int64_t head_stride_query;
+  if (query.dim() == 3 && query.size(1) == num_heads && query.size(2) == head_size) {
+    head_stride_query = query.stride(1);
+  } else {
+    head_stride_query = head_size;
+  }
+
+  int64_t head_stride_key = head_size;
+  if (key.has_value()) {
+    if (key->dim() == 3 && key->size(1) == num_kv_heads && key->size(2) == head_size) {
+      head_stride_key = key->stride(1);
+    } else {
+      head_stride_key = head_size;
+    }
+  }
+
+  const int embed_dim_for_rotation = interleaved ? rot_dim_from_cache : (rot_dim_from_cache / 2);
+  TORCH_CHECK(embed_dim_for_rotation > 0, "embed_dim_for_rotation must be > 0");
+
+  const int max_pairs_to_rotate_per_token =
+      std::max(num_heads * embed_dim_for_rotation, num_kv_heads * embed_dim_for_rotation);
+
+  const at::cuda::OptionalCUDAGuard device_guard(device_of(query));
+  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+
+  AT_DISPATCH_FLOATING_TYPES_AND2(
+      at::ScalarType::Half, at::ScalarType::BFloat16, query.scalar_type(), "rotary_embedding_cos_sin", [&] {
+        using torch_scalar_t = scalar_t;
+        using cuda_scalar_t = typename std::conditional<
+            std::is_same<torch_scalar_t, at::Half>::value,
+            nv_half,
+            typename std::conditional<std::is_same<torch_scalar_t, at::BFloat16>::value, nv_bfloat16, torch_scalar_t>::
+                type>::type;
+
+        // Constants for vectorization
+        constexpr int kVecBytes = 16;
+        constexpr int kElePerVec = kVecBytes / sizeof(torch_scalar_t);
+
+        // Determine how many pairs one thread handles in one vector step
+        const int pairs_per_step = interleaved ? (kElePerVec / 2) : kElePerVec;
+
+        // Logic Split: Decouple "Vectorized Compute" from "Aligned Memory Access"
+        // Check if we can compute in vectors (Dimensions & Cos/Sin alignment)
+        bool can_vec_compute = true;
+        if (embed_dim_for_rotation % pairs_per_step != 0) can_vec_compute = false;
+
+        // Check Cos/Sin Alignment (Base ptr and Row Stride)
+        if (reinterpret_cast<uintptr_t>(cos.data_ptr()) % kVecBytes != 0) can_vec_compute = false;
+        if (reinterpret_cast<uintptr_t>(sin.data_ptr()) % kVecBytes != 0) can_vec_compute = false;
+        // The stride of the cache usually equals rot_dim. Ensure row starts are aligned.
+        if ((rot_dim_from_cache * sizeof(torch_scalar_t)) % kVecBytes != 0) can_vec_compute = false;
+
+        // Check if Q/K are 16-byte aligned (affects load/store instructions only)
+        bool qk_aligned16 = true;
+
+        if (reinterpret_cast<uintptr_t>(query.data_ptr()) % kVecBytes != 0) qk_aligned16 = false;
+        if ((query_token_stride * sizeof(torch_scalar_t)) % kVecBytes != 0) qk_aligned16 = false;
+        if ((head_stride_query * sizeof(torch_scalar_t)) % kVecBytes != 0) qk_aligned16 = false;
+
+        if (key.has_value()) {
+          if (reinterpret_cast<uintptr_t>(key->data_ptr()) % kVecBytes != 0) qk_aligned16 = false;
+          if ((key_token_stride * sizeof(torch_scalar_t)) % kVecBytes != 0) qk_aligned16 = false;
+          if ((head_stride_key * sizeof(torch_scalar_t)) % kVecBytes != 0) qk_aligned16 = false;
+        }
+
+        // Launch Configuration
+        const bool use_vec = can_vec_compute;
+        const int launch_pairs_per_thread = use_vec ? pairs_per_step : 1;
+        const int total_threads_needed =
+            (max_pairs_to_rotate_per_token + launch_pairs_per_thread - 1) / launch_pairs_per_thread;
+        auto round_up32 = [](int x) { return ((x + 31) / 32) * 32; };
+
+        // Case 1: 2D Grid
+        const int threads_per_block_2d =
+            std::min<int>(512, std::max(128, round_up32(std::min(total_threads_needed, 512))));
+        const int blocks_per_token_2d = (total_threads_needed + threads_per_block_2d - 1) / threads_per_block_2d;
+        const bool use_grid_2d = (num_tokens <= 4) && (blocks_per_token_2d > 1);
+        // Case 2: 1D Grid
+        const int threads_per_block_1d = std::min<int>(512, std::max(128, round_up32(total_threads_needed)));
+        // Final Config
+        const int threads_per_block = use_grid_2d ? threads_per_block_2d : threads_per_block_1d;
+        const int blocks_per_token = use_grid_2d ? blocks_per_token_2d : 1;
+
+        dim3 block(threads_per_block);
+        dim3 grid_2d((int)num_tokens, std::max(1, blocks_per_token));
+        dim3 grid_1d((int)num_tokens);
+        size_t smem_size = 0;
+
+        // Kernel Dispatch
+        DISPATCH_ROT_DIM(
+            embed_dim_for_rotation,
+            auto launch_kernel =
+                [&](bool vec_compute, bool aligned_qk) {
+                  if (interleaved) {
+                    if (use_grid_2d) {
+                      if (vec_compute) {
+                        if (aligned_qk) {
+                          rotary_embedding_kernel_2d<cuda_scalar_t, true, true, true, ROT_EMBED_DIM>
+                              <<<grid_2d, block, smem_size, stream>>>(
+                                  reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
+                                  reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
+                                  reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
+                                  key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>())
+                                                  : nullptr,
+                                  rot_dim_from_cache,
+                                  embed_dim_for_rotation,
+                                  query_token_stride,
+                                  key_token_stride,
+                                  head_stride_query,
+                                  head_stride_key,
+                                  num_heads,
+                                  num_kv_heads,
+                                  (int)head_size,
+                                  blocks_per_token);
+                        } else {
+                          // Vectorized Compute + Unaligned Load/Store
+                          rotary_embedding_kernel_2d<cuda_scalar_t, true, true, false, ROT_EMBED_DIM>
+                              <<<grid_2d, block, smem_size, stream>>>(
+                                  reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
+                                  reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
+                                  reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
+                                  key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>())
+                                                  : nullptr,
+                                  rot_dim_from_cache,
+                                  embed_dim_for_rotation,
+                                  query_token_stride,
+                                  key_token_stride,
+                                  head_stride_query,
+                                  head_stride_key,
+                                  num_heads,
+                                  num_kv_heads,
+                                  (int)head_size,
+                                  blocks_per_token);
+                        }
+                      } else {
+                        // Scalar fallback (aligned_qk param ignored or set to true/false arbitrarily)
+                        rotary_embedding_kernel_2d<cuda_scalar_t, true, false, true, ROT_EMBED_DIM>
+                            <<<grid_2d, block, smem_size, stream>>>(
+                                reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
+                                reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
+                                reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
+                                key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>())
+                                                : nullptr,
+                                rot_dim_from_cache,
+                                embed_dim_for_rotation,
+                                query_token_stride,
+                                key_token_stride,
+                                head_stride_query,
+                                head_stride_key,
+                                num_heads,
+                                num_kv_heads,
+                                (int)head_size,
+                                blocks_per_token);
+                      }
+                    } else {  // 1D Grid
+                      if (vec_compute) {
+                        if (aligned_qk) {
+                          rotary_embedding_kernel_1d<cuda_scalar_t, true, true, true, ROT_EMBED_DIM>
+                              <<<grid_1d, block, smem_size, stream>>>(
+                                  reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
+                                  reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
+                                  reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
+                                  key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>())
+                                                  : nullptr,
+                                  rot_dim_from_cache,
+                                  embed_dim_for_rotation,
+                                  query_token_stride,
+                                  key_token_stride,
+                                  head_stride_query,
+                                  head_stride_key,
+                                  num_heads,
+                                  num_kv_heads,
+                                  (int)head_size);
+                        } else {
+                          rotary_embedding_kernel_1d<cuda_scalar_t, true, true, false, ROT_EMBED_DIM>
+                              <<<grid_1d, block, smem_size, stream>>>(
+                                  reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
+                                  reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
+                                  reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
+                                  key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>())
+                                                  : nullptr,
+                                  rot_dim_from_cache,
+                                  embed_dim_for_rotation,
+                                  query_token_stride,
+                                  key_token_stride,
+                                  head_stride_query,
+                                  head_stride_key,
+                                  num_heads,
+                                  num_kv_heads,
+                                  (int)head_size);
+                        }
+                      } else {
+                        rotary_embedding_kernel_1d<cuda_scalar_t, true, false, true, ROT_EMBED_DIM>
+                            <<<grid_1d, block, smem_size, stream>>>(
+                                reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
+                                reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
+                                reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
+                                key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>())
+                                                : nullptr,
+                                rot_dim_from_cache,
+                                embed_dim_for_rotation,
+                                query_token_stride,
+                                key_token_stride,
+                                head_stride_query,
+                                head_stride_key,
+                                num_heads,
+                                num_kv_heads,
+                                (int)head_size);
+                      }
+                    }
+                  } else {  // Non-Interleaved
+                    if (use_grid_2d) {
+                      if (vec_compute) {
+                        if (aligned_qk) {
+                          rotary_embedding_kernel_2d<cuda_scalar_t, false, true, true, ROT_EMBED_DIM>
+                              <<<grid_2d, block, smem_size, stream>>>(
+                                  reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
+                                  reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
+                                  reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
+                                  key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>())
+                                                  : nullptr,
+                                  rot_dim_from_cache,
+                                  embed_dim_for_rotation,
+                                  query_token_stride,
+                                  key_token_stride,
+                                  head_stride_query,
+                                  head_stride_key,
+                                  num_heads,
+                                  num_kv_heads,
+                                  (int)head_size,
+                                  blocks_per_token);
+                        } else {
+                          rotary_embedding_kernel_2d<cuda_scalar_t, false, true, false, ROT_EMBED_DIM>
+                              <<<grid_2d, block, smem_size, stream>>>(
+                                  reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
+                                  reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
+                                  reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
+                                  key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>())
+                                                  : nullptr,
+                                  rot_dim_from_cache,
+                                  embed_dim_for_rotation,
+                                  query_token_stride,
+                                  key_token_stride,
+                                  head_stride_query,
+                                  head_stride_key,
+                                  num_heads,
+                                  num_kv_heads,
+                                  (int)head_size,
+                                  blocks_per_token);
+                        }
+                      } else {
+                        rotary_embedding_kernel_2d<cuda_scalar_t, false, false, true, ROT_EMBED_DIM>
+                            <<<grid_2d, block, smem_size, stream>>>(
+                                reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
+                                reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
+                                reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
+                                key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>())
+                                                : nullptr,
+                                rot_dim_from_cache,
+                                embed_dim_for_rotation,
+                                query_token_stride,
+                                key_token_stride,
+                                head_stride_query,
+                                head_stride_key,
+                                num_heads,
+                                num_kv_heads,
+                                (int)head_size,
+                                blocks_per_token);
+                      }
+                    } else {  // 1D Grid
+                      if (vec_compute) {
+                        if (aligned_qk) {
+                          rotary_embedding_kernel_1d<cuda_scalar_t, false, true, true, ROT_EMBED_DIM>
+                              <<<grid_1d, block, smem_size, stream>>>(
+                                  reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
+                                  reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
+                                  reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
+                                  key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>())
+                                                  : nullptr,
+                                  rot_dim_from_cache,
+                                  embed_dim_for_rotation,
+                                  query_token_stride,
+                                  key_token_stride,
+                                  head_stride_query,
+                                  head_stride_key,
+                                  num_heads,
+                                  num_kv_heads,
+                                  (int)head_size);
+                        } else {
+                          rotary_embedding_kernel_1d<cuda_scalar_t, false, true, false, ROT_EMBED_DIM>
+                              <<<grid_1d, block, smem_size, stream>>>(
+                                  reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
+                                  reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
+                                  reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
+                                  key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>())
+                                                  : nullptr,
+                                  rot_dim_from_cache,
+                                  embed_dim_for_rotation,
+                                  query_token_stride,
+                                  key_token_stride,
+                                  head_stride_query,
+                                  head_stride_key,
+                                  num_heads,
+                                  num_kv_heads,
+                                  (int)head_size);
+                        }
+                      } else {
+                        rotary_embedding_kernel_1d<cuda_scalar_t, false, false, true, ROT_EMBED_DIM>
+                            <<<grid_1d, block, smem_size, stream>>>(
+                                reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
+                                reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
+                                reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
+                                key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>())
+                                                : nullptr,
+                                rot_dim_from_cache,
+                                embed_dim_for_rotation,
+                                query_token_stride,
+                                key_token_stride,
+                                head_stride_query,
+                                head_stride_key,
+                                num_heads,
+                                num_kv_heads,
+                                (int)head_size);
+                      }
+                    }
+                  }
+                };
+
+            // Execute the launch
+            launch_kernel(use_vec, qk_aligned16););
+      });
+
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+}
\ No newline at end of file
diff --git a/sgl-kernel/include/utils.h b/sgl-kernel/include/utils.h
index c5a517cea019..3caf71cd6e50 100644
--- a/sgl-kernel/include/utils.h
+++ b/sgl-kernel/include/utils.h
@@ -459,11 +459,63 @@ inline uint32_t next_pow2(uint32_t x) noexcept {
   return 1u << (32 - __builtin_clz(x - 1));
 }
 
-/*
- * LDG Support
- */
+
 #ifndef USE_ROCM
 #define SGLANG_LDG(arg) __ldg(arg)
 #else
 #define SGLANG_LDG(arg) *(arg)
 #endif
+
+
+// Define dispatch macro for Rotation Dimension (Pair Count)
+// Case 32 -> head_size 64
+// Case 64 -> head_size 128
+// Case 128 -> head_size 256
+#define DISPATCH_ROT_DIM(embed_dim, ...)  \
+  switch (embed_dim) {                    \
+    case 16: {                            \
+      constexpr int ROT_EMBED_DIM = 16;   \
+      __VA_ARGS__;                        \
+      break;                              \
+    }                                     \
+    case 32: {                            \
+      constexpr int ROT_EMBED_DIM = 32;   \
+      __VA_ARGS__;                        \
+      break;                              \
+    }                                     \
+    case 64: {                            \
+      constexpr int ROT_EMBED_DIM = 64;   \
+      __VA_ARGS__;                        \
+      break;                              \
+    }                                     \
+    case 96: {                            \
+      constexpr int ROT_EMBED_DIM = 96;   \
+      __VA_ARGS__;                        \
+      break;                              \
+    }                                     \
+    case 128: {                           \
+      constexpr int ROT_EMBED_DIM = 128;  \
+      __VA_ARGS__;                        \
+      break;                              \
+    }                                     \
+    case 256: {                           \
+      constexpr int ROT_EMBED_DIM = 256;  \
+      __VA_ARGS__;                        \
+      break;                              \
+    }                                     \
+    case 512: {                           \
+      constexpr int ROT_EMBED_DIM = 512;  \
+      __VA_ARGS__;                        \
+      break;                              \
+    }                                     \
+    case 1024: {                          \
+      constexpr int ROT_EMBED_DIM = 1024; \
+      __VA_ARGS__;                        \
+      break;                              \
+    }                                     \
+    default: {                            \
+      constexpr int ROT_EMBED_DIM = 0;    \
+      __VA_ARGS__;                        \
+      break;                              \
+    }                                     \
+  }

From 6d1e0f8e381edc5ad778a7b2cd5e78be81ed33ba Mon Sep 17 00:00:00 2001
From: RubiaCx <1084281732@qq.com>
Date: Thu, 18 Dec 2025 22:04:24 -0800
Subject: [PATCH 18/37] Fix formatting issues from pre-commit hooks

---
 .../benchmark/bench_mm_rotary_embedding.py    | 95 ++++++++++++++-----
 .../sgl_diffusion/rope/rotary_embedding.cu    | 19 ++--
 sgl-kernel/include/utils.h                    |  2 -
 3 files changed, 82 insertions(+), 34 deletions(-)

diff --git a/sgl-kernel/benchmark/bench_mm_rotary_embedding.py b/sgl-kernel/benchmark/bench_mm_rotary_embedding.py
index 8927e7e98abf..b3d833447e33 100755
--- a/sgl-kernel/benchmark/bench_mm_rotary_embedding.py
+++ b/sgl-kernel/benchmark/bench_mm_rotary_embedding.py
@@ -38,10 +38,14 @@ def _expand_neox_full_repeat(half: torch.Tensor) -> torch.Tensor:
     out[:, 1::2] = half
     return out
 
+
 def _expand_llama_full_cat(half: torch.Tensor) -> torch.Tensor:
     return torch.cat([half, half], dim=-1)
 
-def _make_misaligned_contiguous_like(t: torch.Tensor, offset_elems: int = 1) -> torch.Tensor:
+
+def _make_misaligned_contiguous_like(
+    t: torch.Tensor, offset_elems: int = 1
+) -> torch.Tensor:
     # Create a contiguous tensor with same shape/dtype/device but misaligned base pointer.
     numel = t.numel()
     buf = torch.empty((numel + offset_elems,), device=t.device, dtype=t.dtype)
@@ -51,7 +55,9 @@ def _make_misaligned_contiguous_like(t: torch.Tensor, offset_elems: int = 1) ->
     return out
 
 
-def _make_tensor(shape: Tuple[int, ...], *, dtype: torch.dtype, device: str, misalign: bool) -> torch.Tensor:
+def _make_tensor(
+    shape: Tuple[int, ...], *, dtype: torch.dtype, device: str, misalign: bool
+) -> torch.Tensor:
     numel = int(np.prod(shape))
     if not misalign:
         return torch.randn(*shape, dtype=dtype, device=device)
@@ -194,14 +200,22 @@ def _predict_use_grid_2d(
     pairs_per_step = (kElePerVec // 2) if interleaved else kElePerVec
     launch_pairs_per_thread = pairs_per_step if use_vec_compute else 1
 
-    max_pairs = max(num_heads * embed_dim_for_rotation, num_kv_heads * embed_dim_for_rotation)
-    total_threads_needed = (max_pairs + launch_pairs_per_thread - 1) // launch_pairs_per_thread
+    max_pairs = max(
+        num_heads * embed_dim_for_rotation, num_kv_heads * embed_dim_for_rotation
+    )
+    total_threads_needed = (
+        max_pairs + launch_pairs_per_thread - 1
+    ) // launch_pairs_per_thread
 
     def _round_up32(x: int) -> int:
         return ((x + 31) // 32) * 32
 
-    threads_per_block_2d = min(512, max(128, _round_up32(min(total_threads_needed, 512))))
-    blocks_per_token_2d = (total_threads_needed + threads_per_block_2d - 1) // threads_per_block_2d
+    threads_per_block_2d = min(
+        512, max(128, _round_up32(min(total_threads_needed, 512)))
+    )
+    blocks_per_token_2d = (
+        total_threads_needed + threads_per_block_2d - 1
+    ) // threads_per_block_2d
     use_grid_2d = (num_tokens <= 4) and (blocks_per_token_2d > 1)
     return use_grid_2d, blocks_per_token_2d
 
@@ -488,13 +502,20 @@ def benchmark_mm_rotary_embedding(args: argparse.Namespace) -> None:
                     raise ValueError(f"Unknown layout={case.layout}")
 
                 query = _make_tensor(
-                    query_shape, dtype=case.dtype, device=device, misalign=(case.misalign == "q")
+                    query_shape,
+                    dtype=case.dtype,
+                    device=device,
+                    misalign=(case.misalign == "q"),
+                )
+                key = _make_tensor(
+                    key_shape, dtype=case.dtype, device=device, misalign=False
                 )
-                key = _make_tensor(key_shape, dtype=case.dtype, device=device, misalign=False)
                 if case.key_mode == "no_k":
                     key = None
 
-                positions = torch.arange(seq_len, device=device, dtype=torch.int64).repeat(case.batch_size)
+                positions = torch.arange(
+                    seq_len, device=device, dtype=torch.int64
+                ).repeat(case.batch_size)
                 cos = cos_cache[positions]
                 sin = sin_cache[positions]
 
@@ -552,7 +573,9 @@ def fn_sgl_cos_sin() -> None:
                 k = None if key is None else key.clone()
                 sgl_rotary_cos_sin(cos, sin, q, k, case.head_size, case.interleaved)
 
-            ms, _, _ = triton.testing.do_bench(fn_sgl_cos_sin, quantiles=[0.5, 0.2, 0.8])
+            ms, _, _ = triton.testing.do_bench(
+                fn_sgl_cos_sin, quantiles=[0.5, 0.2, 0.8]
+            )
             sgl_cos_sin_time = 1000 * ms
             row_str += f" | {sgl_cos_sin_time:16.4f}"
 
@@ -569,7 +592,9 @@ def fn_sgl_pos() -> None:
                 )
 
             try:
-                ms, _, _ = triton.testing.do_bench(fn_sgl_pos, quantiles=[0.5, 0.2, 0.8])
+                ms, _, _ = triton.testing.do_bench(
+                    fn_sgl_pos, quantiles=[0.5, 0.2, 0.8]
+                )
                 sgl_pos_time = 1000 * ms
                 row_str += f" | {sgl_pos_time:12.4f}"
             except Exception:
@@ -578,7 +603,9 @@ def fn_sgl_pos() -> None:
 
             # Prediction flags (updated for optimized kernel)
             rot_dim_from_cache = int(cos.shape[1])
-            embed_dim_for_rotation = rot_dim_from_cache if case.interleaved else (rot_dim_from_cache // 2)
+            embed_dim_for_rotation = (
+                rot_dim_from_cache if case.interleaved else (rot_dim_from_cache // 2)
+            )
 
             if query.dim() == 2:
                 query_hidden = int(query.size(1))
@@ -652,7 +679,9 @@ def fn_vllm() -> None:
                     vllm_rope.forward_cuda(positions, q, k)
 
                 try:
-                    ms, _, _ = triton.testing.do_bench(fn_vllm, quantiles=[0.5, 0.2, 0.8])
+                    ms, _, _ = triton.testing.do_bench(
+                        fn_vllm, quantiles=[0.5, 0.2, 0.8]
+                    )
                     vllm_time = 1000 * ms
                     row_str += f" | {vllm_time:10.4f}"
                 except Exception:
@@ -661,7 +690,9 @@ def fn_vllm() -> None:
 
             if HAS_FLASH_ATTN:
                 try:
-                    from flash_attn.layers.rotary import RotaryEmbedding as FlashRotaryEmbedding  # noqa
+                    from flash_attn.layers.rotary import (  # noqa
+                        RotaryEmbedding as FlashRotaryEmbedding,
+                    )
 
                     flash_rotary = FlashRotaryEmbedding(case.rotary_dim, device=device)
                     qkv = torch.randn(
@@ -678,14 +709,20 @@ def fn_flash_attn() -> None:
                         qkv_fa = qkv.clone()
                         flash_rotary(qkv_fa, seqlen_offset=0)
 
-                    ms, _, _ = triton.testing.do_bench(fn_flash_attn, quantiles=[0.5, 0.2, 0.8])
+                    ms, _, _ = triton.testing.do_bench(
+                        fn_flash_attn, quantiles=[0.5, 0.2, 0.8]
+                    )
                     fa_time = 1000 * ms
                     row_str += f" | {fa_time:14.4f}"
                 except Exception:
                     row_str += f" | {'ERROR':>14}"
                     fa_time = None
 
-            if (sgl_pos_time is not None) and (sgl_cos_sin_time is not None) and (sgl_cos_sin_time > 0):
+            if (
+                (sgl_pos_time is not None)
+                and (sgl_cos_sin_time is not None)
+                and (sgl_cos_sin_time > 0)
+            ):
                 speedup = sgl_pos_time / sgl_cos_sin_time
                 row_str += f" | {speedup:9.2f}x"
             else:
@@ -706,21 +743,35 @@ def fn_flash_attn() -> None:
 
         native_vals = [r["native"] for r in results if r["native"] is not None]
         naive_vals = [r["naive"] for r in results if r["naive"] is not None]
-        sgl_cos_sin_vals = [r["sgl_cos_sin"] for r in results if r["sgl_cos_sin"] is not None]
+        sgl_cos_sin_vals = [
+            r["sgl_cos_sin"] for r in results if r["sgl_cos_sin"] is not None
+        ]
         sgl_pos_vals = [r["sgl_pos"] for r in results if r["sgl_pos"] is not None]
         vllm_vals = [r["vllm"] for r in results if r["vllm"] is not None]
         fa_vals = [r["flash_attn"] for r in results if r["flash_attn"] is not None]
 
         avg_row = f"{'AVG':>8}"
         if case.bench_native:
-            avg_row += f" | {np.mean(native_vals):12.4f}" if native_vals else f" | {'N/A':>12}"
-        avg_row += f" | {np.mean(naive_vals):11.4f}" if naive_vals else f" | {'N/A':>11}"
-        avg_row += f" | {np.mean(sgl_cos_sin_vals):16.4f}" if sgl_cos_sin_vals else f" | {'N/A':>16}"
-        avg_row += f" | {np.mean(sgl_pos_vals):12.4f}" if sgl_pos_vals else f" | {'N/A':>12}"
+            avg_row += (
+                f" | {np.mean(native_vals):12.4f}" if native_vals else f" | {'N/A':>12}"
+            )
+        avg_row += (
+            f" | {np.mean(naive_vals):11.4f}" if naive_vals else f" | {'N/A':>11}"
+        )
+        avg_row += (
+            f" | {np.mean(sgl_cos_sin_vals):16.4f}"
+            if sgl_cos_sin_vals
+            else f" | {'N/A':>16}"
+        )
+        avg_row += (
+            f" | {np.mean(sgl_pos_vals):12.4f}" if sgl_pos_vals else f" | {'N/A':>12}"
+        )
         avg_row += f" | {'-':>6} | {'-':>4} | {'-':>5}"
 
         if HAS_VLLM:
-            avg_row += f" | {np.mean(vllm_vals):10.4f}" if vllm_vals else f" | {'N/A':>10}"
+            avg_row += (
+                f" | {np.mean(vllm_vals):10.4f}" if vllm_vals else f" | {'N/A':>10}"
+            )
         if HAS_FLASH_ATTN:
             avg_row += f" | {np.mean(fa_vals):14.4f}" if fa_vals else f" | {'N/A':>14}"
         if sgl_pos_vals and sgl_cos_sin_vals and (np.mean(sgl_cos_sin_vals) > 0):
diff --git a/sgl-kernel/csrc/sgl_diffusion/rope/rotary_embedding.cu b/sgl-kernel/csrc/sgl_diffusion/rope/rotary_embedding.cu
index 7b8c57d964c8..0bed41a8f09e 100644
--- a/sgl-kernel/csrc/sgl_diffusion/rope/rotary_embedding.cu
+++ b/sgl-kernel/csrc/sgl_diffusion/rope/rotary_embedding.cu
@@ -116,10 +116,7 @@ __device__ __forceinline__ RotaryVecData<scalar_t, interleaved> load_rotary_vec(
 // Compute RoPE on the loaded vector tile, then store back to q/k.
 template <typename scalar_t, bool interleaved, bool aligned_qk>
 __device__ __forceinline__ void compute_store_rotary_vec(
-    scalar_t* __restrict__ arr, 
-    const RotaryVecData<scalar_t, interleaved>& data, 
-    int rot_offset, 
-    int embed_dim) {
+    scalar_t* __restrict__ arr, const RotaryVecData<scalar_t, interleaved>& data, int rot_offset, int embed_dim) {
   constexpr int kVecBytes = 16;
   constexpr int kElePerVec = kVecBytes / sizeof(scalar_t);
 
@@ -211,7 +208,7 @@ inline __device__ void apply_token_rotary_embedding(
     const scalar_t* __restrict__ cos_ptr,  // [rot_dim]
     const scalar_t* __restrict__ sin_ptr,  // [rot_dim]
     int rot_offset,
-    int embed_dim) {  // for non-interleaved: half dim
+    int embed_dim) {            // for non-interleaved: half dim
   if constexpr (interleaved) {  // NeoX-style: interleaved layout [x0, y0, x1, y1, ...].
     const int x_index = 2 * rot_offset;
     const int y_index = x_index + 1;
@@ -224,7 +221,7 @@ inline __device__ void apply_token_rotary_embedding(
     const float y = static_cast<float>(arr[y_index]);
     arr[x_index] = static_cast<scalar_t>(x * cos_val - y * sin_val);
     arr[y_index] = static_cast<scalar_t>(y * cos_val + x * sin_val);
-  } else { // GPT-J / LLaMA style: layout [x0, x1, ..., y0, y1, ...]
+  } else {  // GPT-J / LLaMA style: layout [x0, x1, ..., y0, y1, ...]
     const int x_index = rot_offset;
     const int y_index = rot_offset + embed_dim;
 
@@ -290,7 +287,7 @@ inline __device__ void apply_token_rotary_embedding_vec(
 
     *reinterpret_cast<vec_t*>(arr + rot_offset * 2) = data.vec;
 
-  } else { // Non-interleaved
+  } else {  // Non-interleaved
     VecUnion data_x, data_y;
     data_x.vec = *reinterpret_cast<const vec_t*>(arr + rot_offset);
     data_y.vec = *reinterpret_cast<const vec_t*>(arr + rot_offset + embed_dim);
@@ -462,7 +459,8 @@ __global__ void rotary_embedding_kernel_2d(
         int head_idx = k_i / embed_dim_for_rotation;
         int rot_offset = k_i % embed_dim_for_rotation;
         scalar_t* ptr = key_for_token + head_idx * (int)head_stride_key;
-        compute_store_rotary_vec<scalar_t, interleaved, aligned_qk>(ptr, curr_data_k, rot_offset, embed_dim_for_rotation);
+        compute_store_rotary_vec<scalar_t, interleaved, aligned_qk>(
+            ptr, curr_data_k, rot_offset, embed_dim_for_rotation);
 
         // SHIFT
         curr_data_k = next_data_k;
@@ -614,7 +612,8 @@ __launch_bounds__(512) __global__ void rotary_embedding_kernel_1d(
           int head_idx = k_i / embed_dim_for_rotation;
           int rot_offset = k_i % embed_dim_for_rotation;
           scalar_t* ptr = key_for_token + head_idx * (int)head_stride_key;
-          compute_store_rotary_vec<scalar_t, interleaved, aligned_qk>(ptr, curr_data_k, rot_offset, embed_dim_for_rotation);
+          compute_store_rotary_vec<scalar_t, interleaved, aligned_qk>(
+              ptr, curr_data_k, rot_offset, embed_dim_for_rotation);
         }
 
         // SHIFT
@@ -1052,4 +1051,4 @@ void rotary_embedding_cos_sin(
       });
 
   C10_CUDA_KERNEL_LAUNCH_CHECK();
-}
\ No newline at end of file
+}
diff --git a/sgl-kernel/include/utils.h b/sgl-kernel/include/utils.h
index 3caf71cd6e50..c55d8feb2d69 100644
--- a/sgl-kernel/include/utils.h
+++ b/sgl-kernel/include/utils.h
@@ -459,14 +459,12 @@ inline uint32_t next_pow2(uint32_t x) noexcept {
   return 1u << (32 - __builtin_clz(x - 1));
 }
 
-
 #ifndef USE_ROCM
 #define SGLANG_LDG(arg) __ldg(arg)
 #else
 #define SGLANG_LDG(arg) *(arg)
 #endif
 
-
 // Define dispatch macro for Rotation Dimension (Pair Count)
 // Case 32 -> head_size 64
 // Case 64 -> head_size 128

From 745a55fc1cbd87c45a8000d762d7e28dd3fb7645 Mon Sep 17 00:00:00 2001
From: RubiaCx <1084281732@qq.com>
Date: Sat, 20 Dec 2025 06:14:41 +0000
Subject: [PATCH 19/37] Fix typo and fix RoPE rotary embedding kernel/interface
 compatibility.

---
 .../sgl_diffusion/rope/rotary_embedding.cu    |  2 +-
 sgl-kernel/include/utils.h                    |  2 +-
 .../python/sgl_kernel/rotary_embedding.py     | 37 ++-----------------
 3 files changed, 6 insertions(+), 35 deletions(-)

diff --git a/sgl-kernel/csrc/sgl_diffusion/rope/rotary_embedding.cu b/sgl-kernel/csrc/sgl_diffusion/rope/rotary_embedding.cu
index 0bed41a8f09e..31fdce6ed056 100644
--- a/sgl-kernel/csrc/sgl_diffusion/rope/rotary_embedding.cu
+++ b/sgl-kernel/csrc/sgl_diffusion/rope/rotary_embedding.cu
@@ -807,7 +807,7 @@ void rotary_embedding_cos_sin(
         size_t smem_size = 0;
 
         // Kernel Dispatch
-        DISPATCH_ROT_DIM(
+        _DIFFUSION_ROT_EMBED_DIM(
             embed_dim_for_rotation,
             auto launch_kernel =
                 [&](bool vec_compute, bool aligned_qk) {
diff --git a/sgl-kernel/include/utils.h b/sgl-kernel/include/utils.h
index c55d8feb2d69..be5ba0726780 100644
--- a/sgl-kernel/include/utils.h
+++ b/sgl-kernel/include/utils.h
@@ -469,7 +469,7 @@ inline uint32_t next_pow2(uint32_t x) noexcept {
 // Case 32 -> head_size 64
 // Case 64 -> head_size 128
 // Case 128 -> head_size 256
-#define DISPATCH_ROT_DIM(embed_dim, ...)  \
+#define _DIFFUSION_ROT_EMBED_DIM(embed_dim, ...)  \
   switch (embed_dim) {                    \
     case 16: {                            \
       constexpr int ROT_EMBED_DIM = 16;   \
diff --git a/sgl-kernel/python/sgl_kernel/rotary_embedding.py b/sgl-kernel/python/sgl_kernel/rotary_embedding.py
index ed890b759096..307c6b6b82ce 100644
--- a/sgl-kernel/python/sgl_kernel/rotary_embedding.py
+++ b/sgl-kernel/python/sgl_kernel/rotary_embedding.py
@@ -5,7 +5,7 @@
 import torch
 
 _HAS_COS_SIN = hasattr(torch.ops.sgl_kernel, "rotary_embedding_cos_sin")
-_HAS_GENERIC = hasattr(torch.ops.sgl_kernel, "rotary_embedding")
+_HAS_POSITION = hasattr(torch.ops.sgl_kernel, "rotary_embedding")
 
 
 def _resolve_interleaved(interleaved: Optional[bool], is_neox: Optional[bool]) -> bool:
@@ -36,16 +36,6 @@ def apply_rotary_embedding(
     positions: Optional[torch.Tensor] = None,
     cos_sin_cache: Optional[torch.Tensor] = None,
 ) -> None:
-    """内部统一的 rotary dispatch：
-
-    - mode="cos_sin": 使用 (cos, sin) 直接旋转
-    - mode="positions": 使用 (positions, cos_sin_cache) 旋转
-
-    interleaved/is_neox:
-      - True (NeoX): [x0, y0, x1, y1, ...]
-      - False (GPT-J/LLaMA): [x0, x1, ..., y0, y1, ...]
-    """
-
     effective_interleaved = _resolve_interleaved(interleaved, is_neox)
 
     if mode == "cos_sin":
@@ -63,8 +53,7 @@ def apply_rotary_embedding(
             )
             return
 
-        # Some older builds overload `rotary_embedding` directly with cos/sin signature.
-        if _HAS_GENERIC:
+        if _HAS_POSITION:
             torch.ops.sgl_kernel.rotary_embedding(
                 cos,
                 sin,
@@ -83,27 +72,12 @@ def apply_rotary_embedding(
         if positions is None or cos_sin_cache is None:
             raise ValueError("mode='positions' requires positions and cos_sin_cache")
 
-        if not _HAS_GENERIC:
+        if not _HAS_POSITION:
             raise RuntimeError(
                 "No positions rotary embedding kernel is available in torch.ops.sgl_kernel"
             )
 
-        # Prefer keyword call if supported.
-        try:
-            torch.ops.sgl_kernel.rotary_embedding(
-                positions=positions,
-                query=query,
-                key=key if key is not None else None,
-                head_size=head_size,
-                cos_sin_cache=cos_sin_cache,
-                is_neox=effective_interleaved,
-            )
-            return
-        except Exception:
-            pass
-
-        # Fallback to positional signature used by some builds/tests:
-        # rotary_embedding(positions, query, key, head_size, cos_sin_cache, is_neox)
+        # Use positional args for maximum compatibility across builds.
         torch.ops.sgl_kernel.rotary_embedding(
             positions,
             query,
@@ -114,9 +88,6 @@ def apply_rotary_embedding(
         )
         return
 
-    raise ValueError(f"Unknown mode: {mode}")
-
-
 def rotary_embedding_cos_sin(
     cos: torch.Tensor,
     sin: torch.Tensor,

From cb0724ba2965c2b0f4aa492e62bfe670ad5b6e80 Mon Sep 17 00:00:00 2001
From: RubiaCx <1084281732@qq.com>
Date: Mon, 29 Dec 2025 23:40:11 -0800
Subject: [PATCH 20/37] Align with JIT kernels guide; keep behavior consistent
 with the AOT implementation.

---
 .../benchmark/bench_rotary_embedding.py       |  420 +++++++
 .../csrc/rotary_embedding_cos_sin.cuh         | 1106 +++++++++++++++++
 python/sglang/jit_kernel/rotary_embedding.py  |  142 +++
 .../jit_kernel/tests/test_rotary_embedding.py |  203 +++
 .../runtime/layers/rotary_embedding.py        |  102 ++
 .../runtime/models/dits/qwen_image.py         |    7 +-
 sgl-kernel/CMakeLists.txt                     |    1 -
 .../benchmark/bench_mm_rotary_embedding.py    |  795 ------------
 sgl-kernel/csrc/common_extension.cc           |   21 +-
 .../sgl_diffusion/rope/rotary_embedding.cu    | 1054 ----------------
 sgl-kernel/include/sgl_kernel_ops.h           |   21 -
 sgl-kernel/include/utils.h                    |   53 -
 sgl-kernel/python/sgl_kernel/__init__.py      |    1 -
 .../python/sgl_kernel/rotary_embedding.py     |  132 --
 sgl-kernel/tests/test_mm_rotary_embedding.py  |  198 ---
 15 files changed, 1980 insertions(+), 2276 deletions(-)
 create mode 100644 python/sglang/jit_kernel/benchmark/bench_rotary_embedding.py
 create mode 100644 python/sglang/jit_kernel/csrc/rotary_embedding_cos_sin.cuh
 create mode 100644 python/sglang/jit_kernel/rotary_embedding.py
 create mode 100644 python/sglang/jit_kernel/tests/test_rotary_embedding.py
 delete mode 100755 sgl-kernel/benchmark/bench_mm_rotary_embedding.py
 delete mode 100644 sgl-kernel/csrc/sgl_diffusion/rope/rotary_embedding.cu
 delete mode 100644 sgl-kernel/python/sgl_kernel/rotary_embedding.py
 delete mode 100755 sgl-kernel/tests/test_mm_rotary_embedding.py

diff --git a/python/sglang/jit_kernel/benchmark/bench_rotary_embedding.py b/python/sglang/jit_kernel/benchmark/bench_rotary_embedding.py
new file mode 100644
index 000000000000..820a47c20b57
--- /dev/null
+++ b/python/sglang/jit_kernel/benchmark/bench_rotary_embedding.py
@@ -0,0 +1,420 @@
+import itertools
+import os
+from functools import lru_cache
+from typing import Optional, Tuple
+import torch
+import triton
+import triton.testing
+
+
+IS_CI = (
+    os.getenv("CI", "false").lower() == "true"
+    or os.getenv("GITHUB_ACTIONS", "false").lower() == "true"
+)
+
+IS_QUICK = os.getenv("SGL_BENCH_QUICK", "0").lower() in ("1", "true", "yes", "y")
+ONLY_PROVIDER = os.getenv("SGL_BENCH_PROVIDER", "").strip() or None
+SHOW_SPEEDUP = os.getenv("SGL_BENCH_SPEEDUP", "1").lower() in ("1", "true", "yes", "y")
+
+DEVICE = "cuda"
+MAX_SEQ_LEN = 8192
+NUM_Q_HEADS = 32
+NUM_KV_HEADS = 8
+DTYPE = torch.bfloat16
+
+try:
+    from vllm.model_executor.layers.rotary_embedding import (
+        RotaryEmbedding as vLLMRotaryEmbedding,
+    )
+
+    HAS_VLLM = True
+except Exception:
+    vLLMRotaryEmbedding = None
+    HAS_VLLM = False
+
+if IS_CI or IS_QUICK:
+    BS_RANGE = [16]
+    SEQ_RANGE = [1, 128]
+    HEAD_SIZE_RANGE = [128]
+    INTERLEAVED_RANGE = [True]
+else:
+    BS_RANGE = [1, 8, 64]
+    SEQ_RANGE = [1, 4, 128, 2048]
+    HEAD_SIZE_RANGE = [64, 96, 128, 256]
+    INTERLEAVED_RANGE = [True, False]
+
+
+def _compute_cos_sin_cache_half(
+    max_seq_len: int,
+    rotary_dim: int,
+    base: float = 10000.0,
+    dtype: torch.dtype = torch.float32,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    inv_freq = 1.0 / (
+        base ** (torch.arange(0, rotary_dim, 2, dtype=torch.float32) / rotary_dim)
+    )
+    t = torch.arange(max_seq_len, dtype=torch.float32)
+    freqs = torch.einsum("i,j->ij", t, inv_freq)
+    cos = freqs.cos().to(dtype)
+    sin = freqs.sin().to(dtype)
+    return cos, sin
+
+
+@lru_cache(maxsize=None)
+def _get_cos_sin_cache_half_cuda(rotary_dim: int) -> Tuple[torch.Tensor, torch.Tensor]:
+    cos, sin = _compute_cos_sin_cache_half(MAX_SEQ_LEN, rotary_dim, dtype=DTYPE)
+    return cos.to(device=DEVICE, dtype=DTYPE), sin.to(device=DEVICE, dtype=DTYPE)
+
+
+@torch.no_grad()
+def torch_impl_rotary_fp32(
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    q: torch.Tensor,
+    k: torch.Tensor,
+    head_size: int,
+    interleaved: bool,
+) -> None:
+    orig_dtype = q.dtype
+    if interleaved and cos.shape[1] == head_size:
+        half = head_size // 2
+        cos = cos.view(cos.shape[0], half, 2).select(2, 0).contiguous()
+        sin = sin.view(sin.shape[0], half, 2).select(2, 1).contiguous()
+
+    cos_f, sin_f = cos.float(), sin.float()
+    q_f, k_f = q.float(), k.float()
+
+    if interleaved:
+        embed_dim = int(cos_f.shape[1])
+        rot_dim = embed_dim * 2
+        cos_b = cos_f[:, None, :embed_dim]
+        sin_b = sin_f[:, None, :embed_dim]
+
+        def _apply(x: torch.Tensor) -> None:
+            xr = x[..., :rot_dim]
+            xr2 = xr.view(xr.shape[0], xr.shape[1], embed_dim, 2)
+            x0 = xr2[..., 0].clone()
+            x1 = xr2[..., 1].clone()
+            xr2[..., 0].copy_(x0 * cos_b - x1 * sin_b)
+            xr2[..., 1].copy_(x1 * cos_b + x0 * sin_b)
+    else:
+        if cos_f.shape[1] == head_size // 2:
+            embed_dim = int(cos_f.shape[1])
+            rot_dim = embed_dim * 2
+            cos_x, sin_x = cos_f[:, None, :], sin_f[:, None, :]
+            cos_y, sin_y = cos_x, sin_x
+        else:
+            embed_dim = int(cos_f.shape[1]) // 2
+            rot_dim = embed_dim * 2
+            cos_x, sin_x = cos_f[:, None, :embed_dim], sin_f[:, None, :embed_dim]
+            cos_y, sin_y = (
+                cos_f[:, None, embed_dim:rot_dim],
+                sin_f[:, None, embed_dim:rot_dim],
+            )
+
+        def _apply(x: torch.Tensor) -> None:
+            xr = x[..., :rot_dim]
+            x0 = xr[..., :embed_dim].clone()
+            x1 = xr[..., embed_dim:rot_dim].clone()
+            xr[..., :embed_dim].copy_(x0 * cos_x - x1 * sin_x)
+            xr[..., embed_dim:rot_dim].copy_(x1 * cos_y + x0 * sin_y)
+
+    _apply(q_f)
+    _apply(k_f)
+    q.copy_(q_f.to(orig_dtype))
+    k.copy_(k_f.to(orig_dtype))
+
+
+def sglang_aot_rotary_positions(
+    positions: torch.Tensor,
+    q: torch.Tensor,
+    k: torch.Tensor,
+    head_size: int,
+    interleaved: bool,
+    cos_sin_cache: torch.Tensor,
+) -> None:
+    from sgl_kernel.rotary_embedding import rotary_embedding
+
+    rotary_embedding(
+        positions=positions,
+        query=q,
+        key=k,
+        head_size=head_size,
+        is_neox=not interleaved,
+        cos_sin_cache=cos_sin_cache,
+    )
+
+
+def sglang_jit_rotary_cos_sin(
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    q: torch.Tensor,
+    k: torch.Tensor,
+    head_size: int,
+    interleaved: bool,
+) -> None:
+    from sglang.jit_kernel.rotary_embedding import rotary_embedding_cos_sin
+
+    rotary_embedding_cos_sin(cos, sin, q, k, head_size, interleaved)
+
+
+BASE_LINE_VALS = ["jit_cos_sin", "aot_pos"]
+BASE_LINE_NAMES = ["SGL JIT (cos/sin)", "SGL AOT (positions)"]
+BASE_STYLES = [("blue", "-"), ("orange", "--")]
+
+LINE_VALS = list(BASE_LINE_VALS)
+LINE_NAMES = list(BASE_LINE_NAMES)
+STYLES = list(BASE_STYLES)
+
+if HAS_VLLM:
+    LINE_VALS.append("vllm_pos")
+    LINE_NAMES.append("vLLM (positions)")
+    STYLES.append(("green", "-."))
+
+if ONLY_PROVIDER is not None:
+    if ONLY_PROVIDER not in LINE_VALS:
+        raise ValueError(
+            f"Unknown SGL_BENCH_PROVIDER={ONLY_PROVIDER}. Allowed: {LINE_VALS}"
+        )
+    idx = LINE_VALS.index(ONLY_PROVIDER)
+    LINE_VALS = [ONLY_PROVIDER]
+    LINE_NAMES = [LINE_NAMES[idx]]
+    STYLES = [STYLES[idx]]
+
+configs = list(
+    itertools.product(BS_RANGE, SEQ_RANGE, HEAD_SIZE_RANGE, INTERLEAVED_RANGE)
+)
+_SANITY_DONE = False
+
+
+def _assert_close_gpu(
+    actual: torch.Tensor,
+    expected: torch.Tensor,
+    *,
+    atol: float,
+    rtol: float,
+    name: str,
+) -> None:
+    if actual.shape != expected.shape:
+        raise AssertionError(
+            f"{name}: shape mismatch {tuple(actual.shape)} vs {tuple(expected.shape)}"
+        )
+    if actual.dtype != expected.dtype:
+        raise AssertionError(
+            f"{name}: dtype mismatch {actual.dtype} vs {expected.dtype}"
+        )
+    if actual.device != expected.device:
+        raise AssertionError(
+            f"{name}: device mismatch {actual.device} vs {expected.device}"
+        )
+
+    diff = (actual - expected).abs()
+    max_abs = float(diff.max().item())
+    denom = atol + rtol * expected.abs()
+    max_rel = float((diff / denom).max().item())
+    if max_abs > atol and max_rel > 1.0:
+        raise AssertionError(
+            f"{name}: not close (atol={atol} rtol={rtol}) max_abs={max_abs:.6g} max_rel={max_rel:.6g}"
+        )
+
+
+@lru_cache(maxsize=None)
+def _get_vllm_rope(
+    head_size: int, rotary_dim: int, interleaved: bool, dtype: torch.dtype
+):
+    if not HAS_VLLM or vLLMRotaryEmbedding is None:
+        raise RuntimeError("vLLM not available")
+    return vLLMRotaryEmbedding(
+        head_size=head_size,
+        rotary_dim=rotary_dim,
+        max_position_embeddings=MAX_SEQ_LEN,
+        base=10000,
+        is_neox_style=interleaved,
+        dtype=dtype,
+    ).cuda()
+
+
+@triton.testing.perf_report(
+    triton.testing.Benchmark(
+        x_names=["batch_size", "seq_len", "head_size", "interleaved"],
+        x_vals=configs,
+        line_arg="provider",
+        line_vals=LINE_VALS,
+        line_names=LINE_NAMES,
+        styles=STYLES,
+        ylabel="us",
+        plot_name="rotary-embedding-performance",
+        args={},
+    )
+)
+def benchmark(
+    batch_size: int,
+    seq_len: int,
+    head_size: int,
+    interleaved: bool,
+    provider: str,
+) -> Tuple[float, float, float]:
+    if DEVICE == "cuda" and not torch.cuda.is_available():
+        raise RuntimeError("CUDA not available")
+
+    num_tokens = batch_size * seq_len
+    rotary_dim = head_size
+
+    cos_cache_half, sin_cache_half = _get_cos_sin_cache_half_cuda(rotary_dim)
+    cos_sin_cache_aot = torch.cat([cos_cache_half, sin_cache_half], dim=-1).contiguous()
+
+    positions = torch.arange(seq_len, device=DEVICE, dtype=torch.int64).repeat(
+        batch_size
+    )
+    cos_half = cos_cache_half[positions].contiguous()
+    sin_half = sin_cache_half[positions].contiguous()
+
+    if interleaved:
+        cos, sin = cos_half, sin_half
+    else:
+        cos = torch.cat([cos_half, cos_half], dim=-1).contiguous()
+        sin = torch.cat([sin_half, sin_half], dim=-1).contiguous()
+
+    q = torch.randn(num_tokens, NUM_Q_HEADS, head_size, device=DEVICE, dtype=DTYPE)
+    k = torch.randn(num_tokens, NUM_KV_HEADS, head_size, device=DEVICE, dtype=DTYPE)
+
+    global _SANITY_DONE
+    if (
+        (not _SANITY_DONE)
+        and provider == "aot_pos"
+        and batch_size == BS_RANGE[0]
+        and seq_len == SEQ_RANGE[0]
+        and head_size == HEAD_SIZE_RANGE[0]
+        and interleaved == INTERLEAVED_RANGE[0]
+    ):
+        q_ref = q.clone()
+        k_ref = k.clone()
+        torch_impl_rotary_fp32(cos, sin, q_ref, k_ref, head_size, interleaved)
+        q_out = q.clone()
+        k_out = k.clone()
+        sglang_aot_rotary_positions(
+            positions, q_out, k_out, head_size, interleaved, cos_sin_cache_aot
+        )
+        ref_atol = 2e-2 if DTYPE == torch.bfloat16 else 2e-3
+        ref_rtol = 2e-2 if DTYPE == torch.bfloat16 else 2e-3
+        _assert_close_gpu(q_out, q_ref, atol=ref_atol, rtol=ref_rtol, name="AOT(q)")
+        _assert_close_gpu(k_out, k_ref, atol=ref_atol, rtol=ref_rtol, name="AOT(k)")
+        _SANITY_DONE = True
+
+    FN_MAP = {
+        "aot_pos": lambda: sglang_aot_rotary_positions(
+            positions, q, k, head_size, interleaved, cos_sin_cache_aot
+        ),
+        "jit_cos_sin": lambda: sglang_jit_rotary_cos_sin(
+            cos, sin, q, k, head_size, interleaved
+        ),
+        "torch_fp32": lambda: torch_impl_rotary_fp32(
+            cos, sin, q, k, head_size, interleaved
+        ),
+    }
+    if HAS_VLLM:
+        vllm_rope = _get_vllm_rope(head_size, rotary_dim, interleaved, DTYPE)
+        FN_MAP["vllm_pos"] = lambda: vllm_rope.forward_cuda(positions, q, k)
+
+    fn = FN_MAP[provider]
+    quantiles = [0.5, 0.2, 0.8]
+    ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(fn, quantiles=quantiles)  # type: ignore
+    return 1000 * ms, 1000 * max_ms, 1000 * min_ms
+
+
+SPEEDUP_LINE_VALS = [p for p in LINE_VALS if p != "jit_cos_sin"]
+SPEEDUP_LINE_NAMES = [LINE_NAMES[LINE_VALS.index(p)] for p in SPEEDUP_LINE_VALS]
+SPEEDUP_STYLES = [STYLES[LINE_VALS.index(p)] for p in SPEEDUP_LINE_VALS]
+
+
+@triton.testing.perf_report(
+    triton.testing.Benchmark(
+        x_names=["batch_size", "seq_len", "head_size", "interleaved"],
+        x_vals=configs,
+        line_arg="provider",
+        line_vals=SPEEDUP_LINE_VALS,
+        line_names=SPEEDUP_LINE_NAMES,
+        styles=SPEEDUP_STYLES,
+        ylabel="speedup (x) vs SGL JIT (cos/sin)",
+        plot_name="rotary-embedding-speedup-vs-jit",
+        args={},
+    )
+)
+def benchmark_speedup_vs_jit(
+    batch_size: int,
+    seq_len: int,
+    head_size: int,
+    interleaved: bool,
+    provider: str,
+) -> Tuple[float, float, float]:
+    if provider == "jit_cos_sin":
+        return 1.0, 1.0, 1.0
+    if DEVICE == "cuda" and not torch.cuda.is_available():
+        raise RuntimeError("CUDA not available")
+
+    num_tokens = batch_size * seq_len
+    rotary_dim = head_size
+
+    cos_cache_half, sin_cache_half = _get_cos_sin_cache_half_cuda(rotary_dim)
+    cos_sin_cache_aot = torch.cat([cos_cache_half, sin_cache_half], dim=-1).contiguous()
+
+    positions = torch.arange(seq_len, device=DEVICE, dtype=torch.int64).repeat(
+        batch_size
+    )
+    cos_half = cos_cache_half[positions].contiguous()
+    sin_half = sin_cache_half[positions].contiguous()
+
+    if interleaved:
+        cos, sin = cos_half, sin_half
+    else:
+        cos = torch.cat([cos_half, cos_half], dim=-1).contiguous()
+        sin = torch.cat([sin_half, sin_half], dim=-1).contiguous()
+
+    q_base = torch.randn(num_tokens, NUM_Q_HEADS, head_size, device=DEVICE, dtype=DTYPE)
+    k_base = torch.randn(
+        num_tokens, NUM_KV_HEADS, head_size, device=DEVICE, dtype=DTYPE
+    )
+
+    q_jit = q_base.clone()
+    k_jit = k_base.clone()
+    q_p = q_base.clone()
+    k_p = k_base.clone()
+
+    FN_MAP = {
+        "jit_cos_sin": lambda: sglang_jit_rotary_cos_sin(
+            cos, sin, q_jit, k_jit, head_size, interleaved
+        ),
+        "aot_pos": lambda: sglang_aot_rotary_positions(
+            positions, q_p, k_p, head_size, interleaved, cos_sin_cache_aot
+        ),
+        "torch_fp32": lambda: torch_impl_rotary_fp32(
+            cos, sin, q_p, k_p, head_size, interleaved
+        ),
+    }
+    if HAS_VLLM:
+        vllm_rope = _get_vllm_rope(head_size, rotary_dim, interleaved, DTYPE)
+        FN_MAP["vllm_pos"] = lambda: vllm_rope.forward_cuda(positions, q_p, k_p)
+
+    quantiles = [0.5, 0.2, 0.8]
+    jit_ms, jit_min_ms, jit_max_ms = triton.testing.do_bench_cudagraph(
+        FN_MAP["jit_cos_sin"], quantiles=quantiles
+    )  # type: ignore
+    p_ms, p_min_ms, p_max_ms = triton.testing.do_bench_cudagraph(
+        FN_MAP[provider], quantiles=quantiles
+    )  # type: ignore
+
+    speed_med = float(jit_ms / p_ms)
+    speed_max = float(jit_max_ms / p_min_ms)
+    speed_min = float(jit_min_ms / p_max_ms)
+    return speed_med, speed_max, speed_min
+
+
+if __name__ == "__main__":
+    benchmark.run(print_data=True)
+    if (
+        SHOW_SPEEDUP
+        and ("jit_cos_sin" in LINE_VALS)
+        and (len(LINE_VALS) > 1)
+        and (ONLY_PROVIDER is None)
+    ):
+        benchmark_speedup_vs_jit.run(print_data=True)
diff --git a/python/sglang/jit_kernel/csrc/rotary_embedding_cos_sin.cuh b/python/sglang/jit_kernel/csrc/rotary_embedding_cos_sin.cuh
new file mode 100644
index 000000000000..0e3e1df7383a
--- /dev/null
+++ b/python/sglang/jit_kernel/csrc/rotary_embedding_cos_sin.cuh
@@ -0,0 +1,1106 @@
+#include <sgl_kernel/tensor.h>
+#include <sgl_kernel/utils.cuh>
+#include <sgl_kernel/utils.h>
+
+#include <cuda_bf16.h>
+#include <cuda_fp16.h>
+#include <cuda_runtime.h>
+#include <dlpack/dlpack.h>
+#include <tvm/ffi/container/tensor.h>
+
+#include <algorithm>
+#include <cstdint>
+#include <type_traits>
+
+namespace {
+
+// Convert scalar types to float for computation.
+template <typename T>
+__device__ __forceinline__ float to_float(T v) {
+  return static_cast<float>(v);
+}
+template <>
+__device__ __forceinline__ float to_float<half>(half v) {
+  return __half2float(v);
+}
+template <>
+__device__ __forceinline__ float to_float<nv_bfloat16>(nv_bfloat16 v) {
+  return __bfloat162float(v);
+}
+
+// Convert float back to scalar types (round-to-nearest for fp16/bf16).
+template <typename T>
+__device__ __forceinline__ T from_float(float v) {
+  return static_cast<T>(v);
+}
+template <>
+__device__ __forceinline__ half from_float<half>(float v) {
+  return __float2half_rn(v);
+}
+template <>
+__device__ __forceinline__ nv_bfloat16 from_float<nv_bfloat16>(float v) {
+  return __float2bfloat16_rn(v);
+}
+
+// Vectorize loads/stores in 16 bytes (float4) regardless of scalar type.
+template <typename scalar_t, bool interleaved>
+struct RotaryVecData;
+
+template <typename scalar_t>
+struct RotaryVecData<scalar_t, true> {
+  float4 vec;
+  float2 cos_vec;
+  float2 sin_vec;
+};
+
+template <typename scalar_t>
+struct RotaryVecData<scalar_t, false> {
+  float4 vec_x;
+  float4 vec_y;
+  float4 cos_x;
+  float4 sin_x;
+  float4 cos_y;
+  float4 sin_y;
+};
+
+// Vector load helper:
+//   - load a 16B tile from q/k (either aligned float4 or scalar gather)
+//   - always vector-load cos/sin (aligned by launch-time checks)
+//
+// rot_offset is the "pair index":
+//   - interleaved:     pair i -> q[2*i], q[2*i+1]
+//   - non-interleaved: pair i -> q[i], q[i + embed_dim]
+// embed_dim is the number of pairs per head.
+template <typename scalar_t, bool interleaved, bool aligned_qk>
+__device__ __forceinline__ RotaryVecData<scalar_t, interleaved> load_rotary_vec(
+    const scalar_t* __restrict__ arr,
+    const scalar_t* __restrict__ cos_ptr,
+    const scalar_t* __restrict__ sin_ptr,
+    int rot_offset,
+    int embed_dim) {
+  RotaryVecData<scalar_t, interleaved> data;
+
+  constexpr int kVecBytes = 16;
+  constexpr int kElePerVec = kVecBytes / sizeof(scalar_t);
+
+  if constexpr (interleaved) {
+    const int base = rot_offset * 2;
+
+    if constexpr (aligned_qk) {
+      data.vec = *reinterpret_cast<const float4*>(arr + base);
+    } else {
+      union VecU {
+        float4 v;
+        scalar_t e[kElePerVec];
+      } tmp;
+#pragma unroll
+      for (int i = 0; i < kElePerVec; ++i)
+        tmp.e[i] = arr[base + i];
+      data.vec = tmp.v;
+    }
+
+    data.cos_vec = *reinterpret_cast<const float2*>(cos_ptr + rot_offset);
+    data.sin_vec = *reinterpret_cast<const float2*>(sin_ptr + rot_offset);
+  } else {
+    const int bx = rot_offset;
+    const int by = rot_offset + embed_dim;
+
+    if constexpr (aligned_qk) {
+      data.vec_x = *reinterpret_cast<const float4*>(arr + bx);
+      data.vec_y = *reinterpret_cast<const float4*>(arr + by);
+    } else {
+      union VecU {
+        float4 v;
+        scalar_t e[kElePerVec];
+      } tx, ty;
+#pragma unroll
+      for (int i = 0; i < kElePerVec; ++i) {
+        tx.e[i] = arr[bx + i];
+        ty.e[i] = arr[by + i];
+      }
+      data.vec_x = tx.v;
+      data.vec_y = ty.v;
+    }
+
+    data.cos_x = *reinterpret_cast<const float4*>(cos_ptr + bx);
+    data.sin_x = *reinterpret_cast<const float4*>(sin_ptr + bx);
+    data.cos_y = *reinterpret_cast<const float4*>(cos_ptr + by);
+    data.sin_y = *reinterpret_cast<const float4*>(sin_ptr + by);
+  }
+
+  return data;
+}
+
+// Vector compute + store:
+// Apply RoPE to the loaded 16B tile and write back to q/k.
+//
+// Interleaved math (per pair):
+//   x' = x*cos - y*sin
+//   y' = y*cos + x*sin
+//
+// Non-interleaved math (general form, allows distinct cos/sin for x and y halves):
+//   x' = x*cos_x - y*sin_x
+//   y' = y*cos_y + x*sin_y
+template <typename scalar_t, bool interleaved, bool aligned_qk>
+__device__ __forceinline__ void compute_store_rotary_vec(
+    scalar_t* __restrict__ arr, const RotaryVecData<scalar_t, interleaved>& data, int rot_offset, int embed_dim) {
+  constexpr int kVecBytes = 16;
+  constexpr int kElePerVec = kVecBytes / sizeof(scalar_t);
+
+  if constexpr (interleaved) {
+    union VecU {
+      float4 v;
+      scalar_t e[kElePerVec];
+    } v;
+    v.v = data.vec;
+
+    union CSU {
+      float2 v;
+      scalar_t e[kElePerVec / 2];
+    } c, s;
+    c.v = data.cos_vec;
+    s.v = data.sin_vec;
+
+#pragma unroll
+    for (int i = 0; i < kElePerVec; i += 2) {
+      const int idx = i / 2;
+      const float cos_val = to_float(c.e[idx]);
+      const float sin_val = to_float(s.e[idx]);
+      const float x = to_float(v.e[i]);
+      const float y = to_float(v.e[i + 1]);
+      v.e[i] = from_float<scalar_t>(x * cos_val - y * sin_val);
+      v.e[i + 1] = from_float<scalar_t>(y * cos_val + x * sin_val);
+    }
+
+    const int base = rot_offset * 2;
+    if constexpr (aligned_qk) {
+      *reinterpret_cast<float4*>(arr + base) = v.v;
+    } else {
+#pragma unroll
+      for (int i = 0; i < kElePerVec; ++i)
+        arr[base + i] = v.e[i];
+    }
+  } else {
+    union VecU {
+      float4 v;
+      scalar_t e[kElePerVec];
+    } vx, vy;
+    vx.v = data.vec_x;
+    vy.v = data.vec_y;
+
+    union CSU {
+      float4 v;
+      scalar_t e[kElePerVec];
+    } cx, sx, cy, sy;
+    cx.v = data.cos_x;
+    sx.v = data.sin_x;
+    cy.v = data.cos_y;
+    sy.v = data.sin_y;
+
+#pragma unroll
+    for (int i = 0; i < kElePerVec; ++i) {
+      const float cos_x = to_float(cx.e[i]);
+      const float sin_x = to_float(sx.e[i]);
+      const float cos_y = to_float(cy.e[i]);
+      const float sin_y = to_float(sy.e[i]);
+
+      const float x = to_float(vx.e[i]);
+      const float y = to_float(vy.e[i]);
+
+      vx.e[i] = from_float<scalar_t>(x * cos_x - y * sin_x);
+      vy.e[i] = from_float<scalar_t>(y * cos_y + x * sin_y);
+    }
+
+    const int bx = rot_offset;
+    const int by = rot_offset + embed_dim;
+    if constexpr (aligned_qk) {
+      *reinterpret_cast<float4*>(arr + bx) = vx.v;
+      *reinterpret_cast<float4*>(arr + by) = vy.v;
+    } else {
+#pragma unroll
+      for (int i = 0; i < kElePerVec; ++i) {
+        arr[bx + i] = vx.e[i];
+        arr[by + i] = vy.e[i];
+      }
+    }
+  }
+}
+
+// Scalar fallback: Apply RoPE for exactly one (x,y) pair per iteration.
+template <typename scalar_t, bool interleaved>
+inline __device__ void apply_token_rotary_embedding(
+    scalar_t* __restrict__ arr,            // [head_size]
+    const scalar_t* __restrict__ cos_ptr,  // [rot_dim]
+    const scalar_t* __restrict__ sin_ptr,  // [rot_dim]
+    int rot_offset,
+    int embed_dim) {
+  if constexpr (interleaved) {
+    const int x_index = 2 * rot_offset;
+    const int y_index = x_index + 1;
+
+    const float cos_val = to_float(cos_ptr[rot_offset]);
+    const float sin_val = to_float(sin_ptr[rot_offset]);
+    const float x = to_float(arr[x_index]);
+    const float y = to_float(arr[y_index]);
+
+    arr[x_index] = from_float<scalar_t>(x * cos_val - y * sin_val);
+    arr[y_index] = from_float<scalar_t>(y * cos_val + x * sin_val);
+  } else {
+    const int x_index = rot_offset;
+    const int y_index = rot_offset + embed_dim;
+
+    const float cos_val_x = to_float(cos_ptr[rot_offset]);
+    const float sin_val_x = to_float(sin_ptr[rot_offset]);
+    const float cos_val_y = to_float(cos_ptr[rot_offset + embed_dim]);
+    const float sin_val_y = to_float(sin_ptr[rot_offset + embed_dim]);
+
+    const float x = to_float(arr[x_index]);
+    const float y = to_float(arr[y_index]);
+
+    arr[x_index] = from_float<scalar_t>(x * cos_val_x - y * sin_val_x);
+    arr[y_index] = from_float<scalar_t>(y * cos_val_y + x * sin_val_y);
+  }
+}
+
+// 2D-grid kernel:
+//   blockIdx.x -> token index
+//   blockIdx.y -> "sub-block" index within the token (tile along pairs dimension)
+//
+// For very small T (few tokens) but large per-token work, using multiple blocks
+// per token can improve occupancy/throughput compared to one-block-per-token.
+template <typename scalar_t, bool interleaved, bool vectorized, bool aligned_qk, int ROT_EMBED_DIM>
+__global__ void rotary_embedding_kernel_2d(
+    const scalar_t* __restrict__ cos_data,  // [num_tokens, rot_dim_arg]
+    const scalar_t* __restrict__ sin_data,  // [num_tokens, rot_dim_arg]
+    scalar_t* __restrict__ query_total,     // [num_tokens, num_heads, head_size] contiguous
+    scalar_t* __restrict__ key_total,       // [num_tokens, num_kv_heads, head_size] contiguous or nullptr
+    const int rot_dim_arg,
+    const int embed_dim_for_rotation_arg,
+    const int64_t query_token_stride,
+    const int64_t key_token_stride,
+    const int64_t head_stride_query,
+    const int64_t head_stride_key,
+    const int num_heads,
+    const int num_kv_heads,
+    const int head_size_arg,
+    const int blocks_per_token) {
+  const int token_idx = blockIdx.x;
+  if (token_idx >= gridDim.x) return;
+
+  const scalar_t* current_token_cos_ptr = cos_data + token_idx * rot_dim_arg;
+  const scalar_t* current_token_sin_ptr = sin_data + token_idx * rot_dim_arg;
+
+  scalar_t* query_for_token = query_total + token_idx * (int)query_token_stride;
+  scalar_t* key_for_token = (key_total != nullptr) ? (key_total + token_idx * (int)key_token_stride) : nullptr;
+
+  const int local_block_idx = blockIdx.y;
+  const int embed_dim_for_rotation = (ROT_EMBED_DIM > 0) ? ROT_EMBED_DIM : embed_dim_for_rotation_arg;
+
+  if constexpr (vectorized) {
+    constexpr int kVecBytes = 16;
+    constexpr int kElePerVec = kVecBytes / sizeof(scalar_t);
+    constexpr int pairs_per_step = interleaved ? (kElePerVec / 2) : kElePerVec;
+
+    const int pair_stride = blockDim.x * blocks_per_token * pairs_per_step;
+    int i = (local_block_idx * blockDim.x + threadIdx.x) * pairs_per_step;
+    const int nq_pairs = num_heads * embed_dim_for_rotation;
+
+    RotaryVecData<scalar_t, interleaved> curr_data;
+    if (i < nq_pairs) {
+      const int head_idx = i / embed_dim_for_rotation;
+      const int rot_offset = i % embed_dim_for_rotation;
+      scalar_t* ptr = query_for_token + head_idx * (int)head_stride_query;
+      curr_data = load_rotary_vec<scalar_t, interleaved, aligned_qk>(
+          ptr, current_token_cos_ptr, current_token_sin_ptr, rot_offset, embed_dim_for_rotation);
+    }
+
+    int next_i = i + pair_stride;
+    for (; i < nq_pairs; i += pair_stride, next_i += pair_stride) {
+      RotaryVecData<scalar_t, interleaved> next_data;
+      const bool active_next = (next_i < nq_pairs);
+
+      if (active_next) {
+        const int head_idx_next = next_i / embed_dim_for_rotation;
+        const int rot_offset_next = next_i % embed_dim_for_rotation;
+        scalar_t* ptr_next = query_for_token + head_idx_next * (int)head_stride_query;
+        next_data = load_rotary_vec<scalar_t, interleaved, aligned_qk>(
+            ptr_next, current_token_cos_ptr, current_token_sin_ptr, rot_offset_next, embed_dim_for_rotation);
+      }
+
+      const int head_idx = i / embed_dim_for_rotation;
+      const int rot_offset = i % embed_dim_for_rotation;
+      scalar_t* ptr = query_for_token + head_idx * (int)head_stride_query;
+      compute_store_rotary_vec<scalar_t, interleaved, aligned_qk>(ptr, curr_data, rot_offset, embed_dim_for_rotation);
+
+      if (active_next) curr_data = next_data;
+    }
+
+    if (key_for_token != nullptr) {
+      const int nk_pairs = num_kv_heads * embed_dim_for_rotation;
+      int k_i = (local_block_idx * blockDim.x + threadIdx.x) * pairs_per_step;
+
+      RotaryVecData<scalar_t, interleaved> curr_data_k;
+      if (k_i < nk_pairs) {
+        const int head_idx = k_i / embed_dim_for_rotation;
+        const int rot_offset = k_i % embed_dim_for_rotation;
+        scalar_t* ptr = key_for_token + head_idx * (int)head_stride_key;
+        curr_data_k = load_rotary_vec<scalar_t, interleaved, aligned_qk>(
+            ptr, current_token_cos_ptr, current_token_sin_ptr, rot_offset, embed_dim_for_rotation);
+      }
+
+      int next_k_i = k_i + pair_stride;
+      for (; k_i < nk_pairs; k_i += pair_stride, next_k_i += pair_stride) {
+        RotaryVecData<scalar_t, interleaved> next_data_k;
+        const bool active_next = (next_k_i < nk_pairs);
+        if (active_next) {
+          const int head_idx_next = next_k_i / embed_dim_for_rotation;
+          const int rot_offset_next = next_k_i % embed_dim_for_rotation;
+          scalar_t* ptr_next = key_for_token + head_idx_next * (int)head_stride_key;
+          next_data_k = load_rotary_vec<scalar_t, interleaved, aligned_qk>(
+              ptr_next, current_token_cos_ptr, current_token_sin_ptr, rot_offset_next, embed_dim_for_rotation);
+        }
+
+        const int head_idx = k_i / embed_dim_for_rotation;
+        const int rot_offset = k_i % embed_dim_for_rotation;
+        scalar_t* ptr = key_for_token + head_idx * (int)head_stride_key;
+        compute_store_rotary_vec<scalar_t, interleaved, aligned_qk>(
+            ptr, curr_data_k, rot_offset, embed_dim_for_rotation);
+
+        if (active_next) curr_data_k = next_data_k;
+      }
+    }
+  } else {
+    // Fallback to scalar implementation
+    const int pair_stride = blockDim.x * blocks_per_token;
+    const int thread_pair_offset = local_block_idx * blockDim.x + threadIdx.x;
+
+    const int nq_pairs = num_heads * embed_dim_for_rotation;
+    for (int i = thread_pair_offset; i < nq_pairs; i += pair_stride) {
+      const int head_idx = i / embed_dim_for_rotation;
+      const int rot_offset = i % embed_dim_for_rotation;
+      scalar_t* query_for_token_head = query_for_token + head_idx * (int)head_stride_query;
+      apply_token_rotary_embedding<scalar_t, interleaved>(
+          query_for_token_head, current_token_cos_ptr, current_token_sin_ptr, rot_offset, embed_dim_for_rotation);
+    }
+
+    if (key_for_token != nullptr) {
+      const int nk_pairs = num_kv_heads * embed_dim_for_rotation;
+      for (int i = thread_pair_offset; i < nk_pairs; i += pair_stride) {
+        const int head_idx = i / embed_dim_for_rotation;
+        const int rot_offset = i % embed_dim_for_rotation;
+        scalar_t* key_for_token_head = key_for_token + head_idx * (int)head_stride_key;
+        apply_token_rotary_embedding<scalar_t, interleaved>(
+            key_for_token_head, current_token_cos_ptr, current_token_sin_ptr, rot_offset, embed_dim_for_rotation);
+      }
+    }
+  }
+}
+
+// 1D-grid kernel:
+//   blockIdx.x -> token index
+//   exactly one block per token
+//
+// The default path for most sizes.
+template <typename scalar_t, bool interleaved, bool vectorized, bool aligned_qk, int ROT_EMBED_DIM>
+__launch_bounds__(512) __global__ void rotary_embedding_kernel_1d(
+    const scalar_t* __restrict__ cos_data,  // [num_tokens, rot_dim_arg]
+    const scalar_t* __restrict__ sin_data,  // [num_tokens, rot_dim_arg]
+    scalar_t* __restrict__ query_total,
+    scalar_t* __restrict__ key_total,
+    const int rot_dim_arg,
+    const int embed_dim_for_rotation_arg,
+    const int64_t query_token_stride,
+    const int64_t key_token_stride,
+    const int64_t head_stride_query,
+    const int64_t head_stride_key,
+    const int num_heads,
+    const int num_kv_heads,
+    const int head_size_arg) {
+  const int token_idx = blockIdx.x;
+  if (token_idx >= gridDim.x) return;
+
+  const scalar_t* current_token_cos_ptr = cos_data + token_idx * rot_dim_arg;
+  const scalar_t* current_token_sin_ptr = sin_data + token_idx * rot_dim_arg;
+
+  scalar_t* query_for_token = query_total + token_idx * (int)query_token_stride;
+  scalar_t* key_for_token = (key_total != nullptr) ? (key_total + token_idx * (int)key_token_stride) : nullptr;
+
+  const int embed_dim_for_rotation = (ROT_EMBED_DIM > 0) ? ROT_EMBED_DIM : embed_dim_for_rotation_arg;
+
+  if constexpr (vectorized) {
+    constexpr int kVecBytes = 16;
+    constexpr int kElePerVec = kVecBytes / sizeof(scalar_t);
+    constexpr int pairs_per_step = interleaved ? (kElePerVec / 2) : kElePerVec;
+
+    const int nq_pairs = num_heads * embed_dim_for_rotation;
+    const int stride = blockDim.x * pairs_per_step;
+    int i = threadIdx.x * pairs_per_step;
+
+    RotaryVecData<scalar_t, interleaved> curr_data;
+    if (i < nq_pairs) {
+      const int head_idx = i / embed_dim_for_rotation;
+      const int rot_offset = i % embed_dim_for_rotation;
+      scalar_t* ptr = query_for_token + head_idx * (int)head_stride_query;
+      curr_data = load_rotary_vec<scalar_t, interleaved, aligned_qk>(
+          ptr, current_token_cos_ptr, current_token_sin_ptr, rot_offset, embed_dim_for_rotation);
+    }
+
+    int next_i = i + stride;
+    for (; i < nq_pairs; i += stride, next_i += stride) {
+      RotaryVecData<scalar_t, interleaved> next_data;
+      if (next_i < nq_pairs) {
+        const int head_idx = next_i / embed_dim_for_rotation;
+        const int rot_offset = next_i % embed_dim_for_rotation;
+        scalar_t* ptr = query_for_token + head_idx * (int)head_stride_query;
+        next_data = load_rotary_vec<scalar_t, interleaved, aligned_qk>(
+            ptr, current_token_cos_ptr, current_token_sin_ptr, rot_offset, embed_dim_for_rotation);
+      }
+
+      if (i < nq_pairs) {
+        const int head_idx = i / embed_dim_for_rotation;
+        const int rot_offset = i % embed_dim_for_rotation;
+        scalar_t* ptr = query_for_token + head_idx * (int)head_stride_query;
+        compute_store_rotary_vec<scalar_t, interleaved, aligned_qk>(ptr, curr_data, rot_offset, embed_dim_for_rotation);
+      }
+      curr_data = next_data;
+    }
+
+    if (key_for_token != nullptr) {
+      const int nk_pairs = num_kv_heads * embed_dim_for_rotation;
+      int k_i = threadIdx.x * pairs_per_step;
+      RotaryVecData<scalar_t, interleaved> curr_data_k;
+
+      if (k_i < nk_pairs) {
+        const int head_idx = k_i / embed_dim_for_rotation;
+        const int rot_offset = k_i % embed_dim_for_rotation;
+        scalar_t* ptr = key_for_token + head_idx * (int)head_stride_key;
+        curr_data_k = load_rotary_vec<scalar_t, interleaved, aligned_qk>(
+            ptr, current_token_cos_ptr, current_token_sin_ptr, rot_offset, embed_dim_for_rotation);
+      }
+
+      int next_k_i = k_i + stride;
+      for (; k_i < nk_pairs; k_i += stride, next_k_i += stride) {
+        RotaryVecData<scalar_t, interleaved> next_data_k;
+        if (next_k_i < nk_pairs) {
+          const int head_idx = next_k_i / embed_dim_for_rotation;
+          const int rot_offset = next_k_i % embed_dim_for_rotation;
+          scalar_t* ptr = key_for_token + head_idx * (int)head_stride_key;
+          next_data_k = load_rotary_vec<scalar_t, interleaved, aligned_qk>(
+              ptr, current_token_cos_ptr, current_token_sin_ptr, rot_offset, embed_dim_for_rotation);
+        }
+
+        if (k_i < nk_pairs) {
+          const int head_idx = k_i / embed_dim_for_rotation;
+          const int rot_offset = k_i % embed_dim_for_rotation;
+          scalar_t* ptr = key_for_token + head_idx * (int)head_stride_key;
+          compute_store_rotary_vec<scalar_t, interleaved, aligned_qk>(
+              ptr, curr_data_k, rot_offset, embed_dim_for_rotation);
+        }
+        curr_data_k = next_data_k;
+      }
+    }
+  } else {
+    const int nq_pairs = num_heads * embed_dim_for_rotation;
+    for (int i = threadIdx.x; i < nq_pairs; i += blockDim.x) {
+      const int head_idx = i / embed_dim_for_rotation;
+      const int rot_offset = i % embed_dim_for_rotation;
+      scalar_t* query_for_token_head = query_for_token + head_idx * (int)head_stride_query;
+      apply_token_rotary_embedding<scalar_t, interleaved>(
+          query_for_token_head, current_token_cos_ptr, current_token_sin_ptr, rot_offset, embed_dim_for_rotation);
+    }
+
+    if (key_for_token != nullptr) {
+      const int nk_pairs = num_kv_heads * embed_dim_for_rotation;
+      for (int i = threadIdx.x; i < nk_pairs; i += blockDim.x) {
+        const int head_idx = i / embed_dim_for_rotation;
+        const int rot_offset = i % embed_dim_for_rotation;
+        scalar_t* key_for_token_head = key_for_token + head_idx * (int)head_stride_key;
+        apply_token_rotary_embedding<scalar_t, interleaved>(
+            key_for_token_head, current_token_cos_ptr, current_token_sin_ptr, rot_offset, embed_dim_for_rotation);
+      }
+    }
+  }
+}
+
+// Kernel variant dispatcher:
+// Select one of:
+//   - grid shape: 2D (multi-block per token) vs 1D (one-block per token)
+//   - layout: interleaved vs non-interleaved
+//   - compute: vectorized 16B tiles vs scalar
+//   - memory: q/k aligned 16B load/store vs scalar gather/scatter
+template <int ROT, typename scalar_t>
+__forceinline__ void dispatch_rotary_launch(
+    bool use_grid_2d,
+    dim3 grid2d,
+    dim3 grid1d,
+    dim3 block,
+    const decltype(host::LaunchKernel::resolve_device(std::declval<DLDevice>())) stream,
+    bool interleaved,
+    bool use_vec,
+    bool qk_aligned16,
+    const scalar_t* cos_ptr,
+    const scalar_t* sin_ptr,
+    scalar_t* q_ptr,
+    scalar_t* k_ptr,
+    int rot_dim_from_cache,
+    int embed_dim_for_rotation,
+    int64_t query_token_stride,
+    int64_t key_token_stride,
+    int64_t head_stride_query,
+    int64_t head_stride_key,
+    int num_heads,
+    int num_kv_heads,
+    int head_size,
+    int blocks_per_token) {
+  // 2D grid path
+  if (use_grid_2d) {
+    if (interleaved) {
+      if (use_vec) {
+        if (qk_aligned16) {
+          host::LaunchKernel(grid2d, block, stream)(
+              rotary_embedding_kernel_2d<scalar_t, true, true, true, ROT>,
+              cos_ptr,
+              sin_ptr,
+              q_ptr,
+              k_ptr,
+              rot_dim_from_cache,
+              embed_dim_for_rotation,
+              query_token_stride,
+              key_token_stride,
+              head_stride_query,
+              head_stride_key,
+              num_heads,
+              num_kv_heads,
+              head_size,
+              blocks_per_token);
+        } else {
+          host::LaunchKernel(grid2d, block, stream)(
+              rotary_embedding_kernel_2d<scalar_t, true, true, false, ROT>,
+              cos_ptr,
+              sin_ptr,
+              q_ptr,
+              k_ptr,
+              rot_dim_from_cache,
+              embed_dim_for_rotation,
+              query_token_stride,
+              key_token_stride,
+              head_stride_query,
+              head_stride_key,
+              num_heads,
+              num_kv_heads,
+              head_size,
+              blocks_per_token);
+        }
+      } else {
+        host::LaunchKernel(grid2d, block, stream)(
+            rotary_embedding_kernel_2d<scalar_t, true, false, true, ROT>,
+            cos_ptr,
+            sin_ptr,
+            q_ptr,
+            k_ptr,
+            rot_dim_from_cache,
+            embed_dim_for_rotation,
+            query_token_stride,
+            key_token_stride,
+            head_stride_query,
+            head_stride_key,
+            num_heads,
+            num_kv_heads,
+            head_size,
+            blocks_per_token);
+      }
+    } else {
+      if (use_vec) {
+        if (qk_aligned16) {
+          host::LaunchKernel(grid2d, block, stream)(
+              rotary_embedding_kernel_2d<scalar_t, false, true, true, ROT>,
+              cos_ptr,
+              sin_ptr,
+              q_ptr,
+              k_ptr,
+              rot_dim_from_cache,
+              embed_dim_for_rotation,
+              query_token_stride,
+              key_token_stride,
+              head_stride_query,
+              head_stride_key,
+              num_heads,
+              num_kv_heads,
+              head_size,
+              blocks_per_token);
+        } else {
+          host::LaunchKernel(grid2d, block, stream)(
+              rotary_embedding_kernel_2d<scalar_t, false, true, false, ROT>,
+              cos_ptr,
+              sin_ptr,
+              q_ptr,
+              k_ptr,
+              rot_dim_from_cache,
+              embed_dim_for_rotation,
+              query_token_stride,
+              key_token_stride,
+              head_stride_query,
+              head_stride_key,
+              num_heads,
+              num_kv_heads,
+              head_size,
+              blocks_per_token);
+        }
+      } else {
+        host::LaunchKernel(grid2d, block, stream)(
+            rotary_embedding_kernel_2d<scalar_t, false, false, true, ROT>,
+            cos_ptr,
+            sin_ptr,
+            q_ptr,
+            k_ptr,
+            rot_dim_from_cache,
+            embed_dim_for_rotation,
+            query_token_stride,
+            key_token_stride,
+            head_stride_query,
+            head_stride_key,
+            num_heads,
+            num_kv_heads,
+            head_size,
+            blocks_per_token);
+      }
+    }
+    return;
+  }
+
+  // 1D grid path
+  if (interleaved) {
+    if (use_vec) {
+      if (qk_aligned16) {
+        host::LaunchKernel(grid1d, block, stream)(
+            rotary_embedding_kernel_1d<scalar_t, true, true, true, ROT>,
+            cos_ptr,
+            sin_ptr,
+            q_ptr,
+            k_ptr,
+            rot_dim_from_cache,
+            embed_dim_for_rotation,
+            query_token_stride,
+            key_token_stride,
+            head_stride_query,
+            head_stride_key,
+            num_heads,
+            num_kv_heads,
+            head_size);
+      } else {
+        host::LaunchKernel(grid1d, block, stream)(
+            rotary_embedding_kernel_1d<scalar_t, true, true, false, ROT>,
+            cos_ptr,
+            sin_ptr,
+            q_ptr,
+            k_ptr,
+            rot_dim_from_cache,
+            embed_dim_for_rotation,
+            query_token_stride,
+            key_token_stride,
+            head_stride_query,
+            head_stride_key,
+            num_heads,
+            num_kv_heads,
+            head_size);
+      }
+    } else {
+      host::LaunchKernel(grid1d, block, stream)(
+          rotary_embedding_kernel_1d<scalar_t, true, false, true, ROT>,
+          cos_ptr,
+          sin_ptr,
+          q_ptr,
+          k_ptr,
+          rot_dim_from_cache,
+          embed_dim_for_rotation,
+          query_token_stride,
+          key_token_stride,
+          head_stride_query,
+          head_stride_key,
+          num_heads,
+          num_kv_heads,
+          head_size);
+    }
+  } else {
+    if (use_vec) {
+      if (qk_aligned16) {
+        host::LaunchKernel(grid1d, block, stream)(
+            rotary_embedding_kernel_1d<scalar_t, false, true, true, ROT>,
+            cos_ptr,
+            sin_ptr,
+            q_ptr,
+            k_ptr,
+            rot_dim_from_cache,
+            embed_dim_for_rotation,
+            query_token_stride,
+            key_token_stride,
+            head_stride_query,
+            head_stride_key,
+            num_heads,
+            num_kv_heads,
+            head_size);
+      } else {
+        host::LaunchKernel(grid1d, block, stream)(
+            rotary_embedding_kernel_1d<scalar_t, false, true, false, ROT>,
+            cos_ptr,
+            sin_ptr,
+            q_ptr,
+            k_ptr,
+            rot_dim_from_cache,
+            embed_dim_for_rotation,
+            query_token_stride,
+            key_token_stride,
+            head_stride_query,
+            head_stride_key,
+            num_heads,
+            num_kv_heads,
+            head_size);
+      }
+    } else {
+      host::LaunchKernel(grid1d, block, stream)(
+          rotary_embedding_kernel_1d<scalar_t, false, false, true, ROT>,
+          cos_ptr,
+          sin_ptr,
+          q_ptr,
+          k_ptr,
+          rot_dim_from_cache,
+          embed_dim_for_rotation,
+          query_token_stride,
+          key_token_stride,
+          head_stride_query,
+          head_stride_key,
+          num_heads,
+          num_kv_heads,
+          head_size);
+    }
+  }
+}
+
+template <typename scalar_t>
+inline void launch_rotary(
+    const tvm::ffi::TensorView cos,
+    const tvm::ffi::TensorView sin,
+    const tvm::ffi::TensorView q,
+    const tvm::ffi::TensorView* k,
+    int64_t head_size,
+    bool interleaved) {
+  using namespace host;
+
+  auto T = SymbolicSize{"T"};
+  auto Hq = SymbolicSize{"Hq"};
+  auto Hk = SymbolicSize{"Hk"};
+  auto D = SymbolicSize{"D"};
+  auto R = SymbolicSize{"R"};
+  auto dtype = SymbolicDType{};
+  auto device = SymbolicDevice{};
+
+  TensorMatcher({T, R}).with_dtype<float, half, nv_bfloat16>(dtype).with_device<kDLCUDA>(device).verify(cos).verify(
+      sin);
+
+  TensorMatcher({T, Hq, D}).with_dtype<float, half, nv_bfloat16>(dtype).with_device<kDLCUDA>(device).verify(q);
+
+  RuntimeCheck(D.unwrap() == head_size, "head_size mismatch: got ", D.unwrap(), " expected ", head_size);
+  RuntimeCheck(cos.size(1) == sin.size(1), "cos/sin dim mismatch");
+
+  const int64_t t = T.unwrap();
+  const int64_t hq = Hq.unwrap();
+  const int64_t d = D.unwrap();
+  const int64_t r = R.unwrap();
+
+  RuntimeCheck(t > 0 && hq > 0 && d > 0 && r > 0, "invalid shape");
+  if (!interleaved) {
+    RuntimeCheck(r % 2 == 0, "non-interleaved requires even R, got ", r);
+  }
+  const int embed_dim_for_rotation = interleaved ? (int)r : (int)(r / 2);
+  RuntimeCheck(embed_dim_for_rotation > 0, "embed_dim_for_rotation must be > 0");
+  RuntimeCheck(2LL * embed_dim_for_rotation <= head_size, "rotate dim exceeds head_size");
+
+  // Strides: JIT wrapper guarantees contiguous [T, H, D] tensors.
+  const int64_t query_token_stride = hq * d;
+  const int64_t head_stride_query = d;
+
+  int64_t key_token_stride = 0;
+  int64_t head_stride_key = d;
+  int hk = 0;
+  if (k != nullptr) {
+    TensorMatcher({T, Hk, D}).with_dtype<float, half, nv_bfloat16>(dtype).with_device<kDLCUDA>(device).verify(*k);
+    hk = (int)Hk.unwrap();
+    RuntimeCheck(hk > 0, "invalid key shape");
+    key_token_stride = (int64_t)hk * d;
+  }
+
+  const int max_pairs_to_rotate_per_token =
+      (k == nullptr) ? ((int)hq * embed_dim_for_rotation)
+                     : std::max((int)hq * embed_dim_for_rotation, (int)hk * embed_dim_for_rotation);
+
+  constexpr int kVecBytes = 16;
+  const int elem_bytes = (int)sizeof(scalar_t);
+  const int kElePerVec = kVecBytes / elem_bytes;
+  const int pairs_per_step = interleaved ? (kElePerVec / 2) : kElePerVec;
+
+  bool can_vec_compute = true;
+  if ((embed_dim_for_rotation % pairs_per_step) != 0) can_vec_compute = false;
+  if ((reinterpret_cast<uintptr_t>(cos.data_ptr()) % kVecBytes) != 0) can_vec_compute = false;
+  if ((reinterpret_cast<uintptr_t>(sin.data_ptr()) % kVecBytes) != 0) can_vec_compute = false;
+  if (((r * elem_bytes) % kVecBytes) != 0) can_vec_compute = false;
+
+  bool qk_aligned16 = true;
+  if ((reinterpret_cast<uintptr_t>(q.data_ptr()) % kVecBytes) != 0) qk_aligned16 = false;
+  if (((query_token_stride * elem_bytes) % kVecBytes) != 0) qk_aligned16 = false;
+  if (((head_stride_query * elem_bytes) % kVecBytes) != 0) qk_aligned16 = false;
+  if (k != nullptr) {
+    if ((reinterpret_cast<uintptr_t>(k->data_ptr()) % kVecBytes) != 0) qk_aligned16 = false;
+    if (((key_token_stride * elem_bytes) % kVecBytes) != 0) qk_aligned16 = false;
+    if (((head_stride_key * elem_bytes) % kVecBytes) != 0) qk_aligned16 = false;
+  }
+
+  const bool use_vec = can_vec_compute;
+  const int launch_pairs_per_thread = use_vec ? pairs_per_step : 1;
+  const int total_threads_needed =
+      (max_pairs_to_rotate_per_token + launch_pairs_per_thread - 1) / launch_pairs_per_thread;
+
+  auto round_up32 = [](int x) { return ((x + 31) / 32) * 32; };
+
+  // Case 1: 2D Grid (only when num_tokens<=4 and needs multiple blocks per token)
+  const int threads_per_block_2d = std::min<int>(512, std::max(128, round_up32(std::min(total_threads_needed, 512))));
+  const int blocks_per_token_2d = (total_threads_needed + threads_per_block_2d - 1) / threads_per_block_2d;
+  const bool use_grid_2d = ((int)t <= 4) && (blocks_per_token_2d > 1);
+
+  // Case 2: 1D Grid
+  const int threads_per_block_1d = std::min<int>(512, std::max(128, round_up32(total_threads_needed)));
+
+  const int threads_per_block = use_grid_2d ? threads_per_block_2d : threads_per_block_1d;
+  const int blocks_per_token = use_grid_2d ? blocks_per_token_2d : 1;
+
+  const auto stream = LaunchKernel::resolve_device(device.unwrap());
+  const scalar_t* cos_ptr = static_cast<const scalar_t*>(cos.data_ptr());
+  const scalar_t* sin_ptr = static_cast<const scalar_t*>(sin.data_ptr());
+  scalar_t* q_ptr = static_cast<scalar_t*>(q.data_ptr());
+  scalar_t* k_ptr = (k != nullptr) ? static_cast<scalar_t*>(k->data_ptr()) : nullptr;
+
+  const dim3 grid2d((int)t, std::max(1, blocks_per_token));
+  const dim3 grid1d((int)t);
+  const dim3 block(threads_per_block);
+
+  // Compile-time specializations for common embed dims; fallback to runtime (ROT=0).
+  switch (embed_dim_for_rotation) {
+    case 32:
+      dispatch_rotary_launch<32, scalar_t>(
+          use_grid_2d,
+          grid2d,
+          grid1d,
+          block,
+          stream,
+          interleaved,
+          use_vec,
+          qk_aligned16,
+          cos_ptr,
+          sin_ptr,
+          q_ptr,
+          k_ptr,
+          (int)r,
+          embed_dim_for_rotation,
+          query_token_stride,
+          key_token_stride,
+          head_stride_query,
+          head_stride_key,
+          (int)hq,
+          (int)hk,
+          (int)head_size,
+          blocks_per_token);
+      break;
+    case 40:
+      dispatch_rotary_launch<40, scalar_t>(
+          use_grid_2d,
+          grid2d,
+          grid1d,
+          block,
+          stream,
+          interleaved,
+          use_vec,
+          qk_aligned16,
+          cos_ptr,
+          sin_ptr,
+          q_ptr,
+          k_ptr,
+          (int)r,
+          embed_dim_for_rotation,
+          query_token_stride,
+          key_token_stride,
+          head_stride_query,
+          head_stride_key,
+          (int)hq,
+          (int)hk,
+          (int)head_size,
+          blocks_per_token);
+      break;
+    case 64:
+      dispatch_rotary_launch<64, scalar_t>(
+          use_grid_2d,
+          grid2d,
+          grid1d,
+          block,
+          stream,
+          interleaved,
+          use_vec,
+          qk_aligned16,
+          cos_ptr,
+          sin_ptr,
+          q_ptr,
+          k_ptr,
+          (int)r,
+          embed_dim_for_rotation,
+          query_token_stride,
+          key_token_stride,
+          head_stride_query,
+          head_stride_key,
+          (int)hq,
+          (int)hk,
+          (int)head_size,
+          blocks_per_token);
+      break;
+    case 80:
+      dispatch_rotary_launch<80, scalar_t>(
+          use_grid_2d,
+          grid2d,
+          grid1d,
+          block,
+          stream,
+          interleaved,
+          use_vec,
+          qk_aligned16,
+          cos_ptr,
+          sin_ptr,
+          q_ptr,
+          k_ptr,
+          (int)r,
+          embed_dim_for_rotation,
+          query_token_stride,
+          key_token_stride,
+          head_stride_query,
+          head_stride_key,
+          (int)hq,
+          (int)hk,
+          (int)head_size,
+          blocks_per_token);
+      break;
+    case 128:
+      dispatch_rotary_launch<128, scalar_t>(
+          use_grid_2d,
+          grid2d,
+          grid1d,
+          block,
+          stream,
+          interleaved,
+          use_vec,
+          qk_aligned16,
+          cos_ptr,
+          sin_ptr,
+          q_ptr,
+          k_ptr,
+          (int)r,
+          embed_dim_for_rotation,
+          query_token_stride,
+          key_token_stride,
+          head_stride_query,
+          head_stride_key,
+          (int)hq,
+          (int)hk,
+          (int)head_size,
+          blocks_per_token);
+      break;
+    case 160:
+      dispatch_rotary_launch<160, scalar_t>(
+          use_grid_2d,
+          grid2d,
+          grid1d,
+          block,
+          stream,
+          interleaved,
+          use_vec,
+          qk_aligned16,
+          cos_ptr,
+          sin_ptr,
+          q_ptr,
+          k_ptr,
+          (int)r,
+          embed_dim_for_rotation,
+          query_token_stride,
+          key_token_stride,
+          head_stride_query,
+          head_stride_key,
+          (int)hq,
+          (int)hk,
+          (int)head_size,
+          blocks_per_token);
+      break;
+    default:
+      dispatch_rotary_launch<0, scalar_t>(
+          use_grid_2d,
+          grid2d,
+          grid1d,
+          block,
+          stream,
+          interleaved,
+          use_vec,
+          qk_aligned16,
+          cos_ptr,
+          sin_ptr,
+          q_ptr,
+          k_ptr,
+          (int)r,
+          embed_dim_for_rotation,
+          query_token_stride,
+          key_token_stride,
+          head_stride_query,
+          head_stride_key,
+          (int)hq,
+          (int)hk,
+          (int)head_size,
+          blocks_per_token);
+      break;
+  }
+
+  RuntimeDeviceCheck();
+}
+
+}  // namespace
+
+struct RotaryEmbeddingCosSinKernel {
+  static void run_q(
+      const tvm::ffi::TensorView cos,
+      const tvm::ffi::TensorView sin,
+      const tvm::ffi::TensorView query,
+      int64_t head_size,
+      bool interleaved) {
+    const auto dt = query.dtype();
+    if (host::is_type<half>(dt)) {
+      launch_rotary<half>(cos, sin, query, nullptr, head_size, interleaved);
+    } else if (host::is_type<nv_bfloat16>(dt)) {
+      launch_rotary<nv_bfloat16>(cos, sin, query, nullptr, head_size, interleaved);
+    } else if (host::is_type<float>(dt)) {
+      launch_rotary<float>(cos, sin, query, nullptr, head_size, interleaved);
+    } else {
+      host::Panic("Unsupported dtype for rotary_embedding_cos_sin");
+    }
+  }
+
+  static void run_qk(
+      const tvm::ffi::TensorView cos,
+      const tvm::ffi::TensorView sin,
+      const tvm::ffi::TensorView query,
+      const tvm::ffi::TensorView key,
+      int64_t head_size,
+      bool interleaved) {
+    const auto dt = query.dtype();
+    if (host::is_type<half>(dt)) {
+      launch_rotary<half>(cos, sin, query, &key, head_size, interleaved);
+    } else if (host::is_type<nv_bfloat16>(dt)) {
+      launch_rotary<nv_bfloat16>(cos, sin, query, &key, head_size, interleaved);
+    } else if (host::is_type<float>(dt)) {
+      launch_rotary<float>(cos, sin, query, &key, head_size, interleaved);
+    } else {
+      host::Panic("Unsupported dtype for rotary_embedding_cos_sin");
+    }
+  }
+};
diff --git a/python/sglang/jit_kernel/rotary_embedding.py b/python/sglang/jit_kernel/rotary_embedding.py
new file mode 100644
index 000000000000..620e61e40d34
--- /dev/null
+++ b/python/sglang/jit_kernel/rotary_embedding.py
@@ -0,0 +1,142 @@
+from __future__ import annotations
+from typing import Optional
+import torch
+from sglang.jit_kernel.utils import cache_once, load_jit
+
+
+def _resolve_interleaved(interleaved: Optional[bool], is_neox: Optional[bool]) -> bool:
+    if interleaved is not None and is_neox is not None and interleaved != is_neox:
+        raise ValueError(
+            f"is_neox({is_neox}) and interleaved({interleaved}) mismatch; keep only one or make them equal."
+        )
+    if is_neox is not None:
+        return is_neox
+    if interleaved is not None:
+        return interleaved
+    return True
+
+
+@cache_once
+def _jit_rotary_embedding_cos_sin_module():
+    return load_jit(
+        "rotary_embedding_cos_sin",
+        cuda_files=["rotary_embedding_cos_sin.cuh"],
+        cuda_wrappers=[
+            ("rotary_embedding_cos_sin_q", "RotaryEmbeddingCosSinKernel::run_q"),
+            ("rotary_embedding_cos_sin_qk", "RotaryEmbeddingCosSinKernel::run_qk"),
+        ],
+    )
+
+
+def rotary_embedding_cos_sin(
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    query: torch.Tensor,
+    key: Optional[torch.Tensor] = None,
+    head_size: int = 0,
+    interleaved: Optional[bool] = None,
+    *,
+    is_neox: Optional[bool] = None,
+) -> None:
+    effective_interleaved = _resolve_interleaved(interleaved, is_neox)
+
+    if query.device.type != "cuda":
+        raise ValueError("query must be a CUDA tensor")
+    if key is not None and key.device.type != "cuda":
+        raise ValueError("key must be a CUDA tensor")
+    if cos.device != query.device or sin.device != query.device:
+        raise ValueError("cos/sin must be on the same device as query")
+
+    if query.dtype != cos.dtype or query.dtype != sin.dtype:
+        raise ValueError("cos/sin dtype must match query dtype")
+    if key is not None and key.dtype != query.dtype:
+        raise ValueError("key dtype must match query dtype")
+
+    if cos.dim() != 2 or sin.dim() != 2:
+        raise ValueError("cos/sin must be 2D tensors: [num_tokens, rot_dim]")
+    if cos.shape != sin.shape:
+        raise ValueError("cos/sin shape mismatch")
+    if cos.shape[0] != query.shape[0]:
+        raise ValueError(
+            f"cos/sin num_tokens mismatch with query: cos.shape[0]={cos.shape[0]} vs query.shape[0]={query.shape[0]}"
+        )
+
+    if head_size == 0:
+        if query.dim() == 3:
+            head_size = int(query.shape[-1])
+        else:
+            raise ValueError("head_size must be provided when query is 2D")
+
+    if not query.is_contiguous():
+        raise ValueError("query must be contiguous (in-place kernel)")
+    if key is not None and (not key.is_contiguous()):
+        raise ValueError("key must be contiguous (in-place kernel)")
+
+    # Optional downsample for interleaved caches that are stored as full-dim (repeat format).
+    # Keep this consistent with AOT kernel behavior.
+    if effective_interleaved and cos.shape[1] == head_size:
+        if head_size % 2 != 0:
+            raise ValueError("interleaved layout requires even head_size")
+        half = head_size // 2
+        cos_to_use = cos.view(cos.shape[0], half, 2).select(2, 0).contiguous()
+        sin_to_use = sin.view(sin.shape[0], half, 2).select(2, 1).contiguous()
+    else:
+        cos_to_use = cos.contiguous()
+        sin_to_use = sin.contiguous()
+
+    rot_dim = int(cos_to_use.shape[1])
+    if rot_dim <= 0:
+        raise ValueError("cos/sin rot_dim must be > 0")
+    if effective_interleaved:
+        # Pairwise layout: rotates 2*rot_dim elements.
+        if 2 * rot_dim > head_size:
+            raise ValueError(
+                f"rotate dim exceeds head_size for interleaved=True: 2*rot_dim={2*rot_dim} > head_size={head_size}"
+            )
+    else:
+        # Split-halves layout: kernel interprets cos/sin as [rot_dim] where embed_dim=rot_dim/2.
+        if rot_dim % 2 != 0:
+            raise ValueError(
+                f"non-interleaved requires even rot_dim (cos.shape[1]), got rot_dim={rot_dim}"
+            )
+        if rot_dim > head_size:
+            raise ValueError(
+                f"rotate dim exceeds head_size for interleaved=False: rot_dim={rot_dim} > head_size={head_size}"
+            )
+
+    def _as_3d(x: torch.Tensor, h: int) -> torch.Tensor:
+        if x.dim() == 3:
+            if x.shape[-1] != head_size:
+                raise ValueError("head_size mismatch with query/key last dim")
+            return x
+        if x.dim() != 2:
+            raise ValueError("query/key must be 2D or 3D tensors")
+        if x.shape[1] % head_size != 0:
+            raise ValueError("hidden_size is not divisible by head_size")
+        return x.view(x.shape[0], h, head_size)
+
+    if query.dim() == 3:
+        num_heads = int(query.shape[1])
+    else:
+        num_heads = int(query.shape[1] // head_size)
+    q3 = _as_3d(query, num_heads)
+    if not q3.is_contiguous():
+        # The CUDA kernel assumes contiguous [T, H, D] layout for fixed-stride addressing.
+        raise ValueError("query must be contiguous (view must stay contiguous)")
+
+    k3 = None
+    if key is not None:
+        if key.dim() == 3:
+            num_kv_heads = int(key.shape[1])
+        else:
+            num_kv_heads = int(key.shape[1] // head_size)
+        k3 = _as_3d(key, num_kv_heads)
+        if not k3.is_contiguous():
+            raise ValueError("key must be contiguous (view must stay contiguous)")
+
+    module = _jit_rotary_embedding_cos_sin_module()
+    if k3 is None:
+        module.rotary_embedding_cos_sin_q(cos_to_use, sin_to_use, q3, head_size, effective_interleaved)
+    else:
+        module.rotary_embedding_cos_sin_qk(cos_to_use, sin_to_use, q3, k3, head_size, effective_interleaved)
+
diff --git a/python/sglang/jit_kernel/tests/test_rotary_embedding.py b/python/sglang/jit_kernel/tests/test_rotary_embedding.py
new file mode 100644
index 000000000000..f061e00039eb
--- /dev/null
+++ b/python/sglang/jit_kernel/tests/test_rotary_embedding.py
@@ -0,0 +1,203 @@
+# pyright: reportMissingImports=false
+
+import torch
+import triton
+
+
+def _compute_cos_sin_cache_half(
+    max_seq_len: int,
+    rotary_dim: int,
+    base: float = 10000.0,
+    dtype: torch.dtype = torch.float32,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    inv_freq = 1.0 / (
+        base ** (torch.arange(0, rotary_dim, 2, dtype=torch.float32) / rotary_dim)
+    )
+    t = torch.arange(max_seq_len, dtype=torch.float32)
+    freqs = torch.einsum("i,j->ij", t, inv_freq)
+    cos = freqs.cos().to(dtype)
+    sin = freqs.sin().to(dtype)
+    return cos, sin
+
+
+def sglang_aot_rotary_positions(
+    positions: torch.Tensor,
+    q: torch.Tensor,
+    k: torch.Tensor,
+    head_size: int,
+    interleaved: bool,
+    cos_sin_cache: torch.Tensor,
+) -> None:
+    # NOTE: `sgl_kernel` positions-kernel uses `is_neox` flag internally:
+    # - is_neox=True  => x_index=i, y_index=embed_dim+i   (split halves)
+    # - is_neox=False => x_index=2*i, y_index=2*i+1       (interleaved pairs)
+    from sgl_kernel.rotary_embedding import rotary_embedding
+
+    rotary_embedding(
+        positions=positions,
+        query=q,
+        key=k,
+        head_size=head_size,
+        is_neox=not interleaved,
+        cos_sin_cache=cos_sin_cache,
+    )
+
+
+def sglang_jit_rotary_cos_sin(
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    q: torch.Tensor,
+    k: torch.Tensor,
+    head_size: int,
+    interleaved: bool,
+) -> None:
+    from sglang.jit_kernel.rotary_embedding import rotary_embedding_cos_sin
+
+    rotary_embedding_cos_sin(cos, sin, q, k, head_size, interleaved)
+
+
+@torch.no_grad()
+def torch_impl_rotary_fp32(
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    q: torch.Tensor,
+    k: torch.Tensor,
+    head_size: int,
+    interleaved: bool,
+) -> None:
+    orig_dtype = q.dtype
+    if interleaved and cos.shape[1] == head_size:
+        half = head_size // 2
+        cos = cos.view(cos.shape[0], half, 2).select(2, 0).contiguous()
+        sin = sin.view(sin.shape[0], half, 2).select(2, 1).contiguous()
+
+    cos_f, sin_f = cos.float(), sin.float()
+    q_f, k_f = q.float(), k.float()
+
+    if interleaved:
+        embed_dim = int(cos_f.shape[1])
+        rot_dim = embed_dim * 2
+        cos_b = cos_f[:, None, :embed_dim]
+        sin_b = sin_f[:, None, :embed_dim]
+
+        def _apply(x: torch.Tensor) -> None:
+            xr = x[..., :rot_dim]
+            xr2 = xr.view(xr.shape[0], xr.shape[1], embed_dim, 2)
+            x0 = xr2[..., 0].clone()
+            x1 = xr2[..., 1].clone()
+            xr2[..., 0].copy_(x0 * cos_b - x1 * sin_b)
+            xr2[..., 1].copy_(x1 * cos_b + x0 * sin_b)
+    else:
+        if cos_f.shape[1] == head_size // 2:
+            embed_dim = int(cos_f.shape[1])
+            rot_dim = embed_dim * 2
+            cos_x, sin_x = cos_f[:, None, :], sin_f[:, None, :]
+            cos_y, sin_y = cos_x, sin_x
+        else:
+            embed_dim = int(cos_f.shape[1]) // 2
+            rot_dim = embed_dim * 2
+            cos_x, sin_x = cos_f[:, None, :embed_dim], sin_f[:, None, :embed_dim]
+            cos_y, sin_y = (
+                cos_f[:, None, embed_dim:rot_dim],
+                sin_f[:, None, embed_dim:rot_dim],
+            )
+
+        def _apply(x: torch.Tensor) -> None:
+            xr = x[..., :rot_dim]
+            x0 = xr[..., :embed_dim].clone()
+            x1 = xr[..., embed_dim:rot_dim].clone()
+            xr[..., :embed_dim].copy_(x0 * cos_x - x1 * sin_x)
+            xr[..., embed_dim:rot_dim].copy_(x1 * cos_y + x0 * sin_y)
+
+    _apply(q_f)
+    _apply(k_f)
+    q.copy_(q_f.to(orig_dtype))
+    k.copy_(k_f.to(orig_dtype))
+
+
+def main():
+    DEVICE = "cuda"
+    if not torch.cuda.is_available():
+        raise SystemExit("CUDA not available")
+
+    NUM_Q_HEADS = 32
+    NUM_KV_HEADS = 8
+    BS_LIST = [2**n for n in range(0, 13)]
+    BS_LIST += [x + 1 + i for i, x in enumerate(BS_LIST)]
+
+    for DTYPE in [torch.bfloat16, torch.float16]:
+        for HEAD_SIZE in [64, 80, 96, 128, 256, 320]:
+            rotary_dim = HEAD_SIZE
+            max_seq_len = 8192
+            cos_cache, sin_cache = _compute_cos_sin_cache_half(
+                max_seq_len, rotary_dim, dtype=DTYPE
+            )
+            cos_cache, sin_cache = cos_cache.to(DEVICE), sin_cache.to(DEVICE)
+
+            for INTERLEAVED in [True, False]:
+                aot_cos_sin_cache = torch.cat(
+                    [cos_cache, sin_cache], dim=-1
+                ).contiguous()
+
+                for BS in BS_LIST:
+                    positions = (
+                        torch.arange(BS, device=DEVICE, dtype=torch.int64) % max_seq_len
+                    )
+                    cos_half = cos_cache[positions].contiguous()
+                    sin_half = sin_cache[positions].contiguous()
+
+                    if INTERLEAVED:
+                        cos, sin = cos_half, sin_half
+                    else:
+                        cos = torch.cat([cos_half, cos_half], dim=-1).contiguous()
+                        sin = torch.cat([sin_half, sin_half], dim=-1).contiguous()
+
+                    q = torch.randn(
+                        BS, NUM_Q_HEADS, HEAD_SIZE, device=DEVICE, dtype=DTYPE
+                    )
+                    k = torch.randn(
+                        BS, NUM_KV_HEADS, HEAD_SIZE, device=DEVICE, dtype=DTYPE
+                    )
+
+                    q_ref_fp32, k_ref_fp32 = q.clone(), k.clone()
+                    torch_impl_rotary_fp32(
+                        cos, sin, q_ref_fp32, k_ref_fp32, HEAD_SIZE, INTERLEAVED
+                    )
+
+                    q_k_aot = (q.clone(), k.clone())
+                    q_k_jit = (q.clone(), k.clone())
+                    sglang_aot_rotary_positions(
+                        positions,
+                        q_k_aot[0],
+                        q_k_aot[1],
+                        HEAD_SIZE,
+                        INTERLEAVED,
+                        aot_cos_sin_cache,
+                    )
+                    sglang_jit_rotary_cos_sin(
+                        cos, sin, q_k_jit[0], q_k_jit[1], HEAD_SIZE, INTERLEAVED
+                    )
+
+                    ref_atol = 2e-2 if DTYPE == torch.bfloat16 else 2e-3
+                    ref_rtol = 2e-2 if DTYPE == torch.bfloat16 else 2e-3
+
+                    triton.testing.assert_close(
+                        q_ref_fp32, q_k_aot[0], atol=ref_atol, rtol=ref_rtol
+                    )
+                    triton.testing.assert_close(
+                        k_ref_fp32, q_k_aot[1], atol=ref_atol, rtol=ref_rtol
+                    )
+                    triton.testing.assert_close(
+                        q_ref_fp32, q_k_jit[0], atol=ref_atol, rtol=ref_rtol
+                    )
+                    triton.testing.assert_close(
+                        k_ref_fp32, q_k_jit[1], atol=ref_atol, rtol=ref_rtol
+                    )
+
+                    print(
+                        f"HEAD_SIZE={HEAD_SIZE} interleaved={INTERLEAVED} dtype={DTYPE} passed."
+                    )
+
+
+if __name__ == "__main__":
+    main()
diff --git a/python/sglang/multimodal_gen/runtime/layers/rotary_embedding.py b/python/sglang/multimodal_gen/runtime/layers/rotary_embedding.py
index ac5e8ed0e091..f5c6e91da7f3 100644
--- a/python/sglang/multimodal_gen/runtime/layers/rotary_embedding.py
+++ b/python/sglang/multimodal_gen/runtime/layers/rotary_embedding.py
@@ -105,6 +105,108 @@ def apply_flashinfer_rope_qk_inplace(
     return q_flat.view(bsz, seqlen, nheads, d), k_flat.view(bsz, seqlen, nheads, d)
 
 
+def _split_cos_sin_from_cache(
+    cache: torch.Tensor, *, dtype: torch.dtype
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Convert a 2D RoPE cache into (cos, sin) tensors."""
+    if cache.is_complex():
+        cos = cache.real.to(dtype)
+        sin = cache.imag.to(dtype)
+        return cos, sin
+    if cache.shape[1] % 2 != 0:
+        raise ValueError(
+            f"Expected complex freqs_cis or cat([cos,sin]) cache; got real cache with odd last dim={cache.shape[1]}"
+        )
+    half = cache.shape[1] // 2
+    cos = cache[:, :half].to(dtype)
+    sin = cache[:, half:].to(dtype)
+    return cos, sin
+
+
+def apply_sglang_jit_rope_qk_inplace(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    cos_sin_cache: torch.Tensor,
+    *,
+    head_size: Optional[int] = None,
+    is_neox: bool = False,
+    positions: Optional[torch.Tensor] = None,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Apply SGLang JIT RoPE on GPU using `sglang.jit_kernel.rotary_embedding_cos_sin`."""
+    if q.dim() != 4 or k.dim() != 4:
+        raise ValueError(
+            f"Expected q/k to be 4D [bsz, seqlen, nheads, head_size], "
+            f"got q:{tuple(q.shape)} k:{tuple(k.shape)}"
+        )
+    if q.shape != k.shape:
+        raise ValueError(f"q and k must have the same shape, got {q.shape} vs {k.shape}")
+
+    if not (isinstance(cos_sin_cache, torch.Tensor) and cos_sin_cache.dim() == 2):
+        raise ValueError("cos_sin_cache must be a 2D torch.Tensor")
+
+    bsz, seqlen, nheads, d = q.shape
+    if head_size is None:
+        head_size = d
+    if head_size != d:
+        raise ValueError(f"head_size mismatch: inferred {d}, but head_size={head_size}")
+
+    try:
+        from sglang.jit_kernel.rotary_embedding import (
+            rotary_embedding_cos_sin as sglang_jit_rotary_embedding_cos_sin,
+        )
+    except Exception as e:
+        raise RuntimeError(
+            "SGLang JIT RoPE is required for apply_sglang_jit_rope_qk_inplace."
+        ) from e
+
+    num_tokens = bsz * seqlen
+
+    if positions is None:
+        pos_1d = torch.arange(seqlen, device="cpu", dtype=torch.long)
+        positions = pos_1d if bsz == 1 else pos_1d.repeat(bsz)
+    else:
+        if not (
+            isinstance(positions, torch.Tensor)
+            and positions.dtype == torch.long
+            and positions.dim() == 1
+        ):
+            raise ValueError("positions must be a 1D torch.long Tensor")
+        if positions.numel() != num_tokens:
+            raise ValueError(
+                f"positions length must be bsz*seqlen={num_tokens}, got {positions.numel()}"
+            )
+    positions = positions.to(q.device, non_blocking=True)
+
+    if cos_sin_cache.shape[0] == seqlen:
+        cache_tok = cos_sin_cache[None, :, :].expand(bsz, seqlen, cos_sin_cache.shape[1]).reshape(
+            num_tokens, cos_sin_cache.shape[1]
+        )
+    elif cos_sin_cache.shape[0] == num_tokens:
+        cache_tok = cos_sin_cache
+    else:
+        cache_tok = cos_sin_cache.index_select(0, positions)
+
+    q3 = q.reshape(num_tokens, nheads, d).contiguous()
+    k3 = k.reshape(num_tokens, nheads, d).contiguous()
+
+    cos_half, sin_half = _split_cos_sin_from_cache(cache_tok.contiguous(), dtype=q.dtype)
+
+    interleaved = not is_neox
+    if interleaved:
+        cos = cos_half.contiguous()
+        sin = sin_half.contiguous()
+    else:
+        if cos_half.shape[1] * 2 == head_size:
+            cos = torch.cat([cos_half, cos_half], dim=-1).contiguous()
+            sin = torch.cat([sin_half, sin_half], dim=-1).contiguous()
+        else:
+            cos = cos_half.contiguous()
+            sin = sin_half.contiguous()
+
+    sglang_jit_rotary_embedding_cos_sin(cos, sin, q3, k3, head_size, interleaved)
+    return q3.view(bsz, seqlen, nheads, d), k3.view(bsz, seqlen, nheads, d)
+
+
 def _rotate_neox(x: torch.Tensor) -> torch.Tensor:
     x1 = x[..., : x.shape[-1] // 2]
     x2 = x[..., x.shape[-1] // 2 :]
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 21bf354818dc..052d7ea97bea 100644
--- a/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py
+++ b/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py
@@ -20,7 +20,7 @@
 from sglang.multimodal_gen.runtime.layers.layernorm import LayerNorm, RMSNorm
 from sglang.multimodal_gen.runtime.layers.linear import ReplicatedLinear
 from sglang.multimodal_gen.runtime.layers.rotary_embedding import (
-    apply_flashinfer_rope_qk_inplace,
+    apply_sglang_jit_rope_qk_inplace,
 )
 from sglang.multimodal_gen.runtime.layers.triton_ops import (
     fuse_scale_shift_gate_select01_kernel,
@@ -565,11 +565,10 @@ def forward(
                 raise RuntimeError("image_rotary_emb must be cos_sin_cache tensors")
 
             img_cache, txt_cache = image_rotary_emb
-
-            img_query, img_key = apply_flashinfer_rope_qk_inplace(
+            img_query, img_key = apply_sglang_jit_rope_qk_inplace(
                 img_query, img_key, img_cache, is_neox=False
             )
-            txt_query, txt_key = apply_flashinfer_rope_qk_inplace(
+            txt_query, txt_key = apply_sglang_jit_rope_qk_inplace(
                 txt_query, txt_key, txt_cache, is_neox=False
             )
 
diff --git a/sgl-kernel/CMakeLists.txt b/sgl-kernel/CMakeLists.txt
index ccfd7d47b24a..1862e8707d5e 100644
--- a/sgl-kernel/CMakeLists.txt
+++ b/sgl-kernel/CMakeLists.txt
@@ -316,7 +316,6 @@ set(SOURCES
     "csrc/kvcacheio/transfer.cu"
     "csrc/mamba/causal_conv1d.cu"
     "csrc/memory/store.cu"
-    "csrc/sgl_diffusion/rope/rotary_embedding.cu"
     "csrc/memory/weak_ref_tensor.cpp"
 
     "csrc/moe/cutlass_moe/w4a8/scaled_mm_entry.cu"
diff --git a/sgl-kernel/benchmark/bench_mm_rotary_embedding.py b/sgl-kernel/benchmark/bench_mm_rotary_embedding.py
deleted file mode 100755
index b3d833447e33..000000000000
--- a/sgl-kernel/benchmark/bench_mm_rotary_embedding.py
+++ /dev/null
@@ -1,795 +0,0 @@
-#!/usr/bin/env python
-# -*- coding: utf-8 -*-
-
-from __future__ import annotations
-
-import argparse
-from dataclasses import dataclass
-from typing import List, Optional, Tuple
-
-import numpy as np
-import torch
-import triton
-from sgl_kernel.rotary_embedding import rotary_embedding_cos_sin as sgl_rotary_cos_sin
-from sgl_kernel.testing.rotary_embedding import RotaryEmbedding as NativeRotaryEmbedding
-
-
-def compute_cos_sin_cache_half(
-    max_seq_len: int,
-    rotary_dim: int,
-    base: float = 10000.0,
-    dtype: torch.dtype = torch.float32,
-) -> Tuple[torch.Tensor, torch.Tensor]:
-    inv_freq = 1.0 / (
-        base ** (torch.arange(0, rotary_dim, 2, dtype=torch.float32) / rotary_dim)
-    )
-    t = torch.arange(max_seq_len, dtype=torch.float32)
-    freqs = torch.einsum("i,j->ij", t, inv_freq)
-    cos = freqs.cos().to(dtype)
-    sin = freqs.sin().to(dtype)
-    return cos, sin
-
-
-def _expand_neox_full_repeat(half: torch.Tensor) -> torch.Tensor:
-    out = torch.empty(
-        (half.shape[0], half.shape[1] * 2), dtype=half.dtype, device=half.device
-    )
-    out[:, 0::2] = half
-    out[:, 1::2] = half
-    return out
-
-
-def _expand_llama_full_cat(half: torch.Tensor) -> torch.Tensor:
-    return torch.cat([half, half], dim=-1)
-
-
-def _make_misaligned_contiguous_like(
-    t: torch.Tensor, offset_elems: int = 1
-) -> torch.Tensor:
-    # Create a contiguous tensor with same shape/dtype/device but misaligned base pointer.
-    numel = t.numel()
-    buf = torch.empty((numel + offset_elems,), device=t.device, dtype=t.dtype)
-    out = buf[offset_elems : offset_elems + numel].view_as(t)
-    out.copy_(t)
-    assert out.is_contiguous()
-    return out
-
-
-def _make_tensor(
-    shape: Tuple[int, ...], *, dtype: torch.dtype, device: str, misalign: bool
-) -> torch.Tensor:
-    numel = int(np.prod(shape))
-    if not misalign:
-        return torch.randn(*shape, dtype=dtype, device=device)
-    buf = torch.randn((numel + 1,), dtype=dtype, device=device)
-    return buf[1 : 1 + numel].view(shape)
-
-
-def _torch_naive_rope_inplace(
-    *,
-    cos: torch.Tensor,
-    sin: torch.Tensor,
-    q: torch.Tensor,
-    k: Optional[torch.Tensor],
-    num_heads: int,
-    num_kv_heads: int,
-    head_size: int,
-    interleaved: bool,
-) -> None:
-    """Naive RoPE on GPU using PyTorch tensor ops (in-place on q/k)."""
-
-    def _as_3d(x: torch.Tensor, h: int) -> torch.Tensor:
-        if x.dim() == 3:
-            return x
-        return x.view(x.shape[0], h, head_size)
-
-    if interleaved and cos.shape[1] == head_size:
-        half = head_size // 2
-        cos = cos.view(cos.shape[0], half, 2).select(2, 0).contiguous()
-        sin = sin.view(sin.shape[0], half, 2).select(2, 1).contiguous()
-
-    if interleaved:
-        embed_dim = int(cos.shape[1])
-        rot_dim = embed_dim * 2
-    else:
-        embed_dim = int(cos.shape[1]) // 2
-        rot_dim = embed_dim * 2
-
-    cos_b = cos[:, None, :embed_dim]
-    sin_b = sin[:, None, :embed_dim]
-
-    def _apply(x: torch.Tensor, h: int) -> None:
-        x3 = _as_3d(x, h)
-        xr = x3[..., :rot_dim]
-        if interleaved:
-            xr2 = xr.view(xr.shape[0], xr.shape[1], embed_dim, 2)
-            x0 = xr2[..., 0]
-            x1 = xr2[..., 1]
-            out0 = x0 * cos_b - x1 * sin_b
-            out1 = x1 * cos_b + x0 * sin_b
-            xr2[..., 0].copy_(out0)
-            xr2[..., 1].copy_(out1)
-        else:
-            x0 = xr[..., :embed_dim]
-            x1 = xr[..., embed_dim:rot_dim]
-            cos_x = cos[:, None, :embed_dim]
-            sin_x = sin[:, None, :embed_dim]
-            cos_y = cos[:, None, embed_dim:rot_dim]
-            sin_y = sin[:, None, embed_dim:rot_dim]
-            out0 = x0 * cos_x - x1 * sin_x
-            out1 = x1 * cos_y + x0 * sin_y
-            xr[..., :embed_dim].copy_(out0)
-            xr[..., embed_dim:rot_dim].copy_(out1)
-
-    _apply(q, num_heads)
-    if k is not None:
-        _apply(k, num_kv_heads)
-
-
-def _predict_vec_flags(
-    *,
-    dtype: torch.dtype,
-    interleaved: bool,
-    embed_dim_for_rotation: int,
-    query: torch.Tensor,
-    key: Optional[torch.Tensor],
-    cos: torch.Tensor,
-    sin: torch.Tensor,
-    query_token_stride_elems: int,
-    key_token_stride_elems: int,
-    head_stride_query_elems: int,
-    head_stride_key_elems: int,
-) -> Tuple[bool, bool]:
-    """Predict flags for the *optimized* kernel dispatch:
-    - use_vec_compute: vectorized compute loop (pairs_per_step) is allowed
-    - qk_aligned16: Q/K supports 16B vec load/store (otherwise use unaligned Q/K path)
-    """
-    kVecBytes = 16
-    elem_bytes = torch.tensor([], dtype=dtype).element_size()
-    kElePerVec = kVecBytes // elem_bytes
-    pairs_per_step = (kElePerVec // 2) if interleaved else kElePerVec
-
-    # A) vectorized compute viability (cos/sin alignment + embed_dim multiple)
-    use_vec_compute = True
-    if embed_dim_for_rotation % pairs_per_step != 0:
-        use_vec_compute = False
-
-    # cos/sin alignment & row stride for vector loads (float2/float4 reinterprets)
-    if (cos.data_ptr() % kVecBytes) != 0:
-        use_vec_compute = False
-    if (sin.data_ptr() % kVecBytes) != 0:
-        use_vec_compute = False
-    if ((cos.shape[1] * elem_bytes) % kVecBytes) != 0:
-        use_vec_compute = False
-    if ((sin.shape[1] * elem_bytes) % kVecBytes) != 0:
-        use_vec_compute = False
-
-    # B) Q/K 16B alignment (only affects aligned_qk template)
-    qk_aligned16 = True
-    if (query.data_ptr() % kVecBytes) != 0:
-        qk_aligned16 = False
-    if (query_token_stride_elems * elem_bytes) % kVecBytes != 0:
-        qk_aligned16 = False
-    if (head_stride_query_elems * elem_bytes) % kVecBytes != 0:
-        qk_aligned16 = False
-
-    if key is not None:
-        if (key.data_ptr() % kVecBytes) != 0:
-            qk_aligned16 = False
-        if (key_token_stride_elems * elem_bytes) % kVecBytes != 0:
-            qk_aligned16 = False
-        if (head_stride_key_elems * elem_bytes) % kVecBytes != 0:
-            qk_aligned16 = False
-
-    return use_vec_compute, qk_aligned16
-
-
-def _predict_use_grid_2d(
-    *,
-    num_tokens: int,
-    num_heads: int,
-    num_kv_heads: int,
-    embed_dim_for_rotation: int,
-    use_vec_compute: bool,
-    dtype: torch.dtype,
-    interleaved: bool,
-) -> Tuple[bool, int]:
-    kVecBytes = 16
-    elem_bytes = torch.tensor([], dtype=dtype).element_size()
-    kElePerVec = kVecBytes // elem_bytes
-    pairs_per_step = (kElePerVec // 2) if interleaved else kElePerVec
-    launch_pairs_per_thread = pairs_per_step if use_vec_compute else 1
-
-    max_pairs = max(
-        num_heads * embed_dim_for_rotation, num_kv_heads * embed_dim_for_rotation
-    )
-    total_threads_needed = (
-        max_pairs + launch_pairs_per_thread - 1
-    ) // launch_pairs_per_thread
-
-    def _round_up32(x: int) -> int:
-        return ((x + 31) // 32) * 32
-
-    threads_per_block_2d = min(
-        512, max(128, _round_up32(min(total_threads_needed, 512)))
-    )
-    blocks_per_token_2d = (
-        total_threads_needed + threads_per_block_2d - 1
-    ) // threads_per_block_2d
-    use_grid_2d = (num_tokens <= 4) and (blocks_per_token_2d > 1)
-    return use_grid_2d, blocks_per_token_2d
-
-
-@dataclass(frozen=True)
-class BenchCase:
-    name: str
-    batch_size: int
-    num_heads: int
-    num_kv_heads: int
-    head_size: int
-    rotary_dim: int
-    interleaved: bool
-    dtype: torch.dtype
-    layout: str  # "2d" or "3d"
-    key_mode: str  # "with_k" or "no_k"
-    cache_format: str  # "half" | "neox_full_repeat" | "llama_full_cat"
-    misalign: str  # "none" | "q" | "cos"
-    seq_lens: List[int]
-    bench_native: bool = True
-
-
-def benchmark_mm_rotary_embedding(args: argparse.Namespace) -> None:
-    try:
-        from vllm.model_executor.layers.rotary_embedding import (
-            RotaryEmbedding as vLLMRotaryEmbedding,
-        )
-
-        HAS_VLLM = True
-    except ImportError:
-        vLLMRotaryEmbedding = None
-        HAS_VLLM = False
-        print("vLLM not available")
-
-    try:
-        from flash_attn.layers.rotary import RotaryEmbedding as FlashRotaryEmbedding
-
-        HAS_FLASH_ATTN = True
-    except ImportError:
-        FlashRotaryEmbedding = None
-        HAS_FLASH_ATTN = False
-        print("flash_attn not available")
-
-    device = "cuda"
-    max_seq_len = 65536
-
-    print(
-        "Legend:\n"
-        "  - vec: vec16 => predicted 16B-aligned vector path; vec16u => vectorized compute but unaligned Q/K load/store; scalar => scalar/fallback\n"
-        "  - grid: 1D => one block per token; 2D => split one token across multiple blocks (only when num_tokens<=4)\n"
-        "  - bpt2d: estimated blocks-per-token if 2D were used (useful to judge 'small grid but many blocks per token')\n"
-    )
-
-    cases: List[BenchCase] = [
-        BenchCase(
-            name="base_bf16_int2d_k_halfcache",
-            batch_size=1,
-            num_heads=32,
-            num_kv_heads=8,
-            head_size=128,
-            rotary_dim=128,
-            interleaved=True,
-            dtype=torch.bfloat16,
-            layout="2d",
-            key_mode="with_k",
-            cache_format="half",
-            misalign="none",
-            seq_lens=[1, 2, 4, 32, 256, 3015, 8192],
-            bench_native=True,
-        ),
-        BenchCase(
-            name="base_bf16_nonint2d_k_llamacache",
-            batch_size=1,
-            num_heads=32,
-            num_kv_heads=8,
-            head_size=128,
-            rotary_dim=128,
-            interleaved=False,
-            dtype=torch.bfloat16,
-            layout="2d",
-            key_mode="with_k",
-            cache_format="llama_full_cat",
-            misalign="none",
-            seq_lens=[1, 4, 256, 8192],
-            bench_native=False,
-        ),
-        BenchCase(
-            name="base_bf16_int2d_k_halfcache_misalignQ",
-            batch_size=1,
-            num_heads=32,
-            num_kv_heads=8,
-            head_size=128,
-            rotary_dim=128,
-            interleaved=True,
-            dtype=torch.bfloat16,
-            layout="2d",
-            key_mode="with_k",
-            cache_format="half",
-            misalign="q",
-            seq_lens=[1, 4, 256, 8192],
-            bench_native=False,
-        ),
-        BenchCase(
-            name="base_bf16_int2d_noK_halfcache",
-            batch_size=1,
-            num_heads=32,
-            num_kv_heads=8,
-            head_size=128,
-            rotary_dim=128,
-            interleaved=True,
-            dtype=torch.bfloat16,
-            layout="2d",
-            key_mode="no_k",
-            cache_format="half",
-            misalign="none",
-            seq_lens=[1, 4, 256, 8192],
-            bench_native=False,
-        ),
-        BenchCase(
-            name="base_bf16_int3d_k_halfcache",
-            batch_size=1,
-            num_heads=32,
-            num_kv_heads=8,
-            head_size=128,
-            rotary_dim=128,
-            interleaved=True,
-            dtype=torch.bfloat16,
-            layout="3d",
-            key_mode="with_k",
-            cache_format="half",
-            misalign="none",
-            seq_lens=[1, 4, 256, 8192],
-            bench_native=False,
-        ),
-        BenchCase(
-            name="base_fp16_int2d_k_halfcache",
-            batch_size=1,
-            num_heads=32,
-            num_kv_heads=8,
-            head_size=128,
-            rotary_dim=128,
-            interleaved=True,
-            dtype=torch.float16,
-            layout="2d",
-            key_mode="with_k",
-            cache_format="half",
-            misalign="none",
-            seq_lens=[1, 4, 256, 8192],
-            bench_native=False,
-        ),
-        BenchCase(
-            name="qwen_image_bf16_int2d_k_halfcache",
-            batch_size=1,
-            num_heads=24,
-            num_kv_heads=24,
-            head_size=128,
-            rotary_dim=128,
-            interleaved=True,
-            dtype=torch.bfloat16,
-            layout="2d",
-            key_mode="with_k",
-            cache_format="half",
-            misalign="none",
-            seq_lens=[1, 4, 3015],
-            bench_native=False,
-        ),
-        BenchCase(
-            name="grid2d_stress_bf16_int2d_k_neoxfullrepeat",
-            batch_size=1,
-            num_heads=128,
-            num_kv_heads=128,
-            head_size=256,
-            rotary_dim=256,
-            interleaved=True,
-            dtype=torch.bfloat16,
-            layout="2d",
-            key_mode="with_k",
-            cache_format="neox_full_repeat",
-            misalign="none",
-            seq_lens=[1, 2, 4],
-            bench_native=False,
-        ),
-        BenchCase(
-            name="batch32_bf16_int2d_k_halfcache",
-            batch_size=32,
-            num_heads=32,
-            num_kv_heads=8,
-            head_size=64,
-            rotary_dim=64,
-            interleaved=True,
-            dtype=torch.bfloat16,
-            layout="2d",
-            key_mode="with_k",
-            cache_format="half",
-            misalign="none",
-            seq_lens=[1, 4, 128, 1024],
-            bench_native=False,
-        ),
-    ]
-
-    if args.case is not None:
-        allow = set(args.case)
-        cases = [c for c in cases if c.name in allow]
-        missing = allow.difference({c.name for c in cases})
-        if missing:
-            print(f"WARNING: unknown --case entries ignored: {sorted(missing)}")
-
-    for case in cases:
-        try:
-            torch.cuda.synchronize()
-        except torch.AcceleratorError:
-            torch.cuda.empty_cache()
-            pass
-
-        try:
-            cos_half, sin_half = compute_cos_sin_cache_half(
-                max_seq_len, case.rotary_dim, dtype=torch.float32
-            )
-            cos_half = cos_half.to(device=device, dtype=case.dtype)
-            sin_half = sin_half.to(device=device, dtype=case.dtype)
-
-            if case.cache_format == "half":
-                cos_cache = cos_half
-                sin_cache = sin_half
-            elif case.cache_format == "neox_full_repeat":
-                cos_cache = _expand_neox_full_repeat(cos_half)
-                sin_cache = _expand_neox_full_repeat(sin_half)
-            elif case.cache_format == "llama_full_cat":
-                cos_cache = _expand_llama_full_cat(cos_half)
-                sin_cache = _expand_llama_full_cat(sin_half)
-            else:
-                raise ValueError(f"Unknown cache_format={case.cache_format}")
-
-            cos_sin_cache = torch.cat([cos_cache, sin_cache], dim=-1)
-
-            native_rope = NativeRotaryEmbedding(
-                head_size=case.head_size,
-                rotary_dim=case.rotary_dim,
-                max_position_embeddings=max_seq_len,
-                base=10000,
-                is_neox_style=case.interleaved,
-                dtype=case.dtype,
-            ).to(device)
-
-            print(
-                f"\nCase: {case.name} | "
-                f"bs={case.batch_size} heads={case.num_heads}/{case.num_kv_heads} "
-                f"hs={case.head_size} rd={case.rotary_dim} "
-                f"dtype={case.dtype} interleaved={case.interleaved} layout={case.layout} "
-                f"key={case.key_mode} cache={case.cache_format} misalign={case.misalign}"
-            )
-            print("-" * 100)
-        except Exception as e:
-            print(f"\nSkipping case {case.name}: {e}")
-            continue
-
-        header = f"{'seq_len':>8}"
-        if case.bench_native:
-            header += f" | {'native (ms)':>12}"
-        header += f" | {'naive (ms)':>11}"
-        header += f" | {'sgl_cos_sin (ms)':>16}"
-        header += f" | {'sgl_pos (ms)':>12}"
-        header += f" | {'vec':>6} | {'grid':>4} | {'bpt2d':>5}"
-        if HAS_VLLM:
-            header += f" | {'vLLM (ms)':>10}"
-        if HAS_FLASH_ATTN:
-            header += f" | {'flash_attn (ms)':>14}"
-        header += f" | {'pos/cos':>9}"
-        print(header)
-        print("-" * 100)
-
-        results = []
-
-        for seq_len in case.seq_lens:
-            try:
-                num_tokens = case.batch_size * seq_len
-                if case.layout == "2d":
-                    query_shape = (num_tokens, case.num_heads * case.head_size)
-                    key_shape = (num_tokens, case.num_kv_heads * case.head_size)
-                elif case.layout == "3d":
-                    query_shape = (num_tokens, case.num_heads, case.head_size)
-                    key_shape = (num_tokens, case.num_kv_heads, case.head_size)
-                else:
-                    raise ValueError(f"Unknown layout={case.layout}")
-
-                query = _make_tensor(
-                    query_shape,
-                    dtype=case.dtype,
-                    device=device,
-                    misalign=(case.misalign == "q"),
-                )
-                key = _make_tensor(
-                    key_shape, dtype=case.dtype, device=device, misalign=False
-                )
-                if case.key_mode == "no_k":
-                    key = None
-
-                positions = torch.arange(
-                    seq_len, device=device, dtype=torch.int64
-                ).repeat(case.batch_size)
-                cos = cos_cache[positions]
-                sin = sin_cache[positions]
-
-                if case.misalign == "cos":
-                    cos = _make_misaligned_contiguous_like(cos)
-                    sin = _make_misaligned_contiguous_like(sin)
-
-                row_str = f"{seq_len:8d}"
-                native_time = None
-                naive_time = None
-                sgl_cos_sin_time = None
-                sgl_pos_time = None
-                vllm_time = None
-                fa_time = None
-            except (RuntimeError, torch.AcceleratorError):
-                print(f"{seq_len:8d} | SKIP (CUDA error from previous iteration)")
-                torch.cuda.synchronize()
-                continue
-
-            def fn_native() -> None:
-                q = query.clone()
-                if key is None:
-                    raise RuntimeError("native requires key")
-                k = key.clone()
-                native_rope.forward_native(positions, q, k)
-
-            if case.bench_native and key is not None:
-                ms, _, _ = triton.testing.do_bench(fn_native, quantiles=[0.5, 0.2, 0.8])
-                native_time = 1000 * ms
-                row_str += f" | {native_time:12.4f}"
-            elif case.bench_native:
-                row_str += f" | {'N/A':>12}"
-
-            @torch.no_grad()
-            def fn_naive() -> None:
-                qn = query.clone()
-                kn = None if key is None else key.clone()
-                _torch_naive_rope_inplace(
-                    cos=cos,
-                    sin=sin,
-                    q=qn,
-                    k=kn,
-                    num_heads=case.num_heads,
-                    num_kv_heads=case.num_kv_heads,
-                    head_size=case.head_size,
-                    interleaved=case.interleaved,
-                )
-
-            ms, _, _ = triton.testing.do_bench(fn_naive, quantiles=[0.5, 0.2, 0.8])
-            naive_time = 1000 * ms
-            row_str += f" | {naive_time:11.4f}"
-
-            def fn_sgl_cos_sin() -> None:
-                q = query.clone()
-                k = None if key is None else key.clone()
-                sgl_rotary_cos_sin(cos, sin, q, k, case.head_size, case.interleaved)
-
-            ms, _, _ = triton.testing.do_bench(
-                fn_sgl_cos_sin, quantiles=[0.5, 0.2, 0.8]
-            )
-            sgl_cos_sin_time = 1000 * ms
-            row_str += f" | {sgl_cos_sin_time:16.4f}"
-
-            def fn_sgl_pos() -> None:
-                q = query.clone()
-                k = None if key is None else key.clone()
-                torch.ops.sgl_kernel.rotary_embedding(
-                    positions,
-                    q,
-                    k,
-                    case.head_size,
-                    cos_sin_cache,
-                    case.interleaved,
-                )
-
-            try:
-                ms, _, _ = triton.testing.do_bench(
-                    fn_sgl_pos, quantiles=[0.5, 0.2, 0.8]
-                )
-                sgl_pos_time = 1000 * ms
-                row_str += f" | {sgl_pos_time:12.4f}"
-            except Exception:
-                row_str += f" | {'ERROR':>12}"
-                sgl_pos_time = None
-
-            # Prediction flags (updated for optimized kernel)
-            rot_dim_from_cache = int(cos.shape[1])
-            embed_dim_for_rotation = (
-                rot_dim_from_cache if case.interleaved else (rot_dim_from_cache // 2)
-            )
-
-            if query.dim() == 2:
-                query_hidden = int(query.size(1))
-                query_token_stride = query_hidden
-                head_stride_query = case.head_size
-            else:
-                query_hidden = int(query.size(1) * query.size(2))
-                query_token_stride = query_hidden
-                head_stride_query = int(query.stride(1))
-
-            if key is None:
-                key_token_stride = 0
-                head_stride_key = case.head_size
-            else:
-                if key.dim() == 2:
-                    key_hidden = int(key.size(1))
-                    key_token_stride = key_hidden
-                    head_stride_key = case.head_size
-                else:
-                    key_hidden = int(key.size(1) * key.size(2))
-                    key_token_stride = key_hidden
-                    head_stride_key = int(key.stride(1))
-
-            use_vec_compute, qk_aligned16 = _predict_vec_flags(
-                dtype=case.dtype,
-                interleaved=case.interleaved,
-                embed_dim_for_rotation=embed_dim_for_rotation,
-                query=query,
-                key=key,
-                cos=cos,
-                sin=sin,
-                query_token_stride_elems=query_token_stride,
-                key_token_stride_elems=key_token_stride,
-                head_stride_query_elems=head_stride_query,
-                head_stride_key_elems=head_stride_key,
-            )
-
-            use_grid2d, bpt2d = _predict_use_grid_2d(
-                num_tokens=num_tokens,
-                num_heads=case.num_heads,
-                num_kv_heads=(0 if key is None else case.num_kv_heads),
-                embed_dim_for_rotation=embed_dim_for_rotation,
-                use_vec_compute=use_vec_compute,
-                dtype=case.dtype,
-                interleaved=case.interleaved,
-            )
-
-            if use_vec_compute:
-                vec_str = "vec16" if qk_aligned16 else "vec16u"
-            else:
-                vec_str = "scalar"
-
-            grid_str = "2D" if use_grid2d else "1D"
-            row_str += f" | {vec_str:>6} | {grid_str:>4} | {str(bpt2d):>5}"
-
-            if HAS_VLLM:
-                vllm_rope = vLLMRotaryEmbedding(
-                    head_size=case.head_size,
-                    rotary_dim=case.rotary_dim,
-                    max_position_embeddings=max_seq_len,
-                    base=10000,
-                    is_neox_style=case.interleaved,
-                    dtype=case.dtype,
-                ).cuda()
-
-                def fn_vllm() -> None:
-                    q = query.clone()
-                    if key is None:
-                        raise RuntimeError("vLLM requires key")
-                    k = key.clone()
-                    vllm_rope.forward_cuda(positions, q, k)
-
-                try:
-                    ms, _, _ = triton.testing.do_bench(
-                        fn_vllm, quantiles=[0.5, 0.2, 0.8]
-                    )
-                    vllm_time = 1000 * ms
-                    row_str += f" | {vllm_time:10.4f}"
-                except Exception:
-                    row_str += f" | {'ERROR':>10}"
-                    vllm_time = None
-
-            if HAS_FLASH_ATTN:
-                try:
-                    from flash_attn.layers.rotary import (  # noqa
-                        RotaryEmbedding as FlashRotaryEmbedding,
-                    )
-
-                    flash_rotary = FlashRotaryEmbedding(case.rotary_dim, device=device)
-                    qkv = torch.randn(
-                        case.batch_size,
-                        seq_len,
-                        3,
-                        case.num_heads,
-                        case.head_size,
-                        dtype=case.dtype,
-                        device=device,
-                    )
-
-                    def fn_flash_attn() -> None:
-                        qkv_fa = qkv.clone()
-                        flash_rotary(qkv_fa, seqlen_offset=0)
-
-                    ms, _, _ = triton.testing.do_bench(
-                        fn_flash_attn, quantiles=[0.5, 0.2, 0.8]
-                    )
-                    fa_time = 1000 * ms
-                    row_str += f" | {fa_time:14.4f}"
-                except Exception:
-                    row_str += f" | {'ERROR':>14}"
-                    fa_time = None
-
-            if (
-                (sgl_pos_time is not None)
-                and (sgl_cos_sin_time is not None)
-                and (sgl_cos_sin_time > 0)
-            ):
-                speedup = sgl_pos_time / sgl_cos_sin_time
-                row_str += f" | {speedup:9.2f}x"
-            else:
-                row_str += f" | {'N/A':>9}"
-
-            print(row_str)
-            results.append(
-                {
-                    "seq_len": seq_len,
-                    "native": native_time,
-                    "naive": naive_time,
-                    "sgl_cos_sin": sgl_cos_sin_time,
-                    "sgl_pos": sgl_pos_time,
-                    "vllm": vllm_time,
-                    "flash_attn": fa_time,
-                }
-            )
-
-        native_vals = [r["native"] for r in results if r["native"] is not None]
-        naive_vals = [r["naive"] for r in results if r["naive"] is not None]
-        sgl_cos_sin_vals = [
-            r["sgl_cos_sin"] for r in results if r["sgl_cos_sin"] is not None
-        ]
-        sgl_pos_vals = [r["sgl_pos"] for r in results if r["sgl_pos"] is not None]
-        vllm_vals = [r["vllm"] for r in results if r["vllm"] is not None]
-        fa_vals = [r["flash_attn"] for r in results if r["flash_attn"] is not None]
-
-        avg_row = f"{'AVG':>8}"
-        if case.bench_native:
-            avg_row += (
-                f" | {np.mean(native_vals):12.4f}" if native_vals else f" | {'N/A':>12}"
-            )
-        avg_row += (
-            f" | {np.mean(naive_vals):11.4f}" if naive_vals else f" | {'N/A':>11}"
-        )
-        avg_row += (
-            f" | {np.mean(sgl_cos_sin_vals):16.4f}"
-            if sgl_cos_sin_vals
-            else f" | {'N/A':>16}"
-        )
-        avg_row += (
-            f" | {np.mean(sgl_pos_vals):12.4f}" if sgl_pos_vals else f" | {'N/A':>12}"
-        )
-        avg_row += f" | {'-':>6} | {'-':>4} | {'-':>5}"
-
-        if HAS_VLLM:
-            avg_row += (
-                f" | {np.mean(vllm_vals):10.4f}" if vllm_vals else f" | {'N/A':>10}"
-            )
-        if HAS_FLASH_ATTN:
-            avg_row += f" | {np.mean(fa_vals):14.4f}" if fa_vals else f" | {'N/A':>14}"
-        if sgl_pos_vals and sgl_cos_sin_vals and (np.mean(sgl_cos_sin_vals) > 0):
-            avg_speedup = float(np.mean(sgl_pos_vals) / np.mean(sgl_cos_sin_vals))
-            avg_row += f" | {avg_speedup:9.2f}x"
-        else:
-            avg_row += f" | {'N/A':>9}"
-
-        print("-" * 100)
-        print(avg_row)
-
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser(description="Benchmark rotary embedding ")
-    parser.add_argument(
-        "--case",
-        action="append",
-        default=None,
-        help="Run only the named case. If omitted, runs all cases.",
-    )
-    benchmark_mm_rotary_embedding(parser.parse_args())
diff --git a/sgl-kernel/csrc/common_extension.cc b/sgl-kernel/csrc/common_extension.cc
index 49a98836b628..c82c8d66321f 100644
--- a/sgl-kernel/csrc/common_extension.cc
+++ b/sgl-kernel/csrc/common_extension.cc
@@ -91,23 +91,10 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) {
   m.impl("apply_rope_pos_ids_cos_sin_cache", torch::kCUDA, &apply_rope_pos_ids_cos_sin_cache);
 
   m.def(
-      "rotary_embedding(Tensor positions, Tensor! query,Tensor!? key, int head_size, Tensor cos_sin_cache, bool "
-      "is_neox) -> ()");
-  m.impl(
-      "rotary_embedding",
-      torch::kCUDA,
-      static_cast<void (*)(
-          torch::Tensor&, torch::Tensor&, std::optional<torch::Tensor>, int64_t, torch::Tensor&, bool)>(
-          &rotary_embedding));
-
-  m.def(
-      "rotary_embedding_cos_sin(Tensor cos, Tensor sin, Tensor(a!) query, Tensor? key, int head_size, bool "
-      "interleaved) -> ()");
-  m.impl(
-      "rotary_embedding_cos_sin",
-      torch::kCUDA,
-      static_cast<void (*)(at::Tensor&, at::Tensor&, at::Tensor&, const std::optional<at::Tensor>&, int64_t, bool)>(
-          &rotary_embedding_cos_sin));
+    "rotary_embedding(Tensor positions, Tensor! query,"
+    "                 Tensor!? key, int head_size,"
+    "                 Tensor cos_sin_cache, bool is_neox) -> ()");
+  m.impl("rotary_embedding", torch::kCUDA, &rotary_embedding);
 
   m.def(
       "downcast_fp8(Tensor k, Tensor v, Tensor k_out, Tensor v_out, Tensor k_scale, Tensor v_scale, Tensor loc, "
diff --git a/sgl-kernel/csrc/sgl_diffusion/rope/rotary_embedding.cu b/sgl-kernel/csrc/sgl_diffusion/rope/rotary_embedding.cu
deleted file mode 100644
index 31fdce6ed056..000000000000
--- a/sgl-kernel/csrc/sgl_diffusion/rope/rotary_embedding.cu
+++ /dev/null
@@ -1,1054 +0,0 @@
-/*
- * Copyright (c) 2025 by SGLang team.
- * Copyright (c) 2025 by FlashInfer team.
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *   http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-#include <ATen/cuda/CUDAContext.h>
-#include <c10/cuda/CUDAException.h>
-#include <c10/cuda/CUDAGuard.h>
-#include <c10/cuda/CUDAStream.h>
-#include <torch/all.h>
-
-#include <cmath>
-#include <vector>
-
-#include "utils.h"
-
-// Packed vector payload per step (16B compute tile)
-template <typename scalar_t, bool interleaved>
-struct RotaryVecData;
-
-template <typename scalar_t>
-struct RotaryVecData<scalar_t, true> {
-  float4 vec;
-  float2 cos_vec;
-  float2 sin_vec;
-};
-
-template <typename scalar_t>
-struct RotaryVecData<scalar_t, false> {
-  float4 vec_x;
-  float4 vec_y;
-  float4 cos_x;
-  float4 sin_x;
-  float4 cos_y;
-  float4 sin_y;
-};
-
-// Vector load (q/k aligned or scalar-gather) + always vector-load cos/sin.
-// rot_offset is "pair index":
-//   - interleaved: pair i maps to q[2*i], q[2*i+1]
-//   - non-interleaved: pair i maps to q[i] and q[i+embed_dim]
-// embed_dim is "half dimension" (d).
-template <typename scalar_t, bool interleaved, bool aligned_qk>
-__device__ __forceinline__ RotaryVecData<scalar_t, interleaved> load_rotary_vec(
-    const scalar_t* __restrict__ arr,
-    const scalar_t* __restrict__ cos_ptr,
-    const scalar_t* __restrict__ sin_ptr,
-    int rot_offset,
-    int embed_dim) {
-  RotaryVecData<scalar_t, interleaved> data;
-
-  constexpr int kVecBytes = 16;
-  constexpr int kElePerVec = kVecBytes / sizeof(scalar_t);
-
-  if constexpr (interleaved) {
-    const int base = rot_offset * 2;
-
-    if constexpr (aligned_qk) {
-      data.vec = *reinterpret_cast<const float4*>(arr + base);
-    } else {
-      union VecU {
-        float4 v;
-        scalar_t e[kElePerVec];
-      } tmp;
-#pragma unroll
-      for (int i = 0; i < kElePerVec; ++i)
-        tmp.e[i] = arr[base + i];
-      data.vec = tmp.v;
-    }
-
-    data.cos_vec = *reinterpret_cast<const float2*>(cos_ptr + rot_offset);
-    data.sin_vec = *reinterpret_cast<const float2*>(sin_ptr + rot_offset);
-
-  } else {
-    const int bx = rot_offset;
-    const int by = rot_offset + embed_dim;
-
-    if constexpr (aligned_qk) {
-      data.vec_x = *reinterpret_cast<const float4*>(arr + bx);
-      data.vec_y = *reinterpret_cast<const float4*>(arr + by);
-    } else {
-      union VecU {
-        float4 v;
-        scalar_t e[kElePerVec];
-      } tx, ty;
-#pragma unroll
-      for (int i = 0; i < kElePerVec; ++i) {
-        tx.e[i] = arr[bx + i];
-        ty.e[i] = arr[by + i];
-      }
-      data.vec_x = tx.v;
-      data.vec_y = ty.v;
-    }
-
-    data.cos_x = *reinterpret_cast<const float4*>(cos_ptr + bx);
-    data.sin_x = *reinterpret_cast<const float4*>(sin_ptr + bx);
-    data.cos_y = *reinterpret_cast<const float4*>(cos_ptr + by);
-    data.sin_y = *reinterpret_cast<const float4*>(sin_ptr + by);
-  }
-
-  return data;
-}
-
-// Compute RoPE on the loaded vector tile, then store back to q/k.
-template <typename scalar_t, bool interleaved, bool aligned_qk>
-__device__ __forceinline__ void compute_store_rotary_vec(
-    scalar_t* __restrict__ arr, const RotaryVecData<scalar_t, interleaved>& data, int rot_offset, int embed_dim) {
-  constexpr int kVecBytes = 16;
-  constexpr int kElePerVec = kVecBytes / sizeof(scalar_t);
-
-  if constexpr (interleaved) {
-    union VecU {
-      float4 v;
-      scalar_t e[kElePerVec];
-    } v;
-    v.v = data.vec;
-
-    union CSU {
-      float2 v;
-      scalar_t e[kElePerVec / 2];
-    } c, s;
-    c.v = data.cos_vec;
-    s.v = data.sin_vec;
-
-#pragma unroll
-    for (int i = 0; i < kElePerVec; i += 2) {
-      int idx = i / 2;
-      float cos_val = static_cast<float>(c.e[idx]);
-      float sin_val = static_cast<float>(s.e[idx]);
-      float x = static_cast<float>(v.e[i]);
-      float y = static_cast<float>(v.e[i + 1]);
-      v.e[i] = static_cast<scalar_t>(x * cos_val - y * sin_val);
-      v.e[i + 1] = static_cast<scalar_t>(y * cos_val + x * sin_val);
-    }
-
-    const int base = rot_offset * 2;
-    if constexpr (aligned_qk) {
-      *reinterpret_cast<float4*>(arr + base) = v.v;
-    } else {
-#pragma unroll
-      for (int i = 0; i < kElePerVec; ++i)
-        arr[base + i] = v.e[i];
-    }
-
-  } else {
-    union VecU {
-      float4 v;
-      scalar_t e[kElePerVec];
-    } vx, vy;
-    vx.v = data.vec_x;
-    vy.v = data.vec_y;
-
-    union CSU {
-      float4 v;
-      scalar_t e[kElePerVec];
-    } cx, sx, cy, sy;
-    cx.v = data.cos_x;
-    sx.v = data.sin_x;
-    cy.v = data.cos_y;
-    sy.v = data.sin_y;
-
-#pragma unroll
-    for (int i = 0; i < kElePerVec; ++i) {
-      float cos_x = static_cast<float>(cx.e[i]);
-      float sin_x = static_cast<float>(sx.e[i]);
-      float cos_y = static_cast<float>(cy.e[i]);
-      float sin_y = static_cast<float>(sy.e[i]);
-
-      float x = static_cast<float>(vx.e[i]);
-      float y = static_cast<float>(vy.e[i]);
-
-      vx.e[i] = static_cast<scalar_t>(x * cos_x - y * sin_x);
-      vy.e[i] = static_cast<scalar_t>(y * cos_y + x * sin_y);
-    }
-
-    const int bx = rot_offset;
-    const int by = rot_offset + embed_dim;
-
-    if constexpr (aligned_qk) {
-      *reinterpret_cast<float4*>(arr + bx) = vx.v;
-      *reinterpret_cast<float4*>(arr + by) = vy.v;
-    } else {
-#pragma unroll
-      for (int i = 0; i < kElePerVec; ++i) {
-        arr[bx + i] = vx.e[i];
-        arr[by + i] = vy.e[i];
-      }
-    }
-  }
-}
-
-// Scalar fallback: one (x,y) pair per thread-iteration.
-template <typename scalar_t, bool interleaved>
-inline __device__ void apply_token_rotary_embedding(
-    scalar_t* __restrict__ arr,            // [head_size]
-    const scalar_t* __restrict__ cos_ptr,  // [rot_dim]
-    const scalar_t* __restrict__ sin_ptr,  // [rot_dim]
-    int rot_offset,
-    int embed_dim) {            // for non-interleaved: half dim
-  if constexpr (interleaved) {  // NeoX-style: interleaved layout [x0, y0, x1, y1, ...].
-    const int x_index = 2 * rot_offset;
-    const int y_index = x_index + 1;
-
-    // Directly access Global Memory
-    const float cos_val = static_cast<float>(cos_ptr[rot_offset]);
-    const float sin_val = static_cast<float>(sin_ptr[rot_offset]);
-
-    const float x = static_cast<float>(arr[x_index]);
-    const float y = static_cast<float>(arr[y_index]);
-    arr[x_index] = static_cast<scalar_t>(x * cos_val - y * sin_val);
-    arr[y_index] = static_cast<scalar_t>(y * cos_val + x * sin_val);
-  } else {  // GPT-J / LLaMA style: layout [x0, x1, ..., y0, y1, ...]
-    const int x_index = rot_offset;
-    const int y_index = rot_offset + embed_dim;
-
-    // Directly access Global Memory
-    const float cos_val_x = static_cast<float>(cos_ptr[rot_offset]);
-    const float sin_val_x = static_cast<float>(sin_ptr[rot_offset]);
-    const float cos_val_y = static_cast<float>(cos_ptr[rot_offset + embed_dim]);
-    const float sin_val_y = static_cast<float>(sin_ptr[rot_offset + embed_dim]);
-
-    const float x = static_cast<float>(arr[x_index]);
-    const float y = static_cast<float>(arr[y_index]);
-    arr[x_index] = static_cast<scalar_t>(x * cos_val_x - y * sin_val_x);
-    arr[y_index] = static_cast<scalar_t>(y * cos_val_y + x * sin_val_y);
-  }
-}
-
-template <typename scalar_t, bool interleaved>
-inline __device__ void apply_token_rotary_embedding_vec(
-    scalar_t* __restrict__ arr,            // [head_size]
-    const scalar_t* __restrict__ cos_ptr,  // [rot_dim]
-    const scalar_t* __restrict__ sin_ptr,  // [rot_dim]
-    int rot_offset,
-    int embed_dim) {
-  using vec_t = float4;
-  constexpr int kVecBytes = sizeof(vec_t);
-  constexpr int kScalarBytes = sizeof(scalar_t);
-  constexpr int kElePerVec = kVecBytes / kScalarBytes;
-
-  // Union for type punning to avoid strict aliasing issues with reinterpret_cast
-  union VecUnion {
-    vec_t vec;
-    scalar_t elems[kElePerVec];
-  };
-
-  if constexpr (interleaved) {
-    VecUnion data;
-    data.vec = *reinterpret_cast<const vec_t*>(arr + rot_offset * 2);
-
-    // Vectorized load for cos/sin: always 8 bytes (half of vector size, float2)
-    float2 cos_vec = *reinterpret_cast<const float2*>(cos_ptr + rot_offset);
-    float2 sin_vec = *reinterpret_cast<const float2*>(sin_ptr + rot_offset);
-
-    union CosSinVec {
-      float2 vec;
-      scalar_t elems[kElePerVec / 2];
-    };
-    CosSinVec cos_u, sin_u;
-    cos_u.vec = cos_vec;
-    sin_u.vec = sin_vec;
-
-#pragma unroll
-    for (int i = 0; i < kElePerVec; i += 2) {
-      int idx = i / 2;
-      float cos_val = static_cast<float>(cos_u.elems[idx]);
-      float sin_val = static_cast<float>(sin_u.elems[idx]);
-
-      float x = static_cast<float>(data.elems[i]);
-      float y = static_cast<float>(data.elems[i + 1]);
-
-      data.elems[i] = static_cast<scalar_t>(x * cos_val - y * sin_val);
-      data.elems[i + 1] = static_cast<scalar_t>(y * cos_val + x * sin_val);
-    }
-
-    *reinterpret_cast<vec_t*>(arr + rot_offset * 2) = data.vec;
-
-  } else {  // Non-interleaved
-    VecUnion data_x, data_y;
-    data_x.vec = *reinterpret_cast<const vec_t*>(arr + rot_offset);
-    data_y.vec = *reinterpret_cast<const vec_t*>(arr + rot_offset + embed_dim);
-
-    // Vectorized load for cos/sin: always 16 bytes (full vector size, float4)
-    float4 cos_vec_x = *reinterpret_cast<const float4*>(cos_ptr + rot_offset);
-    float4 sin_vec_x = *reinterpret_cast<const float4*>(sin_ptr + rot_offset);
-    float4 cos_vec_y = *reinterpret_cast<const float4*>(cos_ptr + rot_offset + embed_dim);
-    float4 sin_vec_y = *reinterpret_cast<const float4*>(sin_ptr + rot_offset + embed_dim);
-
-    union CosSinVec {
-      float4 vec;
-      scalar_t elems[kElePerVec];
-    };
-    CosSinVec cos_u_x, sin_u_x, cos_u_y, sin_u_y;
-    cos_u_x.vec = cos_vec_x;
-    sin_u_x.vec = sin_vec_x;
-    cos_u_y.vec = cos_vec_y;
-    sin_u_y.vec = sin_vec_y;
-
-#pragma unroll
-    for (int i = 0; i < kElePerVec; ++i) {
-      float cos_val_x = static_cast<float>(cos_u_x.elems[i]);
-      float sin_val_x = static_cast<float>(sin_u_x.elems[i]);
-      float cos_val_y = static_cast<float>(cos_u_y.elems[i]);
-      float sin_val_y = static_cast<float>(sin_u_y.elems[i]);
-
-      float x = static_cast<float>(data_x.elems[i]);
-      float y = static_cast<float>(data_y.elems[i]);
-
-      data_x.elems[i] = static_cast<scalar_t>(x * cos_val_x - y * sin_val_x);
-      data_y.elems[i] = static_cast<scalar_t>(y * cos_val_y + x * sin_val_y);
-    }
-
-    *reinterpret_cast<vec_t*>(arr + rot_offset) = data_x.vec;
-    *reinterpret_cast<vec_t*>(arr + rot_offset + embed_dim) = data_y.vec;
-  }
-}
-
-template <>
-inline __device__ void apply_token_rotary_embedding<float, true>(
-    float* __restrict__ arr,
-    const float* __restrict__ cos_ptr,
-    const float* __restrict__ sin_ptr,
-    int rot_offset,
-    int /*embed_dim*/) {
-  float2 xy = *reinterpret_cast<const float2*>(arr + rot_offset * 2);
-
-  const float cos_val = static_cast<float>(cos_ptr[rot_offset]);
-  const float sin_val = static_cast<float>(sin_ptr[rot_offset]);
-
-  float2 out;
-  out.x = xy.x * cos_val - xy.y * sin_val;
-  out.y = xy.y * cos_val + xy.x * sin_val;
-
-  *reinterpret_cast<float2*>(arr + rot_offset * 2) = out;
-}
-
-// 2D grid kernel: blockIdx.x = token, blockIdx.y = tile within token
-// Each thread processes "pairs_per_step" pairs per iteration when vectorized.
-template <typename scalar_t, bool interleaved, bool vectorized, bool aligned_qk, int ROT_EMBED_DIM>
-__global__ void rotary_embedding_kernel_2d(
-    const scalar_t* __restrict__ cos_data,  // [num_tokens, rot_dim_arg]
-    const scalar_t* __restrict__ sin_data,  // [num_tokens, rot_dim_arg]
-    scalar_t* __restrict__ query_total,
-    scalar_t* __restrict__ key_total,
-    const int rot_dim_arg,
-    const int embed_dim_for_rotation_arg,
-    const int64_t query_token_stride,
-    const int64_t key_token_stride,
-    const int64_t head_stride_query,
-    const int64_t head_stride_key,
-    const int num_heads,
-    const int num_kv_heads,
-    const int head_size_arg,
-    const int blocks_per_token) {
-  const int token_idx = blockIdx.x;
-  if (token_idx >= gridDim.x) {
-    return;
-  }
-
-  // Pointers to Global Memory for the current token
-  const scalar_t* current_token_cos_ptr = cos_data + token_idx * rot_dim_arg;
-  const scalar_t* current_token_sin_ptr = sin_data + token_idx * rot_dim_arg;
-
-  constexpr int kVecBytes = 16;
-  const int scalar_size = sizeof(scalar_t);
-
-  scalar_t* query_for_token = query_total + token_idx * (int)query_token_stride;
-  scalar_t* key_for_token = (key_total != nullptr) ? (key_total + token_idx * (int)key_token_stride) : nullptr;
-
-  const int local_block_idx = blockIdx.y;
-
-  // Otherwise fallback to runtime argument
-  const int embed_dim_for_rotation = (ROT_EMBED_DIM > 0) ? ROT_EMBED_DIM : embed_dim_for_rotation_arg;
-
-  if constexpr (vectorized) {
-    using vec_t = float4;
-    constexpr int kElePerVec = kVecBytes / sizeof(scalar_t);
-    constexpr int pairs_per_step = interleaved ? (kElePerVec / 2) : kElePerVec;
-
-    const int pair_stride = blockDim.x * blocks_per_token * pairs_per_step;
-    int i = (local_block_idx * blockDim.x + threadIdx.x) * pairs_per_step;
-    const int nq_pairs = num_heads * embed_dim_for_rotation;
-
-    // PROLOGUE
-    RotaryVecData<scalar_t, interleaved> curr_data;
-    if (i < nq_pairs) {
-      int head_idx = i / embed_dim_for_rotation;
-      int rot_offset = i % embed_dim_for_rotation;
-      scalar_t* ptr = query_for_token + head_idx * (int)head_stride_query;
-      curr_data = load_rotary_vec<scalar_t, interleaved, aligned_qk>(
-          ptr, current_token_cos_ptr, current_token_sin_ptr, rot_offset, embed_dim_for_rotation);
-    }
-
-    int next_i = i + pair_stride;
-    // MAIN LOOP
-    for (; i < nq_pairs; i += pair_stride, next_i += pair_stride) {
-      // PREFETCH NEXT
-      RotaryVecData<scalar_t, interleaved> next_data;
-      bool active_next = (next_i < nq_pairs);
-
-      if (active_next) {
-        int head_idx_next = next_i / embed_dim_for_rotation;
-        int rot_offset_next = next_i % embed_dim_for_rotation;
-        scalar_t* ptr_next = query_for_token + head_idx_next * (int)head_stride_query;
-        next_data = load_rotary_vec<scalar_t, interleaved, aligned_qk>(
-            ptr_next, current_token_cos_ptr, current_token_sin_ptr, rot_offset_next, embed_dim_for_rotation);
-      }
-      // COMPUTE CURRENT
-      int head_idx = i / embed_dim_for_rotation;
-      int rot_offset = i % embed_dim_for_rotation;
-      scalar_t* ptr = query_for_token + head_idx * (int)head_stride_query;
-      compute_store_rotary_vec<scalar_t, interleaved, aligned_qk>(ptr, curr_data, rot_offset, embed_dim_for_rotation);
-      // SHIFT
-      if (active_next) curr_data = next_data;
-    }
-
-    if (key_for_token != nullptr) {
-      const int nk_pairs = num_kv_heads * embed_dim_for_rotation;
-      int k_i = (local_block_idx * blockDim.x + threadIdx.x) * pairs_per_step;
-
-      // PROLOGUE
-      RotaryVecData<scalar_t, interleaved> curr_data_k;
-      if (k_i < nk_pairs) {
-        int head_idx = k_i / embed_dim_for_rotation;
-        int rot_offset = k_i % embed_dim_for_rotation;
-        scalar_t* ptr = key_for_token + head_idx * (int)head_stride_key;
-        curr_data_k = load_rotary_vec<scalar_t, interleaved, aligned_qk>(
-            ptr, current_token_cos_ptr, current_token_sin_ptr, rot_offset, embed_dim_for_rotation);
-      }
-
-      int next_k_i = k_i + pair_stride;
-
-      // MAIN LOOP
-      for (; k_i < nk_pairs; k_i += pair_stride, next_k_i += pair_stride) {
-        // PREFETCH NEXT
-        RotaryVecData<scalar_t, interleaved> next_data_k;
-        bool active_next = (next_k_i < nk_pairs);
-        if (active_next) {
-          int head_idx_next = next_k_i / embed_dim_for_rotation;
-          int rot_offset_next = next_k_i % embed_dim_for_rotation;
-          scalar_t* ptr_next = key_for_token + head_idx_next * (int)head_stride_key;
-          next_data_k = load_rotary_vec<scalar_t, interleaved, aligned_qk>(
-              ptr_next, current_token_cos_ptr, current_token_sin_ptr, rot_offset_next, embed_dim_for_rotation);
-        }
-
-        // COMPUTE CURRENT
-        int head_idx = k_i / embed_dim_for_rotation;
-        int rot_offset = k_i % embed_dim_for_rotation;
-        scalar_t* ptr = key_for_token + head_idx * (int)head_stride_key;
-        compute_store_rotary_vec<scalar_t, interleaved, aligned_qk>(
-            ptr, curr_data_k, rot_offset, embed_dim_for_rotation);
-
-        // SHIFT
-        curr_data_k = next_data_k;
-      }
-    }
-  } else {
-    // Fallback to scalar implementation
-    const int pair_stride = blockDim.x * blocks_per_token;
-    const int thread_pair_offset = local_block_idx * blockDim.x + threadIdx.x;
-
-    const int nq_pairs = num_heads * embed_dim_for_rotation;
-    for (int i = thread_pair_offset; i < nq_pairs; i += pair_stride) {
-      const int head_idx = i / embed_dim_for_rotation;
-      const int rot_offset = i % embed_dim_for_rotation;
-
-      scalar_t* query_for_token_head = query_for_token + head_idx * (int)head_stride_query;
-      apply_token_rotary_embedding<scalar_t, interleaved>(
-          query_for_token_head, current_token_cos_ptr, current_token_sin_ptr, rot_offset, embed_dim_for_rotation);
-    }
-
-    if (key_for_token != nullptr) {
-      const int nk_pairs = num_kv_heads * embed_dim_for_rotation;
-      for (int i = thread_pair_offset; i < nk_pairs; i += pair_stride) {
-        const int head_idx = i / embed_dim_for_rotation;
-        const int rot_offset = i % embed_dim_for_rotation;
-
-        scalar_t* key_for_token_head = key_for_token + head_idx * (int)head_stride_key;
-        apply_token_rotary_embedding<scalar_t, interleaved>(
-            key_for_token_head, current_token_cos_ptr, current_token_sin_ptr, rot_offset, embed_dim_for_rotation);
-      }
-    }
-  }
-}
-
-// 1D grid kernel: one block per token (blockIdx.x = token)
-template <typename scalar_t, bool interleaved, bool vectorized, bool aligned_qk, int ROT_EMBED_DIM>
-__launch_bounds__(512) __global__ void rotary_embedding_kernel_1d(
-    const scalar_t* __restrict__ cos_data,  // [num_tokens, rot_dim_arg]
-    const scalar_t* __restrict__ sin_data,  // [num_tokens, rot_dim_arg]
-    scalar_t* __restrict__ query_total,
-    scalar_t* __restrict__ key_total,
-    const int rot_dim_arg,
-    const int embed_dim_for_rotation_arg,
-    const int64_t query_token_stride,
-    const int64_t key_token_stride,
-    const int64_t head_stride_query,
-    const int64_t head_stride_key,
-    const int num_heads,
-    const int num_kv_heads,
-    const int head_size_arg) {
-  const int token_idx = blockIdx.x;
-  if (token_idx >= gridDim.x) return;
-
-  const scalar_t* current_token_cos_ptr = cos_data + token_idx * rot_dim_arg;
-  const scalar_t* current_token_sin_ptr = sin_data + token_idx * rot_dim_arg;
-
-  constexpr int kVecBytes = 16;
-  const int scalar_size = sizeof(scalar_t);
-
-  scalar_t* query_for_token = query_total + token_idx * (int)query_token_stride;
-  scalar_t* key_for_token = (key_total != nullptr) ? (key_total + token_idx * (int)key_token_stride) : nullptr;
-
-  // Otherwise fallback to runtime argument
-  const int embed_dim_for_rotation = (ROT_EMBED_DIM > 0) ? ROT_EMBED_DIM : embed_dim_for_rotation_arg;
-
-  if constexpr (vectorized) {
-    constexpr int kVecBytes = 16;
-    constexpr int kElePerVec = kVecBytes / sizeof(scalar_t);
-    // Interleaved: 1 vec = kElePerVec elements = kElePerVec/2 pairs
-    // Non-Interleaved: 1 vec (X) + 1 vec (Y) = kElePerVec elements = kElePerVec pairs
-    constexpr int pairs_per_step = interleaved ? (kElePerVec / 2) : kElePerVec;
-    const int nq_pairs = num_heads * embed_dim_for_rotation;
-
-    // Original Stride
-    const int stride = blockDim.x * pairs_per_step;
-
-    int i = threadIdx.x * pairs_per_step;
-
-    // PROLOGUE
-    RotaryVecData<scalar_t, interleaved> curr_data;
-
-    if (i < nq_pairs) {
-      int head_idx = i / embed_dim_for_rotation;
-      int rot_offset = i % embed_dim_for_rotation;
-      scalar_t* ptr = query_for_token + head_idx * (int)head_stride_query;
-      curr_data = load_rotary_vec<scalar_t, interleaved, aligned_qk>(
-          ptr, current_token_cos_ptr, current_token_sin_ptr, rot_offset, embed_dim_for_rotation);
-    }
-
-    int next_i = i + stride;
-
-    // MAIN LOOP
-    for (; i < nq_pairs; i += stride, next_i += stride) {
-      // PREFETCH NEXT
-      RotaryVecData<scalar_t, interleaved> next_data;
-
-      if (next_i < nq_pairs) {
-        int head_idx = next_i / embed_dim_for_rotation;
-        int rot_offset = next_i % embed_dim_for_rotation;
-        scalar_t* ptr = query_for_token + head_idx * (int)head_stride_query;
-        next_data = load_rotary_vec<scalar_t, interleaved, aligned_qk>(
-            ptr, current_token_cos_ptr, current_token_sin_ptr, rot_offset, embed_dim_for_rotation);
-      }
-
-      // COMPUTE CURRENT
-      if (i < nq_pairs) {
-        int head_idx = i / embed_dim_for_rotation;
-        int rot_offset = i % embed_dim_for_rotation;
-        scalar_t* ptr = query_for_token + head_idx * (int)head_stride_query;
-        compute_store_rotary_vec<scalar_t, interleaved, aligned_qk>(ptr, curr_data, rot_offset, embed_dim_for_rotation);
-      }
-
-      // SHIFT
-      curr_data = next_data;
-    }
-
-    if (key_for_token != nullptr) {
-      const int nk_pairs = num_kv_heads * embed_dim_for_rotation;
-      int k_i = threadIdx.x * pairs_per_step;
-
-      // PROLOGUE
-      RotaryVecData<scalar_t, interleaved> curr_data_k;
-
-      if (k_i < nk_pairs) {
-        int head_idx = k_i / embed_dim_for_rotation;
-        int rot_offset = k_i % embed_dim_for_rotation;
-        scalar_t* ptr = key_for_token + head_idx * (int)head_stride_key;
-        curr_data_k = load_rotary_vec<scalar_t, interleaved, aligned_qk>(
-            ptr, current_token_cos_ptr, current_token_sin_ptr, rot_offset, embed_dim_for_rotation);
-      }
-
-      int next_k_i = k_i + stride;
-
-      // MAIN LOOP
-      for (; k_i < nk_pairs; k_i += stride, next_k_i += stride) {
-        // PREFETCH NEXT
-        RotaryVecData<scalar_t, interleaved> next_data_k;
-
-        if (next_k_i < nk_pairs) {
-          int head_idx = next_k_i / embed_dim_for_rotation;
-          int rot_offset = next_k_i % embed_dim_for_rotation;
-          scalar_t* ptr = key_for_token + head_idx * (int)head_stride_key;
-          next_data_k = load_rotary_vec<scalar_t, interleaved, aligned_qk>(
-              ptr, current_token_cos_ptr, current_token_sin_ptr, rot_offset, embed_dim_for_rotation);
-        }
-
-        // COMPUTE CURRENT
-        if (k_i < nk_pairs) {
-          int head_idx = k_i / embed_dim_for_rotation;
-          int rot_offset = k_i % embed_dim_for_rotation;
-          scalar_t* ptr = key_for_token + head_idx * (int)head_stride_key;
-          compute_store_rotary_vec<scalar_t, interleaved, aligned_qk>(
-              ptr, curr_data_k, rot_offset, embed_dim_for_rotation);
-        }
-
-        // SHIFT
-        curr_data_k = next_data_k;
-      }
-    }
-  } else {
-    // Fallback to scalar implementation
-    const int nq_pairs = num_heads * embed_dim_for_rotation;
-    for (int i = threadIdx.x; i < nq_pairs; i += blockDim.x) {
-      const int head_idx = i / embed_dim_for_rotation;
-      const int rot_offset = i % embed_dim_for_rotation;
-      scalar_t* query_for_token_head = query_for_token + head_idx * (int)head_stride_query;
-
-      apply_token_rotary_embedding<scalar_t, interleaved>(
-          query_for_token_head, current_token_cos_ptr, current_token_sin_ptr, rot_offset, embed_dim_for_rotation);
-    }
-
-    if (key_for_token != nullptr) {
-      const int nk_pairs = num_kv_heads * embed_dim_for_rotation;
-      for (int i = threadIdx.x; i < nk_pairs; i += blockDim.x) {
-        const int head_idx = i / embed_dim_for_rotation;
-        const int rot_offset = i % embed_dim_for_rotation;
-
-        scalar_t* key_for_token_head = key_for_token + head_idx * (int)head_stride_key;
-        apply_token_rotary_embedding<scalar_t, interleaved>(
-            key_for_token_head, current_token_cos_ptr, current_token_sin_ptr, rot_offset, embed_dim_for_rotation);
-      }
-    }
-  }
-}
-
-void rotary_embedding_cos_sin(
-    at::Tensor& cos,
-    at::Tensor& sin,
-    at::Tensor& query,
-    const std::optional<at::Tensor>& key,
-    int64_t head_size,
-    bool interleaved) {
-  TORCH_CHECK(
-      query.dim() == 2 || query.dim() == 3,
-      "query must be in shape [num_tokens, hidden_size] or [num_tokens, num_heads, head_size]");
-  if (key.has_value()) {
-    TORCH_CHECK(
-        key->dim() == 2 || key->dim() == 3,
-        "key must be in shape [num_tokens, hidden_size] or [num_tokens, num_kv_heads, head_size]");
-  }
-
-  const int64_t num_tokens = query.size(0);
-
-  TORCH_CHECK(cos.dim() == 2, "cos must be in shape [num_tokens, D_cos]");
-  TORCH_CHECK(sin.dim() == 2, "sin must be in shape [num_tokens, D_sin]");
-  TORCH_CHECK(cos.size(0) == num_tokens, "cos num_tokens mismatch with query");
-  TORCH_CHECK(sin.size(0) == num_tokens, "sin num_tokens mismatch with query");
-  TORCH_CHECK(cos.size(1) == sin.size(1), "cos and sin D_cos/D_sin mismatch");
-
-  TORCH_CHECK(cos.scalar_type() == query.scalar_type(), "cos dtype mismatch with query");
-  TORCH_CHECK(sin.scalar_type() == query.scalar_type(), "sin dtype mismatch with query");
-  TORCH_CHECK(cos.is_cuda() && sin.is_cuda() && query.is_cuda(), "cos/sin/query must be CUDA tensors");
-  TORCH_CHECK(query.is_contiguous(), "query must be contiguous");
-  if (key.has_value()) {
-    TORCH_CHECK(key->is_cuda(), "key must be CUDA tensor if provided");
-    TORCH_CHECK(key->scalar_type() == query.scalar_type(), "key dtype mismatch with query");
-    TORCH_CHECK(key->is_contiguous(), "key must be contiguous");
-  }
-
-  int query_hidden_size_calculated;
-  if (query.dim() == 2) {
-    query_hidden_size_calculated = (int)query.size(1);
-  } else {
-    query_hidden_size_calculated = (int)query.size(1) * (int)query.size(2);
-    TORCH_CHECK(query.size(2) == head_size, "query head_size mismatch in 3D tensor");
-  }
-  TORCH_CHECK(query_hidden_size_calculated % head_size == 0, "query_hidden_size not divisible by head_size");
-  int num_heads = query_hidden_size_calculated / (int)head_size;
-
-  int key_hidden_size_calculated = 0;
-  int num_kv_heads = num_heads;
-  if (key.has_value()) {
-    TORCH_CHECK((int)key->size(0) == num_tokens, "key num_tokens mismatch");
-    if (key->dim() == 2) {
-      key_hidden_size_calculated = (int)key->size(1);
-    } else {
-      key_hidden_size_calculated = (int)key->size(1) * (int)key->size(2);
-      TORCH_CHECK((int)key->size(2) == head_size, "key head_size mismatch in 3D tensor");
-    }
-    TORCH_CHECK(key_hidden_size_calculated % head_size == 0, "key_hidden_size not divisible by head_size");
-    num_kv_heads = key_hidden_size_calculated / (int)head_size;
-  }
-  TORCH_CHECK(num_heads % num_kv_heads == 0, "num_heads must be divisible by num_kv_heads");
-
-  // NeoX interleaved: if cos/sin are full-dim (head_size), downsample once.
-  if (interleaved && cos.size(1) == head_size) {
-    TORCH_CHECK(head_size % 2 == 0, "interleaved layout requires even head_size");
-    const int64_t half = head_size / 2;
-    std::vector<int64_t> new_shape = {cos.size(0), half, 2};
-    cos = cos.view(new_shape).select(2, 0).contiguous();
-    sin = sin.view(new_shape).select(2, 1).contiguous();
-  }
-
-  const int rot_dim_from_cache = (int)cos.size(1);
-  const int64_t query_token_stride = query_hidden_size_calculated;
-  const int64_t key_token_stride = key.has_value() ? key_hidden_size_calculated : 0;
-
-  int64_t head_stride_query;
-  if (query.dim() == 3 && query.size(1) == num_heads && query.size(2) == head_size) {
-    head_stride_query = query.stride(1);
-  } else {
-    head_stride_query = head_size;
-  }
-
-  int64_t head_stride_key = head_size;
-  if (key.has_value()) {
-    if (key->dim() == 3 && key->size(1) == num_kv_heads && key->size(2) == head_size) {
-      head_stride_key = key->stride(1);
-    } else {
-      head_stride_key = head_size;
-    }
-  }
-
-  const int embed_dim_for_rotation = interleaved ? rot_dim_from_cache : (rot_dim_from_cache / 2);
-  TORCH_CHECK(embed_dim_for_rotation > 0, "embed_dim_for_rotation must be > 0");
-
-  const int max_pairs_to_rotate_per_token =
-      std::max(num_heads * embed_dim_for_rotation, num_kv_heads * embed_dim_for_rotation);
-
-  const at::cuda::OptionalCUDAGuard device_guard(device_of(query));
-  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
-
-  AT_DISPATCH_FLOATING_TYPES_AND2(
-      at::ScalarType::Half, at::ScalarType::BFloat16, query.scalar_type(), "rotary_embedding_cos_sin", [&] {
-        using torch_scalar_t = scalar_t;
-        using cuda_scalar_t = typename std::conditional<
-            std::is_same<torch_scalar_t, at::Half>::value,
-            nv_half,
-            typename std::conditional<std::is_same<torch_scalar_t, at::BFloat16>::value, nv_bfloat16, torch_scalar_t>::
-                type>::type;
-
-        // Constants for vectorization
-        constexpr int kVecBytes = 16;
-        constexpr int kElePerVec = kVecBytes / sizeof(torch_scalar_t);
-
-        // Determine how many pairs one thread handles in one vector step
-        const int pairs_per_step = interleaved ? (kElePerVec / 2) : kElePerVec;
-
-        // Logic Split: Decouple "Vectorized Compute" from "Aligned Memory Access"
-        // Check if we can compute in vectors (Dimensions & Cos/Sin alignment)
-        bool can_vec_compute = true;
-        if (embed_dim_for_rotation % pairs_per_step != 0) can_vec_compute = false;
-
-        // Check Cos/Sin Alignment (Base ptr and Row Stride)
-        if (reinterpret_cast<uintptr_t>(cos.data_ptr()) % kVecBytes != 0) can_vec_compute = false;
-        if (reinterpret_cast<uintptr_t>(sin.data_ptr()) % kVecBytes != 0) can_vec_compute = false;
-        // The stride of the cache usually equals rot_dim. Ensure row starts are aligned.
-        if ((rot_dim_from_cache * sizeof(torch_scalar_t)) % kVecBytes != 0) can_vec_compute = false;
-
-        // Check if Q/K are 16-byte aligned (affects load/store instructions only)
-        bool qk_aligned16 = true;
-
-        if (reinterpret_cast<uintptr_t>(query.data_ptr()) % kVecBytes != 0) qk_aligned16 = false;
-        if ((query_token_stride * sizeof(torch_scalar_t)) % kVecBytes != 0) qk_aligned16 = false;
-        if ((head_stride_query * sizeof(torch_scalar_t)) % kVecBytes != 0) qk_aligned16 = false;
-
-        if (key.has_value()) {
-          if (reinterpret_cast<uintptr_t>(key->data_ptr()) % kVecBytes != 0) qk_aligned16 = false;
-          if ((key_token_stride * sizeof(torch_scalar_t)) % kVecBytes != 0) qk_aligned16 = false;
-          if ((head_stride_key * sizeof(torch_scalar_t)) % kVecBytes != 0) qk_aligned16 = false;
-        }
-
-        // Launch Configuration
-        const bool use_vec = can_vec_compute;
-        const int launch_pairs_per_thread = use_vec ? pairs_per_step : 1;
-        const int total_threads_needed =
-            (max_pairs_to_rotate_per_token + launch_pairs_per_thread - 1) / launch_pairs_per_thread;
-        auto round_up32 = [](int x) { return ((x + 31) / 32) * 32; };
-
-        // Case 1: 2D Grid
-        const int threads_per_block_2d =
-            std::min<int>(512, std::max(128, round_up32(std::min(total_threads_needed, 512))));
-        const int blocks_per_token_2d = (total_threads_needed + threads_per_block_2d - 1) / threads_per_block_2d;
-        const bool use_grid_2d = (num_tokens <= 4) && (blocks_per_token_2d > 1);
-        // Case 2: 1D Grid
-        const int threads_per_block_1d = std::min<int>(512, std::max(128, round_up32(total_threads_needed)));
-        // Final Config
-        const int threads_per_block = use_grid_2d ? threads_per_block_2d : threads_per_block_1d;
-        const int blocks_per_token = use_grid_2d ? blocks_per_token_2d : 1;
-
-        dim3 block(threads_per_block);
-        dim3 grid_2d((int)num_tokens, std::max(1, blocks_per_token));
-        dim3 grid_1d((int)num_tokens);
-        size_t smem_size = 0;
-
-        // Kernel Dispatch
-        _DIFFUSION_ROT_EMBED_DIM(
-            embed_dim_for_rotation,
-            auto launch_kernel =
-                [&](bool vec_compute, bool aligned_qk) {
-                  if (interleaved) {
-                    if (use_grid_2d) {
-                      if (vec_compute) {
-                        if (aligned_qk) {
-                          rotary_embedding_kernel_2d<cuda_scalar_t, true, true, true, ROT_EMBED_DIM>
-                              <<<grid_2d, block, smem_size, stream>>>(
-                                  reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
-                                  reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
-                                  reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
-                                  key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>())
-                                                  : nullptr,
-                                  rot_dim_from_cache,
-                                  embed_dim_for_rotation,
-                                  query_token_stride,
-                                  key_token_stride,
-                                  head_stride_query,
-                                  head_stride_key,
-                                  num_heads,
-                                  num_kv_heads,
-                                  (int)head_size,
-                                  blocks_per_token);
-                        } else {
-                          // Vectorized Compute + Unaligned Load/Store
-                          rotary_embedding_kernel_2d<cuda_scalar_t, true, true, false, ROT_EMBED_DIM>
-                              <<<grid_2d, block, smem_size, stream>>>(
-                                  reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
-                                  reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
-                                  reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
-                                  key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>())
-                                                  : nullptr,
-                                  rot_dim_from_cache,
-                                  embed_dim_for_rotation,
-                                  query_token_stride,
-                                  key_token_stride,
-                                  head_stride_query,
-                                  head_stride_key,
-                                  num_heads,
-                                  num_kv_heads,
-                                  (int)head_size,
-                                  blocks_per_token);
-                        }
-                      } else {
-                        // Scalar fallback (aligned_qk param ignored or set to true/false arbitrarily)
-                        rotary_embedding_kernel_2d<cuda_scalar_t, true, false, true, ROT_EMBED_DIM>
-                            <<<grid_2d, block, smem_size, stream>>>(
-                                reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
-                                reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
-                                reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
-                                key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>())
-                                                : nullptr,
-                                rot_dim_from_cache,
-                                embed_dim_for_rotation,
-                                query_token_stride,
-                                key_token_stride,
-                                head_stride_query,
-                                head_stride_key,
-                                num_heads,
-                                num_kv_heads,
-                                (int)head_size,
-                                blocks_per_token);
-                      }
-                    } else {  // 1D Grid
-                      if (vec_compute) {
-                        if (aligned_qk) {
-                          rotary_embedding_kernel_1d<cuda_scalar_t, true, true, true, ROT_EMBED_DIM>
-                              <<<grid_1d, block, smem_size, stream>>>(
-                                  reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
-                                  reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
-                                  reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
-                                  key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>())
-                                                  : nullptr,
-                                  rot_dim_from_cache,
-                                  embed_dim_for_rotation,
-                                  query_token_stride,
-                                  key_token_stride,
-                                  head_stride_query,
-                                  head_stride_key,
-                                  num_heads,
-                                  num_kv_heads,
-                                  (int)head_size);
-                        } else {
-                          rotary_embedding_kernel_1d<cuda_scalar_t, true, true, false, ROT_EMBED_DIM>
-                              <<<grid_1d, block, smem_size, stream>>>(
-                                  reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
-                                  reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
-                                  reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
-                                  key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>())
-                                                  : nullptr,
-                                  rot_dim_from_cache,
-                                  embed_dim_for_rotation,
-                                  query_token_stride,
-                                  key_token_stride,
-                                  head_stride_query,
-                                  head_stride_key,
-                                  num_heads,
-                                  num_kv_heads,
-                                  (int)head_size);
-                        }
-                      } else {
-                        rotary_embedding_kernel_1d<cuda_scalar_t, true, false, true, ROT_EMBED_DIM>
-                            <<<grid_1d, block, smem_size, stream>>>(
-                                reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
-                                reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
-                                reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
-                                key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>())
-                                                : nullptr,
-                                rot_dim_from_cache,
-                                embed_dim_for_rotation,
-                                query_token_stride,
-                                key_token_stride,
-                                head_stride_query,
-                                head_stride_key,
-                                num_heads,
-                                num_kv_heads,
-                                (int)head_size);
-                      }
-                    }
-                  } else {  // Non-Interleaved
-                    if (use_grid_2d) {
-                      if (vec_compute) {
-                        if (aligned_qk) {
-                          rotary_embedding_kernel_2d<cuda_scalar_t, false, true, true, ROT_EMBED_DIM>
-                              <<<grid_2d, block, smem_size, stream>>>(
-                                  reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
-                                  reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
-                                  reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
-                                  key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>())
-                                                  : nullptr,
-                                  rot_dim_from_cache,
-                                  embed_dim_for_rotation,
-                                  query_token_stride,
-                                  key_token_stride,
-                                  head_stride_query,
-                                  head_stride_key,
-                                  num_heads,
-                                  num_kv_heads,
-                                  (int)head_size,
-                                  blocks_per_token);
-                        } else {
-                          rotary_embedding_kernel_2d<cuda_scalar_t, false, true, false, ROT_EMBED_DIM>
-                              <<<grid_2d, block, smem_size, stream>>>(
-                                  reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
-                                  reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
-                                  reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
-                                  key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>())
-                                                  : nullptr,
-                                  rot_dim_from_cache,
-                                  embed_dim_for_rotation,
-                                  query_token_stride,
-                                  key_token_stride,
-                                  head_stride_query,
-                                  head_stride_key,
-                                  num_heads,
-                                  num_kv_heads,
-                                  (int)head_size,
-                                  blocks_per_token);
-                        }
-                      } else {
-                        rotary_embedding_kernel_2d<cuda_scalar_t, false, false, true, ROT_EMBED_DIM>
-                            <<<grid_2d, block, smem_size, stream>>>(
-                                reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
-                                reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
-                                reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
-                                key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>())
-                                                : nullptr,
-                                rot_dim_from_cache,
-                                embed_dim_for_rotation,
-                                query_token_stride,
-                                key_token_stride,
-                                head_stride_query,
-                                head_stride_key,
-                                num_heads,
-                                num_kv_heads,
-                                (int)head_size,
-                                blocks_per_token);
-                      }
-                    } else {  // 1D Grid
-                      if (vec_compute) {
-                        if (aligned_qk) {
-                          rotary_embedding_kernel_1d<cuda_scalar_t, false, true, true, ROT_EMBED_DIM>
-                              <<<grid_1d, block, smem_size, stream>>>(
-                                  reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
-                                  reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
-                                  reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
-                                  key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>())
-                                                  : nullptr,
-                                  rot_dim_from_cache,
-                                  embed_dim_for_rotation,
-                                  query_token_stride,
-                                  key_token_stride,
-                                  head_stride_query,
-                                  head_stride_key,
-                                  num_heads,
-                                  num_kv_heads,
-                                  (int)head_size);
-                        } else {
-                          rotary_embedding_kernel_1d<cuda_scalar_t, false, true, false, ROT_EMBED_DIM>
-                              <<<grid_1d, block, smem_size, stream>>>(
-                                  reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
-                                  reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
-                                  reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
-                                  key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>())
-                                                  : nullptr,
-                                  rot_dim_from_cache,
-                                  embed_dim_for_rotation,
-                                  query_token_stride,
-                                  key_token_stride,
-                                  head_stride_query,
-                                  head_stride_key,
-                                  num_heads,
-                                  num_kv_heads,
-                                  (int)head_size);
-                        }
-                      } else {
-                        rotary_embedding_kernel_1d<cuda_scalar_t, false, false, true, ROT_EMBED_DIM>
-                            <<<grid_1d, block, smem_size, stream>>>(
-                                reinterpret_cast<const cuda_scalar_t*>(cos.data_ptr<torch_scalar_t>()),
-                                reinterpret_cast<const cuda_scalar_t*>(sin.data_ptr<torch_scalar_t>()),
-                                reinterpret_cast<cuda_scalar_t*>(query.data_ptr<torch_scalar_t>()),
-                                key.has_value() ? reinterpret_cast<cuda_scalar_t*>(key->data_ptr<torch_scalar_t>())
-                                                : nullptr,
-                                rot_dim_from_cache,
-                                embed_dim_for_rotation,
-                                query_token_stride,
-                                key_token_stride,
-                                head_stride_query,
-                                head_stride_key,
-                                num_heads,
-                                num_kv_heads,
-                                (int)head_size);
-                      }
-                    }
-                  }
-                };
-
-            // Execute the launch
-            launch_kernel(use_vec, qk_aligned16););
-      });
-
-  C10_CUDA_KERNEL_LAUNCH_CHECK();
-}
diff --git a/sgl-kernel/include/sgl_kernel_ops.h b/sgl-kernel/include/sgl_kernel_ops.h
index 5fba50fe0545..5e3cf24f9036 100644
--- a/sgl-kernel/include/sgl_kernel_ops.h
+++ b/sgl-kernel/include/sgl_kernel_ops.h
@@ -149,7 +149,6 @@ void apply_rope_pos_ids_cos_sin_cache(
     const std::optional<at::Tensor>& v_buffer,
     const std::optional<at::Tensor>& kv_cache_loc);
 
-// RoPE with position ids + cos/sin cache (used by SRT)
 void rotary_embedding(
     torch::Tensor& positions,
     torch::Tensor& query,
@@ -158,15 +157,6 @@ void rotary_embedding(
     torch::Tensor& cos_sin_cache,
     bool is_neox);
 
-// RoPE with explicit cos/sin tensors (used by diffusion / Qwen image)
-void rotary_embedding_cos_sin(
-    at::Tensor& cos,
-    at::Tensor& sin,
-    at::Tensor& query,
-    const std::optional<at::Tensor>& key,
-    int64_t head_size,
-    bool interleaved);
-
 void downcast_fp8(
     at::Tensor& k,
     at::Tensor& v,
@@ -414,17 +404,6 @@ void fused_qk_norm_rope(
     double attention_factor,
     int64_t rotary_dim);
 
-/*
- * From csrc/multimodal/rotary_embedding
- */
-void rotary_embedding(
-    at::Tensor& cos,
-    at::Tensor& sin,
-    at::Tensor& query,
-    const std::optional<at::Tensor>& key,
-    int64_t head_size,
-    bool is_neox);
-
 void cutlass_fp4_group_mm(
     torch::Tensor& output,
     const torch::Tensor& a,
diff --git a/sgl-kernel/include/utils.h b/sgl-kernel/include/utils.h
index be5ba0726780..031cded62596 100644
--- a/sgl-kernel/include/utils.h
+++ b/sgl-kernel/include/utils.h
@@ -464,56 +464,3 @@ inline uint32_t next_pow2(uint32_t x) noexcept {
 #else
 #define SGLANG_LDG(arg) *(arg)
 #endif
-
-// Define dispatch macro for Rotation Dimension (Pair Count)
-// Case 32 -> head_size 64
-// Case 64 -> head_size 128
-// Case 128 -> head_size 256
-#define _DIFFUSION_ROT_EMBED_DIM(embed_dim, ...)  \
-  switch (embed_dim) {                    \
-    case 16: {                            \
-      constexpr int ROT_EMBED_DIM = 16;   \
-      __VA_ARGS__;                        \
-      break;                              \
-    }                                     \
-    case 32: {                            \
-      constexpr int ROT_EMBED_DIM = 32;   \
-      __VA_ARGS__;                        \
-      break;                              \
-    }                                     \
-    case 64: {                            \
-      constexpr int ROT_EMBED_DIM = 64;   \
-      __VA_ARGS__;                        \
-      break;                              \
-    }                                     \
-    case 96: {                            \
-      constexpr int ROT_EMBED_DIM = 96;   \
-      __VA_ARGS__;                        \
-      break;                              \
-    }                                     \
-    case 128: {                           \
-      constexpr int ROT_EMBED_DIM = 128;  \
-      __VA_ARGS__;                        \
-      break;                              \
-    }                                     \
-    case 256: {                           \
-      constexpr int ROT_EMBED_DIM = 256;  \
-      __VA_ARGS__;                        \
-      break;                              \
-    }                                     \
-    case 512: {                           \
-      constexpr int ROT_EMBED_DIM = 512;  \
-      __VA_ARGS__;                        \
-      break;                              \
-    }                                     \
-    case 1024: {                          \
-      constexpr int ROT_EMBED_DIM = 1024; \
-      __VA_ARGS__;                        \
-      break;                              \
-    }                                     \
-    default: {                            \
-      constexpr int ROT_EMBED_DIM = 0;    \
-      __VA_ARGS__;                        \
-      break;                              \
-    }                                     \
-  }
diff --git a/sgl-kernel/python/sgl_kernel/__init__.py b/sgl-kernel/python/sgl_kernel/__init__.py
index d4968b501870..1b97ef94f02b 100644
--- a/sgl-kernel/python/sgl_kernel/__init__.py
+++ b/sgl-kernel/python/sgl_kernel/__init__.py
@@ -108,7 +108,6 @@
     ggml_mul_mat_a8,
     ggml_mul_mat_vec_a8,
 )
-from sgl_kernel.rotary_embedding import rotary_embedding_cos_sin
 from sgl_kernel.sampling import (
     min_p_sampling_from_probs,
     top_k_mask_logits,
diff --git a/sgl-kernel/python/sgl_kernel/rotary_embedding.py b/sgl-kernel/python/sgl_kernel/rotary_embedding.py
deleted file mode 100644
index 307c6b6b82ce..000000000000
--- a/sgl-kernel/python/sgl_kernel/rotary_embedding.py
+++ /dev/null
@@ -1,132 +0,0 @@
-from __future__ import annotations
-
-from typing import Literal, Optional
-
-import torch
-
-_HAS_COS_SIN = hasattr(torch.ops.sgl_kernel, "rotary_embedding_cos_sin")
-_HAS_POSITION = hasattr(torch.ops.sgl_kernel, "rotary_embedding")
-
-
-def _resolve_interleaved(interleaved: Optional[bool], is_neox: Optional[bool]) -> bool:
-    if interleaved is not None and is_neox is not None and interleaved != is_neox:
-        raise ValueError(
-            f"is_neox({is_neox}) and interleaved({interleaved}) mismatch; keep only one or make them equal."
-        )
-    if is_neox is not None:
-        return is_neox
-    if interleaved is not None:
-        return interleaved
-    # Historical default in this wrapper: assume NeoX-style if not specified.
-    return True
-
-
-def apply_rotary_embedding(
-    *,
-    mode: Literal["cos_sin", "positions"],
-    query: torch.Tensor,
-    key: Optional[torch.Tensor] = None,
-    head_size: int = 0,
-    interleaved: Optional[bool] = None,
-    is_neox: Optional[bool] = None,
-    # cos/sin mode
-    cos: Optional[torch.Tensor] = None,
-    sin: Optional[torch.Tensor] = None,
-    # positions mode
-    positions: Optional[torch.Tensor] = None,
-    cos_sin_cache: Optional[torch.Tensor] = None,
-) -> None:
-    effective_interleaved = _resolve_interleaved(interleaved, is_neox)
-
-    if mode == "cos_sin":
-        if cos is None or sin is None:
-            raise ValueError("mode='cos_sin' requires cos and sin")
-
-        if _HAS_COS_SIN:
-            torch.ops.sgl_kernel.rotary_embedding_cos_sin(
-                cos,
-                sin,
-                query,
-                key if key is not None else None,
-                head_size,
-                effective_interleaved,
-            )
-            return
-
-        if _HAS_POSITION:
-            torch.ops.sgl_kernel.rotary_embedding(
-                cos,
-                sin,
-                query,
-                key if key is not None else None,
-                head_size,
-                effective_interleaved,
-            )
-            return
-
-        raise RuntimeError(
-            "No cos/sin rotary embedding kernel is available in torch.ops.sgl_kernel"
-        )
-
-    if mode == "positions":
-        if positions is None or cos_sin_cache is None:
-            raise ValueError("mode='positions' requires positions and cos_sin_cache")
-
-        if not _HAS_POSITION:
-            raise RuntimeError(
-                "No positions rotary embedding kernel is available in torch.ops.sgl_kernel"
-            )
-
-        # Use positional args for maximum compatibility across builds.
-        torch.ops.sgl_kernel.rotary_embedding(
-            positions,
-            query,
-            key if key is not None else None,
-            head_size,
-            cos_sin_cache,
-            effective_interleaved,
-        )
-        return
-
-def rotary_embedding_cos_sin(
-    cos: torch.Tensor,
-    sin: torch.Tensor,
-    query: torch.Tensor,
-    key: Optional[torch.Tensor] = None,
-    head_size: int = 0,
-    interleaved: Optional[bool] = None,
-    *,
-    is_neox: Optional[bool] = None,
-) -> None:
-    apply_rotary_embedding(
-        mode="cos_sin",
-        cos=cos,
-        sin=sin,
-        query=query,
-        key=key,
-        head_size=head_size,
-        interleaved=interleaved,
-        is_neox=is_neox,
-    )
-
-
-def rotary_embedding(
-    positions: torch.Tensor,
-    query: torch.Tensor,
-    key: Optional[torch.Tensor] = None,
-    head_size: int = 0,
-    interleaved: Optional[bool] = None,
-    *,
-    is_neox: Optional[bool] = None,
-    cos_sin_cache: torch.Tensor,
-) -> None:
-    apply_rotary_embedding(
-        mode="positions",
-        positions=positions,
-        cos_sin_cache=cos_sin_cache,
-        query=query,
-        key=key,
-        head_size=head_size,
-        interleaved=interleaved,
-        is_neox=is_neox,
-    )
diff --git a/sgl-kernel/tests/test_mm_rotary_embedding.py b/sgl-kernel/tests/test_mm_rotary_embedding.py
deleted file mode 100755
index 4fc454fa9b44..000000000000
--- a/sgl-kernel/tests/test_mm_rotary_embedding.py
+++ /dev/null
@@ -1,198 +0,0 @@
-#!/usr/bin/env python
-# -*- coding: utf-8 -*-
-
-import os
-import sys
-from dataclasses import dataclass
-from typing import List, Tuple
-
-import torch
-
-try:
-    import pytest
-
-    HAS_PYTEST = True
-except ImportError:
-    HAS_PYTEST = False
-
-try:
-    # Cos/sin-based rotary kernel from sgl_diffusion
-    from sgl_kernel.rotary_embedding import (
-        rotary_embedding_cos_sin as rotary_emb_module,
-    )
-
-    HAS_ROTARY_EMBEDDING = True
-except ImportError:
-    rotary_emb_module = None
-    HAS_ROTARY_EMBEDDING = False
-    print("sgl_kernel.rotary_embedding_cos_sin not available")
-
-
-@dataclass
-class RotaryTestResult:
-    name: str
-    passed: bool
-    q_diff: float = 0.0
-    k_diff: float = 0.0
-    details: str = ""
-
-
-def compute_cos_sin_cache(
-    max_seq_len: int,
-    rotary_dim: int,
-    base: float = 10000.0,
-    dtype: torch.dtype = torch.float32,
-) -> Tuple[torch.Tensor, torch.Tensor]:
-    """
-    Compute separate cos and sin caches.
-    """
-    inv_freq = 1.0 / (
-        base ** (torch.arange(0, rotary_dim, 2, dtype=torch.float32) / rotary_dim)
-    )
-    t = torch.arange(max_seq_len, dtype=torch.float32)
-    freqs = torch.einsum("i,j->ij", t, inv_freq)
-    cos = freqs.cos().to(dtype)
-    sin = freqs.sin().to(dtype)
-    return cos, sin
-
-
-def reference_rotary_neox(
-    query: torch.Tensor,  # (num_tokens, num_heads * head_size)
-    key: torch.Tensor,  # (num_tokens, num_kv_heads * head_size)
-    cos: torch.Tensor,  # (num_tokens, rotary_dim / 2)
-    sin: torch.Tensor,  # (num_tokens, rotary_dim / 2)
-    head_size: int,
-) -> Tuple[torch.Tensor, torch.Tensor]:
-    """
-    Reference implementation of NeoX-style rotary embedding,
-    using explicit cos/sin caches and Q/K in flattened (num_tokens, heads * dim) layout.
-    """
-    num_tokens = query.size(0)
-    num_heads = query.size(1) // head_size
-    num_kv_heads = key.size(1) // head_size
-    rotary_dim = cos.size(1) * 2
-
-    q = query.view(num_tokens, num_heads, head_size).float()
-    k = key.view(num_tokens, num_kv_heads, head_size).float()
-
-    q_rot = q[..., :rotary_dim]
-    q_pass = q[..., rotary_dim:]
-    k_rot = k[..., :rotary_dim]
-    k_pass = k[..., rotary_dim:]
-
-    q_rot = q_rot.view(num_tokens, num_heads, rotary_dim // 2, 2)
-    k_rot = k_rot.view(num_tokens, num_kv_heads, rotary_dim // 2, 2)
-
-    cos_expanded = cos.float().unsqueeze(1)  # (tokens, 1, rotary_dim/2)
-    sin_expanded = sin.float().unsqueeze(1)
-
-    # Apply rotary: x' = x*cos - y*sin, y' = y*cos + x*sin
-    q_x = q_rot[..., 0]
-    q_y = q_rot[..., 1]
-    q_rot_out = torch.stack(
-        [
-            q_x * cos_expanded - q_y * sin_expanded,
-            q_y * cos_expanded + q_x * sin_expanded,
-        ],
-        dim=-1,
-    )
-
-    k_x = k_rot[..., 0]
-    k_y = k_rot[..., 1]
-    k_rot_out = torch.stack(
-        [
-            k_x * cos_expanded - k_y * sin_expanded,
-            k_y * cos_expanded + k_x * sin_expanded,
-        ],
-        dim=-1,
-    )
-
-    q_rot_out = q_rot_out.view(num_tokens, num_heads, rotary_dim)
-    k_rot_out = k_rot_out.view(num_tokens, num_kv_heads, rotary_dim)
-
-    q_out = torch.cat([q_rot_out, q_pass], dim=-1)
-    k_out = torch.cat([k_rot_out, k_pass], dim=-1)
-
-    q_out = q_out.view(num_tokens, num_heads * head_size).to(query.dtype)
-    k_out = k_out.view(num_tokens, num_kv_heads * head_size).to(key.dtype)
-
-    return q_out, k_out
-
-
-def _check_rotary_correctness(
-    batch_size: int = 2,
-    seq_len: int = 128,
-    num_heads: int = 32,
-    num_kv_heads: int = 8,
-    head_size: int = 128,
-    dtype: torch.dtype = torch.bfloat16,
-    tol: float = 1e-2,
-    device: str = "cuda",
-) -> RotaryTestResult:
-    if not HAS_ROTARY_EMBEDDING:
-        return RotaryTestResult(
-            "basic (skipped: rotary_embedding not built)",
-            True,
-            details="rotary_embedding extension not found",
-        )
-
-    rotary_dim = head_size
-    max_seq_len = 8192
-    num_tokens = batch_size * seq_len
-
-    query = torch.randn(num_tokens, num_heads * head_size, dtype=dtype, device=device)
-    key = torch.randn(num_tokens, num_kv_heads * head_size, dtype=dtype, device=device)
-    cos_cache, sin_cache = compute_cos_sin_cache(max_seq_len, rotary_dim, dtype=dtype)
-    cos_cache = cos_cache.to(device)
-    sin_cache = sin_cache.to(device)
-    positions = torch.arange(seq_len, device=device).repeat(batch_size)
-    cos = cos_cache[positions]
-    sin = sin_cache[positions]
-
-    q_ref, k_ref = reference_rotary_neox(
-        query.clone(), key.clone(), cos, sin, head_size
-    )
-
-    q_out = query.clone()
-    k_out = key.clone()
-    rotary_emb_module(cos, sin, q_out, k_out, head_size, True)  # is_neox=True
-
-    q_diff = (q_out - q_ref).abs().max().item()
-    k_diff = (k_out - k_ref).abs().max().item()
-    passed = q_diff < tol and k_diff < tol
-
-    name = f"basic [bs={batch_size}, seq={seq_len}, heads={num_heads}/{num_kv_heads}, dim={head_size}, {dtype}]"
-    return RotaryTestResult(name, passed, q_diff, k_diff)
-
-
-@pytest.mark.parametrize("batch_size", [2, 32, 1])
-@pytest.mark.parametrize("seq_len", [1, 128, 512, 2048])
-@pytest.mark.parametrize("num_heads, num_kv_heads", [(32, 8), (64, 8), (8, 8), (32, 1)])
-@pytest.mark.parametrize("head_size", [64, 128, 256, 80, 320])
-@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
-def test_rotary_embedding_correctness(
-    batch_size: int,
-    seq_len: int,
-    num_heads: int,
-    num_kv_heads: int,
-    head_size: int,
-    dtype: torch.dtype,
-) -> None:
-    if not HAS_ROTARY_EMBEDDING:
-        pytest.skip("sgl_kernel.rotary_embedding not available")
-
-    result = _check_rotary_correctness(
-        batch_size=batch_size,
-        seq_len=seq_len,
-        num_heads=num_heads,
-        num_kv_heads=num_kv_heads,
-        head_size=head_size,
-        dtype=dtype,
-    )
-    assert (
-        result.passed
-    ), f"{result.name} failed: Q={result.q_diff:.2e}, K={result.k_diff:.2e}"
-
-
-if __name__ == "__main__":
-    pytest.main([__file__, "-v"])

From b76e0fad62195ab2c5c99321243871313773f75b Mon Sep 17 00:00:00 2001
From: RubiaCx <1084281732@qq.com>
Date: Mon, 29 Dec 2025 23:44:54 -0800
Subject: [PATCH 21/37] Align with JIT kernels guide; keep behavior consistent
 with the AOT implementation.

---
 .../jit_kernel/benchmark/bench_rotary_embedding.py |  5 +++--
 python/sglang/jit_kernel/rotary_embedding.py       | 12 +++++++++---
 .../jit_kernel/tests/test_rotary_embedding.py      |  1 +
 .../runtime/layers/rotary_embedding.py             | 14 ++++++++++----
 sgl-kernel/csrc/common_extension.cc                |  6 +++---
 sgl-kernel/include/utils.h                         |  3 +++
 6 files changed, 29 insertions(+), 12 deletions(-)

diff --git a/python/sglang/jit_kernel/benchmark/bench_rotary_embedding.py b/python/sglang/jit_kernel/benchmark/bench_rotary_embedding.py
index 820a47c20b57..745d136347d8 100644
--- a/python/sglang/jit_kernel/benchmark/bench_rotary_embedding.py
+++ b/python/sglang/jit_kernel/benchmark/bench_rotary_embedding.py
@@ -1,12 +1,12 @@
 import itertools
 import os
 from functools import lru_cache
-from typing import Optional, Tuple
+from typing import Tuple
+
 import torch
 import triton
 import triton.testing
 
-
 IS_CI = (
     os.getenv("CI", "false").lower() == "true"
     or os.getenv("GITHUB_ACTIONS", "false").lower() == "true"
@@ -97,6 +97,7 @@ def _apply(x: torch.Tensor) -> None:
             x1 = xr2[..., 1].clone()
             xr2[..., 0].copy_(x0 * cos_b - x1 * sin_b)
             xr2[..., 1].copy_(x1 * cos_b + x0 * sin_b)
+
     else:
         if cos_f.shape[1] == head_size // 2:
             embed_dim = int(cos_f.shape[1])
diff --git a/python/sglang/jit_kernel/rotary_embedding.py b/python/sglang/jit_kernel/rotary_embedding.py
index 620e61e40d34..8c958f226eaf 100644
--- a/python/sglang/jit_kernel/rotary_embedding.py
+++ b/python/sglang/jit_kernel/rotary_embedding.py
@@ -1,6 +1,9 @@
 from __future__ import annotations
+
 from typing import Optional
+
 import torch
+
 from sglang.jit_kernel.utils import cache_once, load_jit
 
 
@@ -136,7 +139,10 @@ def _as_3d(x: torch.Tensor, h: int) -> torch.Tensor:
 
     module = _jit_rotary_embedding_cos_sin_module()
     if k3 is None:
-        module.rotary_embedding_cos_sin_q(cos_to_use, sin_to_use, q3, head_size, effective_interleaved)
+        module.rotary_embedding_cos_sin_q(
+            cos_to_use, sin_to_use, q3, head_size, effective_interleaved
+        )
     else:
-        module.rotary_embedding_cos_sin_qk(cos_to_use, sin_to_use, q3, k3, head_size, effective_interleaved)
-
+        module.rotary_embedding_cos_sin_qk(
+            cos_to_use, sin_to_use, q3, k3, head_size, effective_interleaved
+        )
diff --git a/python/sglang/jit_kernel/tests/test_rotary_embedding.py b/python/sglang/jit_kernel/tests/test_rotary_embedding.py
index f061e00039eb..02f686095c93 100644
--- a/python/sglang/jit_kernel/tests/test_rotary_embedding.py
+++ b/python/sglang/jit_kernel/tests/test_rotary_embedding.py
@@ -87,6 +87,7 @@ def _apply(x: torch.Tensor) -> None:
             x1 = xr2[..., 1].clone()
             xr2[..., 0].copy_(x0 * cos_b - x1 * sin_b)
             xr2[..., 1].copy_(x1 * cos_b + x0 * sin_b)
+
     else:
         if cos_f.shape[1] == head_size // 2:
             embed_dim = int(cos_f.shape[1])
diff --git a/python/sglang/multimodal_gen/runtime/layers/rotary_embedding.py b/python/sglang/multimodal_gen/runtime/layers/rotary_embedding.py
index f5c6e91da7f3..8bc5e031834f 100644
--- a/python/sglang/multimodal_gen/runtime/layers/rotary_embedding.py
+++ b/python/sglang/multimodal_gen/runtime/layers/rotary_embedding.py
@@ -139,7 +139,9 @@ def apply_sglang_jit_rope_qk_inplace(
             f"got q:{tuple(q.shape)} k:{tuple(k.shape)}"
         )
     if q.shape != k.shape:
-        raise ValueError(f"q and k must have the same shape, got {q.shape} vs {k.shape}")
+        raise ValueError(
+            f"q and k must have the same shape, got {q.shape} vs {k.shape}"
+        )
 
     if not (isinstance(cos_sin_cache, torch.Tensor) and cos_sin_cache.dim() == 2):
         raise ValueError("cos_sin_cache must be a 2D torch.Tensor")
@@ -178,8 +180,10 @@ def apply_sglang_jit_rope_qk_inplace(
     positions = positions.to(q.device, non_blocking=True)
 
     if cos_sin_cache.shape[0] == seqlen:
-        cache_tok = cos_sin_cache[None, :, :].expand(bsz, seqlen, cos_sin_cache.shape[1]).reshape(
-            num_tokens, cos_sin_cache.shape[1]
+        cache_tok = (
+            cos_sin_cache[None, :, :]
+            .expand(bsz, seqlen, cos_sin_cache.shape[1])
+            .reshape(num_tokens, cos_sin_cache.shape[1])
         )
     elif cos_sin_cache.shape[0] == num_tokens:
         cache_tok = cos_sin_cache
@@ -189,7 +193,9 @@ def apply_sglang_jit_rope_qk_inplace(
     q3 = q.reshape(num_tokens, nheads, d).contiguous()
     k3 = k.reshape(num_tokens, nheads, d).contiguous()
 
-    cos_half, sin_half = _split_cos_sin_from_cache(cache_tok.contiguous(), dtype=q.dtype)
+    cos_half, sin_half = _split_cos_sin_from_cache(
+        cache_tok.contiguous(), dtype=q.dtype
+    )
 
     interleaved = not is_neox
     if interleaved:
diff --git a/sgl-kernel/csrc/common_extension.cc b/sgl-kernel/csrc/common_extension.cc
index c82c8d66321f..463b924fa79b 100644
--- a/sgl-kernel/csrc/common_extension.cc
+++ b/sgl-kernel/csrc/common_extension.cc
@@ -91,9 +91,9 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) {
   m.impl("apply_rope_pos_ids_cos_sin_cache", torch::kCUDA, &apply_rope_pos_ids_cos_sin_cache);
 
   m.def(
-    "rotary_embedding(Tensor positions, Tensor! query,"
-    "                 Tensor!? key, int head_size,"
-    "                 Tensor cos_sin_cache, bool is_neox) -> ()");
+      "rotary_embedding(Tensor positions, Tensor! query,"
+      "                 Tensor!? key, int head_size,"
+      "                 Tensor cos_sin_cache, bool is_neox) -> ()");
   m.impl("rotary_embedding", torch::kCUDA, &rotary_embedding);
 
   m.def(
diff --git a/sgl-kernel/include/utils.h b/sgl-kernel/include/utils.h
index 031cded62596..c5a517cea019 100644
--- a/sgl-kernel/include/utils.h
+++ b/sgl-kernel/include/utils.h
@@ -459,6 +459,9 @@ inline uint32_t next_pow2(uint32_t x) noexcept {
   return 1u << (32 - __builtin_clz(x - 1));
 }
 
+/*
+ * LDG Support
+ */
 #ifndef USE_ROCM
 #define SGLANG_LDG(arg) __ldg(arg)
 #else

From c6637394f9fce650806bc9f18ffc22ad254592da Mon Sep 17 00:00:00 2001
From: RubiaCx <1084281732@qq.com>
Date: Tue, 30 Dec 2025 00:07:53 -0800
Subject: [PATCH 22/37] Update rotary embedding test/bench to import AOT RoPE
 kernel

---
 .../benchmark/bench_rotary_embedding.py       | 36 ++++++++----
 .../jit_kernel/tests/test_rotary_embedding.py | 58 ++++++++++---------
 2 files changed, 56 insertions(+), 38 deletions(-)

diff --git a/python/sglang/jit_kernel/benchmark/bench_rotary_embedding.py b/python/sglang/jit_kernel/benchmark/bench_rotary_embedding.py
index 745d136347d8..8b44581647b2 100644
--- a/python/sglang/jit_kernel/benchmark/bench_rotary_embedding.py
+++ b/python/sglang/jit_kernel/benchmark/bench_rotary_embedding.py
@@ -7,6 +7,13 @@
 import triton
 import triton.testing
 
+try:
+    import sgl_kernel  # noqa: F401
+except Exception:
+    sgl_kernel = None  # type: ignore
+
+HAS_SGL_POS = hasattr(torch.ops.sgl_kernel, "rotary_embedding")
+
 IS_CI = (
     os.getenv("CI", "false").lower() == "true"
     or os.getenv("GITHUB_ACTIONS", "false").lower() == "true"
@@ -134,15 +141,15 @@ def sglang_aot_rotary_positions(
     interleaved: bool,
     cos_sin_cache: torch.Tensor,
 ) -> None:
-    from sgl_kernel.rotary_embedding import rotary_embedding
-
-    rotary_embedding(
-        positions=positions,
-        query=q,
-        key=k,
-        head_size=head_size,
-        is_neox=not interleaved,
-        cos_sin_cache=cos_sin_cache,
+    if not HAS_SGL_POS:
+        raise RuntimeError("torch.ops.sgl_kernel.rotary_embedding is not available")
+    torch.ops.sgl_kernel.rotary_embedding(
+        positions,
+        q,
+        k,
+        head_size,
+        cos_sin_cache,
+        not interleaved,  # sgl-kernel positions op uses is_neox (split-halves) flag
     )
 
 
@@ -159,14 +166,19 @@ def sglang_jit_rotary_cos_sin(
     rotary_embedding_cos_sin(cos, sin, q, k, head_size, interleaved)
 
 
-BASE_LINE_VALS = ["jit_cos_sin", "aot_pos"]
-BASE_LINE_NAMES = ["SGL JIT (cos/sin)", "SGL AOT (positions)"]
-BASE_STYLES = [("blue", "-"), ("orange", "--")]
+BASE_LINE_VALS = ["jit_cos_sin"]
+BASE_LINE_NAMES = ["SGL JIT (cos/sin)"]
+BASE_STYLES = [("blue", "-")]
 
 LINE_VALS = list(BASE_LINE_VALS)
 LINE_NAMES = list(BASE_LINE_NAMES)
 STYLES = list(BASE_STYLES)
 
+if HAS_SGL_POS:
+    LINE_VALS.append("aot_pos")
+    LINE_NAMES.append("SGL AOT (positions)")
+    STYLES.append(("orange", "--"))
+
 if HAS_VLLM:
     LINE_VALS.append("vllm_pos")
     LINE_NAMES.append("vLLM (positions)")
diff --git a/python/sglang/jit_kernel/tests/test_rotary_embedding.py b/python/sglang/jit_kernel/tests/test_rotary_embedding.py
index 02f686095c93..15b520301c04 100644
--- a/python/sglang/jit_kernel/tests/test_rotary_embedding.py
+++ b/python/sglang/jit_kernel/tests/test_rotary_embedding.py
@@ -3,6 +3,13 @@
 import torch
 import triton
 
+try:
+    import sgl_kernel  # noqa: F401
+except Exception:
+    sgl_kernel = None  # type: ignore
+
+HAS_SGL_POS = hasattr(torch.ops.sgl_kernel, "rotary_embedding")
+
 
 def _compute_cos_sin_cache_half(
     max_seq_len: int,
@@ -28,18 +35,15 @@ def sglang_aot_rotary_positions(
     interleaved: bool,
     cos_sin_cache: torch.Tensor,
 ) -> None:
-    # NOTE: `sgl_kernel` positions-kernel uses `is_neox` flag internally:
-    # - is_neox=True  => x_index=i, y_index=embed_dim+i   (split halves)
-    # - is_neox=False => x_index=2*i, y_index=2*i+1       (interleaved pairs)
-    from sgl_kernel.rotary_embedding import rotary_embedding
-
-    rotary_embedding(
-        positions=positions,
-        query=q,
-        key=k,
-        head_size=head_size,
-        is_neox=not interleaved,
-        cos_sin_cache=cos_sin_cache,
+    if not HAS_SGL_POS:
+        raise RuntimeError("torch.ops.sgl_kernel.rotary_embedding is not available")
+    torch.ops.sgl_kernel.rotary_embedding(
+        positions,
+        q,
+        k,
+        head_size,
+        cos_sin_cache,
+        not interleaved,  # sgl-kernel positions op uses is_neox (split-halves) flag
     )
 
 
@@ -167,14 +171,15 @@ def main():
 
                     q_k_aot = (q.clone(), k.clone())
                     q_k_jit = (q.clone(), k.clone())
-                    sglang_aot_rotary_positions(
-                        positions,
-                        q_k_aot[0],
-                        q_k_aot[1],
-                        HEAD_SIZE,
-                        INTERLEAVED,
-                        aot_cos_sin_cache,
-                    )
+                    if HAS_SGL_POS:
+                        sglang_aot_rotary_positions(
+                            positions,
+                            q_k_aot[0],
+                            q_k_aot[1],
+                            HEAD_SIZE,
+                            INTERLEAVED,
+                            aot_cos_sin_cache,
+                        )
                     sglang_jit_rotary_cos_sin(
                         cos, sin, q_k_jit[0], q_k_jit[1], HEAD_SIZE, INTERLEAVED
                     )
@@ -182,12 +187,13 @@ def main():
                     ref_atol = 2e-2 if DTYPE == torch.bfloat16 else 2e-3
                     ref_rtol = 2e-2 if DTYPE == torch.bfloat16 else 2e-3
 
-                    triton.testing.assert_close(
-                        q_ref_fp32, q_k_aot[0], atol=ref_atol, rtol=ref_rtol
-                    )
-                    triton.testing.assert_close(
-                        k_ref_fp32, q_k_aot[1], atol=ref_atol, rtol=ref_rtol
-                    )
+                    if HAS_SGL_POS:
+                        triton.testing.assert_close(
+                            q_ref_fp32, q_k_aot[0], atol=ref_atol, rtol=ref_rtol
+                        )
+                        triton.testing.assert_close(
+                            k_ref_fp32, q_k_aot[1], atol=ref_atol, rtol=ref_rtol
+                        )
                     triton.testing.assert_close(
                         q_ref_fp32, q_k_jit[0], atol=ref_atol, rtol=ref_rtol
                     )

From 47e52a0535845446a199e3e5d3a51e2d80f1e86a Mon Sep 17 00:00:00 2001
From: RubiaCx <1084281732@qq.com>
Date: Tue, 30 Dec 2025 05:47:18 -0800
Subject: [PATCH 23/37] fix fallback for qk_norm

---
 .../runtime/layers/layernorm.py               | 24 +++++--
 .../runtime/models/dits/qwen_image.py         | 68 ++++++-------------
 2 files changed, 37 insertions(+), 55 deletions(-)

diff --git a/python/sglang/multimodal_gen/runtime/layers/layernorm.py b/python/sglang/multimodal_gen/runtime/layers/layernorm.py
index 82fbb76828fe..cb66bc877fd0 100644
--- a/python/sglang/multimodal_gen/runtime/layers/layernorm.py
+++ b/python/sglang/multimodal_gen/runtime/layers/layernorm.py
@@ -434,17 +434,19 @@ def apply_qk_norm(
 ) -> Tuple[torch.Tensor, torch.Tensor]:
     """Apply QK normalization for query and key tensors.
 
-    Minimal multimodal_gen-only implementation: only the JIT fused inplace
-    QK-norm kernel path is supported (no fallback).
+    Prefer the fused inplace QK-norm kernel when applicable;
+    otherwise fall back to PyTorch norms.
     """
 
     batch_size = q.size(0)
     q_eps = q_norm.variance_epsilon
     k_eps = k_norm.variance_epsilon
-    # Only try fused path on CUDA and when it won't introduce implicit copies.
     if (
         q.is_cuda
+        and (torch.version.cuda is not None)
         and allow_inplace
+        and q.is_contiguous()
+        and k.is_contiguous()
         and (q_eps == k_eps)
         and can_use_fused_inplace_qknorm(head_dim)
     ):
@@ -458,7 +460,15 @@ def apply_qk_norm(
         )
         return q, k
 
-    raise RuntimeError(
-        "apply_qk_norm: fused inplace QK-norm is not applicable "
-        "(expected CUDA, contiguous q/k, matching eps, and supported head_dim)"
-    )
+    # Fallback: apply the PyTorch norms.
+    q_out = q_norm(q)
+    k_out = k_norm(k)
+
+    if allow_inplace and q.is_contiguous() and q_out.shape == q.shape:
+        q.copy_(q_out)
+        q_out = q
+    if allow_inplace and k.is_contiguous() and k_out.shape == k.shape:
+        k.copy_(k_out)
+        k_out = k
+
+    return q_out, k_out
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 6c7a37b0cbdd..052d7ea97bea 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
+from math import prod
 from typing import Any, Dict, List, Optional, Tuple, Union
 
 import numpy as np
@@ -15,13 +16,8 @@
 from diffusers.models.normalization import AdaLayerNormContinuous
 
 from sglang.multimodal_gen.configs.models.dits.qwenimage import QwenImageDitConfig
-from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
 from sglang.multimodal_gen.runtime.layers.attention import USPAttention
-from sglang.multimodal_gen.runtime.layers.layernorm import (
-    LayerNorm,
-    RMSNorm,
-    apply_qk_norm,
-)
+from sglang.multimodal_gen.runtime.layers.layernorm import LayerNorm, RMSNorm
 from sglang.multimodal_gen.runtime.layers.linear import ReplicatedLinear
 from sglang.multimodal_gen.runtime.layers.rotary_embedding import (
     apply_sglang_jit_rope_qk_inplace,
@@ -551,23 +547,14 @@ def forward(
         txt_value = txt_value.unflatten(-1, (self.num_heads, -1))
 
         # Apply QK normalization
-        if self.qk_norm:
-            img_query, img_key = apply_qk_norm(
-                q=img_query,
-                k=img_key,
-                q_norm=self.norm_q,
-                k_norm=self.norm_k,
-                head_dim=img_query.shape[-1],
-                allow_inplace=True,
-            )
-            txt_query, txt_key = apply_qk_norm(
-                q=txt_query,
-                k=txt_key,
-                q_norm=self.norm_added_q,
-                k_norm=self.norm_added_k,
-                head_dim=txt_query.shape[-1],
-                allow_inplace=True,
-            )
+        if self.norm_q is not None:
+            img_query = self.norm_q(img_query)
+        if self.norm_k is not None:
+            img_key = self.norm_k(img_key)
+        if self.norm_added_q is not None:
+            txt_query = self.norm_added_q(txt_query)
+        if self.norm_added_k is not None:
+            txt_key = self.norm_added_k(txt_key)
 
         # Apply RoPE
         if image_rotary_emb is not None:
@@ -791,12 +778,6 @@ def forward(
         return encoder_hidden_states, hidden_states
 
 
-def to_hashable(obj):
-    if isinstance(obj, list):
-        return tuple(to_hashable(x) for x in obj)
-    return obj
-
-
 class QwenImageTransformer2DModel(CachableDiT):
     """
     The Transformer model introduced in Qwen.
@@ -873,22 +854,6 @@ def __init__(
             self.inner_dim, patch_size * patch_size * self.out_channels, bias=True
         )
 
-        self.timestep_zero = torch.zeros(
-            (1,), dtype=torch.int, device=get_local_torch_device()
-        )
-
-    @functools.lru_cache(maxsize=50)
-    def build_modulate_index(self, img_shapes: tuple[int, int, int], device):
-        modulate_index_list = []
-        for sample in img_shapes:
-            first_size = sample[0][0] * sample[0][1] * sample[0][2]
-            total_size = sum(s[0] * s[1] * s[2] for s in sample)
-            idx = (torch.arange(total_size, device=device) >= first_size).int()
-            modulate_index_list.append(idx)
-
-        modulate_index = torch.stack(modulate_index_list)
-        return modulate_index
-
     def forward(
         self,
         hidden_states: torch.Tensor,
@@ -944,9 +909,16 @@ def forward(
         timestep = (timestep / 1000).to(hidden_states.dtype)
 
         if self.zero_cond_t:
-            timestep = torch.cat([timestep, self.timestep_zero], dim=0)
-            device = timestep.device
-            modulate_index = self.build_modulate_index(to_hashable(img_shapes), device)
+            timestep = torch.cat([timestep, timestep * 0], dim=0)
+            # Use torch operations for GPU efficiency
+            modulate_index = torch.tensor(
+                [
+                    [0] * prod(sample[0]) + [1] * sum([prod(s) for s in sample[1:]])
+                    for sample in img_shapes
+                ],
+                device=timestep.device,
+                dtype=torch.int,
+            )
         else:
             modulate_index = None
 

From fc2448c4bd6136eca3f9e0d34bdaee2bb6bcf0b3 Mon Sep 17 00:00:00 2001
From: RubiaCx <1084281732@qq.com>
Date: Tue, 30 Dec 2025 21:21:27 -0800
Subject: [PATCH 24/37] Relax QK-norm inplace contiguity + avoid fallback copy

---
 .../runtime/layers/layernorm.py               | 28 +++++++++++++------
 1 file changed, 19 insertions(+), 9 deletions(-)

diff --git a/python/sglang/multimodal_gen/runtime/layers/layernorm.py b/python/sglang/multimodal_gen/runtime/layers/layernorm.py
index cb66bc877fd0..eef1b211ee80 100644
--- a/python/sglang/multimodal_gen/runtime/layers/layernorm.py
+++ b/python/sglang/multimodal_gen/runtime/layers/layernorm.py
@@ -441,13 +441,30 @@ def apply_qk_norm(
     batch_size = q.size(0)
     q_eps = q_norm.variance_epsilon
     k_eps = k_norm.variance_epsilon
+
+    def can_view_as_bnhd(x: torch.Tensor) -> bool:
+        """Whether `x` can be viewed as [batch, *, head_dim] without a copy."""
+        if (
+            x.dim() < 2
+            or x.size(0) != batch_size
+            or x.size(-1) != head_dim
+            or x.stride(-1) != 1
+            or x.stride(-2) != head_dim
+        ):
+            return False
+        try:
+            x.view(batch_size, -1, head_dim)
+            return True
+        except RuntimeError:
+            return False
+
     if (
         q.is_cuda
         and (torch.version.cuda is not None)
         and allow_inplace
-        and q.is_contiguous()
-        and k.is_contiguous()
         and (q_eps == k_eps)
+        and can_view_as_bnhd(q)
+        and can_view_as_bnhd(k)
         and can_use_fused_inplace_qknorm(head_dim)
     ):
         fused_inplace_qknorm(
@@ -464,11 +481,4 @@ def apply_qk_norm(
     q_out = q_norm(q)
     k_out = k_norm(k)
 
-    if allow_inplace and q.is_contiguous() and q_out.shape == q.shape:
-        q.copy_(q_out)
-        q_out = q
-    if allow_inplace and k.is_contiguous() and k_out.shape == k.shape:
-        k.copy_(k_out)
-        k_out = k
-
     return q_out, k_out

From d54d9ce3579c6add74b8a042011e475232ea7e6e Mon Sep 17 00:00:00 2001
From: RubiaCx <1084281732@qq.com>
Date: Fri, 2 Jan 2026 21:55:57 -0800
Subject: [PATCH 25/37] Add FlashInfer RoPE to bench and use it as speedup
 baseline

---
 .../benchmark/bench_rotary_embedding.py       | 187 +++++++++++++++---
 1 file changed, 160 insertions(+), 27 deletions(-)

diff --git a/python/sglang/jit_kernel/benchmark/bench_rotary_embedding.py b/python/sglang/jit_kernel/benchmark/bench_rotary_embedding.py
index 8b44581647b2..f852ed44b1e1 100644
--- a/python/sglang/jit_kernel/benchmark/bench_rotary_embedding.py
+++ b/python/sglang/jit_kernel/benchmark/bench_rotary_embedding.py
@@ -19,7 +19,6 @@
     or os.getenv("GITHUB_ACTIONS", "false").lower() == "true"
 )
 
-IS_QUICK = os.getenv("SGL_BENCH_QUICK", "0").lower() in ("1", "true", "yes", "y")
 ONLY_PROVIDER = os.getenv("SGL_BENCH_PROVIDER", "").strip() or None
 SHOW_SPEEDUP = os.getenv("SGL_BENCH_SPEEDUP", "1").lower() in ("1", "true", "yes", "y")
 
@@ -39,7 +38,24 @@
     vLLMRotaryEmbedding = None
     HAS_VLLM = False
 
-if IS_CI or IS_QUICK:
+try:
+    from sglang.multimodal_gen.runtime.layers.rotary_embedding import (
+        apply_flashinfer_rope_qk_inplace,
+    )
+
+    try:
+        from flashinfer.rope import (  # noqa: F401
+            apply_rope_with_cos_sin_cache_inplace as _,
+        )
+
+        HAS_FLASHINFER = True
+    except Exception:
+        HAS_FLASHINFER = False
+except Exception:
+    apply_flashinfer_rope_qk_inplace = None  # type: ignore
+    HAS_FLASHINFER = False
+
+if IS_CI:
     BS_RANGE = [16]
     SEQ_RANGE = [1, 128]
     HEAD_SIZE_RANGE = [128]
@@ -51,6 +67,42 @@
     INTERLEAVED_RANGE = [True, False]
 
 
+def _is_flashinfer_unsupported_shape_error(e: BaseException) -> bool:
+    msg = str(e)
+    return ("Unsupported head_dim" in msg) or ("cos_sin_cache should be float32" in msg)
+
+
+def _bench(
+    fn,
+    *,
+    quantiles: list[float],
+    use_cudagraph: bool,
+) -> tuple[float, float, float]:
+    if use_cudagraph:
+        return triton.testing.do_bench_cudagraph(fn, quantiles=quantiles)  # type: ignore
+    return triton.testing.do_bench(fn, quantiles=quantiles)  # type: ignore
+
+
+def _bench_provider(
+    provider: str,
+    fn,
+    *,
+    quantiles: list[float],
+    use_cudagraph: bool,
+) -> tuple[float, float, float]:
+    """Bench a provider and return (median, min, max) in ms.
+
+    For FlashInfer, gracefully skip unsupported shapes by returning NaNs.
+    """
+    try:
+        return _bench(fn, quantiles=quantiles, use_cudagraph=use_cudagraph)
+    except Exception as e:  # noqa: BLE001
+        if provider == "flashinfer_rope" and _is_flashinfer_unsupported_shape_error(e):
+            nan = float("nan")
+            return nan, nan, nan
+        raise
+
+
 def _compute_cos_sin_cache_half(
     max_seq_len: int,
     rotary_dim: int,
@@ -166,14 +218,28 @@ def sglang_jit_rotary_cos_sin(
     rotary_embedding_cos_sin(cos, sin, q, k, head_size, interleaved)
 
 
-BASE_LINE_VALS = ["jit_cos_sin"]
-BASE_LINE_NAMES = ["SGL JIT (cos/sin)"]
-BASE_STYLES = [("blue", "-")]
+if HAS_FLASHINFER:
+    BASE_PROVIDER = "flashinfer_rope"
+    BASE_NAME = "FlashInfer RoPE"
+    BASE_STYLE = ("red", "-")
+else:
+    BASE_PROVIDER = "jit_cos_sin"
+    BASE_NAME = "SGL JIT (cos/sin)"
+    BASE_STYLE = ("blue", "-")
+
+BASE_LINE_VALS = [BASE_PROVIDER]
+BASE_LINE_NAMES = [BASE_NAME]
+BASE_STYLES = [BASE_STYLE]
 
 LINE_VALS = list(BASE_LINE_VALS)
 LINE_NAMES = list(BASE_LINE_NAMES)
 STYLES = list(BASE_STYLES)
 
+if "jit_cos_sin" not in LINE_VALS:
+    LINE_VALS.append("jit_cos_sin")
+    LINE_NAMES.append("SGL JIT (cos/sin)")
+    STYLES.append(("blue", "-"))
+
 if HAS_SGL_POS:
     LINE_VALS.append("aot_pos")
     LINE_NAMES.append("SGL AOT (positions)")
@@ -184,6 +250,11 @@ def sglang_jit_rotary_cos_sin(
     LINE_NAMES.append("vLLM (positions)")
     STYLES.append(("green", "-."))
 
+if HAS_FLASHINFER and "flashinfer_rope" not in LINE_VALS:
+    LINE_VALS.append("flashinfer_rope")
+    LINE_NAMES.append("FlashInfer RoPE")
+    STYLES.append(("red", "--"))
+
 if ONLY_PROVIDER is not None:
     if ONLY_PROVIDER not in LINE_VALS:
         raise ValueError(
@@ -275,6 +346,16 @@ def benchmark(
 
     cos_cache_half, sin_cache_half = _get_cos_sin_cache_half_cuda(rotary_dim)
     cos_sin_cache_aot = torch.cat([cos_cache_half, sin_cache_half], dim=-1).contiguous()
+    # FlashInfer requires cos_sin_cache to be float32.
+    if HAS_FLASHINFER:
+        cos_f32, sin_f32 = _compute_cos_sin_cache_half(
+            MAX_SEQ_LEN, rotary_dim, dtype=torch.float32
+        )
+        cos_sin_cache_flashinfer = (
+            torch.cat([cos_f32, sin_f32], dim=-1)
+            .to(device=DEVICE, dtype=torch.float32)
+            .contiguous()
+        )
 
     positions = torch.arange(seq_len, device=DEVICE, dtype=torch.int64).repeat(
         batch_size
@@ -288,8 +369,16 @@ def benchmark(
         cos = torch.cat([cos_half, cos_half], dim=-1).contiguous()
         sin = torch.cat([sin_half, sin_half], dim=-1).contiguous()
 
-    q = torch.randn(num_tokens, NUM_Q_HEADS, head_size, device=DEVICE, dtype=DTYPE)
-    k = torch.randn(num_tokens, NUM_KV_HEADS, head_size, device=DEVICE, dtype=DTYPE)
+    # Keep a canonical 4D layout for providers that require it (e.g. FlashInfer).
+    # View as 3D for providers that accept [num_tokens, nheads, head_size].
+    q4 = torch.randn(
+        batch_size, seq_len, NUM_Q_HEADS, head_size, device=DEVICE, dtype=DTYPE
+    ).contiguous()
+    k4 = torch.randn(
+        batch_size, seq_len, NUM_Q_HEADS, head_size, device=DEVICE, dtype=DTYPE
+    ).contiguous()
+    q = q4.view(num_tokens, NUM_Q_HEADS, head_size)
+    k = k4.view(num_tokens, NUM_Q_HEADS, head_size)
 
     global _SANITY_DONE
     if (
@@ -325,17 +414,28 @@ def benchmark(
             cos, sin, q, k, head_size, interleaved
         ),
     }
+    if HAS_FLASHINFER and apply_flashinfer_rope_qk_inplace is not None:
+        FN_MAP["flashinfer_rope"] = lambda: apply_flashinfer_rope_qk_inplace(
+            q4,
+            k4,
+            cos_sin_cache_flashinfer,
+            is_neox=not interleaved,
+        )
     if HAS_VLLM:
         vllm_rope = _get_vllm_rope(head_size, rotary_dim, interleaved, DTYPE)
         FN_MAP["vllm_pos"] = lambda: vllm_rope.forward_cuda(positions, q, k)
 
     fn = FN_MAP[provider]
     quantiles = [0.5, 0.2, 0.8]
-    ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(fn, quantiles=quantiles)  # type: ignore
+    # If FlashInfer is available, unify timing across providers by using do_bench for all.
+    use_cudagraph = not HAS_FLASHINFER
+    ms, min_ms, max_ms = _bench_provider(
+        provider, fn, quantiles=quantiles, use_cudagraph=use_cudagraph
+    )
     return 1000 * ms, 1000 * max_ms, 1000 * min_ms
 
 
-SPEEDUP_LINE_VALS = [p for p in LINE_VALS if p != "jit_cos_sin"]
+SPEEDUP_LINE_VALS = [p for p in LINE_VALS if p != BASE_PROVIDER]
 SPEEDUP_LINE_NAMES = [LINE_NAMES[LINE_VALS.index(p)] for p in SPEEDUP_LINE_VALS]
 SPEEDUP_STYLES = [STYLES[LINE_VALS.index(p)] for p in SPEEDUP_LINE_VALS]
 
@@ -348,8 +448,8 @@ def benchmark(
         line_vals=SPEEDUP_LINE_VALS,
         line_names=SPEEDUP_LINE_NAMES,
         styles=SPEEDUP_STYLES,
-        ylabel="speedup (x) vs SGL JIT (cos/sin)",
-        plot_name="rotary-embedding-speedup-vs-jit",
+        ylabel=f"speedup (x) vs {BASE_NAME}",
+        plot_name=f"rotary-embedding-speedup-vs-{BASE_PROVIDER}",
         args={},
     )
 )
@@ -360,7 +460,7 @@ def benchmark_speedup_vs_jit(
     interleaved: bool,
     provider: str,
 ) -> Tuple[float, float, float]:
-    if provider == "jit_cos_sin":
+    if provider == BASE_PROVIDER:
         return 1.0, 1.0, 1.0
     if DEVICE == "cuda" and not torch.cuda.is_available():
         raise RuntimeError("CUDA not available")
@@ -370,6 +470,15 @@ def benchmark_speedup_vs_jit(
 
     cos_cache_half, sin_cache_half = _get_cos_sin_cache_half_cuda(rotary_dim)
     cos_sin_cache_aot = torch.cat([cos_cache_half, sin_cache_half], dim=-1).contiguous()
+    if HAS_FLASHINFER:
+        cos_f32, sin_f32 = _compute_cos_sin_cache_half(
+            MAX_SEQ_LEN, rotary_dim, dtype=torch.float32
+        )
+        cos_sin_cache_flashinfer = (
+            torch.cat([cos_f32, sin_f32], dim=-1)
+            .to(device=DEVICE, dtype=torch.float32)
+            .contiguous()
+        )
 
     positions = torch.arange(seq_len, device=DEVICE, dtype=torch.int64).repeat(
         batch_size
@@ -383,10 +492,14 @@ def benchmark_speedup_vs_jit(
         cos = torch.cat([cos_half, cos_half], dim=-1).contiguous()
         sin = torch.cat([sin_half, sin_half], dim=-1).contiguous()
 
-    q_base = torch.randn(num_tokens, NUM_Q_HEADS, head_size, device=DEVICE, dtype=DTYPE)
-    k_base = torch.randn(
-        num_tokens, NUM_KV_HEADS, head_size, device=DEVICE, dtype=DTYPE
-    )
+    q_base4 = torch.randn(
+        batch_size, seq_len, NUM_Q_HEADS, head_size, device=DEVICE, dtype=DTYPE
+    ).contiguous()
+    k_base4 = torch.randn(
+        batch_size, seq_len, NUM_Q_HEADS, head_size, device=DEVICE, dtype=DTYPE
+    ).contiguous()
+    q_base = q_base4.view(num_tokens, NUM_Q_HEADS, head_size)
+    k_base = k_base4.view(num_tokens, NUM_Q_HEADS, head_size)
 
     q_jit = q_base.clone()
     k_jit = k_base.clone()
@@ -404,21 +517,41 @@ def benchmark_speedup_vs_jit(
             cos, sin, q_p, k_p, head_size, interleaved
         ),
     }
+    if HAS_FLASHINFER and apply_flashinfer_rope_qk_inplace is not None:
+        q_f = q_base4.clone()
+        k_f = k_base4.clone()
+        FN_MAP["flashinfer_rope"] = lambda: apply_flashinfer_rope_qk_inplace(
+            q_f,
+            k_f,
+            cos_sin_cache_flashinfer,
+            is_neox=not interleaved,
+        )
     if HAS_VLLM:
         vllm_rope = _get_vllm_rope(head_size, rotary_dim, interleaved, DTYPE)
         FN_MAP["vllm_pos"] = lambda: vllm_rope.forward_cuda(positions, q_p, k_p)
 
     quantiles = [0.5, 0.2, 0.8]
-    jit_ms, jit_min_ms, jit_max_ms = triton.testing.do_bench_cudagraph(
-        FN_MAP["jit_cos_sin"], quantiles=quantiles
-    )  # type: ignore
-    p_ms, p_min_ms, p_max_ms = triton.testing.do_bench_cudagraph(
-        FN_MAP[provider], quantiles=quantiles
-    )  # type: ignore
-
-    speed_med = float(jit_ms / p_ms)
-    speed_max = float(jit_max_ms / p_min_ms)
-    speed_min = float(jit_min_ms / p_max_ms)
+    # If FlashInfer is available, use do_bench for all providers to keep methodology
+    # consistent and avoid CUDA-graph sensitivity.
+    use_cudagraph = not HAS_FLASHINFER
+
+    base_ms, base_min_ms, base_max_ms = _bench_provider(
+        BASE_PROVIDER,
+        FN_MAP[BASE_PROVIDER],
+        quantiles=quantiles,
+        use_cudagraph=use_cudagraph,
+    )
+    if base_ms != base_ms:  # NaN
+        nan = float("nan")
+        return nan, nan, nan
+
+    p_ms, p_min_ms, p_max_ms = _bench_provider(
+        provider, FN_MAP[provider], quantiles=quantiles, use_cudagraph=use_cudagraph
+    )
+
+    speed_med = float(base_ms / p_ms)
+    speed_max = float(base_max_ms / p_min_ms)
+    speed_min = float(base_min_ms / p_max_ms)
     return speed_med, speed_max, speed_min
 
 
@@ -426,7 +559,7 @@ def benchmark_speedup_vs_jit(
     benchmark.run(print_data=True)
     if (
         SHOW_SPEEDUP
-        and ("jit_cos_sin" in LINE_VALS)
+        and (BASE_PROVIDER in LINE_VALS)
         and (len(LINE_VALS) > 1)
         and (ONLY_PROVIDER is None)
     ):

From 36c39e9290d4900cfe13069ed4c5352283c350f8 Mon Sep 17 00:00:00 2001
From: RubiaCx <1084281732@qq.com>
Date: Sat, 3 Jan 2026 00:05:41 -0800
Subject: [PATCH 26/37] Move the ROT specialization to the Python JIT layer

---
 .../benchmark/bench_rotary_embedding.py       |  23 +-
 .../csrc/rotary_embedding_cos_sin.cuh         | 225 +++---------------
 python/sglang/jit_kernel/rotary_embedding.py  | 104 ++++++--
 3 files changed, 127 insertions(+), 225 deletions(-)

diff --git a/python/sglang/jit_kernel/benchmark/bench_rotary_embedding.py b/python/sglang/jit_kernel/benchmark/bench_rotary_embedding.py
index f852ed44b1e1..24be5a60aa29 100644
--- a/python/sglang/jit_kernel/benchmark/bench_rotary_embedding.py
+++ b/python/sglang/jit_kernel/benchmark/bench_rotary_embedding.py
@@ -90,10 +90,7 @@ def _bench_provider(
     quantiles: list[float],
     use_cudagraph: bool,
 ) -> tuple[float, float, float]:
-    """Bench a provider and return (median, min, max) in ms.
-
-    For FlashInfer, gracefully skip unsupported shapes by returning NaNs.
-    """
+    """Bench a provider and return (median, min, max) in ms."""
     try:
         return _bench(fn, quantiles=quantiles, use_cudagraph=use_cudagraph)
     except Exception as e:  # noqa: BLE001
@@ -218,14 +215,9 @@ def sglang_jit_rotary_cos_sin(
     rotary_embedding_cos_sin(cos, sin, q, k, head_size, interleaved)
 
 
-if HAS_FLASHINFER:
-    BASE_PROVIDER = "flashinfer_rope"
-    BASE_NAME = "FlashInfer RoPE"
-    BASE_STYLE = ("red", "-")
-else:
-    BASE_PROVIDER = "jit_cos_sin"
-    BASE_NAME = "SGL JIT (cos/sin)"
-    BASE_STYLE = ("blue", "-")
+BASE_PROVIDER = "jit_cos_sin"
+BASE_NAME = "SGL JIT (cos/sin)"
+BASE_STYLE = ("blue", "-")
 
 BASE_LINE_VALS = [BASE_PROVIDER]
 BASE_LINE_NAMES = [BASE_NAME]
@@ -235,11 +227,6 @@ def sglang_jit_rotary_cos_sin(
 LINE_NAMES = list(BASE_LINE_NAMES)
 STYLES = list(BASE_STYLES)
 
-if "jit_cos_sin" not in LINE_VALS:
-    LINE_VALS.append("jit_cos_sin")
-    LINE_NAMES.append("SGL JIT (cos/sin)")
-    STYLES.append(("blue", "-"))
-
 if HAS_SGL_POS:
     LINE_VALS.append("aot_pos")
     LINE_NAMES.append("SGL AOT (positions)")
@@ -250,7 +237,7 @@ def sglang_jit_rotary_cos_sin(
     LINE_NAMES.append("vLLM (positions)")
     STYLES.append(("green", "-."))
 
-if HAS_FLASHINFER and "flashinfer_rope" not in LINE_VALS:
+if HAS_FLASHINFER:
     LINE_VALS.append("flashinfer_rope")
     LINE_NAMES.append("FlashInfer RoPE")
     STYLES.append(("red", "--"))
diff --git a/python/sglang/jit_kernel/csrc/rotary_embedding_cos_sin.cuh b/python/sglang/jit_kernel/csrc/rotary_embedding_cos_sin.cuh
index 0e3e1df7383a..9dcff91c8f59 100644
--- a/python/sglang/jit_kernel/csrc/rotary_embedding_cos_sin.cuh
+++ b/python/sglang/jit_kernel/csrc/rotary_embedding_cos_sin.cuh
@@ -776,7 +776,7 @@ __forceinline__ void dispatch_rotary_launch(
   }
 }
 
-template <typename scalar_t>
+template <int ROT, typename scalar_t>
 inline void launch_rotary(
     const tvm::ffi::TensorView cos,
     const tvm::ffi::TensorView sin,
@@ -813,9 +813,16 @@ inline void launch_rotary(
   }
   const int embed_dim_for_rotation = interleaved ? (int)r : (int)(r / 2);
   RuntimeCheck(embed_dim_for_rotation > 0, "embed_dim_for_rotation must be > 0");
+  if constexpr (ROT > 0) {
+    RuntimeCheck(
+        embed_dim_for_rotation == ROT,
+        "embed_dim_for_rotation mismatch: got ",
+        embed_dim_for_rotation,
+        " expected ROT=",
+        ROT);
+  }
   RuntimeCheck(2LL * embed_dim_for_rotation <= head_size, "rotate dim exceeds head_size");
 
-  // Strides: JIT wrapper guarantees contiguous [T, H, D] tensors.
   const int64_t query_token_stride = hq * d;
   const int64_t head_stride_query = d;
 
@@ -882,190 +889,36 @@ inline void launch_rotary(
   const dim3 grid1d((int)t);
   const dim3 block(threads_per_block);
 
-  // Compile-time specializations for common embed dims; fallback to runtime (ROT=0).
-  switch (embed_dim_for_rotation) {
-    case 32:
-      dispatch_rotary_launch<32, scalar_t>(
-          use_grid_2d,
-          grid2d,
-          grid1d,
-          block,
-          stream,
-          interleaved,
-          use_vec,
-          qk_aligned16,
-          cos_ptr,
-          sin_ptr,
-          q_ptr,
-          k_ptr,
-          (int)r,
-          embed_dim_for_rotation,
-          query_token_stride,
-          key_token_stride,
-          head_stride_query,
-          head_stride_key,
-          (int)hq,
-          (int)hk,
-          (int)head_size,
-          blocks_per_token);
-      break;
-    case 40:
-      dispatch_rotary_launch<40, scalar_t>(
-          use_grid_2d,
-          grid2d,
-          grid1d,
-          block,
-          stream,
-          interleaved,
-          use_vec,
-          qk_aligned16,
-          cos_ptr,
-          sin_ptr,
-          q_ptr,
-          k_ptr,
-          (int)r,
-          embed_dim_for_rotation,
-          query_token_stride,
-          key_token_stride,
-          head_stride_query,
-          head_stride_key,
-          (int)hq,
-          (int)hk,
-          (int)head_size,
-          blocks_per_token);
-      break;
-    case 64:
-      dispatch_rotary_launch<64, scalar_t>(
-          use_grid_2d,
-          grid2d,
-          grid1d,
-          block,
-          stream,
-          interleaved,
-          use_vec,
-          qk_aligned16,
-          cos_ptr,
-          sin_ptr,
-          q_ptr,
-          k_ptr,
-          (int)r,
-          embed_dim_for_rotation,
-          query_token_stride,
-          key_token_stride,
-          head_stride_query,
-          head_stride_key,
-          (int)hq,
-          (int)hk,
-          (int)head_size,
-          blocks_per_token);
-      break;
-    case 80:
-      dispatch_rotary_launch<80, scalar_t>(
-          use_grid_2d,
-          grid2d,
-          grid1d,
-          block,
-          stream,
-          interleaved,
-          use_vec,
-          qk_aligned16,
-          cos_ptr,
-          sin_ptr,
-          q_ptr,
-          k_ptr,
-          (int)r,
-          embed_dim_for_rotation,
-          query_token_stride,
-          key_token_stride,
-          head_stride_query,
-          head_stride_key,
-          (int)hq,
-          (int)hk,
-          (int)head_size,
-          blocks_per_token);
-      break;
-    case 128:
-      dispatch_rotary_launch<128, scalar_t>(
-          use_grid_2d,
-          grid2d,
-          grid1d,
-          block,
-          stream,
-          interleaved,
-          use_vec,
-          qk_aligned16,
-          cos_ptr,
-          sin_ptr,
-          q_ptr,
-          k_ptr,
-          (int)r,
-          embed_dim_for_rotation,
-          query_token_stride,
-          key_token_stride,
-          head_stride_query,
-          head_stride_key,
-          (int)hq,
-          (int)hk,
-          (int)head_size,
-          blocks_per_token);
-      break;
-    case 160:
-      dispatch_rotary_launch<160, scalar_t>(
-          use_grid_2d,
-          grid2d,
-          grid1d,
-          block,
-          stream,
-          interleaved,
-          use_vec,
-          qk_aligned16,
-          cos_ptr,
-          sin_ptr,
-          q_ptr,
-          k_ptr,
-          (int)r,
-          embed_dim_for_rotation,
-          query_token_stride,
-          key_token_stride,
-          head_stride_query,
-          head_stride_key,
-          (int)hq,
-          (int)hk,
-          (int)head_size,
-          blocks_per_token);
-      break;
-    default:
-      dispatch_rotary_launch<0, scalar_t>(
-          use_grid_2d,
-          grid2d,
-          grid1d,
-          block,
-          stream,
-          interleaved,
-          use_vec,
-          qk_aligned16,
-          cos_ptr,
-          sin_ptr,
-          q_ptr,
-          k_ptr,
-          (int)r,
-          embed_dim_for_rotation,
-          query_token_stride,
-          key_token_stride,
-          head_stride_query,
-          head_stride_key,
-          (int)hq,
-          (int)hk,
-          (int)head_size,
-          blocks_per_token);
-      break;
-  }
+  dispatch_rotary_launch<ROT, scalar_t>(
+      use_grid_2d,
+      grid2d,
+      grid1d,
+      block,
+      stream,
+      interleaved,
+      use_vec,
+      qk_aligned16,
+      cos_ptr,
+      sin_ptr,
+      q_ptr,
+      k_ptr,
+      (int)r,
+      embed_dim_for_rotation,
+      query_token_stride,
+      key_token_stride,
+      head_stride_query,
+      head_stride_key,
+      (int)hq,
+      (int)hk,
+      (int)head_size,
+      blocks_per_token);
 
   RuntimeDeviceCheck();
 }
 
 }  // namespace
 
+template <int ROT>
 struct RotaryEmbeddingCosSinKernel {
   static void run_q(
       const tvm::ffi::TensorView cos,
@@ -1075,11 +928,11 @@ struct RotaryEmbeddingCosSinKernel {
       bool interleaved) {
     const auto dt = query.dtype();
     if (host::is_type<half>(dt)) {
-      launch_rotary<half>(cos, sin, query, nullptr, head_size, interleaved);
+      launch_rotary<ROT, half>(cos, sin, query, nullptr, head_size, interleaved);
     } else if (host::is_type<nv_bfloat16>(dt)) {
-      launch_rotary<nv_bfloat16>(cos, sin, query, nullptr, head_size, interleaved);
+      launch_rotary<ROT, nv_bfloat16>(cos, sin, query, nullptr, head_size, interleaved);
     } else if (host::is_type<float>(dt)) {
-      launch_rotary<float>(cos, sin, query, nullptr, head_size, interleaved);
+      launch_rotary<ROT, float>(cos, sin, query, nullptr, head_size, interleaved);
     } else {
       host::Panic("Unsupported dtype for rotary_embedding_cos_sin");
     }
@@ -1094,11 +947,11 @@ struct RotaryEmbeddingCosSinKernel {
       bool interleaved) {
     const auto dt = query.dtype();
     if (host::is_type<half>(dt)) {
-      launch_rotary<half>(cos, sin, query, &key, head_size, interleaved);
+      launch_rotary<ROT, half>(cos, sin, query, &key, head_size, interleaved);
     } else if (host::is_type<nv_bfloat16>(dt)) {
-      launch_rotary<nv_bfloat16>(cos, sin, query, &key, head_size, interleaved);
+      launch_rotary<ROT, nv_bfloat16>(cos, sin, query, &key, head_size, interleaved);
     } else if (host::is_type<float>(dt)) {
-      launch_rotary<float>(cos, sin, query, &key, head_size, interleaved);
+      launch_rotary<ROT, float>(cos, sin, query, &key, head_size, interleaved);
     } else {
       host::Panic("Unsupported dtype for rotary_embedding_cos_sin");
     }
diff --git a/python/sglang/jit_kernel/rotary_embedding.py b/python/sglang/jit_kernel/rotary_embedding.py
index 8c958f226eaf..36b694780c2e 100644
--- a/python/sglang/jit_kernel/rotary_embedding.py
+++ b/python/sglang/jit_kernel/rotary_embedding.py
@@ -1,11 +1,76 @@
 from __future__ import annotations
 
-from typing import Optional
+import logging
+from typing import TYPE_CHECKING, Optional
 
 import torch
 
 from sglang.jit_kernel.utils import cache_once, load_jit
 
+if TYPE_CHECKING:
+    from tvm_ffi.module import Module
+
+
+@cache_once
+def _jit_rotary_embedding_cos_sin_module(rot: int) -> Module:
+    # Compile one specialized template per rotary embed dim (ROT).
+    return load_jit(
+        "rotary_embedding_cos_sin",
+        str(rot),
+        cuda_files=["rotary_embedding_cos_sin.cuh"],
+        cuda_wrappers=[
+            (
+                "rotary_embedding_cos_sin_q",
+                f"RotaryEmbeddingCosSinKernel<{rot}>::run_q",
+            ),
+            (
+                "rotary_embedding_cos_sin_qk",
+                f"RotaryEmbeddingCosSinKernel<{rot}>::run_qk",
+            ),
+        ],
+    )
+
+
+@cache_once
+def can_use_rotary_embedding_cos_sin(rot: int) -> bool:
+    logger = logging.getLogger(__name__)
+    if rot <= 0:
+        logger.warning(f"Invalid rot={rot} for rotary_embedding_cos_sin")
+        return False
+    try:
+        _jit_rotary_embedding_cos_sin_module(rot)
+        return True
+    except Exception as e:
+        logger.warning(f"Failed to load JIT rotary_embedding_cos_sin (rot={rot}): {e}")
+        return False
+
+
+def rotary_embedding_cos_sin_q(
+    *,
+    rot: int,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    query: torch.Tensor,
+    head_size: int,
+    interleaved: bool,
+) -> None:
+    module = _jit_rotary_embedding_cos_sin_module(rot)
+    module.rotary_embedding_cos_sin_q(cos, sin, query, head_size, interleaved)
+
+
+def rotary_embedding_cos_sin_qk(
+    *,
+    rot: int,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    head_size: int,
+    interleaved: bool,
+) -> None:
+    module = _jit_rotary_embedding_cos_sin_module(rot)
+    module.rotary_embedding_cos_sin_qk(cos, sin, query, key, head_size, interleaved)
+
 
 def _resolve_interleaved(interleaved: Optional[bool], is_neox: Optional[bool]) -> bool:
     if interleaved is not None and is_neox is not None and interleaved != is_neox:
@@ -19,18 +84,6 @@ def _resolve_interleaved(interleaved: Optional[bool], is_neox: Optional[bool]) -
     return True
 
 
-@cache_once
-def _jit_rotary_embedding_cos_sin_module():
-    return load_jit(
-        "rotary_embedding_cos_sin",
-        cuda_files=["rotary_embedding_cos_sin.cuh"],
-        cuda_wrappers=[
-            ("rotary_embedding_cos_sin_q", "RotaryEmbeddingCosSinKernel::run_q"),
-            ("rotary_embedding_cos_sin_qk", "RotaryEmbeddingCosSinKernel::run_qk"),
-        ],
-    )
-
-
 def rotary_embedding_cos_sin(
     cos: torch.Tensor,
     sin: torch.Tensor,
@@ -91,13 +144,11 @@ def rotary_embedding_cos_sin(
     if rot_dim <= 0:
         raise ValueError("cos/sin rot_dim must be > 0")
     if effective_interleaved:
-        # Pairwise layout: rotates 2*rot_dim elements.
         if 2 * rot_dim > head_size:
             raise ValueError(
                 f"rotate dim exceeds head_size for interleaved=True: 2*rot_dim={2*rot_dim} > head_size={head_size}"
             )
     else:
-        # Split-halves layout: kernel interprets cos/sin as [rot_dim] where embed_dim=rot_dim/2.
         if rot_dim % 2 != 0:
             raise ValueError(
                 f"non-interleaved requires even rot_dim (cos.shape[1]), got rot_dim={rot_dim}"
@@ -107,6 +158,8 @@ def rotary_embedding_cos_sin(
                 f"rotate dim exceeds head_size for interleaved=False: rot_dim={rot_dim} > head_size={head_size}"
             )
 
+    embed_dim_for_rotation = rot_dim if effective_interleaved else (rot_dim // 2)
+
     def _as_3d(x: torch.Tensor, h: int) -> torch.Tensor:
         if x.dim() == 3:
             if x.shape[-1] != head_size:
@@ -124,7 +177,6 @@ def _as_3d(x: torch.Tensor, h: int) -> torch.Tensor:
         num_heads = int(query.shape[1] // head_size)
     q3 = _as_3d(query, num_heads)
     if not q3.is_contiguous():
-        # The CUDA kernel assumes contiguous [T, H, D] layout for fixed-stride addressing.
         raise ValueError("query must be contiguous (view must stay contiguous)")
 
     k3 = None
@@ -137,12 +189,22 @@ def _as_3d(x: torch.Tensor, h: int) -> torch.Tensor:
         if not k3.is_contiguous():
             raise ValueError("key must be contiguous (view must stay contiguous)")
 
-    module = _jit_rotary_embedding_cos_sin_module()
     if k3 is None:
-        module.rotary_embedding_cos_sin_q(
-            cos_to_use, sin_to_use, q3, head_size, effective_interleaved
+        rotary_embedding_cos_sin_q(
+            rot=embed_dim_for_rotation,
+            cos=cos_to_use,
+            sin=sin_to_use,
+            query=q3,
+            head_size=head_size,
+            interleaved=effective_interleaved,
         )
     else:
-        module.rotary_embedding_cos_sin_qk(
-            cos_to_use, sin_to_use, q3, k3, head_size, effective_interleaved
+        rotary_embedding_cos_sin_qk(
+            rot=embed_dim_for_rotation,
+            cos=cos_to_use,
+            sin=sin_to_use,
+            query=q3,
+            key=k3,
+            head_size=head_size,
+            interleaved=effective_interleaved,
         )

From 09e71354d65484f012c37242f94e4246fe2b2095 Mon Sep 17 00:00:00 2001
From: RubiaCx <1084281732@qq.com>
Date: Sat, 3 Jan 2026 07:09:22 -0800
Subject: [PATCH 27/37] cache RoPE split/cast

---
 python/sglang/jit_kernel/rotary_embedding.py  |  27 ++++-
 .../runtime/layers/rotary_embedding.py        | 104 ++++++++++++------
 .../runtime/models/dits/qwen_image.py         |  58 ++++++----
 3 files changed, 124 insertions(+), 65 deletions(-)

diff --git a/python/sglang/jit_kernel/rotary_embedding.py b/python/sglang/jit_kernel/rotary_embedding.py
index 36b694780c2e..e1383a8b7e71 100644
--- a/python/sglang/jit_kernel/rotary_embedding.py
+++ b/python/sglang/jit_kernel/rotary_embedding.py
@@ -112,13 +112,19 @@ def rotary_embedding_cos_sin(
         raise ValueError("cos/sin must be 2D tensors: [num_tokens, rot_dim]")
     if cos.shape != sin.shape:
         raise ValueError("cos/sin shape mismatch")
-    if cos.shape[0] != query.shape[0]:
+    if query.dim() == 4:
+        expected_tokens = int(query.shape[0] * query.shape[1])
+    elif query.dim() in (2, 3):
+        expected_tokens = int(query.shape[0])
+    else:
+        raise ValueError("query must be a 2D, 3D or 4D tensor")
+    if cos.shape[0] != expected_tokens:
         raise ValueError(
-            f"cos/sin num_tokens mismatch with query: cos.shape[0]={cos.shape[0]} vs query.shape[0]={query.shape[0]}"
+            f"cos/sin num_tokens mismatch with query: cos.shape[0]={cos.shape[0]} vs expected_tokens={expected_tokens}"
         )
 
     if head_size == 0:
-        if query.dim() == 3:
+        if query.dim() in (3, 4):
             head_size = int(query.shape[-1])
         else:
             raise ValueError("head_size must be provided when query is 2D")
@@ -161,17 +167,24 @@ def rotary_embedding_cos_sin(
     embed_dim_for_rotation = rot_dim if effective_interleaved else (rot_dim // 2)
 
     def _as_3d(x: torch.Tensor, h: int) -> torch.Tensor:
+        if x.dim() == 4:
+            # [bsz, seqlen, num_heads, head_dim] -> [tokens, num_heads, head_dim]
+            if x.shape[-1] != head_size:
+                raise ValueError("head_size mismatch with query/key last dim")
+            return x.flatten(0, 1)
         if x.dim() == 3:
             if x.shape[-1] != head_size:
                 raise ValueError("head_size mismatch with query/key last dim")
             return x
         if x.dim() != 2:
-            raise ValueError("query/key must be 2D or 3D tensors")
+            raise ValueError("query/key must be 2D, 3D or 4D tensors")
         if x.shape[1] % head_size != 0:
             raise ValueError("hidden_size is not divisible by head_size")
         return x.view(x.shape[0], h, head_size)
 
-    if query.dim() == 3:
+    if query.dim() == 4:
+        num_heads = int(query.shape[2])
+    elif query.dim() == 3:
         num_heads = int(query.shape[1])
     else:
         num_heads = int(query.shape[1] // head_size)
@@ -181,7 +194,9 @@ def _as_3d(x: torch.Tensor, h: int) -> torch.Tensor:
 
     k3 = None
     if key is not None:
-        if key.dim() == 3:
+        if key.dim() == 4:
+            num_kv_heads = int(key.shape[2])
+        elif key.dim() == 3:
             num_kv_heads = int(key.shape[1])
         else:
             num_kv_heads = int(key.shape[1] // head_size)
diff --git a/python/sglang/multimodal_gen/runtime/layers/rotary_embedding.py b/python/sglang/multimodal_gen/runtime/layers/rotary_embedding.py
index 96a4650d853d..f2f5ba6c41e4 100644
--- a/python/sglang/multimodal_gen/runtime/layers/rotary_embedding.py
+++ b/python/sglang/multimodal_gen/runtime/layers/rotary_embedding.py
@@ -38,6 +38,10 @@
 
 logger = init_logger(__name__)
 
+_JIT_ROPE_SPLIT_CAST_CACHE: "OrderedDict[tuple, tuple[torch.Tensor, torch.Tensor]]" = (
+    OrderedDict()
+)
+
 
 def apply_flashinfer_rope_qk_inplace(
     q: torch.Tensor,
@@ -128,16 +132,24 @@ def _split_cos_sin_from_cache(
 ) -> tuple[torch.Tensor, torch.Tensor]:
     """Convert a 2D RoPE cache into (cos, sin) tensors."""
     if cache.is_complex():
-        cos = cache.real.to(dtype)
-        sin = cache.imag.to(dtype)
+        cos = cache.real
+        sin = cache.imag
+        if cos.dtype != dtype:
+            cos = cos.to(dtype)
+        if sin.dtype != dtype:
+            sin = sin.to(dtype)
         return cos, sin
     if cache.shape[1] % 2 != 0:
         raise ValueError(
             f"Expected complex freqs_cis or cat([cos,sin]) cache; got real cache with odd last dim={cache.shape[1]}"
         )
     half = cache.shape[1] // 2
-    cos = cache[:, :half].to(dtype)
-    sin = cache[:, half:].to(dtype)
+    cos = cache[:, :half]
+    sin = cache[:, half:]
+    if cos.dtype != dtype:
+        cos = cos.to(dtype)
+    if sin.dtype != dtype:
+        sin = sin.to(dtype)
     return cos, sin
 
 
@@ -181,51 +193,71 @@ def apply_sglang_jit_rope_qk_inplace(
 
     num_tokens = bsz * seqlen
 
-    if positions is None:
-        pos_1d = torch.arange(seqlen, device="cpu", dtype=torch.long)
-        positions = pos_1d if bsz == 1 else pos_1d.repeat(bsz)
-    else:
-        if not (
-            isinstance(positions, torch.Tensor)
-            and positions.dtype == torch.long
-            and positions.dim() == 1
-        ):
-            raise ValueError("positions must be a 1D torch.long Tensor")
-        if positions.numel() != num_tokens:
-            raise ValueError(
-                f"positions length must be bsz*seqlen={num_tokens}, got {positions.numel()}"
-            )
-    positions = positions.to(q.device, non_blocking=True)
-
-    if cos_sin_cache.shape[0] == seqlen:
-        cache_tok = (
-            cos_sin_cache[None, :, :]
-            .expand(bsz, seqlen, cos_sin_cache.shape[1])
-            .reshape(num_tokens, cos_sin_cache.shape[1])
-        )
-    elif cos_sin_cache.shape[0] == num_tokens:
+    cache_tok: torch.Tensor
+    if cos_sin_cache.shape[0] == num_tokens:
+        cache_tok = cos_sin_cache
+    elif cos_sin_cache.shape[0] == seqlen and bsz == 1:
         cache_tok = cos_sin_cache
+    elif cos_sin_cache.shape[0] >= seqlen and positions is None and bsz == 1:
+        # Treat the cache as [max_seq, ...] and slice for this seqlen.
+        cache_tok = cos_sin_cache[:seqlen]
     else:
+        if positions is None:
+            pos_1d = torch.arange(seqlen, device="cpu", dtype=torch.long)
+            positions = pos_1d if bsz == 1 else pos_1d.repeat(bsz)
+        else:
+            if not (
+                isinstance(positions, torch.Tensor)
+                and positions.dtype == torch.long
+                and positions.dim() == 1
+            ):
+                raise ValueError("positions must be a 1D torch.long Tensor")
+            if positions.numel() != num_tokens:
+                raise ValueError(
+                    f"positions length must be bsz*seqlen={num_tokens}, got {positions.numel()}"
+                )
+        positions = positions.to(q.device, non_blocking=True)
         cache_tok = cos_sin_cache.index_select(0, positions)
 
-    q3 = q.reshape(num_tokens, nheads, d).contiguous()
-    k3 = k.reshape(num_tokens, nheads, d).contiguous()
-
-    cos_half, sin_half = _split_cos_sin_from_cache(
-        cache_tok.contiguous(), dtype=q.dtype
+    q3 = q.reshape(num_tokens, nheads, d)
+    if not q3.is_contiguous():
+        q3 = q3.contiguous()
+    k3 = k.reshape(num_tokens, nheads, d)
+    if not k3.is_contiguous():
+        k3 = k3.contiguous()
+
+    cache_key = (
+        int(cache_tok.data_ptr()),
+        int(cache_tok.shape[0]),
+        int(cache_tok.shape[1]),
+        str(cache_tok.device),
+        cache_tok.dtype,
+        q.dtype,
     )
+    cached = _JIT_ROPE_SPLIT_CAST_CACHE.get(cache_key)
+    if cached is not None:
+        cos_half, sin_half = cached
+    else:
+        cos_half, sin_half = _split_cos_sin_from_cache(cache_tok, dtype=q.dtype)
+        if not cos_half.is_contiguous():
+            cos_half = cos_half.contiguous()
+        if not sin_half.is_contiguous():
+            sin_half = sin_half.contiguous()
+        _JIT_ROPE_SPLIT_CAST_CACHE[cache_key] = (cos_half, sin_half)
+        if len(_JIT_ROPE_SPLIT_CAST_CACHE) > 64:
+            _JIT_ROPE_SPLIT_CAST_CACHE.popitem(last=False)
 
     interleaved = not is_neox
     if interleaved:
-        cos = cos_half.contiguous()
-        sin = sin_half.contiguous()
+        cos = cos_half if cos_half.is_contiguous() else cos_half.contiguous()
+        sin = sin_half if sin_half.is_contiguous() else sin_half.contiguous()
     else:
         if cos_half.shape[1] * 2 == head_size:
             cos = torch.cat([cos_half, cos_half], dim=-1).contiguous()
             sin = torch.cat([sin_half, sin_half], dim=-1).contiguous()
         else:
-            cos = cos_half.contiguous()
-            sin = sin_half.contiguous()
+            cos = cos_half if cos_half.is_contiguous() else cos_half.contiguous()
+            sin = sin_half if sin_half.is_contiguous() else sin_half.contiguous()
 
     sglang_jit_rotary_embedding_cos_sin(cos, sin, q3, k3, head_size, interleaved)
     return q3.view(bsz, seqlen, nheads, d), k3.view(bsz, seqlen, nheads, d)
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 188d6eb63e1f..73d478dcb217 100644
--- a/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py
+++ b/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py
@@ -3,7 +3,6 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import functools
-from math import prod
 from typing import Any, Dict, List, Optional, Tuple, Union
 
 import numpy as np
@@ -18,10 +17,14 @@
 from sglang.multimodal_gen.configs.models.dits.qwenimage import QwenImageDitConfig
 from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
 from sglang.multimodal_gen.runtime.layers.attention import USPAttention
-from sglang.multimodal_gen.runtime.layers.layernorm import LayerNorm, RMSNorm
+from sglang.multimodal_gen.runtime.layers.layernorm import (
+    LayerNorm,
+    RMSNorm,
+    apply_qk_norm,
+)
 from sglang.multimodal_gen.runtime.layers.linear import ReplicatedLinear
 from sglang.multimodal_gen.runtime.layers.rotary_embedding import (
-    apply_sglang_jit_rope_qk_inplace,
+    apply_flashinfer_rope_qk_inplace,
 )
 from sglang.multimodal_gen.runtime.layers.triton_ops import (
     fuse_scale_shift_gate_select01_kernel,
@@ -549,14 +552,23 @@ def forward(
         txt_value = txt_value.unflatten(-1, (self.num_heads, -1))
 
         # Apply QK normalization
-        if self.norm_q is not None:
-            img_query = self.norm_q(img_query)
-        if self.norm_k is not None:
-            img_key = self.norm_k(img_key)
-        if self.norm_added_q is not None:
-            txt_query = self.norm_added_q(txt_query)
-        if self.norm_added_k is not None:
-            txt_key = self.norm_added_k(txt_key)
+        if self.qk_norm:
+            img_query, img_key = apply_qk_norm(
+                q=img_query,
+                k=img_key,
+                q_norm=self.norm_q,
+                k_norm=self.norm_k,
+                head_dim=img_query.shape[-1],
+                allow_inplace=True,
+            )
+            txt_query, txt_key = apply_qk_norm(
+                q=txt_query,
+                k=txt_key,
+                q_norm=self.norm_added_q,
+                k_norm=self.norm_added_k,
+                head_dim=txt_query.shape[-1],
+                allow_inplace=True,
+            )
 
         # Apply RoPE
         if image_rotary_emb is not None:
@@ -567,10 +579,11 @@ def forward(
                 raise RuntimeError("image_rotary_emb must be cos_sin_cache tensors")
 
             img_cache, txt_cache = image_rotary_emb
-            img_query, img_key = apply_sglang_jit_rope_qk_inplace(
+
+            img_query, img_key = apply_flashinfer_rope_qk_inplace(
                 img_query, img_key, img_cache, is_neox=False
             )
-            txt_query, txt_key = apply_sglang_jit_rope_qk_inplace(
+            txt_query, txt_key = apply_flashinfer_rope_qk_inplace(
                 txt_query, txt_key, txt_cache, is_neox=False
             )
 
@@ -780,6 +793,12 @@ def forward(
         return encoder_hidden_states, hidden_states
 
 
+def to_hashable(obj):
+    if isinstance(obj, list):
+        return tuple(to_hashable(x) for x in obj)
+    return obj
+
+
 class QwenImageTransformer2DModel(CachableDiT, OffloadableDiTMixin):
     """
     The Transformer model introduced in Qwen.
@@ -929,16 +948,9 @@ def forward(
         timestep = (timestep / 1000).to(hidden_states.dtype)
 
         if self.zero_cond_t:
-            timestep = torch.cat([timestep, timestep * 0], dim=0)
-            # Use torch operations for GPU efficiency
-            modulate_index = torch.tensor(
-                [
-                    [0] * prod(sample[0]) + [1] * sum([prod(s) for s in sample[1:]])
-                    for sample in img_shapes
-                ],
-                device=timestep.device,
-                dtype=torch.int,
-            )
+            timestep = torch.cat([timestep, self.timestep_zero], dim=0)
+            device = timestep.device
+            modulate_index = self.build_modulate_index(to_hashable(img_shapes), device)
         else:
             modulate_index = None
 

From 2ccb1d3d89b3992ce7f7d5f00ca3aebc3de6b02c Mon Sep 17 00:00:00 2001
From: Ther-LF <2639852836@qq.com>
Date: Sun, 4 Jan 2026 15:31:20 +0800
Subject: [PATCH 28/37] jit: fuse rope cache gathering into kernel to reduce
 overhead

Signed-off-by: Ther-LF <2639852836@qq.com>
---
 .../csrc/rotary_embedding_cos_sin.cuh         | 218 +++++++++++++++---
 python/sglang/jit_kernel/rotary_embedding.py  |  13 +-
 .../runtime/layers/rotary_embedding.py        |  99 ++++----
 3 files changed, 250 insertions(+), 80 deletions(-)

diff --git a/python/sglang/jit_kernel/csrc/rotary_embedding_cos_sin.cuh b/python/sglang/jit_kernel/csrc/rotary_embedding_cos_sin.cuh
index 9dcff91c8f59..8a37a75aaaf8 100644
--- a/python/sglang/jit_kernel/csrc/rotary_embedding_cos_sin.cuh
+++ b/python/sglang/jit_kernel/csrc/rotary_embedding_cos_sin.cuh
@@ -274,6 +274,7 @@ __global__ void rotary_embedding_kernel_2d(
     const scalar_t* __restrict__ sin_data,  // [num_tokens, rot_dim_arg]
     scalar_t* __restrict__ query_total,     // [num_tokens, num_heads, head_size] contiguous
     scalar_t* __restrict__ key_total,       // [num_tokens, num_kv_heads, head_size] contiguous or nullptr
+    const int64_t* __restrict__ positions,  // [num_tokens] or nullptr
     const int rot_dim_arg,
     const int embed_dim_for_rotation_arg,
     const int64_t query_token_stride,
@@ -287,8 +288,9 @@ __global__ void rotary_embedding_kernel_2d(
   const int token_idx = blockIdx.x;
   if (token_idx >= gridDim.x) return;
 
-  const scalar_t* current_token_cos_ptr = cos_data + token_idx * rot_dim_arg;
-  const scalar_t* current_token_sin_ptr = sin_data + token_idx * rot_dim_arg;
+  const int64_t rot_idx = (positions != nullptr) ? positions[token_idx] : token_idx;
+  const scalar_t* current_token_cos_ptr = cos_data + rot_idx * rot_dim_arg;
+  const scalar_t* current_token_sin_ptr = sin_data + rot_idx * rot_dim_arg;
 
   scalar_t* query_for_token = query_total + token_idx * (int)query_token_stride;
   scalar_t* key_for_token = (key_total != nullptr) ? (key_total + token_idx * (int)key_token_stride) : nullptr;
@@ -407,6 +409,7 @@ __launch_bounds__(512) __global__ void rotary_embedding_kernel_1d(
     const scalar_t* __restrict__ sin_data,  // [num_tokens, rot_dim_arg]
     scalar_t* __restrict__ query_total,
     scalar_t* __restrict__ key_total,
+    const int64_t* __restrict__ positions,  // [num_tokens] or nullptr
     const int rot_dim_arg,
     const int embed_dim_for_rotation_arg,
     const int64_t query_token_stride,
@@ -419,8 +422,9 @@ __launch_bounds__(512) __global__ void rotary_embedding_kernel_1d(
   const int token_idx = blockIdx.x;
   if (token_idx >= gridDim.x) return;
 
-  const scalar_t* current_token_cos_ptr = cos_data + token_idx * rot_dim_arg;
-  const scalar_t* current_token_sin_ptr = sin_data + token_idx * rot_dim_arg;
+  const int64_t rot_idx = (positions != nullptr) ? positions[token_idx] : token_idx;
+  const scalar_t* current_token_cos_ptr = cos_data + rot_idx * rot_dim_arg;
+  const scalar_t* current_token_sin_ptr = sin_data + rot_idx * rot_dim_arg;
 
   scalar_t* query_for_token = query_total + token_idx * (int)query_token_stride;
   scalar_t* key_for_token = (key_total != nullptr) ? (key_total + token_idx * (int)key_token_stride) : nullptr;
@@ -542,6 +546,7 @@ __forceinline__ void dispatch_rotary_launch(
     const scalar_t* sin_ptr,
     scalar_t* q_ptr,
     scalar_t* k_ptr,
+    const int64_t* positions_ptr,
     int rot_dim_from_cache,
     int embed_dim_for_rotation,
     int64_t query_token_stride,
@@ -563,6 +568,7 @@ __forceinline__ void dispatch_rotary_launch(
               sin_ptr,
               q_ptr,
               k_ptr,
+              positions_ptr,
               rot_dim_from_cache,
               embed_dim_for_rotation,
               query_token_stride,
@@ -580,6 +586,7 @@ __forceinline__ void dispatch_rotary_launch(
               sin_ptr,
               q_ptr,
               k_ptr,
+              positions_ptr,
               rot_dim_from_cache,
               embed_dim_for_rotation,
               query_token_stride,
@@ -598,6 +605,7 @@ __forceinline__ void dispatch_rotary_launch(
             sin_ptr,
             q_ptr,
             k_ptr,
+            positions_ptr,
             rot_dim_from_cache,
             embed_dim_for_rotation,
             query_token_stride,
@@ -618,6 +626,7 @@ __forceinline__ void dispatch_rotary_launch(
               sin_ptr,
               q_ptr,
               k_ptr,
+              positions_ptr,
               rot_dim_from_cache,
               embed_dim_for_rotation,
               query_token_stride,
@@ -635,6 +644,7 @@ __forceinline__ void dispatch_rotary_launch(
               sin_ptr,
               q_ptr,
               k_ptr,
+              positions_ptr,
               rot_dim_from_cache,
               embed_dim_for_rotation,
               query_token_stride,
@@ -653,6 +663,7 @@ __forceinline__ void dispatch_rotary_launch(
             sin_ptr,
             q_ptr,
             k_ptr,
+            positions_ptr,
             rot_dim_from_cache,
             embed_dim_for_rotation,
             query_token_stride,
@@ -678,6 +689,7 @@ __forceinline__ void dispatch_rotary_launch(
             sin_ptr,
             q_ptr,
             k_ptr,
+            positions_ptr,
             rot_dim_from_cache,
             embed_dim_for_rotation,
             query_token_stride,
@@ -694,6 +706,7 @@ __forceinline__ void dispatch_rotary_launch(
             sin_ptr,
             q_ptr,
             k_ptr,
+            positions_ptr,
             rot_dim_from_cache,
             embed_dim_for_rotation,
             query_token_stride,
@@ -711,6 +724,7 @@ __forceinline__ void dispatch_rotary_launch(
           sin_ptr,
           q_ptr,
           k_ptr,
+          positions_ptr,
           rot_dim_from_cache,
           embed_dim_for_rotation,
           query_token_stride,
@@ -730,6 +744,7 @@ __forceinline__ void dispatch_rotary_launch(
             sin_ptr,
             q_ptr,
             k_ptr,
+            positions_ptr,
             rot_dim_from_cache,
             embed_dim_for_rotation,
             query_token_stride,
@@ -746,6 +761,7 @@ __forceinline__ void dispatch_rotary_launch(
             sin_ptr,
             q_ptr,
             k_ptr,
+            positions_ptr,
             rot_dim_from_cache,
             embed_dim_for_rotation,
             query_token_stride,
@@ -763,15 +779,16 @@ __forceinline__ void dispatch_rotary_launch(
           sin_ptr,
           q_ptr,
           k_ptr,
+          positions_ptr,
           rot_dim_from_cache,
-          embed_dim_for_rotation,
-          query_token_stride,
-          key_token_stride,
-          head_stride_query,
-          head_stride_key,
-          num_heads,
-          num_kv_heads,
-          head_size);
+            embed_dim_for_rotation,
+            query_token_stride,
+            key_token_stride,
+            head_stride_query,
+            head_stride_key,
+            num_heads,
+            num_kv_heads,
+            head_size);
     }
   }
 }
@@ -782,30 +799,59 @@ inline void launch_rotary(
     const tvm::ffi::TensorView sin,
     const tvm::ffi::TensorView q,
     const tvm::ffi::TensorView* k,
+    const tvm::ffi::TensorView* positions,
     int64_t head_size,
     bool interleaved) {
   using namespace host;
 
-  auto T = SymbolicSize{"T"};
   auto Hq = SymbolicSize{"Hq"};
   auto Hk = SymbolicSize{"Hk"};
   auto D = SymbolicSize{"D"};
-  auto R = SymbolicSize{"R"};
   auto dtype = SymbolicDType{};
   auto device = SymbolicDevice{};
 
-  TensorMatcher({T, R}).with_dtype<float, half, nv_bfloat16>(dtype).with_device<kDLCUDA>(device).verify(cos).verify(
-      sin);
-
-  TensorMatcher({T, Hq, D}).with_dtype<float, half, nv_bfloat16>(dtype).with_device<kDLCUDA>(device).verify(q);
+  // Verify q first to establish device/dtype
+  // q: [T, Hq, D]
+  TensorMatcher q_matcher = TensorMatcher();
+  q_matcher.with_dtype<float, half, nv_bfloat16>(dtype).with_device<kDLCUDA>(device).verify(q);
+  
+  if (q.ndim() != 3) {
+      host::Panic("q must be 3D [T, Hq, D]");
+  }
+  const int64_t t = q.size(0);
+  const int64_t hq = q.size(1);
+  const int64_t d = q.size(2);
+  
+  // Verify cos/sin
+  // cos: [T_cache, R]
+  // sin: [T_cache, R]
+  TensorMatcher c_matcher = TensorMatcher();
+  c_matcher.with_dtype(dtype).with_device(device).verify(cos).verify(sin);
+  if (cos.ndim() != 2 || sin.ndim() != 2) {
+      host::Panic("cos/sin must be 2D");
+  }
+  const int64_t r = cos.size(1);
+  if (cos.size(0) != sin.size(0) || cos.size(1) != sin.size(1)) {
+      host::Panic("cos/sin shape mismatch");
+  }
 
-  RuntimeCheck(D.unwrap() == head_size, "head_size mismatch: got ", D.unwrap(), " expected ", head_size);
-  RuntimeCheck(cos.size(1) == sin.size(1), "cos/sin dim mismatch");
+  // Handle positions
+  const int64_t* positions_ptr = nullptr;
+  if (positions != nullptr) {
+    TensorMatcher p_matcher = TensorMatcher();
+    p_matcher.with_dtype<int64_t>().with_device(device).verify(*positions);
+    if (positions->ndim() != 1 || positions->size(0) != t) {
+         host::Panic("positions must be 1D [T]");
+    }
+    positions_ptr = static_cast<const int64_t*>(positions->data_ptr());
+  } else {
+    // If positions is null, cos/sin length must equal t
+    if (cos.size(0) != t) {
+        host::Panic("cos/sin length mismatch (expected ", t, ", got ", cos.size(0), ") when positions is None");
+    }
+  }
 
-  const int64_t t = T.unwrap();
-  const int64_t hq = Hq.unwrap();
-  const int64_t d = D.unwrap();
-  const int64_t r = R.unwrap();
+  RuntimeCheck(d == head_size, "head_size mismatch: got ", d, " expected ", head_size);
 
   RuntimeCheck(t > 0 && hq > 0 && d > 0 && r > 0, "invalid shape");
   if (!interleaved) {
@@ -826,6 +872,116 @@ inline void launch_rotary(
   const int64_t query_token_stride = hq * d;
   const int64_t head_stride_query = d;
 
+  int64_t key_token_stride = 0;
+  int64_t head_stride_key = d;
+  int hk = 0;
+  if (k != nullptr) {
+    // k: [T, Hk, D]
+    TensorMatcher k_matcher = TensorMatcher();
+    k_matcher.with_dtype(dtype).with_device(device).verify(*k);
+    if (k->ndim() != 3 || k->size(0) != t || k->size(2) != d) {
+        host::Panic("key shape mismatch");
+    }
+    hk = (int)k->size(1);
+    RuntimeCheck(hk > 0, "invalid key shape");
+    key_token_stride = (int64_t)hk * d;
+  }
+
+  const int max_pairs_to_rotate_per_token =
+      (k == nullptr) ? ((int)hq * embed_dim_for_rotation)
+                     : std::max((int)hq * embed_dim_for_rotation, (int)hk * embed_dim_for_rotation);
+
+  constexpr int kVecBytes = 16;
+  const int elem_bytes = (int)sizeof(scalar_t);
+  const int kElePerVec = kVecBytes / elem_bytes;
+  const int pairs_per_step = interleaved ? (kElePerVec / 2) : kElePerVec;
+
+  bool can_vec_compute = true;
+  if ((embed_dim_for_rotation % pairs_per_step) != 0) can_vec_compute = false;
+  if ((reinterpret_cast<uintptr_t>(cos.data_ptr()) % kVecBytes) != 0) can_vec_compute = false;
+  if ((reinterpret_cast<uintptr_t>(sin.data_ptr()) % kVecBytes) != 0) can_vec_compute = false;
+  if (((r * elem_bytes) % kVecBytes) != 0) can_vec_compute = false;
+
+  bool qk_aligned16 = true;
+  if ((reinterpret_cast<uintptr_t>(q.data_ptr()) % kVecBytes) != 0) qk_aligned16 = false;
+  if (((query_token_stride * elem_bytes) % kVecBytes) != 0) qk_aligned16 = false;
+  if (((head_stride_query * elem_bytes) % kVecBytes) != 0) qk_aligned16 = false;
+  if (k != nullptr) {
+    if ((reinterpret_cast<uintptr_t>(k->data_ptr()) % kVecBytes) != 0) qk_aligned16 = false;
+    if (((key_token_stride * elem_bytes) % kVecBytes) != 0) qk_aligned16 = false;
+    if (((head_stride_key * elem_bytes) % kVecBytes) != 0) qk_aligned16 = false;
+  }
+
+  const bool use_vec = can_vec_compute;
+  const int launch_pairs_per_thread = use_vec ? pairs_per_step : 1;
+  const int total_threads_needed =
+      (max_pairs_to_rotate_per_token + launch_pairs_per_thread - 1) / launch_pairs_per_thread;
+
+  auto round_up32 = [](int x) { return ((x + 31) / 32) * 32; };
+
+  // Case 1: 2D Grid (only when num_tokens<=4 and needs multiple blocks per token)
+  const int threads_per_block_2d = std::min<int>(512, std::max(128, round_up32(std::min(total_threads_needed, 512))));
+  const int blocks_per_token_2d = (total_threads_needed + threads_per_block_2d - 1) / threads_per_block_2d;
+  const bool use_grid_2d = ((int)t <= 4) && (blocks_per_token_2d > 1);
+
+  // Case 2: 1D Grid
+  const int threads_per_block_1d = std::min<int>(512, std::max(128, round_up32(total_threads_needed)));
+
+  const int threads_per_block = use_grid_2d ? threads_per_block_2d : threads_per_block_1d;
+  const int blocks_per_token = use_grid_2d ? blocks_per_token_2d : 1;
+
+  const auto stream = LaunchKernel::resolve_device(device.unwrap());
+  const scalar_t* cos_ptr = static_cast<const scalar_t*>(cos.data_ptr());
+  const scalar_t* sin_ptr = static_cast<const scalar_t*>(sin.data_ptr());
+  scalar_t* q_ptr = static_cast<scalar_t*>(q.data_ptr());
+  scalar_t* k_ptr = (k != nullptr) ? static_cast<scalar_t*>(k->data_ptr()) : nullptr;
+
+  const dim3 grid2d((int)t, std::max(1, blocks_per_token));
+  const dim3 grid1d((int)t);
+  const dim3 block(threads_per_block);
+
+  dispatch_rotary_launch<ROT, scalar_t>(
+      use_grid_2d,
+      grid2d,
+      grid1d,
+      block,
+      stream,
+      interleaved,
+      use_vec,
+      qk_aligned16,
+      cos_ptr,
+      sin_ptr,
+      q_ptr,
+      k_ptr,
+      positions_ptr,
+      (int)r,
+      embed_dim_for_rotation,
+      query_token_stride,
+      key_token_stride,
+      head_stride_query,
+      head_stride_key,
+      (int)hq,
+      (int)hk,
+      (int)head_size,
+      blocks_per_token);
+
+  RuntimeDeviceCheck();
+}
+  const int embed_dim_for_rotation = interleaved ? (int)r : (int)(r / 2);
+  RuntimeCheck(embed_dim_for_rotation > 0, "embed_dim_for_rotation must be > 0");
+  if constexpr (ROT > 0) {
+    RuntimeCheck(
+        embed_dim_for_rotation == ROT,
+        "embed_dim_for_rotation mismatch: got ",
+        embed_dim_for_rotation,
+        " expected ROT=",
+        ROT);
+  }
+  RuntimeCheck(2LL * embed_dim_for_rotation <= head_size, "rotate dim exceeds head_size");
+
+  const int64_t query_token_stride = hq * d;
+  const int64_t head_stride_query = d;
+
   int64_t key_token_stride = 0;
   int64_t head_stride_key = d;
   int hk = 0;
@@ -924,15 +1080,16 @@ struct RotaryEmbeddingCosSinKernel {
       const tvm::ffi::TensorView cos,
       const tvm::ffi::TensorView sin,
       const tvm::ffi::TensorView query,
+      const tvm::ffi::TensorView* positions,
       int64_t head_size,
       bool interleaved) {
     const auto dt = query.dtype();
     if (host::is_type<half>(dt)) {
-      launch_rotary<ROT, half>(cos, sin, query, nullptr, head_size, interleaved);
+      launch_rotary<ROT, half>(cos, sin, query, nullptr, positions, head_size, interleaved);
     } else if (host::is_type<nv_bfloat16>(dt)) {
-      launch_rotary<ROT, nv_bfloat16>(cos, sin, query, nullptr, head_size, interleaved);
+      launch_rotary<ROT, nv_bfloat16>(cos, sin, query, nullptr, positions, head_size, interleaved);
     } else if (host::is_type<float>(dt)) {
-      launch_rotary<ROT, float>(cos, sin, query, nullptr, head_size, interleaved);
+      launch_rotary<ROT, float>(cos, sin, query, nullptr, positions, head_size, interleaved);
     } else {
       host::Panic("Unsupported dtype for rotary_embedding_cos_sin");
     }
@@ -943,15 +1100,16 @@ struct RotaryEmbeddingCosSinKernel {
       const tvm::ffi::TensorView sin,
       const tvm::ffi::TensorView query,
       const tvm::ffi::TensorView key,
+      const tvm::ffi::TensorView* positions,
       int64_t head_size,
       bool interleaved) {
     const auto dt = query.dtype();
     if (host::is_type<half>(dt)) {
-      launch_rotary<ROT, half>(cos, sin, query, &key, head_size, interleaved);
+      launch_rotary<ROT, half>(cos, sin, query, &key, positions, head_size, interleaved);
     } else if (host::is_type<nv_bfloat16>(dt)) {
-      launch_rotary<ROT, nv_bfloat16>(cos, sin, query, &key, head_size, interleaved);
+      launch_rotary<ROT, nv_bfloat16>(cos, sin, query, &key, positions, head_size, interleaved);
     } else if (host::is_type<float>(dt)) {
-      launch_rotary<ROT, float>(cos, sin, query, &key, head_size, interleaved);
+      launch_rotary<ROT, float>(cos, sin, query, &key, positions, head_size, interleaved);
     } else {
       host::Panic("Unsupported dtype for rotary_embedding_cos_sin");
     }
diff --git a/python/sglang/jit_kernel/rotary_embedding.py b/python/sglang/jit_kernel/rotary_embedding.py
index e1383a8b7e71..8959c20e097b 100644
--- a/python/sglang/jit_kernel/rotary_embedding.py
+++ b/python/sglang/jit_kernel/rotary_embedding.py
@@ -53,9 +53,10 @@ def rotary_embedding_cos_sin_q(
     query: torch.Tensor,
     head_size: int,
     interleaved: bool,
+    positions: Optional[torch.Tensor] = None,
 ) -> None:
     module = _jit_rotary_embedding_cos_sin_module(rot)
-    module.rotary_embedding_cos_sin_q(cos, sin, query, head_size, interleaved)
+    module.rotary_embedding_cos_sin_q(cos, sin, query, positions, head_size, interleaved)
 
 
 def rotary_embedding_cos_sin_qk(
@@ -67,9 +68,12 @@ def rotary_embedding_cos_sin_qk(
     key: torch.Tensor,
     head_size: int,
     interleaved: bool,
+    positions: Optional[torch.Tensor] = None,
 ) -> None:
     module = _jit_rotary_embedding_cos_sin_module(rot)
-    module.rotary_embedding_cos_sin_qk(cos, sin, query, key, head_size, interleaved)
+    module.rotary_embedding_cos_sin_qk(
+        cos, sin, query, key, positions, head_size, interleaved
+    )
 
 
 def _resolve_interleaved(interleaved: Optional[bool], is_neox: Optional[bool]) -> bool:
@@ -91,6 +95,7 @@ def rotary_embedding_cos_sin(
     key: Optional[torch.Tensor] = None,
     head_size: int = 0,
     interleaved: Optional[bool] = None,
+    positions: Optional[torch.Tensor] = None,
     *,
     is_neox: Optional[bool] = None,
 ) -> None:
@@ -118,7 +123,7 @@ def rotary_embedding_cos_sin(
         expected_tokens = int(query.shape[0])
     else:
         raise ValueError("query must be a 2D, 3D or 4D tensor")
-    if cos.shape[0] != expected_tokens:
+    if positions is None and cos.shape[0] != expected_tokens:
         raise ValueError(
             f"cos/sin num_tokens mismatch with query: cos.shape[0]={cos.shape[0]} vs expected_tokens={expected_tokens}"
         )
@@ -212,6 +217,7 @@ def _as_3d(x: torch.Tensor, h: int) -> torch.Tensor:
             query=q3,
             head_size=head_size,
             interleaved=effective_interleaved,
+            positions=positions,
         )
     else:
         rotary_embedding_cos_sin_qk(
@@ -222,4 +228,5 @@ def _as_3d(x: torch.Tensor, h: int) -> torch.Tensor:
             key=k3,
             head_size=head_size,
             interleaved=effective_interleaved,
+            positions=positions,
         )
diff --git a/python/sglang/multimodal_gen/runtime/layers/rotary_embedding.py b/python/sglang/multimodal_gen/runtime/layers/rotary_embedding.py
index f2f5ba6c41e4..19a0b5b5b643 100644
--- a/python/sglang/multimodal_gen/runtime/layers/rotary_embedding.py
+++ b/python/sglang/multimodal_gen/runtime/layers/rotary_embedding.py
@@ -194,30 +194,31 @@ def apply_sglang_jit_rope_qk_inplace(
     num_tokens = bsz * seqlen
 
     cache_tok: torch.Tensor
-    if cos_sin_cache.shape[0] == num_tokens:
-        cache_tok = cos_sin_cache
-    elif cos_sin_cache.shape[0] == seqlen and bsz == 1:
-        cache_tok = cos_sin_cache
-    elif cos_sin_cache.shape[0] >= seqlen and positions is None and bsz == 1:
-        # Treat the cache as [max_seq, ...] and slice for this seqlen.
-        cache_tok = cos_sin_cache[:seqlen]
-    else:
-        if positions is None:
-            pos_1d = torch.arange(seqlen, device="cpu", dtype=torch.long)
-            positions = pos_1d if bsz == 1 else pos_1d.repeat(bsz)
+    if cos_sin_cache.shape[0] < seqlen:
+         # Should not happen if cache is sufficient
+         pass
+
+    if positions is None:
+        if bsz > 1:
+            # For bsz > 1, we need explicit positions [0..seqlen, 0..seqlen]
+            # Create directly on device to avoid CPU overhead
+            pos_1d = torch.arange(seqlen, device=q.device, dtype=torch.long)
+            positions = pos_1d.repeat(bsz)
         else:
-            if not (
-                isinstance(positions, torch.Tensor)
-                and positions.dtype == torch.long
-                and positions.dim() == 1
-            ):
-                raise ValueError("positions must be a 1D torch.long Tensor")
-            if positions.numel() != num_tokens:
-                raise ValueError(
-                    f"positions length must be bsz*seqlen={num_tokens}, got {positions.numel()}"
-                )
+            # For bsz == 1, positions=None implies 0..seqlen, which Kernel supports natively
+            pass
+    else:
+        if not (
+            isinstance(positions, torch.Tensor)
+            and positions.dtype == torch.long
+            and positions.dim() == 1
+        ):
+            raise ValueError("positions must be a 1D torch.long Tensor")
+        if positions.numel() != num_tokens:
+            raise ValueError(
+                f"positions length must be bsz*seqlen={num_tokens}, got {positions.numel()}"
+            )
         positions = positions.to(q.device, non_blocking=True)
-        cache_tok = cos_sin_cache.index_select(0, positions)
 
     q3 = q.reshape(num_tokens, nheads, d)
     if not q3.is_contiguous():
@@ -226,40 +227,44 @@ def apply_sglang_jit_rope_qk_inplace(
     if not k3.is_contiguous():
         k3 = k3.contiguous()
 
+    # Use the FULL cache (no index_select)
+    cache_full = cos_sin_cache
+
+    # Include is_neox in key because it affects how we process (cat vs contiguous)
+    interleaved = not is_neox
     cache_key = (
-        int(cache_tok.data_ptr()),
-        int(cache_tok.shape[0]),
-        int(cache_tok.shape[1]),
-        str(cache_tok.device),
-        cache_tok.dtype,
+        int(cache_full.data_ptr()),
+        int(cache_full.shape[0]),
+        int(cache_full.shape[1]),
+        str(cache_full.device),
+        cache_full.dtype,
         q.dtype,
+        interleaved,
     )
     cached = _JIT_ROPE_SPLIT_CAST_CACHE.get(cache_key)
     if cached is not None:
-        cos_half, sin_half = cached
+        cos, sin = cached
     else:
-        cos_half, sin_half = _split_cos_sin_from_cache(cache_tok, dtype=q.dtype)
-        if not cos_half.is_contiguous():
-            cos_half = cos_half.contiguous()
-        if not sin_half.is_contiguous():
-            sin_half = sin_half.contiguous()
-        _JIT_ROPE_SPLIT_CAST_CACHE[cache_key] = (cos_half, sin_half)
+        # Split and Cast the FULL cache
+        cos_full, sin_full = _split_cos_sin_from_cache(cache_full, dtype=q.dtype)
+        
+        # Process layout (Contiguous or Cat)
+        if interleaved:
+            cos = cos_full.contiguous()
+            sin = sin_full.contiguous()
+        else:
+            if cos_full.shape[1] * 2 == head_size:
+                cos = torch.cat([cos_full, cos_full], dim=-1).contiguous()
+                sin = torch.cat([sin_full, sin_full], dim=-1).contiguous()
+            else:
+                cos = cos_full.contiguous()
+                sin = sin_full.contiguous()
+        
+        _JIT_ROPE_SPLIT_CAST_CACHE[cache_key] = (cos, sin)
         if len(_JIT_ROPE_SPLIT_CAST_CACHE) > 64:
             _JIT_ROPE_SPLIT_CAST_CACHE.popitem(last=False)
 
-    interleaved = not is_neox
-    if interleaved:
-        cos = cos_half if cos_half.is_contiguous() else cos_half.contiguous()
-        sin = sin_half if sin_half.is_contiguous() else sin_half.contiguous()
-    else:
-        if cos_half.shape[1] * 2 == head_size:
-            cos = torch.cat([cos_half, cos_half], dim=-1).contiguous()
-            sin = torch.cat([sin_half, sin_half], dim=-1).contiguous()
-        else:
-            cos = cos_half if cos_half.is_contiguous() else cos_half.contiguous()
-            sin = sin_half if sin_half.is_contiguous() else sin_half.contiguous()
-
-    sglang_jit_rotary_embedding_cos_sin(cos, sin, q3, k3, head_size, interleaved)
+    sglang_jit_rotary_embedding_cos_sin(cos, sin, q3, k3, head_size, interleaved, positions=positions)
     return q3.view(bsz, seqlen, nheads, d), k3.view(bsz, seqlen, nheads, d)
 
 

From e19f5a9a394fb61de69455d0afd92649cfb23ff4 Mon Sep 17 00:00:00 2001
From: Ther-LF <2639852836@qq.com>
Date: Sun, 4 Jan 2026 15:34:56 +0800
Subject: [PATCH 29/37] jit: fuse rope cache gathering into kernel to reduce
 overhead

Signed-off-by: Ther-LF <2639852836@qq.com>
---
 .../csrc/rotary_embedding_cos_sin.cuh         | 234 +++++++++---------
 python/sglang/jit_kernel/rotary_embedding.py  |   4 +-
 .../runtime/layers/rotary_embedding.py        |  12 +-
 3 files changed, 127 insertions(+), 123 deletions(-)

diff --git a/python/sglang/jit_kernel/csrc/rotary_embedding_cos_sin.cuh b/python/sglang/jit_kernel/csrc/rotary_embedding_cos_sin.cuh
index 8a37a75aaaf8..0def1c03788b 100644
--- a/python/sglang/jit_kernel/csrc/rotary_embedding_cos_sin.cuh
+++ b/python/sglang/jit_kernel/csrc/rotary_embedding_cos_sin.cuh
@@ -781,14 +781,14 @@ __forceinline__ void dispatch_rotary_launch(
           k_ptr,
           positions_ptr,
           rot_dim_from_cache,
-            embed_dim_for_rotation,
-            query_token_stride,
-            key_token_stride,
-            head_stride_query,
-            head_stride_key,
-            num_heads,
-            num_kv_heads,
-            head_size);
+          embed_dim_for_rotation,
+          query_token_stride,
+          key_token_stride,
+          head_stride_query,
+          head_stride_key,
+          num_heads,
+          num_kv_heads,
+          head_size);
     }
   }
 }
@@ -814,25 +814,25 @@ inline void launch_rotary(
   // q: [T, Hq, D]
   TensorMatcher q_matcher = TensorMatcher();
   q_matcher.with_dtype<float, half, nv_bfloat16>(dtype).with_device<kDLCUDA>(device).verify(q);
-  
+
   if (q.ndim() != 3) {
-      host::Panic("q must be 3D [T, Hq, D]");
+    host::Panic("q must be 3D [T, Hq, D]");
   }
   const int64_t t = q.size(0);
   const int64_t hq = q.size(1);
   const int64_t d = q.size(2);
-  
+
   // Verify cos/sin
   // cos: [T_cache, R]
   // sin: [T_cache, R]
   TensorMatcher c_matcher = TensorMatcher();
   c_matcher.with_dtype(dtype).with_device(device).verify(cos).verify(sin);
   if (cos.ndim() != 2 || sin.ndim() != 2) {
-      host::Panic("cos/sin must be 2D");
+    host::Panic("cos/sin must be 2D");
   }
   const int64_t r = cos.size(1);
   if (cos.size(0) != sin.size(0) || cos.size(1) != sin.size(1)) {
-      host::Panic("cos/sin shape mismatch");
+    host::Panic("cos/sin shape mismatch");
   }
 
   // Handle positions
@@ -841,13 +841,13 @@ inline void launch_rotary(
     TensorMatcher p_matcher = TensorMatcher();
     p_matcher.with_dtype<int64_t>().with_device(device).verify(*positions);
     if (positions->ndim() != 1 || positions->size(0) != t) {
-         host::Panic("positions must be 1D [T]");
+      host::Panic("positions must be 1D [T]");
     }
     positions_ptr = static_cast<const int64_t*>(positions->data_ptr());
   } else {
     // If positions is null, cos/sin length must equal t
     if (cos.size(0) != t) {
-        host::Panic("cos/sin length mismatch (expected ", t, ", got ", cos.size(0), ") when positions is None");
+      host::Panic("cos/sin length mismatch (expected ", t, ", got ", cos.size(0), ") when positions is None");
     }
   }
 
@@ -880,7 +880,7 @@ inline void launch_rotary(
     TensorMatcher k_matcher = TensorMatcher();
     k_matcher.with_dtype(dtype).with_device(device).verify(*k);
     if (k->ndim() != 3 || k->size(0) != t || k->size(2) != d) {
-        host::Panic("key shape mismatch");
+      host::Panic("key shape mismatch");
     }
     hk = (int)k->size(1);
     RuntimeCheck(hk > 0, "invalid key shape");
@@ -967,109 +967,109 @@ inline void launch_rotary(
 
   RuntimeDeviceCheck();
 }
-  const int embed_dim_for_rotation = interleaved ? (int)r : (int)(r / 2);
-  RuntimeCheck(embed_dim_for_rotation > 0, "embed_dim_for_rotation must be > 0");
-  if constexpr (ROT > 0) {
-    RuntimeCheck(
-        embed_dim_for_rotation == ROT,
-        "embed_dim_for_rotation mismatch: got ",
-        embed_dim_for_rotation,
-        " expected ROT=",
-        ROT);
-  }
-  RuntimeCheck(2LL * embed_dim_for_rotation <= head_size, "rotate dim exceeds head_size");
-
-  const int64_t query_token_stride = hq * d;
-  const int64_t head_stride_query = d;
-
-  int64_t key_token_stride = 0;
-  int64_t head_stride_key = d;
-  int hk = 0;
-  if (k != nullptr) {
-    TensorMatcher({T, Hk, D}).with_dtype<float, half, nv_bfloat16>(dtype).with_device<kDLCUDA>(device).verify(*k);
-    hk = (int)Hk.unwrap();
-    RuntimeCheck(hk > 0, "invalid key shape");
-    key_token_stride = (int64_t)hk * d;
-  }
-
-  const int max_pairs_to_rotate_per_token =
-      (k == nullptr) ? ((int)hq * embed_dim_for_rotation)
-                     : std::max((int)hq * embed_dim_for_rotation, (int)hk * embed_dim_for_rotation);
-
-  constexpr int kVecBytes = 16;
-  const int elem_bytes = (int)sizeof(scalar_t);
-  const int kElePerVec = kVecBytes / elem_bytes;
-  const int pairs_per_step = interleaved ? (kElePerVec / 2) : kElePerVec;
-
-  bool can_vec_compute = true;
-  if ((embed_dim_for_rotation % pairs_per_step) != 0) can_vec_compute = false;
-  if ((reinterpret_cast<uintptr_t>(cos.data_ptr()) % kVecBytes) != 0) can_vec_compute = false;
-  if ((reinterpret_cast<uintptr_t>(sin.data_ptr()) % kVecBytes) != 0) can_vec_compute = false;
-  if (((r * elem_bytes) % kVecBytes) != 0) can_vec_compute = false;
-
-  bool qk_aligned16 = true;
-  if ((reinterpret_cast<uintptr_t>(q.data_ptr()) % kVecBytes) != 0) qk_aligned16 = false;
-  if (((query_token_stride * elem_bytes) % kVecBytes) != 0) qk_aligned16 = false;
-  if (((head_stride_query * elem_bytes) % kVecBytes) != 0) qk_aligned16 = false;
-  if (k != nullptr) {
-    if ((reinterpret_cast<uintptr_t>(k->data_ptr()) % kVecBytes) != 0) qk_aligned16 = false;
-    if (((key_token_stride * elem_bytes) % kVecBytes) != 0) qk_aligned16 = false;
-    if (((head_stride_key * elem_bytes) % kVecBytes) != 0) qk_aligned16 = false;
-  }
-
-  const bool use_vec = can_vec_compute;
-  const int launch_pairs_per_thread = use_vec ? pairs_per_step : 1;
-  const int total_threads_needed =
-      (max_pairs_to_rotate_per_token + launch_pairs_per_thread - 1) / launch_pairs_per_thread;
-
-  auto round_up32 = [](int x) { return ((x + 31) / 32) * 32; };
-
-  // Case 1: 2D Grid (only when num_tokens<=4 and needs multiple blocks per token)
-  const int threads_per_block_2d = std::min<int>(512, std::max(128, round_up32(std::min(total_threads_needed, 512))));
-  const int blocks_per_token_2d = (total_threads_needed + threads_per_block_2d - 1) / threads_per_block_2d;
-  const bool use_grid_2d = ((int)t <= 4) && (blocks_per_token_2d > 1);
-
-  // Case 2: 1D Grid
-  const int threads_per_block_1d = std::min<int>(512, std::max(128, round_up32(total_threads_needed)));
-
-  const int threads_per_block = use_grid_2d ? threads_per_block_2d : threads_per_block_1d;
-  const int blocks_per_token = use_grid_2d ? blocks_per_token_2d : 1;
-
-  const auto stream = LaunchKernel::resolve_device(device.unwrap());
-  const scalar_t* cos_ptr = static_cast<const scalar_t*>(cos.data_ptr());
-  const scalar_t* sin_ptr = static_cast<const scalar_t*>(sin.data_ptr());
-  scalar_t* q_ptr = static_cast<scalar_t*>(q.data_ptr());
-  scalar_t* k_ptr = (k != nullptr) ? static_cast<scalar_t*>(k->data_ptr()) : nullptr;
-
-  const dim3 grid2d((int)t, std::max(1, blocks_per_token));
-  const dim3 grid1d((int)t);
-  const dim3 block(threads_per_block);
-
-  dispatch_rotary_launch<ROT, scalar_t>(
-      use_grid_2d,
-      grid2d,
-      grid1d,
-      block,
-      stream,
-      interleaved,
-      use_vec,
-      qk_aligned16,
-      cos_ptr,
-      sin_ptr,
-      q_ptr,
-      k_ptr,
-      (int)r,
+const int embed_dim_for_rotation = interleaved ? (int)r : (int)(r / 2);
+RuntimeCheck(embed_dim_for_rotation > 0, "embed_dim_for_rotation must be > 0");
+if constexpr (ROT > 0) {
+  RuntimeCheck(
+      embed_dim_for_rotation == ROT,
+      "embed_dim_for_rotation mismatch: got ",
       embed_dim_for_rotation,
-      query_token_stride,
-      key_token_stride,
-      head_stride_query,
-      head_stride_key,
-      (int)hq,
-      (int)hk,
-      (int)head_size,
-      blocks_per_token);
+      " expected ROT=",
+      ROT);
+}
+RuntimeCheck(2LL * embed_dim_for_rotation <= head_size, "rotate dim exceeds head_size");
+
+const int64_t query_token_stride = hq * d;
+const int64_t head_stride_query = d;
+
+int64_t key_token_stride = 0;
+int64_t head_stride_key = d;
+int hk = 0;
+if (k != nullptr) {
+  TensorMatcher({T, Hk, D}).with_dtype<float, half, nv_bfloat16>(dtype).with_device<kDLCUDA>(device).verify(*k);
+  hk = (int)Hk.unwrap();
+  RuntimeCheck(hk > 0, "invalid key shape");
+  key_token_stride = (int64_t)hk * d;
+}
 
-  RuntimeDeviceCheck();
+const int max_pairs_to_rotate_per_token =
+    (k == nullptr) ? ((int)hq * embed_dim_for_rotation)
+                   : std::max((int)hq * embed_dim_for_rotation, (int)hk* embed_dim_for_rotation);
+
+constexpr int kVecBytes = 16;
+const int elem_bytes = (int)sizeof(scalar_t);
+const int kElePerVec = kVecBytes / elem_bytes;
+const int pairs_per_step = interleaved ? (kElePerVec / 2) : kElePerVec;
+
+bool can_vec_compute = true;
+if ((embed_dim_for_rotation % pairs_per_step) != 0) can_vec_compute = false;
+if ((reinterpret_cast<uintptr_t>(cos.data_ptr()) % kVecBytes) != 0) can_vec_compute = false;
+if ((reinterpret_cast<uintptr_t>(sin.data_ptr()) % kVecBytes) != 0) can_vec_compute = false;
+if (((r * elem_bytes) % kVecBytes) != 0) can_vec_compute = false;
+
+bool qk_aligned16 = true;
+if ((reinterpret_cast<uintptr_t>(q.data_ptr()) % kVecBytes) != 0) qk_aligned16 = false;
+if (((query_token_stride * elem_bytes) % kVecBytes) != 0) qk_aligned16 = false;
+if (((head_stride_query * elem_bytes) % kVecBytes) != 0) qk_aligned16 = false;
+if (k != nullptr) {
+  if ((reinterpret_cast<uintptr_t>(k->data_ptr()) % kVecBytes) != 0) qk_aligned16 = false;
+  if (((key_token_stride * elem_bytes) % kVecBytes) != 0) qk_aligned16 = false;
+  if (((head_stride_key * elem_bytes) % kVecBytes) != 0) qk_aligned16 = false;
+}
+
+const bool use_vec = can_vec_compute;
+const int launch_pairs_per_thread = use_vec ? pairs_per_step : 1;
+const int total_threads_needed =
+    (max_pairs_to_rotate_per_token + launch_pairs_per_thread - 1) / launch_pairs_per_thread;
+
+auto round_up32 = [](int x) { return ((x + 31) / 32) * 32; };
+
+// Case 1: 2D Grid (only when num_tokens<=4 and needs multiple blocks per token)
+const int threads_per_block_2d = std::min<int>(512, std::max(128, round_up32(std::min(total_threads_needed, 512))));
+const int blocks_per_token_2d = (total_threads_needed + threads_per_block_2d - 1) / threads_per_block_2d;
+const bool use_grid_2d = ((int)t <= 4) && (blocks_per_token_2d > 1);
+
+// Case 2: 1D Grid
+const int threads_per_block_1d = std::min<int>(512, std::max(128, round_up32(total_threads_needed)));
+
+const int threads_per_block = use_grid_2d ? threads_per_block_2d : threads_per_block_1d;
+const int blocks_per_token = use_grid_2d ? blocks_per_token_2d : 1;
+
+const auto stream = LaunchKernel::resolve_device(device.unwrap());
+const scalar_t* cos_ptr = static_cast<const scalar_t*>(cos.data_ptr());
+const scalar_t* sin_ptr = static_cast<const scalar_t*>(sin.data_ptr());
+scalar_t* q_ptr = static_cast<scalar_t*>(q.data_ptr());
+scalar_t* k_ptr = (k != nullptr) ? static_cast<scalar_t*>(k->data_ptr()) : nullptr;
+
+const dim3 grid2d((int)t, std::max(1, blocks_per_token));
+const dim3 grid1d((int)t);
+const dim3 block(threads_per_block);
+
+dispatch_rotary_launch<ROT, scalar_t>(
+    use_grid_2d,
+    grid2d,
+    grid1d,
+    block,
+    stream,
+    interleaved,
+    use_vec,
+    qk_aligned16,
+    cos_ptr,
+    sin_ptr,
+    q_ptr,
+    k_ptr,
+    (int)r,
+    embed_dim_for_rotation,
+    query_token_stride,
+    key_token_stride,
+    head_stride_query,
+    head_stride_key,
+    (int)hq,
+    (int)hk,
+    (int)head_size,
+    blocks_per_token);
+
+RuntimeDeviceCheck();
 }
 
 }  // namespace
diff --git a/python/sglang/jit_kernel/rotary_embedding.py b/python/sglang/jit_kernel/rotary_embedding.py
index 8959c20e097b..2d7402966639 100644
--- a/python/sglang/jit_kernel/rotary_embedding.py
+++ b/python/sglang/jit_kernel/rotary_embedding.py
@@ -56,7 +56,9 @@ def rotary_embedding_cos_sin_q(
     positions: Optional[torch.Tensor] = None,
 ) -> None:
     module = _jit_rotary_embedding_cos_sin_module(rot)
-    module.rotary_embedding_cos_sin_q(cos, sin, query, positions, head_size, interleaved)
+    module.rotary_embedding_cos_sin_q(
+        cos, sin, query, positions, head_size, interleaved
+    )
 
 
 def rotary_embedding_cos_sin_qk(
diff --git a/python/sglang/multimodal_gen/runtime/layers/rotary_embedding.py b/python/sglang/multimodal_gen/runtime/layers/rotary_embedding.py
index 19a0b5b5b643..f1a59bcf15b7 100644
--- a/python/sglang/multimodal_gen/runtime/layers/rotary_embedding.py
+++ b/python/sglang/multimodal_gen/runtime/layers/rotary_embedding.py
@@ -195,8 +195,8 @@ def apply_sglang_jit_rope_qk_inplace(
 
     cache_tok: torch.Tensor
     if cos_sin_cache.shape[0] < seqlen:
-         # Should not happen if cache is sufficient
-         pass
+        # Should not happen if cache is sufficient
+        pass
 
     if positions is None:
         if bsz > 1:
@@ -247,7 +247,7 @@ def apply_sglang_jit_rope_qk_inplace(
     else:
         # Split and Cast the FULL cache
         cos_full, sin_full = _split_cos_sin_from_cache(cache_full, dtype=q.dtype)
-        
+
         # Process layout (Contiguous or Cat)
         if interleaved:
             cos = cos_full.contiguous()
@@ -259,12 +259,14 @@ def apply_sglang_jit_rope_qk_inplace(
             else:
                 cos = cos_full.contiguous()
                 sin = sin_full.contiguous()
-        
+
         _JIT_ROPE_SPLIT_CAST_CACHE[cache_key] = (cos, sin)
         if len(_JIT_ROPE_SPLIT_CAST_CACHE) > 64:
             _JIT_ROPE_SPLIT_CAST_CACHE.popitem(last=False)
 
-    sglang_jit_rotary_embedding_cos_sin(cos, sin, q3, k3, head_size, interleaved, positions=positions)
+    sglang_jit_rotary_embedding_cos_sin(
+        cos, sin, q3, k3, head_size, interleaved, positions=positions
+    )
     return q3.view(bsz, seqlen, nheads, d), k3.view(bsz, seqlen, nheads, d)
 
 

From b53470ab09999fbf39c014a5dd470875c2c8e26a Mon Sep 17 00:00:00 2001
From: Ther-LF <2639852836@qq.com>
Date: Sun, 4 Jan 2026 15:48:51 +0800
Subject: [PATCH 30/37] jit: fuse rope cache gathering into kernel to reduce
 overhead

Signed-off-by: Ther-LF <2639852836@qq.com>
---
 .../csrc/rotary_embedding_cos_sin.cuh         | 129 +++---------------
 1 file changed, 17 insertions(+), 112 deletions(-)

diff --git a/python/sglang/jit_kernel/csrc/rotary_embedding_cos_sin.cuh b/python/sglang/jit_kernel/csrc/rotary_embedding_cos_sin.cuh
index 0def1c03788b..3332b739ebc7 100644
--- a/python/sglang/jit_kernel/csrc/rotary_embedding_cos_sin.cuh
+++ b/python/sglang/jit_kernel/csrc/rotary_embedding_cos_sin.cuh
@@ -812,8 +812,10 @@ inline void launch_rotary(
 
   // Verify q first to establish device/dtype
   // q: [T, Hq, D]
-  TensorMatcher q_matcher = TensorMatcher();
-  q_matcher.with_dtype<float, half, nv_bfloat16>(dtype).with_device<kDLCUDA>(device).verify(q);
+  TensorMatcher()
+      .with_dtype<float, half, nv_bfloat16>(dtype)
+      .with_device<kDLCUDA>(device)
+      .verify(q);
 
   if (q.ndim() != 3) {
     host::Panic("q must be 3D [T, Hq, D]");
@@ -825,8 +827,11 @@ inline void launch_rotary(
   // Verify cos/sin
   // cos: [T_cache, R]
   // sin: [T_cache, R]
-  TensorMatcher c_matcher = TensorMatcher();
-  c_matcher.with_dtype(dtype).with_device(device).verify(cos).verify(sin);
+  TensorMatcher()
+      .with_dtype<float, half, nv_bfloat16>(dtype)
+      .with_device<kDLCUDA>(device)
+      .verify(cos)
+      .verify(sin);
   if (cos.ndim() != 2 || sin.ndim() != 2) {
     host::Panic("cos/sin must be 2D");
   }
@@ -838,8 +843,10 @@ inline void launch_rotary(
   // Handle positions
   const int64_t* positions_ptr = nullptr;
   if (positions != nullptr) {
-    TensorMatcher p_matcher = TensorMatcher();
-    p_matcher.with_dtype<int64_t>().with_device(device).verify(*positions);
+    TensorMatcher()
+        .with_dtype<int64_t>()
+        .with_device<kDLCUDA>(device)
+        .verify(*positions);
     if (positions->ndim() != 1 || positions->size(0) != t) {
       host::Panic("positions must be 1D [T]");
     }
@@ -877,8 +884,10 @@ inline void launch_rotary(
   int hk = 0;
   if (k != nullptr) {
     // k: [T, Hk, D]
-    TensorMatcher k_matcher = TensorMatcher();
-    k_matcher.with_dtype(dtype).with_device(device).verify(*k);
+    TensorMatcher()
+        .with_dtype<float, half, nv_bfloat16>(dtype)
+        .with_device<kDLCUDA>(device)
+        .verify(*k);
     if (k->ndim() != 3 || k->size(0) != t || k->size(2) != d) {
       host::Panic("key shape mismatch");
     }
@@ -967,110 +976,6 @@ inline void launch_rotary(
 
   RuntimeDeviceCheck();
 }
-const int embed_dim_for_rotation = interleaved ? (int)r : (int)(r / 2);
-RuntimeCheck(embed_dim_for_rotation > 0, "embed_dim_for_rotation must be > 0");
-if constexpr (ROT > 0) {
-  RuntimeCheck(
-      embed_dim_for_rotation == ROT,
-      "embed_dim_for_rotation mismatch: got ",
-      embed_dim_for_rotation,
-      " expected ROT=",
-      ROT);
-}
-RuntimeCheck(2LL * embed_dim_for_rotation <= head_size, "rotate dim exceeds head_size");
-
-const int64_t query_token_stride = hq * d;
-const int64_t head_stride_query = d;
-
-int64_t key_token_stride = 0;
-int64_t head_stride_key = d;
-int hk = 0;
-if (k != nullptr) {
-  TensorMatcher({T, Hk, D}).with_dtype<float, half, nv_bfloat16>(dtype).with_device<kDLCUDA>(device).verify(*k);
-  hk = (int)Hk.unwrap();
-  RuntimeCheck(hk > 0, "invalid key shape");
-  key_token_stride = (int64_t)hk * d;
-}
-
-const int max_pairs_to_rotate_per_token =
-    (k == nullptr) ? ((int)hq * embed_dim_for_rotation)
-                   : std::max((int)hq * embed_dim_for_rotation, (int)hk* embed_dim_for_rotation);
-
-constexpr int kVecBytes = 16;
-const int elem_bytes = (int)sizeof(scalar_t);
-const int kElePerVec = kVecBytes / elem_bytes;
-const int pairs_per_step = interleaved ? (kElePerVec / 2) : kElePerVec;
-
-bool can_vec_compute = true;
-if ((embed_dim_for_rotation % pairs_per_step) != 0) can_vec_compute = false;
-if ((reinterpret_cast<uintptr_t>(cos.data_ptr()) % kVecBytes) != 0) can_vec_compute = false;
-if ((reinterpret_cast<uintptr_t>(sin.data_ptr()) % kVecBytes) != 0) can_vec_compute = false;
-if (((r * elem_bytes) % kVecBytes) != 0) can_vec_compute = false;
-
-bool qk_aligned16 = true;
-if ((reinterpret_cast<uintptr_t>(q.data_ptr()) % kVecBytes) != 0) qk_aligned16 = false;
-if (((query_token_stride * elem_bytes) % kVecBytes) != 0) qk_aligned16 = false;
-if (((head_stride_query * elem_bytes) % kVecBytes) != 0) qk_aligned16 = false;
-if (k != nullptr) {
-  if ((reinterpret_cast<uintptr_t>(k->data_ptr()) % kVecBytes) != 0) qk_aligned16 = false;
-  if (((key_token_stride * elem_bytes) % kVecBytes) != 0) qk_aligned16 = false;
-  if (((head_stride_key * elem_bytes) % kVecBytes) != 0) qk_aligned16 = false;
-}
-
-const bool use_vec = can_vec_compute;
-const int launch_pairs_per_thread = use_vec ? pairs_per_step : 1;
-const int total_threads_needed =
-    (max_pairs_to_rotate_per_token + launch_pairs_per_thread - 1) / launch_pairs_per_thread;
-
-auto round_up32 = [](int x) { return ((x + 31) / 32) * 32; };
-
-// Case 1: 2D Grid (only when num_tokens<=4 and needs multiple blocks per token)
-const int threads_per_block_2d = std::min<int>(512, std::max(128, round_up32(std::min(total_threads_needed, 512))));
-const int blocks_per_token_2d = (total_threads_needed + threads_per_block_2d - 1) / threads_per_block_2d;
-const bool use_grid_2d = ((int)t <= 4) && (blocks_per_token_2d > 1);
-
-// Case 2: 1D Grid
-const int threads_per_block_1d = std::min<int>(512, std::max(128, round_up32(total_threads_needed)));
-
-const int threads_per_block = use_grid_2d ? threads_per_block_2d : threads_per_block_1d;
-const int blocks_per_token = use_grid_2d ? blocks_per_token_2d : 1;
-
-const auto stream = LaunchKernel::resolve_device(device.unwrap());
-const scalar_t* cos_ptr = static_cast<const scalar_t*>(cos.data_ptr());
-const scalar_t* sin_ptr = static_cast<const scalar_t*>(sin.data_ptr());
-scalar_t* q_ptr = static_cast<scalar_t*>(q.data_ptr());
-scalar_t* k_ptr = (k != nullptr) ? static_cast<scalar_t*>(k->data_ptr()) : nullptr;
-
-const dim3 grid2d((int)t, std::max(1, blocks_per_token));
-const dim3 grid1d((int)t);
-const dim3 block(threads_per_block);
-
-dispatch_rotary_launch<ROT, scalar_t>(
-    use_grid_2d,
-    grid2d,
-    grid1d,
-    block,
-    stream,
-    interleaved,
-    use_vec,
-    qk_aligned16,
-    cos_ptr,
-    sin_ptr,
-    q_ptr,
-    k_ptr,
-    (int)r,
-    embed_dim_for_rotation,
-    query_token_stride,
-    key_token_stride,
-    head_stride_query,
-    head_stride_key,
-    (int)hq,
-    (int)hk,
-    (int)head_size,
-    blocks_per_token);
-
-RuntimeDeviceCheck();
-}
 
 }  // namespace
 

From e37a0fa9bb626e6d9773cc4fc881c4269421003b Mon Sep 17 00:00:00 2001
From: Ther-LF <2639852836@qq.com>
Date: Sun, 4 Jan 2026 15:55:19 +0800
Subject: [PATCH 31/37] jit: fuse rope cache gathering into kernel to reduce
 overhead

Signed-off-by: Ther-LF <2639852836@qq.com>
---
 .../csrc/rotary_embedding_cos_sin.cuh         | 31 +++++++++++++++----
 1 file changed, 25 insertions(+), 6 deletions(-)

diff --git a/python/sglang/jit_kernel/csrc/rotary_embedding_cos_sin.cuh b/python/sglang/jit_kernel/csrc/rotary_embedding_cos_sin.cuh
index 3332b739ebc7..5a372170b9bc 100644
--- a/python/sglang/jit_kernel/csrc/rotary_embedding_cos_sin.cuh
+++ b/python/sglang/jit_kernel/csrc/rotary_embedding_cos_sin.cuh
@@ -6,6 +6,7 @@
 #include <cuda_fp16.h>
 #include <cuda_runtime.h>
 #include <dlpack/dlpack.h>
+#include <tvm/ffi/any.h>
 #include <tvm/ffi/container/tensor.h>
 
 #include <algorithm>
@@ -812,7 +813,7 @@ inline void launch_rotary(
 
   // Verify q first to establish device/dtype
   // q: [T, Hq, D]
-  TensorMatcher()
+  TensorMatcher({})
       .with_dtype<float, half, nv_bfloat16>(dtype)
       .with_device<kDLCUDA>(device)
       .verify(q);
@@ -827,7 +828,7 @@ inline void launch_rotary(
   // Verify cos/sin
   // cos: [T_cache, R]
   // sin: [T_cache, R]
-  TensorMatcher()
+  TensorMatcher({})
       .with_dtype<float, half, nv_bfloat16>(dtype)
       .with_device<kDLCUDA>(device)
       .verify(cos)
@@ -843,7 +844,7 @@ inline void launch_rotary(
   // Handle positions
   const int64_t* positions_ptr = nullptr;
   if (positions != nullptr) {
-    TensorMatcher()
+    TensorMatcher({})
         .with_dtype<int64_t>()
         .with_device<kDLCUDA>(device)
         .verify(*positions);
@@ -884,7 +885,7 @@ inline void launch_rotary(
   int hk = 0;
   if (k != nullptr) {
     // k: [T, Hk, D]
-    TensorMatcher()
+    TensorMatcher({})
         .with_dtype<float, half, nv_bfloat16>(dtype)
         .with_device<kDLCUDA>(device)
         .verify(*k);
@@ -985,9 +986,18 @@ struct RotaryEmbeddingCosSinKernel {
       const tvm::ffi::TensorView cos,
       const tvm::ffi::TensorView sin,
       const tvm::ffi::TensorView query,
-      const tvm::ffi::TensorView* positions,
+      const tvm::ffi::AnyView positions_any,
       int64_t head_size,
       bool interleaved) {
+    // Convert AnyView to TensorView pointer
+    const tvm::ffi::TensorView* positions = nullptr;
+    if (!positions_any.is_none()) {
+      auto positions_opt = positions_any.try_cast<tvm::ffi::TensorView>();
+      if (positions_opt.has_value()) {
+        positions = &positions_opt.value();
+      }
+    }
+    
     const auto dt = query.dtype();
     if (host::is_type<half>(dt)) {
       launch_rotary<ROT, half>(cos, sin, query, nullptr, positions, head_size, interleaved);
@@ -1005,9 +1015,18 @@ struct RotaryEmbeddingCosSinKernel {
       const tvm::ffi::TensorView sin,
       const tvm::ffi::TensorView query,
       const tvm::ffi::TensorView key,
-      const tvm::ffi::TensorView* positions,
+      const tvm::ffi::AnyView positions_any,
       int64_t head_size,
       bool interleaved) {
+    // Convert AnyView to TensorView pointer
+    const tvm::ffi::TensorView* positions = nullptr;
+    if (!positions_any.is_none()) {
+      auto positions_opt = positions_any.try_cast<tvm::ffi::TensorView>();
+      if (positions_opt.has_value()) {
+        positions = &positions_opt.value();
+      }
+    }
+    
     const auto dt = query.dtype();
     if (host::is_type<half>(dt)) {
       launch_rotary<ROT, half>(cos, sin, query, &key, positions, head_size, interleaved);

From 8185631948e1379350c3aff8d6ad48eef4566c0e Mon Sep 17 00:00:00 2001
From: Ther-LF <2639852836@qq.com>
Date: Sun, 4 Jan 2026 15:57:29 +0800
Subject: [PATCH 32/37] jit: fuse rope cache gathering into kernel to reduce
 overhead

Signed-off-by: Ther-LF <2639852836@qq.com>
---
 .../jit_kernel/csrc/rotary_embedding_cos_sin.cuh | 16 ++++++----------
 1 file changed, 6 insertions(+), 10 deletions(-)

diff --git a/python/sglang/jit_kernel/csrc/rotary_embedding_cos_sin.cuh b/python/sglang/jit_kernel/csrc/rotary_embedding_cos_sin.cuh
index 5a372170b9bc..e4c0677ed472 100644
--- a/python/sglang/jit_kernel/csrc/rotary_embedding_cos_sin.cuh
+++ b/python/sglang/jit_kernel/csrc/rotary_embedding_cos_sin.cuh
@@ -991,11 +991,9 @@ struct RotaryEmbeddingCosSinKernel {
       bool interleaved) {
     // Convert AnyView to TensorView pointer
     const tvm::ffi::TensorView* positions = nullptr;
-    if (!positions_any.is_none()) {
-      auto positions_opt = positions_any.try_cast<tvm::ffi::TensorView>();
-      if (positions_opt.has_value()) {
-        positions = &positions_opt.value();
-      }
+    auto positions_opt = positions_any.try_cast<tvm::ffi::TensorView>();
+    if (positions_opt.has_value()) {
+      positions = &positions_opt.value();
     }
     
     const auto dt = query.dtype();
@@ -1020,11 +1018,9 @@ struct RotaryEmbeddingCosSinKernel {
       bool interleaved) {
     // Convert AnyView to TensorView pointer
     const tvm::ffi::TensorView* positions = nullptr;
-    if (!positions_any.is_none()) {
-      auto positions_opt = positions_any.try_cast<tvm::ffi::TensorView>();
-      if (positions_opt.has_value()) {
-        positions = &positions_opt.value();
-      }
+    auto positions_opt = positions_any.try_cast<tvm::ffi::TensorView>();
+    if (positions_opt.has_value()) {
+      positions = &positions_opt.value();
     }
     
     const auto dt = query.dtype();

From 29ca17d42ca8dafeb94d95008574e652c34f668c Mon Sep 17 00:00:00 2001
From: Ther-LF <2639852836@qq.com>
Date: Sun, 4 Jan 2026 16:00:08 +0800
Subject: [PATCH 33/37] jit: fuse rope cache gathering into kernel to reduce
 overhead

Signed-off-by: Ther-LF <2639852836@qq.com>
---
 .../csrc/rotary_embedding_cos_sin.cuh         | 43 +++++++------------
 1 file changed, 16 insertions(+), 27 deletions(-)

diff --git a/python/sglang/jit_kernel/csrc/rotary_embedding_cos_sin.cuh b/python/sglang/jit_kernel/csrc/rotary_embedding_cos_sin.cuh
index e4c0677ed472..2afed22d9e89 100644
--- a/python/sglang/jit_kernel/csrc/rotary_embedding_cos_sin.cuh
+++ b/python/sglang/jit_kernel/csrc/rotary_embedding_cos_sin.cuh
@@ -805,58 +805,50 @@ inline void launch_rotary(
     bool interleaved) {
   using namespace host;
 
+  auto T = SymbolicSize{"T"};
   auto Hq = SymbolicSize{"Hq"};
   auto Hk = SymbolicSize{"Hk"};
   auto D = SymbolicSize{"D"};
+  auto R = SymbolicSize{"R"};
   auto dtype = SymbolicDType{};
   auto device = SymbolicDevice{};
 
   // Verify q first to establish device/dtype
   // q: [T, Hq, D]
-  TensorMatcher({})
+  TensorMatcher({T, Hq, D})
       .with_dtype<float, half, nv_bfloat16>(dtype)
       .with_device<kDLCUDA>(device)
       .verify(q);
 
-  if (q.ndim() != 3) {
-    host::Panic("q must be 3D [T, Hq, D]");
-  }
-  const int64_t t = q.size(0);
-  const int64_t hq = q.size(1);
-  const int64_t d = q.size(2);
+  const int64_t t = T.unwrap();
+  const int64_t hq = Hq.unwrap();
+  const int64_t d = D.unwrap();
 
   // Verify cos/sin
   // cos: [T_cache, R]
   // sin: [T_cache, R]
-  TensorMatcher({})
+  auto T_cache = SymbolicSize{"T_cache"};
+  TensorMatcher({T_cache, R})
       .with_dtype<float, half, nv_bfloat16>(dtype)
       .with_device<kDLCUDA>(device)
       .verify(cos)
       .verify(sin);
-  if (cos.ndim() != 2 || sin.ndim() != 2) {
-    host::Panic("cos/sin must be 2D");
-  }
-  const int64_t r = cos.size(1);
-  if (cos.size(0) != sin.size(0) || cos.size(1) != sin.size(1)) {
-    host::Panic("cos/sin shape mismatch");
-  }
+
+  const int64_t r = R.unwrap();
 
   // Handle positions
   const int64_t* positions_ptr = nullptr;
   if (positions != nullptr) {
-    TensorMatcher({})
+    auto T_pos = SymbolicSize{"T_pos"};
+    TensorMatcher({T_pos})
         .with_dtype<int64_t>()
         .with_device<kDLCUDA>(device)
         .verify(*positions);
-    if (positions->ndim() != 1 || positions->size(0) != t) {
-      host::Panic("positions must be 1D [T]");
-    }
+    RuntimeCheck(T_pos.unwrap() == t, "positions length mismatch");
     positions_ptr = static_cast<const int64_t*>(positions->data_ptr());
   } else {
     // If positions is null, cos/sin length must equal t
-    if (cos.size(0) != t) {
-      host::Panic("cos/sin length mismatch (expected ", t, ", got ", cos.size(0), ") when positions is None");
-    }
+    RuntimeCheck(T_cache.unwrap() == t, "cos/sin length mismatch (expected ", t, ", got ", T_cache.unwrap(), ") when positions is None");
   }
 
   RuntimeCheck(d == head_size, "head_size mismatch: got ", d, " expected ", head_size);
@@ -885,14 +877,11 @@ inline void launch_rotary(
   int hk = 0;
   if (k != nullptr) {
     // k: [T, Hk, D]
-    TensorMatcher({})
+    TensorMatcher({T, Hk, D})
         .with_dtype<float, half, nv_bfloat16>(dtype)
         .with_device<kDLCUDA>(device)
         .verify(*k);
-    if (k->ndim() != 3 || k->size(0) != t || k->size(2) != d) {
-      host::Panic("key shape mismatch");
-    }
-    hk = (int)k->size(1);
+    hk = (int)Hk.unwrap();
     RuntimeCheck(hk > 0, "invalid key shape");
     key_token_stride = (int64_t)hk * d;
   }

From 2315e410195b3497df75c4787e1724524141b049 Mon Sep 17 00:00:00 2001
From: Ther-LF <2639852836@qq.com>
Date: Sun, 4 Jan 2026 16:18:37 +0800
Subject: [PATCH 34/37] bench: jit fused overhead

Signed-off-by: Ther-LF <2639852836@qq.com>
---
 .../benchmark/bench_rotary_embedding.py       | 160 +++++++++++++++---
 1 file changed, 140 insertions(+), 20 deletions(-)

diff --git a/python/sglang/jit_kernel/benchmark/bench_rotary_embedding.py b/python/sglang/jit_kernel/benchmark/bench_rotary_embedding.py
index 24be5a60aa29..c99b93e26215 100644
--- a/python/sglang/jit_kernel/benchmark/bench_rotary_embedding.py
+++ b/python/sglang/jit_kernel/benchmark/bench_rotary_embedding.py
@@ -1,3 +1,23 @@
+"""
+Benchmark for Rotary Embedding implementations.
+
+Key test scenarios:
+- jit_fused: JIT kernel with fused gathering (positions + full cache) - OPTIMIZED
+  * Avoids explicit index_select in Python
+  * Gathering happens inside the kernel
+  
+- jit_unfused: JIT kernel with explicit gathering (positions + full cache) - BASELINE
+  * Includes index_select overhead in Python
+  * Shows the cost we're trying to avoid
+  
+- jit_cos_sin: JIT kernel with pre-gathered cos/sin (backward compatibility)
+  * Traditional approach: cos/sin already gathered before kernel call
+  * Does not include gathering overhead in timing
+  
+- aot_pos: AOT kernel with positions
+- vllm_pos: vLLM's implementation
+- flashinfer_rope: FlashInfer's RoPE implementation
+"""
 import itertools
 import os
 from functools import lru_cache
@@ -215,8 +235,51 @@ def sglang_jit_rotary_cos_sin(
     rotary_embedding_cos_sin(cos, sin, q, k, head_size, interleaved)
 
 
-BASE_PROVIDER = "jit_cos_sin"
-BASE_NAME = "SGL JIT (cos/sin)"
+def sglang_jit_rotary_with_positions(
+    positions: torch.Tensor,
+    cos_cache: torch.Tensor,
+    sin_cache: torch.Tensor,
+    q: torch.Tensor,
+    k: torch.Tensor,
+    head_size: int,
+    interleaved: bool,
+) -> None:
+    """
+    JIT kernel with fused gathering: positions + full cache.
+    This is the optimized version that avoids explicit index_select in Python.
+    """
+    from sglang.jit_kernel.rotary_embedding import rotary_embedding_cos_sin
+
+    rotary_embedding_cos_sin(
+        cos_cache, sin_cache, q, k, head_size, interleaved, positions=positions
+    )
+
+
+def sglang_jit_rotary_with_gather(
+    positions: torch.Tensor,
+    cos_cache: torch.Tensor,
+    sin_cache: torch.Tensor,
+    q: torch.Tensor,
+    k: torch.Tensor,
+    head_size: int,
+    interleaved: bool,
+) -> None:
+    """
+    JIT kernel WITHOUT fused gathering: explicit index_select in Python.
+    This is the baseline version that includes the gathering overhead.
+    """
+    from sglang.jit_kernel.rotary_embedding import rotary_embedding_cos_sin
+
+    # Explicit gathering in Python (the overhead we want to avoid)
+    cos_gathered = cos_cache[positions].contiguous()
+    sin_gathered = sin_cache[positions].contiguous()
+    rotary_embedding_cos_sin(
+        cos_gathered, sin_gathered, q, k, head_size, interleaved
+    )
+
+
+BASE_PROVIDER = "jit_fused"
+BASE_NAME = "SGL JIT (fused)"
 BASE_STYLE = ("blue", "-")
 
 BASE_LINE_VALS = [BASE_PROVIDER]
@@ -227,6 +290,11 @@ def sglang_jit_rotary_cos_sin(
 LINE_NAMES = list(BASE_LINE_NAMES)
 STYLES = list(BASE_STYLES)
 
+# Add unfused version for comparison
+LINE_VALS.append("jit_unfused")
+LINE_NAMES.append("SGL JIT (unfused)")
+STYLES.append(("purple", ":"))
+
 if HAS_SGL_POS:
     LINE_VALS.append("aot_pos")
     LINE_NAMES.append("SGL AOT (positions)")
@@ -331,8 +399,10 @@ def benchmark(
     num_tokens = batch_size * seq_len
     rotary_dim = head_size
 
+    # Prepare full cache (for fused gathering tests)
     cos_cache_half, sin_cache_half = _get_cos_sin_cache_half_cuda(rotary_dim)
     cos_sin_cache_aot = torch.cat([cos_cache_half, sin_cache_half], dim=-1).contiguous()
+    
     # FlashInfer requires cos_sin_cache to be float32.
     if HAS_FLASHINFER:
         cos_f32, sin_f32 = _compute_cos_sin_cache_half(
@@ -344,17 +414,27 @@ def benchmark(
             .contiguous()
         )
 
+    # Prepare positions (simulate real decoding scenario)
+    # Use arange for simplicity, but in real scenario this could be non-contiguous
     positions = torch.arange(seq_len, device=DEVICE, dtype=torch.int64).repeat(
         batch_size
     )
+    
+    # Pre-gathered cos/sin for unfused tests (traditional approach)
     cos_half = cos_cache_half[positions].contiguous()
     sin_half = sin_cache_half[positions].contiguous()
 
     if interleaved:
-        cos, sin = cos_half, sin_half
+        cos_gathered, sin_gathered = cos_half, sin_half
+        # For fused tests, prepare cache in interleaved format
+        cos_cache_for_fused = cos_cache_half
+        sin_cache_for_fused = sin_cache_half
     else:
-        cos = torch.cat([cos_half, cos_half], dim=-1).contiguous()
-        sin = torch.cat([sin_half, sin_half], dim=-1).contiguous()
+        cos_gathered = torch.cat([cos_half, cos_half], dim=-1).contiguous()
+        sin_gathered = torch.cat([sin_half, sin_half], dim=-1).contiguous()
+        # For fused tests, prepare cache in non-interleaved format
+        cos_cache_for_fused = torch.cat([cos_cache_half, cos_cache_half], dim=-1).contiguous()
+        sin_cache_for_fused = torch.cat([sin_cache_half, sin_cache_half], dim=-1).contiguous()
 
     # Keep a canonical 4D layout for providers that require it (e.g. FlashInfer).
     # View as 3D for providers that accept [num_tokens, nheads, head_size].
@@ -370,35 +450,60 @@ def benchmark(
     global _SANITY_DONE
     if (
         (not _SANITY_DONE)
-        and provider == "aot_pos"
+        and provider == "jit_fused"
         and batch_size == BS_RANGE[0]
         and seq_len == SEQ_RANGE[0]
         and head_size == HEAD_SIZE_RANGE[0]
         and interleaved == INTERLEAVED_RANGE[0]
     ):
+        # Reference implementation
         q_ref = q.clone()
         k_ref = k.clone()
-        torch_impl_rotary_fp32(cos, sin, q_ref, k_ref, head_size, interleaved)
-        q_out = q.clone()
-        k_out = k.clone()
-        sglang_aot_rotary_positions(
-            positions, q_out, k_out, head_size, interleaved, cos_sin_cache_aot
+        torch_impl_rotary_fp32(cos_gathered, sin_gathered, q_ref, k_ref, head_size, interleaved)
+        
+        # Test fused version
+        q_fused_test = q.clone()
+        k_fused_test = k.clone()
+        sglang_jit_rotary_with_positions(
+            positions, cos_cache_for_fused, sin_cache_for_fused, 
+            q_fused_test, k_fused_test, head_size, interleaved
         )
+        
+        # Test unfused version
+        q_unfused_test = q.clone()
+        k_unfused_test = k.clone()
+        sglang_jit_rotary_with_gather(
+            positions, cos_cache_for_fused, sin_cache_for_fused,
+            q_unfused_test, k_unfused_test, head_size, interleaved
+        )
+        
         ref_atol = 2e-2 if DTYPE == torch.bfloat16 else 2e-3
         ref_rtol = 2e-2 if DTYPE == torch.bfloat16 else 2e-3
-        _assert_close_gpu(q_out, q_ref, atol=ref_atol, rtol=ref_rtol, name="AOT(q)")
-        _assert_close_gpu(k_out, k_ref, atol=ref_atol, rtol=ref_rtol, name="AOT(k)")
+        _assert_close_gpu(q_fused_test, q_ref, atol=ref_atol, rtol=ref_rtol, name="JIT-fused(q)")
+        _assert_close_gpu(k_fused_test, k_ref, atol=ref_atol, rtol=ref_rtol, name="JIT-fused(k)")
+        _assert_close_gpu(q_unfused_test, q_ref, atol=ref_atol, rtol=ref_rtol, name="JIT-unfused(q)")
+        _assert_close_gpu(k_unfused_test, k_ref, atol=ref_atol, rtol=ref_rtol, name="JIT-unfused(k)")
+        print("✅ Sanity check passed: fused and unfused versions match reference!")
         _SANITY_DONE = True
 
     FN_MAP = {
         "aot_pos": lambda: sglang_aot_rotary_positions(
             positions, q, k, head_size, interleaved, cos_sin_cache_aot
         ),
+        # NEW: Fused gathering - positions + full cache (OPTIMIZED)
+        "jit_fused": lambda: sglang_jit_rotary_with_positions(
+            positions, cos_cache_for_fused, sin_cache_for_fused, q, k, head_size, interleaved
+        ),
+        # NEW: Unfused - explicit gathering in Python (BASELINE)
+        "jit_unfused": lambda: sglang_jit_rotary_with_gather(
+            positions, cos_cache_for_fused, sin_cache_for_fused, q, k, head_size, interleaved
+        ),
+        # OLD: Pre-gathered cos/sin (for backward compatibility)
         "jit_cos_sin": lambda: sglang_jit_rotary_cos_sin(
-            cos, sin, q, k, head_size, interleaved
+            cos_gathered, sin_gathered, q, k, head_size, interleaved
         ),
         "torch_fp32": lambda: torch_impl_rotary_fp32(
-            cos, sin, q, k, head_size, interleaved
+            cos_gathered, sin_gathered, q, k, head_size, interleaved
         ),
     }
     if HAS_FLASHINFER and apply_flashinfer_rope_qk_inplace is not None:
@@ -455,6 +560,7 @@ def benchmark_speedup_vs_jit(
     num_tokens = batch_size * seq_len
     rotary_dim = head_size
 
+    # Prepare full cache
     cos_cache_half, sin_cache_half = _get_cos_sin_cache_half_cuda(rotary_dim)
     cos_sin_cache_aot = torch.cat([cos_cache_half, sin_cache_half], dim=-1).contiguous()
     if HAS_FLASHINFER:
@@ -474,10 +580,14 @@ def benchmark_speedup_vs_jit(
     sin_half = sin_cache_half[positions].contiguous()
 
     if interleaved:
-        cos, sin = cos_half, sin_half
+        cos_gathered, sin_gathered = cos_half, sin_half
+        cos_cache_for_fused = cos_cache_half
+        sin_cache_for_fused = sin_cache_half
     else:
-        cos = torch.cat([cos_half, cos_half], dim=-1).contiguous()
-        sin = torch.cat([sin_half, sin_half], dim=-1).contiguous()
+        cos_gathered = torch.cat([cos_half, cos_half], dim=-1).contiguous()
+        sin_gathered = torch.cat([sin_half, sin_half], dim=-1).contiguous()
+        cos_cache_for_fused = torch.cat([cos_cache_half, cos_cache_half], dim=-1).contiguous()
+        sin_cache_for_fused = torch.cat([sin_cache_half, sin_cache_half], dim=-1).contiguous()
 
     q_base4 = torch.randn(
         batch_size, seq_len, NUM_Q_HEADS, head_size, device=DEVICE, dtype=DTYPE
@@ -488,20 +598,30 @@ def benchmark_speedup_vs_jit(
     q_base = q_base4.view(num_tokens, NUM_Q_HEADS, head_size)
     k_base = k_base4.view(num_tokens, NUM_Q_HEADS, head_size)
 
+    q_fused = q_base.clone()
+    k_fused = k_base.clone()
+    q_unfused = q_base.clone()
+    k_unfused = k_base.clone()
     q_jit = q_base.clone()
     k_jit = k_base.clone()
     q_p = q_base.clone()
     k_p = k_base.clone()
 
     FN_MAP = {
+        "jit_fused": lambda: sglang_jit_rotary_with_positions(
+            positions, cos_cache_for_fused, sin_cache_for_fused, q_fused, k_fused, head_size, interleaved
+        ),
+        "jit_unfused": lambda: sglang_jit_rotary_with_gather(
+            positions, cos_cache_for_fused, sin_cache_for_fused, q_unfused, k_unfused, head_size, interleaved
+        ),
         "jit_cos_sin": lambda: sglang_jit_rotary_cos_sin(
-            cos, sin, q_jit, k_jit, head_size, interleaved
+            cos_gathered, sin_gathered, q_jit, k_jit, head_size, interleaved
         ),
         "aot_pos": lambda: sglang_aot_rotary_positions(
             positions, q_p, k_p, head_size, interleaved, cos_sin_cache_aot
         ),
         "torch_fp32": lambda: torch_impl_rotary_fp32(
-            cos, sin, q_p, k_p, head_size, interleaved
+            cos_gathered, sin_gathered, q_p, k_p, head_size, interleaved
         ),
     }
     if HAS_FLASHINFER and apply_flashinfer_rope_qk_inplace is not None:

From 1a4b8741c05f78caf4b8df62afd2492827da7bf9 Mon Sep 17 00:00:00 2001
From: RubiaCx <1084281732@qq.com>
Date: Mon, 5 Jan 2026 06:04:40 -0800
Subject: [PATCH 35/37] clean up; fix 1D prefetch UB; align RoPE
 benchmarks/tests with positions+cache

---
 .../benchmark/bench_rotary_embedding.py       | 784 +++++++-----------
 .../csrc/rotary_embedding_cos_sin.cuh         | 344 ++------
 python/sglang/jit_kernel/rotary_embedding.py  | 221 ++---
 .../jit_kernel/tests/test_rotary_embedding.py |  42 +-
 4 files changed, 436 insertions(+), 955 deletions(-)

diff --git a/python/sglang/jit_kernel/benchmark/bench_rotary_embedding.py b/python/sglang/jit_kernel/benchmark/bench_rotary_embedding.py
index c99b93e26215..6d5621cb8699 100644
--- a/python/sglang/jit_kernel/benchmark/bench_rotary_embedding.py
+++ b/python/sglang/jit_kernel/benchmark/bench_rotary_embedding.py
@@ -1,53 +1,51 @@
 """
-Benchmark for Rotary Embedding implementations.
-
-Key test scenarios:
-- jit_fused: JIT kernel with fused gathering (positions + full cache) - OPTIMIZED
-  * Avoids explicit index_select in Python
-  * Gathering happens inside the kernel
-  
-- jit_unfused: JIT kernel with explicit gathering (positions + full cache) - BASELINE
-  * Includes index_select overhead in Python
-  * Shows the cost we're trying to avoid
-  
-- jit_cos_sin: JIT kernel with pre-gathered cos/sin (backward compatibility)
-  * Traditional approach: cos/sin already gathered before kernel call
-  * Does not include gathering overhead in timing
-  
-- aot_pos: AOT kernel with positions
-- vllm_pos: vLLM's implementation
-- flashinfer_rope: FlashInfer's RoPE implementation
+Lean Rotary Embedding benchmark.
+
+Providers:
+- jit_fused   : SGL JIT with fused gather (positions + full cache)
+- jit_unfused : SGL JIT with Python gather (indexing cos/sin in Python)
+- aot_pos     : SGL AOT with fused gather
+- vllm_pos    : vLLM
+- flashinfer  : FlashInfer
+- torch_fp32  : Torch
 """
+
 import itertools
 import os
+from dataclasses import dataclass
 from functools import lru_cache
-from typing import Tuple
+from typing import Callable, Dict, Tuple
 
 import torch
-import triton
 import triton.testing
 
+DEVICE = "cuda"
+DTYPE = torch.bfloat16
+MAX_SEQ_LEN = 8192
+NUM_HEADS = 32
+BS_RANGE = [1, 64]
+SEQ_RANGE = [1, 2048]
+HEAD_SIZE_RANGE = [128]
+INTERLEAVED_RANGE = [True, False]
+
+ONLY_PROVIDER = os.getenv("SGL_BENCH_PROVIDER", "").strip() or None
+INCLUDE_TORCH_REF = os.getenv("SGL_BENCH_INCLUDE_TORCH", "0").lower() in (
+    "1",
+    "true",
+    "yes",
+    "y",
+)
+SHOW_SPEEDUP = os.getenv("SGL_BENCH_SPEEDUP", "1").lower() in ("1", "true", "yes", "y")
+
 try:
     import sgl_kernel  # noqa: F401
 except Exception:
     sgl_kernel = None  # type: ignore
 
-HAS_SGL_POS = hasattr(torch.ops.sgl_kernel, "rotary_embedding")
-
-IS_CI = (
-    os.getenv("CI", "false").lower() == "true"
-    or os.getenv("GITHUB_ACTIONS", "false").lower() == "true"
+HAS_SGL_POS = hasattr(torch.ops, "sgl_kernel") and hasattr(
+    torch.ops.sgl_kernel, "rotary_embedding"
 )
 
-ONLY_PROVIDER = os.getenv("SGL_BENCH_PROVIDER", "").strip() or None
-SHOW_SPEEDUP = os.getenv("SGL_BENCH_SPEEDUP", "1").lower() in ("1", "true", "yes", "y")
-
-DEVICE = "cuda"
-MAX_SEQ_LEN = 8192
-NUM_Q_HEADS = 32
-NUM_KV_HEADS = 8
-DTYPE = torch.bfloat16
-
 try:
     from vllm.model_executor.layers.rotary_embedding import (
         RotaryEmbedding as vLLMRotaryEmbedding,
@@ -63,94 +61,48 @@
         apply_flashinfer_rope_qk_inplace,
     )
 
-    try:
-        from flashinfer.rope import (  # noqa: F401
-            apply_rope_with_cos_sin_cache_inplace as _,
-        )
-
-        HAS_FLASHINFER = True
-    except Exception:
-        HAS_FLASHINFER = False
+    HAS_FLASHINFER = True
 except Exception:
     apply_flashinfer_rope_qk_inplace = None  # type: ignore
     HAS_FLASHINFER = False
 
-if IS_CI:
-    BS_RANGE = [16]
-    SEQ_RANGE = [1, 128]
-    HEAD_SIZE_RANGE = [128]
-    INTERLEAVED_RANGE = [True]
-else:
-    BS_RANGE = [1, 8, 64]
-    SEQ_RANGE = [1, 4, 128, 2048]
-    HEAD_SIZE_RANGE = [64, 96, 128, 256]
-    INTERLEAVED_RANGE = [True, False]
-
-
-def _is_flashinfer_unsupported_shape_error(e: BaseException) -> bool:
-    msg = str(e)
-    return ("Unsupported head_dim" in msg) or ("cos_sin_cache should be float32" in msg)
-
-
-def _bench(
-    fn,
-    *,
-    quantiles: list[float],
-    use_cudagraph: bool,
-) -> tuple[float, float, float]:
-    if use_cudagraph:
-        return triton.testing.do_bench_cudagraph(fn, quantiles=quantiles)  # type: ignore
-    return triton.testing.do_bench(fn, quantiles=quantiles)  # type: ignore
-
-
-def _bench_provider(
-    provider: str,
-    fn,
-    *,
-    quantiles: list[float],
-    use_cudagraph: bool,
-) -> tuple[float, float, float]:
-    """Bench a provider and return (median, min, max) in ms."""
-    try:
-        return _bench(fn, quantiles=quantiles, use_cudagraph=use_cudagraph)
-    except Exception as e:  # noqa: BLE001
-        if provider == "flashinfer_rope" and _is_flashinfer_unsupported_shape_error(e):
-            nan = float("nan")
-            return nan, nan, nan
-        raise
-
 
-def _compute_cos_sin_cache_half(
+def _compute_cos_sin_cache(
     max_seq_len: int,
     rotary_dim: int,
     base: float = 10000.0,
     dtype: torch.dtype = torch.float32,
-) -> Tuple[torch.Tensor, torch.Tensor]:
+):
     inv_freq = 1.0 / (
         base ** (torch.arange(0, rotary_dim, 2, dtype=torch.float32) / rotary_dim)
     )
     t = torch.arange(max_seq_len, dtype=torch.float32)
     freqs = torch.einsum("i,j->ij", t, inv_freq)
-    cos = freqs.cos().to(dtype)
-    sin = freqs.sin().to(dtype)
-    return cos, sin
+    return freqs.cos().to(dtype), freqs.sin().to(dtype)
 
 
 @lru_cache(maxsize=None)
-def _get_cos_sin_cache_half_cuda(rotary_dim: int) -> Tuple[torch.Tensor, torch.Tensor]:
-    cos, sin = _compute_cos_sin_cache_half(MAX_SEQ_LEN, rotary_dim, dtype=DTYPE)
+def _cos_sin_cache_half_cuda(rotary_dim: int) -> Tuple[torch.Tensor, torch.Tensor]:
+    cos, sin = _compute_cos_sin_cache(MAX_SEQ_LEN, rotary_dim, dtype=DTYPE)
     return cos.to(device=DEVICE, dtype=DTYPE), sin.to(device=DEVICE, dtype=DTYPE)
 
 
+@lru_cache(maxsize=None)
+def _vllm_rope(head_size: int, rotary_dim: int, interleaved: bool, dtype: torch.dtype):
+    if not HAS_VLLM or vLLMRotaryEmbedding is None:
+        raise RuntimeError("vLLM not available")
+    return vLLMRotaryEmbedding(
+        head_size=head_size,
+        rotary_dim=rotary_dim,
+        max_position_embeddings=MAX_SEQ_LEN,
+        base=10000,
+        is_neox_style=interleaved,
+        dtype=dtype,
+    ).cuda()
+
+
 @torch.no_grad()
-def torch_impl_rotary_fp32(
-    cos: torch.Tensor,
-    sin: torch.Tensor,
-    q: torch.Tensor,
-    k: torch.Tensor,
-    head_size: int,
-    interleaved: bool,
-) -> None:
+def torch_impl_rotary_fp32(cos, sin, q, k, head_size: int, interleaved: bool) -> None:
     orig_dtype = q.dtype
     if interleaved and cos.shape[1] == head_size:
         half = head_size // 2
@@ -166,7 +118,7 @@ def torch_impl_rotary_fp32(
         cos_b = cos_f[:, None, :embed_dim]
         sin_b = sin_f[:, None, :embed_dim]
 
-        def _apply(x: torch.Tensor) -> None:
+        def _apply(x):
             xr = x[..., :rot_dim]
             xr2 = xr.view(xr.shape[0], xr.shape[1], embed_dim, 2)
             x0 = xr2[..., 0].clone()
@@ -189,7 +141,7 @@ def _apply(x: torch.Tensor) -> None:
                 sin_f[:, None, embed_dim:rot_dim],
             )
 
-        def _apply(x: torch.Tensor) -> None:
+        def _apply(x):
             xr = x[..., :rot_dim]
             x0 = xr[..., :embed_dim].clone()
             x1 = xr[..., embed_dim:rot_dim].clone()
@@ -202,52 +154,20 @@ def _apply(x: torch.Tensor) -> None:
     k.copy_(k_f.to(orig_dtype))
 
 
-def sglang_aot_rotary_positions(
-    positions: torch.Tensor,
-    q: torch.Tensor,
-    k: torch.Tensor,
-    head_size: int,
-    interleaved: bool,
-    cos_sin_cache: torch.Tensor,
-) -> None:
-    if not HAS_SGL_POS:
-        raise RuntimeError("torch.ops.sgl_kernel.rotary_embedding is not available")
-    torch.ops.sgl_kernel.rotary_embedding(
-        positions,
-        q,
-        k,
-        head_size,
-        cos_sin_cache,
-        not interleaved,  # sgl-kernel positions op uses is_neox (split-halves) flag
-    )
-
+def _assert_close_gpu(actual, expected, atol: float, rtol: float, name: str):
+    diff = (actual - expected).abs()
+    max_abs = float(diff.max().item())
+    denom = atol + rtol * expected.abs()
+    max_rel = float((diff / denom).max().item())
+    if max_abs > atol and max_rel > 1.0:
+        raise AssertionError(
+            f"{name}: not close max_abs={max_abs:.6g} max_rel={max_rel:.6g}"
+        )
 
-def sglang_jit_rotary_cos_sin(
-    cos: torch.Tensor,
-    sin: torch.Tensor,
-    q: torch.Tensor,
-    k: torch.Tensor,
-    head_size: int,
-    interleaved: bool,
-) -> None:
-    from sglang.jit_kernel.rotary_embedding import rotary_embedding_cos_sin
 
-    rotary_embedding_cos_sin(cos, sin, q, k, head_size, interleaved)
-
-
-def sglang_jit_rotary_with_positions(
-    positions: torch.Tensor,
-    cos_cache: torch.Tensor,
-    sin_cache: torch.Tensor,
-    q: torch.Tensor,
-    k: torch.Tensor,
-    head_size: int,
-    interleaved: bool,
-) -> None:
-    """
-    JIT kernel with fused gathering: positions + full cache.
-    This is the optimized version that avoids explicit index_select in Python.
-    """
+def sgl_jit_fused(
+    positions, cos_cache, sin_cache, q, k, head_size: int, interleaved: bool
+):
     from sglang.jit_kernel.rotary_embedding import rotary_embedding_cos_sin
 
     rotary_embedding_cos_sin(
@@ -255,281 +175,197 @@ def sglang_jit_rotary_with_positions(
     )
 
 
-def sglang_jit_rotary_with_gather(
-    positions: torch.Tensor,
-    cos_cache: torch.Tensor,
-    sin_cache: torch.Tensor,
-    q: torch.Tensor,
-    k: torch.Tensor,
-    head_size: int,
-    interleaved: bool,
-) -> None:
-    """
-    JIT kernel WITHOUT fused gathering: explicit index_select in Python.
-    This is the baseline version that includes the gathering overhead.
-    """
+def sgl_jit_unfused(
+    positions, cos_cache, sin_cache, q, k, head_size: int, interleaved: bool
+):
     from sglang.jit_kernel.rotary_embedding import rotary_embedding_cos_sin
 
-    # Explicit gathering in Python (the overhead we want to avoid)
-    cos_gathered = cos_cache[positions].contiguous()
-    sin_gathered = sin_cache[positions].contiguous()
-    rotary_embedding_cos_sin(
-        cos_gathered, sin_gathered, q, k, head_size, interleaved
-    )
-
-
-BASE_PROVIDER = "jit_fused"
-BASE_NAME = "SGL JIT (fused)"
-BASE_STYLE = ("blue", "-")
-
-BASE_LINE_VALS = [BASE_PROVIDER]
-BASE_LINE_NAMES = [BASE_NAME]
-BASE_STYLES = [BASE_STYLE]
-
-LINE_VALS = list(BASE_LINE_VALS)
-LINE_NAMES = list(BASE_LINE_NAMES)
-STYLES = list(BASE_STYLES)
-
-# Add unfused version for comparison
-LINE_VALS.append("jit_unfused")
-LINE_NAMES.append("SGL JIT (unfused)")
-STYLES.append(("purple", ":"))
-
-if HAS_SGL_POS:
-    LINE_VALS.append("aot_pos")
-    LINE_NAMES.append("SGL AOT (positions)")
-    STYLES.append(("orange", "--"))
-
-if HAS_VLLM:
-    LINE_VALS.append("vllm_pos")
-    LINE_NAMES.append("vLLM (positions)")
-    STYLES.append(("green", "-."))
-
-if HAS_FLASHINFER:
-    LINE_VALS.append("flashinfer_rope")
-    LINE_NAMES.append("FlashInfer RoPE")
-    STYLES.append(("red", "--"))
-
-if ONLY_PROVIDER is not None:
-    if ONLY_PROVIDER not in LINE_VALS:
-        raise ValueError(
-            f"Unknown SGL_BENCH_PROVIDER={ONLY_PROVIDER}. Allowed: {LINE_VALS}"
-        )
-    idx = LINE_VALS.index(ONLY_PROVIDER)
-    LINE_VALS = [ONLY_PROVIDER]
-    LINE_NAMES = [LINE_NAMES[idx]]
-    STYLES = [STYLES[idx]]
-
-configs = list(
-    itertools.product(BS_RANGE, SEQ_RANGE, HEAD_SIZE_RANGE, INTERLEAVED_RANGE)
-)
-_SANITY_DONE = False
-
+    cos_g = cos_cache[positions].contiguous()
+    sin_g = sin_cache[positions].contiguous()
+    rotary_embedding_cos_sin(cos_g, sin_g, q, k, head_size, interleaved)
 
-def _assert_close_gpu(
-    actual: torch.Tensor,
-    expected: torch.Tensor,
-    *,
-    atol: float,
-    rtol: float,
-    name: str,
-) -> None:
-    if actual.shape != expected.shape:
-        raise AssertionError(
-            f"{name}: shape mismatch {tuple(actual.shape)} vs {tuple(expected.shape)}"
-        )
-    if actual.dtype != expected.dtype:
-        raise AssertionError(
-            f"{name}: dtype mismatch {actual.dtype} vs {expected.dtype}"
-        )
-    if actual.device != expected.device:
-        raise AssertionError(
-            f"{name}: device mismatch {actual.device} vs {expected.device}"
-        )
 
-    diff = (actual - expected).abs()
-    max_abs = float(diff.max().item())
-    denom = atol + rtol * expected.abs()
-    max_rel = float((diff / denom).max().item())
-    if max_abs > atol and max_rel > 1.0:
-        raise AssertionError(
-            f"{name}: not close (atol={atol} rtol={rtol}) max_abs={max_abs:.6g} max_rel={max_rel:.6g}"
-        )
+def sgl_aot_pos(positions, q, k, head_size: int, interleaved: bool, cos_sin_cache):
+    if not HAS_SGL_POS:
+        raise RuntimeError("torch.ops.sgl_kernel.rotary_embedding not available")
+    torch.ops.sgl_kernel.rotary_embedding(
+        positions, q, k, head_size, cos_sin_cache, not interleaved  # is_neox flag
+    )
 
 
-@lru_cache(maxsize=None)
-def _get_vllm_rope(
-    head_size: int, rotary_dim: int, interleaved: bool, dtype: torch.dtype
-):
-    if not HAS_VLLM or vLLMRotaryEmbedding is None:
-        raise RuntimeError("vLLM not available")
-    return vLLMRotaryEmbedding(
-        head_size=head_size,
-        rotary_dim=rotary_dim,
-        max_position_embeddings=MAX_SEQ_LEN,
-        base=10000,
-        is_neox_style=interleaved,
-        dtype=dtype,
-    ).cuda()
-
+def _is_flashinfer_unsupported(e: BaseException) -> bool:
+    msg = str(e)
+    return ("Unsupported head_dim" in msg) or ("cos_sin_cache should be float32" in msg)
 
-@triton.testing.perf_report(
-    triton.testing.Benchmark(
-        x_names=["batch_size", "seq_len", "head_size", "interleaved"],
-        x_vals=configs,
-        line_arg="provider",
-        line_vals=LINE_VALS,
-        line_names=LINE_NAMES,
-        styles=STYLES,
-        ylabel="us",
-        plot_name="rotary-embedding-performance",
-        args={},
-    )
-)
-def benchmark(
-    batch_size: int,
-    seq_len: int,
-    head_size: int,
-    interleaved: bool,
-    provider: str,
-) -> Tuple[float, float, float]:
-    if DEVICE == "cuda" and not torch.cuda.is_available():
-        raise RuntimeError("CUDA not available")
 
+@dataclass
+class Case:
+    positions: torch.Tensor
+    cos_cache: torch.Tensor
+    sin_cache: torch.Tensor
+    cos_gathered: torch.Tensor
+    sin_gathered: torch.Tensor
+    cos_sin_cache_aot: torch.Tensor
+    q_base: torch.Tensor  # [T, H, D]
+    k_base: torch.Tensor  # [T, H, D]
+    q4_base: torch.Tensor  # [B,S,H,D] for flashinfer
+    k4_base: torch.Tensor
+    cos_sin_cache_flashinfer: torch.Tensor | None
+
+
+def prepare_case(
+    batch_size: int, seq_len: int, head_size: int, interleaved: bool
+) -> Case:
     num_tokens = batch_size * seq_len
     rotary_dim = head_size
 
-    # Prepare full cache (for fused gathering tests)
-    cos_cache_half, sin_cache_half = _get_cos_sin_cache_half_cuda(rotary_dim)
-    cos_sin_cache_aot = torch.cat([cos_cache_half, sin_cache_half], dim=-1).contiguous()
-    
-    # FlashInfer requires cos_sin_cache to be float32.
-    if HAS_FLASHINFER:
-        cos_f32, sin_f32 = _compute_cos_sin_cache_half(
-            MAX_SEQ_LEN, rotary_dim, dtype=torch.float32
-        )
-        cos_sin_cache_flashinfer = (
-            torch.cat([cos_f32, sin_f32], dim=-1)
-            .to(device=DEVICE, dtype=torch.float32)
-            .contiguous()
-        )
-
-    # Prepare positions (simulate real decoding scenario)
-    # Use arange for simplicity, but in real scenario this could be non-contiguous
+    cos_cache_half, sin_cache_half = _cos_sin_cache_half_cuda(rotary_dim)
     positions = torch.arange(seq_len, device=DEVICE, dtype=torch.int64).repeat(
         batch_size
     )
-    
-    # Pre-gathered cos/sin for unfused tests (traditional approach)
+
     cos_half = cos_cache_half[positions].contiguous()
     sin_half = sin_cache_half[positions].contiguous()
 
     if interleaved:
-        cos_gathered, sin_gathered = cos_half, sin_half
-        # For fused tests, prepare cache in interleaved format
-        cos_cache_for_fused = cos_cache_half
-        sin_cache_for_fused = sin_cache_half
+        cos_cache, sin_cache = cos_cache_half, sin_cache_half
+        cos_g, sin_g = cos_half, sin_half
     else:
-        cos_gathered = torch.cat([cos_half, cos_half], dim=-1).contiguous()
-        sin_gathered = torch.cat([sin_half, sin_half], dim=-1).contiguous()
-        # For fused tests, prepare cache in non-interleaved format
-        cos_cache_for_fused = torch.cat([cos_cache_half, cos_cache_half], dim=-1).contiguous()
-        sin_cache_for_fused = torch.cat([sin_cache_half, sin_cache_half], dim=-1).contiguous()
-
-    # Keep a canonical 4D layout for providers that require it (e.g. FlashInfer).
-    # View as 3D for providers that accept [num_tokens, nheads, head_size].
+        # non-interleaved: kernel expects R=2*embed_dim (x and y halves)
+        cos_cache = torch.cat([cos_cache_half, cos_cache_half], dim=-1).contiguous()
+        sin_cache = torch.cat([sin_cache_half, sin_cache_half], dim=-1).contiguous()
+        cos_g = torch.cat([cos_half, cos_half], dim=-1).contiguous()
+        sin_g = torch.cat([sin_half, sin_half], dim=-1).contiguous()
+
+    cos_sin_cache_aot = torch.cat([cos_cache_half, sin_cache_half], dim=-1).contiguous()
+
     q4 = torch.randn(
-        batch_size, seq_len, NUM_Q_HEADS, head_size, device=DEVICE, dtype=DTYPE
+        batch_size, seq_len, NUM_HEADS, head_size, device=DEVICE, dtype=DTYPE
     ).contiguous()
     k4 = torch.randn(
-        batch_size, seq_len, NUM_Q_HEADS, head_size, device=DEVICE, dtype=DTYPE
+        batch_size, seq_len, NUM_HEADS, head_size, device=DEVICE, dtype=DTYPE
     ).contiguous()
-    q = q4.view(num_tokens, NUM_Q_HEADS, head_size)
-    k = k4.view(num_tokens, NUM_Q_HEADS, head_size)
+    q = q4.view(num_tokens, NUM_HEADS, head_size)
+    k = k4.view(num_tokens, NUM_HEADS, head_size)
 
-    global _SANITY_DONE
-    if (
-        (not _SANITY_DONE)
-        and provider == "jit_fused"
-        and batch_size == BS_RANGE[0]
-        and seq_len == SEQ_RANGE[0]
-        and head_size == HEAD_SIZE_RANGE[0]
-        and interleaved == INTERLEAVED_RANGE[0]
-    ):
-        # Reference implementation
-        q_ref = q.clone()
-        k_ref = k.clone()
-        torch_impl_rotary_fp32(cos_gathered, sin_gathered, q_ref, k_ref, head_size, interleaved)
-        
-        # Test fused version
-        q_fused_test = q.clone()
-        k_fused_test = k.clone()
-        sglang_jit_rotary_with_positions(
-            positions, cos_cache_for_fused, sin_cache_for_fused, 
-            q_fused_test, k_fused_test, head_size, interleaved
+    cos_sin_cache_flashinfer = None
+    if HAS_FLASHINFER:
+        cos_f32, sin_f32 = _compute_cos_sin_cache(
+            MAX_SEQ_LEN, rotary_dim, dtype=torch.float32
         )
-        
-        # Test unfused version
-        q_unfused_test = q.clone()
-        k_unfused_test = k.clone()
-        sglang_jit_rotary_with_gather(
-            positions, cos_cache_for_fused, sin_cache_for_fused,
-            q_unfused_test, k_unfused_test, head_size, interleaved
+        cos_sin_cache_flashinfer = (
+            torch.cat([cos_f32, sin_f32], dim=-1)
+            .to(device=DEVICE, dtype=torch.float32)
+            .contiguous()
         )
-        
-        ref_atol = 2e-2 if DTYPE == torch.bfloat16 else 2e-3
-        ref_rtol = 2e-2 if DTYPE == torch.bfloat16 else 2e-3
-        _assert_close_gpu(q_fused_test, q_ref, atol=ref_atol, rtol=ref_rtol, name="JIT-fused(q)")
-        _assert_close_gpu(k_fused_test, k_ref, atol=ref_atol, rtol=ref_rtol, name="JIT-fused(k)")
-        _assert_close_gpu(q_unfused_test, q_ref, atol=ref_atol, rtol=ref_rtol, name="JIT-unfused(q)")
-        _assert_close_gpu(k_unfused_test, k_ref, atol=ref_atol, rtol=ref_rtol, name="JIT-unfused(k)")
-        print("✅ Sanity check passed: fused and unfused versions match reference!")
-        _SANITY_DONE = True
 
-    FN_MAP = {
-        "aot_pos": lambda: sglang_aot_rotary_positions(
-            positions, q, k, head_size, interleaved, cos_sin_cache_aot
-        ),
-        # NEW: Fused gathering - positions + full cache (OPTIMIZED)
-        "jit_fused": lambda: sglang_jit_rotary_with_positions(
-            positions, cos_cache_for_fused, sin_cache_for_fused, q, k, head_size, interleaved
-        ),
-        # NEW: Unfused - explicit gathering in Python (BASELINE)
-        "jit_unfused": lambda: sglang_jit_rotary_with_gather(
-            positions, cos_cache_for_fused, sin_cache_for_fused, q, k, head_size, interleaved
-        ),
-        # OLD: Pre-gathered cos/sin (for backward compatibility)
-        "jit_cos_sin": lambda: sglang_jit_rotary_cos_sin(
-            cos_gathered, sin_gathered, q, k, head_size, interleaved
-        ),
-        "torch_fp32": lambda: torch_impl_rotary_fp32(
-            cos_gathered, sin_gathered, q, k, head_size, interleaved
-        ),
-    }
-    if HAS_FLASHINFER and apply_flashinfer_rope_qk_inplace is not None:
-        FN_MAP["flashinfer_rope"] = lambda: apply_flashinfer_rope_qk_inplace(
+    return Case(
+        positions=positions,
+        cos_cache=cos_cache,
+        sin_cache=sin_cache,
+        cos_gathered=cos_g,
+        sin_gathered=sin_g,
+        cos_sin_cache_aot=cos_sin_cache_aot,
+        q_base=q,
+        k_base=k,
+        q4_base=q4,
+        k4_base=k4,
+        cos_sin_cache_flashinfer=cos_sin_cache_flashinfer,
+    )
+
+
+def make_fn(
+    provider: str, case: Case, head_size: int, interleaved: bool
+) -> Callable[[], None]:
+    # 统一 clone：每次 bench 都从同一份 base 开始，避免“越跑越旋”
+    if provider in ("jit_fused", "jit_unfused", "aot_pos", "torch_fp32", "vllm_pos"):
+        q = case.q_base.clone()
+        k = case.k_base.clone()
+
+        if provider == "jit_fused":
+            return lambda: sgl_jit_fused(
+                case.positions,
+                case.cos_cache,
+                case.sin_cache,
+                q,
+                k,
+                head_size,
+                interleaved,
+            )
+        if provider == "jit_unfused":
+            return lambda: sgl_jit_unfused(
+                case.positions,
+                case.cos_cache,
+                case.sin_cache,
+                q,
+                k,
+                head_size,
+                interleaved,
+            )
+        if provider == "aot_pos":
+            return lambda: sgl_aot_pos(
+                case.positions, q, k, head_size, interleaved, case.cos_sin_cache_aot
+            )
+        if provider == "torch_fp32":
+            return lambda: torch_impl_rotary_fp32(
+                case.cos_gathered, case.sin_gathered, q, k, head_size, interleaved
+            )
+        if provider == "vllm_pos":
+            rope = _vllm_rope(head_size, head_size, interleaved, DTYPE)
+            return lambda: rope.forward_cuda(case.positions, q, k)
+
+    if provider == "flashinfer":
+        if not (HAS_FLASHINFER and apply_flashinfer_rope_qk_inplace is not None):
+            raise RuntimeError("FlashInfer not available")
+        q4 = case.q4_base.clone()
+        k4 = case.k4_base.clone()
+        return lambda: apply_flashinfer_rope_qk_inplace(
             q4,
             k4,
-            cos_sin_cache_flashinfer,
+            case.cos_sin_cache_flashinfer,
             is_neox=not interleaved,
+            positions=case.positions,
         )
+
+    raise ValueError(f"Unknown provider={provider}")
+
+
+def bench(fn: Callable[[], None], use_cudagraph: bool) -> Tuple[float, float, float]:
+    qs = [0.5, 0.2, 0.8]
+    if use_cudagraph:
+        return triton.testing.do_bench_cudagraph(fn, quantiles=qs)  # type: ignore
+    return triton.testing.do_bench(fn, quantiles=qs)  # type: ignore
+
+
+def available_providers() -> Dict[str, str]:
+    ps: Dict[str, str] = {
+        "jit_fused": "SGL JIT (fused gather)",
+        "jit_unfused": "SGL JIT (Python gather)",
+    }
+    if HAS_SGL_POS:
+        ps["aot_pos"] = "SGL AOT (fused gather)"
     if HAS_VLLM:
-        vllm_rope = _get_vllm_rope(head_size, rotary_dim, interleaved, DTYPE)
-        FN_MAP["vllm_pos"] = lambda: vllm_rope.forward_cuda(positions, q, k)
-
-    fn = FN_MAP[provider]
-    quantiles = [0.5, 0.2, 0.8]
-    # If FlashInfer is available, unify timing across providers by using do_bench for all.
-    use_cudagraph = not HAS_FLASHINFER
-    ms, min_ms, max_ms = _bench_provider(
-        provider, fn, quantiles=quantiles, use_cudagraph=use_cudagraph
-    )
-    return 1000 * ms, 1000 * max_ms, 1000 * min_ms
+        ps["vllm_pos"] = "vLLM"
+    if HAS_FLASHINFER:
+        ps["flashinfer"] = "FlashInfer"
+    if INCLUDE_TORCH_REF:
+        ps["torch_fp32"] = "Torch"
+    return ps
 
 
-SPEEDUP_LINE_VALS = [p for p in LINE_VALS if p != BASE_PROVIDER]
-SPEEDUP_LINE_NAMES = [LINE_NAMES[LINE_VALS.index(p)] for p in SPEEDUP_LINE_VALS]
-SPEEDUP_STYLES = [STYLES[LINE_VALS.index(p)] for p in SPEEDUP_LINE_VALS]
+PROVIDERS = available_providers()
+if ONLY_PROVIDER is not None:
+    if ONLY_PROVIDER not in PROVIDERS:
+        raise ValueError(
+            f"Unknown provider={ONLY_PROVIDER}. Allowed: {list(PROVIDERS)}"
+        )
+    PROVIDERS = {ONLY_PROVIDER: PROVIDERS[ONLY_PROVIDER]}
+
+configs = list(
+    itertools.product(BS_RANGE, SEQ_RANGE, HEAD_SIZE_RANGE, INTERLEAVED_RANGE)
+)
+
+_SANITY_DONE = False
 
 
 @triton.testing.perf_report(
@@ -537,137 +373,75 @@ def benchmark(
         x_names=["batch_size", "seq_len", "head_size", "interleaved"],
         x_vals=configs,
         line_arg="provider",
-        line_vals=SPEEDUP_LINE_VALS,
-        line_names=SPEEDUP_LINE_NAMES,
-        styles=SPEEDUP_STYLES,
-        ylabel=f"speedup (x) vs {BASE_NAME}",
-        plot_name=f"rotary-embedding-speedup-vs-{BASE_PROVIDER}",
+        line_vals=list(PROVIDERS.keys()),
+        line_names=list(PROVIDERS.values()),
+        styles=None,
+        ylabel="us",
+        plot_name="rotary-embedding-performance",
         args={},
     )
 )
-def benchmark_speedup_vs_jit(
-    batch_size: int,
-    seq_len: int,
-    head_size: int,
-    interleaved: bool,
-    provider: str,
-) -> Tuple[float, float, float]:
-    if provider == BASE_PROVIDER:
-        return 1.0, 1.0, 1.0
+def benchmark(
+    batch_size: int, seq_len: int, head_size: int, interleaved: bool, provider: str
+):
     if DEVICE == "cuda" and not torch.cuda.is_available():
         raise RuntimeError("CUDA not available")
 
-    num_tokens = batch_size * seq_len
-    rotary_dim = head_size
+    case = prepare_case(batch_size, seq_len, head_size, interleaved)
 
-    # Prepare full cache
-    cos_cache_half, sin_cache_half = _get_cos_sin_cache_half_cuda(rotary_dim)
-    cos_sin_cache_aot = torch.cat([cos_cache_half, sin_cache_half], dim=-1).contiguous()
-    if HAS_FLASHINFER:
-        cos_f32, sin_f32 = _compute_cos_sin_cache_half(
-            MAX_SEQ_LEN, rotary_dim, dtype=torch.float32
-        )
-        cos_sin_cache_flashinfer = (
-            torch.cat([cos_f32, sin_f32], dim=-1)
-            .to(device=DEVICE, dtype=torch.float32)
-            .contiguous()
+    global _SANITY_DONE
+    if not _SANITY_DONE and provider == "jit_fused":
+        # sanity: compare jit_fused / jit_unfused with fp32 ref
+        q_ref = case.q_base.clone()
+        k_ref = case.k_base.clone()
+        torch_impl_rotary_fp32(
+            case.cos_gathered, case.sin_gathered, q_ref, k_ref, head_size, interleaved
         )
 
-    positions = torch.arange(seq_len, device=DEVICE, dtype=torch.int64).repeat(
-        batch_size
-    )
-    cos_half = cos_cache_half[positions].contiguous()
-    sin_half = sin_cache_half[positions].contiguous()
+        for p in ("jit_fused", "jit_unfused"):
+            q = case.q_base.clone()
+            k = case.k_base.clone()
+            if p == "jit_fused":
+                sgl_jit_fused(
+                    case.positions,
+                    case.cos_cache,
+                    case.sin_cache,
+                    q,
+                    k,
+                    head_size,
+                    interleaved,
+                )
+            else:
+                sgl_jit_unfused(
+                    case.positions,
+                    case.cos_cache,
+                    case.sin_cache,
+                    q,
+                    k,
+                    head_size,
+                    interleaved,
+                )
+
+            atol = 2e-2 if DTYPE == torch.bfloat16 else 2e-3
+            rtol = 2e-2 if DTYPE == torch.bfloat16 else 2e-3
+            _assert_close_gpu(q, q_ref, atol=atol, rtol=rtol, name=f"{p}(q)")
+            _assert_close_gpu(k, k_ref, atol=atol, rtol=rtol, name=f"{p}(k)")
 
-    if interleaved:
-        cos_gathered, sin_gathered = cos_half, sin_half
-        cos_cache_for_fused = cos_cache_half
-        sin_cache_for_fused = sin_cache_half
-    else:
-        cos_gathered = torch.cat([cos_half, cos_half], dim=-1).contiguous()
-        sin_gathered = torch.cat([sin_half, sin_half], dim=-1).contiguous()
-        cos_cache_for_fused = torch.cat([cos_cache_half, cos_cache_half], dim=-1).contiguous()
-        sin_cache_for_fused = torch.cat([sin_cache_half, sin_cache_half], dim=-1).contiguous()
+        _SANITY_DONE = True
 
-    q_base4 = torch.randn(
-        batch_size, seq_len, NUM_Q_HEADS, head_size, device=DEVICE, dtype=DTYPE
-    ).contiguous()
-    k_base4 = torch.randn(
-        batch_size, seq_len, NUM_Q_HEADS, head_size, device=DEVICE, dtype=DTYPE
-    ).contiguous()
-    q_base = q_base4.view(num_tokens, NUM_Q_HEADS, head_size)
-    k_base = k_base4.view(num_tokens, NUM_Q_HEADS, head_size)
-
-    q_fused = q_base.clone()
-    k_fused = k_base.clone()
-    q_unfused = q_base.clone()
-    k_unfused = k_base.clone()
-    q_jit = q_base.clone()
-    k_jit = k_base.clone()
-    q_p = q_base.clone()
-    k_p = k_base.clone()
-
-    FN_MAP = {
-        "jit_fused": lambda: sglang_jit_rotary_with_positions(
-            positions, cos_cache_for_fused, sin_cache_for_fused, q_fused, k_fused, head_size, interleaved
-        ),
-        "jit_unfused": lambda: sglang_jit_rotary_with_gather(
-            positions, cos_cache_for_fused, sin_cache_for_fused, q_unfused, k_unfused, head_size, interleaved
-        ),
-        "jit_cos_sin": lambda: sglang_jit_rotary_cos_sin(
-            cos_gathered, sin_gathered, q_jit, k_jit, head_size, interleaved
-        ),
-        "aot_pos": lambda: sglang_aot_rotary_positions(
-            positions, q_p, k_p, head_size, interleaved, cos_sin_cache_aot
-        ),
-        "torch_fp32": lambda: torch_impl_rotary_fp32(
-            cos_gathered, sin_gathered, q_p, k_p, head_size, interleaved
-        ),
-    }
-    if HAS_FLASHINFER and apply_flashinfer_rope_qk_inplace is not None:
-        q_f = q_base4.clone()
-        k_f = k_base4.clone()
-        FN_MAP["flashinfer_rope"] = lambda: apply_flashinfer_rope_qk_inplace(
-            q_f,
-            k_f,
-            cos_sin_cache_flashinfer,
-            is_neox=not interleaved,
-        )
-    if HAS_VLLM:
-        vllm_rope = _get_vllm_rope(head_size, rotary_dim, interleaved, DTYPE)
-        FN_MAP["vllm_pos"] = lambda: vllm_rope.forward_cuda(positions, q_p, k_p)
-
-    quantiles = [0.5, 0.2, 0.8]
-    # If FlashInfer is available, use do_bench for all providers to keep methodology
-    # consistent and avoid CUDA-graph sensitivity.
-    use_cudagraph = not HAS_FLASHINFER
-
-    base_ms, base_min_ms, base_max_ms = _bench_provider(
-        BASE_PROVIDER,
-        FN_MAP[BASE_PROVIDER],
-        quantiles=quantiles,
-        use_cudagraph=use_cudagraph,
-    )
-    if base_ms != base_ms:  # NaN
-        nan = float("nan")
-        return nan, nan, nan
+    fn = make_fn(provider, case, head_size, interleaved)
 
-    p_ms, p_min_ms, p_max_ms = _bench_provider(
-        provider, FN_MAP[provider], quantiles=quantiles, use_cudagraph=use_cudagraph
-    )
+    use_cudagraph = provider != "flashinfer"
+    try:
+        ms, min_ms, max_ms = bench(fn, use_cudagraph=use_cudagraph)
+    except Exception as e:
+        if provider == "flashinfer" and _is_flashinfer_unsupported(e):
+            nan = float("nan")
+            return nan, nan, nan
+        raise
 
-    speed_med = float(base_ms / p_ms)
-    speed_max = float(base_max_ms / p_min_ms)
-    speed_min = float(base_min_ms / p_max_ms)
-    return speed_med, speed_max, speed_min
+    return 1000 * ms, 1000 * min_ms, 1000 * max_ms
 
 
 if __name__ == "__main__":
     benchmark.run(print_data=True)
-    if (
-        SHOW_SPEEDUP
-        and (BASE_PROVIDER in LINE_VALS)
-        and (len(LINE_VALS) > 1)
-        and (ONLY_PROVIDER is None)
-    ):
-        benchmark_speedup_vs_jit.run(print_data=True)
diff --git a/python/sglang/jit_kernel/csrc/rotary_embedding_cos_sin.cuh b/python/sglang/jit_kernel/csrc/rotary_embedding_cos_sin.cuh
index 2afed22d9e89..08a00124cfff 100644
--- a/python/sglang/jit_kernel/csrc/rotary_embedding_cos_sin.cuh
+++ b/python/sglang/jit_kernel/csrc/rotary_embedding_cos_sin.cuh
@@ -64,14 +64,8 @@ struct RotaryVecData<scalar_t, false> {
   float4 sin_y;
 };
 
-// Vector load helper:
-//   - load a 16B tile from q/k (either aligned float4 or scalar gather)
-//   - always vector-load cos/sin (aligned by launch-time checks)
-//
-// rot_offset is the "pair index":
-//   - interleaved:     pair i -> q[2*i], q[2*i+1]
-//   - non-interleaved: pair i -> q[i], q[i + embed_dim]
-// embed_dim is the number of pairs per head.
+// Per-thread temporary storage for a 16B tile and its cos/sin values.
+// Layout depends on interleaved vs split-halves format.
 template <typename scalar_t, bool interleaved, bool aligned_qk>
 __device__ __forceinline__ RotaryVecData<scalar_t, interleaved> load_rotary_vec(
     const scalar_t* __restrict__ arr,
@@ -132,16 +126,7 @@ __device__ __forceinline__ RotaryVecData<scalar_t, interleaved> load_rotary_vec(
   return data;
 }
 
-// Vector compute + store:
-// Apply RoPE to the loaded 16B tile and write back to q/k.
-//
-// Interleaved math (per pair):
-//   x' = x*cos - y*sin
-//   y' = y*cos + x*sin
-//
-// Non-interleaved math (general form, allows distinct cos/sin for x and y halves):
-//   x' = x*cos_x - y*sin_x
-//   y' = y*cos_y + x*sin_y
+// Apply RoPE to the loaded 16B tile and write back to q/k
 template <typename scalar_t, bool interleaved, bool aligned_qk>
 __device__ __forceinline__ void compute_store_rotary_vec(
     scalar_t* __restrict__ arr, const RotaryVecData<scalar_t, interleaved>& data, int rot_offset, int embed_dim) {
@@ -263,12 +248,9 @@ inline __device__ void apply_token_rotary_embedding(
   }
 }
 
-// 2D-grid kernel:
-//   blockIdx.x -> token index
-//   blockIdx.y -> "sub-block" index within the token (tile along pairs dimension)
-//
-// For very small T (few tokens) but large per-token work, using multiple blocks
-// per token can improve occupancy/throughput compared to one-block-per-token.
+// RoPE kernel with 2D grid: (token_idx, sub_block_idx).
+// Each block processes a strided subset of rotary pairs across all heads for one token.
+// Supports optional key tensor (nullptr means Q-only).
 template <typename scalar_t, bool interleaved, bool vectorized, bool aligned_qk, int ROT_EMBED_DIM>
 __global__ void rotary_embedding_kernel_2d(
     const scalar_t* __restrict__ cos_data,  // [num_tokens, rot_dim_arg]
@@ -399,11 +381,7 @@ __global__ void rotary_embedding_kernel_2d(
   }
 }
 
-// 1D-grid kernel:
-//   blockIdx.x -> token index
-//   exactly one block per token
-//
-// The default path for most sizes.
+// RoPE kernel with 1D grid: one block per token.
 template <typename scalar_t, bool interleaved, bool vectorized, bool aligned_qk, int ROT_EMBED_DIM>
 __launch_bounds__(512) __global__ void rotary_embedding_kernel_1d(
     const scalar_t* __restrict__ cos_data,  // [num_tokens, rot_dim_arg]
@@ -452,7 +430,7 @@ __launch_bounds__(512) __global__ void rotary_embedding_kernel_1d(
 
     int next_i = i + stride;
     for (; i < nq_pairs; i += stride, next_i += stride) {
-      RotaryVecData<scalar_t, interleaved> next_data;
+      RotaryVecData<scalar_t, interleaved> next_data = curr_data;
       if (next_i < nq_pairs) {
         const int head_idx = next_i / embed_dim_for_rotation;
         const int rot_offset = next_i % embed_dim_for_rotation;
@@ -485,7 +463,7 @@ __launch_bounds__(512) __global__ void rotary_embedding_kernel_1d(
 
       int next_k_i = k_i + stride;
       for (; k_i < nk_pairs; k_i += stride, next_k_i += stride) {
-        RotaryVecData<scalar_t, interleaved> next_data_k;
+        RotaryVecData<scalar_t, interleaved> next_data_k = curr_data_k;
         if (next_k_i < nk_pairs) {
           const int head_idx = next_k_i / embed_dim_for_rotation;
           const int rot_offset = next_k_i % embed_dim_for_rotation;
@@ -527,19 +505,15 @@ __launch_bounds__(512) __global__ void rotary_embedding_kernel_1d(
   }
 }
 
-// Kernel variant dispatcher:
-// Select one of:
-//   - grid shape: 2D (multi-block per token) vs 1D (one-block per token)
-//   - layout: interleaved vs non-interleaved
-//   - compute: vectorized 16B tiles vs scalar
-//   - memory: q/k aligned 16B load/store vs scalar gather/scatter
+// Dispatch kernel specialization by layout (interleaved), vectorization, and alignment.
+// Selects 1D vs 2D grid based on launch configuration.
 template <int ROT, typename scalar_t>
 __forceinline__ void dispatch_rotary_launch(
     bool use_grid_2d,
     dim3 grid2d,
     dim3 grid1d,
     dim3 block,
-    const decltype(host::LaunchKernel::resolve_device(std::declval<DLDevice>())) stream,
+    cudaStream_t stream,
     bool interleaved,
     bool use_vec,
     bool qk_aligned16,
@@ -558,169 +532,35 @@ __forceinline__ void dispatch_rotary_launch(
     int num_kv_heads,
     int head_size,
     int blocks_per_token) {
-  // 2D grid path
-  if (use_grid_2d) {
+  auto dispatch_all = [&](auto f) {
     if (interleaved) {
       if (use_vec) {
-        if (qk_aligned16) {
-          host::LaunchKernel(grid2d, block, stream)(
-              rotary_embedding_kernel_2d<scalar_t, true, true, true, ROT>,
-              cos_ptr,
-              sin_ptr,
-              q_ptr,
-              k_ptr,
-              positions_ptr,
-              rot_dim_from_cache,
-              embed_dim_for_rotation,
-              query_token_stride,
-              key_token_stride,
-              head_stride_query,
-              head_stride_key,
-              num_heads,
-              num_kv_heads,
-              head_size,
-              blocks_per_token);
-        } else {
-          host::LaunchKernel(grid2d, block, stream)(
-              rotary_embedding_kernel_2d<scalar_t, true, true, false, ROT>,
-              cos_ptr,
-              sin_ptr,
-              q_ptr,
-              k_ptr,
-              positions_ptr,
-              rot_dim_from_cache,
-              embed_dim_for_rotation,
-              query_token_stride,
-              key_token_stride,
-              head_stride_query,
-              head_stride_key,
-              num_heads,
-              num_kv_heads,
-              head_size,
-              blocks_per_token);
-        }
+        if (qk_aligned16)
+          f(std::true_type{}, std::true_type{}, std::true_type{});
+        else
+          f(std::true_type{}, std::true_type{}, std::false_type{});
       } else {
-        host::LaunchKernel(grid2d, block, stream)(
-            rotary_embedding_kernel_2d<scalar_t, true, false, true, ROT>,
-            cos_ptr,
-            sin_ptr,
-            q_ptr,
-            k_ptr,
-            positions_ptr,
-            rot_dim_from_cache,
-            embed_dim_for_rotation,
-            query_token_stride,
-            key_token_stride,
-            head_stride_query,
-            head_stride_key,
-            num_heads,
-            num_kv_heads,
-            head_size,
-            blocks_per_token);
+        f(std::true_type{}, std::false_type{}, std::true_type{});
       }
     } else {
       if (use_vec) {
-        if (qk_aligned16) {
-          host::LaunchKernel(grid2d, block, stream)(
-              rotary_embedding_kernel_2d<scalar_t, false, true, true, ROT>,
-              cos_ptr,
-              sin_ptr,
-              q_ptr,
-              k_ptr,
-              positions_ptr,
-              rot_dim_from_cache,
-              embed_dim_for_rotation,
-              query_token_stride,
-              key_token_stride,
-              head_stride_query,
-              head_stride_key,
-              num_heads,
-              num_kv_heads,
-              head_size,
-              blocks_per_token);
-        } else {
-          host::LaunchKernel(grid2d, block, stream)(
-              rotary_embedding_kernel_2d<scalar_t, false, true, false, ROT>,
-              cos_ptr,
-              sin_ptr,
-              q_ptr,
-              k_ptr,
-              positions_ptr,
-              rot_dim_from_cache,
-              embed_dim_for_rotation,
-              query_token_stride,
-              key_token_stride,
-              head_stride_query,
-              head_stride_key,
-              num_heads,
-              num_kv_heads,
-              head_size,
-              blocks_per_token);
-        }
+        if (qk_aligned16)
+          f(std::false_type{}, std::true_type{}, std::true_type{});
+        else
+          f(std::false_type{}, std::true_type{}, std::false_type{});
       } else {
-        host::LaunchKernel(grid2d, block, stream)(
-            rotary_embedding_kernel_2d<scalar_t, false, false, true, ROT>,
-            cos_ptr,
-            sin_ptr,
-            q_ptr,
-            k_ptr,
-            positions_ptr,
-            rot_dim_from_cache,
-            embed_dim_for_rotation,
-            query_token_stride,
-            key_token_stride,
-            head_stride_query,
-            head_stride_key,
-            num_heads,
-            num_kv_heads,
-            head_size,
-            blocks_per_token);
+        f(std::false_type{}, std::false_type{}, std::true_type{});
       }
     }
-    return;
-  }
+  };
 
-  // 1D grid path
-  if (interleaved) {
-    if (use_vec) {
-      if (qk_aligned16) {
-        host::LaunchKernel(grid1d, block, stream)(
-            rotary_embedding_kernel_1d<scalar_t, true, true, true, ROT>,
-            cos_ptr,
-            sin_ptr,
-            q_ptr,
-            k_ptr,
-            positions_ptr,
-            rot_dim_from_cache,
-            embed_dim_for_rotation,
-            query_token_stride,
-            key_token_stride,
-            head_stride_query,
-            head_stride_key,
-            num_heads,
-            num_kv_heads,
-            head_size);
-      } else {
-        host::LaunchKernel(grid1d, block, stream)(
-            rotary_embedding_kernel_1d<scalar_t, true, true, false, ROT>,
-            cos_ptr,
-            sin_ptr,
-            q_ptr,
-            k_ptr,
-            positions_ptr,
-            rot_dim_from_cache,
-            embed_dim_for_rotation,
-            query_token_stride,
-            key_token_stride,
-            head_stride_query,
-            head_stride_key,
-            num_heads,
-            num_kv_heads,
-            head_size);
-      }
-    } else {
-      host::LaunchKernel(grid1d, block, stream)(
-          rotary_embedding_kernel_1d<scalar_t, true, false, true, ROT>,
+  if (use_grid_2d) {
+    dispatch_all([&](auto i_v, auto v_v, auto a_v) {
+      constexpr bool I = decltype(i_v)::value;
+      constexpr bool V = decltype(v_v)::value;
+      constexpr bool A = decltype(a_v)::value;
+      host::LaunchKernel(grid2d, block, stream)(
+          rotary_embedding_kernel_2d<scalar_t, I, V, A, ROT>,
           cos_ptr,
           sin_ptr,
           q_ptr,
@@ -734,48 +574,16 @@ __forceinline__ void dispatch_rotary_launch(
           head_stride_key,
           num_heads,
           num_kv_heads,
-          head_size);
-    }
+          head_size,
+          blocks_per_token);
+    });
   } else {
-    if (use_vec) {
-      if (qk_aligned16) {
-        host::LaunchKernel(grid1d, block, stream)(
-            rotary_embedding_kernel_1d<scalar_t, false, true, true, ROT>,
-            cos_ptr,
-            sin_ptr,
-            q_ptr,
-            k_ptr,
-            positions_ptr,
-            rot_dim_from_cache,
-            embed_dim_for_rotation,
-            query_token_stride,
-            key_token_stride,
-            head_stride_query,
-            head_stride_key,
-            num_heads,
-            num_kv_heads,
-            head_size);
-      } else {
-        host::LaunchKernel(grid1d, block, stream)(
-            rotary_embedding_kernel_1d<scalar_t, false, true, false, ROT>,
-            cos_ptr,
-            sin_ptr,
-            q_ptr,
-            k_ptr,
-            positions_ptr,
-            rot_dim_from_cache,
-            embed_dim_for_rotation,
-            query_token_stride,
-            key_token_stride,
-            head_stride_query,
-            head_stride_key,
-            num_heads,
-            num_kv_heads,
-            head_size);
-      }
-    } else {
+    dispatch_all([&](auto i_v, auto v_v, auto a_v) {
+      constexpr bool I = decltype(i_v)::value;
+      constexpr bool V = decltype(v_v)::value;
+      constexpr bool A = decltype(a_v)::value;
       host::LaunchKernel(grid1d, block, stream)(
-          rotary_embedding_kernel_1d<scalar_t, false, false, true, ROT>,
+          rotary_embedding_kernel_1d<scalar_t, I, V, A, ROT>,
           cos_ptr,
           sin_ptr,
           q_ptr,
@@ -790,10 +598,13 @@ __forceinline__ void dispatch_rotary_launch(
           num_heads,
           num_kv_heads,
           head_size);
-    }
+    });
   }
 }
 
+// Validate inputs, choose vectorization/alignment path, and launch the RoPE kernel.
+// cos/sin are indexed by either positions[t] (if provided) or token index t.
+// Rotates the first 2*embed_dim_for_rotation elements of each head.
 template <int ROT, typename scalar_t>
 inline void launch_rotary(
     const tvm::ffi::TensorView cos,
@@ -813,20 +624,14 @@ inline void launch_rotary(
   auto dtype = SymbolicDType{};
   auto device = SymbolicDevice{};
 
-  // Verify q first to establish device/dtype
-  // q: [T, Hq, D]
-  TensorMatcher({T, Hq, D})
-      .with_dtype<float, half, nv_bfloat16>(dtype)
-      .with_device<kDLCUDA>(device)
-      .verify(q);
+  // Verify q: [T, Hq, D] first to establish device/dtype
+  TensorMatcher({T, Hq, D}).with_dtype<float, half, nv_bfloat16>(dtype).with_device<kDLCUDA>(device).verify(q);
 
   const int64_t t = T.unwrap();
   const int64_t hq = Hq.unwrap();
   const int64_t d = D.unwrap();
 
-  // Verify cos/sin
-  // cos: [T_cache, R]
-  // sin: [T_cache, R]
+  // Verify cos/sin: [T_cache, R]
   auto T_cache = SymbolicSize{"T_cache"};
   TensorMatcher({T_cache, R})
       .with_dtype<float, half, nv_bfloat16>(dtype)
@@ -836,23 +641,24 @@ inline void launch_rotary(
 
   const int64_t r = R.unwrap();
 
-  // Handle positions
   const int64_t* positions_ptr = nullptr;
   if (positions != nullptr) {
     auto T_pos = SymbolicSize{"T_pos"};
-    TensorMatcher({T_pos})
-        .with_dtype<int64_t>()
-        .with_device<kDLCUDA>(device)
-        .verify(*positions);
+    TensorMatcher({T_pos}).with_dtype<int64_t>().with_device<kDLCUDA>(device).verify(*positions);
     RuntimeCheck(T_pos.unwrap() == t, "positions length mismatch");
     positions_ptr = static_cast<const int64_t*>(positions->data_ptr());
   } else {
-    // If positions is null, cos/sin length must equal t
-    RuntimeCheck(T_cache.unwrap() == t, "cos/sin length mismatch (expected ", t, ", got ", T_cache.unwrap(), ") when positions is None");
+    // If positions is null, cos/sin length must equal to t
+    RuntimeCheck(
+        T_cache.unwrap() == t,
+        "cos/sin length mismatch (expected ",
+        t,
+        ", got ",
+        T_cache.unwrap(),
+        ") when positions is None");
   }
 
   RuntimeCheck(d == head_size, "head_size mismatch: got ", d, " expected ", head_size);
-
   RuntimeCheck(t > 0 && hq > 0 && d > 0 && r > 0, "invalid shape");
   if (!interleaved) {
     RuntimeCheck(r % 2 == 0, "non-interleaved requires even R, got ", r);
@@ -876,11 +682,7 @@ inline void launch_rotary(
   int64_t head_stride_key = d;
   int hk = 0;
   if (k != nullptr) {
-    // k: [T, Hk, D]
-    TensorMatcher({T, Hk, D})
-        .with_dtype<float, half, nv_bfloat16>(dtype)
-        .with_device<kDLCUDA>(device)
-        .verify(*k);
+    TensorMatcher({T, Hk, D}).with_dtype<float, half, nv_bfloat16>(dtype).with_device<kDLCUDA>(device).verify(*k);
     hk = (int)Hk.unwrap();
     RuntimeCheck(hk > 0, "invalid key shape");
     key_token_stride = (int64_t)hk * d;
@@ -975,23 +777,18 @@ struct RotaryEmbeddingCosSinKernel {
       const tvm::ffi::TensorView cos,
       const tvm::ffi::TensorView sin,
       const tvm::ffi::TensorView query,
-      const tvm::ffi::AnyView positions_any,
+      const tvm::ffi::Optional<tvm::ffi::TensorView> positions,
       int64_t head_size,
       bool interleaved) {
-    // Convert AnyView to TensorView pointer
-    const tvm::ffi::TensorView* positions = nullptr;
-    auto positions_opt = positions_any.try_cast<tvm::ffi::TensorView>();
-    if (positions_opt.has_value()) {
-      positions = &positions_opt.value();
-    }
-    
+    const tvm::ffi::TensorView* positions_ptr = positions.has_value() ? &positions.value() : nullptr;
+
     const auto dt = query.dtype();
     if (host::is_type<half>(dt)) {
-      launch_rotary<ROT, half>(cos, sin, query, nullptr, positions, head_size, interleaved);
+      launch_rotary<ROT, half>(cos, sin, query, nullptr, positions_ptr, head_size, interleaved);
     } else if (host::is_type<nv_bfloat16>(dt)) {
-      launch_rotary<ROT, nv_bfloat16>(cos, sin, query, nullptr, positions, head_size, interleaved);
+      launch_rotary<ROT, nv_bfloat16>(cos, sin, query, nullptr, positions_ptr, head_size, interleaved);
     } else if (host::is_type<float>(dt)) {
-      launch_rotary<ROT, float>(cos, sin, query, nullptr, positions, head_size, interleaved);
+      launch_rotary<ROT, float>(cos, sin, query, nullptr, positions_ptr, head_size, interleaved);
     } else {
       host::Panic("Unsupported dtype for rotary_embedding_cos_sin");
     }
@@ -1002,23 +799,18 @@ struct RotaryEmbeddingCosSinKernel {
       const tvm::ffi::TensorView sin,
       const tvm::ffi::TensorView query,
       const tvm::ffi::TensorView key,
-      const tvm::ffi::AnyView positions_any,
+      const tvm::ffi::Optional<tvm::ffi::TensorView> positions,
       int64_t head_size,
       bool interleaved) {
-    // Convert AnyView to TensorView pointer
-    const tvm::ffi::TensorView* positions = nullptr;
-    auto positions_opt = positions_any.try_cast<tvm::ffi::TensorView>();
-    if (positions_opt.has_value()) {
-      positions = &positions_opt.value();
-    }
-    
+    const tvm::ffi::TensorView* positions_ptr = positions.has_value() ? &positions.value() : nullptr;
+
     const auto dt = query.dtype();
     if (host::is_type<half>(dt)) {
-      launch_rotary<ROT, half>(cos, sin, query, &key, positions, head_size, interleaved);
+      launch_rotary<ROT, half>(cos, sin, query, &key, positions_ptr, head_size, interleaved);
     } else if (host::is_type<nv_bfloat16>(dt)) {
-      launch_rotary<ROT, nv_bfloat16>(cos, sin, query, &key, positions, head_size, interleaved);
+      launch_rotary<ROT, nv_bfloat16>(cos, sin, query, &key, positions_ptr, head_size, interleaved);
     } else if (host::is_type<float>(dt)) {
-      launch_rotary<ROT, float>(cos, sin, query, &key, positions, head_size, interleaved);
+      launch_rotary<ROT, float>(cos, sin, query, &key, positions_ptr, head_size, interleaved);
     } else {
       host::Panic("Unsupported dtype for rotary_embedding_cos_sin");
     }
diff --git a/python/sglang/jit_kernel/rotary_embedding.py b/python/sglang/jit_kernel/rotary_embedding.py
index 2d7402966639..da43dd59e32e 100644
--- a/python/sglang/jit_kernel/rotary_embedding.py
+++ b/python/sglang/jit_kernel/rotary_embedding.py
@@ -13,7 +13,6 @@
 
 @cache_once
 def _jit_rotary_embedding_cos_sin_module(rot: int) -> Module:
-    # Compile one specialized template per rotary embed dim (ROT).
     return load_jit(
         "rotary_embedding_cos_sin",
         str(rot),
@@ -33,63 +32,17 @@ def _jit_rotary_embedding_cos_sin_module(rot: int) -> Module:
 
 @cache_once
 def can_use_rotary_embedding_cos_sin(rot: int) -> bool:
-    logger = logging.getLogger(__name__)
     if rot <= 0:
-        logger.warning(f"Invalid rot={rot} for rotary_embedding_cos_sin")
+        logging.getLogger(__name__).warning(f"Invalid rot={rot}")
         return False
     try:
         _jit_rotary_embedding_cos_sin_module(rot)
         return True
     except Exception as e:
-        logger.warning(f"Failed to load JIT rotary_embedding_cos_sin (rot={rot}): {e}")
+        logging.getLogger(__name__).warning(f"JIT compile failed (rot={rot}): {e}")
         return False
 
 
-def rotary_embedding_cos_sin_q(
-    *,
-    rot: int,
-    cos: torch.Tensor,
-    sin: torch.Tensor,
-    query: torch.Tensor,
-    head_size: int,
-    interleaved: bool,
-    positions: Optional[torch.Tensor] = None,
-) -> None:
-    module = _jit_rotary_embedding_cos_sin_module(rot)
-    module.rotary_embedding_cos_sin_q(
-        cos, sin, query, positions, head_size, interleaved
-    )
-
-
-def rotary_embedding_cos_sin_qk(
-    *,
-    rot: int,
-    cos: torch.Tensor,
-    sin: torch.Tensor,
-    query: torch.Tensor,
-    key: torch.Tensor,
-    head_size: int,
-    interleaved: bool,
-    positions: Optional[torch.Tensor] = None,
-) -> None:
-    module = _jit_rotary_embedding_cos_sin_module(rot)
-    module.rotary_embedding_cos_sin_qk(
-        cos, sin, query, key, positions, head_size, interleaved
-    )
-
-
-def _resolve_interleaved(interleaved: Optional[bool], is_neox: Optional[bool]) -> bool:
-    if interleaved is not None and is_neox is not None and interleaved != is_neox:
-        raise ValueError(
-            f"is_neox({is_neox}) and interleaved({interleaved}) mismatch; keep only one or make them equal."
-        )
-    if is_neox is not None:
-        return is_neox
-    if interleaved is not None:
-        return interleaved
-    return True
-
-
 def rotary_embedding_cos_sin(
     cos: torch.Tensor,
     sin: torch.Tensor,
@@ -101,134 +54,78 @@ def rotary_embedding_cos_sin(
     *,
     is_neox: Optional[bool] = None,
 ) -> None:
-    effective_interleaved = _resolve_interleaved(interleaved, is_neox)
-
-    if query.device.type != "cuda":
-        raise ValueError("query must be a CUDA tensor")
-    if key is not None and key.device.type != "cuda":
-        raise ValueError("key must be a CUDA tensor")
-    if cos.device != query.device or sin.device != query.device:
-        raise ValueError("cos/sin must be on the same device as query")
-
-    if query.dtype != cos.dtype or query.dtype != sin.dtype:
-        raise ValueError("cos/sin dtype must match query dtype")
-    if key is not None and key.dtype != query.dtype:
-        raise ValueError("key dtype must match query dtype")
-
-    if cos.dim() != 2 or sin.dim() != 2:
-        raise ValueError("cos/sin must be 2D tensors: [num_tokens, rot_dim]")
-    if cos.shape != sin.shape:
-        raise ValueError("cos/sin shape mismatch")
-    if query.dim() == 4:
-        expected_tokens = int(query.shape[0] * query.shape[1])
-    elif query.dim() in (2, 3):
-        expected_tokens = int(query.shape[0])
-    else:
-        raise ValueError("query must be a 2D, 3D or 4D tensor")
-    if positions is None and cos.shape[0] != expected_tokens:
-        raise ValueError(
-            f"cos/sin num_tokens mismatch with query: cos.shape[0]={cos.shape[0]} vs expected_tokens={expected_tokens}"
-        )
+    if is_neox is not None:
+        if interleaved is not None and interleaved != is_neox:
+            raise ValueError(f"Mismatch: is_neox={is_neox}, interleaved={interleaved}")
+        interleaved = is_neox
+    interleaved = interleaved if interleaved is not None else True
 
     if head_size == 0:
-        if query.dim() in (3, 4):
-            head_size = int(query.shape[-1])
-        else:
+        if query.dim() < 3:
             raise ValueError("head_size must be provided when query is 2D")
+        head_size = query.shape[-1]
+
+    def _prepare(t: torch.Tensor, name: str) -> torch.Tensor:
+        if t.device.type != "cuda":
+            raise ValueError(f"{name} must be CUDA")
+        if not t.is_contiguous():
+            raise ValueError(f"{name} must be contiguous")
+        if t.dtype != query.dtype:
+            raise ValueError(f"{name} dtype mismatch")
+
+        # Reshape to [tokens, heads, head_size]
+        if t.dim() == 2:
+            return t.view(t.shape[0], -1, head_size)
+        if t.dim() == 3:
+            if t.shape[-1] != head_size:
+                raise ValueError(f"{name} head_size mismatch")
+            return t
+        if t.dim() == 4:
+            if t.shape[-1] != head_size:
+                raise ValueError(f"{name} head_size mismatch")
+            return t.flatten(0, 1)
+        raise ValueError(f"{name} must be 2D, 3D or 4D")
+
+    q_3d = _prepare(query, "query")
+    k_3d = _prepare(key, "key") if key is not None else None
 
-    if not query.is_contiguous():
-        raise ValueError("query must be contiguous (in-place kernel)")
-    if key is not None and (not key.is_contiguous()):
-        raise ValueError("key must be contiguous (in-place kernel)")
+    if cos.device != query.device or sin.device != query.device:
+        raise ValueError("cos/sin device mismatch")
+    if cos.dtype != query.dtype or sin.dtype != query.dtype:
+        raise ValueError("cos/sin dtype mismatch")
+    if positions is None and cos.shape[0] != q_3d.shape[0]:
+        raise ValueError(f"cos/sin shape {cos.shape} mismatches tokens {q_3d.shape[0]}")
 
-    # Optional downsample for interleaved caches that are stored as full-dim (repeat format).
-    # Keep this consistent with AOT kernel behavior.
-    if effective_interleaved and cos.shape[1] == head_size:
+    if interleaved and cos.shape[1] == head_size:
         if head_size % 2 != 0:
             raise ValueError("interleaved layout requires even head_size")
         half = head_size // 2
-        cos_to_use = cos.view(cos.shape[0], half, 2).select(2, 0).contiguous()
-        sin_to_use = sin.view(sin.shape[0], half, 2).select(2, 1).contiguous()
+        cos = cos.view(cos.shape[0], half, 2).select(2, 0).contiguous()
+        sin = sin.view(sin.shape[0], half, 2).select(2, 1).contiguous()
     else:
-        cos_to_use = cos.contiguous()
-        sin_to_use = sin.contiguous()
+        cos, sin = cos.contiguous(), sin.contiguous()
 
-    rot_dim = int(cos_to_use.shape[1])
+    rot_dim = cos.shape[1]
     if rot_dim <= 0:
-        raise ValueError("cos/sin rot_dim must be > 0")
-    if effective_interleaved:
+        raise ValueError("rot_dim must be > 0")
+
+    if interleaved:
         if 2 * rot_dim > head_size:
-            raise ValueError(
-                f"rotate dim exceeds head_size for interleaved=True: 2*rot_dim={2*rot_dim} > head_size={head_size}"
-            )
+            raise ValueError(f"rot_dim {rot_dim} too large for interleaved")
+        embed_dim = rot_dim
     else:
         if rot_dim % 2 != 0:
-            raise ValueError(
-                f"non-interleaved requires even rot_dim (cos.shape[1]), got rot_dim={rot_dim}"
-            )
+            raise ValueError("non-interleaved requires even rot_dim")
         if rot_dim > head_size:
-            raise ValueError(
-                f"rotate dim exceeds head_size for interleaved=False: rot_dim={rot_dim} > head_size={head_size}"
-            )
-
-    embed_dim_for_rotation = rot_dim if effective_interleaved else (rot_dim // 2)
-
-    def _as_3d(x: torch.Tensor, h: int) -> torch.Tensor:
-        if x.dim() == 4:
-            # [bsz, seqlen, num_heads, head_dim] -> [tokens, num_heads, head_dim]
-            if x.shape[-1] != head_size:
-                raise ValueError("head_size mismatch with query/key last dim")
-            return x.flatten(0, 1)
-        if x.dim() == 3:
-            if x.shape[-1] != head_size:
-                raise ValueError("head_size mismatch with query/key last dim")
-            return x
-        if x.dim() != 2:
-            raise ValueError("query/key must be 2D, 3D or 4D tensors")
-        if x.shape[1] % head_size != 0:
-            raise ValueError("hidden_size is not divisible by head_size")
-        return x.view(x.shape[0], h, head_size)
-
-    if query.dim() == 4:
-        num_heads = int(query.shape[2])
-    elif query.dim() == 3:
-        num_heads = int(query.shape[1])
-    else:
-        num_heads = int(query.shape[1] // head_size)
-    q3 = _as_3d(query, num_heads)
-    if not q3.is_contiguous():
-        raise ValueError("query must be contiguous (view must stay contiguous)")
-
-    k3 = None
-    if key is not None:
-        if key.dim() == 4:
-            num_kv_heads = int(key.shape[2])
-        elif key.dim() == 3:
-            num_kv_heads = int(key.shape[1])
-        else:
-            num_kv_heads = int(key.shape[1] // head_size)
-        k3 = _as_3d(key, num_kv_heads)
-        if not k3.is_contiguous():
-            raise ValueError("key must be contiguous (view must stay contiguous)")
-
-    if k3 is None:
-        rotary_embedding_cos_sin_q(
-            rot=embed_dim_for_rotation,
-            cos=cos_to_use,
-            sin=sin_to_use,
-            query=q3,
-            head_size=head_size,
-            interleaved=effective_interleaved,
-            positions=positions,
+            raise ValueError(f"rot_dim {rot_dim} too large")
+        embed_dim = rot_dim // 2
+
+    module = _jit_rotary_embedding_cos_sin_module(embed_dim)
+    if k_3d is not None:
+        module.rotary_embedding_cos_sin_qk(
+            cos, sin, q_3d, k_3d, positions, head_size, interleaved
         )
     else:
-        rotary_embedding_cos_sin_qk(
-            rot=embed_dim_for_rotation,
-            cos=cos_to_use,
-            sin=sin_to_use,
-            query=q3,
-            key=k3,
-            head_size=head_size,
-            interleaved=effective_interleaved,
-            positions=positions,
+        module.rotary_embedding_cos_sin_q(
+            cos, sin, q_3d, positions, head_size, interleaved
         )
diff --git a/python/sglang/jit_kernel/tests/test_rotary_embedding.py b/python/sglang/jit_kernel/tests/test_rotary_embedding.py
index 15b520301c04..5cf10bdb430f 100644
--- a/python/sglang/jit_kernel/tests/test_rotary_embedding.py
+++ b/python/sglang/jit_kernel/tests/test_rotary_embedding.py
@@ -47,9 +47,10 @@ def sglang_aot_rotary_positions(
     )
 
 
-def sglang_jit_rotary_cos_sin(
-    cos: torch.Tensor,
-    sin: torch.Tensor,
+def sglang_jit_rotary_with_positions(
+    positions: torch.Tensor,
+    cos_cache: torch.Tensor,
+    sin_cache: torch.Tensor,
     q: torch.Tensor,
     k: torch.Tensor,
     head_size: int,
@@ -57,7 +58,9 @@ def sglang_jit_rotary_cos_sin(
 ) -> None:
     from sglang.jit_kernel.rotary_embedding import rotary_embedding_cos_sin
 
-    rotary_embedding_cos_sin(cos, sin, q, k, head_size, interleaved)
+    rotary_embedding_cos_sin(
+        cos_cache, sin_cache, q, k, head_size, interleaved, positions=positions
+    )
 
 
 @torch.no_grad()
@@ -152,10 +155,19 @@ def main():
                     sin_half = sin_cache[positions].contiguous()
 
                     if INTERLEAVED:
-                        cos, sin = cos_half, sin_half
+                        # interleaved: cache R = embed_dim = head_size/2
+                        cos_g, sin_g = cos_half, sin_half
+                        cos_cache_for_jit, sin_cache_for_jit = cos_cache, sin_cache
                     else:
-                        cos = torch.cat([cos_half, cos_half], dim=-1).contiguous()
-                        sin = torch.cat([sin_half, sin_half], dim=-1).contiguous()
+                        # non-interleaved: kernel expects R = head_size (x/y halves)
+                        cos_g = torch.cat([cos_half, cos_half], dim=-1).contiguous()
+                        sin_g = torch.cat([sin_half, sin_half], dim=-1).contiguous()
+                        cos_cache_for_jit = torch.cat(
+                            [cos_cache, cos_cache], dim=-1
+                        ).contiguous()
+                        sin_cache_for_jit = torch.cat(
+                            [sin_cache, sin_cache], dim=-1
+                        ).contiguous()
 
                     q = torch.randn(
                         BS, NUM_Q_HEADS, HEAD_SIZE, device=DEVICE, dtype=DTYPE
@@ -166,7 +178,7 @@ def main():
 
                     q_ref_fp32, k_ref_fp32 = q.clone(), k.clone()
                     torch_impl_rotary_fp32(
-                        cos, sin, q_ref_fp32, k_ref_fp32, HEAD_SIZE, INTERLEAVED
+                        cos_g, sin_g, q_ref_fp32, k_ref_fp32, HEAD_SIZE, INTERLEAVED
                     )
 
                     q_k_aot = (q.clone(), k.clone())
@@ -180,12 +192,18 @@ def main():
                             INTERLEAVED,
                             aot_cos_sin_cache,
                         )
-                    sglang_jit_rotary_cos_sin(
-                        cos, sin, q_k_jit[0], q_k_jit[1], HEAD_SIZE, INTERLEAVED
+                    sglang_jit_rotary_with_positions(
+                        positions,
+                        cos_cache_for_jit,
+                        sin_cache_for_jit,
+                        q_k_jit[0],
+                        q_k_jit[1],
+                        HEAD_SIZE,
+                        INTERLEAVED,
                     )
 
-                    ref_atol = 2e-2 if DTYPE == torch.bfloat16 else 2e-3
-                    ref_rtol = 2e-2 if DTYPE == torch.bfloat16 else 2e-3
+                    ref_atol = 1e-2 if DTYPE == torch.bfloat16 else 1e-3
+                    ref_rtol = 1e-2 if DTYPE == torch.bfloat16 else 1e-3
 
                     if HAS_SGL_POS:
                         triton.testing.assert_close(

From 4298be8f6e0ca980f5f07b94b787911380b6e4d0 Mon Sep 17 00:00:00 2001
From: RubiaCx <1084281732@qq.com>
Date: Mon, 5 Jan 2026 09:22:50 -0800
Subject: [PATCH 36/37] Refactor RoPE benchmark; remove cache; use view; relax
 test tolerances.

---
 .../benchmark/bench_rotary_embedding.py       | 47 +-------------
 .../jit_kernel/tests/test_rotary_embedding.py |  4 +-
 .../runtime/layers/rotary_embedding.py        | 63 +++++--------------
 3 files changed, 21 insertions(+), 93 deletions(-)

diff --git a/python/sglang/jit_kernel/benchmark/bench_rotary_embedding.py b/python/sglang/jit_kernel/benchmark/bench_rotary_embedding.py
index 6d5621cb8699..6b0eb5081326 100644
--- a/python/sglang/jit_kernel/benchmark/bench_rotary_embedding.py
+++ b/python/sglang/jit_kernel/benchmark/bench_rotary_embedding.py
@@ -24,6 +24,8 @@
 MAX_SEQ_LEN = 8192
 NUM_HEADS = 32
 BS_RANGE = [1, 64]
+if not os.getenv("CI"):
+    BS_RANGE = [1, 2, 4, 8, 16, 32, 64, 128]
 SEQ_RANGE = [1, 2048]
 HEAD_SIZE_RANGE = [128]
 INTERLEAVED_RANGE = [True, False]
@@ -277,7 +279,6 @@ def prepare_case(
 def make_fn(
     provider: str, case: Case, head_size: int, interleaved: bool
 ) -> Callable[[], None]:
-    # 统一 clone：每次 bench 都从同一份 base 开始，避免“越跑越旋”
     if provider in ("jit_fused", "jit_unfused", "aot_pos", "torch_fp32", "vllm_pos"):
         q = case.q_base.clone()
         k = case.k_base.clone()
@@ -384,51 +385,7 @@ def available_providers() -> Dict[str, str]:
 def benchmark(
     batch_size: int, seq_len: int, head_size: int, interleaved: bool, provider: str
 ):
-    if DEVICE == "cuda" and not torch.cuda.is_available():
-        raise RuntimeError("CUDA not available")
-
     case = prepare_case(batch_size, seq_len, head_size, interleaved)
-
-    global _SANITY_DONE
-    if not _SANITY_DONE and provider == "jit_fused":
-        # sanity: compare jit_fused / jit_unfused with fp32 ref
-        q_ref = case.q_base.clone()
-        k_ref = case.k_base.clone()
-        torch_impl_rotary_fp32(
-            case.cos_gathered, case.sin_gathered, q_ref, k_ref, head_size, interleaved
-        )
-
-        for p in ("jit_fused", "jit_unfused"):
-            q = case.q_base.clone()
-            k = case.k_base.clone()
-            if p == "jit_fused":
-                sgl_jit_fused(
-                    case.positions,
-                    case.cos_cache,
-                    case.sin_cache,
-                    q,
-                    k,
-                    head_size,
-                    interleaved,
-                )
-            else:
-                sgl_jit_unfused(
-                    case.positions,
-                    case.cos_cache,
-                    case.sin_cache,
-                    q,
-                    k,
-                    head_size,
-                    interleaved,
-                )
-
-            atol = 2e-2 if DTYPE == torch.bfloat16 else 2e-3
-            rtol = 2e-2 if DTYPE == torch.bfloat16 else 2e-3
-            _assert_close_gpu(q, q_ref, atol=atol, rtol=rtol, name=f"{p}(q)")
-            _assert_close_gpu(k, k_ref, atol=atol, rtol=rtol, name=f"{p}(k)")
-
-        _SANITY_DONE = True
-
     fn = make_fn(provider, case, head_size, interleaved)
 
     use_cudagraph = provider != "flashinfer"
diff --git a/python/sglang/jit_kernel/tests/test_rotary_embedding.py b/python/sglang/jit_kernel/tests/test_rotary_embedding.py
index 5cf10bdb430f..9c01c91fae8b 100644
--- a/python/sglang/jit_kernel/tests/test_rotary_embedding.py
+++ b/python/sglang/jit_kernel/tests/test_rotary_embedding.py
@@ -202,8 +202,8 @@ def main():
                         INTERLEAVED,
                     )
 
-                    ref_atol = 1e-2 if DTYPE == torch.bfloat16 else 1e-3
-                    ref_rtol = 1e-2 if DTYPE == torch.bfloat16 else 1e-3
+                    ref_atol = 2e-2 if DTYPE == torch.bfloat16 else 2e-3
+                    ref_rtol = 2e-2 if DTYPE == torch.bfloat16 else 2e-3
 
                     if HAS_SGL_POS:
                         triton.testing.assert_close(
diff --git a/python/sglang/multimodal_gen/runtime/layers/rotary_embedding.py b/python/sglang/multimodal_gen/runtime/layers/rotary_embedding.py
index f1a59bcf15b7..72fd65d836d3 100644
--- a/python/sglang/multimodal_gen/runtime/layers/rotary_embedding.py
+++ b/python/sglang/multimodal_gen/runtime/layers/rotary_embedding.py
@@ -38,10 +38,6 @@
 
 logger = init_logger(__name__)
 
-_JIT_ROPE_SPLIT_CAST_CACHE: "OrderedDict[tuple, tuple[torch.Tensor, torch.Tensor]]" = (
-    OrderedDict()
-)
-
 
 def apply_flashinfer_rope_qk_inplace(
     q: torch.Tensor,
@@ -182,14 +178,9 @@ def apply_sglang_jit_rope_qk_inplace(
     if head_size != d:
         raise ValueError(f"head_size mismatch: inferred {d}, but head_size={head_size}")
 
-    try:
-        from sglang.jit_kernel.rotary_embedding import (
-            rotary_embedding_cos_sin as sglang_jit_rotary_embedding_cos_sin,
-        )
-    except Exception as e:
-        raise RuntimeError(
-            "SGLang JIT RoPE is required for apply_sglang_jit_rope_qk_inplace."
-        ) from e
+    from sglang.jit_kernel.rotary_embedding import (
+        rotary_embedding_cos_sin as sglang_jit_rotary_embedding_cos_sin,
+    )
 
     num_tokens = bsz * seqlen
 
@@ -220,49 +211,29 @@ def apply_sglang_jit_rope_qk_inplace(
             )
         positions = positions.to(q.device, non_blocking=True)
 
-    q3 = q.reshape(num_tokens, nheads, d)
-    if not q3.is_contiguous():
-        q3 = q3.contiguous()
-    k3 = k.reshape(num_tokens, nheads, d)
-    if not k3.is_contiguous():
-        k3 = k3.contiguous()
+    q3 = q.view(num_tokens, nheads, d)
+    k3 = k.view(num_tokens, nheads, d)
 
     # Use the FULL cache (no index_select)
     cache_full = cos_sin_cache
 
     # Include is_neox in key because it affects how we process (cat vs contiguous)
     interleaved = not is_neox
-    cache_key = (
-        int(cache_full.data_ptr()),
-        int(cache_full.shape[0]),
-        int(cache_full.shape[1]),
-        str(cache_full.device),
-        cache_full.dtype,
-        q.dtype,
-        interleaved,
-    )
-    cached = _JIT_ROPE_SPLIT_CAST_CACHE.get(cache_key)
-    if cached is not None:
-        cos, sin = cached
-    else:
-        # Split and Cast the FULL cache
-        cos_full, sin_full = _split_cos_sin_from_cache(cache_full, dtype=q.dtype)
 
-        # Process layout (Contiguous or Cat)
-        if interleaved:
+    # Split and Cast the FULL cache
+    cos_full, sin_full = _split_cos_sin_from_cache(cache_full, dtype=q.dtype)
+
+    # Process layout (Contiguous or Cat)
+    if interleaved:
+        cos = cos_full.contiguous()
+        sin = sin_full.contiguous()
+    else:
+        if cos_full.shape[1] * 2 == head_size:
+            cos = torch.cat([cos_full, cos_full], dim=-1).contiguous()
+            sin = torch.cat([sin_full, sin_full], dim=-1).contiguous()
+        else:
             cos = cos_full.contiguous()
             sin = sin_full.contiguous()
-        else:
-            if cos_full.shape[1] * 2 == head_size:
-                cos = torch.cat([cos_full, cos_full], dim=-1).contiguous()
-                sin = torch.cat([sin_full, sin_full], dim=-1).contiguous()
-            else:
-                cos = cos_full.contiguous()
-                sin = sin_full.contiguous()
-
-        _JIT_ROPE_SPLIT_CAST_CACHE[cache_key] = (cos, sin)
-        if len(_JIT_ROPE_SPLIT_CAST_CACHE) > 64:
-            _JIT_ROPE_SPLIT_CAST_CACHE.popitem(last=False)
 
     sglang_jit_rotary_embedding_cos_sin(
         cos, sin, q3, k3, head_size, interleaved, positions=positions

From 2d2d8c3f748c7ed1e7b08cc3a4c96a8360a05c69 Mon Sep 17 00:00:00 2001
From: RubiaCx <1084281732@qq.com>
Date: Mon, 5 Jan 2026 09:30:44 -0800
Subject: [PATCH 37/37] Refactor RoPE benchmark.

---
 .../benchmark/bench_rotary_embedding.py        | 18 ++++++++++++++----
 1 file changed, 14 insertions(+), 4 deletions(-)

diff --git a/python/sglang/jit_kernel/benchmark/bench_rotary_embedding.py b/python/sglang/jit_kernel/benchmark/bench_rotary_embedding.py
index 6b0eb5081326..13a3b7681e68 100644
--- a/python/sglang/jit_kernel/benchmark/bench_rotary_embedding.py
+++ b/python/sglang/jit_kernel/benchmark/bench_rotary_embedding.py
@@ -23,11 +23,21 @@
 DTYPE = torch.bfloat16
 MAX_SEQ_LEN = 8192
 NUM_HEADS = 32
-BS_RANGE = [1, 64]
-if not os.getenv("CI"):
+
+IS_CI = (
+    os.getenv("CI", "false").lower() == "true"
+    or os.getenv("GITHUB_ACTIONS", "false").lower() == "true"
+)
+
+if IS_CI:
+    BS_RANGE = [1]
+    SEQ_RANGE = [1]
+    HEAD_SIZE_RANGE = [128]
+else:
     BS_RANGE = [1, 2, 4, 8, 16, 32, 64, 128]
-SEQ_RANGE = [1, 2048]
-HEAD_SIZE_RANGE = [128]
+    SEQ_RANGE = [1, 16, 64, 256, 1024, 2048]
+    HEAD_SIZE_RANGE = [64, 128, 256]
+
 INTERLEAVED_RANGE = [True, False]
 
 ONLY_PROVIDER = os.getenv("SGL_BENCH_PROVIDER", "").strip() or None