chore: import upstream snapshot with attribution

python/sglang/jit_kernel/csrc/elementwise/activation.cuh

@@ -0,0 +1,256 @@
+#include <sgl_kernel/tensor.h>
+#include <sgl_kernel/utils.h>
+
+#include <sgl_kernel/runtime.cuh>
+#include <sgl_kernel/type.cuh>
+#include <sgl_kernel/utils.cuh>
+#include <sgl_kernel/vec.cuh>
+
+#include <tvm/ffi/container/tensor.h>
+
+#include <cmath>
+#include <cstdint>
+#include <limits>
+#include <string>
+
+namespace {
+
+enum class ActivationKind : uint32_t {
+  kSiLU,
+  kGELU,
+  kGELUTanh,
+  kReLU2,
+};
+
+template <ActivationKind kAct>
+SGL_DEVICE float apply_activation_f32(float x_f32) {
+  if constexpr (kAct == ActivationKind::kSiLU) {
+    return x_f32 / (1.0f + expf(-x_f32));
+  } else if constexpr (kAct == ActivationKind::kGELU) {
+    constexpr auto kSqrt1Over2 = 0.7071067811865475f;
+    return x_f32 * (0.5f * (1.0f + erff(x_f32 * kSqrt1Over2)));
+  } else if constexpr (kAct == ActivationKind::kGELUTanh) {
+    constexpr auto kGeluTanhAlpha = 0.044715f;
+    constexpr auto kGeluTanhBeta = 0.7978845608028654f;
+    const float cdf = 0.5f * (1.0f + tanhf(kGeluTanhBeta * (x_f32 + kGeluTanhAlpha * x_f32 * x_f32 * x_f32)));
+    return x_f32 * cdf;
+  } else if constexpr (kAct == ActivationKind::kReLU2) {
+    const float relu = x_f32 > 0.0f ? x_f32 : 0.0f;
+    return relu * relu;
+  } else {
+    static_assert(host::dependent_false_v<decltype(kAct)>, "unsupported activation kind");
+    return 0.0f;
+  }
+}
+
+struct ActivationParams {
+  const void* __restrict__ input;
+  void* __restrict__ out;
+  uint32_t hidden_dim;
+  uint32_t num_tokens;
+  // Optional MoE expert filtering: when expert_ids != nullptr, a token is
+  // skipped if expert_ids[token_id / expert_step] == -1. expert_step is 1
+  // for per-token routing and BLOCK_SIZE_M for sorted/TMA routing.
+  const int32_t* __restrict__ expert_ids;
+  uint32_t expert_step;
+};
+
+template <typename T, ActivationKind kAct, bool kUsePDL, bool kFilterExpert>
+__global__ void act_and_mul_kernel(const __grid_constant__ ActivationParams params) {
+  using namespace device;
+  constexpr auto kVecSize = kMaxVecBytes / sizeof(T);
+  using vec_t = AlignedVector<T, kMaxVecBytes / sizeof(T)>;
+  const auto num_vecs = params.hidden_dim / kVecSize;  // per token
+  const auto tid = blockIdx.x * blockDim.x + threadIdx.x;
+  const auto token_id = tid / num_vecs;
+
+  if (token_id >= params.num_tokens) return;
+  if constexpr (kFilterExpert) {
+    if (params.expert_ids[token_id / params.expert_step] == -1) return;
+  }
+  const auto offset = tid % num_vecs;
+  const auto input_offset = token_id * (num_vecs * 2) + offset;
+  const auto output_offset = tid;
+  PDLWaitPrimary<kUsePDL>();
+  const auto gate = device::load_as<vec_t>(params.input, input_offset);
+  const auto up = device::load_as<vec_t>(params.input, input_offset + num_vecs);
+  vec_t out;
+#pragma unroll
+  for (int i = 0; i < kVecSize; ++i) {
+    const float gate_f32 = device::cast<fp32_t>(gate[i]);
+    const float up_f32 = device::cast<fp32_t>(up[i]);
+    out[i] = device::cast<T>(apply_activation_f32<kAct>(gate_f32) * up_f32);
+  }
+  device::store_as<vec_t>(params.out, out, output_offset);
+  PDLTriggerSecondary<kUsePDL>();
+}
+
+struct UnaryActivationParams {
+  const void* __restrict__ input;
+  void* __restrict__ out;
+  uint32_t num_vecs;
+};
+
+template <typename T, ActivationKind kAct, bool kUsePDL>
+__global__ void act_kernel(const __grid_constant__ UnaryActivationParams params) {
+  using namespace device;
+  constexpr auto kVecSize = kMaxVecBytes / sizeof(T);
+  using vec_t = AlignedVector<T, kMaxVecBytes / sizeof(T)>;
+  const auto vec_id = blockIdx.x * blockDim.x + threadIdx.x;
+  if (vec_id >= params.num_vecs) return;
+  PDLWaitPrimary<kUsePDL>();
+  const auto in = device::load_as<vec_t>(params.input, vec_id);
+  vec_t out;
+#pragma unroll
+  for (int i = 0; i < kVecSize; ++i) {
+    out[i] = device::cast<T>(apply_activation_f32<kAct>(device::cast<fp32_t>(in[i])));
+  }
+  device::store_as<vec_t>(params.out, out, vec_id);
+  PDLTriggerSecondary<kUsePDL>();
+}
+
+template <typename T, bool kUsePDL>
+struct ActivationKernel {
+  static constexpr auto kVecSize = device::kMaxVecBytes / sizeof(T);
+  static constexpr auto kBlockSize = 256u;
+
+  using kernel_fn_t = decltype(&act_and_mul_kernel<T, ActivationKind::kSiLU, kUsePDL, false>);
+  using unary_kernel_fn_t = decltype(&act_kernel<T, ActivationKind::kReLU2, kUsePDL>);
+
+  template <ActivationKind kAct, bool kFilterExpert>
+  static constexpr kernel_fn_t activation_kernel = act_and_mul_kernel<T, kAct, kUsePDL, kFilterExpert>;
+
+  static_assert(device::kMaxVecBytes % sizeof(T) == 0, "unsupported data type");
+
+  template <bool kFilterExpert>
+  static kernel_fn_t select_kernel(const std::string& type) {
+    using namespace host;
+    if (type == "silu") {
+      return activation_kernel<ActivationKind::kSiLU, kFilterExpert>;
+    } else if (type == "gelu") {
+      return activation_kernel<ActivationKind::kGELU, kFilterExpert>;
+    } else if (type == "gelu_tanh") {
+      return activation_kernel<ActivationKind::kGELUTanh, kFilterExpert>;
+    } else {
+      Panic("unsupported activation type: ", type);
+    }
+    return nullptr;
+  }
+
+  static void launch(
+      const tvm::ffi::TensorView& input,
+      const tvm::ffi::TensorView& out,
+      const std::string& type,
+      const int32_t* expert_ids,
+      uint32_t expert_step) {
+    using namespace host;
+
+    auto N = SymbolicSize{"num_tokens"};
+    auto D_in = SymbolicSize{"input_width"};
+    auto D_out = SymbolicSize{"output_width"};
+    auto device_ = SymbolicDevice{};
+    device_.set_options<kDLCUDA>();
+
+    TensorMatcher({N, D_out})  //
+        .with_dtype<T>()
+        .with_device(device_)
+        .verify(out);
+    TensorMatcher({N, D_in})  //
+        .with_dtype<T>()
+        .with_device(device_)
+        .verify(input);
+
+    const auto hidden_size = static_cast<uint32_t>(D_out.unwrap());
+    const auto num_tokens = static_cast<uint32_t>(N.unwrap());
+    const auto device = device_.unwrap();
+    if (num_tokens == 0) return;
+    RuntimeCheck(hidden_size * 2 == D_in.unwrap(), "invalid activation dimension");
+    RuntimeCheck(hidden_size % kVecSize == 0, "hidden size must be divisible by vector size");
+    // only get once to avoid overhead
+    const auto num_total_items = num_tokens * (hidden_size / kVecSize);
+    RuntimeCheck(num_total_items <= std::numeric_limits<uint32_t>::max(), "too many items for 32-bit indexing");
+    const auto num_blocks = div_ceil(static_cast<uint32_t>(num_total_items), kBlockSize);
+    const auto params = ActivationParams{
+        .input = input.data_ptr(),
+        .out = out.data_ptr(),
+        .hidden_dim = hidden_size,
+        .num_tokens = num_tokens,
+        .expert_ids = expert_ids,
+        .expert_step = expert_step,
+    };
+    if (expert_ids != nullptr) {
+      RuntimeCheck(expert_step > 0, "expert_step must be positive");
+      const auto kernel = select_kernel<true>(type);
+      LaunchKernel(num_blocks, kBlockSize, device).enable_pdl(kUsePDL)(kernel, params);
+    } else {
+      const auto kernel = select_kernel<false>(type);
+      LaunchKernel(num_blocks, kBlockSize, device).enable_pdl(kUsePDL)(kernel, params);
+    }
+  }
+
+  static void run_activation(const tvm::ffi::TensorView input, const tvm::ffi::TensorView out, std::string type) {
+    launch(input, out, type, /*expert_ids=*/nullptr, /*expert_step=*/1);
+  }
+
+  static void run_activation_filtered(
+      const tvm::ffi::TensorView input,
+      const tvm::ffi::TensorView out,
+      const tvm::ffi::TensorView expert_ids,
+      int64_t expert_step,
+      std::string type) {
+    using namespace host;
+    RuntimeCheck(is_type<int32_t>(expert_ids.dtype()), "expert_ids must have dtype int32");
+    RuntimeCheck(expert_step >= 1, "expert_step must be positive");
+    launch(input, out, type, static_cast<const int32_t*>(expert_ids.data_ptr()), static_cast<uint32_t>(expert_step));
+  }
+
+  template <ActivationKind kAct>
+  static constexpr auto unary_kernel = act_kernel<T, kAct, kUsePDL>;
+
+  // Use the explicit non-const function-pointer type (mirrors select_kernel's
+  // kernel_fn_t) rather than a trailing `decltype(unary_kernel<...>)` return,
+  // which deduces a const-qualified pointer that clang-HIP (gfx942) refuses to
+  // initialize from an lvalue / nullptr. nvcc accepts both; this form works for
+  // CUDA and ROCm alike.
+  static unary_kernel_fn_t select_unary_kernel(const std::string& type) {
+    using namespace host;
+    if (type == "relu2") {
+      return ActivationKernel::template unary_kernel<ActivationKind::kReLU2>;
+    } else {
+      Panic("unsupported unary activation type: ", type);
+    }
+    return nullptr;
+  }
+
+  static void run_unary_activation(const tvm::ffi::TensorView input, const tvm::ffi::TensorView out, std::string type) {
+    using namespace host;
+
+    auto N = SymbolicSize{"num_tokens"};
+    auto D = SymbolicSize{"hidden"};
+    auto device_ = SymbolicDevice{};
+    device_.set_options<kDLCUDA>();
+
+    TensorMatcher({N, D})  //
+        .with_dtype<T>()
+        .with_device(device_)
+        .verify(out)
+        .verify(input);
+
+    const auto num_elems = static_cast<int64_t>(N.unwrap()) * D.unwrap();
+    const auto device = device_.unwrap();
+    if (num_elems == 0) return;
+    RuntimeCheck(num_elems % kVecSize == 0, "num elements must be divisible by vector size");
+    const auto num_vecs = num_elems / kVecSize;
+    RuntimeCheck(num_vecs <= std::numeric_limits<uint32_t>::max(), "too many items for 32-bit indexing");
+    const auto num_blocks = div_ceil(static_cast<uint32_t>(num_vecs), kBlockSize);
+    const auto params = UnaryActivationParams{
+        .input = input.data_ptr(),
+        .out = out.data_ptr(),
+        .num_vecs = static_cast<uint32_t>(num_vecs),
+    };
+    const auto kernel = select_unary_kernel(type);
+    LaunchKernel(num_blocks, kBlockSize, device).enable_pdl(kUsePDL)(kernel, params);
+  }
+};
+
+}  // namespace

python/sglang/jit_kernel/csrc/elementwise/clamp_position.cuh

@@ -0,0 +1,54 @@
+#include <sgl_kernel/tensor.h>  // For TensorMatcher, SymbolicSize, SymbolicDevice
+#include <sgl_kernel/utils.h>   // For div_ceil
+
+#include <sgl_kernel/utils.cuh>  // For LaunchKernel
+
+#include <dlpack/dlpack.h>
+#include <tvm/ffi/container/tensor.h>
+
+#include <cstddef>
+#include <cstdint>
+
+namespace {
+
+template <typename T>
+__global__ void clamp_position_kernel(T* __restrict__ dst, const T* __restrict__ seq_lens, size_t n) {
+  size_t idx = blockIdx.x * blockDim.x + threadIdx.x;
+  if (idx < n) {
+    T val = seq_lens[idx] - 1;
+    dst[idx] = val < 0 ? 0 : val;
+  }
+}
+
+constexpr size_t kBlockSize = 256;
+
+template <typename T>
+struct ClampPosition {
+  static void run(tvm::ffi::TensorView dst, tvm::ffi::TensorView seq_lens) {
+    using namespace host;
+
+    SymbolicSize N = {"num_elements"};
+    SymbolicDevice device_;
+    device_.set_options<kDLCUDA, kDLROCM>();
+
+    TensorMatcher({N})  //
+        .with_dtype<T>()
+        .with_device(device_)
+        .verify(dst)
+        .verify(seq_lens);
+
+    const size_t num_elements = N.unwrap();
+    if (num_elements == 0) return;
+
+    const size_t grid_size = div_ceil(num_elements, kBlockSize);
+    const DLDevice device = device_.unwrap();
+
+    LaunchKernel(grid_size, kBlockSize, device)(
+        clamp_position_kernel<T>,
+        static_cast<T*>(dst.data_ptr()),
+        static_cast<const T*>(seq_lens.data_ptr()),
+        num_elements);
+  }
+};
+
+}  // namespace

python/sglang/jit_kernel/csrc/elementwise/concat_mla.cuh

@@ -0,0 +1,325 @@
+#include <sgl_kernel/tensor.h>
+#include <sgl_kernel/utils.h>
+
+#include <sgl_kernel/utils.cuh>
+
+#include <tvm/ffi/container/tensor.h>
+
+#include <cuda_bf16.h>
+#include <cuda_runtime.h>
+
+namespace {
+
+// ======================= Memory Utilities =======================
+// Adapted from DeepEP: https://github.com/deepseek-ai/DeepEP/blob/main/csrc/kernels/utils.cuh
+
+SGL_DEVICE int get_lane_id() {
+  int lane_id;
+  asm("mov.s32 %0, %laneid;" : "=r"(lane_id));
+  return lane_id;
+}
+
+SGL_DEVICE void st_na_global_v1(const int* ptr, int v) {
+  asm volatile("st.global.L1::no_allocate.s32 [%0], %1;" ::"l"(ptr), "r"(v) : "memory");
+}
+
+SGL_DEVICE void st_na_global_v2(const int2* ptr, const int2& v) {
+  asm volatile("st.global.L1::no_allocate.v2.s32 [%0], {%1, %2};" ::"l"(ptr), "r"(v.x), "r"(v.y) : "memory");
+}
+
+SGL_DEVICE int ld_na_global_v1(const int* ptr) {
+  int r;
+  asm volatile("ld.global.nc.L1::no_allocate.s32 %0, [%1];" : "=r"(r) : "l"(ptr));
+  return r;
+}
+
+SGL_DEVICE int2 ld_na_global_v2(const int2* ptr) {
+  int2 r;
+  asm volatile("ld.global.nc.L1::no_allocate.v2.s32 {%0, %1}, [%2];" : "=r"(r.x), "=r"(r.y) : "l"(ptr));
+  return r;
+}
+
+SGL_DEVICE void prefetch_L2(const void* p) {
+#if defined(ENABLE_L2_PREFETCH)
+  asm volatile("prefetch.global.L2 [%0];" ::"l"(p));
+#endif
+}
+
+// ======================= concat_mla_k Kernel =======================
+
+constexpr int NUM_LOCAL_HEADS = 128;
+constexpr int QK_NOPE_HEAD_DIM = 128;
+constexpr int QK_ROPE_HEAD_DIM = 64;
+constexpr int K_HEAD_DIM = QK_NOPE_HEAD_DIM + QK_ROPE_HEAD_DIM;
+
+constexpr int HEAD_CHUNK_SIZE = 16;
+constexpr int NUM_HEAD_CHUNKS = NUM_LOCAL_HEADS / HEAD_CHUNK_SIZE;
+
+__global__ void concat_mla_k_kernel(
+    bf16_t* __restrict__ k,
+    const bf16_t* __restrict__ k_nope,
+    const bf16_t* __restrict__ k_rope,
+    const int num_tokens,
+    const int64_t k_stride_0,
+    const int k_stride_1,
+    const int64_t k_nope_stride_0,
+    const int k_nope_stride_1,
+    const int64_t k_rope_stride_0) {
+  const int flat_warp_id = (blockIdx.x * blockDim.x + threadIdx.x) / 32;
+  const int token_id = flat_warp_id / NUM_HEAD_CHUNKS;
+  const int head_chunk_id = flat_warp_id % NUM_HEAD_CHUNKS;
+  const int lane_id = get_lane_id();
+  if (token_id >= num_tokens) return;
+
+  using NopeVec = int2;  // 8B/thread, 32 threads = 256B/row
+  using RopeVec = int;   // 4B/thread, 32 threads = 128B/row
+  static_assert(sizeof(NopeVec) * 32 == QK_NOPE_HEAD_DIM * sizeof(bf16_t), "nope vec mismatch");
+  static_assert(sizeof(RopeVec) * 32 == QK_ROPE_HEAD_DIM * sizeof(bf16_t), "rope vec mismatch");
+
+  const int head_row0 = head_chunk_id * HEAD_CHUNK_SIZE;
+
+  const int2* __restrict__ nope_src =
+      reinterpret_cast<const int2*>(k_nope + token_id * k_nope_stride_0 + head_row0 * k_nope_stride_1) + lane_id;
+
+  int2* __restrict__ nope_dst = reinterpret_cast<int2*>(k + token_id * k_stride_0 + head_row0 * k_stride_1) + lane_id;
+
+  int* __restrict__ rope_dst =
+      reinterpret_cast<int*>(k + token_id * k_stride_0 + head_row0 * k_stride_1 + QK_NOPE_HEAD_DIM) + lane_id;
+
+  const int nope_src_stride_v = (k_nope_stride_1 >> 2);  // int2 covers 4 bf16
+  const int nope_dst_stride_v = (k_stride_1 >> 2);
+  const int rope_dst_stride_v = (k_stride_1 >> 1);  // int covers 2 bf16
+
+  const int* rope_base = reinterpret_cast<const int*>(k_rope + token_id * k_rope_stride_0);
+  const RopeVec rope_val = ld_na_global_v1(rope_base + lane_id);
+
+  prefetch_L2(nope_src);
+  NopeVec cur = ld_na_global_v2(nope_src);
+
+#pragma unroll
+  for (int i = 0; i < HEAD_CHUNK_SIZE; ++i) {
+    NopeVec next;
+    if (i + 1 < HEAD_CHUNK_SIZE) {
+      const int2* next_src = nope_src + nope_src_stride_v;
+      prefetch_L2(next_src);
+      next = ld_na_global_v2(next_src);
+    }
+
+    st_na_global_v2(nope_dst, cur);
+    st_na_global_v1(rope_dst, rope_val);
+
+    nope_src += nope_src_stride_v;
+    nope_dst += nope_dst_stride_v;
+    rope_dst += rope_dst_stride_v;
+
+    cur = next;
+  }
+}
+
+struct ConcatMlaKKernel {
+  static void run(tvm::ffi::TensorView k, tvm::ffi::TensorView k_nope, tvm::ffi::TensorView k_rope) {
+    using namespace host;
+
+    auto N = SymbolicSize{"num_tokens"};
+    auto H = SymbolicSize{"num_heads"};
+    auto D = SymbolicSize{"k_head_dim"};
+    auto D_nope = SymbolicSize{"nope_head_dim"};
+    auto D_rope = SymbolicSize{"rope_head_dim"};
+    auto S0_k = SymbolicSize{"k_stride_0"};
+    auto S1_k = SymbolicSize{"k_stride_1"};
+    auto S0_k_nope = SymbolicSize{"k_nope_stride_0"};
+    auto S1_k_nope = SymbolicSize{"k_nope_stride_1"};
+    auto S0_k_rope = SymbolicSize{"k_rope_stride_0"};
+    auto device = SymbolicDevice{};
+
+    // Set known fixed values
+    H.set_value(NUM_LOCAL_HEADS);
+    D.set_value(K_HEAD_DIM);
+    D_nope.set_value(QK_NOPE_HEAD_DIM);
+    D_rope.set_value(QK_ROPE_HEAD_DIM);
+
+    // Verify k: [num_tokens, num_heads, k_head_dim]
+    TensorMatcher({N, H, D}).with_strides({S0_k, S1_k, 1}).with_dtype<bf16_t>().with_device<kDLCUDA>(device).verify(k);
+
+    // Verify k_nope: [num_tokens, num_heads, nope_head_dim]
+    TensorMatcher({N, H, D_nope})
+        .with_strides({S0_k_nope, S1_k_nope, 1})
+        .with_dtype<bf16_t>()
+        .with_device<kDLCUDA>(device)
+        .verify(k_nope);
+
+    // Verify k_rope: [num_tokens, 1, rope_head_dim]
+    TensorMatcher({N, 1, D_rope})
+        .with_strides({S0_k_rope, -1, 1})
+        .with_dtype<bf16_t>()
+        .with_device<kDLCUDA>(device)
+        .verify(k_rope);
+
+    // Check alignment
+    RuntimeCheck(reinterpret_cast<uintptr_t>(k.data_ptr()) % 16 == 0, "Tensor k must be 16-byte aligned");
+    RuntimeCheck(reinterpret_cast<uintptr_t>(k_nope.data_ptr()) % 16 == 0, "Tensor k_nope must be 16-byte aligned");
+    RuntimeCheck(reinterpret_cast<uintptr_t>(k_rope.data_ptr()) % 16 == 0, "Tensor k_rope must be 16-byte aligned");
+
+    const int num_tokens = static_cast<int>(N.unwrap());
+
+    constexpr int num_warps_per_block = 32;
+    const int grid_size = div_ceil(num_tokens * NUM_HEAD_CHUNKS, num_warps_per_block);
+    const int block_size = num_warps_per_block * 32;
+
+    LaunchKernel(grid_size, block_size, device.unwrap())(
+        concat_mla_k_kernel,
+        static_cast<bf16_t*>(k.data_ptr()),
+        static_cast<const bf16_t*>(k_nope.data_ptr()),
+        static_cast<const bf16_t*>(k_rope.data_ptr()),
+        num_tokens,
+        S0_k.unwrap(),
+        static_cast<int>(S1_k.unwrap()),
+        S0_k_nope.unwrap(),
+        static_cast<int>(S1_k_nope.unwrap()),
+        S0_k_rope.unwrap());
+  }
+};
+
+// ======================= concat_mla_absorb_q Kernel =======================
+
+constexpr int A_LAST_DIM = 512;
+constexpr int B_LAST_DIM = 64;
+constexpr int OUT_LAST_DIM = A_LAST_DIM + B_LAST_DIM;
+
+__global__ void concat_mla_absorb_q_kernel(
+    bf16_t* a,
+    bf16_t* b,
+    bf16_t* out,
+    const int num_items,
+    const int dim_1,
+    const int64_t a_stride_0,
+    const int a_stride_1,
+    const int64_t b_stride_0,
+    const int b_stride_1,
+    const int64_t out_stride_0,
+    const int out_stride_1) {
+  const int flat_warp_id = (blockIdx.x * blockDim.x + threadIdx.x) / 32;
+  const int lane_id = get_lane_id();
+
+  const int idx_0 = flat_warp_id / dim_1;
+  const int idx_1 = flat_warp_id % dim_1;
+
+  if (flat_warp_id >= num_items) {
+    return;
+  }
+
+  using ABufType = int4;
+  constexpr int A_NUM_UNROLL = 2;
+  static_assert(sizeof(ABufType) * A_NUM_UNROLL == A_LAST_DIM * sizeof(a[0]) / 32);
+  ABufType a_buf[A_NUM_UNROLL];
+
+  using BBufType = int;
+  constexpr int B_NUM_UNROLL = 1;
+  static_assert(sizeof(BBufType) * B_NUM_UNROLL == B_LAST_DIM * sizeof(b[0]) / 32);
+  BBufType b_buf;
+
+  {
+    const BBufType* base_addr = reinterpret_cast<BBufType*>(b + idx_0 * b_stride_0 + idx_1 * b_stride_1);
+    b_buf = *(base_addr + lane_id);
+  }
+
+#pragma unroll
+  for (int i = 0; i < A_NUM_UNROLL; ++i) {
+    const ABufType* base_addr = reinterpret_cast<ABufType*>(a + idx_0 * a_stride_0 + idx_1 * a_stride_1);
+    a_buf[i] = *(base_addr + i * 32 + lane_id);
+  }
+
+  {
+    BBufType* base_addr = reinterpret_cast<BBufType*>(out + idx_0 * out_stride_0 + idx_1 * out_stride_1 + A_LAST_DIM);
+    *(base_addr + lane_id) = b_buf;
+  }
+
+#pragma unroll
+  for (int i = 0; i < A_NUM_UNROLL; ++i) {
+    ABufType* base_addr = reinterpret_cast<ABufType*>(out + idx_0 * out_stride_0 + idx_1 * out_stride_1);
+    *(base_addr + i * 32 + lane_id) = a_buf[i];
+  }
+}
+
+struct ConcatMlaAbsorbQKernel {
+  static void run(tvm::ffi::TensorView a, tvm::ffi::TensorView b, tvm::ffi::TensorView out) {
+    using namespace host;
+
+    auto N0_a = SymbolicSize{"a_dim_0"};
+    auto N1_a = SymbolicSize{"a_dim_1"};
+    auto D_a = SymbolicSize{"a_last_dim"};
+    auto N0_b = SymbolicSize{"b_dim_0"};
+    auto N1_b = SymbolicSize{"b_dim_1"};
+    auto D_b = SymbolicSize{"b_last_dim"};
+    auto N0_out = SymbolicSize{"out_dim_0"};
+    auto N1_out = SymbolicSize{"out_dim_1"};
+    auto D_out = SymbolicSize{"out_last_dim"};
+    auto S0_a = SymbolicSize{"a_stride_0"};
+    auto S1_a = SymbolicSize{"a_stride_1"};
+    auto S0_b = SymbolicSize{"b_stride_0"};
+    auto S1_b = SymbolicSize{"b_stride_1"};
+    auto S0_out = SymbolicSize{"out_stride_0"};
+    auto S1_out = SymbolicSize{"out_stride_1"};
+    auto device = SymbolicDevice{};
+
+    // Set known fixed values
+    D_a.set_value(A_LAST_DIM);
+    D_b.set_value(B_LAST_DIM);
+    D_out.set_value(OUT_LAST_DIM);
+
+    // Verify a: [dim_0, dim_1, A_LAST_DIM]
+    TensorMatcher({N0_a, N1_a, D_a})
+        .with_strides({S0_a, S1_a, 1})
+        .with_dtype<bf16_t>()
+        .with_device<kDLCUDA>(device)
+        .verify(a);
+
+    // Verify b: [dim_0, dim_1, B_LAST_DIM]
+    TensorMatcher({N0_b, N1_b, D_b})
+        .with_strides({S0_b, S1_b, 1})
+        .with_dtype<bf16_t>()
+        .with_device<kDLCUDA>(device)
+        .verify(b);
+
+    // Verify out: [dim_0, dim_1, OUT_LAST_DIM]
+    TensorMatcher({N0_out, N1_out, D_out})
+        .with_strides({S0_out, S1_out, 1})
+        .with_dtype<bf16_t>()
+        .with_device<kDLCUDA>(device)
+        .verify(out);
+
+    // Check alignment
+    RuntimeCheck(reinterpret_cast<uintptr_t>(a.data_ptr()) % 16 == 0, "Tensor a must be 16-byte aligned");
+    RuntimeCheck(reinterpret_cast<uintptr_t>(b.data_ptr()) % 16 == 0, "Tensor b must be 16-byte aligned");
+    RuntimeCheck(reinterpret_cast<uintptr_t>(out.data_ptr()) % 16 == 0, "Tensor out must be 16-byte aligned");
+
+    // Verify dimensions match: a.size(0) * a.size(1) == b.size(0) * b.size(1)
+    RuntimeCheck(
+        N0_a.unwrap() * N1_a.unwrap() == N0_b.unwrap() * N1_b.unwrap(),
+        "Dimension mismatch: a.size(0) * a.size(1) must equal b.size(0) * b.size(1)");
+    RuntimeCheck(N1_a.unwrap() == N1_b.unwrap(), "Dimension mismatch: a.size(1) must equal b.size(1)");
+
+    const int num_items = static_cast<int>(N0_a.unwrap() * N1_a.unwrap());
+    const int dim_1 = static_cast<int>(N1_a.unwrap());
+
+    constexpr int num_warps_per_block = 32;
+    const int grid_size = div_ceil(num_items, num_warps_per_block);
+    const int block_size = num_warps_per_block * 32;
+
+    LaunchKernel(grid_size, block_size, device.unwrap())(
+        concat_mla_absorb_q_kernel,
+        static_cast<bf16_t*>(a.data_ptr()),
+        static_cast<bf16_t*>(b.data_ptr()),
+        static_cast<bf16_t*>(out.data_ptr()),
+        num_items,
+        dim_1,
+        S0_a.unwrap(),
+        static_cast<int>(S1_a.unwrap()),
+        S0_b.unwrap(),
+        static_cast<int>(S1_b.unwrap()),
+        S0_out.unwrap(),
+        static_cast<int>(S1_out.unwrap()));
+  }
+};
+
+}  // namespace

