chore: import upstream snapshot with attribution

python/sglang/jit_kernel/csrc/diffusion/causal_conv3d_cat_pad.cuh

@@ -0,0 +1,253 @@
+// Native CUDA fast path for Cosmos3 VAE causal-Conv3D cat/pad copy.
+//
+// The op writes the output of:
+//   pad(cat(cache_x, x, dim=T), (Wl, Wr, Ht, Hb, Dl - cache_t, Dr))
+// for 5D NCTHW tensors. It is a memory-bound copy/zero-fill kernel and is only
+// entered for contiguous CUDA tensors; unsupported cases fall back to Triton in
+// the Python caller.
+//
+// Developed with MIT HAN Lab Kernel Design Agents:
+// https://github.com/mit-han-lab/kernel-design-agents
+
+#pragma once
+
+#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 <sgl_kernel/vec.cuh>    // For device::AlignedVector
+
+#include <cstdint>
+
+namespace sglang_causal_conv3d_cat_pad {
+
+namespace {
+
+constexpr int kBlockSize = 256;
+
+template <typename ET, int kVec>
+__global__ void __launch_bounds__(kBlockSize) cat_pad_flat_kernel(
+    const ET* __restrict__ x,
+    const ET* __restrict__ cache,
+    ET* __restrict__ out,
+    int64_t total_vecs,
+    int64_t channels,
+    int64_t t_size,
+    int64_t h_size,
+    int64_t w_size,
+    int64_t cache_t,
+    int64_t out_t,
+    int64_t out_h,
+    int64_t out_w,
+    int64_t pad_d_left,
+    int64_t pad_h_top,
+    int64_t pad_w_left) {
+  using Pack = device::AlignedVector<ET, kVec>;
+
+  const int64_t nthreads = static_cast<int64_t>(gridDim.x) * blockDim.x;
+  for (int64_t vid = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x; vid < total_vecs; vid += nthreads) {
+    int64_t base = vid * kVec;
+    int64_t ow = base % out_w;
+    int64_t tmp = base / out_w;
+    int64_t oh = tmp % out_h;
+    tmp /= out_h;
+    int64_t od = tmp % out_t;
+    tmp /= out_t;
+    int64_t oc = tmp % channels;
+    int64_t ob = tmp / channels;
+
+    int64_t ih = oh - pad_h_top;
+    int64_t src_t = od - pad_d_left;
+    bool interior = ih >= 0 && ih < h_size && src_t >= 0 && src_t < cache_t + t_size;
+
+    const ET* src = nullptr;
+    if (interior) {
+      if (src_t < cache_t) {
+        src = cache + (((ob * channels + oc) * cache_t + src_t) * h_size + ih) * w_size;
+      } else {
+        src = x + (((ob * channels + oc) * t_size + (src_t - cache_t)) * h_size + ih) * w_size;
+      }
+    }
+
+    Pack pack;
+#pragma unroll
+    for (int i = 0; i < kVec; ++i) {
+      ET value = ET(0);
+      if (interior) {
+        const int64_t iw = ow - pad_w_left;
+        if (iw >= 0 && iw < w_size) {
+          value = SGLANG_LDG(src + iw);
+        }
+      }
+      pack[i] = value;
+
+      if (++ow == out_w) {
+        ow = 0;
+        if (++oh == out_h) {
+          oh = 0;
+          if (++od == out_t) {
+            od = 0;
+            if (++oc == channels) {
+              oc = 0;
+              ++ob;
+            }
+          }
+        }
+        ih = oh - pad_h_top;
+        src_t = od - pad_d_left;
+        interior = ih >= 0 && ih < h_size && src_t >= 0 && src_t < cache_t + t_size;
+        if (interior) {
+          if (src_t < cache_t) {
+            src = cache + (((ob * channels + oc) * cache_t + src_t) * h_size + ih) * w_size;
+          } else {
+            src = x + (((ob * channels + oc) * t_size + (src_t - cache_t)) * h_size + ih) * w_size;
+          }
+        } else {
+          src = nullptr;
+        }
+      }
+    }
+
+    pack.store(out, vid);
+  }
+}
+
+template <typename ET, int kVec>
+void launch_cat_pad_flat(
+    const void* x,
+    const void* cache,
+    void* out,
+    int64_t total,
+    int64_t channels,
+    int64_t t_size,
+    int64_t h_size,
+    int64_t w_size,
+    int64_t cache_t,
+    int64_t out_t,
+    int64_t out_h,
+    int64_t out_w,
+    int64_t depth_left,
+    int64_t pad_h_top,
+    int64_t pad_w_left,
+    DLDevice device) {
+  const int64_t total_vecs = total / kVec;
+  const uint32_t grid = static_cast<uint32_t>(host::div_ceil(total_vecs, static_cast<int64_t>(kBlockSize)));
+  host::LaunchKernel(grid, kBlockSize, device)(
+      cat_pad_flat_kernel<ET, kVec>,
+      static_cast<const ET*>(x),
+      static_cast<const ET*>(cache),
+      static_cast<ET*>(out),
+      total_vecs,
+      channels,
+      t_size,
+      h_size,
+      w_size,
+      cache_t,
+      out_t,
+      out_h,
+      out_w,
+      depth_left,
+      pad_h_top,
+      pad_w_left);
+}
+
+}  // namespace
+
+template <typename T>
+struct CausalConv3dCatPadKernel {
+  static void
+  run(tvm::ffi::TensorView out,
+      tvm::ffi::TensorView x,
+      tvm::ffi::TensorView cache,
+      int64_t pad_w_left,
+      int64_t pad_w_right,
+      int64_t pad_h_top,
+      int64_t pad_h_bottom,
+      int64_t pad_d_left,
+      int64_t pad_d_right) {
+    using namespace host;
+
+    auto bsz = SymbolicSize{"batch"};
+    auto channels = SymbolicSize{"channels"};
+    auto t_size = SymbolicSize{"t_size"};
+    auto h_size = SymbolicSize{"h_size"};
+    auto w_size = SymbolicSize{"w_size"};
+    auto cache_t = SymbolicSize{"cache_t"};
+    auto out_t = SymbolicSize{"out_t"};
+    auto out_h = SymbolicSize{"out_h"};
+    auto out_w = SymbolicSize{"out_w"};
+    auto device = SymbolicDevice{};
+    device.set_options<kDLGPU>();
+
+    TensorMatcher({bsz, channels, t_size, h_size, w_size})
+        .with_dtype<T>()
+        .template with_device<kDLGPU>(device)
+        .verify(x);
+    TensorMatcher({bsz, channels, cache_t, h_size, w_size})
+        .with_dtype<T>()
+        .template with_device<kDLGPU>(device)
+        .verify(cache);
+    TensorMatcher({bsz, channels, out_t, out_h, out_w})
+        .with_dtype<T>()
+        .template with_device<kDLGPU>(device)
+        .verify(out);
+
+    const int64_t depth_left = pad_d_left - cache_t.unwrap();
+    RuntimeCheck(depth_left >= 0, "pad_d_left must be >= cache_t");
+    RuntimeCheck(pad_d_right == 0, "pad_d_right must be 0");
+    RuntimeCheck(pad_w_left == pad_w_right, "width padding must be symmetric");
+    RuntimeCheck(pad_h_top == pad_h_bottom, "height padding must be symmetric");
+    RuntimeCheck(out_t.unwrap() == t_size.unwrap() + cache_t.unwrap() + depth_left + pad_d_right, "out_t mismatch");
+    RuntimeCheck(out_h.unwrap() == h_size.unwrap() + pad_h_top + pad_h_bottom, "out_h mismatch");
+    RuntimeCheck(out_w.unwrap() == w_size.unwrap() + pad_w_left + pad_w_right, "out_w mismatch");
+
+    const int64_t total = bsz.unwrap() * channels.unwrap() * out_t.unwrap() * out_h.unwrap() * out_w.unwrap();
+    if (total == 0) {
+      return;
+    }
+
+    constexpr int kVec = 16 / sizeof(T);
+    RuntimeCheck(total % kVec == 0, "output element count must be divisible by vector width");
+    RuntimeCheck(reinterpret_cast<uintptr_t>(out.data_ptr()) % 16 == 0, "output pointer must be 16-byte aligned");
+
+    if constexpr (sizeof(T) == 2) {
+      launch_cat_pad_flat<uint16_t, kVec>(
+          x.data_ptr(),
+          cache.data_ptr(),
+          out.data_ptr(),
+          total,
+          channels.unwrap(),
+          t_size.unwrap(),
+          h_size.unwrap(),
+          w_size.unwrap(),
+          cache_t.unwrap(),
+          out_t.unwrap(),
+          out_h.unwrap(),
+          out_w.unwrap(),
+          depth_left,
+          pad_h_top,
+          pad_w_left,
+          device.unwrap());
+    } else {
+      launch_cat_pad_flat<uint32_t, kVec>(
+          x.data_ptr(),
+          cache.data_ptr(),
+          out.data_ptr(),
+          total,
+          channels.unwrap(),
+          t_size.unwrap(),
+          h_size.unwrap(),
+          w_size.unwrap(),
+          cache_t.unwrap(),
+          out_t.unwrap(),
+          out_h.unwrap(),
+          out_w.unwrap(),
+          depth_left,
+          pad_h_top,
+          pad_w_left,
+          device.unwrap());
+    }
+  }
+};
+
+}  // namespace sglang_causal_conv3d_cat_pad

python/sglang/jit_kernel/csrc/diffusion/ltx2_qknorm_split_rope.cuh

@@ -0,0 +1,276 @@
+// CUDA fast path for LTX2 Q/K RMSNorm + split RoPE.
+//
+// Developed with MIT HAN Lab Kernel Design Agents:
+// https://github.com/mit-han-lab/kernel-design-agents
+//
+// This mirrors the LTX2 eager oracle: RMSNorm and split RoPE both run in
+// fp32, rounding to bf16 only once at the final attention input.
+
+#pragma once
+
+#include <sgl_kernel/tensor.h>  // For TensorMatcher, SymbolicSize, SymbolicDevice
+#include <sgl_kernel/utils.h>   // For RuntimeCheck
+
+#include <sgl_kernel/utils.cuh>  // For LaunchKernel and CUDA dtype aliases
+
+#include <cstdint>
+#include <cuda_bf16.h>
+
+namespace sglang_ltx2_qknorm_split_rope {
+
+namespace {
+
+constexpr int kThreads = 128;
+
+inline const char* data_ptr(const tvm::ffi::TensorView& t) {
+  return static_cast<const char*>(t.data_ptr()) + t.byte_offset();
+}
+
+inline char* mutable_data_ptr(const tvm::ffi::TensorView& t) {
+  return static_cast<char*>(t.data_ptr()) + t.byte_offset();
+}
+
+SGL_DEVICE float compute_rstd(
+    const bf16_t* __restrict__ xrow,
+    int64_t hidden_size,
+    float eps,
+    int tid,
+    int lane,
+    int warp_id,
+    float* warp_sum,
+    float* s_rstd) {
+  float local = 0.f;
+  const int64_t n_vec = hidden_size >> 2;
+  for (int64_t i = tid; i < n_vec; i += kThreads) {
+    const int64_t base = i << 2;
+    const float v0 = __bfloat162float(xrow[base + 0]);
+    const float v1 = __bfloat162float(xrow[base + 1]);
+    const float v2 = __bfloat162float(xrow[base + 2]);
+    const float v3 = __bfloat162float(xrow[base + 3]);
+    local = fmaf(v0, v0, local);
+    local = fmaf(v1, v1, local);
+    local = fmaf(v2, v2, local);
+    local = fmaf(v3, v3, local);
+  }
+
+#pragma unroll
+  for (int offset = 16; offset > 0; offset >>= 1) {
+    local += __shfl_down_sync(0xffffffffu, local, offset);
+  }
+  if (lane == 0) {
+    warp_sum[warp_id] = local;
+  }
+  __syncthreads();
+
+  if (tid == 0) {
+    const float total = (warp_sum[0] + warp_sum[2]) + (warp_sum[1] + warp_sum[3]);
+    *s_rstd = rsqrtf(total / static_cast<float>(hidden_size) + eps);
+  }
+  __syncthreads();
+  return *s_rstd;
+}
+
+SGL_DEVICE float norm_value(float x, float weight, float rstd) {
+  return weight * (rstd * x);
+}
+
+SGL_DEVICE void rope_pair(float x0, float x1, float cos, float sin, float& y0, float& y1) {
+  const float p0 = x0 * cos;
+  const float p1 = x1 * cos;
+  y0 = fmaf(-sin, x1, p0);
+  y1 = fmaf(sin, x0, p1);
+}
+
+__global__ void ltx2_qknorm_split_rope_kernel(
+    const bf16_t* __restrict__ x,
+    const bf16_t* __restrict__ cos,
+    const bf16_t* __restrict__ sin,
+    const bf16_t* __restrict__ weight,
+    bf16_t* __restrict__ out,
+    float eps,
+    int64_t seq_len,
+    int64_t num_heads,
+    int64_t head_dim,
+    int64_t stride_cos_b,
+    int64_t stride_cos_h,
+    int64_t stride_cos_t,
+    int64_t stride_sin_b,
+    int64_t stride_sin_h,
+    int64_t stride_sin_t) {
+  const int64_t row = static_cast<int64_t>(blockIdx.x);
+  const int64_t batch = row / seq_len;
+  const int64_t token = row - batch * seq_len;
+  const int64_t hidden_size = num_heads * head_dim;
+  const int64_t half_dim = head_dim >> 1;
+  const auto* __restrict__ xrow = x + row * hidden_size;
+  auto* __restrict__ outrow = out + row * hidden_size;
+  const int tid = threadIdx.x + threadIdx.y * 32;
+  const int lane = threadIdx.x;
+  const int warp_id = threadIdx.y;
+
+  __shared__ float warp_sum[4];
+  __shared__ float s_rstd;
+  const float rstd = compute_rstd(xrow, hidden_size, eps, tid, lane, warp_id, warp_sum, &s_rstd);
+
+  const int64_t num_pairs = num_heads * half_dim;
+  for (int64_t pair = tid; pair < num_pairs; pair += kThreads) {
+    const int64_t head = pair / half_dim;
+    const int64_t offset = pair - head * half_dim;
+    const int64_t idx0 = head * head_dim + offset;
+    const int64_t idx1 = idx0 + half_dim;
+    const float n0 = norm_value(__bfloat162float(xrow[idx0]), __bfloat162float(weight[idx0]), rstd);
+    const float n1 = norm_value(__bfloat162float(xrow[idx1]), __bfloat162float(weight[idx1]), rstd);
+    const int64_t cos_offset = batch * stride_cos_b + head * stride_cos_h + token * stride_cos_t + offset;
+    const int64_t sin_offset = batch * stride_sin_b + head * stride_sin_h + token * stride_sin_t + offset;
+
+    float y0;
+    float y1;
+    rope_pair(n0, n1, __bfloat162float(cos[cos_offset]), __bfloat162float(sin[sin_offset]), y0, y1);
+    outrow[idx0] = __float2bfloat16_rn(y0);
+    outrow[idx1] = __float2bfloat16_rn(y1);
+  }
+}
+
+inline void launch_one(
+    const tvm::ffi::TensorView& x,
+    const tvm::ffi::TensorView& cos,
+    const tvm::ffi::TensorView& sin,
+    const tvm::ffi::TensorView& weight,
+    const tvm::ffi::TensorView& out,
+    float eps,
+    int64_t num_rows,
+    int64_t seq_len,
+    int64_t num_heads,
+    int64_t head_dim,
+    int64_t stride_cos_b,
+    int64_t stride_cos_h,
+    int64_t stride_cos_t,
+    int64_t stride_sin_b,
+    int64_t stride_sin_h,
+    int64_t stride_sin_t,
+    DLDevice device) {
+  if (num_rows == 0) {
+    return;
+  }
+  host::RuntimeCheck(num_rows <= static_cast<int64_t>(UINT32_MAX), "LTX2 QKNorm split-RoPE grid is too large");
+  host::LaunchKernel(dim3(static_cast<uint32_t>(num_rows)), dim3(32, 4), device)(
+      ltx2_qknorm_split_rope_kernel,
+      reinterpret_cast<const bf16_t*>(data_ptr(x)),
+      reinterpret_cast<const bf16_t*>(data_ptr(cos)),
+      reinterpret_cast<const bf16_t*>(data_ptr(sin)),
+      reinterpret_cast<const bf16_t*>(data_ptr(weight)),
+      reinterpret_cast<bf16_t*>(mutable_data_ptr(out)),
+      eps,
+      seq_len,
+      num_heads,
+      head_dim,
+      stride_cos_b,
+      stride_cos_h,
+      stride_cos_t,
+      stride_sin_b,
+      stride_sin_h,
+      stride_sin_t);
+}
+
+}  // namespace
+
+struct LTX2QKNormSplitRopeKernel {
+  static void
+  run(tvm::ffi::TensorView q_out,
+      tvm::ffi::TensorView k_out,
+      tvm::ffi::TensorView q,
+      tvm::ffi::TensorView q_cos,
+      tvm::ffi::TensorView q_sin,
+      tvm::ffi::TensorView q_weight,
+      tvm::ffi::TensorView k,
+      tvm::ffi::TensorView k_cos,
+      tvm::ffi::TensorView k_sin,
+      tvm::ffi::TensorView k_weight,
+      double eps,
+      int64_t num_heads,
+      int64_t head_dim) {
+    using namespace host;
+
+    RuntimeCheck(num_heads > 0, "num_heads must be positive");
+    RuntimeCheck(head_dim > 0, "head_dim must be positive");
+    RuntimeCheck(head_dim % 2 == 0, "head_dim must be even");
+    const int64_t hidden_size = num_heads * head_dim;
+    RuntimeCheck(hidden_size % 4 == 0, "hidden size must be divisible by 4");
+
+    auto batch = SymbolicSize{"batch"};
+    auto q_seq_len = SymbolicSize{"q_seq_len"};
+    auto k_seq_len = SymbolicSize{"k_seq_len"};
+    auto heads = SymbolicSize{"num_heads"};
+    auto half_dim = SymbolicSize{"half_dim"};
+    auto device = SymbolicDevice{};
+    heads.set_value(num_heads);
+    half_dim.set_value(head_dim / 2);
+    device.set_options<kDLCUDA>();
+
+    TensorMatcher({batch, q_seq_len, hidden_size}).with_dtype<bf16_t>().with_device(device).verify(q).verify(q_out);
+    TensorMatcher({batch, k_seq_len, hidden_size}).with_dtype<bf16_t>().with_device(device).verify(k).verify(k_out);
+    TensorMatcher({hidden_size}).with_dtype<bf16_t>().with_device(device).verify(q_weight);
+    TensorMatcher({hidden_size}).with_dtype<bf16_t>().with_device(device).verify(k_weight);
+    TensorMatcher({batch, heads, q_seq_len, half_dim})
+        .with_strides({-1, -1, -1, 1})
+        .with_dtype<bf16_t>()
+        .with_device(device)
+        .verify(q_cos);
+    TensorMatcher({batch, heads, q_seq_len, half_dim})
+        .with_strides({-1, -1, -1, 1})
+        .with_dtype<bf16_t>()
+        .with_device(device)
+        .verify(q_sin);
+    TensorMatcher({batch, heads, k_seq_len, half_dim})
+        .with_strides({-1, -1, -1, 1})
+        .with_dtype<bf16_t>()
+        .with_device(device)
+        .verify(k_cos);
+    TensorMatcher({batch, heads, k_seq_len, half_dim})
+        .with_strides({-1, -1, -1, 1})
+        .with_dtype<bf16_t>()
+        .with_device(device)
+        .verify(k_sin);
+
+    const int64_t batch_size = batch.unwrap();
+    const DLDevice dl_device = device.unwrap();
+    launch_one(
+        q,
+        q_cos,
+        q_sin,
+        q_weight,
+        q_out,
+        static_cast<float>(eps),
+        batch_size * q_seq_len.unwrap(),
+        q_seq_len.unwrap(),
+        num_heads,
+        head_dim,
+        q_cos.stride(0),
+        q_cos.stride(1),
+        q_cos.stride(2),
+        q_sin.stride(0),
+        q_sin.stride(1),
+        q_sin.stride(2),
+        dl_device);
+    launch_one(
+        k,
+        k_cos,
+        k_sin,
+        k_weight,
+        k_out,
+        static_cast<float>(eps),
+        batch_size * k_seq_len.unwrap(),
+        k_seq_len.unwrap(),
+        num_heads,
+        head_dim,
+        k_cos.stride(0),
+        k_cos.stride(1),
+        k_cos.stride(2),
+        k_sin.stride(0),
+        k_sin.stride(1),
+        k_sin.stride(2),
+        dl_device);
+  }
+};
+
+}  // namespace sglang_ltx2_qknorm_split_rope