python/sglang/jit_kernel/csrc/elementwise/fused_add_rmsnorm.cuh

@@ -0,0 +1,197 @@
+#include <sgl_kernel/tensor.h>
+#include <sgl_kernel/utils.h>
+
+#include <sgl_kernel/tile.cuh>
+#include <sgl_kernel/type.cuh>
+#include <sgl_kernel/utils.cuh>
+#include <sgl_kernel/vec.cuh>
+
+#include <cooperative_groups/reduce.h>
+#include <tvm/ffi/container/tensor.h>
+
+#include <cooperative_groups.h>
+#include <type_traits>
+
+namespace {
+
+template <typename T, int VEC_SIZE_IN_BYTE>
+struct VecTypeTrait;
+
+template <>
+struct VecTypeTrait<bf16_t, 16> {
+  using packed_t = packed_t<bf16_t>;
+  using vec_t = device::AlignedVector<packed_t, 4>;
+};
+
+template <>
+struct VecTypeTrait<fp16_t, 16> {
+  using packed_t = packed_t<fp16_t>;
+  using vec_t = device::AlignedVector<packed_t, 4>;
+};
+
+template <>
+struct VecTypeTrait<bf16_t, 32> {
+  using packed_t = packed_t<bf16_t>;
+  using vec_t = device::AlignedVector<packed_t, 8>;
+};
+
+template <>
+struct VecTypeTrait<fp16_t, 32> {
+  using packed_t = packed_t<fp16_t>;
+  using vec_t = device::AlignedVector<packed_t, 8>;
+};
+
+template <bool kCastXBeforeOutMul, typename packed_t>
+SGL_DEVICE packed_t rms(float2 valf, packed_t& weight, float rsqrt_square_sum) {
+  float2 weightf = device::cast<fp32x2_t, packed_t>(weight);
+  if constexpr (kCastXBeforeOutMul) {
+    auto rounded = device::cast<packed_t, fp32x2_t>(make_float2(valf.x * rsqrt_square_sum, valf.y * rsqrt_square_sum));
+    valf = device::cast<fp32x2_t, packed_t>(rounded);
+    return device::cast<packed_t, fp32x2_t>(make_float2(valf.x * weightf.x, valf.y * weightf.y));
+  }
+  return device::cast<packed_t, fp32x2_t>(
+      make_float2(valf.x * weightf.x * rsqrt_square_sum, valf.y * weightf.y * rsqrt_square_sum));
+}
+
+template <bool kCastXBeforeOutMul, typename T, int VEC_SIZE_IN_BYTE>
+__global__ void fused_add_rmsnorm_reg_kernel(
+    T* __restrict__ input, T* __restrict__ residual, const T* __restrict__ weight, int vec_hidden_size, float eps) {
+  constexpr int inner_loop = VEC_SIZE_IN_BYTE == 16 ? 4 : 8;
+
+  __shared__ float shared_memory[32];  // Used for CTA reduce
+
+  using vec_t = typename VecTypeTrait<T, VEC_SIZE_IN_BYTE>::vec_t;
+  using packed_t = typename VecTypeTrait<T, VEC_SIZE_IN_BYTE>::packed_t;
+  vec_t v;                           // Save input
+  vec_t v_res;                       // Save residual
+  vec_t v_weight;                    // Save weight
+  vec_t v_out;                       // Save output
+  float2 inp_res_cache[inner_loop];  // fp32 sum cache; only read when kCastXBeforeOutMul=true
+
+  auto token_id = blockIdx.x;
+  float2 acc_square = make_float2(0.0f, 0.0f);  // Sum of squares for each thread
+
+  if (threadIdx.x < vec_hidden_size) {
+    // Compute address
+    vec_t* p = reinterpret_cast<vec_t*>(input) + token_id * vec_hidden_size;
+    vec_t* p_res = reinterpret_cast<vec_t*>(residual) + token_id * vec_hidden_size;
+    const vec_t* p_weight = reinterpret_cast<const vec_t*>(weight);
+
+    // Load data
+    v = p[threadIdx.x];
+    v_res = p_res[threadIdx.x];
+    v_weight = p_weight[threadIdx.x];
+
+    for (int i = 0; i < inner_loop; i++) {
+      float2 val = device::cast<fp32x2_t, packed_t>(v[i]);
+      float2 res = device::cast<fp32x2_t, packed_t>(v_res[i]);
+      float2 inp_res = make_float2(val.x + res.x, val.y + res.y);
+      acc_square.x += inp_res.x * inp_res.x;
+      acc_square.y += inp_res.y * inp_res.y;
+      v[i] = device::cast<packed_t, fp32x2_t>(inp_res);
+      if constexpr (kCastXBeforeOutMul) {
+        inp_res_cache[i] = inp_res;
+      }
+    }
+
+    // Store inp+res to residual
+    p_res[threadIdx.x] = v;
+  }
+
+  // CTA Reduce
+  // Step 0: Warp Reduce
+  auto cg_warp = cooperative_groups::tiled_partition<32>(cooperative_groups::this_thread_block());
+  float warp_sum = cooperative_groups::reduce(cg_warp, acc_square.x + acc_square.y, cooperative_groups::plus<float>());
+
+  float* buffer = shared_memory;
+  if (threadIdx.x % 32 == 0) {
+    buffer[threadIdx.x / 32] = warp_sum;  // Write warp_sum to buffer
+  }
+
+  // Step 1: CTA Reduce
+  __syncthreads();
+  if (threadIdx.x < 32) {
+    float cta_sum = cooperative_groups::reduce(
+        cg_warp, (threadIdx.x < blockDim.x / 32) ? buffer[threadIdx.x] : 0.0f, cooperative_groups::plus<float>());
+    buffer[threadIdx.x] =
+        rsqrtf(eps + cta_sum * (1.0f / static_cast<float>(vec_hidden_size * (VEC_SIZE_IN_BYTE / sizeof(T)))));
+  }
+  __syncthreads();
+
+  // Compute RMSNorm
+  if (threadIdx.x < vec_hidden_size) {
+    float rsqrt_square_sum = buffer[threadIdx.x / 32];  // Read rsqrt from Shared Memory(Broadcast)
+    for (int i = 0; i < inner_loop; i++) {
+      // HF parity needs the full fp32 sum (not the DType-rounded v[i]).
+      float2 valf;
+      if constexpr (kCastXBeforeOutMul) {
+        valf = inp_res_cache[i];
+      } else {
+        valf = device::cast<fp32x2_t, packed_t>(v[i]);
+      }
+      v_out[i] = rms<kCastXBeforeOutMul>(valf, v_weight[i], rsqrt_square_sum);
+    }
+    vec_t* p_out = reinterpret_cast<vec_t*>(input) + token_id * vec_hidden_size;
+    p_out[threadIdx.x] = v_out;
+  }
+}
+
+template <bool kCastXBeforeOutMul, typename DType>
+struct FusedAddRMSNormKernel {
+  static void
+  run(const tvm::ffi::TensorView input,
+      const tvm::ffi::TensorView residual,
+      const tvm::ffi::TensorView weight,
+      float eps) {
+    using namespace host;
+    auto N = SymbolicSize{"num_tokens"};
+    auto D = SymbolicSize{"hidden_size"};
+    auto device = SymbolicDevice{};
+    device.set_options<kDLCUDA>();
+
+    TensorMatcher({N, D})  // input
+        .with_strides({D, 1})
+        .with_dtype<DType>()
+        .with_device(device)
+        .verify(input);
+    TensorMatcher({D})  // weight
+        .with_dtype<DType>()
+        .with_device(device)
+        .verify(weight);
+    TensorMatcher({N, D})  // residual
+        .with_strides({D, 1})
+        .with_dtype<DType>()
+        .with_device(device)
+        .verify(residual);
+
+    int hidden_size = static_cast<int>(D.unwrap());
+    if (hidden_size <= (device::kMaxVecBytes == 32 ? 12288 : 8192)) {
+      int elements_in_vec = device::kMaxVecBytes / sizeof(DType);
+      int vec_hidden_size = hidden_size / elements_in_vec;
+      uint threads = (vec_hidden_size + 31) / 32 * 32;
+
+      // Runtime check
+      host::RuntimeCheck(
+          hidden_size % elements_in_vec == 0,
+          "hidden_size",
+          hidden_size,
+          " can not align to elements_in_vec ",
+          elements_in_vec);
+
+      // Launch kernel
+      auto kernel = fused_add_rmsnorm_reg_kernel<kCastXBeforeOutMul, DType, device::kMaxVecBytes>;
+      LaunchKernel(static_cast<uint>(N.unwrap()), threads, device.unwrap())
+          .enable_pdl(false)(
+              kernel,
+              reinterpret_cast<DType*>(input.data_ptr()),
+              reinterpret_cast<DType*>(residual.data_ptr()),
+              reinterpret_cast<DType*>(weight.data_ptr()),
+              vec_hidden_size,
+              eps);
+    } else {
+      host::RuntimeCheck(false, "Large hidden_sizes are not supported for now.");
+    }
+  }
+};
+
+}  // namespace

python/sglang/jit_kernel/csrc/elementwise/fused_eh_norm.cuh

@@ -0,0 +1,113 @@
+#include <sgl_kernel/tensor.h>
+#include <sgl_kernel/utils.h>
+
+#include <sgl_kernel/runtime.cuh>
+#include <sgl_kernel/tile.cuh>
+#include <sgl_kernel/utils.cuh>
+#include <sgl_kernel/vec.cuh>
+
+#include <sgl_kernel/impl/norm.cuh>
+
+#include <tvm/ffi/container/tensor.h>
+
+namespace {
+
+struct FusedEHNormParams {
+  const void* __restrict__ embeds;
+  const void* __restrict__ previous_hidden;
+  const void* __restrict__ enorm_weight;
+  const void* __restrict__ hnorm_weight;
+  void* __restrict__ output;
+  int64_t embeds_stride;
+  int64_t previous_hidden_stride;
+  int64_t output_stride;
+  float eps;
+};
+
+template <int64_t kHidden, bool kUsePDL, typename T>
+__global__ void fused_eh_norm_kernel(const __grid_constant__ FusedEHNormParams params) {
+  using namespace device;
+  using Storage = norm::StorageType<T, kHidden>;
+
+  constexpr auto kNumThreads = host::norm::get_cta_threads<T, kHidden>();
+  constexpr auto kNumWarps = kNumThreads / kWarpThreads;
+
+  const auto embeds = static_cast<const T*>(pointer::offset<T>(params.embeds, blockIdx.x * params.embeds_stride));
+  const auto previous_hidden =
+      static_cast<const T*>(pointer::offset<T>(params.previous_hidden, blockIdx.x * params.previous_hidden_stride));
+  const auto enorm_weight = static_cast<const T*>(params.enorm_weight);
+  const auto hnorm_weight = static_cast<const T*>(params.hnorm_weight);
+  const auto output = static_cast<T*>(pointer::offset<T>(params.output, blockIdx.x * params.output_stride));
+
+  const auto gmem = tile::Memory<Storage>::cta(kNumThreads);
+  __shared__ float smem[norm::kSmemBufferSize];
+
+  PDLWaitPrimary<kUsePDL>();
+
+  const auto embeds_vec = gmem.load(embeds);
+  const auto enorm_weight_vec = gmem.load(enorm_weight);
+  const auto embeds_output_vec =
+      norm::apply_norm_cta<kHidden>(embeds_vec, enorm_weight_vec, params.eps, smem, kNumWarps);
+  gmem.store(output, embeds_output_vec);
+
+  const auto prev_vec = gmem.load(previous_hidden);
+  const auto hnorm_weight_vec = gmem.load(hnorm_weight);
+  const auto prev_output_vec = norm::apply_norm_cta<kHidden>(prev_vec, hnorm_weight_vec, params.eps, smem, kNumWarps);
+  gmem.store(output + kHidden, prev_output_vec);
+
+  PDLTriggerSecondary<kUsePDL>();
+}
+
+template <int64_t kHidden, bool kUsePDL, typename T>
+struct FusedEHNormKernel {
+  static_assert(host::norm::is_config_supported<T, kHidden>(), "Unsupported norm configuration");
+  static_assert(host::norm::should_use_cta<T, kHidden>(), "fused_eh_norm requires CTA norm");
+  static constexpr auto kernel = fused_eh_norm_kernel<kHidden, kUsePDL, T>;
+  static constexpr uint32_t kBlockSize = host::norm::get_cta_threads<T, kHidden>();
+
+  static void
+  run(const tvm::ffi::TensorView embeds,
+      const tvm::ffi::TensorView previous_hidden,
+      const tvm::ffi::TensorView enorm_weight,
+      const tvm::ffi::TensorView hnorm_weight,
+      const tvm::ffi::TensorView output,
+      float eps) {
+    using namespace host;
+
+    auto N = SymbolicSize{"num_tokens"};
+    auto H = SymbolicSize{"hidden_size"};
+    auto H2 = SymbolicSize{"hidden_size_times_2"};
+    auto SE = SymbolicSize{"embeds_stride"};
+    auto SP = SymbolicSize{"previous_hidden_stride"};
+    auto SO = SymbolicSize{"output_stride"};
+    auto device_ = SymbolicDevice{};
+    H.set_value(kHidden);
+    H2.set_value(kHidden * 2);
+    device_.set_options<kDLCUDA>();
+
+    TensorMatcher({N, H}).with_strides({SE, 1}).with_dtype<T>().with_device(device_).verify(embeds);
+    TensorMatcher({N, H}).with_strides({SP, 1}).with_dtype<T>().with_device(device_).verify(previous_hidden);
+    TensorMatcher({H}).with_dtype<T>().with_device(device_).verify(enorm_weight);
+    TensorMatcher({H}).with_dtype<T>().with_device(device_).verify(hnorm_weight);
+    TensorMatcher({N, H2}).with_strides({SO, 1}).with_dtype<T>().with_device(device_).verify(output);
+
+    const auto num_tokens = static_cast<uint32_t>(N.unwrap());
+    const auto params = FusedEHNormParams{
+        .embeds = embeds.data_ptr(),
+        .previous_hidden = previous_hidden.data_ptr(),
+        .enorm_weight = enorm_weight.data_ptr(),
+        .hnorm_weight = hnorm_weight.data_ptr(),
+        .output = output.data_ptr(),
+        .embeds_stride = SE.unwrap(),
+        .previous_hidden_stride = SP.unwrap(),
+        .output_stride = SO.unwrap(),
+        .eps = eps,
+    };
+
+    const auto num_blocks = num_tokens;
+    LaunchKernel(num_blocks, kBlockSize, device_.unwrap())  //
+        .enable_pdl(kUsePDL)(kernel, params);
+  }
+};
+
+}  // namespace

python/sglang/jit_kernel/csrc/elementwise/fused_metadata_copy.cuh