python/sglang/jit_kernel/csrc/diffusion/norm_scale_shift.cuh

@@ -0,0 +1,216 @@
+// Minimal native-CUDA fast path for Qwen-Image diffusion norm-scale-shift.
+//
+// Supported shape family:
+//   - bf16 activations, B == 1, hidden dim == 3072
+//   - layer norm only, no affine weight/bias
+//   - scale/shift are bf16 row-broadcast tensors ([D], [1,D], or [1,1,D])
+//   - optional residual path uses a bf16 row-broadcast gate
+//
+// All other public-op inputs fall back to the existing CuTe-DSL implementation
+// from the Python dispatcher.
+//
+// Developed with MIT HAN Lab Kernel Design Agents:
+// https://github.com/mit-han-lab/kernel-design-agents
+
+#pragma once
+
+#include <sgl_kernel/tensor.h>  // For TensorMatcher, SymbolicSize, SymbolicDevice
+
+#include <sgl_kernel/math.cuh>   // For device::math::rsqrt
+#include <sgl_kernel/utils.cuh>  // For SGL_DEVICE, bf16_t, LaunchKernel
+#include <sgl_kernel/vec.cuh>    // For AlignedVector
+#include <sgl_kernel/warp.cuh>   // For warp::reduce_sum
+
+#include <cstdint>
+
+namespace sglang_norm_scale_shift {
+
+namespace {
+
+constexpr int kHidden = 3072;
+constexpr int kVecElems = 16;  // 32B/thread for bf16 on Blackwell.
+constexpr int kThreads = kHidden / kVecElems;
+constexpr int kWarps = kThreads / device::kWarpThreads;
+constexpr float kInvHidden = 1.0f / float(kHidden);
+
+static_assert(kThreads == 192);
+static_assert(kWarps == 6);
+
+struct QwenImageNormParams {
+  void* y;
+  void* res_out;
+  const void* x;
+  const void* residual;
+  const void* gate;
+  const void* scale;
+  const void* shift;
+  float eps;
+};
+
+SGL_DEVICE float cta_reduce_sum(float v, int warp, int lane, float* scratch) {
+  v = device::warp::reduce_sum(v);
+  if (lane == 0) {
+    scratch[warp] = v;
+  }
+  __syncthreads();
+
+  if (warp == 0) {
+    float a = lane < kWarps ? scratch[lane] : 0.0f;
+    a = device::warp::reduce_sum(a);
+    if (lane == 0) {
+      scratch[kWarps] = a;
+    }
+  }
+  __syncthreads();
+  return scratch[kWarps];
+}
+
+template <bool kHasResidual>
+__global__ void qwen_image_norm_scale_shift_kernel(const QwenImageNormParams __grid_constant__ params) {
+  using namespace device;
+  using Vec = AlignedVector<bf16_t, kVecElems>;
+
+  const int row = blockIdx.x;
+  const int tid = threadIdx.x;
+  const int lane = tid & int(kWarpThreads - 1);
+  const int warp = tid >> 5;
+  const int row_offset = row * kHidden;
+  const int elem_offset = tid * kVecElems;
+
+  __shared__ float scratch_a[kWarps + 1];
+  __shared__ float scratch_b[kWarps + 1];
+
+  Vec xv;
+  xv.load(static_cast<const bf16_t*>(params.x) + row_offset + elem_offset);
+
+  float v[kVecElems];
+#pragma unroll
+  for (int i = 0; i < kVecElems; ++i) {
+    v[i] = static_cast<float>(xv[i]);
+  }
+
+  if constexpr (kHasResidual) {
+    Vec gv;
+    Vec rv;
+    Vec ro;
+    gv.load(static_cast<const bf16_t*>(params.gate) + elem_offset);
+    rv.load(static_cast<const bf16_t*>(params.residual) + row_offset + elem_offset);
+
+#pragma unroll
+    for (int i = 0; i < kVecElems; ++i) {
+      const bf16_t rounded = static_cast<bf16_t>(v[i] * static_cast<float>(gv[i]) + static_cast<float>(rv[i]));
+      ro[i] = rounded;
+      v[i] = static_cast<float>(rounded);
+    }
+    ro.store(static_cast<bf16_t*>(params.res_out) + row_offset + elem_offset);
+  }
+
+  float sum = 0.0f;
+#pragma unroll
+  for (int i = 0; i < kVecElems; ++i) {
+    sum += v[i];
+  }
+  const float mean = cta_reduce_sum(sum, warp, lane, scratch_a) * kInvHidden;
+
+  float var_sum = 0.0f;
+#pragma unroll
+  for (int i = 0; i < kVecElems; ++i) {
+    const float d = v[i] - mean;
+    var_sum += d * d;
+  }
+  const float var = cta_reduce_sum(var_sum, warp, lane, scratch_b) * kInvHidden;
+  const float factor = math::rsqrt(var + params.eps);
+
+  Vec scv;
+  Vec shv;
+  Vec yv;
+  scv.load(static_cast<const bf16_t*>(params.scale) + elem_offset);
+  shv.load(static_cast<const bf16_t*>(params.shift) + elem_offset);
+
+#pragma unroll
+  for (int i = 0; i < kVecElems; ++i) {
+    const float norm = static_cast<float>(static_cast<bf16_t>((v[i] - mean) * factor));
+    yv[i] = static_cast<bf16_t>(norm * (1.0f + static_cast<float>(scv[i])) + static_cast<float>(shv[i]));
+  }
+  yv.store(static_cast<bf16_t*>(params.y) + row_offset + elem_offset);
+}
+
+inline uint32_t verify_qwen_geometry(host::SymbolicSize& num_rows) {
+  using namespace host;
+  RuntimeCheck(num_rows.unwrap() > 0, "num_rows must be positive");
+  RuntimeCheck(num_rows.unwrap() <= int64_t(UINT32_MAX), "num_rows out of range");
+  return static_cast<uint32_t>(num_rows.unwrap());
+}
+
+}  // namespace
+
+struct QwenImageNormScaleShiftKernel {
+  static void
+  run(tvm::ffi::TensorView y,
+      tvm::ffi::TensorView x,
+      tvm::ffi::TensorView scale,
+      tvm::ffi::TensorView shift,
+      double eps) {
+    using namespace host;
+    auto N = SymbolicSize{"num_rows"};
+    auto device = SymbolicDevice{};
+    device.set_options<kDLCUDA>();
+
+    TensorMatcher({N, kHidden}).with_dtype<bf16_t>().with_device(device).verify(x).verify(y);
+    TensorMatcher({kHidden}).with_dtype<bf16_t>().with_device(device).verify(scale).verify(shift);
+
+    const uint32_t grid = verify_qwen_geometry(N);
+    const auto params = QwenImageNormParams{
+        .y = y.data_ptr(),
+        .res_out = nullptr,
+        .x = x.data_ptr(),
+        .residual = nullptr,
+        .gate = nullptr,
+        .scale = scale.data_ptr(),
+        .shift = shift.data_ptr(),
+        .eps = static_cast<float>(eps),
+    };
+    LaunchKernel(grid, kThreads, device.unwrap())(qwen_image_norm_scale_shift_kernel<false>, params);
+  }
+};
+
+struct QwenImageScaleResidualNormScaleShiftKernel {
+  static void
+  run(tvm::ffi::TensorView y,
+      tvm::ffi::TensorView res_out,
+      tvm::ffi::TensorView residual,
+      tvm::ffi::TensorView x,
+      tvm::ffi::TensorView gate,
+      tvm::ffi::TensorView scale,
+      tvm::ffi::TensorView shift,
+      double eps) {
+    using namespace host;
+    auto N = SymbolicSize{"num_rows"};
+    auto device = SymbolicDevice{};
+    device.set_options<kDLCUDA>();
+
+    TensorMatcher({N, kHidden})
+        .with_dtype<bf16_t>()
+        .with_device(device)
+        .verify(x)
+        .verify(residual)
+        .verify(y)
+        .verify(res_out);
+    TensorMatcher({kHidden}).with_dtype<bf16_t>().with_device(device).verify(gate).verify(scale).verify(shift);
+
+    const uint32_t grid = verify_qwen_geometry(N);
+    const auto params = QwenImageNormParams{
+        .y = y.data_ptr(),
+        .res_out = res_out.data_ptr(),
+        .x = x.data_ptr(),
+        .residual = residual.data_ptr(),
+        .gate = gate.data_ptr(),
+        .scale = scale.data_ptr(),
+        .shift = shift.data_ptr(),
+        .eps = static_cast<float>(eps),
+    };
+    LaunchKernel(grid, kThreads, device.unwrap())(qwen_image_norm_scale_shift_kernel<true>, params);
+  }
+};
+
+}  // namespace sglang_norm_scale_shift