@@ -0,0 +1,726 @@
+/*
+ * Fused metadata copy kernel for DSA backend CUDA graph replay.
+ * JIT-compiled version for python/sglang/jit_kernel.
+ *
+ * OVERVIEW:
+ * This kernel fuses multiple tensor copy operations (cache_seqlens, cu_seqlens_k,
+ * page_table, dsa metadata, and optional FlashMLA metadata) into single kernel
+ * launches, significantly reducing kernel launch overhead and improving CUDA
+ * graph replay performance during inference.
+ *
+ * PERFORMANCE BENEFITS:
+ * - Single kernel launch vs. multiple separate copies (3-10x faster)
+ * - Optimized memory coalescing and SM utilization
+ * - __grid_constant__ parameter passing via constant memory
+ * - Especially beneficial in CUDA graph replay scenarios
+ *
+ * DESIGN:
+ * - Unified kernel supporting all forward modes (DECODE, TARGET_VERIFY, DRAFT_EXTEND)
+ * - Structured parameter passing (SourcePointers/DestinationPointers) for clarity
+ * - Template parameters (HAS_REAL_PAGE_TABLE, HAS_FLASHMLA) for compile-time optimization
+ * - Multi-backend variant copies to 3 destinations in one kernel (for speculative decoding)
+ *
+ * USAGE:
+ * This header is included by JIT compilation system. The FusedMetadataCopyKernel
+ * and FusedMetadataCopyMultiKernel wrapper structs provide the Python-accessible interface.
+ */
+
+#pragma once
+
+#include <sgl_kernel/tensor.h>
+#include <sgl_kernel/utils.h>
+
+#include <sgl_kernel/utils.cuh>
+
+#include <tvm/ffi/container/tensor.h>
+
+#include <algorithm>  // for std::min
+#ifndef USE_ROCM
+#include <cuda_runtime.h>
+#else
+#include <hip/hip_runtime.h>
+#endif
+
+// Forward mode enum (must match Python ForwardMode in sglang/srt/layers/attention/dsa_backend.py)
+enum ForwardModeEnum { DECODE = 0, TARGET_VERIFY = 1, DRAFT_EXTEND = 2 };
+
+/**
+ * Source pointers for metadata copy operations.
+ * Groups all source tensor pointers for cleaner parameter passing.
+ * Some pointers may be nullptr depending on forward mode and feature flags.
+ */
+struct SourcePointers {
+  const int32_t* __restrict__ cache_seqlens;        // [bs] sequence lengths in cache
+  const int32_t* __restrict__ cu_seqlens_k;         // [bs+1] cumulative sequence lengths
+  const int32_t* __restrict__ page_indices;         // page table indices
+  const int32_t* __restrict__ dsa_cache_seqlens;    // DSA-specific cache lengths
+  const int32_t* __restrict__ seqlens_expanded;     // expanded sequence lengths (TARGET_VERIFY/DRAFT_EXTEND only)
+  const int32_t* __restrict__ dsa_cu_seqlens_k;     // DSA cumulative sequence lengths
+  const int32_t* __restrict__ real_page_table;      // optional real page table
+  const int32_t* __restrict__ flashmla_num_splits;  // optional FlashMLA split counts
+  const int32_t* __restrict__ flashmla_metadata;    // optional FlashMLA metadata
+};
+
+/**
+ * Destination pointers for metadata copy operations.
+ * Groups all destination tensor pointers for cleaner parameter passing.
+ * Layout matches SourcePointers for consistency.
+ */
+struct DestinationPointers {
+  int32_t* __restrict__ cache_seqlens;        // [bs] sequence lengths in cache
+  int32_t* __restrict__ cu_seqlens_k;         // [bs+1] cumulative sequence lengths
+  int32_t* __restrict__ page_table_1;         // page table (note: different name from source)
+  int32_t* __restrict__ dsa_cache_seqlens;    // DSA-specific cache lengths
+  int32_t* __restrict__ seqlens_expanded;     // expanded sequence lengths (TARGET_VERIFY/DRAFT_EXTEND only)
+  int32_t* __restrict__ dsa_cu_seqlens_k;     // DSA cumulative sequence lengths
+  int32_t* __restrict__ real_page_table;      // optional real page table
+  int32_t* __restrict__ flashmla_num_splits;  // optional FlashMLA split counts
+  int32_t* __restrict__ flashmla_metadata;    // optional FlashMLA metadata
+};
+
+/**
+ * Parameter structure for single-backend fused metadata copy kernel.
+ * Passed via __grid_constant__ for efficient constant memory access.
+ */
+struct FusedMetadataCopyParams {
+  SourcePointers src;       // Source tensor pointers
+  DestinationPointers dst;  // Destination tensor pointers
+
+  // Kernel parameters
+  int forward_mode;                // 0=DECODE, 1=TARGET_VERIFY, 2=DRAFT_EXTEND
+  int bs;                          // Batch size
+  int max_len;                     // Max length for DECODE mode
+  int max_seqlen_k;                // Max sequence length for TARGET_VERIFY/DRAFT_EXTEND
+  int seqlens_expanded_size;       // Size of expanded sequence lengths
+  int page_indices_rows;           // Number of rows in page_indices
+  int page_table_1_stride;         // Stride for page_table_1
+  int real_page_table_cols;        // Columns in real_page_table
+  int real_page_table_dst_stride;  // Stride for destination real_page_table
+  int flashmla_metadata_size;      // Size of FlashMLA metadata
+};
+
+/**
+ * Parameter structure for multi-backend fused metadata copy kernel.
+ * Enables copying from one source to three destinations in a single kernel launch.
+ * Used for speculative decoding with multiple draft backends.
+ */
+struct FusedMetadataCopyMultiParams {
+  SourcePointers src;        // Source pointers (shared across all backends)
+  DestinationPointers dst0;  // Backend 0 destination pointers
+  DestinationPointers dst1;  // Backend 1 destination pointers
+  DestinationPointers dst2;  // Backend 2 destination pointers
+
+  // Kernel parameters
+  int bs;                          // Batch size
+  int max_len;                     // Max length (DECODE mode only)
+  int seqlens_expanded_size;       // Size of expanded sequence lengths
+  int page_table_1_stride;         // Stride for page_table_1
+  int real_page_table_cols;        // Columns in real_page_table
+  int real_page_table_dst_stride;  // Stride for destination real_page_table
+  int flashmla_metadata_size;      // Size of FlashMLA metadata
+};
+
+/**
+ * Unified kernel for all forward modes (DECODE, TARGET_VERIFY, DRAFT_EXTEND).
+ * Uses runtime branches for mode selection, with template parameters for
+ * compile-time optimization of optional features.
+ *
+ * DESIGN:
+ * - Runtime branches (forward_mode) handle mode-specific logic
+ * - Template parameters (HAS_*) eliminate unused feature code at compile time
+ * - Structured parameters (SourcePointers/DestinationPointers) passed via constant memory
+ *
+ * Used by FusedMetadataCopyKernel for single-backend metadata copy.
+ *
+ * @tparam HAS_REAL_PAGE_TABLE Compile-time flag for real_page_table support
+ * @tparam HAS_FLASHMLA Compile-time flag for FlashMLA metadata support
+ */
+template <bool HAS_REAL_PAGE_TABLE, bool HAS_FLASHMLA>
+__global__ void fused_metadata_copy_kernel(const FusedMetadataCopyParams __grid_constant__ params) {
+  int tid = blockIdx.x * blockDim.x + threadIdx.x;
+  int total_threads = gridDim.x * blockDim.x;
+
+  // Unpack parameters for readability
+  const auto& src = params.src;
+  const auto& dst = params.dst;
+  const int forward_mode = params.forward_mode;
+  const int bs = params.bs;
+  const int max_len = params.max_len;
+  const int max_seqlen_k = params.max_seqlen_k;
+  const int seqlens_expanded_size = params.seqlens_expanded_size;
+  const int page_indices_rows = params.page_indices_rows;
+  const int page_table_1_stride = params.page_table_1_stride;
+  const int real_page_table_cols = params.real_page_table_cols;
+  const int real_page_table_dst_stride = params.real_page_table_dst_stride;
+  const int flashmla_metadata_size = params.flashmla_metadata_size;
+
+  // Copy cache_seqlens (bs elements) - common to all modes
+#pragma unroll 8
+  for (int i = tid; i < bs; i += total_threads) {
+    dst.cache_seqlens[i] = src.cache_seqlens[i];
+  }
+
+  // Copy cu_seqlens_k (skip first element) - common to all modes
+#pragma unroll 8
+  for (int i = tid; i < bs; i += total_threads) {
+    dst.cu_seqlens_k[i + 1] = src.cu_seqlens_k[i + 1];
+  }
+
+  // Branch 1: page_table copy (different dimensions per mode)
+  if (forward_mode == 0) {  // DECODE
+    int page_table_elements = bs * max_len;
+#pragma unroll 4
+    for (int i = tid; i < page_table_elements; i += total_threads) {
+      int row = i / max_len;
+      int col = i % max_len;
+      dst.page_table_1[row * page_table_1_stride + col] = src.page_indices[i];
+    }
+  } else {  // TARGET_VERIFY or DRAFT_EXTEND
+    int page_table_elements = page_indices_rows * max_seqlen_k;
+#pragma unroll 4
+    for (int i = tid; i < page_table_elements; i += total_threads) {
+      int row = i / max_seqlen_k;
+      int col = i % max_seqlen_k;
+      dst.page_table_1[row * page_table_1_stride + col] = src.page_indices[i];
+    }
+  }
+
+  // Branch 2: seqlens_expanded copy (only for TARGET_VERIFY/DRAFT_EXTEND)
+  if (forward_mode != 0) {  // TARGET_VERIFY or DRAFT_EXTEND
+#pragma unroll 4
+    for (int i = tid; i < seqlens_expanded_size; i += total_threads) {
+      dst.seqlens_expanded[i] = src.seqlens_expanded[i];
+    }
+  }
+
+  // Branch 3: DSA metadata copy (different loop sizes per mode)
+  if (forward_mode == 0) {  // DECODE
+#pragma unroll 8
+    for (int i = tid; i < bs; i += total_threads) {
+      dst.dsa_cache_seqlens[i] = src.dsa_cache_seqlens[i];
+    }
+
+#pragma unroll 8
+    for (int i = tid; i < bs; i += total_threads) {
+      dst.dsa_cu_seqlens_k[i + 1] = src.dsa_cu_seqlens_k[i + 1];
+    }
+  } else {  // TARGET_VERIFY or DRAFT_EXTEND
+#pragma unroll 4
+    for (int i = tid; i < seqlens_expanded_size; i += total_threads) {
+      dst.dsa_cache_seqlens[i] = src.dsa_cache_seqlens[i];
+    }
+
+#pragma unroll 4
+    for (int i = tid; i < seqlens_expanded_size; i += total_threads) {
+      dst.dsa_cu_seqlens_k[i + 1] = src.dsa_cu_seqlens_k[i + 1];
+    }
+  }
+
+  // Copy real page table - compile-time branch
+  if constexpr (HAS_REAL_PAGE_TABLE) {
+    int real_table_elements = (forward_mode == 0 ? bs : page_indices_rows) * real_page_table_cols;
+#pragma unroll 2
+    for (int i = tid; i < real_table_elements; i += total_threads) {
+      int row = i / real_page_table_cols;
+      int col = i % real_page_table_cols;
+      dst.real_page_table[row * real_page_table_dst_stride + col] =
+          src.real_page_table[row * real_page_table_cols + col];
+    }
+  }
+
+  // Branch 4: FlashMLA metadata copy (different sizes per mode)
+  if constexpr (HAS_FLASHMLA) {
+    int flashmla_size = (forward_mode == 0) ? (bs + 1) : (seqlens_expanded_size + 1);
+
+    if (forward_mode == 0) {
+#pragma unroll 8
+      for (int i = tid; i < flashmla_size; i += total_threads) {
+        dst.flashmla_num_splits[i] = src.flashmla_num_splits[i];
+      }
+    } else {
+#pragma unroll 4
+      for (int i = tid; i < flashmla_size; i += total_threads) {
+        dst.flashmla_num_splits[i] = src.flashmla_num_splits[i];
+      }
+    }
+
+#pragma unroll 2
+    for (int i = tid; i < flashmla_metadata_size; i += total_threads) {
+      dst.flashmla_metadata[i] = src.flashmla_metadata[i];
+    }
+  }
+}
+
+/**
+ * Multi-backend kernel for DECODE mode.
+ * Copies from one source to THREE destinations in a single kernel launch.
+ *
+ * PERFORMANCE: 3x faster than three separate kernel launches due to:
+ * - Reduced kernel launch overhead (1 launch instead of 3)
+ * - Improved memory coalescing (source read once, written to 3 destinations)
+ * - Better instruction-level parallelism
+ *
+ * Used by FusedMetadataCopyMultiKernel for speculative decoding scenarios.
+ *
+ * @tparam HAS_REAL_PAGE_TABLE Compile-time flag for real_page_table support
+ * @tparam HAS_FLASHMLA Compile-time flag for FlashMLA metadata support
+ */
+template <bool HAS_REAL_PAGE_TABLE, bool HAS_FLASHMLA>
+__global__ void fused_metadata_copy_multi_kernel(const FusedMetadataCopyMultiParams __grid_constant__ params) {
+  int tid = blockIdx.x * blockDim.x + threadIdx.x;
+  int total_threads = gridDim.x * blockDim.x;
+
+  // Unpack parameters for readability
+  const auto& src = params.src;
+  const auto& dst0 = params.dst0;
+  const auto& dst1 = params.dst1;
+  const auto& dst2 = params.dst2;
+  const int bs = params.bs;
+  const int max_len = params.max_len;
+  const int seqlens_expanded_size = params.seqlens_expanded_size;
+  const int page_table_1_stride = params.page_table_1_stride;
+  const int real_page_table_cols = params.real_page_table_cols;
+  const int real_page_table_dst_stride = params.real_page_table_dst_stride;
+  const int flashmla_metadata_size = params.flashmla_metadata_size;
+
+  // Copy cache_seqlens to all 3 backends
+#pragma unroll 8
+  for (int i = tid; i < bs; i += total_threads) {
+    int32_t val = src.cache_seqlens[i];
+    dst0.cache_seqlens[i] = val;
+    dst1.cache_seqlens[i] = val;
+    dst2.cache_seqlens[i] = val;
+  }
+
+  // Copy cu_seqlens_k to all 3 backends (skip first element)
+#pragma unroll 8
+  for (int i = tid; i < bs; i += total_threads) {
+    int32_t val = src.cu_seqlens_k[i + 1];
+    dst0.cu_seqlens_k[i + 1] = val;
+    dst1.cu_seqlens_k[i + 1] = val;
+    dst2.cu_seqlens_k[i + 1] = val;
+  }
+
+  // DECODE mode: copy page_table_1 to all 3 backends
+  int page_table_elements = bs * max_len;
+#pragma unroll 4
+  for (int i = tid; i < page_table_elements; i += total_threads) {
+    int row = i / max_len;
+    int col = i % max_len;
+    int32_t val = src.page_indices[i];
+    dst0.page_table_1[row * page_table_1_stride + col] = val;
+    dst1.page_table_1[row * page_table_1_stride + col] = val;
+    dst2.page_table_1[row * page_table_1_stride + col] = val;
+  }
+
+  // Copy dsa_cache_seqlens to all 3 backends
+#pragma unroll 8
+  for (int i = tid; i < bs; i += total_threads) {
+    int32_t val = src.dsa_cache_seqlens[i];
+    dst0.dsa_cache_seqlens[i] = val;
+    dst1.dsa_cache_seqlens[i] = val;
+    dst2.dsa_cache_seqlens[i] = val;
+  }
+
+  // Copy DSA cu_seqlens to all 3 backends
+#pragma unroll 8
+  for (int i = tid; i < bs; i += total_threads) {
+    int32_t val = src.dsa_cu_seqlens_k[i + 1];
+    dst0.dsa_cu_seqlens_k[i + 1] = val;
+    dst1.dsa_cu_seqlens_k[i + 1] = val;
+    dst2.dsa_cu_seqlens_k[i + 1] = val;
+  }
+
+  // Copy real page table to all 3 backends
+  if (src.real_page_table != nullptr && dst0.real_page_table != nullptr) {
+    int real_table_elements = bs * real_page_table_cols;
+#pragma unroll 2
+    for (int i = tid; i < real_table_elements; i += total_threads) {
+      int row = i / real_page_table_cols;
+      int col = i % real_page_table_cols;
+      int src_idx = row * real_page_table_cols + col;
+      int dst_idx = row * real_page_table_dst_stride + col;
+      int32_t val = src.real_page_table[src_idx];
+      dst0.real_page_table[dst_idx] = val;
+      dst1.real_page_table[dst_idx] = val;
+      dst2.real_page_table[dst_idx] = val;
+    }
+  }
+
+  // Copy FlashMLA metadata to all 3 backends
+  if constexpr (HAS_FLASHMLA) {
+    int flashmla_size = bs + 1;
+#pragma unroll 8
+    for (int i = tid; i < flashmla_size; i += total_threads) {
+      int32_t val = src.flashmla_num_splits[i];
+      dst0.flashmla_num_splits[i] = val;
+      dst1.flashmla_num_splits[i] = val;
+      dst2.flashmla_num_splits[i] = val;
+    }
+
+#pragma unroll 2
+    for (int i = tid; i < flashmla_metadata_size; i += total_threads) {
+      int32_t val = src.flashmla_metadata[i];
+      dst0.flashmla_metadata[i] = val;
+      dst1.flashmla_metadata[i] = val;
+      dst2.flashmla_metadata[i] = val;
+    }
+  }
+}
+
+// ============================================================================
+// Host-side launcher wrappers for JIT compilation
+// ============================================================================
+
+namespace {
+
+// Launch configuration constants
+constexpr int THREADS_PER_BLOCK = 256;
+constexpr int MAX_GRID_SIZE = 1024;  // Limit to prevent excessive resource usage
+
+/**
+ * Helper function to extract a typed data pointer from a TensorView.
+ * Performs runtime type checking and returns the properly cast pointer.
+ *
+ * @tparam T The expected element type (e.g., int32_t)
+ * @param tensor The TensorView to extract the pointer from
+ * @param name The name of the tensor (for error reporting)
+ * @return Typed pointer to the tensor data
+ */
+template <typename T>
+inline const T* unwrap_data_ptr(const tvm::ffi::TensorView& tensor, const char* name) {
+  using namespace host;
+  if (tensor.data_ptr()) {
+    RuntimeCheck(is_type<T>(tensor.dtype()), "Tensor ", name, " must have dtype int32");
+  }
+  return static_cast<const T*>(tensor.data_ptr());
+}
+
+/**
+ * Helper function to extract a typed mutable data pointer from a TensorView.
+ * Performs runtime type checking and returns the properly cast pointer.
+ *
+ * @tparam T The expected element type (e.g., int32_t)
+ * @param tensor The TensorView to extract the pointer from
+ * @param name The name of the tensor (for error reporting)
+ * @return Typed mutable pointer to the tensor data
+ */
+template <typename T>
+inline T* unwrap_data_ptr_mut(const tvm::ffi::TensorView& tensor, const char* name) {
+  using namespace host;
+  if (tensor.data_ptr()) {
+    RuntimeCheck(is_type<T>(tensor.dtype()), "Tensor ", name, " must have dtype int32");
+  }
+  return static_cast<T*>(tensor.data_ptr());
+}
+
+/**
+ * Helper function to extract a typed data pointer from an Optional TensorView.
+ * Returns nullptr if the optional has no value, otherwise performs type checking.
+ *
+ * @tparam T The expected element type (e.g., int32_t)
+ * @param optional_tensor The Optional TensorView to extract the pointer from
+ * @param name The name of the tensor (for error reporting)
+ * @return Typed pointer to the tensor data, or nullptr if optional has no value
+ */
+template <typename T>
+inline const T*
+unwrap_optional_data_ptr(const tvm::ffi::Optional<tvm::ffi::TensorView>& optional_tensor, const char* name) {
+  using namespace host;
+  if (!optional_tensor.has_value()) {
+    return nullptr;
+  }
+  const auto& tensor = optional_tensor.value();
+  RuntimeCheck(is_type<T>(tensor.dtype()), "Tensor ", name, " must have dtype int32");
+  return static_cast<const T*>(tensor.data_ptr());
+}
+
+/**
+ * Helper function to extract a typed mutable data pointer from an Optional TensorView.
+ * Returns nullptr if the optional has no value, otherwise performs type checking.
+ *
+ * @tparam T The expected element type (e.g., int32_t)
+ * @param optional_tensor The Optional TensorView to extract the pointer from
+ * @param name The name of the tensor (for error reporting)
+ * @return Typed mutable pointer to the tensor data, or nullptr if optional has no value
+ */
+template <typename T>
+inline T*
+unwrap_optional_data_ptr_mut(const tvm::ffi::Optional<tvm::ffi::TensorView>& optional_tensor, const char* name) {
+  using namespace host;
+  if (!optional_tensor.has_value()) {
+    return nullptr;
+  }
+  const auto& tensor = optional_tensor.value();
+  RuntimeCheck(is_type<T>(tensor.dtype()), "Tensor ", name, " must have dtype int32");
+  return static_cast<T*>(tensor.data_ptr());
+}
+
+/**
+ * Calculate kernel launch configuration.
+ *
+ * @param total_work Total number of work items
+ * @param threads_per_block Threads per block (default: THREADS_PER_BLOCK)
+ * @return Grid dimension for kernel launch
+ */
+inline dim3 get_launch_config(int total_work, int threads_per_block = THREADS_PER_BLOCK) {
+  int num_blocks = (total_work + threads_per_block - 1) / threads_per_block;
+  // Limit grid size to prevent excessive resource usage while ensuring coverage
+  num_blocks = std::min(num_blocks, MAX_GRID_SIZE);
+  return dim3(num_blocks);
+}
+
+/**
+ * JIT wrapper for single-backend fused metadata copy kernel.
+ *
+ * This struct provides a unified interface for launching the fused metadata copy
+ * kernel with different forward modes. It constructs the parameter struct and
+ * launches the unified kernel.
+ *
+ * IMPLEMENTATION:
+ * - Extracts raw pointers from TensorView objects
+ * - Constructs FusedMetadataCopyParams with nested SourcePointers/DestinationPointers
+ * - Calculates grid configuration based on maximum work size
+ * - Launches fused_metadata_copy_kernel with __grid_constant__ parameters
+ *
+ * @tparam FORWARD_MODE Forward mode: 0=DECODE, 1=TARGET_VERIFY, 2=DRAFT_EXTEND
+ * @tparam HAS_REAL_PAGE_TABLE Whether real_page_table tensors are present
+ * @tparam HAS_FLASHMLA Whether FlashMLA metadata tensors are present
+ */
+template <int FORWARD_MODE, bool HAS_REAL_PAGE_TABLE, bool HAS_FLASHMLA>
+struct FusedMetadataCopyKernel {
+  static_assert(
+      FORWARD_MODE >= 0 && FORWARD_MODE <= 2,
+      "FORWARD_MODE must be 0 (DECODE), 1 (TARGET_VERIFY), or 2 (DRAFT_EXTEND)");
+
+  static void
+  run(const tvm::ffi::TensorView cache_seqlens_src,
+      const tvm::ffi::TensorView cu_seqlens_k_src,
+      const tvm::ffi::TensorView page_indices_src,
+      const tvm::ffi::TensorView dsa_cache_seqlens_src,
+      const tvm::ffi::Optional<tvm::ffi::TensorView> seqlens_expanded_src,
+      const tvm::ffi::TensorView dsa_cu_seqlens_k_src,
+      const tvm::ffi::Optional<tvm::ffi::TensorView> real_page_table_src,
+      const tvm::ffi::Optional<tvm::ffi::TensorView> flashmla_num_splits_src,
+      const tvm::ffi::Optional<tvm::ffi::TensorView> flashmla_metadata_src,
+      const tvm::ffi::TensorView cache_seqlens_dst,
+      const tvm::ffi::TensorView cu_seqlens_k_dst,
+      const tvm::ffi::TensorView page_table_1_dst,
+      const tvm::ffi::TensorView dsa_cache_seqlens_dst,
+      const tvm::ffi::Optional<tvm::ffi::TensorView> seqlens_expanded_dst,
+      const tvm::ffi::TensorView dsa_cu_seqlens_k_dst,
+      const tvm::ffi::Optional<tvm::ffi::TensorView> real_page_table_dst,
+      const tvm::ffi::Optional<tvm::ffi::TensorView> flashmla_num_splits_dst,
+      const tvm::ffi::Optional<tvm::ffi::TensorView> flashmla_metadata_dst,
+      int bs,
+      int max_len,
+      int max_seqlen_k,
+      int seqlens_expanded_size) {
+    using namespace host;
+
+    // Build parameter struct with nested source/destination pointers
+    // unwrap_data_ptr and unwrap_optional_data_ptr perform dtype validation
+    const auto params = FusedMetadataCopyParams{
+        .src =
+            {
+                .cache_seqlens = unwrap_data_ptr<int32_t>(cache_seqlens_src, "cache_seqlens_src"),
+                .cu_seqlens_k = unwrap_data_ptr<int32_t>(cu_seqlens_k_src, "cu_seqlens_k_src"),
+                .page_indices = unwrap_data_ptr<int32_t>(page_indices_src, "page_indices_src"),
+                .dsa_cache_seqlens = unwrap_data_ptr<int32_t>(dsa_cache_seqlens_src, "dsa_cache_seqlens_src"),
+                .seqlens_expanded = unwrap_optional_data_ptr<int32_t>(seqlens_expanded_src, "seqlens_expanded_src"),
+                .dsa_cu_seqlens_k = unwrap_data_ptr<int32_t>(dsa_cu_seqlens_k_src, "dsa_cu_seqlens_k_src"),
+                .real_page_table = unwrap_optional_data_ptr<int32_t>(real_page_table_src, "real_page_table_src"),
+                .flashmla_num_splits =
+                    unwrap_optional_data_ptr<int32_t>(flashmla_num_splits_src, "flashmla_num_splits_src"),
+                .flashmla_metadata = unwrap_optional_data_ptr<int32_t>(flashmla_metadata_src, "flashmla_metadata_src"),
+            },
+        .dst =
+            {
+                .cache_seqlens = unwrap_data_ptr_mut<int32_t>(cache_seqlens_dst, "cache_seqlens_dst"),
+                .cu_seqlens_k = unwrap_data_ptr_mut<int32_t>(cu_seqlens_k_dst, "cu_seqlens_k_dst"),
+                .page_table_1 = unwrap_data_ptr_mut<int32_t>(page_table_1_dst, "page_table_1_dst"),
+                .dsa_cache_seqlens = unwrap_data_ptr_mut<int32_t>(dsa_cache_seqlens_dst, "dsa_cache_seqlens_dst"),
+                .seqlens_expanded = unwrap_optional_data_ptr_mut<int32_t>(seqlens_expanded_dst, "seqlens_expanded_dst"),
+                .dsa_cu_seqlens_k = unwrap_data_ptr_mut<int32_t>(dsa_cu_seqlens_k_dst, "dsa_cu_seqlens_k_dst"),
+                .real_page_table = unwrap_optional_data_ptr_mut<int32_t>(real_page_table_dst, "real_page_table_dst"),
+                .flashmla_num_splits =
+                    unwrap_optional_data_ptr_mut<int32_t>(flashmla_num_splits_dst, "flashmla_num_splits_dst"),
+                .flashmla_metadata =
+                    unwrap_optional_data_ptr_mut<int32_t>(flashmla_metadata_dst, "flashmla_metadata_dst"),
+            },
+        .forward_mode = FORWARD_MODE,
+        .bs = bs,
+        .max_len = max_len,
+        .max_seqlen_k = max_seqlen_k,
+        .seqlens_expanded_size = seqlens_expanded_size,
+        .page_indices_rows = static_cast<int>(page_indices_src.shape()[0]),
+        .page_table_1_stride = static_cast<int>(page_table_1_dst.shape()[1]),
+        .real_page_table_cols =
+            real_page_table_src.has_value() ? static_cast<int>(real_page_table_src.value().shape()[1]) : 0,
+        .real_page_table_dst_stride =
+            real_page_table_dst.has_value() ? static_cast<int>(real_page_table_dst.value().stride(0)) : 0,
+        .flashmla_metadata_size =
+            flashmla_metadata_src.has_value() ? static_cast<int>(flashmla_metadata_src.value().numel()) : 0,
+    };
+
+    // Calculate grid configuration
+    int max_elements = std::max(
+        {bs,
+         params.page_indices_rows * max_seqlen_k,
+         seqlens_expanded_size,
+         HAS_FLASHMLA ? (seqlens_expanded_size + 1) : 0,
+         HAS_FLASHMLA ? params.flashmla_metadata_size : 0});
+
+    dim3 grid = get_launch_config(max_elements);
+    dim3 block(THREADS_PER_BLOCK);
+    DLDevice device = cache_seqlens_src.device();
+
+    // Launch unified kernel with params struct
+    host::LaunchKernel(grid, block, device)(fused_metadata_copy_kernel<HAS_REAL_PAGE_TABLE, HAS_FLASHMLA>, params);
+  }
+};
+
+/**
+ * JIT wrapper for multi-backend fused metadata copy kernel.
+ *
+ * This kernel optimizes the common case where metadata needs to be copied from
+ * one source to THREE destination backends in a single kernel launch. This is
+ * 3x faster than launching three separate kernels due to:
+ * - Reduced kernel launch overhead (1 launch instead of 3)
+ * - Improved memory coalescing (source read once, written to 3 destinations)
+ * - Better GPU occupancy and instruction-level parallelism
+ *
+ * USAGE: Primarily for speculative decoding with multiple draft models, where
+ * the same source metadata needs to be replicated to multiple backend contexts.
+ *
+ * LIMITATION: Currently only supports DECODE mode, which is the most frequently
+ * used mode in speculative decoding scenarios.
+ *
+ * IMPLEMENTATION:
+ * - Constructs FusedMetadataCopyMultiParams with 1 SourcePointers + 3 DestinationPointers
+ * - Launches fused_metadata_copy_multi_kernel with __grid_constant__ parameters
+ *
+ * @tparam HAS_REAL_PAGE_TABLE Whether real_page_table tensors are present
+ * @tparam HAS_FLASHMLA Whether FlashMLA metadata tensors are present
+ */
+template <bool HAS_REAL_PAGE_TABLE, bool HAS_FLASHMLA>
+struct FusedMetadataCopyMultiKernel {
+  static void
+  run(const tvm::ffi::TensorView cache_seqlens_src,
+      const tvm::ffi::TensorView cu_seqlens_k_src,
+      const tvm::ffi::TensorView page_indices_src,
+      const tvm::ffi::TensorView dsa_cache_seqlens_src,
+      const tvm::ffi::TensorView dsa_cu_seqlens_k_src,
+      const tvm::ffi::Optional<tvm::ffi::TensorView> real_page_table_src,
+      const tvm::ffi::Optional<tvm::ffi::TensorView> flashmla_num_splits_src,
+      const tvm::ffi::Optional<tvm::ffi::TensorView> flashmla_metadata_src,
+      const tvm::ffi::TensorView cache_seqlens_dst0,
+      const tvm::ffi::TensorView cu_seqlens_k_dst0,
+      const tvm::ffi::TensorView page_table_1_dst0,
+      const tvm::ffi::TensorView dsa_cache_seqlens_dst0,
+      const tvm::ffi::TensorView dsa_cu_seqlens_k_dst0,
+      const tvm::ffi::Optional<tvm::ffi::TensorView> real_page_table_dst0,
+      const tvm::ffi::Optional<tvm::ffi::TensorView> flashmla_num_splits_dst0,
+      const tvm::ffi::Optional<tvm::ffi::TensorView> flashmla_metadata_dst0,
+      const tvm::ffi::TensorView cache_seqlens_dst1,
+      const tvm::ffi::TensorView cu_seqlens_k_dst1,
+      const tvm::ffi::TensorView page_table_1_dst1,
+      const tvm::ffi::TensorView dsa_cache_seqlens_dst1,
+      const tvm::ffi::TensorView dsa_cu_seqlens_k_dst1,
+      const tvm::ffi::Optional<tvm::ffi::TensorView> real_page_table_dst1,
+      const tvm::ffi::Optional<tvm::ffi::TensorView> flashmla_num_splits_dst1,
+      const tvm::ffi::Optional<tvm::ffi::TensorView> flashmla_metadata_dst1,
+      const tvm::ffi::TensorView cache_seqlens_dst2,
+      const tvm::ffi::TensorView cu_seqlens_k_dst2,
+      const tvm::ffi::TensorView page_table_1_dst2,
+      const tvm::ffi::TensorView dsa_cache_seqlens_dst2,
+      const tvm::ffi::TensorView dsa_cu_seqlens_k_dst2,
+      const tvm::ffi::Optional<tvm::ffi::TensorView> real_page_table_dst2,
+      const tvm::ffi::Optional<tvm::ffi::TensorView> flashmla_num_splits_dst2,
+      const tvm::ffi::Optional<tvm::ffi::TensorView> flashmla_metadata_dst2,
+      int bs,
+      int max_len,
+      int seqlens_expanded_size) {
+    using namespace host;
+
+    // Build parameter struct with nested source/destination pointers
+    // unwrap_data_ptr and unwrap_optional_data_ptr perform dtype validation
+    const auto params = FusedMetadataCopyMultiParams{
+        .src =
+            {
+                .cache_seqlens = unwrap_data_ptr<int32_t>(cache_seqlens_src, "cache_seqlens_src"),
+                .cu_seqlens_k = unwrap_data_ptr<int32_t>(cu_seqlens_k_src, "cu_seqlens_k_src"),
+                .page_indices = unwrap_data_ptr<int32_t>(page_indices_src, "page_indices_src"),
+                .dsa_cache_seqlens = unwrap_data_ptr<int32_t>(dsa_cache_seqlens_src, "dsa_cache_seqlens_src"),
+                .seqlens_expanded = nullptr,  // Not used in multi-backend DECODE mode
+                .dsa_cu_seqlens_k = unwrap_data_ptr<int32_t>(dsa_cu_seqlens_k_src, "dsa_cu_seqlens_k_src"),
+                .real_page_table = unwrap_optional_data_ptr<int32_t>(real_page_table_src, "real_page_table_src"),
+                .flashmla_num_splits =
+                    unwrap_optional_data_ptr<int32_t>(flashmla_num_splits_src, "flashmla_num_splits_src"),
+                .flashmla_metadata = unwrap_optional_data_ptr<int32_t>(flashmla_metadata_src, "flashmla_metadata_src"),
+            },
+        .dst0 =
+            {
+                .cache_seqlens = unwrap_data_ptr_mut<int32_t>(cache_seqlens_dst0, "cache_seqlens_dst0"),
+                .cu_seqlens_k = unwrap_data_ptr_mut<int32_t>(cu_seqlens_k_dst0, "cu_seqlens_k_dst0"),
+                .page_table_1 = unwrap_data_ptr_mut<int32_t>(page_table_1_dst0, "page_table_1_dst0"),
+                .dsa_cache_seqlens = unwrap_data_ptr_mut<int32_t>(dsa_cache_seqlens_dst0, "dsa_cache_seqlens_dst0"),
+                .seqlens_expanded = nullptr,
+                .dsa_cu_seqlens_k = unwrap_data_ptr_mut<int32_t>(dsa_cu_seqlens_k_dst0, "dsa_cu_seqlens_k_dst0"),
+                .real_page_table = unwrap_optional_data_ptr_mut<int32_t>(real_page_table_dst0, "real_page_table_dst0"),
+                .flashmla_num_splits =
+                    unwrap_optional_data_ptr_mut<int32_t>(flashmla_num_splits_dst0, "flashmla_num_splits_dst0"),
+                .flashmla_metadata =
+                    unwrap_optional_data_ptr_mut<int32_t>(flashmla_metadata_dst0, "flashmla_metadata_dst0"),
+            },
+        .dst1 =
+            {
+                .cache_seqlens = unwrap_data_ptr_mut<int32_t>(cache_seqlens_dst1, "cache_seqlens_dst1"),
+                .cu_seqlens_k = unwrap_data_ptr_mut<int32_t>(cu_seqlens_k_dst1, "cu_seqlens_k_dst1"),
+                .page_table_1 = unwrap_data_ptr_mut<int32_t>(page_table_1_dst1, "page_table_1_dst1"),
+                .dsa_cache_seqlens = unwrap_data_ptr_mut<int32_t>(dsa_cache_seqlens_dst1, "dsa_cache_seqlens_dst1"),
+                .seqlens_expanded = nullptr,
+                .dsa_cu_seqlens_k = unwrap_data_ptr_mut<int32_t>(dsa_cu_seqlens_k_dst1, "dsa_cu_seqlens_k_dst1"),
+                .real_page_table = unwrap_optional_data_ptr_mut<int32_t>(real_page_table_dst1, "real_page_table_dst1"),
+                .flashmla_num_splits =
+                    unwrap_optional_data_ptr_mut<int32_t>(flashmla_num_splits_dst1, "flashmla_num_splits_dst1"),
+                .flashmla_metadata =
+                    unwrap_optional_data_ptr_mut<int32_t>(flashmla_metadata_dst1, "flashmla_metadata_dst1"),
+            },
+        .dst2 =
+            {
+                .cache_seqlens = unwrap_data_ptr_mut<int32_t>(cache_seqlens_dst2, "cache_seqlens_dst2"),
+                .cu_seqlens_k = unwrap_data_ptr_mut<int32_t>(cu_seqlens_k_dst2, "cu_seqlens_k_dst2"),
+                .page_table_1 = unwrap_data_ptr_mut<int32_t>(page_table_1_dst2, "page_table_1_dst2"),
+                .dsa_cache_seqlens = unwrap_data_ptr_mut<int32_t>(dsa_cache_seqlens_dst2, "dsa_cache_seqlens_dst2"),
+                .seqlens_expanded = nullptr,
+                .dsa_cu_seqlens_k = unwrap_data_ptr_mut<int32_t>(dsa_cu_seqlens_k_dst2, "dsa_cu_seqlens_k_dst2"),
+                .real_page_table = unwrap_optional_data_ptr_mut<int32_t>(real_page_table_dst2, "real_page_table_dst2"),
+                .flashmla_num_splits =
+                    unwrap_optional_data_ptr_mut<int32_t>(flashmla_num_splits_dst2, "flashmla_num_splits_dst2"),
+                .flashmla_metadata =
+                    unwrap_optional_data_ptr_mut<int32_t>(flashmla_metadata_dst2, "flashmla_metadata_dst2"),
+            },
+        .bs = bs,
+        .max_len = max_len,
+        .seqlens_expanded_size = seqlens_expanded_size,
+        .page_table_1_stride = static_cast<int>(page_table_1_dst0.shape()[1]),
+        .real_page_table_cols =
+            real_page_table_src.has_value() ? static_cast<int>(real_page_table_src.value().shape()[1]) : 0,
+        .real_page_table_dst_stride =
+            real_page_table_dst0.has_value() ? static_cast<int>(real_page_table_dst0.value().stride(0)) : 0,
+        .flashmla_metadata_size =
+            flashmla_metadata_src.has_value() ? static_cast<int>(flashmla_metadata_src.value().numel()) : 0,
+    };
+
+    dim3 grid = get_launch_config(bs * max_len);
+    dim3 block(THREADS_PER_BLOCK);
+    DLDevice device = cache_seqlens_src.device();
+
+    // Launch multi-backend kernel with params struct
+    host::LaunchKernel(grid, block, device)(
+        fused_metadata_copy_multi_kernel<HAS_REAL_PAGE_TABLE, HAS_FLASHMLA>, params);
+  }
+};
+
+}  // namespace