python/sglang/jit_kernel/csrc/diffusion/qknorm_rope.cuh

@@ -0,0 +1,246 @@
+#include <sgl_kernel/tensor.h>
+
+#include <sgl_kernel/runtime.cuh>
+#include <sgl_kernel/type.cuh>
+#include <sgl_kernel/utils.cuh>
+#include <sgl_kernel/vec.cuh>
+#include <sgl_kernel/warp.cuh>
+
+#include <dlpack/dlpack.h>
+
+#include <cstdint>
+#include <type_traits>
+
+namespace {
+
+struct QKNormRopeParams {
+  void* __restrict__ q_ptr;
+  void* __restrict__ k_ptr;  // pre-offset by -num_qo_heads * head_stride_bytes
+  const void* __restrict__ q_weight_ptr;
+  const void* __restrict__ k_weight_ptr;
+  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;
+  float eps;
+};
+
+constexpr uint32_t kThreadsPerBlock = 256;
+constexpr uint32_t kWarpsPerBlock = kThreadsPerBlock / device::kWarpThreads;
+
+template <uint32_t kLaneCount>
+constexpr uint32_t active_mask() {
+  static_assert(kLaneCount <= device::kWarpThreads, "active_mask lane count must not exceed warp size");
+  if constexpr (kLaneCount == device::kWarpThreads) {
+    return 0xffffffffu;
+  } else {
+    return (1u << kLaneCount) - 1u;
+  }
+}
+
+SGL_DEVICE float load_cache_value(const float* ptr, int64_t idx) {
+#ifdef USE_ROCM
+  return ptr[idx];
+#else
+  return __ldg(ptr + idx);
+#endif
+}
+
+template <int64_t kHeadDim, int64_t kRopeDim, bool kIsNeox, bool kUsePDL, typename DType, typename IdType>
+__global__ void fused_qknorm_rope_warp(const QKNormRopeParams __grid_constant__ params) {
+  using namespace device;
+
+  static_assert(std::is_same_v<DType, fp16_t> || std::is_same_v<DType, bf16_t>);
+  static_assert(kHeadDim <= 256, "Only warp-level fused qknorm+rope is supported");
+  static_assert(kHeadDim % kWarpThreads == 0, "head_dim must be divisible by warp size");
+
+  constexpr uint32_t kElemsPerThread = kHeadDim / kWarpThreads;
+  constexpr uint32_t kVecSize = kElemsPerThread / 2;
+  constexpr uint32_t kRotaryLanes = kRopeDim / kElemsPerThread;
+  constexpr uint32_t kHalfRotaryLanes = kRotaryLanes / 2;
+  constexpr uint32_t kActiveMask = active_mask<kRotaryLanes>();
+  constexpr int64_t kCosSinStrideBytes = kRopeDim * sizeof(float);
+
+  static_assert(kElemsPerThread % 2 == 0, "Each lane must own an even number of elements");
+  static_assert(kRopeDim > 0 && kRopeDim <= kHeadDim, "Invalid rope dimension");
+  static_assert(kRopeDim % kElemsPerThread == 0, "rope_dim must align with per-lane vector width");
+  static_assert(
+      !kIsNeox || (kRotaryLanes >= 2 && ((kRotaryLanes & (kRotaryLanes - 1)) == 0)),
+      "NeoX fused qknorm+rope requires rotary lane count to be a power of 2");
+
+  using Packed = packed_t<DType>;
+  using Storage = AlignedVector<Packed, kVecSize>;
+
+  const auto& [q_ptr, k_ptr, q_weight_ptr, k_weight_ptr, cos_sin_cache_ptr, positions, q_stride_bytes, k_stride_bytes, head_stride_bytes, num_qo_heads, num_kv_heads, num_tokens, eps] =
+      params;
+
+  const uint32_t lane_id = threadIdx.x % kWarpThreads;
+  const uint32_t warp_id = threadIdx.x / kWarpThreads;
+  const uint32_t start_worker_id = blockIdx.x * kWarpsPerBlock + warp_id;
+  const uint32_t num_workers = gridDim.x * kWarpsPerBlock;
+  const uint32_t num_qk_heads = num_qo_heads + num_kv_heads;
+  const uint32_t num_works = num_qk_heads * num_tokens;
+
+  PDLWaitPrimary<kUsePDL>();
+
+  for (uint32_t idx = start_worker_id; idx < num_works; idx += num_workers) {
+    const uint32_t token_id = idx / num_qk_heads;
+    const uint32_t head_id = idx % num_qk_heads;
+    const bool load_q = head_id < num_qo_heads;
+    const void* input = load_q ? pointer::offset(q_ptr, token_id * q_stride_bytes, head_id * head_stride_bytes)
+                               : pointer::offset(k_ptr, token_id * k_stride_bytes, head_id * head_stride_bytes);
+    const void* weight_ptr = load_q ? q_weight_ptr : k_weight_ptr;
+
+    auto input_vec = load_as<Storage>(input, lane_id);
+    const auto weight_vec = load_as<Storage>(weight_ptr, lane_id);
+
+    float elems[kElemsPerThread];
+    float sum_of_squares = 0.0f;
+
+#pragma unroll
+    for (uint32_t j = 0; j < kVecSize; ++j) {
+      const auto [x0, x1] = cast<fp32x2_t>(input_vec[j]);
+      elems[2 * j] = x0;
+      elems[2 * j + 1] = x1;
+      sum_of_squares += x0 * x0 + x1 * x1;
+    }
+
+    sum_of_squares = warp::reduce_sum(sum_of_squares);
+    const float norm_factor = math::rsqrt(sum_of_squares / static_cast<float>(kHeadDim) + eps);
+
+#pragma unroll
+    for (uint32_t j = 0; j < kVecSize; ++j) {
+      const auto [w0, w1] = cast<fp32x2_t>(weight_vec[j]);
+      elems[2 * j] *= norm_factor * w0;
+      elems[2 * j + 1] *= norm_factor * w1;
+    }
+
+    if constexpr (kIsNeox) {
+      if (lane_id < kRotaryLanes) {
+        const auto pos = static_cast<int64_t>(static_cast<const IdType*>(positions)[token_id]);
+        const auto cos_ptr = static_cast<const float*>(pointer::offset(cos_sin_cache_ptr, pos * kCosSinStrideBytes));
+        const auto sin_ptr = cos_ptr + kRopeDim / 2;
+
+#pragma unroll
+        for (uint32_t i = 0; i < kElemsPerThread; ++i) {
+          float swapped = __shfl_xor_sync(kActiveMask, elems[i], kHalfRotaryLanes);
+          if (lane_id < kHalfRotaryLanes) {
+            swapped = -swapped;
+          }
+          int dim_idx = static_cast<int>(lane_id * kElemsPerThread + i);
+          dim_idx = (dim_idx * 2) % kRopeDim;
+          const int half_idx = dim_idx / 2;
+          const float cos = load_cache_value(cos_ptr, half_idx);
+          const float sin = load_cache_value(sin_ptr, half_idx);
+          elems[i] = elems[i] * cos + swapped * sin;
+        }
+      }
+    } else {
+      if (lane_id < kRotaryLanes) {
+        const auto pos = static_cast<int64_t>(static_cast<const IdType*>(positions)[token_id]);
+        const auto cos_ptr = static_cast<const float*>(pointer::offset(cos_sin_cache_ptr, pos * kCosSinStrideBytes));
+        const auto sin_ptr = cos_ptr + kRopeDim / 2;
+
+#pragma unroll
+        for (uint32_t i = 0; i < kElemsPerThread; i += 2) {
+          const float x = elems[i];
+          const float y = elems[i + 1];
+          const int half_idx = static_cast<int>(lane_id * kElemsPerThread + i) / 2;
+          const float cos = load_cache_value(cos_ptr, half_idx);
+          const float sin = load_cache_value(sin_ptr, half_idx);
+          elems[i] = x * cos - y * sin;
+          elems[i + 1] = y * cos + x * sin;
+        }
+      }
+    }
+
+#pragma unroll
+    for (uint32_t j = 0; j < kVecSize; ++j) {
+      input_vec[j] = cast<Packed, fp32x2_t>({elems[2 * j], elems[2 * j + 1]});
+    }
+    store_as<Storage>(const_cast<void*>(input), input_vec, lane_id);
+  }
+
+  PDLTriggerSecondary<kUsePDL>();
+}
+
+template <int64_t kHeadDim, int64_t kRopeDim, bool kIsNeox, bool kUsePDL, typename DType>
+struct QKNormRopeKernel {
+  static_assert(kHeadDim <= 256, "Only head_dim <= 256 is supported");
+  template <typename IdType>
+  static constexpr auto kernel = fused_qknorm_rope_warp<kHeadDim, kRopeDim, kIsNeox, kUsePDL, DType, IdType>;
+
+  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,
+      const tvm::ffi::TensorView cos_sin_cache,
+      const tvm::ffi::TensorView positions,
+      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 R = 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(kHeadDim);
+    R.set_value(kRopeDim);
+    device.set_options<kDLCUDA>();
+
+    TensorMatcher({N, Q, D}).with_strides({Dq, Dd, 1}).with_dtype<DType>().with_device(device).verify(q);
+    TensorMatcher({N, K, D}).with_strides({Dk, Dd, 1}).with_dtype<DType>().with_device(device).verify(k);
+    TensorMatcher({D}).with_dtype<DType>().with_device(device).verify(q_weight).verify(k_weight);
+    TensorMatcher({-1, R}).with_dtype<float>().with_device(device).verify(cos_sin_cache);
+    TensorMatcher({N}).with_dtype<int32_t, int64_t>(id_type).with_device(device).verify(positions);
+
+    const auto num_tokens = static_cast<uint32_t>(N.unwrap());
+    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));
+
+    const int64_t k_offset = static_cast<int64_t>(num_qo_heads) * head_stride_bytes;
+    const auto params = QKNormRopeParams{
+        .q_ptr = q.data_ptr(),
+        .k_ptr = pointer::offset(k.data_ptr(), -k_offset),
+        .q_weight_ptr = q_weight.data_ptr(),
+        .k_weight_ptr = k_weight.data_ptr(),
+        .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,
+        .eps = eps,
+    };
+
+    const auto is_int32 = id_type.is_type<int32_t>();
+    const auto selected_kernel = is_int32 ? kernel<int32_t> : kernel<int64_t>;
+    const uint32_t kNumSM = runtime::get_sm_count(device.unwrap().device_id);
+    static const uint32_t kOccupancyTable[2] = {
+        runtime::get_blocks_per_sm(kernel<int32_t>, kThreadsPerBlock),
+        runtime::get_blocks_per_sm(kernel<int64_t>, kThreadsPerBlock),
+    };
+    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, kWarpsPerBlock);
+    const auto num_blocks = std::min(max_blocks, needed_blocks);
+    LaunchKernel(num_blocks, kThreadsPerBlock, device.unwrap()).enable_pdl(kUsePDL)(selected_kernel, params);
+  }
+};
+
+}  // namespace