python/sglang/jit_kernel/csrc/elementwise/fused_qknorm_rope.cuh

@@ -0,0 +1,346 @@
+/*
+ * Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * 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.
+ */
+// Adapted from
+// https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/tensorrt_llm/kernels/fusedQKNormRopeKernel.cu
+
+#include <sgl_kernel/tensor.h>
+#include <sgl_kernel/utils.h>
+
+#include <sgl_kernel/type.cuh>
+#include <sgl_kernel/utils.cuh>
+#include <sgl_kernel/vec.cuh>
+#include <sgl_kernel/warp.cuh>
+
+#include <tvm/ffi/container/tensor.h>
+
+#include <cmath>
+#include <cuda_bf16.h>
+#include <cuda_runtime.h>
+
+namespace {
+
+// ---------------------------------------------------------------------------
+// YaRN-aware frequency computation
+//
+// When factor == 1.0, reduces to standard RoPE: base^(-2*half_dim/rotary_dim)
+// When factor != 1.0, blends interpolated and extrapolated frequencies.
+// ---------------------------------------------------------------------------
+
+template <bool yarn>
+__device__ inline float compute_freq(float base, int rotary_dim, int half_dim, float factor, float low, float high) {
+  float freq = powf(base, -2.0f * half_dim / static_cast<float>(rotary_dim));
+
+  if constexpr (yarn) {
+    float inv_freq_extrapolation = freq;
+    float inv_freq_interpolation = freq / factor;
+
+    float high_adj = high;
+    if (fabsf(low - high_adj) <= 1e-6f) {
+      high_adj += 0.001f;
+    }
+
+    float linear_func = (static_cast<float>(half_dim) - low) / (high_adj - low);
+    float ramp_func = fminf(fmaxf(linear_func, 0.0f), 1.0f);
+    float inv_freq_extrapolation_factor = 1.0f - ramp_func;
+
+    freq = inv_freq_interpolation * (1.0f - inv_freq_extrapolation_factor) +
+           inv_freq_extrapolation * inv_freq_extrapolation_factor;
+  }
+
+  return freq;
+}
+
+// ---------------------------------------------------------------------------
+// Fused QK-Norm + RoPE kernel
+//
+// Each warp processes one (token, head) pair.
+//   head_dim:   compile-time head dimension (64, 128, or 256)
+//   interleave: true  → interleave / GPT-J style RoPE (!is_neox)
+//               false → NeoX style RoPE (is_neox)
+// ---------------------------------------------------------------------------
+
+// interleave (GPT-J) pairs (2k,2k+1) share the same freq/theta,
+//   so sin/cos is computed once per pair and copied to the odd element,
+//   halving powf + __sincosf calls vs a naive per-element approach.
+template <int head_dim, bool interleave, bool yarn>
+__global__ void fusedQKNormRopeKernel(
+    __nv_bfloat16* qkv,  // [num_tokens, (nq+nk+nv)*head_dim], in-place
+    int const num_heads_q,
+    int const num_heads_k,
+    int const num_heads_v,
+    float const eps,
+    __nv_bfloat16 const* q_weight,  // [head_dim]
+    __nv_bfloat16 const* k_weight,  // [head_dim]
+    float const base,
+    int const* position_ids,  // [num_tokens]
+    int const num_tokens,
+    float factor,
+    float low,
+    float high,
+    float attention_factor,
+    int const rotary_dim) {
+  int const warpsPerBlock = blockDim.x / 32;
+  int const warpId = threadIdx.x / 32;
+  int const laneId = threadIdx.x % 32;
+
+  int const globalWarpIdx = blockIdx.x * warpsPerBlock + warpId;
+  int const total_qk_heads = num_heads_q + num_heads_k;
+
+  int const tokenIdx = globalWarpIdx / total_qk_heads;
+  int const localHeadIdx = globalWarpIdx % total_qk_heads;
+
+  if (tokenIdx >= num_tokens) return;
+
+  bool const isQ = localHeadIdx < num_heads_q;
+  int const headIdx = isQ ? localHeadIdx : localHeadIdx - num_heads_q;
+  int const num_heads = num_heads_q + num_heads_k + num_heads_v;
+
+  static_assert(head_dim % (32 * 2) == 0, "head_dim must be divisible by 64 (each warp handles one head)");
+  constexpr int numElemsPerThread = head_dim / 32;
+  float elements[numElemsPerThread];
+  using vec_T = device::AlignedVector<bf16_t, numElemsPerThread>;
+
+  // Compute flat offset of this warp's head in qkv
+  int offsetWarp;
+  if (isQ) {
+    offsetWarp = tokenIdx * num_heads * head_dim + headIdx * head_dim;
+  } else {
+    offsetWarp = tokenIdx * num_heads * head_dim + num_heads_q * head_dim + headIdx * head_dim;
+  }
+  int offsetThread = offsetWarp + laneId * numElemsPerThread;
+
+  // -------------------------------------------------------------------
+  // Load and compute sum-of-squares for RMSNorm
+  // -------------------------------------------------------------------
+  float sumOfSquares = 0.0f;
+  {
+    vec_T vec;
+    vec.load(qkv + offsetThread);
+    for (int i = 0; i < numElemsPerThread; i++) {
+      float val = device::cast<float>(vec[i]);
+      sumOfSquares += val * val;
+      elements[i] = val;
+    }
+  }
+
+  sumOfSquares = device::warp::reduce_sum(sumOfSquares);
+
+  // -------------------------------------------------------------------
+  // Apply RMSNorm
+  // -------------------------------------------------------------------
+  float rms_rcp = rsqrtf(sumOfSquares / static_cast<float>(head_dim) + eps);
+  {
+    vec_T wvec;
+    wvec.load((isQ ? q_weight : k_weight) + offsetThread - offsetWarp);
+    for (int i = 0; i < numElemsPerThread; i++) {
+      elements[i] *= rms_rcp * device::cast<float>(wvec[i]);
+    }
+  }
+
+  // -------------------------------------------------------------------
+  // Apply RoPE to the first rotary_dim elements
+  // -------------------------------------------------------------------
+  float pos_id = static_cast<float>(position_ids[tokenIdx]);
+  int const rotary_lanes = rotary_dim / numElemsPerThread;
+  bool const applyRotary = (laneId < rotary_lanes);
+
+  if (applyRotary) {
+    if constexpr (interleave) {
+      // Pairs (2k, 2k+1) share the same half_dim → same freq/theta.
+      // numElemsPerThread is always even (head_dim/32, head_dim in {64,128,256}),
+      // so we step by 2 and handle both elements of each pair per iteration.
+      //
+      // freq follows a geometric series across pairs: freq[k] = freq[0] * ratio^k,
+      // where ratio = base^(-2/rotary_dim). Pre-compute both outside the loop to
+      // replace all but the first powf call with a single multiply per iteration.
+      //
+      // sin/cos are applied immediately to e0/e1, eliminating the elements2,
+      // cos_vals, sin_vals intermediate arrays and reducing register pressure.
+      int const half_dim_start = laneId * numElemsPerThread / 2;
+      float freq = powf(base, -2.0f * static_cast<float>(half_dim_start) / static_cast<float>(rotary_dim));
+      float const freq_ratio = powf(base, -2.0f / static_cast<float>(rotary_dim));
+
+      for (int i = 0; i < numElemsPerThread; i += 2) {
+        float e0 = elements[i];
+        float e1 = elements[i + 1];
+
+        float f = freq;
+        if constexpr (yarn) {
+          int half_dim = half_dim_start + i / 2;
+          float inv_freq_interpolation = freq / factor;
+          float high_adj = (fabsf(low - high) <= 1e-6f) ? high + 0.001f : high;
+          float linear_func = (static_cast<float>(half_dim) - low) / (high_adj - low);
+          float ramp_func = fminf(fmaxf(linear_func, 0.0f), 1.0f);
+          float extrap_factor = 1.0f - ramp_func;
+          f = inv_freq_interpolation * (1.0f - extrap_factor) + freq * extrap_factor;
+        }
+
+        float s, c;
+        __sincosf(pos_id * f, &s, &c);
+        elements[i] = (e0 * c - e1 * s) * attention_factor;
+        elements[i + 1] = (e1 * c + e0 * s) * attention_factor;
+
+        freq *= freq_ratio;
+      }
+    } else {
+      // NeoX style: first and second halves of the rotary region are paired
+      float elements2[numElemsPerThread];
+      float cos_vals[numElemsPerThread];
+      float sin_vals[numElemsPerThread];
+
+      __syncwarp();
+      int const half_rotary_lanes = rotary_lanes / 2;
+      // Avoid UB from (1u << 32) when rotary_lanes == 32
+      unsigned int active_mask = 0xffffffffu >> (32 - rotary_lanes);
+      for (int i = 0; i < numElemsPerThread; i++) {
+        elements2[i] = __shfl_xor_sync(active_mask, elements[i], half_rotary_lanes);
+        if (laneId < half_rotary_lanes) {
+          elements2[i] = -elements2[i];
+        }
+
+        int dim_idx = laneId * numElemsPerThread + i;
+        // Remap so that both halves use the same set of frequencies
+        dim_idx = (dim_idx * 2) % rotary_dim;
+        int half_dim = dim_idx / 2;
+        float freq = compute_freq<yarn>(base, rotary_dim, half_dim, factor, low, high);
+        float theta = pos_id * freq;
+        __sincosf(theta, &sin_vals[i], &cos_vals[i]);
+      }
+      __syncwarp();
+
+      for (int i = 0; i < numElemsPerThread; i++) {
+        elements[i] = (elements[i] * cos_vals[i] + elements2[i] * sin_vals[i]) * attention_factor;
+      }
+    }
+  }
+
+  // -------------------------------------------------------------------
+  // Store (all elements: rotated + pass-through normalized)
+  // -------------------------------------------------------------------
+  {
+    vec_T vec;
+    for (int i = 0; i < numElemsPerThread; i++) {
+      vec[i] = device::cast<bf16_t>(elements[i]);
+    }
+    vec.store(qkv + offsetThread);
+  }
+}
+
+// ---------------------------------------------------------------------------
+// Host-side tvm-ffi entry point
+// ---------------------------------------------------------------------------
+
+void fused_qk_norm_rope(
+    tvm::ffi::TensorView qkv,           // [num_tokens, (nq+nk+nv)*head_dim] bf16
+    tvm::ffi::TensorView q_weight,      // [head_dim] bf16
+    tvm::ffi::TensorView k_weight,      // [head_dim] bf16
+    tvm::ffi::TensorView position_ids,  // [num_tokens] int32
+    int num_heads_q,
+    int num_heads_k,
+    int num_heads_v,
+    int head_dim,
+    float eps,
+    float base,
+    int is_neox,  // 0 = interleave style, 1 = NeoX style
+    float factor,
+    float low,
+    float high,
+    float attention_factor,
+    int rotary_dim) {
+  using namespace host;
+
+  RuntimeCheck(qkv.device().device_type == kDLCUDA, "qkv must be a CUDA tensor");
+  RuntimeCheck(qkv.is_contiguous(), "qkv must be contiguous");
+  RuntimeCheck(qkv.dtype().code == kDLBfloat && qkv.dtype().bits == 16, "qkv must be bfloat16");
+  RuntimeCheck(qkv.ndim() == 2, "qkv must be 2D: [num_tokens, (nq+nk+nv)*head_dim]");
+
+  RuntimeCheck(q_weight.is_contiguous(), "q_weight must be contiguous");
+  RuntimeCheck(q_weight.dtype().code == kDLBfloat && q_weight.dtype().bits == 16, "q_weight must be bfloat16");
+  RuntimeCheck(
+      q_weight.ndim() == 1 && static_cast<int>(q_weight.size(0)) == head_dim, "q_weight must be 1D of size head_dim");
+
+  RuntimeCheck(k_weight.is_contiguous(), "k_weight must be contiguous");
+  RuntimeCheck(k_weight.dtype().code == kDLBfloat && k_weight.dtype().bits == 16, "k_weight must be bfloat16");
+  RuntimeCheck(
+      k_weight.ndim() == 1 && static_cast<int>(k_weight.size(0)) == head_dim, "k_weight must be 1D of size head_dim");
+
+  RuntimeCheck(position_ids.device().device_type == kDLCUDA, "position_ids must be a CUDA tensor");
+  RuntimeCheck(position_ids.is_contiguous(), "position_ids must be contiguous");
+  RuntimeCheck(position_ids.dtype().code == kDLInt && position_ids.dtype().bits == 32, "position_ids must be int32");
+  RuntimeCheck(position_ids.ndim() == 1, "position_ids must be 1D: [num_tokens]");
+
+  int num_tokens = static_cast<int>(qkv.size(0));
+  int total_heads = num_heads_q + num_heads_k + num_heads_v;
+  RuntimeCheck(
+      static_cast<int>(qkv.size(1)) == total_heads * head_dim, "qkv.size(1) must equal (nq + nk + nv) * head_dim");
+  RuntimeCheck(static_cast<int>(position_ids.size(0)) == num_tokens, "position_ids must have num_tokens elements");
+
+  static_assert(
+      JIT_HEAD_DIM == 64 || JIT_HEAD_DIM == 128 || JIT_HEAD_DIM == 256, "JIT_HEAD_DIM must be 64, 128, or 256");
+  static_assert(JIT_INTERLEAVE == 0 || JIT_INTERLEAVE == 1, "JIT_INTERLEAVE must be 0 or 1");
+  static_assert(JIT_YARN == 0 || JIT_YARN == 1, "JIT_YARN must be 0 or 1");
+  RuntimeCheck(head_dim == JIT_HEAD_DIM, "head_dim mismatch with JIT-compiled kernel");
+
+  int numElemsPerThread = head_dim / 32;
+  RuntimeCheck(rotary_dim % numElemsPerThread == 0, "rotary_dim must be divisible by (head_dim / 32)");
+
+  bool neox = static_cast<bool>(is_neox);
+  if (neox) {
+    // NeoX uses __shfl_xor_sync which requires half_rotary_lanes to be a power of 2
+    int rotary_lanes = rotary_dim / numElemsPerThread;
+    int half_rotary_lanes = rotary_lanes / 2;
+    bool is_pow2 = (half_rotary_lanes >= 1) && ((half_rotary_lanes & (half_rotary_lanes - 1)) == 0);
+    RuntimeCheck(is_pow2, "half_rotary_lanes must be a power of 2 for NeoX style RoPE");
+  }
+
+  bool interleave = !neox;
+  RuntimeCheck(interleave == static_cast<bool>(JIT_INTERLEAVE), "interleave mismatch with JIT-compiled kernel");
+  bool use_yarn = (factor != 1.0f);
+  RuntimeCheck(use_yarn == static_cast<bool>(JIT_YARN), "yarn mismatch with JIT-compiled kernel");
+
+  cudaStream_t stream = LaunchKernel::resolve_device(qkv.device());
+
+  constexpr int blockSize = 256;
+  int warpsPerBlock = blockSize / 32;
+  int totalQKHeads = num_heads_q + num_heads_k;
+  int totalWarps = num_tokens * totalQKHeads;
+  int gridSize = div_ceil(totalWarps, warpsPerBlock);
+
+  auto* qkv_ptr = reinterpret_cast<__nv_bfloat16*>(qkv.data_ptr());
+  auto const* qw_ptr = reinterpret_cast<__nv_bfloat16 const*>(q_weight.data_ptr());
+  auto const* kw_ptr = reinterpret_cast<__nv_bfloat16 const*>(k_weight.data_ptr());
+  auto const* pos_ptr = reinterpret_cast<int const*>(position_ids.data_ptr());
+
+  fusedQKNormRopeKernel<JIT_HEAD_DIM, static_cast<bool>(JIT_INTERLEAVE), static_cast<bool>(JIT_YARN)>
+      <<<gridSize, blockSize, 0, stream>>>(
+          qkv_ptr,
+          num_heads_q,
+          num_heads_k,
+          num_heads_v,
+          eps,
+          qw_ptr,
+          kw_ptr,
+          base,
+          pos_ptr,
+          num_tokens,
+          factor,
+          low,
+          high,
+          attention_factor,
+          rotary_dim);
+}
+
+}  // namespace

python/sglang/jit_kernel/csrc/elementwise/kvcache.cuh

@@ -0,0 +1,209 @@
+#include <sgl_kernel/tensor.h>
+#include <sgl_kernel/utils.h>
+
+#include <sgl_kernel/tile.cuh>
+#include <sgl_kernel/utils.cuh>
+#include <sgl_kernel/vec.cuh>
+
+#include <dlpack/dlpack.h>
+#include <tvm/ffi/container/tensor.h>
+
+#include <cassert>
+#include <cstdint>
+
+namespace {
+
+struct StoreKVCacheParams {
+  const void* __restrict__ k;
+  const void* __restrict__ v;
+  void* __restrict__ k_cache;
+  void* __restrict__ v_cache;
+  const void* __restrict__ indices;
+  int64_t stride_k_bytes;
+  int64_t stride_v_bytes;
+  int64_t stride_cache_bytes;
+  int64_t stride_indices;
+  uint32_t batch_size;
+  int64_t size_limit;
+};
+
+constexpr uint32_t kNumWarps = 4;
+constexpr uint32_t kThreadsPerBlock = kNumWarps * device::kWarpThreads;
+
+/**
+ * \brief Use a single warp to copy key and value data from source to destination.
+ * Each thread in the warp copies a portion of the data in a coalesced manner.
+ * \tparam kElementBytes The size of each key/value element in bytes.
+ * \param k_src Pointer to the source key data.
+ * \param v_src Pointer to the source value data.
+ * \param k_dst Pointer to the destination key data.
+ * \param v_dst Pointer to the destination value data.
+ */
+template <int64_t kElementBytes>
+SGL_DEVICE void copy_kv_warp(
+    const void* __restrict__ k_src,
+    const void* __restrict__ v_src,
+    void* __restrict__ k_dst,
+    void* __restrict__ v_dst) {
+  using namespace device;
+  constexpr int64_t kAlignment = (kElementBytes % (16 * kWarpThreads) == 0) ? 16
+                                 : kElementBytes % (8 * kWarpThreads) == 0  ? 8
+                                 : kElementBytes % (4 * kWarpThreads) == 0  ? 4
+                                 : kElementBytes % 4 == 0                   ? 4
+                                                                            : 0;
+
+  static_assert(kAlignment > 0, "Element size must be multiple of 4 bytes");
+
+  using vec_t = AlignedStorage<uint32_t, kAlignment / 4>;
+  constexpr auto kLoopBytes = sizeof(vec_t) * kWarpThreads;
+  constexpr auto kLoopCount = kElementBytes / kLoopBytes;
+
+  const auto gmem = tile::Memory<vec_t>::warp();
+
+#pragma unroll kLoopCount
+  for (int64_t i = 0; i < kLoopCount; ++i) {
+    const auto k = gmem.load(k_src, i);
+    const auto v = gmem.load(v_src, i);
+    gmem.store(k_dst, k, i);
+    gmem.store(v_dst, v, i);
+  }
+
+  // handle the epilogue if any
+  if constexpr (kLoopCount * kLoopBytes < kElementBytes) {
+    if (gmem.in_bound(kElementBytes / sizeof(vec_t), kLoopCount)) {
+      const auto k = gmem.load(k_src, kLoopCount);
+      const auto v = gmem.load(v_src, kLoopCount);
+      gmem.store(k_dst, k, kLoopCount);
+      gmem.store(v_dst, v, kLoopCount);
+    }
+  }
+}
+
+/**
+ * \brief Kernel to store key-value pairs into the KV cache.
+ * Each element is split into multiple parts to allow parallel memory copy.
+ * \tparam kElementBytes The size of each key/value element in bytes.
+ * \tparam kSplit The number of warps that handle each element.
+ * \tparam kUsePDL Whether to use PDL feature.
+ * \tparam T The data type of the indices (`int32_t` or `int64_t`).
+ */
+template <int64_t kElementBytes, int kSplit, bool kUsePDL, typename T>
+__global__ void store_kvcache(const __grid_constant__ StoreKVCacheParams params) {
+  using namespace device;
+  constexpr auto kSplitSize = kElementBytes / kSplit;
+  const uint32_t warp_id = blockIdx.x * kNumWarps + threadIdx.x / kWarpThreads;
+  const uint32_t item_id = warp_id / kSplit;
+  const uint32_t split_id = warp_id % kSplit;
+  const auto& [
+    k_input, v_input, k_cache, v_cache, indices, // ptr
+    stride_k, stride_v, stride_cache, stride_indices, batch_size, // size
+    size_limit // bound
+  ] = params;
+  if (item_id >= batch_size) return;
+
+  const auto index_ptr = static_cast<const T*>(indices) + item_id * stride_indices;
+  PDLWaitPrimary<kUsePDL>();
+
+  const auto index = *index_ptr;
+  // A stale/OOB slot id would cause an illegal memory access in the store below;
+  // fail fast at the culprit instead. always-on (kvcache JIT compiles without NDEBUG).
+  assert(index >= 0 && index < size_limit);
+  const auto k_src = pointer::offset(k_input, item_id * stride_k, split_id * kSplitSize);
+  const auto v_src = pointer::offset(v_input, item_id * stride_v, split_id * kSplitSize);
+  const auto k_dst = pointer::offset(k_cache, index * stride_cache, split_id * kSplitSize);
+  const auto v_dst = pointer::offset(v_cache, index * stride_cache, split_id * kSplitSize);
+
+  copy_kv_warp<kSplitSize>(k_src, v_src, k_dst, v_dst);
+  PDLTriggerSecondary<kUsePDL>();
+}
+
+template <int64_t kElementBytes, bool kUsePDL>
+struct StoreKVCacheKernel {
+  static_assert(kElementBytes > 0 && kElementBytes % 4 == 0);
+
+  template <int kSplit, typename T>
+  static constexpr auto store_kernel = store_kvcache<kElementBytes, kSplit, kUsePDL, T>;
+
+  template <typename T>
+  static auto get_kernel(const int num_split) {
+    using namespace host;
+    // only apply split optimization when element size is aligned
+    if constexpr (kElementBytes % (4 * 128) == 0) {
+      if (num_split == 4) return store_kernel<4, T>;
+    }
+    if constexpr (kElementBytes % (2 * 128) == 0) {
+      if (num_split == 2) return store_kernel<2, T>;
+    }
+    if (num_split == 1) return store_kernel<1, T>;
+    Panic("Unsupported num_split {} for element size {}", num_split, kElementBytes);
+  }
+
+  static void
+  run(const tvm::ffi::TensorView k,
+      const tvm::ffi::TensorView v,
+      const tvm::ffi::TensorView k_cache,
+      const tvm::ffi::TensorView v_cache,
+      const tvm::ffi::TensorView indices,
+      const int num_split,
+      const int64_t size_limit) {
+    using namespace host;
+    auto B = SymbolicSize{"batch_size"};
+    auto D = SymbolicSize{"element_size"};
+    auto KS = SymbolicSize{"k_stride"};
+    auto VS = SymbolicSize{"v_stride"};
+    auto S = SymbolicSize{"cache_stride"};
+    auto I = SymbolicSize{"indices_stride"};
+    auto dtype = SymbolicDType{};
+    auto device = SymbolicDevice{};
+    auto indice_dtype = SymbolicDType{};
+    device.set_options<kDLCUDA, kDLROCM>();
+
+    TensorMatcher({B, D})  //
+        .with_strides({KS, 1})
+        .with_dtype(dtype)
+        .with_device(device)
+        .verify(k);
+    TensorMatcher({B, D})  //
+        .with_strides({VS, 1})
+        .with_dtype(dtype)
+        .with_device(device)
+        .verify(v);
+    TensorMatcher({-1, D})  //
+        .with_strides({S, 1})
+        .with_dtype(dtype)
+        .with_device(device)
+        .verify(k_cache)
+        .verify(v_cache);
+    TensorMatcher({B})  //
+        .with_strides({I})
+        .with_dtype<int32_t, int64_t>(indice_dtype)
+        .with_device(device)
+        .verify(indices);
+
+    const int64_t dtype_size = dtype_bytes(dtype.unwrap());
+    const uint32_t num_elements = static_cast<uint32_t>(B.unwrap());
+    RuntimeCheck(kElementBytes == dtype_size * D.unwrap());
+
+    const auto params = StoreKVCacheParams{
+        .k = k.data_ptr(),
+        .v = v.data_ptr(),
+        .k_cache = k_cache.data_ptr(),
+        .v_cache = v_cache.data_ptr(),
+        .indices = indices.data_ptr(),
+        .stride_k_bytes = KS.unwrap() * dtype_size,
+        .stride_v_bytes = VS.unwrap() * dtype_size,
+        .stride_cache_bytes = S.unwrap() * dtype_size,
+        .stride_indices = I.unwrap(),
+        .batch_size = static_cast<uint32_t>(B.unwrap()),
+        .size_limit = size_limit,
+    };
+    // select kernel and update num_split if needed
+    const auto use_int32 = indice_dtype.is_type<int32_t>();
+    const auto kernel = use_int32 ? get_kernel<int32_t>(num_split) : get_kernel<int64_t>(num_split);
+    const auto num_blocks = div_ceil(num_elements * num_split, kNumWarps);
+    LaunchKernel(num_blocks, kThreadsPerBlock, device.unwrap())  //
+        .enable_pdl(kUsePDL)(kernel, params);
+  }
+};
+
+}  // namespace

python/sglang/jit_kernel/csrc/elementwise/pos_enc.cuh

@@ -0,0 +1,313 @@
+// Adapted from
+// https://github.com/vllm-project/vllm/blob/014ece97c7aa49084a1119dca792af081a18dbc1/csrc/pos_encoding_kernels.cu
+
+#include <sgl_kernel/tensor.h>
+#include <sgl_kernel/utils.h>
+
+#include <sgl_kernel/utils.cuh>
+
+#include <tvm/ffi/container/tensor.h>
+
+#include <cuda_fp16.h>
+#include <cuda_runtime.h>
+
+namespace {
+
+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;
+  scalar_t cos, sin;
+  if (IS_NEOX) {
+    // GPT-NeoX style rotary embedding.
+    x_index = rot_offset;
+    y_index = embed_dim + rot_offset;
+    cos = SGLANG_LDG(cos_ptr + x_index);
+    sin = SGLANG_LDG(sin_ptr + x_index);
+  } else {
+    // GPT-J style rotary embedding.
+    x_index = 2 * rot_offset;
+    y_index = 2 * rot_offset + 1;
+    cos = SGLANG_LDG(cos_ptr + x_index / 2);
+    sin = SGLANG_LDG(sin_ptr + x_index / 2);
+  }
+
+  const scalar_t x = arr[x_index];
+  const scalar_t y = arr[y_index];
+  arr[x_index] = x * cos - y * sin;
+  arr[y_index] = y * cos + x * sin;
+}
+
+template <typename scalar_t, bool IS_NEOX>
+inline __device__ void apply_rotary_embedding(
+    scalar_t* __restrict__ query,  // [batch_size, seq_len, num_heads,
+                                   // head_size] or [num_tokens, num_heads,
+                                   // head_size]
+    scalar_t* __restrict__ key,    // nullptr or
+                                   // [batch_size, seq_len, num_kv_heads,
+                                   // head_size] or [num_tokens, num_kv_heads,
+                                   // head_size]
+    const scalar_t* cache_ptr,
+    const int head_size,
+    const int num_heads,
+    const int num_kv_heads,
+    const int rot_dim,
+    const int token_idx,
+    const int64_t query_stride,
+    const int64_t key_stride,
+    const int64_t head_stride) {
+  const int embed_dim = rot_dim / 2;
+  const scalar_t* cos_ptr = cache_ptr;
+  const scalar_t* sin_ptr = cache_ptr + embed_dim;
+
+  const int nq = num_heads * embed_dim;
+  for (int i = threadIdx.x; i < nq; i += blockDim.x) {
+    const int head_idx = i / embed_dim;
+    const int64_t token_head = token_idx * query_stride + head_idx * head_stride;
+    const int rot_offset = i % embed_dim;
+    apply_token_rotary_embedding<scalar_t, IS_NEOX>(query + token_head, cos_ptr, sin_ptr, rot_offset, embed_dim);
+  }
+
+  if (key != nullptr) {
+    const int nk = num_kv_heads * embed_dim;
+    for (int i = threadIdx.x; i < nk; i += blockDim.x) {
+      const int head_idx = i / embed_dim;
+      const int64_t token_head = token_idx * key_stride + head_idx * head_stride;
+      const int rot_offset = i % embed_dim;
+      apply_token_rotary_embedding<scalar_t, IS_NEOX>(key + token_head, cos_ptr, sin_ptr, rot_offset, embed_dim);
+    }
+  }
+}
+
+template <typename scalar_t, bool IS_NEOX>
+__global__ void rotary_embedding_kernel(
+    const int64_t* __restrict__ positions,       // [batch_size, seq_len] or
+                                                 // [num_tokens]
+    scalar_t* __restrict__ query,                // [batch_size, seq_len, num_heads,
+                                                 // head_size] or [num_tokens, num_heads,
+                                                 // head_size]
+    scalar_t* __restrict__ key,                  // nullptr or
+                                                 // [batch_size, seq_len, num_kv_heads,
+                                                 // head_size] or [num_tokens, num_kv_heads,
+                                                 // head_size]
+    const scalar_t* __restrict__ cos_sin_cache,  // [max_position, 2, rot_dim //
+                                                 // 2]
+    const int rot_dim,
+    const int64_t query_stride,
+    const int64_t key_stride,
+    const int64_t head_stride,
+    const int num_heads,
+    const int num_kv_heads,
+    const int head_size) {
+  // Each thread block is responsible for one token.
+  const int token_idx = blockIdx.x;
+  int64_t pos = positions[token_idx];
+  const scalar_t* cache_ptr = cos_sin_cache + pos * rot_dim;
+
+  apply_rotary_embedding<scalar_t, IS_NEOX>(
+      query,
+      key,
+      cache_ptr,
+      head_size,
+      num_heads,
+      num_kv_heads,
+      rot_dim,
+      token_idx,
+      query_stride,
+      key_stride,
+      head_stride);
+}
+
+// Helper struct to launch kernel
+template <typename scalar_t, bool IS_NEOX>
+void launch_kernel(
+    const int64_t* positions_data_ptr,
+    void* query_ptr,
+    void* key_ptr,
+    const void* cos_sin_cache_ptr,
+    int rot_dim,
+    int64_t query_stride,
+    int64_t key_stride,
+    int64_t head_stride,
+    int num_heads,
+    int num_kv_heads,
+    int head_size,
+    dim3 grid,
+    dim3 block,
+    const cudaStream_t stream) {
+  rotary_embedding_kernel<scalar_t, IS_NEOX><<<grid, block, 0, stream>>>(
+      positions_data_ptr,
+      static_cast<scalar_t*>(query_ptr),
+      static_cast<scalar_t*>(key_ptr),
+      static_cast<const scalar_t*>(cos_sin_cache_ptr),
+      rot_dim,
+      query_stride,
+      key_stride,
+      head_stride,
+      num_heads,
+      num_kv_heads,
+      head_size);
+};
+
+// Helper macro to reduce repetition
+#define DISPATCH_DTYPE(DTYPE_CODE, DTYPE_BITS, IS_NEOX, ...)                                                      \
+  if (DTYPE_CODE == kDLFloat && DTYPE_BITS == 32) {                                                               \
+    launch_kernel<float, IS_NEOX>(__VA_ARGS__);                                                                   \
+  } else if (DTYPE_CODE == kDLFloat && DTYPE_BITS == 16) {                                                        \
+    launch_kernel<half, IS_NEOX>(__VA_ARGS__);                                                                    \
+  } else if (DTYPE_CODE == kDLBfloat && DTYPE_BITS == 16) {                                                       \
+    launch_kernel<nv_bfloat16, IS_NEOX>(__VA_ARGS__);                                                             \
+  } else {                                                                                                        \
+    RuntimeCheck(                                                                                                 \
+        false, "Unsupported data type for rotary embedding. Only float32, float16, and bfloat16 are supported."); \
+  }
+
+// Helper function to dispatch based on data type
+template <bool IS_NEOX>
+void dispatch_by_dtype(
+    const int64_t* positions_data_ptr,
+    DLDataType query_dtype,
+    void* query_ptr,
+    void* key_ptr,
+    void* cos_sin_cache_ptr,
+    int rot_dim,
+    int64_t query_stride,
+    int64_t key_stride,
+    int64_t head_stride,
+    int num_heads,
+    int num_kv_heads,
+    int head_size,
+    dim3 grid,
+    dim3 block,
+    const cudaStream_t stream) {
+  using namespace host;
+  DISPATCH_DTYPE(
+      query_dtype.code,
+      query_dtype.bits,
+      IS_NEOX,
+      positions_data_ptr,
+      query_ptr,
+      key_ptr,
+      cos_sin_cache_ptr,
+      rot_dim,
+      query_stride,
+      key_stride,
+      head_stride,
+      num_heads,
+      num_kv_heads,
+      head_size,
+      grid,
+      block,
+      stream);
+}
+
+struct RotaryEmbeddingKernel {
+  static void
+  run(tvm::ffi::TensorView positions,  // [batch_size, seq_len] or [num_tokens]
+      tvm::ffi::TensorView query,      // [batch_size, seq_len, num_heads * head_size] or
+                                       // [num_tokens, num_heads * head_size] or
+                                       // [batch_size, seq_len, num_heads, head_size] or
+                                       // [num_tokens, num_heads, head_size]
+      tvm::ffi::Optional<tvm::ffi::TensorView> key,
+      // null or
+      // [batch_size, seq_len, num_kv_heads * head_size] or
+      // [num_tokens, num_kv_heads * head_size] or
+      // [batch_size, seq_len, num_heads, head_size] or
+      // [num_tokens, num_heads, head_size]
+      int64_t head_size,
+      tvm::ffi::TensorView cos_sin_cache,  // [max_position, rot_dim]
+      bool is_neox) {
+    using namespace host;
+
+    // num_tokens = batch_size * seq_len
+    int64_t num_tokens = positions.numel();
+    int32_t positions_ndim = positions.ndim();
+
+    // Make sure num_tokens dim is consistent across positions, query, and key
+    RuntimeCheck(
+        positions_ndim == 1 || positions_ndim == 2, "positions must have shape [num_tokens] or [batch_size, seq_len]");
+    if (positions_ndim == 1) {
+      RuntimeCheck(
+          query.size(0) == positions.size(0) && (!key.has_value() || key.value().size(0) == positions.size(0)),
+          "query, key and positions must have the same number of tokens");
+    }
+    if (positions_ndim == 2) {
+      RuntimeCheck(
+          query.size(0) == positions.size(0) && (!key.has_value() || key.value().size(0) == positions.size(0)) &&
+              query.size(1) == positions.size(1) && (!key.has_value() || key.value().size(1) == positions.size(1)),
+          "query, key and positions must have the same batch_size and seq_len");
+    }
+
+    // Make sure head_size is valid for query and key
+    // hidden_size = num_heads * head_size
+    int query_hidden_size = query.numel() / num_tokens;
+    int key_hidden_size = key.has_value() ? key.value().numel() / num_tokens : 0;
+    RuntimeCheck(query_hidden_size % head_size == 0);
+    RuntimeCheck(key_hidden_size % head_size == 0);
+
+    // Make sure query and key have consistent number of heads
+    int num_heads = query_hidden_size / head_size;
+    int num_kv_heads = key.has_value() ? key_hidden_size / head_size : num_heads;
+    RuntimeCheck(num_heads % num_kv_heads == 0);
+
+    int rot_dim = cos_sin_cache.size(1);
+    int seq_dim_idx = positions_ndim - 1;
+    int64_t query_stride = query.stride(seq_dim_idx);
+    int64_t key_stride = key.has_value() ? key.value().stride(seq_dim_idx) : 0;
+    // Determine head stride: for [*, heads, head_size] use stride of last dim;
+    // for flat [*, heads*head_size], heads blocks are contiguous of size
+    // head_size
+    int query_ndim = query.dim();
+    int64_t head_stride = (query_ndim == positions_ndim + 2) ? query.stride(-2) : head_size;
+
+    dim3 grid(num_tokens);
+    dim3 block(std::min<int64_t>(num_heads * rot_dim / 2, 512));
+
+    auto device = query.device();
+    const cudaStream_t stream = LaunchKernel::resolve_device(device);
+
+    auto positions_data_ptr = static_cast<const int64_t*>(positions.data_ptr());
+
+    if (is_neox) {
+      dispatch_by_dtype<true>(
+          positions_data_ptr,
+          query.dtype(),
+          query.data_ptr(),
+          key.has_value() ? key.value().data_ptr() : nullptr,
+          cos_sin_cache.data_ptr(),
+          rot_dim,
+          query_stride,
+          key_stride,
+          head_stride,
+          num_heads,
+          num_kv_heads,
+          head_size,
+          grid,
+          block,
+          stream);
+    } else {
+      dispatch_by_dtype<false>(
+          positions_data_ptr,
+          query.dtype(),
+          query.data_ptr(),
+          key.has_value() ? key.value().data_ptr() : nullptr,
+          cos_sin_cache.data_ptr(),
+          rot_dim,
+          query_stride,
+          key_stride,
+          head_stride,
+          num_heads,
+          num_kv_heads,
+          head_size,
+          grid,
+          block,
+          stream);
+    }
+  }
+};
+
+}  // namespace

python/sglang/jit_kernel/csrc/elementwise/qknorm.cuh

@@ -0,0 +1,257 @@
+#include <sgl_kernel/tensor.h>
+#include <sgl_kernel/utils.h>
+
+#include <sgl_kernel/runtime.cuh>
+#include <sgl_kernel/tile.cuh>
+#include <sgl_kernel/utils.cuh>
+
+#include <sgl_kernel/impl/norm.cuh>
+
+#include <dlpack/dlpack.h>
+#include <tvm/ffi/container/tensor.h>
+
+#include <cstdint>
+#include <cuda_bf16.h>
+#include <cuda_fp16.h>
+#include <type_traits>
+
+namespace {
+
+struct QKNormParams {
+  void* __restrict__ q;
+  void* __restrict__ k;  // k is offset by (-num_qo_heads * head_dim) elements
+  int64_t q_stride;
+  int64_t k_stride;
+  uint32_t num_qo_heads;
+  uint32_t num_kv_heads;
+  float eps;
+  const void* __restrict__ q_weight;
+  const void* __restrict__ k_weight;
+  uint32_t num_tokens;
+};
+
+constexpr uint32_t kWarpsPerBlock = 4;
+constexpr uint32_t kThreadsPerBlock = kWarpsPerBlock * device::kWarpThreads;
+
+// Warp-level kernel for head_dim <= 256
+template <int64_t kHeadDim, bool kUsePDL, typename Float>
+__global__ void fused_qknorm_warp(const QKNormParams __grid_constant__ params) {
+  using namespace device;
+  using Storage = norm::StorageType<Float, kHeadDim>;
+
+  static_assert(sizeof(Float) == 2, "Only support FP16/BF16");
+  const auto& [q, k, q_stride, k_stride, num_qo_heads, num_kv_heads, eps, q_weight, k_weight, num_tokens] = params;
+
+  const auto num_blks = gridDim.x;
+  const auto num_workers = num_blks * kWarpsPerBlock;
+  const auto num_q_and_k_heads = num_qo_heads + num_kv_heads;
+  const auto num_works = num_q_and_k_heads * num_tokens;
+  const auto start_worker_id = blockIdx.x * kWarpsPerBlock + threadIdx.x / kWarpThreads;
+  const auto gmem = tile::Memory<Storage>::warp();
+
+  PDLWaitPrimary<kUsePDL>();  // wait for primary kernel
+
+  for (auto idx = start_worker_id; idx < num_works; idx += num_workers) {
+    const int64_t token_id = idx / num_q_and_k_heads;
+    const int64_t head_id = idx % num_q_and_k_heads;
+    const auto load_q = head_id < num_qo_heads;
+    const auto input = load_q ? pointer::offset(q, 2 * (token_id * q_stride + head_id * kHeadDim))
+                              : pointer::offset(k, 2 * (token_id * k_stride + head_id * kHeadDim));
+    const auto weight = load_q ? q_weight : k_weight;
+    const auto input_vec = gmem.load(input);
+    const auto weight_vec = gmem.load(weight);
+    const auto output_vec = norm::apply_norm_warp<kHeadDim>(input_vec, weight_vec, eps);
+    gmem.store(input, output_vec);
+  }
+
+  PDLTriggerSecondary<kUsePDL>();  // launch secondary kernel
+}
+
+// For CTA level, used for head_dim > 256 (512,1024)
+template <int64_t kHeadDim, bool kUsePDL, typename Float>
+__global__ void fused_qknorm_cta(const QKNormParams __grid_constant__ params) {
+  using namespace device;
+  using Storage = norm::StorageType<Float, kHeadDim>;
+
+  constexpr auto kNumThreads = host::norm::get_cta_threads<Float, kHeadDim>();
+  constexpr auto kNumWarps = kNumThreads / kWarpThreads;
+
+  static_assert(sizeof(Float) == 2, "Only support FP16/BF16");
+  const auto& [q, k, q_stride, k_stride, num_qo_heads, num_kv_heads, eps, q_weight, k_weight, num_tokens] = params;
+
+  const auto num_q_and_k_heads = num_qo_heads + num_kv_heads;
+  const auto num_works = num_q_and_k_heads * num_tokens;
+  const auto gmem = tile::Memory<Storage>::cta(kNumThreads);
+  __shared__ float smem[norm::kSmemBufferSize];
+
+  PDLWaitPrimary<kUsePDL>();  // wait for primary kernel
+
+  for (auto idx = blockIdx.x; idx < num_works; idx += gridDim.x) {
+    const int64_t token_id = idx / num_q_and_k_heads;
+    const int64_t head_id = idx % num_q_and_k_heads;
+    const auto load_q = head_id < num_qo_heads;
+    const auto input = load_q ? pointer::offset(q, 2 * (token_id * q_stride + head_id * kHeadDim))
+                              : pointer::offset(k, 2 * (token_id * k_stride + head_id * kHeadDim));
+    const auto weight = load_q ? q_weight : k_weight;
+    const auto input_vec = gmem.load(input);
+    const auto weight_vec = gmem.load(weight);
+    const auto output_vec = norm::apply_norm_cta<kHeadDim>(input_vec, weight_vec, eps, smem, kNumWarps);
+    gmem.store(input, output_vec);
+  }
+
+  PDLTriggerSecondary<kUsePDL>();  // launch secondary kernel
+}
+
+// Warp-level kernel struct for head_dim <= 256
+template <int64_t kHeadDim, bool kUsePDL, typename DType>
+struct QKNormKernelWarp {
+  static_assert(std::is_same_v<DType, fp16_t> || std::is_same_v<DType, bf16_t>);
+  static_assert(!host::norm::should_use_cta<DType, kHeadDim>(), "Use QKNormKernelCTA for head_dim > 256");
+  static constexpr auto kernel = fused_qknorm_warp<kHeadDim, kUsePDL, DType>;
+
+  static void
+  run(const tvm::ffi::TensorView q,
+      const tvm::ffi::TensorView k,
+      const tvm::ffi::TensorView q_weight,
+      const tvm::ffi::TensorView k_weight,
+      float eps) {
+    using namespace host;
+
+    auto N = SymbolicSize{"num_tokens"};
+    auto Q = SymbolicSize{"num_qo_heads"};
+    auto K = SymbolicSize{"num_kv_heads"};
+    auto D = SymbolicSize{"head_dim"};
+    auto Sq = SymbolicSize{"q_stride"};
+    auto Sk = SymbolicSize{"k_stride"};
+    auto device = SymbolicDevice{};
+    D.set_value(kHeadDim);
+    device.set_options<kDLCUDA>();
+
+    TensorMatcher({N, Q, D})  // q input
+        .with_strides({Sq, D, 1})
+        .with_dtype<DType>()
+        .with_device(device)
+        .verify(q);
+    TensorMatcher({N, K, D})  // k input
+        .with_strides({Sk, D, 1})
+        .with_dtype<DType>()
+        .with_device(device)
+        .verify(k);
+    TensorMatcher({D})  // weight
+        .with_dtype<DType>()
+        .with_device(device)
+        .verify(q_weight)
+        .verify(k_weight);
+
+    const auto num_tokens = static_cast<uint32_t>(N.unwrap());
+    const auto num_qo_heads = static_cast<uint32_t>(Q.unwrap());
+    const auto num_kv_heads = static_cast<uint32_t>(K.unwrap());
+
+    // NOTE: we offset the k here to reduce computation cost in the kernel
+    const auto params = QKNormParams{
+        .q = q.data_ptr(),
+        .k = pointer::offset(k.data_ptr(), -2 * static_cast<int64_t>(num_qo_heads) * kHeadDim),
+        .q_stride = static_cast<int64_t>(Sq.unwrap()),
+        .k_stride = static_cast<int64_t>(Sk.unwrap()),
+        .num_qo_heads = num_qo_heads,
+        .num_kv_heads = num_kv_heads,
+        .eps = eps,
+        .q_weight = q_weight.data_ptr(),
+        .k_weight = k_weight.data_ptr(),
+        .num_tokens = num_tokens,
+    };
+
+    static const uint32_t max_occupancy = runtime::get_blocks_per_sm(kernel, kThreadsPerBlock);
+    static const uint32_t kNumSM = runtime::get_sm_count(device.unwrap().device_id);
+
+    // choose kernel based on dtype
+    const auto num_works = (num_qo_heads + num_kv_heads) * num_tokens;
+    const auto needed_blocks = div_ceil(num_works, kWarpsPerBlock);
+
+    // we use persistent kernel, which limit the number of blocks to reduce overhead
+    const auto num_blocks = std::min(kNumSM * max_occupancy, needed_blocks);
+    LaunchKernel(num_blocks, kThreadsPerBlock, device.unwrap())  //
+        .enable_pdl(kUsePDL)(kernel, params);
+  }
+};
+
+// This goes with fused_qknorm_cta
+template <int64_t kHeadDim, bool kUsePDL, typename DType>
+struct QKNormKernelCTA {
+  static_assert(std::is_same_v<DType, fp16_t> || std::is_same_v<DType, bf16_t>);
+  static_assert(host::norm::should_use_cta<DType, kHeadDim>(), "Use QKNormKernelWarp for head_dim <= 256");
+  static constexpr auto kernel = fused_qknorm_cta<kHeadDim, kUsePDL, DType>;
+  static constexpr auto kNumThreads = host::norm::get_cta_threads<DType, kHeadDim>();
+
+  static void
+  run(const tvm::ffi::TensorView q,
+      const tvm::ffi::TensorView k,
+      const tvm::ffi::TensorView q_weight,
+      const tvm::ffi::TensorView k_weight,
+      float eps) {
+    using namespace host;
+
+    auto N = SymbolicSize{"num_tokens"};
+    auto Q = SymbolicSize{"num_qo_heads"};
+    auto K = SymbolicSize{"num_kv_heads"};
+    auto D = SymbolicSize{"head_dim"};
+    auto Sq = SymbolicSize{"q_stride"};
+    auto Sk = SymbolicSize{"k_stride"};
+    auto device = SymbolicDevice{};
+    D.set_value(kHeadDim);
+    device.set_options<kDLCUDA>();
+
+    TensorMatcher({N, Q, D})  // q input
+        .with_strides({Sq, D, 1})
+        .with_dtype<DType>()
+        .with_device(device)
+        .verify(q);
+    TensorMatcher({N, K, D})  // k input
+        .with_strides({Sk, D, 1})
+        .with_dtype<DType>()
+        .with_device(device)
+        .verify(k);
+    TensorMatcher({D})  // weight
+        .with_dtype<DType>()
+        .with_device(device)
+        .verify(q_weight)
+        .verify(k_weight);
+
+    const auto num_tokens = static_cast<uint32_t>(N.unwrap());
+    const auto num_qo_heads = static_cast<uint32_t>(Q.unwrap());
+    const auto num_kv_heads = static_cast<uint32_t>(K.unwrap());
+
+    // NOTE: we offset the k here to reduce computation cost in the kernel
+    const auto params = QKNormParams{
+        .q = q.data_ptr(),
+        .k = pointer::offset(k.data_ptr(), -2 * static_cast<int64_t>(num_qo_heads) * kHeadDim),
+        .q_stride = static_cast<int64_t>(Sq.unwrap()),
+        .k_stride = static_cast<int64_t>(Sk.unwrap()),
+        .num_qo_heads = num_qo_heads,
+        .num_kv_heads = num_kv_heads,
+        .eps = eps,
+        .q_weight = q_weight.data_ptr(),
+        .k_weight = k_weight.data_ptr(),
+        .num_tokens = num_tokens,
+    };
+
+    static const uint32_t max_occupancy = runtime::get_blocks_per_sm(kernel, kNumThreads);
+    static const uint32_t kNumSM = runtime::get_sm_count(device.unwrap().device_id);
+
+    const auto num_works = (num_qo_heads + num_kv_heads) * num_tokens;
+
+    // we use persistent kernel, which limit the number of blocks to reduce overhead
+    const auto num_blocks = std::min<uint32_t>(num_works, max_occupancy * kNumSM);
+    LaunchKernel(num_blocks, kNumThreads, device.unwrap())  //
+        .enable_pdl(kUsePDL)(kernel, params);
+  }
+};
+
+// Unified dispatch: select warp or CTA kernel based on head_dim
+template <int64_t kHeadDim, bool kUsePDL, typename DType>
+using QKNormKernel = std::conditional_t<
+    host::norm::should_use_cta<DType, kHeadDim>(),
+    QKNormKernelCTA<kHeadDim, kUsePDL, DType>,
+    QKNormKernelWarp<kHeadDim, kUsePDL, DType>>;
+
+}  // namespace