python/sglang/jit_kernel/csrc/diffusion/residual_gate_add.cuh

@@ -0,0 +1,317 @@
+// CUDA fast path for diffusion residual-gate elementwise updates.
+//
+// Implements:
+//   out = residual + update * gate
+//
+// The production shapes come from LTX-2.3 HQ residual/gate updates.  This is
+// intentionally narrow: contiguous residual/update/out tensors, with either a
+// full contiguous gate or a row-broadcast [1, 1, D] gate.
+//
+// Developed with MIT HAN Lab Kernel Design Agents:
+// https://github.com/mit-han-lab/kernel-design-agents
+
+#pragma once
+
+#include <sgl_kernel/tensor.h>  // For host dtype helpers and TensorView metadata
+#include <sgl_kernel/utils.h>   // For RuntimeCheck and div_ceil
+
+#include <sgl_kernel/type.cuh>   // For dtype_trait conversions
+#include <sgl_kernel/utils.cuh>  // For LaunchKernel and CUDA dtype aliases
+#include <sgl_kernel/vec.cuh>    // For device::AlignedVector
+
+#include <cstdint>
+
+namespace sglang_residual_gate_add {
+
+namespace {
+
+constexpr int kBlockSize = 256;
+constexpr int kBcastRowsPerBlock = 4;
+constexpr int kBcastColsVecPerBlock = 256;
+constexpr int64_t kMaxGrid = 65535;
+
+enum class GateMode : int { kFull = 0, kBcastRow = 1 };
+
+inline const char* data_ptr(const tvm::ffi::TensorView& t) {
+  return static_cast<const char*>(t.data_ptr()) + t.byte_offset();
+}
+
+inline char* mutable_data_ptr(const tvm::ffi::TensorView& t) {
+  return static_cast<char*>(t.data_ptr()) + t.byte_offset();
+}
+
+inline bool aligned16(const void* p) {
+  return (reinterpret_cast<uintptr_t>(p) & 0xF) == 0;
+}
+
+inline int64_t numel(const tvm::ffi::TensorView& t) {
+  int64_t n = 1;
+  for (int i = 0; i < t.ndim(); ++i) {
+    n *= t.size(i);
+  }
+  return n;
+}
+
+inline int64_t grid_for(int64_t total) {
+  int64_t grid = host::div_ceil(total, static_cast<int64_t>(kBlockSize));
+  if (grid < 1) {
+    grid = 1;
+  }
+  if (grid > kMaxGrid) {
+    grid = kMaxGrid;
+  }
+  return grid;
+}
+
+inline bool is_dense_contiguous(const tvm::ffi::TensorView& t) {
+  int64_t expected = 1;
+  for (int i = t.ndim() - 1; i >= 0; --i) {
+    if (t.size(i) == 1) {
+      continue;
+    }
+    if (t.stride(i) != expected) {
+      return false;
+    }
+    expected *= t.size(i);
+  }
+  return true;
+}
+
+template <typename T>
+inline void check_dtype(const tvm::ffi::TensorView& t) {
+  host::RuntimeCheck(host::is_type<T>(t.dtype()), "unexpected dtype for residual_gate_add");
+}
+
+template <typename T>
+__device__ __forceinline__ float to_float(T v) {
+  return static_cast<float>(v);
+}
+
+template <>
+__device__ __forceinline__ float to_float<fp16_t>(fp16_t v) {
+  return __half2float(v);
+}
+
+template <>
+__device__ __forceinline__ float to_float<bf16_t>(bf16_t v) {
+  return __bfloat162float(v);
+}
+
+template <typename T>
+__device__ __forceinline__ T residual_gate_value(T residual, T update, T gate) {
+  const T product = dtype_trait<T>::from(to_float(update) * to_float(gate));
+  return dtype_trait<T>::from(to_float(residual) + to_float(product));
+}
+
+template <typename T, int kVec>
+__global__ void residual_gate_add_vec_kernel(
+    const T* __restrict__ residual,
+    const T* __restrict__ update,
+    const T* __restrict__ gate,
+    T* __restrict__ out,
+    int64_t n_vec) {
+  using Vec = device::AlignedVector<T, kVec>;
+  const int64_t stride = static_cast<int64_t>(gridDim.x) * blockDim.x;
+  for (int64_t v = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x; v < n_vec; v += stride) {
+    Vec r, u, g, o;
+    r.load(residual, v);
+    u.load(update, v);
+    g.load(gate, v);
+#pragma unroll
+    for (int i = 0; i < kVec; ++i) {
+      o[i] = residual_gate_value(r[i], u[i], g[i]);
+    }
+    o.store(out, v);
+  }
+}
+
+template <typename T, int kVec>
+__global__ void residual_gate_add_bcast_row_tile_kernel(
+    const T* __restrict__ residual,
+    const T* __restrict__ update,
+    const T* __restrict__ gate,
+    T* __restrict__ out,
+    int64_t rows,
+    int64_t row_vec) {
+  using Vec = device::AlignedVector<T, kVec>;
+  const int64_t col_vec = static_cast<int64_t>(blockIdx.x) * kBcastColsVecPerBlock + threadIdx.x;
+  if (col_vec >= row_vec) {
+    return;
+  }
+
+  Vec g;
+  g.load(gate, col_vec);
+
+  // Grid-stride over row tiles so the launch stays valid even when the number
+  // of row tiles exceeds the gridDim.y hardware limit.
+  const int64_t row_tile_stride = static_cast<int64_t>(gridDim.y) * kBcastRowsPerBlock;
+  for (int64_t row_base = static_cast<int64_t>(blockIdx.y) * kBcastRowsPerBlock; row_base < rows;
+       row_base += row_tile_stride) {
+#pragma unroll
+    for (int row_offset = 0; row_offset < kBcastRowsPerBlock; ++row_offset) {
+      const int64_t row = row_base + row_offset;
+      if (row < rows) {
+        const int64_t v = row * row_vec + col_vec;
+        Vec r, u, o;
+        r.load(residual, v);
+        u.load(update, v);
+#pragma unroll
+        for (int i = 0; i < kVec; ++i) {
+          o[i] = residual_gate_value(r[i], u[i], g[i]);
+        }
+        o.store(out, v);
+      }
+    }
+  }
+}
+
+template <typename T, GateMode kGate>
+__global__ void residual_gate_add_scalar_kernel(
+    const T* __restrict__ residual,
+    const T* __restrict__ update,
+    const T* __restrict__ gate,
+    T* __restrict__ out,
+    int64_t begin,
+    int64_t total,
+    int64_t D) {
+  const int64_t stride = static_cast<int64_t>(gridDim.x) * blockDim.x;
+  for (int64_t i = begin + static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x; i < total; i += stride) {
+    const T gate_value = kGate == GateMode::kFull ? gate[i] : SGLANG_LDG(gate + (i % D));
+    out[i] = residual_gate_value(residual[i], update[i], gate_value);
+  }
+}
+
+template <typename T>
+inline void launch_residual_gate_add(
+    const tvm::ffi::TensorView& out,
+    const tvm::ffi::TensorView& residual,
+    const tvm::ffi::TensorView& update,
+    const tvm::ffi::TensorView& gate,
+    GateMode mode) {
+  const int64_t total = numel(residual);
+  if (total == 0) {
+    return;
+  }
+
+  const int64_t D = residual.size(residual.ndim() - 1);
+  const T* residual_ptr = reinterpret_cast<const T*>(data_ptr(residual));
+  const T* update_ptr = reinterpret_cast<const T*>(data_ptr(update));
+  const T* gate_ptr = reinterpret_cast<const T*>(data_ptr(gate));
+  T* out_ptr = reinterpret_cast<T*>(mutable_data_ptr(out));
+  constexpr int kVec = 16 / sizeof(T);
+
+  const bool vec_ok = aligned16(residual_ptr) && aligned16(update_ptr) && aligned16(gate_ptr) && aligned16(out_ptr) &&
+                      (D % kVec == 0) && (mode == GateMode::kBcastRow || total % kVec == 0);
+
+  int64_t done = 0;
+  if (vec_ok) {
+    const int64_t n_vec = total / kVec;
+    const int64_t row_vec = D / kVec;
+    if (mode == GateMode::kFull) {
+      host::LaunchKernel(static_cast<uint32_t>(grid_for(n_vec)), kBlockSize, out.device())(
+          residual_gate_add_vec_kernel<T, kVec>, residual_ptr, update_ptr, gate_ptr, out_ptr, n_vec);
+    } else {
+      const int64_t rows = total / D;
+      const int64_t col_blocks = host::div_ceil(row_vec, static_cast<int64_t>(kBcastColsVecPerBlock));
+      const int64_t row_tiles = host::div_ceil(rows, static_cast<int64_t>(kBcastRowsPerBlock));
+      const int64_t row_blocks = row_tiles > kMaxGrid ? kMaxGrid : row_tiles;
+      host::LaunchKernel(
+          dim3(static_cast<uint32_t>(col_blocks), static_cast<uint32_t>(row_blocks)),
+          dim3(kBcastColsVecPerBlock),
+          out.device())(
+          residual_gate_add_bcast_row_tile_kernel<T, kVec>, residual_ptr, update_ptr, gate_ptr, out_ptr, rows, row_vec);
+    }
+    done = n_vec * kVec;
+  }
+
+  if (done < total) {
+    if (mode == GateMode::kFull) {
+      host::LaunchKernel(static_cast<uint32_t>(grid_for(total - done)), kBlockSize, out.device())(
+          residual_gate_add_scalar_kernel<T, GateMode::kFull>,
+          residual_ptr,
+          update_ptr,
+          gate_ptr,
+          out_ptr,
+          done,
+          total,
+          D);
+    } else {
+      host::LaunchKernel(static_cast<uint32_t>(grid_for(total - done)), kBlockSize, out.device())(
+          residual_gate_add_scalar_kernel<T, GateMode::kBcastRow>,
+          residual_ptr,
+          update_ptr,
+          gate_ptr,
+          out_ptr,
+          done,
+          total,
+          D);
+    }
+  }
+}
+
+template <typename T>
+inline GateMode validate_residual_gate_add(
+    const tvm::ffi::TensorView& out,
+    const tvm::ffi::TensorView& residual,
+    const tvm::ffi::TensorView& update,
+    const tvm::ffi::TensorView& gate) {
+  check_dtype<T>(out);
+  check_dtype<T>(residual);
+  check_dtype<T>(update);
+  check_dtype<T>(gate);
+  host::RuntimeCheck(residual.device().device_type == kDLCUDA, "residual must be CUDA");
+  host::RuntimeCheck(update.device().device_type == kDLCUDA, "update must be CUDA");
+  host::RuntimeCheck(gate.device().device_type == kDLCUDA, "gate must be CUDA");
+  host::RuntimeCheck(out.device().device_type == kDLCUDA, "out must be CUDA");
+  host::RuntimeCheck(
+      residual.device().device_id == update.device().device_id &&
+          residual.device().device_id == gate.device().device_id &&
+          residual.device().device_id == out.device().device_id,
+      "residual/update/gate/out must be on the same CUDA device");
+  host::RuntimeCheck(residual.ndim() >= 2, "residual must be at least 2D");
+  host::RuntimeCheck(update.ndim() == residual.ndim(), "update rank must match residual");
+  host::RuntimeCheck(out.ndim() == residual.ndim(), "out rank must match residual");
+  for (int i = 0; i < residual.ndim(); ++i) {
+    host::RuntimeCheck(update.size(i) == residual.size(i), "update shape must match residual");
+    host::RuntimeCheck(out.size(i) == residual.size(i), "out shape must match residual");
+  }
+  host::RuntimeCheck(is_dense_contiguous(residual), "residual must be contiguous");
+  host::RuntimeCheck(is_dense_contiguous(update), "update must be contiguous");
+  host::RuntimeCheck(is_dense_contiguous(out), "out must be contiguous");
+  host::RuntimeCheck(is_dense_contiguous(gate), "gate must be contiguous");
+  host::RuntimeCheck(data_ptr(out) != data_ptr(residual), "out must not alias residual");
+  host::RuntimeCheck(data_ptr(out) != data_ptr(update), "out must not alias update");
+  host::RuntimeCheck(data_ptr(out) != data_ptr(gate), "out must not alias gate");
+
+  const int D_dim = residual.ndim() - 1;
+  const int row_dim = residual.ndim() - 2;
+  host::RuntimeCheck(gate.ndim() == residual.ndim(), "gate rank must match residual");
+  host::RuntimeCheck(gate.size(D_dim) == residual.size(D_dim), "gate last dim must match residual");
+
+  bool full_gate = true;
+  for (int i = 0; i < residual.ndim(); ++i) {
+    full_gate = full_gate && gate.size(i) == residual.size(i);
+  }
+  if (full_gate) {
+    return GateMode::kFull;
+  }
+
+  host::RuntimeCheck(gate.size(row_dim) == 1, "broadcast gate row dim must be 1");
+  for (int i = 0; i < D_dim; ++i) {
+    host::RuntimeCheck(gate.size(i) == 1, "broadcast gate leading dims must be 1");
+  }
+  return GateMode::kBcastRow;
+}
+
+}  // namespace
+
+template <typename T>
+struct ResidualGateAddKernel {
+  static void
+  run(tvm::ffi::TensorView out, tvm::ffi::TensorView residual, tvm::ffi::TensorView update, tvm::ffi::TensorView gate) {
+    const GateMode mode = validate_residual_gate_add<T>(out, residual, update, gate);
+    launch_residual_gate_add<T>(out, residual, update, gate, mode);
+  }
+};
+
+}  // namespace sglang_residual_gate_add

python/sglang/jit_kernel/csrc/diffusion/timestep_embedding.cuh

@@ -0,0 +1,154 @@
+#pragma once
+
+#include <sgl_kernel/tensor.h>
+#include <sgl_kernel/utils.h>
+
+#include <sgl_kernel/math.cuh>
+#include <sgl_kernel/type.cuh>
+#include <sgl_kernel/utils.cuh>
+#include <sgl_kernel/vec.cuh>  // For device::AlignedVector
+
+#include <dlpack/dlpack.h>
+#include <tvm/ffi/container/tensor.h>
+
+#include <algorithm>
+#include <cmath>
+#include <cstdint>
+#include <cuda_runtime.h>
+#include <type_traits>
+
+namespace sglang_timestep_embedding {
+
+namespace {
+
+constexpr int kVec = 4;  // 16B float vector store
+
+template <bool kFlipSinToCos, typename TIn>
+__global__ void timestep_embedding_kernel(
+    const TIn* __restrict__ t_ptr,
+    float* __restrict__ output_ptr,
+    int dim,
+    float neg_log_max_period,
+    float scale,
+    int batch_size) {
+  using Vec = device::AlignedVector<float, kVec>;
+
+  int row_idx = static_cast<int>(blockIdx.x * blockDim.y + threadIdx.y);
+  if (row_idx >= batch_size) {
+    return;
+  }
+
+  float t_val = device::cast<float>(t_ptr[row_idx]);
+  float* output_batch_base_ptr = output_ptr + row_idx * dim;
+
+  int half_dim = dim / 2;
+  int thread_offset = static_cast<int>(threadIdx.x);
+  while (thread_offset * kVec < half_dim) {
+    // !flip: output is [sin | cos]; flip: output is [cos | sin].
+    float* cos_dst;
+    float* sin_dst;
+    if constexpr (!kFlipSinToCos) {
+      sin_dst = output_batch_base_ptr + thread_offset * kVec;
+      cos_dst = output_batch_base_ptr + half_dim + thread_offset * kVec;
+    } else {
+      cos_dst = output_batch_base_ptr + thread_offset * kVec;
+      sin_dst = output_batch_base_ptr + half_dim + thread_offset * kVec;
+    }
+
+    Vec cos_vec;
+    Vec sin_vec;
+#pragma unroll
+    for (int i = 0; i < kVec; ++i) {
+      const float angle =
+          scale * t_val * device::math::exp(neg_log_max_period * __int2float_rn(thread_offset * kVec + i));
+      cos_vec[i] = device::math::cos(angle);
+      sin_vec[i] = device::math::sin(angle);
+    }
+    cos_vec.store(cos_dst);
+    sin_vec.store(sin_dst);
+
+    thread_offset += static_cast<int>(blockDim.x);
+  }
+}
+
+template <typename TIn>
+inline void launch_timestep_embedding(
+    const tvm::ffi::TensorView t,
+    const tvm::ffi::TensorView output,
+    int dim,
+    bool flip_sin_to_cos,
+    float downscale_freq_shift,
+    float scale,
+    int max_period) {
+  using namespace host;
+
+  const int batch_size = static_cast<int>(t.shape()[0]);
+  const int half_dim = dim / 2;
+
+  constexpr int kMaxThreadsPerBlock = 1024;
+  constexpr int kMinThreadsPerBlock = 128;
+
+  const int num_threads_per_row = std::min(kMaxThreadsPerBlock, half_dim / 4);
+  const int num_rows = (kMinThreadsPerBlock + num_threads_per_row - 1) / num_threads_per_row;
+
+  dim3 grid((batch_size + num_rows - 1) / num_rows);
+  dim3 block(num_threads_per_row, num_rows);
+
+  const float neg_log_max_period =
+      std::log(static_cast<float>(max_period)) * (-1.0f) / (static_cast<float>(half_dim) - downscale_freq_shift);
+
+  const DLDevice device = output.device();
+
+  if (flip_sin_to_cos) {
+    LaunchKernel(grid, block, device)(
+        timestep_embedding_kernel<true, TIn>,
+        static_cast<const TIn*>(t.data_ptr()),
+        static_cast<float*>(output.data_ptr()),
+        dim,
+        neg_log_max_period,
+        scale,
+        batch_size);
+  } else {
+    LaunchKernel(grid, block, device)(
+        timestep_embedding_kernel<false, TIn>,
+        static_cast<const TIn*>(t.data_ptr()),
+        static_cast<float*>(output.data_ptr()),
+        dim,
+        neg_log_max_period,
+        scale,
+        batch_size);
+  }
+}
+
+}  // namespace
+
+template <typename TIn>
+void timestep_embedding(
+    tvm::ffi::TensorView input,
+    tvm::ffi::TensorView output,
+    int dim,
+    bool flip_sin_to_cos,
+    float downscale_freq_shift,
+    float scale,
+    int max_period) {
+  using namespace host;
+
+  auto B = SymbolicSize{"batch_size"};
+  auto D = SymbolicSize{"dim"};
+  auto device = SymbolicDevice{};
+
+  TensorMatcher({B})  // input
+      .with_strides({1})
+      .with_dtype<TIn>()
+      .template with_device<kDLCUDA>(device)
+      .verify(input);
+
+  TensorMatcher({B, D}).with_strides({D, 1}).with_dtype<float>().template with_device<kDLCUDA>(device).verify(output);
+
+  RuntimeCheck(D.unwrap() == dim, "Output dim mismatch: ", D.unwrap(), " vs ", dim);
+  RuntimeCheck(dim % 8 == 0, "dim must align to 8, got ", dim);
+
+  launch_timestep_embedding<TIn>(input, output, dim, flip_sin_to_cos, downscale_freq_shift, scale, max_period);
+}
+
+}  // namespace sglang_timestep_embedding