python/sglang/jit_kernel/csrc/elementwise/qknorm_across_heads.cuh

@@ -0,0 +1,179 @@
+#include <sgl_kernel/tensor.h>
+#include <sgl_kernel/utils.h>
+
+#include <sgl_kernel/tile.cuh>
+#include <sgl_kernel/type.cuh>
+#include <sgl_kernel/utils.cuh>
+#include <sgl_kernel/vec.cuh>
+
+#include <cooperative_groups/reduce.h>
+#include <tvm/ffi/container/tensor.h>
+
+#include <cooperative_groups.h>
+#include <type_traits>
+
+namespace {
+
+template <typename T, int VEC_SIZE_IN_BYTE>
+struct VecTypeTrait;
+
+template <>
+struct VecTypeTrait<bf16_t, 16> {
+  using packed_t = packed_t<bf16_t>;
+  using vec_t = device::AlignedVector<packed_t, 4>;
+};
+
+template <>
+struct VecTypeTrait<fp16_t, 16> {
+  using packed_t = packed_t<fp16_t>;
+  using vec_t = device::AlignedVector<packed_t, 4>;
+};
+
+template <>
+struct VecTypeTrait<bf16_t, 32> {
+  using packed_t = packed_t<bf16_t>;
+  using vec_t = device::AlignedVector<packed_t, 8>;
+};
+
+template <>
+struct VecTypeTrait<fp16_t, 32> {
+  using packed_t = packed_t<fp16_t>;
+  using vec_t = device::AlignedVector<packed_t, 8>;
+};
+
+template <typename packed_t>
+SGL_DEVICE packed_t rms(const packed_t& val, const packed_t& weight, float rsqrt_square_sum) {
+  float2 valf = device::cast<fp32x2_t, packed_t>(val);
+  float2 weightf = device::cast<fp32x2_t, packed_t>(weight);
+  return device::cast<packed_t, fp32x2_t>(
+      make_float2(valf.x * weightf.x * rsqrt_square_sum, valf.y * weightf.y * rsqrt_square_sum));
+}
+
+template <typename T, int VEC_SIZE_IN_BYTE>
+__global__ void qknorm_across_heads_reg_kernel(
+    T* __restrict__ q,
+    T* __restrict__ k,
+    const T* __restrict__ q_weight,
+    const T* __restrict__ k_weight,
+    int vec_hidden_size,
+    float eps) {
+  constexpr int inner_loop = VEC_SIZE_IN_BYTE == 16 ? 4 : 8;
+
+  __shared__ float shared_memory[32];
+
+  using vec_t = typename VecTypeTrait<T, VEC_SIZE_IN_BYTE>::vec_t;
+  using packed_t = typename VecTypeTrait<T, VEC_SIZE_IN_BYTE>::packed_t;
+  vec_t v_data;
+  vec_t v_weight;
+  const int warp_id = threadIdx.x >> 5;
+  const int lane_id = threadIdx.x & 31;
+  const int warp_count = (blockDim.x + 31) >> 5;
+  const float inv_hidden_size = 1.0f / static_cast<float>(vec_hidden_size * (VEC_SIZE_IN_BYTE / sizeof(T)));
+  const bool is_q = blockIdx.y == 0;
+
+  const auto token_id = blockIdx.x;
+  float2 acc_square = make_float2(0.0f, 0.0f);
+  vec_t* data = reinterpret_cast<vec_t*>(is_q ? q : k) + token_id * vec_hidden_size;
+  const vec_t* weight = reinterpret_cast<const vec_t*>(is_q ? q_weight : k_weight);
+
+  if (threadIdx.x < vec_hidden_size) {
+    v_data = data[threadIdx.x];
+    v_weight = weight[threadIdx.x];
+    for (int i = 0; i < inner_loop; i++) {
+      float2 val = device::cast<fp32x2_t, packed_t>(v_data[i]);
+      acc_square.x += val.x * val.x;
+      acc_square.y += val.y * val.y;
+    }
+  }
+
+  auto cg_warp = cooperative_groups::tiled_partition<32>(cooperative_groups::this_thread_block());
+  float* buffer = shared_memory;
+  float warp_sum = cooperative_groups::reduce(cg_warp, acc_square.x + acc_square.y, cooperative_groups::plus<float>());
+  if (lane_id == 0) {
+    buffer[warp_id] = warp_sum;
+  }
+
+  __syncthreads();
+  if (threadIdx.x < 32) {
+    float cta_sum = cooperative_groups::reduce(
+        cg_warp, (threadIdx.x < warp_count) ? buffer[threadIdx.x] : 0.0f, cooperative_groups::plus<float>());
+    if (threadIdx.x == 0) {
+      buffer[0] = rsqrtf(eps + cta_sum * inv_hidden_size);
+    }
+  }
+  __syncthreads();
+
+  if (threadIdx.x < vec_hidden_size) {
+    float rsqrt_val = buffer[0];
+    for (int i = 0; i < inner_loop; i++) {
+      v_data[i] = rms(v_data[i], v_weight[i], rsqrt_val);
+    }
+    data[threadIdx.x] = v_data;
+  }
+}
+
+template <typename DType>
+struct QKNormAcrossHeadsKernel {
+  static void
+  run(const tvm::ffi::TensorView q,
+      const tvm::ffi::TensorView k,
+      const tvm::ffi::TensorView q_weight,
+      const tvm::ffi::TensorView k_weight,
+      float eps) {
+    using namespace host;
+    auto N = SymbolicSize{"num_tokens"};
+    auto D = SymbolicSize{"hidden_size"};
+    auto device = SymbolicDevice{};
+    device.set_options<kDLCUDA>();
+
+    TensorMatcher({N, D})  // q
+        .with_strides({D, 1})
+        .with_dtype<DType>()
+        .with_device(device)
+        .verify(q);
+    TensorMatcher({N, D})  // k
+        .with_strides({D, 1})
+        .with_dtype<DType>()
+        .with_device(device)
+        .verify(k);
+    TensorMatcher({D})  // q_weight
+        .with_dtype<DType>()
+        .with_device(device)
+        .verify(q_weight);
+    TensorMatcher({D})  // k_weight
+        .with_dtype<DType>()
+        .with_device(device)
+        .verify(k_weight);
+
+    int hidden_size = static_cast<int>(D.unwrap());
+    if (hidden_size <= (device::kMaxVecBytes == 32 ? 12288 : 8192)) {
+      int elements_in_vec = device::kMaxVecBytes / sizeof(DType);
+      int vec_hidden_size = hidden_size / elements_in_vec;
+      uint threads = (vec_hidden_size + 31) / 32 * 32;
+
+      // Runtime check
+      host::RuntimeCheck(
+          hidden_size % elements_in_vec == 0,
+          "hidden_size",
+          hidden_size,
+          " can not align to elements_in_vec ",
+          elements_in_vec);
+
+      auto kernel = qknorm_across_heads_reg_kernel<DType, device::kMaxVecBytes>;
+
+      LaunchKernel(dim3(static_cast<uint>(N.unwrap()), 2), threads, device.unwrap())
+          .enable_pdl(false)(
+              kernel,
+              reinterpret_cast<DType*>(q.data_ptr()),
+              reinterpret_cast<DType*>(k.data_ptr()),
+              reinterpret_cast<DType*>(q_weight.data_ptr()),
+              reinterpret_cast<DType*>(k_weight.data_ptr()),
+              vec_hidden_size,
+              eps);
+    } else {
+      host::RuntimeCheck(false, "Large hidden_sizes are not supported for now.");
+    }
+  }
+};
+
+}  // namespace

python/sglang/jit_kernel/csrc/elementwise/resolve_future_token_ids.cuh

@@ -0,0 +1,57 @@
+#include <sgl_kernel/tensor.h>  // For TensorMatcher, SymbolicSize, SymbolicDevice
+#include <sgl_kernel/utils.h>   // For RuntimeCheck, div_ceil
+
+#include <sgl_kernel/utils.cuh>  // For LaunchKernel
+
+#include <dlpack/dlpack.h>
+#include <tvm/ffi/container/tensor.h>
+
+#include <cstddef>
+#include <cstdint>
+
+namespace {
+
+template <typename T>
+__global__ void resolve_future_token_ids_kernel(T* __restrict__ input_ids, const T* __restrict__ future_map, size_t n) {
+  size_t idx = blockIdx.x * blockDim.x + threadIdx.x;
+  if (idx < n) {
+    T val = input_ids[idx];
+    if (val < 0) {
+      T key = -val;
+      if (key < 0) key = 0;  // clamp for overflow
+      input_ids[idx] = future_map[key];
+    }
+  }
+}
+
+constexpr size_t kBlockSize = 256;
+
+template <typename T>
+struct ResolveFutureTokenIds {
+  static void run(tvm::ffi::TensorView input_ids, tvm::ffi::TensorView future_map) {
+    using namespace host;
+
+    SymbolicSize N = {"num_tokens"};
+    SymbolicSize M = {"map_size"};
+    SymbolicDevice device_;
+    device_.set_options<kDLCUDA, kDLROCM>();
+
+    TensorMatcher({N}).with_dtype<T>().with_device(device_).verify(input_ids);
+
+    TensorMatcher({M}).with_dtype<T>().with_device(device_).verify(future_map);
+
+    const size_t num_tokens = N.unwrap();
+    if (num_tokens == 0) return;
+
+    const size_t grid_size = div_ceil(num_tokens, kBlockSize);
+    const DLDevice device = device_.unwrap();
+
+    LaunchKernel(grid_size, kBlockSize, device)(
+        resolve_future_token_ids_kernel<T>,
+        static_cast<T*>(input_ids.data_ptr()),
+        static_cast<const T*>(future_map.data_ptr()),
+        num_tokens);
+  }
+};
+
+}  // namespace

python/sglang/jit_kernel/csrc/elementwise/rmsnorm.cuh

@@ -0,0 +1,371 @@
+#include <sgl_kernel/tensor.h>
+#include <sgl_kernel/utils.h>
+
+#include <sgl_kernel/runtime.cuh>
+#include <sgl_kernel/tile.cuh>
+#include <sgl_kernel/utils.cuh>
+#include <sgl_kernel/vec.cuh>
+
+#include <sgl_kernel/impl/norm.cuh>
+
+#include <tvm/ffi/container/tensor.h>
+
+namespace {
+
+struct RMSNormParams {
+  const void* input;
+  const void* __restrict__ weight;
+  void* output;
+  int64_t input_stride;
+  int64_t output_stride;
+  uint32_t num_tokens;
+  float eps;
+};
+
+template <int64_t kDim, bool kUsePDL, typename Float>
+__global__ void rmsnorm_cta(const RMSNormParams __grid_constant__ params) {
+  using namespace device;
+  using Storage = norm::StorageType<Float, kDim>;
+
+  constexpr auto kNumThreads = host::norm::get_cta_threads<Float, kDim>();
+  constexpr auto kNumWarps = kNumThreads / kWarpThreads;
+
+  const auto& [input, weight_ptr, output, input_stride, output_stride, num_tokens, eps] = params;
+  const auto gmem = tile::Memory<Storage>::cta(kNumThreads);
+  __shared__ float smem[norm::kSmemBufferSize];
+
+  PDLWaitPrimary<kUsePDL>();  // wait for primary kernel
+
+  for (uint32_t i = blockIdx.x; i < num_tokens; i += gridDim.x) {
+    const auto input_ptr = pointer::offset<Float>(input, i * input_stride);
+    const auto output_ptr = pointer::offset<Float>(output, i * output_stride);
+    const auto input_vec = gmem.load(input_ptr);
+    const auto weight_vec = gmem.load(weight_ptr);
+    const auto output_vec = norm::apply_norm_cta<kDim>(input_vec, weight_vec, eps, smem, kNumWarps);
+    gmem.store(output_ptr, output_vec);
+  }
+
+  PDLTriggerSecondary<kUsePDL>();  // launch secondary kernel
+}
+
+// Pre-Blackwell: 16B vector, each thread loads/stores twice
+template <int64_t kDim, bool kUsePDL, typename Float>
+__global__ __launch_bounds__(kDim / 16) void rmsnorm_cta_double(const RMSNormParams __grid_constant__ params) {
+  using namespace device;
+  using Float2 = packed_t<Float>;
+  using Storage = AlignedVector<Float2, 4>;
+
+  constexpr auto kNumThreads = kDim / 16;
+  constexpr auto kNumWarps = kNumThreads / kWarpThreads;
+
+  const auto& [input, weight_ptr, output, input_stride, output_stride, num_tokens, eps] = params;
+  const auto gmem = tile::Memory<Storage>::cta(kNumThreads);
+  __shared__ float smem[32];
+
+  PDLWaitPrimary<kUsePDL>();
+
+  const auto input_ptr = pointer::offset<Float>(input, blockIdx.x * input_stride);
+  const auto output_ptr = pointer::offset<Float>(output, blockIdx.x * output_stride);
+
+  const auto input_first = gmem.load(input_ptr, 0);
+  const auto input_second = gmem.load(input_ptr, 1);
+  const auto weight_first = gmem.load(weight_ptr, 0);
+  const auto weight_second = gmem.load(weight_ptr, 1);
+
+  float sum_of_squares = 0.0f;
+#pragma unroll
+  for (auto j = 0u; j < 4u; ++j) {
+    const auto [x, y] = cast<fp32x2_t>(input_first[j]);
+    sum_of_squares += x * x + y * y;
+  }
+#pragma unroll
+  for (auto j = 0u; j < 4u; ++j) {
+    const auto [x, y] = cast<fp32x2_t>(input_second[j]);
+    sum_of_squares += x * x + y * y;
+  }
+
+  sum_of_squares = warp::reduce_sum(sum_of_squares);
+  float norm_factor;
+  if constexpr (kNumWarps == 1) {
+    norm_factor = math::rsqrt(sum_of_squares / kDim + eps);
+  } else {
+    const auto warp_id = threadIdx.x / kWarpThreads;
+    smem[warp_id] = sum_of_squares;
+    __syncthreads();
+    if (warp_id == 0) {
+      const auto tx = threadIdx.x;
+      const auto local_sum = tx < kNumWarps ? smem[tx] : 0.0f;
+      sum_of_squares = warp::reduce_sum(local_sum);
+      smem[tx] = math::rsqrt(sum_of_squares / kDim + eps);
+    }
+    __syncthreads();
+    norm_factor = smem[warp_id];
+  }
+
+  Storage output_first, output_second;
+#pragma unroll
+  for (auto j = 0u; j < 4u; ++j) {
+    const auto [ix, iy] = cast<fp32x2_t>(input_first[j]);
+    const auto [wx, wy] = cast<fp32x2_t>(weight_first[j]);
+    output_first[j] = cast<Float2>(fp32x2_t{ix * norm_factor * wx, iy * norm_factor * wy});
+  }
+#pragma unroll
+  for (auto j = 0u; j < 4u; ++j) {
+    const auto [ix, iy] = cast<fp32x2_t>(input_second[j]);
+    const auto [wx, wy] = cast<fp32x2_t>(weight_second[j]);
+    output_second[j] = cast<Float2>(fp32x2_t{ix * norm_factor * wx, iy * norm_factor * wy});
+  }
+
+  gmem.store(output_ptr, output_first, 0);
+  gmem.store(output_ptr, output_second, 1);
+
+  PDLTriggerSecondary<kUsePDL>();
+}
+
+// Blackwell: 32B vector, each thread loads/stores once
+template <int64_t kDim, bool kUsePDL, typename Float>
+__global__ __launch_bounds__(kDim / 16) void rmsnorm_cta_wide(const RMSNormParams __grid_constant__ params) {
+  using namespace device;
+  using Float2 = packed_t<Float>;
+  using Storage = AlignedVector<Float2, 8>;
+
+  constexpr auto kNumThreads = kDim / 16;
+  constexpr auto kNumWarps = kNumThreads / kWarpThreads;
+
+  const auto& [input, weight_ptr, output, input_stride, output_stride, num_tokens, eps] = params;
+  const auto gmem = tile::Memory<Storage>::cta(kNumThreads);
+  __shared__ float smem[32];
+
+  PDLWaitPrimary<kUsePDL>();
+
+  const auto input_ptr = pointer::offset<Float>(input, blockIdx.x * input_stride);
+  const auto output_ptr = pointer::offset<Float>(output, blockIdx.x * output_stride);
+
+  const auto input_vec = gmem.load(input_ptr);
+  const auto weight_vec = gmem.load(weight_ptr);
+
+  float sum_of_squares = 0.0f;
+#pragma unroll
+  for (auto j = 0u; j < 8u; ++j) {
+    const auto [x, y] = cast<fp32x2_t>(input_vec[j]);
+    sum_of_squares += x * x + y * y;
+  }
+
+  sum_of_squares = warp::reduce_sum(sum_of_squares);
+  float norm_factor;
+  if constexpr (kNumWarps == 1) {
+    norm_factor = math::rsqrt(sum_of_squares / kDim + eps);
+  } else {
+    const auto warp_id = threadIdx.x / kWarpThreads;
+    smem[warp_id] = sum_of_squares;
+    __syncthreads();
+    if (warp_id == 0) {
+      const auto tx = threadIdx.x;
+      const auto local_sum = tx < kNumWarps ? smem[tx] : 0.0f;
+      sum_of_squares = warp::reduce_sum(local_sum);
+      smem[tx] = math::rsqrt(sum_of_squares / kDim + eps);
+    }
+    __syncthreads();
+    norm_factor = smem[warp_id];
+  }
+
+  Storage output_vec;
+#pragma unroll
+  for (auto j = 0u; j < 8u; ++j) {
+    const auto [ix, iy] = cast<fp32x2_t>(input_vec[j]);
+    const auto [wx, wy] = cast<fp32x2_t>(weight_vec[j]);
+    output_vec[j] = cast<Float2>(fp32x2_t{ix * norm_factor * wx, iy * norm_factor * wy});
+  }
+
+  gmem.store(output_ptr, output_vec);
+
+  PDLTriggerSecondary<kUsePDL>();
+}
+
+template <int64_t kDim, bool kUsePDL, typename Float>
+__global__ void rmsnorm_warp(const RMSNormParams __grid_constant__ params) {
+  using namespace device;
+  using Storage = norm::StorageType<Float, kDim>;
+
+  const auto& [input, weight_ptr, output, input_stride, output_stride, num_tokens, eps] = params;
+  const auto gmem = tile::Memory<Storage>::warp();
+
+  PDLWaitPrimary<kUsePDL>();  // wait for primary kernel
+
+  for (uint32_t i = blockIdx.x; i < num_tokens; i += gridDim.x) {
+    const auto input_ptr = pointer::offset<Float>(input, i * input_stride);
+    const auto output_ptr = pointer::offset<Float>(output, i * output_stride);
+    const auto input_vec = gmem.load(input_ptr);
+    const auto weight_vec = gmem.load(weight_ptr);
+    const auto output_vec = norm::apply_norm_warp<kDim>(input_vec, weight_vec, eps);
+    gmem.store(output_ptr, output_vec);
+  }
+
+  PDLTriggerSecondary<kUsePDL>();  // launch secondary kernel
+}
+
+template <int64_t kDim, bool kUsePDL, typename DType>
+struct RMSNormWarpKernel {
+  static_assert(host::norm::is_config_supported<DType, kDim>(), "Unsupported norm configuration");
+  static_assert(kDim <= 256, "Use RMSNormKernel for hidden sizes > 256");
+  static constexpr auto kernel = rmsnorm_warp<kDim, kUsePDL, DType>;
+
+  static void
+  run(const tvm::ffi::TensorView input,
+      const tvm::ffi::TensorView weight,
+      const tvm::ffi::TensorView output,
+      float eps) {
+    using namespace host;
+    auto N = SymbolicSize{"num_tokens"};
+    auto D = SymbolicSize{"hidden_size"};
+    auto SI = SymbolicSize{"input_stride"};
+    auto SO = SymbolicSize{"output_stride"};
+    auto device = SymbolicDevice{};
+    D.set_value(kDim);
+    device.set_options<kDLCUDA>();
+
+    TensorMatcher({N, D})  // input
+        .with_strides({SI, 1})
+        .with_dtype<DType>()
+        .with_device(device)
+        .verify(input);
+    TensorMatcher({D})  // weight
+        .with_dtype<DType>()
+        .with_device(device)
+        .verify(weight);
+    TensorMatcher({N, D})  // output
+        .with_strides({SO, 1})
+        .with_dtype<DType>()
+        .with_device(device)
+        .verify(output);
+
+    const auto num_tokens = static_cast<uint32_t>(N.unwrap());
+    const auto params = RMSNormParams{
+        .input = input.data_ptr(),
+        .weight = weight.data_ptr(),
+        .output = output.data_ptr(),
+        .input_stride = SI.unwrap(),
+        .output_stride = SO.unwrap(),
+        .num_tokens = num_tokens,
+        .eps = eps,
+    };
+
+    static constexpr uint32_t kNumThreads = device::kWarpThreads;
+    static const uint32_t max_occupancy = runtime::get_blocks_per_sm(kernel, kNumThreads);
+    static const uint32_t kNumSM = runtime::get_sm_count(device.unwrap().device_id);
+    const auto num_blocks = std::min<uint32_t>(num_tokens, max_occupancy * kNumSM);
+    LaunchKernel(num_blocks, kNumThreads, device.unwrap())  //
+        .enable_pdl(kUsePDL)(kernel, params);
+  }
+};
+
+template <int64_t kDim, bool kUsePDL, typename DType>
+struct RMSNormKernel {
+  static_assert(host::norm::should_use_cta<DType, kDim>(), "Hidden size invalid for RMSNorm");
+  static constexpr auto kernel = rmsnorm_cta<kDim, kUsePDL, DType>;
+
+  static void
+  run(const tvm::ffi::TensorView input,
+      const tvm::ffi::TensorView weight,
+      const tvm::ffi::TensorView output,
+      float eps) {
+    using namespace host;
+    auto N = SymbolicSize{"num_tokens"};
+    auto D = SymbolicSize{"hidden_size"};
+    auto SI = SymbolicSize{"input_stride"};
+    auto SO = SymbolicSize{"output_stride"};
+    auto device = SymbolicDevice{};
+    D.set_value(kDim);
+    device.set_options<kDLCUDA>();
+
+    TensorMatcher({N, D})  // input
+        .with_strides({SI, 1})
+        .with_dtype<DType>()
+        .with_device(device)
+        .verify(input);
+    TensorMatcher({D})  // weight
+        .with_dtype<DType>()
+        .with_device(device)
+        .verify(weight);
+    TensorMatcher({N, D})  // output
+        .with_strides({SO, 1})
+        .with_dtype<DType>()
+        .with_device(device)
+        .verify(output);
+
+    const auto num_tokens = static_cast<uint32_t>(N.unwrap());
+    const auto params = RMSNormParams{
+        .input = input.data_ptr(),
+        .weight = weight.data_ptr(),
+        .output = output.data_ptr(),
+        .input_stride = SI.unwrap(),
+        .output_stride = SO.unwrap(),
+        .num_tokens = num_tokens,
+        .eps = eps,
+    };
+
+    static constexpr auto kNumThreads = norm::get_cta_threads<DType, kDim>();
+    static const uint32_t max_occupancy = runtime::get_blocks_per_sm(kernel, kNumThreads);
+    static const uint32_t kNumSM = runtime::get_sm_count(device.unwrap().device_id);
+    const auto num_blocks = std::min<uint32_t>(num_tokens, max_occupancy * kNumSM);
+    LaunchKernel(num_blocks, kNumThreads, device.unwrap())  //
+        .enable_pdl(kUsePDL)(kernel, params);
+  }
+};
+
+template <int64_t kDim, bool kUsePDL, typename DType>
+struct RMSNormHalfKernel {
+  static_assert(kDim % 512 == 0 && sizeof(DType) == 2);
+#if SGL_ARCH_BLACKWELL_OR_GREATER
+  static constexpr auto kernel = rmsnorm_cta_wide<kDim, kUsePDL, DType>;
+#else
+  static constexpr auto kernel = rmsnorm_cta_double<kDim, kUsePDL, DType>;
+#endif
+  static constexpr auto kBlockSize = static_cast<uint32_t>(kDim / 16);
+
+  static void
+  run(const tvm::ffi::TensorView input,
+      const tvm::ffi::TensorView weight,
+      const tvm::ffi::TensorView output,
+      float eps) {
+    using namespace host;
+    auto N = SymbolicSize{"num_tokens"};
+    auto D = SymbolicSize{"hidden_size"};
+    auto SI = SymbolicSize{"input_stride"};
+    auto SO = SymbolicSize{"output_stride"};
+    auto device = SymbolicDevice{};
+    D.set_value(kDim);
+    device.set_options<kDLCUDA>();
+
+    TensorMatcher({N, D})  // input
+        .with_strides({SI, 1})
+        .with_dtype<DType>()
+        .with_device(device)
+        .verify(input);
+    TensorMatcher({D})  // weight
+        .with_dtype<DType>()
+        .with_device(device)
+        .verify(weight);
+    TensorMatcher({N, D})  // output
+        .with_strides({SO, 1})
+        .with_dtype<DType>()
+        .with_device(device)
+        .verify(output);
+
+    const auto num_tokens = static_cast<uint32_t>(N.unwrap());
+    const auto params = RMSNormParams{
+        .input = input.data_ptr(),
+        .weight = weight.data_ptr(),
+        .output = output.data_ptr(),
+        .input_stride = SI.unwrap(),
+        .output_stride = SO.unwrap(),
+        .num_tokens = num_tokens,
+        .eps = eps,
+    };
+
+    LaunchKernel(num_tokens, kBlockSize, device.unwrap())  //
+        .enable_pdl(kUsePDL)(kernel, params);
+  }
+};
+
+}  // namespace

python/sglang/jit_kernel/csrc/elementwise/rmsnorm_hf.cuh

@@ -0,0 +1,253 @@
+/**
+ * RMSNorm with HuggingFace semantics:
+ *   out[i] = weight[i] * cast_dtype( rsqrt(mean_j(x[j]^2) + eps) * x[i] )
+ *
+ * vs. standard rmsnorm: the normalized x is rounded to the activation dtype
+ * BEFORE the weight multiply (not after). The multiply itself is done in fp32
+ * either way; the load-bearing step is the intermediate rounding. Required
+ * for HF `LlamaRMSNorm` parity under weight-only quantization.
+ *
+ * Two launch configs:
+ *   - Warp kernel: 32 threads/row for small hidden sizes (q/k norms).
+ *   - CTA kernel:  512-thread scalar-strided with register cache (token norms).
+ */
+
+#include <sgl_kernel/tensor.h>  // For TensorMatcher, SymbolicSize, SymbolicDevice
+#include <sgl_kernel/utils.h>   // For RuntimeCheck
+
+#include <sgl_kernel/math.cuh>     // For device::math::rsqrt
+#include <sgl_kernel/runtime.cuh>  // For runtime::get_blocks_per_sm, get_sm_count
+#include <sgl_kernel/utils.cuh>    // For LaunchKernel, SGL_DEVICE, type aliases, PDL, cast
+#include <sgl_kernel/warp.cuh>     // For warp::reduce_sum
+
+#include <tvm/ffi/container/tensor.h>
+
+namespace {
+
+struct RMSNormHFParams {
+  const void* input;
+  const void* __restrict__ weight;
+  void* output;
+  int64_t input_stride;
+  int64_t output_stride;
+  uint32_t num_tokens;
+  float eps;
+};
+
+// ---------------------------------------------------------------------------
+// Warp kernel: one warp per row, for small hidden sizes (e.g. q/k norms at
+// head_dim ∈ {32, 64, 96, 128, 256}). No shared memory, no block reduce —
+// warp reduce is sufficient. Grid-strided over rows.
+// ---------------------------------------------------------------------------
+template <int64_t kDim, bool kUsePDL, typename Float>
+__global__ __launch_bounds__(32) void rmsnorm_hf_warp_kernel(const RMSNormHFParams __grid_constant__ params) {
+  using namespace device;
+  constexpr int kElemsPerThread = kDim / kWarpThreads;
+
+  const auto& [input, weight_ptr, output, input_stride, output_stride, num_tokens, eps] = params;
+  const auto wr = static_cast<const Float*>(weight_ptr);
+
+  PDLWaitPrimary<kUsePDL>();
+
+  for (uint32_t row = blockIdx.x; row < num_tokens; row += gridDim.x) {
+    const auto xr = static_cast<const Float*>(pointer::offset<Float>(input, row * input_stride));
+    const auto yr = static_cast<Float*>(pointer::offset<Float>(output, row * output_stride));
+
+    float xi_cache[kElemsPerThread];
+    float lsq = 0.f;
+#pragma unroll
+    for (int k = 0; k < kElemsPerThread; ++k) {
+      const int i = threadIdx.x + k * kWarpThreads;
+      xi_cache[k] = static_cast<float>(xr[i]);
+      lsq += xi_cache[k] * xi_cache[k];
+    }
+    lsq = warp::reduce_sum(lsq);
+    const float rstd = math::rsqrt(lsq / kDim + eps);
+
+    // HF semantics — round (x*rstd) to dtype, THEN multiply by weight.
+#pragma unroll
+    for (int k = 0; k < kElemsPerThread; ++k) {
+      const int i = threadIdx.x + k * kWarpThreads;
+      const Float xn = cast<Float>(xi_cache[k] * rstd);
+      yr[i] = cast<Float>(static_cast<float>(xn) * static_cast<float>(wr[i]));
+    }
+  }
+
+  PDLTriggerSecondary<kUsePDL>();
+}
+
+// ---------------------------------------------------------------------------
+// Kernel: 512-thread scalar-strided RMSNorm with HF semantics + register cache.
+//
+// Pass 1: each thread loads its strided elements, caches them in registers,
+//         and accumulates the fp32 sum-of-squares. Warp + block reduction
+//         yields `rstd = rsqrt(mean(x^2) + eps)`.
+// Pass 2: reuse cached fp32 values — no second global read of `x`. Per-elem:
+//             xn = cast_to_dtype(x_fp32 * rstd)   <- HF's cast-before-mul
+//             y  = cast_to_dtype(float(xn) * float(w))
+// ---------------------------------------------------------------------------
+template <int64_t kDim, bool kUsePDL, typename Float>
+__global__ __launch_bounds__(512) void rmsnorm_hf_scalar_kernel(const RMSNormHFParams __grid_constant__ params) {
+  using namespace device;
+  constexpr int kNumThreads = 512;
+  constexpr int kNumWarps = kNumThreads / kWarpThreads;
+  // For kDim=4096: kElemsPerThread = 8 (32 bytes of fp32 cache per thread).
+  constexpr int kElemsPerThread = (kDim + kNumThreads - 1) / kNumThreads;
+
+  const auto& [input, weight_ptr, output, input_stride, output_stride, num_tokens, eps] = params;
+  const auto xr = static_cast<const Float*>(pointer::offset<Float>(input, blockIdx.x * input_stride));
+  const auto yr = static_cast<Float*>(pointer::offset<Float>(output, blockIdx.x * output_stride));
+  const auto wr = static_cast<const Float*>(weight_ptr);
+
+  PDLWaitPrimary<kUsePDL>();
+
+  // Pass 1: load, square, accumulate; cache fp32 values in registers.
+  float xi_cache[kElemsPerThread];
+  float lsq = 0.f;
+#pragma unroll
+  for (int k = 0; k < kElemsPerThread; ++k) {
+    const int i = threadIdx.x + k * kNumThreads;
+    xi_cache[k] = static_cast<float>(xr[i]);
+    lsq += xi_cache[k] * xi_cache[k];
+  }
+
+  // Warp reduce.
+  lsq = warp::reduce_sum(lsq);
+
+  // Block reduce via shared memory (32 warps * 1 fp32 each).
+  __shared__ float smem[32];
+  const int warp_id = threadIdx.x / kWarpThreads;
+  const int lane_id = threadIdx.x & (kWarpThreads - 1);
+  if (lane_id == 0) smem[warp_id] = lsq;
+  __syncthreads();
+
+  __shared__ float rstd_s;
+  if (threadIdx.x < kWarpThreads) {
+    float v = (threadIdx.x < kNumWarps) ? smem[threadIdx.x] : 0.f;
+    v = warp::reduce_sum(v);
+    if (threadIdx.x == 0) rstd_s = math::rsqrt(v / kDim + eps);
+  }
+  __syncthreads();
+  const float rstd = rstd_s;
+
+  // Pass 2: HF semantics — round (x*rstd) to dtype, THEN multiply by weight.
+#pragma unroll
+  for (int k = 0; k < kElemsPerThread; ++k) {
+    const int i = threadIdx.x + k * kNumThreads;
+    const Float xn = cast<Float>(xi_cache[k] * rstd);
+    yr[i] = cast<Float>(static_cast<float>(xn) * static_cast<float>(wr[i]));
+  }
+
+  PDLTriggerSecondary<kUsePDL>();
+}
+
+// ---------------------------------------------------------------------------
+// Warp launcher: occupancy-sized grid, 32 threads/block, one warp per row.
+// Targets small hidden sizes (q/k RMSNorms). kDim must be a multiple of 32
+// in [32, 512).
+// ---------------------------------------------------------------------------
+template <int64_t kDim, bool kUsePDL, typename DType>
+struct HFRMSNormWarpKernel {
+  static_assert(sizeof(DType) == 2, "rmsnorm_hf: DType must be fp16_t or bf16_t");
+  static_assert(
+      kDim >= 32 && kDim < 512 && kDim % 32 == 0, "rmsnorm_hf_warp: kDim must be a multiple of 32, in [32, 512)");
+  static constexpr auto kernel = rmsnorm_hf_warp_kernel<kDim, kUsePDL, DType>;
+  static constexpr uint32_t kBlockSize = device::kWarpThreads;
+
+  static void
+  run(const tvm::ffi::TensorView input,
+      const tvm::ffi::TensorView weight,
+      const tvm::ffi::TensorView output,
+      float eps) {
+    using namespace host;
+    auto N = SymbolicSize{"num_tokens"};
+    auto D = SymbolicSize{"hidden_size"};
+    auto SI = SymbolicSize{"input_stride"};
+    auto SO = SymbolicSize{"output_stride"};
+    auto device_ = SymbolicDevice{};
+    D.set_value(kDim);
+    device_.set_options<kDLCUDA>();
+
+    TensorMatcher({N, D}).with_strides({SI, 1}).with_dtype<DType>().with_device(device_).verify(input);
+    TensorMatcher({D}).with_dtype<DType>().with_device(device_).verify(weight);
+    TensorMatcher({N, D}).with_strides({SO, 1}).with_dtype<DType>().with_device(device_).verify(output);
+
+    const auto num_tokens = static_cast<uint32_t>(N.unwrap());
+    RuntimeCheck(num_tokens > 0, "rmsnorm_hf: num_tokens must be > 0");
+
+    const auto params = RMSNormHFParams{
+        .input = input.data_ptr(),
+        .weight = weight.data_ptr(),
+        .output = output.data_ptr(),
+        .input_stride = SI.unwrap(),
+        .output_stride = SO.unwrap(),
+        .num_tokens = num_tokens,
+        .eps = eps,
+    };
+
+    static const uint32_t max_occupancy = runtime::get_blocks_per_sm(kernel, kBlockSize);
+    static const uint32_t kNumSM = runtime::get_sm_count(device_.unwrap().device_id);
+    const auto num_blocks = std::min<uint32_t>(num_tokens, max_occupancy * kNumSM);
+    LaunchKernel(num_blocks, kBlockSize, device_.unwrap())  //
+        .enable_pdl(kUsePDL)(kernel, params);
+  }
+};
+
+// ---------------------------------------------------------------------------
+// CTA launcher: validates tensors, launches one block per row.
+// ---------------------------------------------------------------------------
+template <int64_t kDim, bool kUsePDL, typename DType>
+struct HFRMSNormKernel {
+  static_assert(sizeof(DType) == 2, "rmsnorm_hf: DType must be fp16_t or bf16_t");
+  static_assert(kDim >= 512 && kDim % 512 == 0, "rmsnorm_hf: kDim must be a multiple of 512");
+  static constexpr auto kernel = rmsnorm_hf_scalar_kernel<kDim, kUsePDL, DType>;
+  static constexpr uint32_t kBlockSize = 512;
+
+  static void
+  run(const tvm::ffi::TensorView input,
+      const tvm::ffi::TensorView weight,
+      const tvm::ffi::TensorView output,
+      float eps) {
+    using namespace host;
+    auto N = SymbolicSize{"num_tokens"};
+    auto D = SymbolicSize{"hidden_size"};
+    auto SI = SymbolicSize{"input_stride"};
+    auto SO = SymbolicSize{"output_stride"};
+    auto device_ = SymbolicDevice{};
+    D.set_value(kDim);
+    device_.set_options<kDLCUDA>();
+
+    TensorMatcher({N, D})  // input
+        .with_strides({SI, 1})
+        .with_dtype<DType>()
+        .with_device(device_)
+        .verify(input);
+    TensorMatcher({D})  // weight
+        .with_dtype<DType>()
+        .with_device(device_)
+        .verify(weight);
+    TensorMatcher({N, D})  // output
+        .with_strides({SO, 1})
+        .with_dtype<DType>()
+        .with_device(device_)
+        .verify(output);
+
+    const auto num_tokens = static_cast<uint32_t>(N.unwrap());
+    RuntimeCheck(num_tokens > 0, "rmsnorm_hf: num_tokens must be > 0");
+
+    const auto params = RMSNormHFParams{
+        .input = input.data_ptr(),
+        .weight = weight.data_ptr(),
+        .output = output.data_ptr(),
+        .input_stride = SI.unwrap(),
+        .output_stride = SO.unwrap(),
+        .num_tokens = num_tokens,
+        .eps = eps,
+    };
+
+    LaunchKernel(num_tokens, kBlockSize, device_.unwrap())  //
+        .enable_pdl(kUsePDL)(kernel, params);
+  }
+};
+
+}  // namespace

python/sglang/jit_kernel/csrc/elementwise/rope.cuh

@@ -0,0 +1,470 @@
+#include <sgl_kernel/tensor.h>
+#include <sgl_kernel/utils.h>
+
+#include <sgl_kernel/runtime.cuh>
+#include <sgl_kernel/tile.cuh>
+#include <sgl_kernel/type.cuh>
+#include <sgl_kernel/utils.cuh>
+#include <sgl_kernel/vec.cuh>
+
+#include <dlpack/dlpack.h>
+
+#include <numeric>
+
+namespace {
+
+struct FusedRopeParams {
+  void* __restrict__ q_ptr;
+  void* __restrict__ k_ptr;  // NOTE: this k is pre-offset in host code to reduce computation in kernel
+  const void* __restrict__ cos_sin_cache_ptr;
+  const void* __restrict__ positions;
+  int64_t q_stride_bytes;
+  int64_t k_stride_bytes;
+  int64_t head_stride_bytes;
+  uint32_t num_qo_heads;
+  uint32_t num_kv_heads;
+  uint32_t num_tokens;
+};
+
+struct FusedRopeStoreParams {
+  FusedRopeParams base_params;
+  void* v_ptr;
+  void* __restrict__ k_cache;
+  void* __restrict__ v_cache;
+  const void* __restrict__ out_loc;
+  int64_t v_stride_bytes;
+  int64_t cache_stride_bytes;
+};
+
+constexpr uint32_t kBlockSize = 128;
+
+[[maybe_unused]]
+constexpr auto next_pow2(uint32_t target, uint32_t factor = 1) {
+  uint32_t power = 1;
+  while (power * factor < target)
+    power *= 2;
+  return power;
+}
+
+template <bool kIsNeox, int64_t kRopeDim, bool kUsePDL, typename DType, typename IdType, uint32_t kWorkThreads>
+__global__ void fused_rope_kernel(const __grid_constant__ FusedRopeParams params) {
+  using namespace device;
+
+  constexpr int64_t kCosSinStrideBytes = kRopeDim * sizeof(float);
+  constexpr int64_t kVecSize = next_pow2(kRopeDim, (2 * kWorkThreads * (1 + kIsNeox)));
+  using DType2 = packed_t<DType>;
+  using InputStorage = AlignedVector<DType2, kVecSize>;
+  constexpr int64_t kDimPerThread = kVecSize * 2 * (1 + kIsNeox);
+  constexpr uint32_t kLaneCount = kRopeDim / kDimPerThread;
+  static_assert(kRopeDim % kDimPerThread == 0 && kLaneCount <= kWorkThreads);
+
+  const auto &[
+    q, k, cos_sin_cache_ptr, positions, // pointers
+    q_stride_bytes, k_stride_bytes, head_stride_bytes,  // strides
+    num_qo_heads, num_kv_heads, num_tokens // dimensions
+  ] = params;
+
+  const auto num_blks = gridDim.x;
+  constexpr auto kWorkersPerBlock = kBlockSize / kWorkThreads;
+  const auto num_workers = num_blks * kWorkersPerBlock;
+  const auto num_q_and_k_heads = num_qo_heads + num_kv_heads;
+  const auto num_works = num_q_and_k_heads * num_tokens;
+  const auto start_worker_id = (blockIdx.x * kBlockSize + threadIdx.x) / kWorkThreads;
+  const auto cos_cache_ptr = cos_sin_cache_ptr;
+  const auto sin_cache_ptr = pointer::offset(cos_sin_cache_ptr, kCosSinStrideBytes / 2);
+
+  uint32_t lane_id = threadIdx.x % kWorkThreads;
+  if constexpr (kLaneCount < kWorkThreads) {
+    if (lane_id >= kLaneCount) return;
+  }
+
+  PDLWaitPrimary<kUsePDL>();
+
+  for (auto idx = start_worker_id; idx < num_works; idx += num_workers) {
+    const int64_t token_id = idx / num_q_and_k_heads;
+    const int64_t head_id = idx % num_q_and_k_heads;
+    const auto pos = static_cast<const IdType*>(positions)[token_id];
+    const auto load_q = head_id < num_qo_heads;
+    const auto input_ = load_q ? pointer::offset(q, token_id * q_stride_bytes)  //
+                               : pointer::offset(k, token_id * k_stride_bytes);
+    const auto input = pointer::offset(input_, head_id * head_stride_bytes);
+    const auto cos_ptr = pointer::offset(cos_cache_ptr, pos * kCosSinStrideBytes);
+    const auto sin_ptr = pointer::offset(sin_cache_ptr, pos * kCosSinStrideBytes);
+    if constexpr (kIsNeox) {
+      using CacheStorage = AlignedVector<fp32x2_t, kVecSize>;
+      const auto input_x = input;
+      const auto input_y = pointer::offset(input, (kRopeDim / 2) * sizeof(DType));
+      auto input_vec_x = load_as<InputStorage>(input_x, lane_id);
+      auto input_vec_y = load_as<InputStorage>(input_y, lane_id);
+      const auto cos_pair = load_as<CacheStorage>(cos_ptr, lane_id);
+      const auto sin_pair = load_as<CacheStorage>(sin_ptr, lane_id);
+#pragma unroll
+      for (int64_t j = 0; j < kVecSize; ++j) {
+        const auto [x0, x1] = cast<fp32x2_t>(input_vec_x[j]);
+        const auto [y0, y1] = cast<fp32x2_t>(input_vec_y[j]);
+        const auto [cos_0, cos_1] = cos_pair[j];
+        const auto [sin_0, sin_1] = sin_pair[j];
+        const auto out_x0 = x0 * cos_0 - y0 * sin_0;
+        const auto out_y0 = x0 * sin_0 + y0 * cos_0;
+        const auto out_x1 = x1 * cos_1 - y1 * sin_1;
+        const auto out_y1 = x1 * sin_1 + y1 * cos_1;
+        input_vec_x[j] = cast<DType2, fp32x2_t>({out_x0, out_x1});
+        input_vec_y[j] = cast<DType2, fp32x2_t>({out_y0, out_y1});
+      }
+      store_as<InputStorage>(input_x, input_vec_x, lane_id);
+      store_as<InputStorage>(input_y, input_vec_y, lane_id);
+    } else {
+      using CacheStorage = AlignedVector<float, kVecSize>;
+      auto input_vec = load_as<InputStorage>(input, lane_id);
+      const auto cos_vec = load_as<CacheStorage>(cos_ptr, lane_id);
+      const auto sin_vec = load_as<CacheStorage>(sin_ptr, lane_id);
+#pragma unroll
+      for (int64_t j = 0; j < kVecSize; ++j) {
+        const auto [x, y] = cast<fp32x2_t>(input_vec[j]);
+        const auto cos = cos_vec[j];
+        const auto sin = sin_vec[j];
+        const auto out_x = x * cos - y * sin;
+        const auto out_y = x * sin + y * cos;
+        input_vec[j] = cast<DType2, fp32x2_t>({out_x, out_y});
+      }
+      store_as<InputStorage>(input, input_vec, lane_id);
+    }
+  }
+
+  PDLTriggerSecondary<kUsePDL>();
+}
+
+template <bool kIsNeox, int64_t kRopeDim, bool kUsePDL, typename DType, typename IdType, uint32_t kWorkThreads>
+__global__ void fused_rope_store_kernel(const __grid_constant__ FusedRopeStoreParams params) {
+  using namespace device;
+
+  constexpr int64_t kCosSinStrideBytes = kRopeDim * sizeof(float);
+  constexpr int64_t kVecSize = kRopeDim / (2 * kWorkThreads * (1 + kIsNeox));
+  using DType2 = packed_t<DType>;
+  using InputStorage = AlignedVector<DType2, kVecSize>;
+  constexpr int64_t kDimPerThread = kVecSize * 2 * (1 + kIsNeox);
+  static_assert(kRopeDim == kDimPerThread * kWorkThreads);
+
+  const auto& [base_params, v_ptr, k_cache, v_cache, out_loc, v_stride_bytes, cache_stride_bytes] = params;
+  const auto &[
+    q, k, cos_sin_cache_ptr, positions, // pointers
+    q_stride_bytes, k_stride_bytes, head_stride_bytes,  // strides
+    num_qo_heads, num_kv_heads, num_tokens // dimensions
+  ] = base_params;
+
+  const auto num_blks = gridDim.x;
+  constexpr auto kWorkersPerBlock = kBlockSize / kWorkThreads;
+  const auto num_workers = num_blks * kWorkersPerBlock;
+  const auto num_q_and_k_heads = num_qo_heads + num_kv_heads;
+  const auto num_works = num_q_and_k_heads * num_tokens;
+  const auto num_extra_works = num_kv_heads * num_tokens;  // rope works + v store works
+  const auto start_worker_id = (blockIdx.x * kBlockSize + threadIdx.x) / kWorkThreads;
+  const auto lane_id = threadIdx.x % kWorkThreads;
+  const auto cos_cache_ptr = cos_sin_cache_ptr;
+  const auto sin_cache_ptr = pointer::offset(cos_sin_cache_ptr, kCosSinStrideBytes / 2);
+
+  auto idx = start_worker_id;
+
+  PDLWaitPrimary<kUsePDL>();
+  // in this case, head_dim = rope_dim must be true
+  __builtin_assume(head_stride_bytes == kRopeDim * sizeof(DType));
+
+  for (; idx < num_works; idx += num_workers) {
+    const int64_t token_id = idx / num_q_and_k_heads;
+    const int64_t head_id = idx % num_q_and_k_heads;
+    const auto pos = static_cast<const IdType*>(positions)[token_id];
+    const auto loc = static_cast<const IdType*>(out_loc)[token_id];
+    const auto load_q = head_id < num_qo_heads;
+    const auto input_ = load_q ? pointer::offset(q, token_id * q_stride_bytes)  //
+                               : pointer::offset(k, token_id * k_stride_bytes);
+    const auto input = pointer::offset(input_, head_id * head_stride_bytes);
+    const auto cos_ptr = pointer::offset(cos_cache_ptr, pos * kCosSinStrideBytes);
+    const auto sin_ptr = pointer::offset(sin_cache_ptr, pos * kCosSinStrideBytes);
+    if constexpr (kIsNeox) {
+      using CacheStorage = AlignedVector<fp32x2_t, kVecSize>;
+      const auto input_x = input;
+      const auto input_y = pointer::offset(input, (kRopeDim / 2) * sizeof(DType));
+      auto input_vec_x = load_as<InputStorage>(input_x, lane_id);
+      auto input_vec_y = load_as<InputStorage>(input_y, lane_id);
+      const auto cos_pair = load_as<CacheStorage>(cos_ptr, lane_id);
+      const auto sin_pair = load_as<CacheStorage>(sin_ptr, lane_id);
+#pragma unroll
+      for (int64_t j = 0; j < kVecSize; ++j) {
+        const auto [x0, x1] = cast<fp32x2_t>(input_vec_x[j]);
+        const auto [y0, y1] = cast<fp32x2_t>(input_vec_y[j]);
+        const auto [cos_0, cos_1] = cos_pair[j];
+        const auto [sin_0, sin_1] = sin_pair[j];
+        const auto out_x0 = x0 * cos_0 - y0 * sin_0;
+        const auto out_y0 = x0 * sin_0 + y0 * cos_0;
+        const auto out_x1 = x1 * cos_1 - y1 * sin_1;
+        const auto out_y1 = x1 * sin_1 + y1 * cos_1;
+        input_vec_x[j] = cast<DType2, fp32x2_t>({out_x0, out_x1});
+        input_vec_y[j] = cast<DType2, fp32x2_t>({out_y0, out_y1});
+      }
+      store_as<InputStorage>(input, input_vec_x, lane_id);
+      const auto input_y_out = pointer::offset(input, (kRopeDim / 2) * sizeof(DType));
+      store_as<InputStorage>(input_y_out, input_vec_y, lane_id);
+      if (!load_q) {
+        const auto k_out = pointer::offset(k_cache, loc * cache_stride_bytes, head_id * head_stride_bytes);
+        store_as<InputStorage>(k_out, input_vec_x, lane_id);
+        const auto k_out_y = pointer::offset(k_out, (kRopeDim / 2) * sizeof(DType));
+        store_as<InputStorage>(k_out_y, input_vec_y, lane_id);
+      }
+    } else {
+      using CacheStorage = AlignedVector<float, kVecSize>;
+      auto input_vec = load_as<InputStorage>(input, lane_id);
+      const auto cos_vec = load_as<CacheStorage>(cos_ptr, lane_id);
+      const auto sin_vec = load_as<CacheStorage>(sin_ptr, lane_id);
+#pragma unroll
+      for (int64_t j = 0; j < kVecSize; ++j) {
+        const auto [x, y] = cast<fp32x2_t>(input_vec[j]);
+        const auto cos = cos_vec[j];
+        const auto sin = sin_vec[j];
+        const auto out_x = x * cos - y * sin;
+        const auto out_y = x * sin + y * cos;
+        input_vec[j] = cast<DType2, fp32x2_t>({out_x, out_y});
+      }
+      store_as<InputStorage>(input, input_vec, lane_id);
+      if (!load_q) {
+        const auto k_out = pointer::offset(k_cache, loc * cache_stride_bytes, head_id * head_stride_bytes);
+        store_as<InputStorage>(k_out, input_vec, lane_id);
+      }
+    }
+  }
+
+  __syncwarp();  // to avoid warp divergence
+  idx -= num_works;
+  for (; idx < num_extra_works; idx += num_workers) {
+    using VStorage = AlignedVector<DType, kRopeDim / kWorkThreads>;
+    const int64_t token_id = idx / num_kv_heads;
+    const int64_t head_id = idx % num_kv_heads;
+    const auto loc = static_cast<const IdType*>(out_loc)[token_id];
+    const auto input = pointer::offset(v_ptr, token_id * v_stride_bytes, head_id * head_stride_bytes);
+    const auto input_vec = load_as<VStorage>(input, lane_id);
+    const auto output = pointer::offset(v_cache, loc * cache_stride_bytes, head_id * head_stride_bytes);
+    store_as<VStorage>(output, input_vec, lane_id);
+  }
+  PDLTriggerSecondary<kUsePDL>();
+}
+
+template <bool kIsNeox, int64_t kRopeDim, bool kUsePDL, typename DType>
+struct FusedRopeKernel {
+  static constexpr uint32_t kDimPerThread = std::gcd(16 / sizeof(DType), kRopeDim);
+  static constexpr uint32_t kWorkThreads = next_pow2(kRopeDim, kDimPerThread);
+  static constexpr bool kSupportFused = kWorkThreads * kDimPerThread == kRopeDim;
+  static_assert(kRopeDim % kDimPerThread == 0);
+  static_assert(kBlockSize % kWorkThreads == 0);
+
+  template <typename IdType>
+  static constexpr auto _kernel_0 = fused_rope_kernel<kIsNeox, kRopeDim, kUsePDL, DType, IdType, kWorkThreads>;
+  template <typename IdType>
+  static constexpr auto _kernel_1 = fused_rope_store_kernel<kIsNeox, kRopeDim, kUsePDL, DType, IdType, kWorkThreads>;
+
+  static auto get_num_sm(DLDevice device) {
+    static const auto kNumSM = host::runtime::get_sm_count(device.device_id);
+    return kNumSM;
+  }
+
+  static void
+  run(const tvm::ffi::TensorView q,
+      const tvm::ffi::TensorView k,
+      const tvm::ffi::TensorView cos_sin_cache,
+      const tvm::ffi::TensorView positions) {
+    using namespace host;
+    auto N = SymbolicSize{"num_tokens"};
+    auto Q = SymbolicSize{"num_qo_heads"};
+    auto K = SymbolicSize{"num_kv_heads"};
+    auto D = SymbolicSize{"rope_dim"};
+    auto Dq = SymbolicSize{"q_stride"};
+    auto Dk = SymbolicSize{"k_stride"};
+    auto Dd = SymbolicSize{"head_stride"};
+    auto device = SymbolicDevice{};
+    auto id_type = SymbolicDType{};
+    D.set_value(kRopeDim);
+    device.set_options<kDLCUDA>();
+    TensorMatcher({N, Q, D})  // q input
+        .with_strides({Dq, Dd, 1})
+        .with_dtype<DType>()
+        .with_device(device)
+        .verify(q);
+    TensorMatcher({N, K, D})  // k input
+        .with_strides({Dk, Dd, 1})
+        .with_dtype<DType>()
+        .with_device(device)
+        .verify(k);
+    TensorMatcher({-1, D})  // cos_sin_cache
+        .with_dtype<float>()
+        .with_device(device)
+        .verify(cos_sin_cache);
+    TensorMatcher({N})  // positions
+        .with_dtype<int32_t, int64_t>(id_type)
+        .with_device(device)
+        .verify(positions);
+
+    const auto num_tokens = static_cast<uint32_t>(N.unwrap());
+    const auto num_qo_heads = static_cast<uint32_t>(Q.unwrap());
+    const auto num_kv_heads = static_cast<uint32_t>(K.unwrap());
+    const auto q_stride_bytes = static_cast<int64_t>(Dq.unwrap() * sizeof(DType));
+    const auto k_stride_bytes = static_cast<int64_t>(Dk.unwrap() * sizeof(DType));
+    const auto head_stride_bytes = static_cast<int64_t>(Dd.unwrap() * sizeof(DType));
+
+    // NOTE: we offset the k here to reduce computation cost in the kernel
+    const int64_t k_offset = static_cast<int64_t>(num_qo_heads) * head_stride_bytes;
+    const auto params = FusedRopeParams{
+        .q_ptr = q.data_ptr(),
+        .k_ptr = pointer::offset(k.data_ptr(), -k_offset),
+        .cos_sin_cache_ptr = cos_sin_cache.data_ptr(),
+        .positions = positions.data_ptr(),
+        .q_stride_bytes = q_stride_bytes,
+        .k_stride_bytes = k_stride_bytes,
+        .head_stride_bytes = head_stride_bytes,
+        .num_qo_heads = num_qo_heads,
+        .num_kv_heads = num_kv_heads,
+        .num_tokens = num_tokens,
+    };
+
+    const auto is_int32 = id_type.is_type<int32_t>();
+    const auto kernel = is_int32 ? _kernel_0<int32_t> : _kernel_0<int64_t>;
+    const uint32_t kNumSM = get_num_sm(device.unwrap());
+    static const uint32_t kOccupancyTable[2] = {
+        runtime::get_blocks_per_sm(_kernel_0<int32_t>, kBlockSize),
+        runtime::get_blocks_per_sm(_kernel_0<int64_t>, kBlockSize),
+    };
+    const auto max_blocks = kOccupancyTable[is_int32 ? 0 : 1] * kNumSM;
+    const auto num_works = (num_qo_heads + num_kv_heads) * num_tokens;
+    const auto needed_blocks = div_ceil(num_works, (kBlockSize / kWorkThreads));
+    const auto num_blocks = std::min(max_blocks, needed_blocks);
+    LaunchKernel(num_blocks, kBlockSize, device.unwrap())  //
+        .enable_pdl(kUsePDL)(kernel, params);
+  }
+
+  static void run_fused(
+      const tvm::ffi::TensorView q,
+      const tvm::ffi::TensorView k,
+      const tvm::ffi::TensorView v,
+      const tvm::ffi::TensorView k_cache,
+      const tvm::ffi::TensorView v_cache,
+      const tvm::ffi::TensorView cos_sin_cache,
+      const tvm::ffi::TensorView positions,
+      const tvm::ffi::TensorView out_loc) {
+    if constexpr (kSupportFused) {
+      return _run_fused_impl(q, k, v, k_cache, v_cache, cos_sin_cache, positions, out_loc);
+    } else {
+      host::Panic("Fused rope + store is not supported for rope_dim ", kRopeDim);
+    }
+  }
+
+  static void _run_fused_impl(
+      const tvm::ffi::TensorView q,
+      const tvm::ffi::TensorView k,
+      const tvm::ffi::TensorView v,
+      const tvm::ffi::TensorView k_cache,
+      const tvm::ffi::TensorView v_cache,
+      const tvm::ffi::TensorView cos_sin_cache,
+      const tvm::ffi::TensorView positions,
+      const tvm::ffi::TensorView out_loc) {
+    using namespace host;
+
+    auto N = SymbolicSize{"num_tokens"};
+    auto Q = SymbolicSize{"num_qo_heads"};
+    auto K = SymbolicSize{"num_kv_heads"};
+    auto D = SymbolicSize{"rope_dim"};
+    auto R = SymbolicSize{"row_size"};
+    auto Dq = SymbolicSize{"q_stride"};
+    auto Dk = SymbolicSize{"k_stride"};
+    auto Dv = SymbolicSize{"v_stride"};
+    auto Dd = SymbolicSize{"head_stride"};
+    auto Dc = SymbolicSize{"cache_stride"};
+    auto device = SymbolicDevice{};
+    auto id_type = SymbolicDType{};
+    D.set_value(kRopeDim);
+    device.set_options<kDLCUDA>();
+
+    TensorMatcher({N, Q, D})  // q input
+        .with_strides({Dq, Dd, 1})
+        .with_dtype<DType>()
+        .with_device(device)
+        .verify(q);
+    TensorMatcher({N, K, D})  // k input
+        .with_strides({Dk, Dd, 1})
+        .with_dtype<DType>()
+        .with_device(device)
+        .verify(k);
+    TensorMatcher({N, K, D})  // v input
+        .with_strides({Dv, Dd, 1})
+        .with_dtype<DType>()
+        .with_device(device)
+        .verify(v);
+    TensorMatcher({-1, D})  // cos_sin_cache
+        .with_dtype<float>()
+        .with_device(device)
+        .verify(cos_sin_cache);
+    TensorMatcher({N})  // positions, out_loc
+        .with_dtype<int32_t, int64_t>(id_type)
+        .with_device(device)
+        .verify(positions)
+        .verify(out_loc);
+    TensorMatcher({-1, R})  // k_cache
+        .with_strides({Dc, 1})
+        .with_dtype<DType>()
+        .with_device(device)
+        .verify(k_cache)
+        .verify(v_cache);
+
+    const auto num_tokens = static_cast<uint32_t>(N.unwrap());
+    const auto num_qo_heads = static_cast<uint32_t>(Q.unwrap());
+    const auto num_kv_heads = static_cast<uint32_t>(K.unwrap());
+    const auto q_stride_bytes = static_cast<int64_t>(Dq.unwrap() * sizeof(DType));
+    const auto k_stride_bytes = static_cast<int64_t>(Dk.unwrap() * sizeof(DType));
+    const auto head_stride = Dd.unwrap();
+    const auto row_dim = R.unwrap();
+    const auto head_stride_bytes = static_cast<int64_t>(Dd.unwrap() * sizeof(DType));
+
+    RuntimeCheck(kRopeDim == head_stride, "rope_dim ", kRopeDim, " should = head_stride ", head_stride);
+    RuntimeCheck(num_kv_heads * kRopeDim == row_dim, "invalid kvcache");
+
+    // NOTE: we offset the k here to reduce computation cost in the kernel
+    const int64_t k_offset = static_cast<int64_t>(num_qo_heads) * head_stride_bytes;
+    const auto params = FusedRopeParams{
+        .q_ptr = q.data_ptr(),
+        .k_ptr = pointer::offset(k.data_ptr(), -k_offset),
+        .cos_sin_cache_ptr = cos_sin_cache.data_ptr(),
+        .positions = positions.data_ptr(),
+        .q_stride_bytes = q_stride_bytes,
+        .k_stride_bytes = k_stride_bytes,
+        .head_stride_bytes = head_stride_bytes,
+        .num_qo_heads = num_qo_heads,
+        .num_kv_heads = num_kv_heads,
+        .num_tokens = num_tokens,
+    };
+
+    const auto v_stride_bytes = static_cast<int64_t>(Dv.unwrap() * sizeof(DType));
+    const auto cache_stride_bytes = static_cast<int64_t>(Dc.unwrap() * sizeof(DType));
+    const auto store_params = FusedRopeStoreParams{
+        .base_params = params,
+        .v_ptr = v.data_ptr(),
+        .k_cache = pointer::offset(k_cache.data_ptr(), -k_offset),
+        .v_cache = v_cache.data_ptr(),
+        .out_loc = out_loc.data_ptr(),
+        .v_stride_bytes = v_stride_bytes,
+        .cache_stride_bytes = cache_stride_bytes,
+    };
+
+    const auto is_int32 = id_type.is_type<int32_t>();
+    const auto kernel = is_int32 ? _kernel_1<int32_t> : _kernel_1<int64_t>;
+    const uint32_t kNumSM = get_num_sm(device.unwrap());
+    static const uint32_t kOccupancyTable[2] = {
+        runtime::get_blocks_per_sm(_kernel_1<int32_t>, kBlockSize),
+        runtime::get_blocks_per_sm(_kernel_1<int64_t>, kBlockSize),
+    };
+    const auto max_blocks = kOccupancyTable[is_int32 ? 0 : 1] * kNumSM;
+    // rope works for q+k heads, plus v store works for kv heads
+    const auto num_total_works = (num_qo_heads + 2 * num_kv_heads) * num_tokens;
+    const auto needed_blocks = div_ceil(num_total_works, (kBlockSize / kWorkThreads));
+    const auto num_blocks = std::min(max_blocks, needed_blocks);
+    LaunchKernel(num_blocks, kBlockSize, device.unwrap())  //
+        .enable_pdl(kUsePDL)(kernel, store_params);
+  }
+};
+
+}  // namespace

python/sglang/jit_kernel/csrc/elementwise/set_mla_kv_buffer.cuh

@@ -0,0 +1,249 @@
+// JIT TMA bulk-store kernel for MLA paged-KV scatter writes.
+//
+// Each warp:
+//   1. Cooperatively loads one item's (nope, rope) row into a per-warp slot in
+//      shared memory via vectorised ld/st.
+//   2. Lane 0 issues a single ``cp.async.bulk.global.shared::cta`` (TMA bulk
+//      store, non-tensor variant) to scatter the row to
+//      ``kv_buffer + loc[item] * stride_buffer``.
+//
+// End-of-CTA: ``cp.async.bulk.commit_group`` + ``wait_group<0>`` ensures all
+// in-flight stores commit before the kernel exits so the writes are visible
+// to subsequent kernels and the host.
+//
+// Two correctness gotchas worth a comment (easy to lose):
+//   - ``fence.proxy.async.shared::cta`` between the smem fill and the TMA
+//     store. The TMA engine reads via the async proxy; without the fence it
+//     observes stale smem under heavy concurrency (manifests as zero rows at
+//     large bs).
+//   - ``wait_group`` not ``wait_group_read`` — the latter only allows early
+//     smem reuse; it does not wait for the gmem store to commit globally.
+
+#pragma once
+
+#include <sgl_kernel/tensor.h>
+#include <sgl_kernel/utils.h>
+
+#include <sgl_kernel/tile.cuh>
+#include <sgl_kernel/utils.cuh>
+#include <sgl_kernel/vec.cuh>
+
+#include <cuda/ptx>
+#include <dlpack/dlpack.h>
+#include <tvm/ffi/container/tensor.h>
+
+#include <cstdint>
+
+namespace {
+
+struct SetMlaKVBufferParams {
+  const void* __restrict__ k_nope;
+  const void* __restrict__ k_rope;
+  void* __restrict__ kv_buffer;
+  const void* __restrict__ loc;
+  int64_t stride_nope_bytes;
+  int64_t stride_rope_bytes;
+  int64_t stride_buffer_bytes;
+  uint32_t batch_size;
+};
+
+// Warp-cooperative gmem -> smem copy. Picks the widest vec width that divides
+// both the per-thread share and the byte total. Caller guarantees src is
+// 16-byte aligned (PyTorch tensors are) and dst is the start of a per-warp
+// smem slot (also 16-byte aligned by ``alignas(16)``).
+template <int64_t kBytes>
+SGL_DEVICE void warp_g2s_copy(const void* __restrict__ src, void* __restrict__ dst) {
+  using namespace device;
+  constexpr int64_t kAlignment = (kBytes % (16 * kWarpThreads) == 0)  ? 16
+                                 : (kBytes % (8 * kWarpThreads) == 0) ? 8
+                                 : (kBytes % (4 * kWarpThreads) == 0) ? 4
+                                 : (kBytes % 4 == 0)                  ? 4
+                                                                      : 0;
+  static_assert(kAlignment > 0, "kBytes must be a multiple of 4");
+
+  using vec_t = AlignedStorage<uint32_t, kAlignment / 4>;
+  constexpr auto kLoopBytes = sizeof(vec_t) * kWarpThreads;
+  constexpr auto kLoopCount = kBytes / kLoopBytes;
+  constexpr int64_t kTailVecs = (kBytes - kLoopCount * kLoopBytes) / sizeof(vec_t);
+
+  const auto gmem = tile::Memory<vec_t>::warp();
+
+#pragma unroll
+  for (int64_t i = 0; i < kLoopCount; ++i) {
+    const auto v = gmem.load(src, i);
+    gmem.store(dst, v, i);
+  }
+  if constexpr (kTailVecs > 0) {
+    if (gmem.in_bound(kLoopCount * kWarpThreads + kTailVecs, kLoopCount)) {
+      const auto v = gmem.load(src, kLoopCount);
+      gmem.store(dst, v, kLoopCount);
+    }
+  }
+}
+
+template <int64_t kNopeBytes, int64_t kRopeBytes, int kNumWarps, bool kUsePDL, typename TLoc>
+__global__ void set_mla_kv_buffer_kernel(const __grid_constant__ SetMlaKVBufferParams params) {
+  using namespace device;
+  static_assert((kNopeBytes + kRopeBytes) % 16 == 0, "TMA bulk store requires total row to be 16-byte aligned");
+
+  constexpr int64_t kRowBytes = kNopeBytes + kRopeBytes;
+
+  // One contiguous smem slot per warp; align to 16 for TMA.
+  __shared__ alignas(16) uint8_t smem[kNumWarps][kRowBytes];
+
+  const uint32_t warp_in_cta = threadIdx.x / kWarpThreads;
+  const uint32_t item_id = blockIdx.x * kNumWarps + warp_in_cta;
+  if (item_id >= params.batch_size) return;
+
+  PDLWaitPrimary<kUsePDL>();
+
+  const int64_t loc = static_cast<int64_t>(static_cast<const TLoc*>(params.loc)[item_id]);
+
+  const auto nope_src = pointer::offset(params.k_nope, item_id * params.stride_nope_bytes);
+  const auto rope_src = pointer::offset(params.k_rope, item_id * params.stride_rope_bytes);
+  void* const gmem_dst = pointer::offset(params.kv_buffer, loc * params.stride_buffer_bytes);
+
+  // Warp-cooperative load (nope, rope) into the per-warp smem slot.
+  warp_g2s_copy<kNopeBytes>(nope_src, &smem[warp_in_cta][0]);
+  warp_g2s_copy<kRopeBytes>(rope_src, &smem[warp_in_cta][kNopeBytes]);
+
+  // Fence required: TMA reads smem via the async proxy, normal sts writes
+  // through the generic proxy. Without this the TMA engine can observe stale
+  // values at large bs.
+  __syncwarp();
+  asm volatile("fence.proxy.async.shared::cta;" ::: "memory");
+
+  // Lane 0 issues one bulk store from the smem slot to the scattered gmem row.
+  if (threadIdx.x % kWarpThreads == 0) {
+    cuda::ptx::cp_async_bulk(
+        cuda::ptx::space_global,
+        cuda::ptx::space_shared,
+        gmem_dst,
+        &smem[warp_in_cta][0],
+        static_cast<uint32_t>(kRowBytes));
+  }
+
+  // Commit and wait for the CTA's bulk-stores to be globally visible before
+  // returning. ``wait_group`` (not ``_read``) is the one that waits for gmem
+  // commit; ``_read`` only releases smem for reuse.
+  cuda::ptx::cp_async_bulk_commit_group();
+  cuda::ptx::cp_async_bulk_wait_group(cuda::ptx::n32_t<0>{});
+
+  PDLTriggerSecondary<kUsePDL>();
+}
+
+template <int64_t kNopeBytes, int64_t kRopeBytes, bool kUsePDL>
+struct SetMlaKVBufferKernel {
+  static_assert(kNopeBytes > 0 && kNopeBytes % 4 == 0, "kNopeBytes must be a positive multiple of 4");
+  static_assert(kRopeBytes > 0 && kRopeBytes % 4 == 0, "kRopeBytes must be a positive multiple of 4");
+  static_assert(
+      (kNopeBytes + kRopeBytes) % 16 == 0, "TMA bulk store requires (kNopeBytes + kRopeBytes) to be a multiple of 16");
+
+  template <int kNumWarps, typename TLoc>
+  static constexpr auto kernel = set_mla_kv_buffer_kernel<kNopeBytes, kRopeBytes, kNumWarps, kUsePDL, TLoc>;
+
+  static void
+  run(tvm::ffi::TensorView kv_buffer,
+      tvm::ffi::TensorView loc,
+      tvm::ffi::TensorView k_nope,
+      tvm::ffi::TensorView k_rope,
+      int64_t num_warps_per_block) {
+    using namespace host;
+
+    auto B = SymbolicSize{"batch_size"};
+    auto D_nope = SymbolicSize{"nope_dim"};
+    auto D_rope = SymbolicSize{"rope_dim"};
+    auto D_buf = SymbolicSize{"buffer_last_dim"};
+    auto S_nope = SymbolicSize{"nope_stride"};
+    auto S_rope = SymbolicSize{"rope_stride"};
+    auto S_buf = SymbolicSize{"buffer_stride"};
+    auto S_loc = SymbolicSize{"loc_stride"};
+    auto dtype = SymbolicDType{};
+    auto loc_dtype = SymbolicDType{};
+    auto device = SymbolicDevice{};
+    device.set_options<kDLCUDA>();
+
+    TensorMatcher({B, D_nope})  //
+        .with_strides({S_nope, 1})
+        .with_dtype(dtype)
+        .with_device(device)
+        .verify(k_nope);
+    TensorMatcher({B, D_rope})  //
+        .with_strides({S_rope, 1})
+        .with_dtype(dtype)
+        .with_device(device)
+        .verify(k_rope);
+    TensorMatcher({-1, D_buf})  //
+        .with_strides({S_buf, 1})
+        .with_dtype(dtype)
+        .with_device(device)
+        .verify(kv_buffer);
+    TensorMatcher({B})  //
+        .with_strides({S_loc})
+        .with_dtype<int32_t, int64_t>(loc_dtype)
+        .with_device(device)
+        .verify(loc);
+
+    const int64_t dtype_size = dtype_bytes(dtype.unwrap());
+    RuntimeCheck(
+        kNopeBytes == dtype_size * D_nope.unwrap(),
+        "kNopeBytes mismatch: expected ",
+        kNopeBytes,
+        ", got ",
+        dtype_size * D_nope.unwrap());
+    RuntimeCheck(
+        kRopeBytes == dtype_size * D_rope.unwrap(),
+        "kRopeBytes mismatch: expected ",
+        kRopeBytes,
+        ", got ",
+        dtype_size * D_rope.unwrap());
+    RuntimeCheck(dtype_size * D_buf.unwrap() >= kNopeBytes + kRopeBytes, "kv_buffer last dim too small");
+    RuntimeCheck(
+        (S_buf.unwrap() * dtype_size) % 16 == 0,
+        "kv_buffer row stride must be a multiple of 16 bytes for TMA bulk store; got ",
+        S_buf.unwrap() * dtype_size);
+
+    const uint32_t batch = static_cast<uint32_t>(B.unwrap());
+    if (batch == 0) return;
+
+    const auto params = SetMlaKVBufferParams{
+        .k_nope = k_nope.data_ptr(),
+        .k_rope = k_rope.data_ptr(),
+        .kv_buffer = kv_buffer.data_ptr(),
+        .loc = loc.data_ptr(),
+        .stride_nope_bytes = S_nope.unwrap() * dtype_size,
+        .stride_rope_bytes = S_rope.unwrap() * dtype_size,
+        .stride_buffer_bytes = S_buf.unwrap() * dtype_size,
+        .batch_size = batch,
+    };
+
+    const auto use_int32 = loc_dtype.is_type<int32_t>();
+
+    auto launch = [&]<int kNW>() {
+      const auto kernel_ptr = use_int32 ? kernel<kNW, int32_t> : kernel<kNW, int64_t>;
+      const uint32_t num_blocks = div_ceil(batch, static_cast<uint32_t>(kNW));
+      const uint32_t threads_per_block = static_cast<uint32_t>(kNW) * device::kWarpThreads;
+      LaunchKernel(num_blocks, threads_per_block, device.unwrap())  //
+          .enable_pdl(kUsePDL)(kernel_ptr, params);
+    };
+
+    switch (num_warps_per_block) {
+      case 1:
+        launch.template operator()<1>();
+        break;
+      case 2:
+        launch.template operator()<2>();
+        break;
+      case 4:
+        launch.template operator()<4>();
+        break;
+      case 8:
+        launch.template operator()<8>();
+        break;
+      default:
+        Panic("Unsupported num_warps_per_block=", num_warps_per_block);
+    }
+  }
+};
+
+}  // namespace