[diffusion] Clean up kernels and shared fast paths (#34085)

python/sglang/kernels/jit/csrc/diffusion/modulate_scale_shift.cuh

@@ -1,108 +1,45 @@
-// CUDA fast path for diffusion adaLN modulate chains.
+// CUDA fast path for diffusion adaLN modulation.
 //
-// Implements, with each intermediate computed in fp32 and rounded to the
+// Reproduces the eager storage-dtype rounding boundaries:
-// storage dtype (the per-op kernel boundaries of the eager aten chain):
+//   out = round(round(x * round(1 + scale)) + shift)
-//
-//   out = (x * (1 + scale)) + shift
-//       = round(round(x * round(1 + scale)) + shift)
-//
-// so the fused kernel is bit-exact vs eager for fp16/bf16.  x is a
-// contiguous [B, L, D] activation; scale/shift are contiguous [B, D]
-// modulation rows.
-//
-// Intentionally narrow: 16-byte aligned tensors, D % kVec == 0 (the Python
-// guard enforces this).
 
 #pragma once
 
-#include <sgl_kernel/tensor.h>  // For host dtype helpers and TensorView metadata
+#include <sgl_kernel/tensor.h>
-#include <sgl_kernel/utils.h>   // For RuntimeCheck and div_ceil
+#include <sgl_kernel/utils.h>
 
-#include <sgl_kernel/type.cuh>   // For DTypeTrait conversions
+#include <sgl_kernel/type.cuh>
-#include <sgl_kernel/utils.cuh>  // For LaunchKernel and CUDA dtype aliases
+#include <sgl_kernel/utils.cuh>
-#include <sgl_kernel/vec.cuh>    // For device::AlignedVector
+#include <sgl_kernel/vec.cuh>
 
+#include <algorithm>
 #include <cstdint>
+#include <type_traits>
 
 namespace sglang {
 
-namespace sglang_modulate_scale_shift {
+namespace modulate_scale_shift {
 
 namespace {
 
-constexpr int kRowsPerBlock = 4;
+constexpr uint32_t kRowsPerBlock = 4;
-constexpr int kColsVecPerBlock = 256;
+constexpr uint32_t kColsVecPerBlock = 256;
-constexpr int64_t kMaxGrid = 65535;
+constexpr uint32_t kMaxGridY = 65535;
-
+constexpr uintptr_t kAlignment = 16;
-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 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) {
+SGL_DEVICE T modulate_value(T x, T scale, T shift) {
-  host::RuntimeCheck(host::is_type<T>(t.dtype()), "unexpected dtype for modulate_scale_shift");
+  const T one_plus_scale = device::cast<T>(1.0f + device::cast<fp32_t>(scale));
-}
+  const T product = device::cast<T>(device::cast<fp32_t>(x) * device::cast<fp32_t>(one_plus_scale));
-
+  return device::cast<T>(device::cast<fp32_t>(product) + device::cast<fp32_t>(shift));
-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 modulate_value(T x, T scale, T shift) {
-  // Round each intermediate back to T (the eager chain's kernel boundaries;
-  // also blocks fmul+fadd FMA contraction).
-  const T one_plus_scale = DTypeTrait<T>::from(1.0f + to_float(scale));
-  const T product = DTypeTrait<T>::from(to_float(x) * to_float(one_plus_scale));
-  return DTypeTrait<T>::from(to_float(product) + to_float(shift));
 }
 
 template <typename T, int kVec>
-__global__ void modulate_scale_shift_vec_kernel(
+__global__ void modulate_scale_shift_kernel(
+    T* __restrict__ out,
     const T* __restrict__ x,
     const T* __restrict__ scale,
     const T* __restrict__ shift,
-    T* __restrict__ out,
     int64_t rows,
     int64_t rows_per_batch,
     int64_t row_vec) {
@@ -112,114 +49,93 @@ __global__ void modulate_scale_shift_vec_kernel(
     return;
   }
 
-  // Grid-stride: the row-tile count can exceed the gridDim.y hardware limit.
+  const int64_t row_stride = static_cast<int64_t>(gridDim.y) * kRowsPerBlock;
-  const int64_t row_tile_stride = static_cast<int64_t>(gridDim.y) * kRowsPerBlock;
+  for (int64_t row_base = static_cast<int64_t>(blockIdx.y) * kRowsPerBlock; row_base < rows; row_base += row_stride) {
-  for (int64_t row_base = static_cast<int64_t>(blockIdx.y) * kRowsPerBlock; row_base < rows;
-       row_base += row_tile_stride) {
 #pragma unroll
-    for (int row_offset = 0; row_offset < kRowsPerBlock; ++row_offset) {
+    for (uint32_t row_offset = 0; row_offset < kRowsPerBlock; ++row_offset) {
       const int64_t row = row_base + row_offset;
-      if (row < rows) {
+      if (row >= rows) {
+        continue;
+      }
+
       const int64_t batch = row / rows_per_batch;
-        const int64_t mod_v = batch * row_vec + col_vec;
+      const int64_t modulation_offset = batch * row_vec + col_vec;
-        const int64_t v = row * row_vec + col_vec;
+      const int64_t activation_offset = row * row_vec + col_vec;
-        Vec xv, s, b, o;
+      Vec x_vec, scale_vec, shift_vec, out_vec;
-        s.load(scale, mod_v);
+      x_vec.load(x, activation_offset);
-        b.load(shift, mod_v);
+      scale_vec.load(scale, modulation_offset);
-        xv.load(x, v);
+      shift_vec.load(shift, modulation_offset);
 #pragma unroll
       for (int i = 0; i < kVec; ++i) {
-          o[i] = modulate_value(xv[i], s[i], b[i]);
+        out_vec[i] = modulate_value(x_vec[i], scale_vec[i], shift_vec[i]);
       }
-        o.store(out, v);
+      out_vec.store(out, activation_offset);
     }
   }
-  }
-}
-
-template <typename T>
-inline void launch_modulate_scale_shift(
-    const tvm::ffi::TensorView& out,
-    const tvm::ffi::TensorView& x,
-    const tvm::ffi::TensorView& scale,
-    const tvm::ffi::TensorView& shift) {
-  const int64_t total = numel(x);
-  if (total == 0) {
-    return;
-  }
-
-  const int64_t D = x.size(x.ndim() - 1);
-  const int64_t rows = total / D;
-  const int64_t batches = scale.size(0);
-  const int64_t rows_per_batch = rows / batches;
-  const T* x_ptr = reinterpret_cast<const T*>(data_ptr(x));
-  const T* scale_ptr = reinterpret_cast<const T*>(data_ptr(scale));
-  const T* shift_ptr = reinterpret_cast<const T*>(data_ptr(shift));
-  T* out_ptr = reinterpret_cast<T*>(mutable_data_ptr(out));
-  constexpr int kVec = 16 / sizeof(T);
-
-  host::RuntimeCheck(
-      aligned16(x_ptr) && aligned16(scale_ptr) && aligned16(shift_ptr) && aligned16(out_ptr),
-      "modulate_scale_shift requires 16-byte aligned tensors");
-  host::RuntimeCheck(D % kVec == 0, "modulate_scale_shift requires D to be a multiple of the vector width");
-
-  const int64_t row_vec = D / kVec;
-  const int64_t col_blocks = host::div_ceil(row_vec, static_cast<int64_t>(kColsVecPerBlock));
-  const int64_t row_tiles = host::div_ceil(rows, static_cast<int64_t>(kRowsPerBlock));
-  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(kColsVecPerBlock), out.device())(
-      modulate_scale_shift_vec_kernel<T, kVec>, x_ptr, scale_ptr, shift_ptr, out_ptr, rows, rows_per_batch, row_vec);
-}
-
-template <typename T>
-inline void validate_modulate_scale_shift(
-    const tvm::ffi::TensorView& out,
-    const tvm::ffi::TensorView& x,
-    const tvm::ffi::TensorView& scale,
-    const tvm::ffi::TensorView& shift) {
-  check_dtype<T>(out);
-  check_dtype<T>(x);
-  check_dtype<T>(scale);
-  check_dtype<T>(shift);
-  host::RuntimeCheck(x.device().device_type == kDLCUDA, "x must be CUDA");
-  host::RuntimeCheck(scale.device().device_type == kDLCUDA, "scale must be CUDA");
-  host::RuntimeCheck(shift.device().device_type == kDLCUDA, "shift must be CUDA");
-  host::RuntimeCheck(out.device().device_type == kDLCUDA, "out must be CUDA");
-  host::RuntimeCheck(
-      x.device().device_id == scale.device().device_id && x.device().device_id == shift.device().device_id &&
-          x.device().device_id == out.device().device_id,
-      "x/scale/shift/out must be on the same CUDA device");
-  host::RuntimeCheck(x.ndim() == 3, "x must be [B, L, D]");
-  host::RuntimeCheck(scale.ndim() == 2, "scale must be [B, D]");
-  host::RuntimeCheck(shift.ndim() == 2, "shift must be [B, D]");
-  host::RuntimeCheck(out.ndim() == x.ndim(), "out rank must match x");
-  for (int i = 0; i < x.ndim(); ++i) {
-    host::RuntimeCheck(out.size(i) == x.size(i), "out shape must match x");
-  }
-  host::RuntimeCheck(scale.size(0) == x.size(0), "scale batch dim must match x");
-  host::RuntimeCheck(scale.size(1) == x.size(2), "scale last dim must match x");
-  host::RuntimeCheck(shift.size(0) == scale.size(0) && shift.size(1) == scale.size(1), "shift shape must match scale");
-  host::RuntimeCheck(is_dense_contiguous(x), "x must be contiguous");
-  host::RuntimeCheck(is_dense_contiguous(scale), "scale must be contiguous");
-  host::RuntimeCheck(is_dense_contiguous(shift), "shift must be contiguous");
-  host::RuntimeCheck(is_dense_contiguous(out), "out must be contiguous");
-  host::RuntimeCheck(data_ptr(out) != data_ptr(x), "out must not alias x");
-  host::RuntimeCheck(data_ptr(out) != data_ptr(scale), "out must not alias scale");
-  host::RuntimeCheck(data_ptr(out) != data_ptr(shift), "out must not alias shift");
 }
 
 }  // namespace
 
+/**
+ * \brief Validate and launch bit-exact diffusion adaLN modulation.
+ *
+ * \tparam T Activation type: fp16_t or bf16_t.
+ */
 template <typename T>
 struct ModulateScaleShiftKernel {
+  static_assert(std::is_same_v<T, fp16_t> || std::is_same_v<T, bf16_t>);
+
+  /**
+   * \param out Output tensor with shape [B, L, D].
+   * \param x Input tensor with shape [B, L, D].
+   * \param scale Scale tensor with shape [B, D].
+   * \param shift Shift tensor with shape [B, D].
+   */
   static void
   run(tvm::ffi::TensorView out, tvm::ffi::TensorView x, tvm::ffi::TensorView scale, tvm::ffi::TensorView shift) {
-    validate_modulate_scale_shift<T>(out, x, scale, shift);
+    using namespace host;
-    launch_modulate_scale_shift<T>(out, x, scale, shift);
+
+    auto B = SymbolicSize{"batch"};
+    auto L = SymbolicSize{"sequence_length"};
+    auto D = SymbolicSize{"hidden_size"};
+    auto device = SymbolicDevice{};
+    device.set_options<kDLCUDA>();
+
+    TensorMatcher({B, L, D}).with_dtype<T>().with_device(device).verify(out).verify(x);
+    TensorMatcher({B, D}).with_dtype<T>().with_device(device).verify(scale).verify(shift);
+
+    const int64_t batch = B.unwrap();
+    const int64_t sequence_length = L.unwrap();
+    const int64_t hidden_size = D.unwrap();
+    const int64_t rows = batch * sequence_length;
+    if (rows == 0 || hidden_size == 0) {
+      return;
+    }
+
+    constexpr int kVec = kAlignment / sizeof(T);
+    CHECK_HOST(hidden_size % kVec == 0) << "hidden size must be a multiple of " << kVec;
+
+    auto* out_ptr = static_cast<T*>(out.data_ptr());
+    const auto* x_ptr = static_cast<const T*>(x.data_ptr());
+    const auto* scale_ptr = static_cast<const T*>(scale.data_ptr());
+    const auto* shift_ptr = static_cast<const T*>(shift.data_ptr());
+    CHECK_HOST(
+        reinterpret_cast<uintptr_t>(out_ptr) % kAlignment == 0 &&
+        reinterpret_cast<uintptr_t>(x_ptr) % kAlignment == 0 &&
+        reinterpret_cast<uintptr_t>(scale_ptr) % kAlignment == 0 &&
+        reinterpret_cast<uintptr_t>(shift_ptr) % kAlignment == 0)
+        << "modulate_scale_shift requires 16-byte aligned tensors";
+    CHECK_HOST(out_ptr != x_ptr && out_ptr != scale_ptr && out_ptr != shift_ptr) << "output must not alias an input";
+
+    const int64_t row_vec = hidden_size / kVec;
+    const auto col_blocks = static_cast<uint32_t>(div_ceil(row_vec, static_cast<int64_t>(kColsVecPerBlock)));
+    const int64_t row_tiles = div_ceil(rows, static_cast<int64_t>(kRowsPerBlock));
+    const auto row_blocks = static_cast<uint32_t>(std::min<int64_t>(row_tiles, kMaxGridY));
+    LaunchKernel(dim3(col_blocks, row_blocks), kColsVecPerBlock, device.unwrap())(
+        modulate_scale_shift_kernel<T, kVec>, out_ptr, x_ptr, scale_ptr, shift_ptr, rows, sequence_length, row_vec);
   }
 };
 
-}  // namespace sglang_modulate_scale_shift
+}  // namespace modulate_scale_shift
 
 }  // namespace sglang

python/sglang/kernels/jit/csrc/diffusion/qknorm_rope.cuh

@@ -277,6 +277,7 @@ struct QKNormRopeKernel {
     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());
+    if (num_tokens == 0 || (num_qo_heads == 0 && num_kv_heads == 0)) return;
     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));

python/sglang/kernels/jit/csrc/diffusion/residual_gate_add.cuh

@@ -1,315 +1,206 @@
-// CUDA fast path for diffusion residual-gate elementwise updates.
+// CUDA fast path for bit-exact diffusion residual-gate 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/tensor.h>
-#include <sgl_kernel/utils.h>   // For RuntimeCheck and div_ceil
+#include <sgl_kernel/utils.h>
 
-#include <sgl_kernel/type.cuh>   // For DTypeTrait conversions
+#include <sgl_kernel/type.cuh>
-#include <sgl_kernel/utils.cuh>  // For LaunchKernel and CUDA dtype aliases
+#include <sgl_kernel/utils.cuh>
-#include <sgl_kernel/vec.cuh>    // For device::AlignedVector
+#include <sgl_kernel/vec.cuh>
 
+#include <algorithm>
 #include <cstdint>
+#include <type_traits>
 
 namespace sglang {
 
 namespace residual_gate_add {
 
-constexpr int kBlockSize = 256;
+namespace {
-constexpr int kBcastRowsPerBlock = 4;
-constexpr int kBcastColsVecPerBlock = 256;
-constexpr int64_t kMaxGrid = 65535;
 
-enum class GateMode : int { kFull = 0, kBcastRow = 1 };
+constexpr uint32_t kBlockSize = 256;
+constexpr uint32_t kBroadcastRowsPerBlock = 4;
+constexpr uint32_t kBroadcastColsPerBlock = 256;
+constexpr uint32_t kMaxGrid = 65535;
+constexpr uintptr_t kAlignment = 16;
 
-inline const char* data_ptr(const tvm::ffi::TensorView& t) {
+enum class GateMode : int { kFull, kBroadcastRow };
-  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) {
+SGL_DEVICE T residual_gate_value(T residual, T update, T gate) {
-  host::RuntimeCheck(host::is_type<T>(t.dtype()), "unexpected dtype for residual_gate_add");
+  const T product = device::cast<T>(device::cast<fp32_t>(update) * device::cast<fp32_t>(gate));
-}
+  return device::cast<T>(device::cast<fp32_t>(residual) + device::cast<fp32_t>(product));
-
-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 = DTypeTrait<T>::from(to_float(update) * to_float(gate));
-  return DTypeTrait<T>::from(to_float(residual) + to_float(product));
 }
 
 template <typename T, int kVec>
 __global__ void residual_gate_add_vec_kernel(
+    T* __restrict__ out,
     const T* __restrict__ residual,
     const T* __restrict__ update,
     const T* __restrict__ gate,
-    T* __restrict__ out,
+    int64_t num_vectors) {
-    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) {
+  for (int64_t vector = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x; vector < num_vectors;
-    Vec r, u, g, o;
+       vector += stride) {
-    r.load(residual, v);
+    Vec residual_vec, update_vec, gate_vec, out_vec;
-    u.load(update, v);
+    residual_vec.load(residual, vector);
-    g.load(gate, v);
+    update_vec.load(update, vector);
+    gate_vec.load(gate, vector);
 #pragma unroll
     for (int i = 0; i < kVec; ++i) {
-      o[i] = residual_gate_value(r[i], u[i], g[i]);
+      out_vec[i] = residual_gate_value(residual_vec[i], update_vec[i], gate_vec[i]);
     }
-    o.store(out, v);
+    out_vec.store(out, vector);
   }
 }
 
 template <typename T, int kVec>
-__global__ void residual_gate_add_bcast_row_tile_kernel(
+__global__ void residual_gate_add_broadcast_kernel(
+    T* __restrict__ out,
     const T* __restrict__ residual,
     const T* __restrict__ update,
     const T* __restrict__ gate,
-    T* __restrict__ out,
     int64_t rows,
-    int64_t row_vec) {
+    int64_t row_vectors) {
   using Vec = device::AlignedVector<T, kVec>;
-  const int64_t col_vec = static_cast<int64_t>(blockIdx.x) * kBcastColsVecPerBlock + threadIdx.x;
+  const int64_t column = static_cast<int64_t>(blockIdx.x) * kBroadcastColsPerBlock + threadIdx.x;
-  if (col_vec >= row_vec) {
+  if (column >= row_vectors) {
     return;
   }
 
-  Vec g;
+  Vec gate_vec;
-  g.load(gate, col_vec);
+  gate_vec.load(gate, column);
-
+  const int64_t row_stride = static_cast<int64_t>(gridDim.y) * kBroadcastRowsPerBlock;
-  // Grid-stride over row tiles so the launch stays valid even when the number
+  for (int64_t row_base = static_cast<int64_t>(blockIdx.y) * kBroadcastRowsPerBlock; row_base < rows;
-  // of row tiles exceeds the gridDim.y hardware limit.
+       row_base += row_stride) {
-  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) {
+    for (uint32_t row_offset = 0; row_offset < kBroadcastRowsPerBlock; ++row_offset) {
       const int64_t row = row_base + row_offset;
-      if (row < rows) {
+      if (row >= rows) {
-        const int64_t v = row * row_vec + col_vec;
+        continue;
-        Vec r, u, o;
+      }
-        r.load(residual, v);
+      const int64_t vector = row * row_vectors + column;
-        u.load(update, v);
+      Vec residual_vec, update_vec, out_vec;
+      residual_vec.load(residual, vector);
+      update_vec.load(update, vector);
 #pragma unroll
       for (int i = 0; i < kVec; ++i) {
-          o[i] = residual_gate_value(r[i], u[i], g[i]);
+        out_vec[i] = residual_gate_value(residual_vec[i], update_vec[i], gate_vec[i]);
-        }
-        o.store(out, v);
       }
+      out_vec.store(out, vector);
     }
   }
 }
 
-template <typename T, GateMode kGate>
+template <typename T, GateMode kGateMode>
 __global__ void residual_gate_add_scalar_kernel(
+    T* __restrict__ out,
     const T* __restrict__ residual,
     const T* __restrict__ update,
     const T* __restrict__ gate,
-    T* __restrict__ out,
+    int64_t numel,
-    int64_t begin,
+    int64_t hidden_size) {
-    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) {
+  for (int64_t index = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x; index < numel; index += stride) {
-    const T gate_value = kGate == GateMode::kFull ? gate[i] : SGLANG_LDG(gate + (i % D));
+    const T gate_value = kGateMode == GateMode::kFull ? gate[index] : SGLANG_LDG(gate + index % hidden_size);
-    out[i] = residual_gate_value(residual[i], update[i], gate_value);
+    out[index] = residual_gate_value(residual[index], update[index], gate_value);
   }
 }
 
+}  // namespace
+
+/**
+ * \brief Validate and launch a bit-exact residual-gate update.
+ *
+ * Python flattens the tensors before dispatch. A broadcast gate contains one
+ * hidden-size row; a full gate has the same number of elements as the inputs.
+ */
 template <typename T>
-inline void launch_residual_gate_add(
+struct ResidualGateAddKernel {
-    const tvm::ffi::TensorView& out,
+  static_assert(std::is_same_v<T, fp16_t> || std::is_same_v<T, bf16_t> || std::is_same_v<T, fp32_t>);
-    const tvm::ffi::TensorView& residual,
+
-    const tvm::ffi::TensorView& update,
+  static void
-    const tvm::ffi::TensorView& gate,
+  run(tvm::ffi::TensorView out,
-    GateMode mode) {
+      tvm::ffi::TensorView residual,
-  const int64_t total = numel(residual);
+      tvm::ffi::TensorView update,
-  if (total == 0) {
+      tvm::ffi::TensorView gate,
+      int64_t hidden_size,
+      bool broadcast_gate) {
+    using namespace host;
+
+    auto N = SymbolicSize{"numel"};
+    auto G = SymbolicSize{"gate_numel"};
+    auto device = SymbolicDevice{};
+    device.set_options<kDLCUDA>();
+    TensorMatcher({N}).with_dtype<T>().with_device(device).verify(out).verify(residual).verify(update);
+    TensorMatcher({G}).with_dtype<T>().with_device(device).verify(gate);
+
+    const int64_t numel = N.unwrap();
+    CHECK_HOST(hidden_size > 0 && numel % hidden_size == 0) << "hidden size must be positive and divide the input size";
+    CHECK_HOST(G.unwrap() == (broadcast_gate ? hidden_size : numel)) << "gate size does not match its mode";
+    if (numel == 0) {
       return;
     }
 
-  const int64_t D = residual.size(residual.ndim() - 1);
+    auto* out_ptr = static_cast<T*>(out.data_ptr());
-  const T* residual_ptr = reinterpret_cast<const T*>(data_ptr(residual));
+    const auto* residual_ptr = static_cast<const T*>(residual.data_ptr());
-  const T* update_ptr = reinterpret_cast<const T*>(data_ptr(update));
+    const auto* update_ptr = static_cast<const T*>(update.data_ptr());
-  const T* gate_ptr = reinterpret_cast<const T*>(data_ptr(gate));
+    const auto* gate_ptr = static_cast<const T*>(gate.data_ptr());
-  T* out_ptr = reinterpret_cast<T*>(mutable_data_ptr(out));
+    CHECK_HOST(out_ptr != residual_ptr && out_ptr != update_ptr && out_ptr != gate_ptr)
-  constexpr int kVec = 16 / sizeof(T);
+        << "output must not alias an input";
 
-  const bool vec_ok = aligned16(residual_ptr) && aligned16(update_ptr) && aligned16(gate_ptr) && aligned16(out_ptr) &&
+    constexpr int kVec = kAlignment / sizeof(T);
-                      (D % kVec == 0) && (mode == GateMode::kBcastRow || total % kVec == 0);
+    const bool aligned = reinterpret_cast<uintptr_t>(out_ptr) % kAlignment == 0 &&
+                         reinterpret_cast<uintptr_t>(residual_ptr) % kAlignment == 0 &&
+                         reinterpret_cast<uintptr_t>(update_ptr) % kAlignment == 0 &&
+                         reinterpret_cast<uintptr_t>(gate_ptr) % kAlignment == 0;
+    const bool vectorized = aligned && hidden_size % kVec == 0;
+    if (vectorized) {
+      const int64_t num_vectors = numel / kVec;
+      if (!broadcast_gate) {
+        const auto blocks =
+            static_cast<uint32_t>(std::min<int64_t>(div_ceil(num_vectors, static_cast<int64_t>(kBlockSize)), kMaxGrid));
+        LaunchKernel(blocks, kBlockSize, device.unwrap())(
+            residual_gate_add_vec_kernel<T, kVec>, out_ptr, residual_ptr, update_ptr, gate_ptr, num_vectors);
+        return;
+      }
 
-  int64_t done = 0;
+      const int64_t rows = numel / hidden_size;
-  if (vec_ok) {
+      const int64_t row_vectors = hidden_size / kVec;
-    const int64_t n_vec = total / kVec;
+      const auto column_blocks =
-    const int64_t row_vec = D / kVec;
+          static_cast<uint32_t>(div_ceil(row_vectors, static_cast<int64_t>(kBroadcastColsPerBlock)));
-    if (mode == GateMode::kFull) {
+      const int64_t row_tiles = div_ceil(rows, static_cast<int64_t>(kBroadcastRowsPerBlock));
-      host::LaunchKernel(static_cast<uint32_t>(grid_for(n_vec)), kBlockSize, out.device())(
+      const auto row_blocks = static_cast<uint32_t>(std::min<int64_t>(row_tiles, kMaxGrid));
-          residual_gate_add_vec_kernel<T, kVec>, residual_ptr, update_ptr, gate_ptr, out_ptr, n_vec);
+      LaunchKernel(dim3(column_blocks, row_blocks), kBroadcastColsPerBlock, device.unwrap())(
+          residual_gate_add_broadcast_kernel<T, kVec>, out_ptr, residual_ptr, update_ptr, gate_ptr, rows, row_vectors);
+      return;
+    }
+
+    const auto blocks =
+        static_cast<uint32_t>(std::min<int64_t>(div_ceil(numel, static_cast<int64_t>(kBlockSize)), kMaxGrid));
+    if (broadcast_gate) {
+      LaunchKernel(blocks, kBlockSize, device.unwrap())(
+          residual_gate_add_scalar_kernel<T, GateMode::kBroadcastRow>,
+          out_ptr,
+          residual_ptr,
+          update_ptr,
+          gate_ptr,
+          numel,
+          hidden_size);
     } else {
-      const int64_t rows = total / D;
+      LaunchKernel(blocks, kBlockSize, device.unwrap())(
-      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>,
+          out_ptr,
           residual_ptr,
           update_ptr,
           gate_ptr,
-          out_ptr,
+          numel,
-          done,
+          hidden_size);
-          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;
-}
-
-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 residual_gate_add

python/sglang/kernels/jit/csrc/diffusion/usp_relayout.cuh

@@ -1,111 +1,59 @@
-// CUDA fast path for the Ulysses sequence-parallel output head merge.
+// CUDA fast path for the Ulysses sequence-parallel output-head merge:
-//
+//   [W, S, B, H, D] -> [B, S, W, H, D]
-//   usp_merge_heads:
-//     x [W, S, B, h_local, D] (contiguous, the output all-to-all result)
-//       -> out [B, S, W, h_local, D] (contiguous)
-//     Replaces `x.permute(2, 1, 0, 3, 4).contiguous()` on the head_dim=2
-//     output path of `_usp_output_all_to_all`.
-//
-// A pure copy (no arithmetic), so it is bit-exact with the eager permute by
-// construction. It exists because ATen's generic permute-copy reaches well
-// under half of HBM bandwidth on the packed-DiT shapes, while a single pass
-// with coalesced vectorized stores runs near roofline.
 
 #pragma once
 
-#include <sgl_kernel/tensor.h>  // For host dtype helpers and TensorView metadata
+#include <sgl_kernel/tensor.h>
-#include <sgl_kernel/utils.h>   // For RuntimeCheck and div_ceil
+#include <sgl_kernel/utils.h>
 
-#include <sgl_kernel/type.cuh>   // For CUDA dtype aliases
+#include <sgl_kernel/type.cuh>
-#include <sgl_kernel/utils.cuh>  // For LaunchKernel
+#include <sgl_kernel/utils.cuh>
-#include <sgl_kernel/vec.cuh>    // For device::AlignedVector
+#include <sgl_kernel/vec.cuh>
 
+#include <algorithm>
 #include <cstdint>
+#include <type_traits>
 
 namespace sglang {
 
 namespace usp_relayout {
 
-constexpr int kBlockSize = 256;
+namespace {
-constexpr int64_t kMaxGrid = 65535;
 
-inline const char* data_ptr(const tvm::ffi::TensorView& t) {
+constexpr uint32_t kBlockSize = 256;
-  return static_cast<const char*>(t.data_ptr()) + t.byte_offset();
+constexpr uint32_t kMaxGrid = 65535;
-}
+constexpr uintptr_t kAlignment = 16;
 
-inline char* mutable_data_ptr(const tvm::ffi::TensorView& t) {
+// out[b, s, w, h, d] = x[w, s, b, h, d]
-  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 usp_merge_heads tensor");
-}
-
-// out[b, s, w, h, c] = x[w, s, b, h, c]
 template <typename T, int kVec>
 __global__ void usp_merge_heads_vec_kernel(
     T* __restrict__ out,
     const T* __restrict__ x,
-    int64_t n_vec,
+    int64_t num_vectors,
-    int64_t d_vec,  // D / kVec
+    int64_t head_vectors,
-    int64_t h_local,
+    int64_t local_heads,
     int64_t batch,
-    int64_t seq,
+    int64_t sequence_length,
-    int64_t world) {
+    int64_t world_size) {
   const int64_t stride = static_cast<int64_t>(gridDim.x) * blockDim.x;
-  for (int64_t i = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x; i < n_vec; i += stride) {
+  for (int64_t index = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x; index < num_vectors;
-    int64_t rest = i;
+       index += stride) {
-    const int64_t c_vec = rest % d_vec;
+    int64_t rest = index;
-    rest /= d_vec;
+    const int64_t head_offset = rest % head_vectors;
-    const int64_t h = rest % h_local;
+    rest /= head_vectors;
-    rest /= h_local;
+    const int64_t head = rest % local_heads;
-    const int64_t w = rest % world;
+    rest /= local_heads;
-    rest /= world;
+    const int64_t rank = rest % world_size;
-    const int64_t s = rest % seq;
+    rest /= world_size;
-    const int64_t b = rest / seq;
+    const int64_t sequence = rest % sequence_length;
+    const int64_t batch_index = rest / sequence_length;
 
-    const int64_t src_vec = ((((w * seq + s) * batch + b) * h_local) + h) * d_vec + c_vec;
+    const int64_t source =
-    device::AlignedVector<T, kVec> val;
+        ((((rank * sequence_length + sequence) * batch + batch_index) * local_heads) + head) * head_vectors +
-    val.load(x, src_vec);
+        head_offset;
-    val.store(out, i);
+    device::AlignedVector<T, kVec> value;
+    value.load(x, source);
+    value.store(out, index);
   }
 }
 
@@ -113,66 +61,83 @@ template <typename T>
 __global__ void usp_merge_heads_scalar_kernel(
     T* __restrict__ out,
     const T* __restrict__ x,
-    int64_t total,
+    int64_t numel,
     int64_t head_dim,
-    int64_t h_local,
+    int64_t local_heads,
     int64_t batch,
-    int64_t seq,
+    int64_t sequence_length,
-    int64_t world) {
+    int64_t world_size) {
   const int64_t stride = static_cast<int64_t>(gridDim.x) * blockDim.x;
-  for (int64_t i = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x; i < total; i += stride) {
+  for (int64_t index = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x; index < numel; index += stride) {
-    int64_t rest = i;
+    int64_t rest = index;
-    const int64_t c = rest % head_dim;
+    const int64_t head_offset = rest % head_dim;
     rest /= head_dim;
-    const int64_t h = rest % h_local;
+    const int64_t head = rest % local_heads;
-    rest /= h_local;
+    rest /= local_heads;
-    const int64_t w = rest % world;
+    const int64_t rank = rest % world_size;
-    rest /= world;
+    rest /= world_size;
-    const int64_t s = rest % seq;
+    const int64_t sequence = rest % sequence_length;
-    const int64_t b = rest / seq;
+    const int64_t batch_index = rest / sequence_length;
 
-    out[i] = x[((((w * seq + s) * batch + b) * h_local) + h) * head_dim + c];
+    out[index] =
+        x[((((rank * sequence_length + sequence) * batch + batch_index) * local_heads) + head) * head_dim +
+          head_offset];
   }
 }
 
+}  // namespace
+
+/** \brief Merge Ulysses output heads with a bit-exact layout copy. */
 template <typename T>
 struct UspMergeHeadsKernel {
-  static void run(tvm::ffi::TensorView out, tvm::ffi::TensorView x) {
+  static_assert(std::is_same_v<T, fp16_t> || std::is_same_v<T, bf16_t> || std::is_same_v<T, fp32_t>);
-    check_dtype<T>(out);
-    check_dtype<T>(x);
-    host::RuntimeCheck(x.ndim() == 5, "x must be [W, S, B, h_local, D]");
-    host::RuntimeCheck(out.ndim() == 5, "out must be [B, S, W, h_local, D]");
-    for (auto* t : {&x, &out}) {
-      host::RuntimeCheck(t->device().device_type == kDLCUDA, "usp_merge_heads tensors must be CUDA");
-      host::RuntimeCheck(is_dense_contiguous(*t), "usp_merge_heads tensors must be contiguous");
-    }
-    const int64_t world = x.size(0);
-    const int64_t seq = x.size(1);
-    const int64_t batch = x.size(2);
-    const int64_t h_local = x.size(3);
-    const int64_t head_dim = x.size(4);
-    host::RuntimeCheck(
-        out.size(0) == batch && out.size(1) == seq && out.size(2) == world && out.size(3) == h_local &&
-            out.size(4) == head_dim,
-        "out must be the [B, S, W, h_local, D] permutation of x");
 
-    const int64_t total = numel(x);
+  static void run(tvm::ffi::TensorView out, tvm::ffi::TensorView x) {
-    if (total == 0) {
+    using namespace host;
+
+    auto W = SymbolicSize{"world_size"};
+    auto S = SymbolicSize{"sequence_length"};
+    auto B = SymbolicSize{"batch"};
+    auto H = SymbolicSize{"local_heads"};
+    auto D = SymbolicSize{"head_dim"};
+    auto device = SymbolicDevice{};
+    device.set_options<kDLCUDA>();
+    TensorMatcher({W, S, B, H, D}).with_dtype<T>().with_device(device).verify(x);
+    TensorMatcher({B, S, W, H, D}).with_dtype<T>().with_device(device).verify(out);
+
+    const int64_t world_size = W.unwrap();
+    const int64_t sequence_length = S.unwrap();
+    const int64_t batch = B.unwrap();
+    const int64_t local_heads = H.unwrap();
+    const int64_t head_dim = D.unwrap();
+    const int64_t numel = world_size * sequence_length * batch * local_heads * head_dim;
+    if (numel == 0) {
       return;
     }
 
-    T* out_ptr = reinterpret_cast<T*>(mutable_data_ptr(out));
+    auto* out_ptr = static_cast<T*>(out.data_ptr());
-    const T* x_ptr = reinterpret_cast<const T*>(data_ptr(x));
+    const auto* x_ptr = static_cast<const T*>(x.data_ptr());
+    CHECK_HOST(out_ptr != x_ptr) << "output must not alias the input";
+    const auto launch = [&](auto kernel, int64_t work_items, auto... args) {
+      const auto blocks =
+          static_cast<uint32_t>(std::min<int64_t>(div_ceil(work_items, static_cast<int64_t>(kBlockSize)), kMaxGrid));
+      LaunchKernel(blocks, kBlockSize, device.unwrap())(kernel, out_ptr, x_ptr, work_items, args...);
+    };
 
-    constexpr int kVec = 16 / sizeof(T);
+    constexpr int kVec = kAlignment / sizeof(T);
-    const bool vec_ok = (head_dim % kVec == 0) && aligned16(out_ptr) && aligned16(x_ptr);
+    const bool vectorized = head_dim % kVec == 0 && reinterpret_cast<uintptr_t>(out_ptr) % kAlignment == 0 &&
-    if (vec_ok) {
+                            reinterpret_cast<uintptr_t>(x_ptr) % kAlignment == 0;
-      const int64_t n_vec = total / kVec;
+    if (vectorized) {
-      host::LaunchKernel(static_cast<uint32_t>(grid_for(n_vec)), kBlockSize, out.device())(
+      launch(
-          usp_merge_heads_vec_kernel<T, kVec>, out_ptr, x_ptr, n_vec, head_dim / kVec, h_local, batch, seq, world);
+          usp_merge_heads_vec_kernel<T, kVec>,
+          numel / kVec,
+          head_dim / kVec,
+          local_heads,
+          batch,
+          sequence_length,
+          world_size);
     } else {
-      host::LaunchKernel(static_cast<uint32_t>(grid_for(total)), kBlockSize, out.device())(
+      launch(usp_merge_heads_scalar_kernel<T>, numel, head_dim, local_heads, batch, sequence_length, world_size);
-          usp_merge_heads_scalar_kernel<T>, out_ptr, x_ptr, total, head_dim, h_local, batch, seq, world);
     }
   }
 };

python/sglang/kernels/ops/diffusion/__init__.py

@@ -36,24 +36,12 @@ register_kernel(
     KernelSpec(
         op="diffusion.residual_gate_add",
         backend=KernelBackend.JIT,
-        target="sglang.kernels.ops.diffusion.residual_gate_add:residual_gate_add_cuda",
+        target="sglang.kernels.ops.diffusion.residual_gate_add:residual_gate_add",
         capabilities=_CUDA,
         format_signature=FormatSignature(description="residual + gate * update"),
         description="Fused residual gate-add (sglang.kernels.jit).",
     )
 )
-register_kernel(
-    KernelSpec(
-        op="diffusion.fused_linear_gelu_tanh",
-        backend=KernelBackend.TORCH,
-        target="sglang.kernels.ops.diffusion.fused_linear_gelu:fused_linear_gelu_tanh",
-        capabilities=_CUDA,
-        format_signature=FormatSignature(
-            description="linear + tanh-GELU via the cublasLt GELU epilogue"
-        ),
-        description="Fused up-proj GEMM + tanh-GELU (torch._addmm_activation).",
-    )
-)
 register_kernel(
     KernelSpec(
         op="diffusion.fused_inplace_qknorm_rope",

python/sglang/kernels/ops/diffusion/cutedsl/scale_residual_norm_scale_shift.py

@@ -195,7 +195,7 @@ def validate_x(t: torch.Tensor, B: int, S: int, D: int):
     if t.shape != (B, S, D):
         raise ValueError(f"Validate failed: unsupported tensor shape: {t.shape}.")
     if t.stride()[-1] != 1:
-        raise ValueError(f"Validate failed: not contiguous on dim D.")
+        raise ValueError("Validate failed: not contiguous on dim D.")
 
 
 def validate_weight_bias(t: Optional[torch.Tensor], D: int):
@@ -206,7 +206,7 @@ def validate_weight_bias(t: Optional[torch.Tensor], D: int):
     if t.shape != (D,):
         raise ValueError(f"Validate failed: unsupported tensor shape: {t.shape}.")
     if t.stride()[-1] != 1:
-        raise ValueError(f"Validate failed: not contiguous on dim D.")
+        raise ValueError("Validate failed: not contiguous on dim D.")
 
 
 def validate_scale_shift(t: torch.Tensor, B: int, S: int, D: int):
@@ -230,7 +230,7 @@ def validate_scale_shift(t: torch.Tensor, B: int, S: int, D: int):
     if failed:
         raise ValueError(f"Validate failed: unsupported tensor shape: {t.shape}.")
     if t.stride()[-1] != 1:
-        raise ValueError(f"Validate failed: not contiguous on dim D.")
+        raise ValueError("Validate failed: not contiguous on dim D.")
 
 
 def validate_gate(t: Union[torch.Tensor, int], B: int, S: int, D: int):
@@ -311,7 +311,7 @@ def fused_norm_scale_shift(
         compiled_fn(*torch_tensors, eps, stream)
         return y
     else:
-        raise ValueError(f'norm_type must be one of "layer" and "rms"')
+        raise ValueError('norm_type must be one of "layer" and "rms"')
 
 
 @fused_norm_scale_shift.register_fake
@@ -401,7 +401,7 @@ def fused_scale_residual_norm_scale_shift(
         compiled_fn(*torch_tensors, eps, stream)
         return y, resi_out
     else:
-        raise ValueError(f'norm_type must be one of "layer" and "rms"')
+        raise ValueError('norm_type must be one of "layer" and "rms"')
 
 
 @fused_scale_residual_norm_scale_shift.register_fake

python/sglang/kernels/ops/diffusion/fused_gate_rmsnorm.py

@@ -1,10 +1,10 @@
-"""Quality-gated fused RMSNorm modulate/gate sites (Z-Image Triton suite reuse).
+"""Quality-gated fused RMSNorm modulate/gate sites.
 
 Adaln-style DiT blocks (Ideogram 4) spend four elementwise chains per block on
 modulate/gate around each RMSNorm: ``RMSNorm(x) * scale`` before
-attention/FFN and ``x + tanh(gate) * RMSNorm(out)`` after. The Z-Image
+attention/FFN and ``x + tanh(gate) * RMSNorm(out)`` after. Shared BF16-native
-bf16-native Triton kernels
+Triton kernels
-(:mod:`sglang.kernels.ops.diffusion.triton.zimage_native_norm`) fuse each
+(:mod:`sglang.kernels.ops.diffusion.triton.native_bf16_rmsnorm`) fuse each
 chain into a single kernel (RMSNorm + tanh + mul + add in one pass).
 
 Z-Image mounts those kernels unconditionally because they reproduce its own
@@ -24,16 +24,23 @@ every site on that transformer stays on the reference path.
 from __future__ import annotations
 
 import logging
-from typing import Iterator
+from importlib import import_module
 
 import torch
 import torch.nn as nn
 
+from sglang.kernels.ops.diffusion.quality_gate import QualityGatedFusion
+
 logger = logging.getLogger(__name__)
 
 # Attributes of the site protocol (set by ``mark_fused_gate_rmsnorm_site``).
 _SITE_NORM_ATTRS = "_sgl_fused_gate_rmsnorm_norm_attrs"
 _SITE_ENABLED_ATTR = "_sgl_fused_gate_rmsnorm_enabled"
+_FUSION = QualityGatedFusion(
+    name="fused gate RMSNorm",
+    marker_attr=_SITE_NORM_ATTRS,
+    enabled_attr=_SITE_ENABLED_ATTR,
+)
 
 # The Triton kernels mask a single block over the hidden dim.
 _MAX_HIDDEN_SIZE = 8192
@@ -43,11 +50,11 @@ def fused_rmsnorm_scale(
     x: torch.Tensor, weight: torch.Tensor, scale: torch.Tensor, eps: float
 ) -> torch.Tensor | None:
     """``RMSNorm(x, weight, eps) * scale`` in one Triton kernel (or None)."""
-    from sglang.kernels.ops.diffusion.triton.zimage_native_norm import (
+    from sglang.kernels.ops.diffusion.triton.native_bf16_rmsnorm import (
-        zimage_rmsnorm_scale,
+        rmsnorm_scale,
     )
 
-    return zimage_rmsnorm_scale(x, weight, scale, eps)
+    return rmsnorm_scale(x, weight, scale, eps)
 
 
 def fused_rmsnorm_tanh_residual(
@@ -58,20 +65,20 @@ def fused_rmsnorm_tanh_residual(
     eps: float,
 ) -> torch.Tensor | None:
     """``residual + tanh(gate) * RMSNorm(x, weight, eps)`` fused (or None)."""
-    from sglang.kernels.ops.diffusion.triton.zimage_native_norm import (
+    from sglang.kernels.ops.diffusion.triton.native_bf16_rmsnorm import (
-        zimage_rmsnorm_tanh_residual,
+        rmsnorm_tanh_residual,
     )
 
-    return zimage_rmsnorm_tanh_residual(x, gate, residual, weight, eps)
+    return rmsnorm_tanh_residual(x, gate, residual, weight, eps)
 
 
 def _static_reject_reason(site: nn.Module) -> str | None:
     """Why ``site`` may never use the fused kernels, or None if it may."""
     try:
-        import triton  # type: ignore # noqa: F401
+        import_module("triton")
     except ImportError:
         return "triton unavailable"
-    for attr in getattr(site, _SITE_NORM_ATTRS, ()):
+    for attr in _FUSION.metadata(site, ()):
         norm = getattr(site, attr, None)
         weight = getattr(norm, "weight", None)
         if weight is None or weight.dim() != 1:
@@ -92,15 +99,12 @@ def mark_fused_gate_rmsnorm_site(module: nn.Module, norm_attrs: tuple[str, ...])
     keep the reference path bit-exact until :func:`mount_fused_gate_rmsnorm`
     enables it.
     """
-    setattr(module, _SITE_NORM_ATTRS, tuple(norm_attrs))
+    _FUSION.mark(module, tuple(norm_attrs))
-    setattr(module, _SITE_ENABLED_ATTR, False)
 
 
-def iter_fused_gate_rmsnorm_sites(root: nn.Module) -> Iterator[nn.Module]:
+def fused_gate_rmsnorm_active(module: nn.Module) -> bool:
-    """Yield every marked site under ``root`` (including ``root``)."""
+    """Whether the quality-gated fused path is mounted on ``module``."""
-    for module in root.modules():
+    return _FUSION.is_enabled(module)
-        if getattr(module, _SITE_NORM_ATTRS, None) is not None:
-            yield module
 
 
 def mount_fused_gate_rmsnorm(root: nn.Module) -> bool:
@@ -110,26 +114,9 @@ def mount_fused_gate_rmsnorm(root: nn.Module) -> bool:
     left (or reset) on the reference path and False is returned. Returns False
     as well when ``root`` has no marked sites.
     """
-    sites = list(iter_fused_gate_rmsnorm_sites(root))
+    return _FUSION.mount(root, reject_reason=_static_reject_reason, logger=logger)
-    if not sites:
-        return False
-    for site in sites:
-        reason = _static_reject_reason(site)
-        if reason is not None:
-            unmount_fused_gate_rmsnorm(root)
-            logger.info(
-                "fused gate RMSNorm: %s site failed static guards (%s); "
-                "keeping the whole model on the reference path",
-                type(site).__name__,
-                reason,
-            )
-            return False
-    for site in sites:
-        setattr(site, _SITE_ENABLED_ATTR, True)
-    return True
 
 
 def unmount_fused_gate_rmsnorm(root: nn.Module) -> None:
     """Reset every marked site under ``root`` to the bit-exact reference path."""
-    for site in iter_fused_gate_rmsnorm_sites(root):
+    _FUSION.unmount(root)
-        setattr(site, _SITE_ENABLED_ATTR, False)

python/sglang/kernels/ops/diffusion/fused_linear_gelu.py

@@ -26,11 +26,12 @@ single opaque op under ``torch.compile`` -- no graph break.
 from __future__ import annotations
 
 import logging
-from typing import Any, Iterator
+from typing import Any
 
 import torch
 import torch.nn as nn
 
+from sglang.kernels.ops.diffusion.quality_gate import QualityGatedFusion
 from sglang.srt.utils.custom_op import register_custom_op
 
 logger = logging.getLogger(__name__)
@@ -43,6 +44,11 @@ _HAS_ADDMM_ACTIVATION = hasattr(torch, "_addmm_activation")
 # Attributes of the site protocol (set by ``mark_fused_gelu_site``).
 _SITE_LINEAR_ATTR = "_sgl_fused_gelu_linear_attr"
 _SITE_ENABLED_ATTR = "_sgl_fused_gelu_enabled"
+_FUSION = QualityGatedFusion(
+    name="fused linear+GELU",
+    marker_attr=_SITE_LINEAR_ATTR,
+    enabled_attr=_SITE_ENABLED_ATTR,
+)
 
 
 def _fused_linear_gelu_tanh_fake(
@@ -138,15 +144,17 @@ def mark_fused_gelu_site(module: nn.Module, linear_attr: str) -> None:
     (``_sgl_fused_gelu_enabled = False``): the module's forward must keep the
     reference path bit-exact until :func:`mount_fused_linear_gelu` enables it.
     """
-    setattr(module, _SITE_LINEAR_ATTR, linear_attr)
+    _FUSION.mark(module, linear_attr)
-    setattr(module, _SITE_ENABLED_ATTR, False)
 
 
-def iter_fused_gelu_sites(root: nn.Module) -> Iterator[nn.Module]:
+def fused_gelu_active(module: nn.Module) -> bool:
-    """Yield every marked fusion site under ``root`` (including ``root``)."""
+    """Whether the quality-gated fused path is mounted on ``module``."""
-    for module in root.modules():
+    return _FUSION.is_enabled(module)
-        if getattr(module, _SITE_LINEAR_ATTR, None) is not None:
+
-            yield module
+
+def _site_reject_reason(site: nn.Module) -> str | None:
+    linear = getattr(site, _FUSION.metadata(site), None)
+    return "missing linear" if linear is None else _static_reject_reason(linear)
 
 
 def mount_fused_linear_gelu(root: nn.Module) -> bool:
@@ -156,27 +164,9 @@ def mount_fused_linear_gelu(root: nn.Module) -> bool:
     left (or reset) on the reference path and False is returned. Returns False
     as well when ``root`` has no marked sites.
     """
-    sites = list(iter_fused_gelu_sites(root))
+    return _FUSION.mount(root, reject_reason=_site_reject_reason, logger=logger)
-    if not sites:
-        return False
-    for site in sites:
-        linear = getattr(site, getattr(site, _SITE_LINEAR_ATTR), None)
-        reason = "missing linear" if linear is None else _static_reject_reason(linear)
-        if reason is not None:
-            unmount_fused_linear_gelu(root)
-            logger.info(
-                "fused linear+GELU: %s site failed static guards (%s); "
-                "keeping the whole model on the reference path",
-                type(site).__name__,
-                reason,
-            )
-            return False
-    for site in sites:
-        setattr(site, _SITE_ENABLED_ATTR, True)
-    return True
 
 
 def unmount_fused_linear_gelu(root: nn.Module) -> None:
     """Reset every marked site under ``root`` to the bit-exact reference path."""
-    for site in iter_fused_gelu_sites(root):
+    _FUSION.unmount(root)
-        setattr(site, _SITE_ENABLED_ATTR, False)

python/sglang/kernels/ops/diffusion/fused_ln_modulate.py

@@ -16,42 +16,38 @@ boundaries for ``quality="high"`` requests.
 
 from __future__ import annotations
 
-from typing import Iterator
-
 import torch
 import torch.nn.functional as F
 from torch import nn
 
+from sglang.kernels.ops.diffusion.quality_gate import QualityGatedFusion
+
+_SITE_MARKER_ATTR = "_sgl_fused_ln_modulate_site"
 _SITE_ENABLED_ATTR = "_sgl_fused_ln_modulate_enabled"
+_FUSION = QualityGatedFusion(
+    name="fused LN+modulate",
+    marker_attr=_SITE_MARKER_ATTR,
+    enabled_attr=_SITE_ENABLED_ATTR,
+)
 
 _SUPPORTED_DTYPES = (torch.float16, torch.bfloat16)
 
 
 def mark_fused_ln_modulate_site(module: nn.Module) -> None:
     """Mark ``module`` as an LN+modulate fusion site (mounted off)."""
-    setattr(module, _SITE_ENABLED_ATTR, False)
+    _FUSION.mark(module)
 
 
 def fused_ln_modulate_active(module: nn.Module) -> bool:
-    return getattr(module, _SITE_ENABLED_ATTR, False)
+    return _FUSION.is_enabled(module)
-
-
-def iter_fused_ln_modulate_sites(root: nn.Module) -> Iterator[nn.Module]:
-    for module in root.modules():
-        if hasattr(module, _SITE_ENABLED_ATTR):
-            yield module
 
 
 def mount_fused_ln_modulate(root: nn.Module) -> bool:
-    sites = list(iter_fused_ln_modulate_sites(root))
+    return _FUSION.mount(root)
-    for site in sites:
-        setattr(site, _SITE_ENABLED_ATTR, True)
-    return bool(sites)
 
 
 def unmount_fused_ln_modulate(root: nn.Module) -> None:
-    for site in iter_fused_ln_modulate_sites(root):
+    _FUSION.unmount(root)
-        setattr(site, _SITE_ENABLED_ATTR, False)
 
 
 def can_fuse_ln_modulate(

python/sglang/kernels/ops/diffusion/modulate_scale_shift.py

@@ -7,6 +7,7 @@ bit-exact vs the eager chain (``torch.equal``) and needs no quality gate.
 
 from __future__ import annotations
 
+import logging
 from typing import TYPE_CHECKING
 
 import torch
@@ -20,10 +21,15 @@ if TYPE_CHECKING:
 
 _SUPPORTED_DTYPES = (torch.float16, torch.bfloat16)
 _ALIGN_BYTES = 16
+_FAILED_RUNTIME_KEYS: set[tuple[int | None, torch.dtype]] = set()
+
+logger = logging.getLogger(__name__)
 
 
 @cache_once
 def _jit_modulate_scale_shift_module(dtype: torch.dtype) -> Module:
+    if dtype not in _SUPPORTED_DTYPES:
+        raise RuntimeError(f"Unsupported modulate_scale_shift dtype: {dtype}")
     args = make_cpp_args(dtype)
     return load_jit(
         "diffusion_modulate_scale_shift",
@@ -32,8 +38,7 @@ def _jit_modulate_scale_shift_module(dtype: torch.dtype) -> Module:
         cuda_wrappers=[
             (
                 "modulate_scale_shift",
-                "sglang_modulate_scale_shift::"
+                f"modulate_scale_shift::ModulateScaleShiftKernel<{args}>::run",
-                f"ModulateScaleShiftKernel<{args}>::run",
             ),
         ],
     )
@@ -93,3 +98,33 @@ def modulate_scale_shift_cuda(
     if not can_use_modulate_scale_shift_cuda(x, scale, shift):
         raise RuntimeError("unsupported input for modulate_scale_shift CUDA")
     return _modulate_scale_shift_custom_op(x, scale, shift)
+
+
+def modulate_scale_shift(
+    x: torch.Tensor, scale: torch.Tensor, shift: torch.Tensor
+) -> torch.Tensor:
+    """Use the bit-exact CUDA fast path when supported, otherwise eager."""
+    runtime_key = (x.device.index, x.dtype)
+    if runtime_key not in _FAILED_RUNTIME_KEYS and can_use_modulate_scale_shift_cuda(
+        x, scale, shift
+    ):
+        try:
+            return modulate_scale_shift_cuda(x, scale, shift)
+        except Exception as exc:
+            if torch.compiler.is_compiling():
+                raise
+            _FAILED_RUNTIME_KEYS.add(runtime_key)
+            logger.warning(
+                "Disabling diffusion modulate CUDA fast path on %s/%s: %s",
+                x.device,
+                x.dtype,
+                exc,
+            )
+    return x * (1 + scale[:, None]) + shift[:, None]
+
+
+__all__ = [
+    "can_use_modulate_scale_shift_cuda",
+    "modulate_scale_shift",
+    "modulate_scale_shift_cuda",
+]

python/sglang/kernels/ops/diffusion/qknorm_rope.py

@@ -19,6 +19,9 @@ if TYPE_CHECKING:
 
 logger = logging.getLogger(__name__)
 
+_SUPPORTED_DTYPES = (torch.float16, torch.bfloat16)
+_SUPPORTED_CACHE_DTYPES = (*_SUPPORTED_DTYPES, torch.float32)
+
 
 @cache_once
 def _jit_qknorm_rope_module(
@@ -56,6 +59,13 @@ def can_use_fused_inplace_qknorm_rope(
     cache_dtype: torch.dtype = torch.float32,
     round_norm_before_rope: bool = False,
 ) -> bool:
+    if dtype not in _SUPPORTED_DTYPES or cache_dtype not in _SUPPORTED_CACHE_DTYPES:
+        logger.warning(
+            "Unsupported dtype pair (%s, %s) for JIT fused QKNorm+RoPE",
+            dtype,
+            cache_dtype,
+        )
+        return False
     if head_dim not in (64, 128, 256):
         logger.warning(f"Unsupported head_dim={head_dim} for JIT fused QKNorm+RoPE")
         return False

python/sglang/kernels/ops/diffusion/quality_gate.py

@@ -0,0 +1,82 @@
+"""Shared module-site protocol for request-scoped diffusion fast paths."""
+
+from __future__ import annotations
+
+import logging
+from collections.abc import Callable, Iterator
+from typing import Any
+
+from torch import nn
+
+RejectReason = Callable[[nn.Module], str | None]
+
+
+class QualityGatedFusion:
+    """Track and toggle one family of opt-in fusion sites.
+
+    The marker metadata describes the model attribute(s) owned by a site. The
+    enabled flag remains a plain module attribute so compiled model forwards
+    can read it without depending on this helper object.
+    """
+
+    __slots__ = ("enabled_attr", "marker_attr", "name")
+
+    def __init__(self, *, name: str, marker_attr: str, enabled_attr: str) -> None:
+        self.name = name
+        self.marker_attr = marker_attr
+        self.enabled_attr = enabled_attr
+
+    def mark(self, module: nn.Module, metadata: Any = True) -> None:
+        setattr(module, self.marker_attr, metadata)
+        setattr(module, self.enabled_attr, False)
+
+    def metadata(self, module: nn.Module, default: Any = None) -> Any:
+        return getattr(module, self.marker_attr, default)
+
+    def is_enabled(self, module: nn.Module) -> bool:
+        return bool(getattr(module, self.enabled_attr, False))
+
+    def iter_sites(self, root: nn.Module) -> Iterator[nn.Module]:
+        for module in root.modules():
+            if hasattr(module, self.marker_attr):
+                yield module
+
+    def mount(
+        self,
+        root: nn.Module,
+        *,
+        reject_reason: RejectReason | None = None,
+        logger: logging.Logger | None = None,
+    ) -> bool:
+        """Enable every eligible site, or leave the whole family disabled."""
+        sites = list(self.iter_sites(root))
+        if not sites:
+            return False
+
+        if reject_reason is not None:
+            for site in sites:
+                reason = reject_reason(site)
+                if reason is None:
+                    continue
+                self._set_enabled(sites, False)
+                if logger is not None:
+                    logger.info(
+                        "%s: %s site failed static guards (%s); keeping the "
+                        "whole model on the reference path",
+                        self.name,
+                        type(site).__name__,
+                        reason,
+                    )
+                return False
+
+        self._set_enabled(sites, True)
+        return True
+
+    def unmount(self, root: nn.Module) -> None:
+        self._set_enabled(self.iter_sites(root), False)
+
+    def _set_enabled(
+        self, sites: Iterator[nn.Module] | list[nn.Module], enabled: bool
+    ) -> None:
+        for site in sites:
+            setattr(site, self.enabled_attr, enabled)

python/sglang/kernels/ops/diffusion/residual_gate_add.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import logging
 from typing import TYPE_CHECKING
 
 import torch
@@ -12,10 +13,16 @@ if TYPE_CHECKING:
 
 
 _SUPPORTED_DTYPES = (torch.float16, torch.bfloat16, torch.float32)
+_BIT_EXACT_DTYPES = (torch.float16, torch.bfloat16)
+_FAILED_RUNTIME_KEYS: set[tuple[int | None, torch.dtype]] = set()
+
+logger = logging.getLogger(__name__)
 
 
 @cache_once
 def _jit_residual_gate_add_module(dtype: torch.dtype) -> Module:
+    if dtype not in _SUPPORTED_DTYPES:
+        raise RuntimeError(f"Unsupported residual_gate_add dtype: {dtype}")
     args = make_cpp_args(dtype)
     return load_jit(
         "diffusion_residual_gate_add",
@@ -46,15 +53,22 @@ def _residual_gate_add_custom_op(
 ) -> torch.Tensor:
     out = torch.empty_like(residual)
     module = _jit_residual_gate_add_module(residual.dtype)
-    module.residual_gate_add(out, residual, update, gate)
+    broadcast_gate = gate.shape != residual.shape
+    module.residual_gate_add(
+        out.view(-1),
+        residual.view(-1),
+        update.view(-1),
+        gate.view(-1),
+        residual.shape[-1],
+        broadcast_gate,
+    )
     return out
 
 
 def _is_row_broadcast_gate(residual: torch.Tensor, gate: torch.Tensor) -> bool:
     if gate.dim() != residual.dim() or gate.shape[-1] != residual.shape[-1]:
         return False
-    row_dim = gate.dim() - 2
+    return all(size == 1 for size in gate.shape[:-1])
-    return gate.shape[row_dim] == 1 and all(size == 1 for size in gate.shape[:-1])
 
 
 def can_use_residual_gate_add_cuda(
@@ -69,6 +83,7 @@ def can_use_residual_gate_add_cuda(
         and gate.is_cuda
         and residual.device == update.device == gate.device
         and residual.dim() >= 2
+        and residual.numel() > 0
         and update.shape == residual.shape
         and (gate.shape == residual.shape or _is_row_broadcast_gate(residual, gate))
         and residual.is_contiguous()
@@ -83,3 +98,39 @@ def residual_gate_add_cuda(
     if not can_use_residual_gate_add_cuda(residual, update, gate):
         raise RuntimeError("unsupported input for residual_gate_add CUDA")
     return _residual_gate_add_custom_op(residual, update, gate)
+
+
+def residual_gate_add(
+    residual: torch.Tensor, update: torch.Tensor, gate: torch.Tensor
+) -> torch.Tensor:
+    """Use the bit-exact CUDA fast path when supported, otherwise eager.
+
+    Runtime build failures are cached per device and dtype so every diffusion
+    model shares one fallback policy instead of maintaining model-local flags.
+    """
+    runtime_key = (residual.device.index, residual.dtype)
+    if (
+        residual.dtype in _BIT_EXACT_DTYPES
+        and runtime_key not in _FAILED_RUNTIME_KEYS
+        and can_use_residual_gate_add_cuda(residual, update, gate)
+    ):
+        try:
+            return residual_gate_add_cuda(residual, update, gate)
+        except Exception as exc:
+            if torch.compiler.is_compiling():
+                raise
+            _FAILED_RUNTIME_KEYS.add(runtime_key)
+            logger.warning(
+                "Disabling diffusion residual-gate CUDA fast path on %s/%s: %s",
+                residual.device,
+                residual.dtype,
+                exc,
+            )
+    return residual + update * gate
+
+
+__all__ = [
+    "can_use_residual_gate_add_cuda",
+    "residual_gate_add",
+    "residual_gate_add_cuda",
+]

python/sglang/kernels/ops/diffusion/triton/group_norm_silu_twopass.py

@@ -133,6 +133,7 @@ def _gn_silu_rows(x3, weight, bias, num_groups, eps, apply_silu):
     nchunks = triton.cdiv(rows, rows_per_prog)
     psum = torch.empty((n_batch, nchunks, c), device=x3.device, dtype=torch.float32)
     psq = torch.empty_like(psum)
+    with torch.get_device_module().device(x3.device):
         _gn_partial_rows_kernel[(nchunks, n_batch)](
             x3, psum, psq, rows, rows_per_prog, C=c, BLOCK_R=block_r, num_warps=4
         )
@@ -161,7 +162,7 @@ def _gn_silu_rows(x3, weight, bias, num_groups, eps, apply_silu):
 
 def _twopass_supported(x, weight, bias, num_groups) -> bool:
     """Tensor-level support check shared by the 4D and rows entry points."""
-    if not (x.is_cuda and not torch.is_grad_enabled()):
+    if not (x.is_cuda and x.numel() > 0 and not torch.is_grad_enabled()):
         return False
     if x.requires_grad or x.dtype not in _SUPPORTED_DTYPES:
         return False
@@ -170,6 +171,13 @@ def _twopass_supported(x, weight, bias, num_groups) -> bool:
     c = x.shape[1] if x.dim() == 4 else x.shape[-1]
     if weight.shape != (c,) or bias.shape != (c,):
         return False
+    if not (
+        weight.device == bias.device == x.device
+        and weight.dtype == bias.dtype == x.dtype
+        and weight.is_contiguous()
+        and bias.is_contiguous()
+    ):
+        return False
     if num_groups < 1 or c % num_groups != 0:
         return False
     # tl.arange needs a power-of-two C; num_groups divides it, so the

python/sglang/kernels/ops/diffusion/triton/indexed_modulation.py

@@ -4,14 +4,7 @@ import torch
 import triton
 import triton.language as tl
 
-
+from sglang.kernels.ops.diffusion.triton.numerics import round_bf16_to_fp32
-@triton.jit
-def _round_bf16_to_fp32(value):
-    # force the eager BF16 kernel boundary so Triton cannot contract the next add
-    bits = value.to(tl.int32, bitcast=True)
-    rounding_bias = 0x7FFF + ((bits >> 16) & 1)
-    rounded_bits = (bits + rounding_bias) & -65536
-    return rounded_bits.to(tl.float32, bitcast=True)
 
 
 @triton.jit
@@ -43,8 +36,8 @@ def _indexed_scale_shift_bf16_kernel(
         scale_ptr + index * stride_scale_row + columns, mask=mask, other=0.0
     ).to(tl.float32)
 
-    one_plus_scale = _round_bf16_to_fp32(1.0 + scale)
+    one_plus_scale = round_bf16_to_fp32(1.0 + scale)
-    scaled = _round_bf16_to_fp32(x * one_plus_scale)
+    scaled = round_bf16_to_fp32(x * one_plus_scale)
     tl.store(
         output_ptr + row * stride_x_row + columns,
         scaled + shift,
@@ -81,7 +74,7 @@ def _indexed_gate_bf16_kernel(
         other_ptr + row * stride_other_row + columns, mask=mask, other=0.0
     ).to(tl.float32)
 
-    gated = _round_bf16_to_fp32(gate * other)
+    gated = round_bf16_to_fp32(gate * other)
     tl.store(
         output_ptr + row * stride_x_row + columns,
         x + gated,

python/sglang/kernels/ops/diffusion/triton/layernorm_modulate.py

@@ -46,20 +46,13 @@ import torch
 import triton  # type: ignore
 import triton.language as tl  # type: ignore
 
+from sglang.kernels.ops.diffusion.triton.numerics import (
+    cuda_rsqrtf,
+    div_rn_f32,
+    round_bf16_to_fp32,
+)
 from sglang.srt.utils.custom_op import register_custom_op
 
-_FLT_MIN = tl.constexpr(1.1754943508222875e-38)
-
-
-@triton.jit
-def _round_bf16_to_fp32(value):
-    # RNE round of an fp32 value to bf16 precision, staying in fp32 registers
-    # (also blocks any fmul+fadd contraction across the boundary).
-    bits = value.to(tl.int32, bitcast=True)
-    rounding_bias = 0x7FFF + ((bits >> 16) & 1)
-    rounded_bits = (bits + rounding_bias) & -65536
-    return rounded_bits.to(tl.float32, bitcast=True)
-
 
 @triton.jit
 def _rcp4(x):
@@ -80,40 +73,6 @@ def _rcp4(x):
     )
 
 
-@triton.jit
-def _div_rn(x, y):
-    # IEEE correctly-rounded fp32 division.
-    return tl.inline_asm_elementwise(
-        asm="div.rn.f32 $0, $1, $2;",
-        constraints="=f,f,f",
-        args=[x, y],
-        dtype=tl.float32,
-        is_pure=True,
-        pack=1,
-    )
-
-
-@triton.jit
-def _rsqrt_approx(x):
-    return tl.inline_asm_elementwise(
-        asm="rsqrt.approx.f32 $0, $1;",
-        constraints="=f,f",
-        args=[x],
-        dtype=tl.float32,
-        is_pure=True,
-        pack=1,
-    )
-
-
-@triton.jit
-def _rsqrtf(x):
-    # CUDA rsqrtf: MUFU.RSQ with a 2^24 / 2^12 rescale for subnormal inputs.
-    p = tl.abs(x) < _FLT_MIN
-    xs = tl.where(p, x * 16777216.0, x)
-    r = _rsqrt_approx(xs)
-    return tl.where(p, r * 4096.0, r)
-
-
 @triton.jit
 def _welford_push(val, mean, m2, cnt, valid, MASKED: tl.constexpr):
     # ``valid`` masks lanes whose aten thread does not execute this
@@ -269,7 +228,7 @@ def _layernorm_modulate_kernel(
     mean, m2, cnt = _fold_halves(mean, m2, cnt, ROWS, 1)
 
     denom = tl.zeros((ROWS, 1), dtype=tl.float32) + D
-    rstd = _rsqrtf(_div_rn(m2, denom) + eps)  # (ROWS, 1)
+    rstd = cuda_rsqrtf(div_rn_f32(m2, denom) + eps)  # (ROWS, 1)
 
     batch = row_offs // seq_len
 
@@ -284,7 +243,7 @@ def _layernorm_modulate_kernel(
             mask=mask,
             other=0.0,
         ).to(tl.float32)
-        y = _round_bf16_to_fp32(rstd * (x - mean))
+        y = round_bf16_to_fp32(rstd * (x - mean))
         sc = tl.load(
             scale_ptr + batch[:, None] * scale_row_stride + cols[None, :],
             mask=mask,
@@ -295,8 +254,8 @@ def _layernorm_modulate_kernel(
             mask=mask,
             other=0.0,
         ).to(tl.float32)
-        one_plus = _round_bf16_to_fp32(1.0 + sc)
+        one_plus = round_bf16_to_fp32(1.0 + sc)
-        y = _round_bf16_to_fp32(y * one_plus) + sh
+        y = round_bf16_to_fp32(y * one_plus) + sh
         tl.store(y_ptr + row_base[:, None] + cols[None, :], y, mask=mask)
 
 
@@ -337,7 +296,7 @@ def _qk_ln_head_one(
     mean, m2, cnt = _welford_combine(mean, m2, cnt, zero, zero, zero)
 
     denom = tl.zeros((ROWS, 1), dtype=tl.float32) + D
-    rstd = _rsqrtf(_div_rn(m2, denom) + eps)
+    rstd = cuda_rsqrtf(div_rn_f32(m2, denom) + eps)
 
     cols2 = tl.arange(0, D_POW2)
     out_mask = row_mask[:, None] & (cols2 < D)[None, :]

python/sglang/kernels/ops/diffusion/triton/mps_fallback.py

@@ -18,8 +18,12 @@ from torch import Tensor
 from sglang.srt.utils.tensor_bridge import mlx_to_torch, torch_to_mlx, use_mlx
 
 from .torch_fallback import (
-    apply_rotary_embedding_native,
+    apply_rotary_embedding_native as apply_rotary_embedding_native,
-    fuse_scale_shift_kernel_native,
+)
+from .torch_fallback import (
+    fuse_scale_shift_kernel_native as fuse_scale_shift_kernel_native,
+)
+from .torch_fallback import (
     norm_infer_native,
     rms_norm_fn_native,
     triton_one_pass_rms_norm_native,
@@ -30,13 +34,6 @@ _use_mlx = use_mlx()
 if _use_mlx:
     import mlx.core as mx
 
-# use the common torch native version form torch_fallback
-fuse_scale_shift_kernel_native = fuse_scale_shift_kernel_native
-apply_rotary_embedding_native = apply_rotary_embedding_native
-norm_infer_native = norm_infer_native
-triton_one_pass_rms_norm_native = triton_one_pass_rms_norm_native
-rms_norm_fn_native = rms_norm_fn_native
-
 # MLX-accelerated norm ops (1.4x–2.9x faster than torch native on MPS)
 # Uses mx.fast.rms_norm / mx.fast.layer_norm — single fused Metal kernels
 # instead of 7+ separate PyTorch MPS kernel launches.

python/sglang/kernels/ops/diffusion/triton/native_bf16_rmsnorm.py

@@ -0,0 +1,208 @@
+# SPDX-License-Identifier: Apache-2.0
+"""BF16-native RMSNorm fusions shared by diffusion transformer models."""
+
+from __future__ import annotations
+
+import torch
+import triton  # type: ignore
+import triton.language as tl  # type: ignore
+
+MAX_HIDDEN_SIZE = 8192
+
+
+@triton.jit
+def _tanh(x):
+    return 2.0 / (1.0 + tl.exp(-2.0 * x)) - 1.0
+
+
+@triton.jit
+def _rmsnorm_scale_kernel(
+    y_ptr,
+    x_ptr,
+    weight_ptr,
+    scale_ptr,
+    x_row_stride,
+    scale_row_stride,
+    seq_len,
+    dim: tl.constexpr,
+    eps: tl.constexpr,
+    block_dim: tl.constexpr,
+):
+    row = tl.program_id(0)
+    offsets = tl.arange(0, block_dim)
+    mask = offsets < dim
+
+    x = tl.load(x_ptr + row * x_row_stride + offsets, mask=mask, other=0.0)
+    square = (x * x).to(tl.bfloat16)
+    mean_square = (tl.sum(square, axis=0) / dim).to(tl.bfloat16)
+    rstd = tl.rsqrt((mean_square + eps).to(tl.bfloat16).to(tl.float32)).to(tl.bfloat16)
+
+    batch = row // seq_len
+    weight = tl.load(weight_ptr + offsets, mask=mask, other=0.0)
+    scale = tl.load(
+        scale_ptr + batch * scale_row_stride + offsets, mask=mask, other=0.0
+    )
+    y = (((x * rstd).to(tl.bfloat16) * weight).to(tl.bfloat16) * scale).to(tl.bfloat16)
+    tl.store(y_ptr + row * dim + offsets, y, mask=mask)
+
+
+@triton.jit
+def _rmsnorm_tanh_residual_kernel(
+    y_ptr,
+    x_ptr,
+    gate_ptr,
+    residual_ptr,
+    weight_ptr,
+    x_row_stride,
+    gate_row_stride,
+    residual_row_stride,
+    seq_len,
+    dim: tl.constexpr,
+    eps: tl.constexpr,
+    block_dim: tl.constexpr,
+):
+    row = tl.program_id(0)
+    offsets = tl.arange(0, block_dim)
+    mask = offsets < dim
+
+    x = tl.load(x_ptr + row * x_row_stride + offsets, mask=mask, other=0.0)
+    square = (x * x).to(tl.bfloat16)
+    mean_square = (tl.sum(square, axis=0) / dim).to(tl.bfloat16)
+    rstd = tl.rsqrt((mean_square + eps).to(tl.bfloat16).to(tl.float32)).to(tl.bfloat16)
+
+    batch = row // seq_len
+    gate = tl.load(gate_ptr + batch * gate_row_stride + offsets, mask=mask, other=0.0)
+    residual = tl.load(
+        residual_ptr + row * residual_row_stride + offsets, mask=mask, other=0.0
+    )
+    weight = tl.load(weight_ptr + offsets, mask=mask, other=0.0)
+    norm = ((x * rstd).to(tl.bfloat16) * weight).to(tl.bfloat16)
+    gated = (_tanh(gate.to(tl.float32)).to(tl.bfloat16) * norm).to(tl.bfloat16)
+    y = (residual + gated).to(tl.bfloat16)
+    tl.store(y_ptr + row * dim + offsets, y, mask=mask)
+
+
+def _flat_row_stride(x: torch.Tensor) -> int | None:
+    if x.dim() < 2 or x.stride(-1) != 1:
+        return None
+    row_stride = x.stride(-2)
+    expected_stride = row_stride * x.shape[-2]
+    for dim in range(x.dim() - 3, -1, -1):
+        if x.stride(dim) != expected_stride:
+            return None
+        expected_stride *= x.shape[dim]
+    return row_stride
+
+
+def _can_use_operand(
+    x: torch.Tensor, weight: torch.Tensor, other: torch.Tensor
+) -> bool:
+    return (
+        x.is_cuda
+        and weight.is_cuda
+        and other.is_cuda
+        and x.device == weight.device == other.device
+        and x.dtype == weight.dtype == other.dtype == torch.bfloat16
+        and x.dim() >= 2
+        and 0 < x.shape[-1] <= MAX_HIDDEN_SIZE
+        and x.numel() > 0
+        and weight.shape == (x.shape[-1],)
+        and weight.is_contiguous()
+        and other.dim() >= 2
+        and other.shape[-1] == x.shape[-1]
+        and other.numel() > 0
+        and _flat_row_stride(x) is not None
+        and _flat_row_stride(other) is not None
+    )
+
+
+def rmsnorm_scale(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    scale: torch.Tensor,
+    eps: float,
+) -> torch.Tensor | None:
+    """Apply BF16-native ``RMSNorm(x) * scale`` or return ``None``."""
+    if not _can_use_operand(x, weight, scale):
+        return None
+
+    dim = x.shape[-1]
+    x_rows = x.numel() // dim
+    scale_rows = scale.numel() // dim
+    if x_rows % scale_rows != 0:
+        return None
+
+    x_row_stride = _flat_row_stride(x)
+    scale_row_stride = _flat_row_stride(scale)
+    if x_row_stride is None or scale_row_stride is None:
+        return None
+
+    out = torch.empty_like(x, memory_format=torch.contiguous_format)
+    with torch.get_device_module().device(x.device):
+        _rmsnorm_scale_kernel[(x_rows,)](
+            out.reshape(-1, dim),
+            x,
+            weight,
+            scale,
+            x_row_stride,
+            scale_row_stride,
+            x_rows // scale_rows,
+            dim,
+            eps,
+            block_dim=triton.next_power_of_2(dim),
+            num_warps=8,
+        )
+    return out
+
+
+def rmsnorm_tanh_residual(
+    x: torch.Tensor,
+    gate: torch.Tensor,
+    residual: torch.Tensor,
+    weight: torch.Tensor,
+    eps: float,
+) -> torch.Tensor | None:
+    """Apply BF16-native gated RMSNorm residual fusion or return ``None``."""
+    if not _can_use_operand(x, weight, gate):
+        return None
+    if (
+        residual.device != x.device
+        or residual.dtype != x.dtype
+        or residual.shape != x.shape
+        or _flat_row_stride(residual) is None
+    ):
+        return None
+
+    dim = x.shape[-1]
+    x_rows = x.numel() // dim
+    gate_rows = gate.numel() // dim
+    if x_rows % gate_rows != 0:
+        return None
+
+    x_row_stride = _flat_row_stride(x)
+    gate_row_stride = _flat_row_stride(gate)
+    residual_row_stride = _flat_row_stride(residual)
+    if x_row_stride is None or gate_row_stride is None or residual_row_stride is None:
+        return None
+
+    out = torch.empty_like(x, memory_format=torch.contiguous_format)
+    with torch.get_device_module().device(x.device):
+        _rmsnorm_tanh_residual_kernel[(x_rows,)](
+            out.reshape(-1, dim),
+            x,
+            gate,
+            residual,
+            weight,
+            x_row_stride,
+            gate_row_stride,
+            residual_row_stride,
+            x_rows // gate_rows,
+            dim,
+            eps,
+            block_dim=triton.next_power_of_2(dim),
+            num_warps=8,
+        )
+    return out
+
+
+__all__ = ["MAX_HIDDEN_SIZE", "rmsnorm_scale", "rmsnorm_tanh_residual"]

python/sglang/kernels/ops/diffusion/triton/numerics.py

@@ -0,0 +1,72 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Numerical primitives shared by bit-exact diffusion Triton kernels."""
+
+import triton  # type: ignore
+import triton.language as tl  # type: ignore
+
+_FLT_MIN = tl.constexpr(1.1754943508222875e-38)
+
+
+@triton.jit
+def round_bf16_to_fp32(value):
+    """RNE-round fp32 to bf16 precision while keeping an fp32 register."""
+    bits = value.to(tl.int32, bitcast=True)
+    rounding_bias = 0x7FFF + ((bits >> 16) & 1)
+    rounded_bits = (bits + rounding_bias) & -65536
+    return rounded_bits.to(tl.float32, bitcast=True)
+
+
+@triton.jit
+def mul_rn_f32(x, y):
+    """Opaque correctly-rounded fp32 multiply that blocks FMA contraction."""
+    return tl.inline_asm_elementwise(
+        asm="mul.rn.f32 $0, $1, $2;",
+        constraints="=f,f,f",
+        args=[x, y],
+        dtype=tl.float32,
+        is_pure=True,
+        pack=1,
+    )
+
+
+@triton.jit
+def div_rn_f32(x, y):
+    """IEEE correctly-rounded fp32 division."""
+    return tl.inline_asm_elementwise(
+        asm="div.rn.f32 $0, $1, $2;",
+        constraints="=f,f,f",
+        args=[x, y],
+        dtype=tl.float32,
+        is_pure=True,
+        pack=1,
+    )
+
+
+@triton.jit
+def rsqrt_approx_f32(x):
+    return tl.inline_asm_elementwise(
+        asm="rsqrt.approx.f32 $0, $1;",
+        constraints="=f,f",
+        args=[x],
+        dtype=tl.float32,
+        is_pure=True,
+        pack=1,
+    )
+
+
+@triton.jit
+def cuda_rsqrtf(x):
+    """Match CUDA ``rsqrtf``, including its subnormal rescaling path."""
+    is_subnormal = tl.abs(x) < _FLT_MIN
+    scaled = tl.where(is_subnormal, x * 16777216.0, x)
+    result = rsqrt_approx_f32(scaled)
+    return tl.where(is_subnormal, result * 4096.0, result)
+
+
+__all__ = [
+    "cuda_rsqrtf",
+    "div_rn_f32",
+    "mul_rn_f32",
+    "round_bf16_to_fp32",
+    "rsqrt_approx_f32",
+]

python/sglang/kernels/ops/diffusion/triton/rmsnorm_scale_shift_bitexact.py

@@ -55,45 +55,14 @@ import torch
 import triton  # type: ignore
 import triton.language as tl  # type: ignore
 
+from sglang.kernels.ops.diffusion.triton.numerics import (
+    mul_rn_f32,
+    round_bf16_to_fp32,
+    rsqrt_approx_f32,
+)
 from sglang.srt.utils.custom_op import register_custom_op
 
 
-@triton.jit
-def _round_bf16_to_fp32(value):
-    # RNE round of an fp32 value to bf16 precision, staying in fp32 registers
-    # (also blocks any fmul+fadd contraction across the boundary).
-    bits = value.to(tl.int32, bitcast=True)
-    rounding_bias = 0x7FFF + ((bits >> 16) & 1)
-    rounded_bits = (bits + rounding_bias) & -65536
-    return rounded_bits.to(tl.float32, bitcast=True)
-
-
-@triton.jit
-def _mul_rn_f32(x, y):
-    # opaque mul.rn.f32: keeps the square a separately rounded fp32 op and
-    # blocks contraction with the following add.
-    return tl.inline_asm_elementwise(
-        asm="mul.rn.f32 $0, $1, $2;",
-        constraints="=f,f,f",
-        args=[x, y],
-        dtype=tl.float32,
-        is_pure=True,
-        pack=1,
-    )
-
-
-@triton.jit
-def _rsqrt_approx_f32(x):
-    return tl.inline_asm_elementwise(
-        asm="rsqrt.approx.f32 $0, $1;",
-        constraints="=f,f",
-        args=[x],
-        dtype=tl.float32,
-        is_pure=True,
-        pack=1,
-    )
-
-
 @triton.jit
 def _fold_adjacent(p, rows: tl.constexpr, width: tl.constexpr):
     # (rows, 2*width) -> (rows, width): add adjacent pairs (even + odd), the
@@ -138,10 +107,10 @@ def _rmsnorm_scale_shift_kernel(
                 uj = tl.load(x_ptr + row_base + col).to(tl.float32)
                 gj = tl.load(gate_ptr + vec_base + col).to(tl.float32)
                 # eager pair: bf16 round after gate*update and after the add
-                xj = _round_bf16_to_fp32(rj + _round_bf16_to_fp32(gj * uj))
+                xj = round_bf16_to_fp32(rj + round_bf16_to_fp32(gj * uj))
             else:
                 xj = tl.load(x_ptr + row_base + col).to(tl.float32)
-            acc = acc + _mul_rn_f32(xj, xj)
+            acc = acc + mul_rn_f32(xj, xj)
 
     # warp butterfly (offsets 1,2,4,8,16) == adjacent-pairs fold tree,
     # then the WPR warp sums are combined the same way.
@@ -154,7 +123,7 @@ def _rmsnorm_scale_shift_kernel(
     s = tl.reshape(p, (1, WPR))
     if WPR == 2:
         s = _fold_adjacent(s, 1, 1)
-    rcp = tl.sum(_rsqrt_approx_f32(s / D + eps))  # single element, exact
+    rcp = tl.sum(rsqrt_approx_f32(s / D + eps))  # single element, exact
 
     # ----- pass 2: normalize + modulate, contiguous chunks -----
     for i in tl.static_range(D // 1024):
@@ -163,16 +132,16 @@ def _rmsnorm_scale_shift_kernel(
             r = tl.load(residual_ptr + row_base + cols).to(tl.float32)
             u = tl.load(x_ptr + row_base + cols).to(tl.float32)
             g = tl.load(gate_ptr + vec_base + cols).to(tl.float32)
-            xin = _round_bf16_to_fp32(r + _round_bf16_to_fp32(g * u))
+            xin = round_bf16_to_fp32(r + round_bf16_to_fp32(g * u))
             tl.store(res_out_ptr + row_base + cols, xin)
         else:
             xin = tl.load(x_ptr + row_base + cols).to(tl.float32)
         w = tl.load(weight_ptr + cols).to(tl.float32)
         sc = tl.load(scale_ptr + vec_base + cols).to(tl.float32)
         sh = tl.load(shift_ptr + vec_base + cols).to(tl.float32)
-        y = _round_bf16_to_fp32(xin * rcp * w)  # (bf16)(x * rstd * w)
+        y = round_bf16_to_fp32(xin * rcp * w)  # (bf16)(x * rstd * w)
-        one_plus = _round_bf16_to_fp32(1.0 + sc)
+        one_plus = round_bf16_to_fp32(1.0 + sc)
-        prod = _round_bf16_to_fp32(y * one_plus)
+        prod = round_bf16_to_fp32(y * one_plus)
         tl.store(out_ptr + row_base + cols, prod + sh)  # store rounds to bf16
 
 

python/sglang/kernels/ops/diffusion/triton/sana_wm_gdn_chunkwise.py

@@ -885,8 +885,13 @@ def phase_b_triton(
 
     load_init = init_state_kv is not None
     dummy = torch.empty(1, device=device, dtype=fdtype)
-    full_M = lambda: torch.empty(BH, F, BLOCK_D, BLOCK_D, device=device, dtype=fdtype)
+
-    full_z = lambda: torch.empty(BH, F, BLOCK_D, device=device, dtype=fdtype)
+    def full_M():
+        return torch.empty(BH, F, BLOCK_D, BLOCK_D, device=device, dtype=fdtype)
+
+    def full_z():
+        return torch.empty(BH, F, BLOCK_D, device=device, dtype=fdtype)
+
     M_fwd = dummy if direction == 2 else full_M()
     z_fwd = dummy if (direction == 2 or skip_z) else full_z()
     # Combined-history reuses M_fwd/z_fwd as M_hist/z_hist; rev outputs are

python/sglang/kernels/ops/diffusion/triton/scale_shift.py

@@ -2,26 +2,10 @@ import torch
 import triton  # type: ignore
 import triton.language as tl  # type: ignore
 
+from sglang.kernels.ops.diffusion.triton.numerics import mul_rn_f32
 from sglang.multimodal_gen.runtime.platforms import current_platform
 
 
-@triton.jit
-def _fp32_mul_add_rn(x, scale, residual):
-    """Match separate CUDA FP32 multiply and add rounding (no FMA)."""
-    return tl.inline_asm_elementwise(
-        asm="""{
-            .reg .f32 product;
-            mul.rn.f32 product, $1, $2;
-            add.rn.f32 $0, $3, product;
-        }""",
-        constraints="=f,f,f,f",
-        args=(x, scale, residual),
-        dtype=tl.float32,
-        is_pure=True,
-        pack=1,
-    )
-
-
 @triton.jit
 def _fused_scaled_residual_add_exact_kernel(
     output_ptr,
@@ -37,7 +21,8 @@ def _fused_scaled_residual_add_exact_kernel(
     x = tl.load(x_ptr + offsets, mask=mask).to(tl.float32)
     scale = tl.load(scale_ptr + offsets % width, mask=mask)
     residual = tl.load(residual_ptr + offsets, mask=mask)
-    output = _fp32_mul_add_rn(x, scale, residual)
+    # The opaque multiply keeps Triton from contracting this into an FMA.
+    output = residual + mul_rn_f32(x, scale)
     tl.store(output_ptr + offsets, output, mask=mask)
 
 

python/sglang/kernels/ops/diffusion/triton/ulysses_qkv.py

@@ -59,12 +59,32 @@ def pack_qkv_destination_major(
     world_size: int,
     out: torch.Tensor | None = None,
 ) -> torch.Tensor:
+    """Pack matching ``[rows, global_heads, head_size]`` Q/K/V tensors."""
+    if q.dim() != 3 or q.shape != k.shape or q.shape != v.shape:
+        raise ValueError("q, k, and v must have the same 3D shape")
+    if not (q.is_cuda and k.is_cuda and v.is_cuda):
+        raise ValueError("q, k, and v must be CUDA tensors")
+    if not (q.device == k.device == v.device and q.dtype == k.dtype == v.dtype):
+        raise ValueError("q, k, and v must have the same device and dtype")
+    if q.stride(-1) != 1 or k.stride(-1) != 1 or v.stride(-1) != 1:
+        raise ValueError("q, k, and v must be contiguous in head_size")
+    if world_size < 1 or q.shape[1] % world_size != 0:
+        raise ValueError("world_size must be positive and divide global_heads")
+
     rows, global_heads, head_size = q.shape
     local_heads = global_heads // world_size
     expected_shape = (world_size, rows, local_heads, 3 * head_size)
     if out is not None:
-        assert out.shape == expected_shape and out.is_contiguous()
+        if not (
-        assert out.dtype == q.dtype and out.device == q.device
+            out.shape == expected_shape
+            and out.is_contiguous()
+            and out.dtype == q.dtype
+            and out.device == q.device
+        ):
+            raise ValueError(
+                "out must be a contiguous tensor with the expected shape, "
+                "device, and dtype"
+            )
         output = out
     else:
         output = torch.empty(
@@ -77,6 +97,7 @@ def pack_qkv_destination_major(
         return output
 
     block_size = 1024
+    with torch.get_device_module().device(q.device):
         _pack_qkv_destination_major_kernel[(triton.cdiv(total_elements, block_size),)](
             output,
             q,
@@ -96,3 +117,6 @@ def pack_qkv_destination_major(
             num_warps=8,
         )
     return output
+
+
+__all__ = ["pack_qkv_destination_major"]

python/sglang/kernels/ops/diffusion/triton/wan_rmsnorm_silu.py

@@ -119,7 +119,7 @@ def _triton_wan_rmsnorm_silu_cuda(
     block_c = triton.next_power_of_2(channels)
     num_warps = 1 if block_c <= 64 else 4 if block_c <= 512 else 8
 
-    with torch.cuda.device(x.device):
+    with torch.get_device_module().device(x.device):
         _wan_rmsnorm_silu_kernel[(bsz * t_size * h_size * w_size,)](
             x,
             gamma,
@@ -161,6 +161,7 @@ def can_use_wan_rmsnorm_silu(
         and not x.requires_grad
         and x.dtype in _SUPPORTED_DTYPES
         and x.ndim == 5
+        and x.numel() > 0
         and 0 < x.shape[1] <= _MAX_CHANNELS
         and x.is_contiguous(memory_format=torch.channels_last_3d)
         and _affine_supported(x, gamma)

python/sglang/kernels/ops/diffusion/triton/zimage_native_norm.py

@@ -1,187 +1,13 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Z-Image-specific bit-exact per-head RMSNorm kernel."""
+
+from __future__ import annotations
+
 import torch
 import triton  # type: ignore
 import triton.language as tl  # type: ignore
 
 
-@triton.jit
-def _tanh(x):
-    return 2.0 / (1.0 + tl.exp(-2.0 * x)) - 1.0
-
-
-@triton.jit
-def _rmsnorm_scale_kernel(
-    y_ptr,
-    x_ptr,
-    weight_ptr,
-    scale_ptr,
-    x_row_stride,
-    scale_row_stride,
-    seq_len,
-    dim: tl.constexpr,
-    eps: tl.constexpr,
-    block_dim: tl.constexpr,
-):
-    row = tl.program_id(0)
-    offsets = tl.arange(0, block_dim)
-    mask = offsets < dim
-
-    x = tl.load(x_ptr + row * x_row_stride + offsets, mask=mask, other=0.0)
-    square = (x * x).to(tl.bfloat16)
-    mean_square = (tl.sum(square, axis=0) / dim).to(tl.bfloat16)
-    rstd = tl.rsqrt((mean_square + eps).to(tl.bfloat16).to(tl.float32)).to(tl.bfloat16)
-
-    batch = row // seq_len
-    weight = tl.load(weight_ptr + offsets, mask=mask, other=0.0)
-    scale = tl.load(
-        scale_ptr + batch * scale_row_stride + offsets, mask=mask, other=0.0
-    )
-    y = (((x * rstd).to(tl.bfloat16) * weight).to(tl.bfloat16) * scale).to(tl.bfloat16)
-    tl.store(y_ptr + row * dim + offsets, y, mask=mask)
-
-
-@triton.jit
-def _rmsnorm_tanh_residual_kernel(
-    y_ptr,
-    x_ptr,
-    gate_ptr,
-    residual_ptr,
-    weight_ptr,
-    x_row_stride,
-    gate_row_stride,
-    residual_row_stride,
-    seq_len,
-    dim: tl.constexpr,
-    eps: tl.constexpr,
-    block_dim: tl.constexpr,
-):
-    row = tl.program_id(0)
-    offsets = tl.arange(0, block_dim)
-    mask = offsets < dim
-
-    x = tl.load(x_ptr + row * x_row_stride + offsets, mask=mask, other=0.0)
-    square = (x * x).to(tl.bfloat16)
-    mean_square = (tl.sum(square, axis=0) / dim).to(tl.bfloat16)
-    rstd = tl.rsqrt((mean_square + eps).to(tl.bfloat16).to(tl.float32)).to(tl.bfloat16)
-
-    batch = row // seq_len
-    gate = tl.load(gate_ptr + batch * gate_row_stride + offsets, mask=mask, other=0.0)
-    residual = tl.load(
-        residual_ptr + row * residual_row_stride + offsets, mask=mask, other=0.0
-    )
-    weight = tl.load(weight_ptr + offsets, mask=mask, other=0.0)
-    norm = ((x * rstd).to(tl.bfloat16) * weight).to(tl.bfloat16)
-    gated = (_tanh(gate.to(tl.float32)).to(tl.bfloat16) * norm).to(tl.bfloat16)
-    y = (residual + gated).to(tl.bfloat16)
-    tl.store(y_ptr + row * dim + offsets, y, mask=mask)
-
-
-def _flat_row_stride(x: torch.Tensor) -> int | None:
-    if x.dim() < 2 or x.stride(-1) != 1:
-        return None
-    row_stride = x.stride(-2)
-    expected_stride = row_stride * x.shape[-2]
-    for dim in range(x.dim() - 3, -1, -1):
-        if x.stride(dim) != expected_stride:
-            return None
-        expected_stride *= x.shape[dim]
-    return row_stride
-
-
-def _can_use(x: torch.Tensor, weight: torch.Tensor, other: torch.Tensor) -> bool:
-    return (
-        x.is_cuda
-        and weight.is_cuda
-        and other.is_cuda
-        and x.dtype == torch.bfloat16
-        and weight.dtype == torch.bfloat16
-        and other.dtype == torch.bfloat16
-        and weight.is_contiguous()
-        and x.shape[-1] <= 8192
-        and _flat_row_stride(x) is not None
-        and _flat_row_stride(other) is not None
-    )
-
-
-def zimage_rmsnorm_scale(
-    x: torch.Tensor,
-    weight: torch.Tensor,
-    scale: torch.Tensor,
-    eps: float,
-) -> torch.Tensor | None:
-    if not _can_use(x, weight, scale):
-        return None
-    shape = x.shape
-    dim = shape[-1]
-    x_rows = x.numel() // dim
-    scale_rows = scale.numel() // dim
-    if x_rows % scale_rows != 0:
-        return None
-    seq_len = x_rows // scale_rows
-    x_row_stride = _flat_row_stride(x)
-    scale_row_stride = _flat_row_stride(scale)
-    if x_row_stride is None or scale_row_stride is None:
-        return None
-    y = torch.empty_like(x, memory_format=torch.contiguous_format)
-    with torch.get_device_module().device(x.device):
-        _rmsnorm_scale_kernel[(x_rows,)](
-            y.reshape(-1, dim),
-            x,
-            weight,
-            scale,
-            x_row_stride,
-            scale_row_stride,
-            seq_len,
-            dim,
-            eps,
-            block_dim=triton.next_power_of_2(dim),
-            num_warps=8,
-        )
-    return y
-
-
-def zimage_rmsnorm_tanh_residual(
-    x: torch.Tensor,
-    gate: torch.Tensor,
-    residual: torch.Tensor,
-    weight: torch.Tensor,
-    eps: float,
-) -> torch.Tensor | None:
-    if not (_can_use(x, weight, gate) and residual.is_cuda):
-        return None
-    if residual.dtype != x.dtype or _flat_row_stride(residual) is None:
-        return None
-    shape = x.shape
-    dim = shape[-1]
-    x_rows = x.numel() // dim
-    gate_rows = gate.numel() // dim
-    if x_rows % gate_rows != 0:
-        return None
-    seq_len = x_rows // gate_rows
-    x_row_stride = _flat_row_stride(x)
-    gate_row_stride = _flat_row_stride(gate)
-    residual_row_stride = _flat_row_stride(residual)
-    if x_row_stride is None or gate_row_stride is None or residual_row_stride is None:
-        return None
-    y = torch.empty_like(x, memory_format=torch.contiguous_format)
-    with torch.get_device_module().device(x.device):
-        _rmsnorm_tanh_residual_kernel[(x_rows,)](
-            y.reshape(-1, dim),
-            x,
-            gate,
-            residual,
-            weight,
-            x_row_stride,
-            gate_row_stride,
-            residual_row_stride,
-            seq_len,
-            dim,
-            eps,
-            block_dim=triton.next_power_of_2(dim),
-            num_warps=8,
-        )
-    return y
-
-
 @triton.jit
 def _qk_rmsnorm_native_kernel(
     y_ptr,
@@ -294,9 +120,11 @@ def can_use_qk_rmsnorm_native(
     return (
         x.is_cuda
         and weight.is_cuda
-        and x.dtype == torch.bfloat16
+        and x.device == weight.device
+        and x.dtype == weight.dtype == torch.bfloat16
+        and x.numel() > 0
         and head_dim == 128
-        and weight.numel() == head_dim
+        and weight.shape == (head_dim,)
         and weight.is_contiguous()
         and _qk_head_token_stride(x, head_dim) is not None
     )
@@ -319,8 +147,6 @@ def zimage_qk_rmsnorm_native(
     token_stride = _qk_head_token_stride(x, head_dim)
     if token_stride is None:
         return None
-    if weight.dtype != x.dtype:
-        weight = weight.to(dtype=x.dtype)
     nheads = x.shape[2]
     n_rows = x.shape[0] * x.shape[1] * nheads
     rows_per_prog = 8
@@ -340,3 +166,6 @@ def zimage_qk_rmsnorm_native(
             num_warps=8,
         )
     return y
+
+
+__all__ = ["can_use_qk_rmsnorm_native", "zimage_qk_rmsnorm_native"]

python/sglang/kernels/ops/diffusion/usp_relayout.py

@@ -16,6 +16,8 @@ _SUPPORTED_DTYPES = (torch.float16, torch.bfloat16, torch.float32)
 
 @cache_once
 def _jit_usp_relayout_module(dtype: torch.dtype) -> Module:
+    if dtype not in _SUPPORTED_DTYPES:
+        raise RuntimeError(f"Unsupported usp_merge_heads dtype: {dtype}")
     args = make_cpp_args(dtype)
     return load_jit(
         "diffusion_usp_relayout",

python/sglang/multimodal_gen/runtime/models/dits/ernie_image.py

@@ -19,10 +19,7 @@ import torch.nn as nn
 import torch.nn.functional as F
 from diffusers.models.embeddings import TimestepEmbedding, Timesteps
 
-from sglang.kernels.ops.diffusion.residual_gate_add import (
+from sglang.kernels.ops.diffusion.residual_gate_add import residual_gate_add
-    can_use_residual_gate_add_cuda,
-    residual_gate_add_cuda,
-)
 from sglang.kernels.ops.diffusion.triton.rmsnorm_scale_shift_bitexact import (
     can_use_fused_rmsnorm_scale_shift,
     can_use_fused_scale_residual_rmsnorm_scale_shift,
@@ -54,37 +51,6 @@ from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 
 logger = init_logger(__name__)
 
-_ERNIE_RESIDUAL_GATE_CUDA_DISABLED = False
-
-
-def _ernie_residual_gate_add(
-    residual: torch.Tensor,
-    update: torch.Tensor,
-    gate: torch.Tensor,
-) -> torch.Tensor:
-    """Single-kernel ``residual + gate * update``, bit-exact vs the eager pair.
-
-    Restricted to half dtypes: there the kernel reproduces the eager pair's
-    two-step rounding exactly (verified by ``torch.equal``), while for fp32 it
-    would contract to an fma (one rounding) and stop being bit-exact.
-    """
-    global _ERNIE_RESIDUAL_GATE_CUDA_DISABLED
-
-    if (
-        not _ERNIE_RESIDUAL_GATE_CUDA_DISABLED
-        and residual.dtype in (torch.float16, torch.bfloat16)
-        and can_use_residual_gate_add_cuda(residual, update, gate)
-    ):
-        try:
-            return residual_gate_add_cuda(residual, update, gate)
-        except Exception as exc:
-            if torch.compiler.is_compiling():
-                raise
-            logger.warning_once(f"Disabling ERNIE residual-gate CUDA fast path: {exc}")
-            _ERNIE_RESIDUAL_GATE_CUDA_DISABLED = True
-
-    return residual + gate * update
-
 
 _ERNIE_FUSED_NORM_DISABLED = False
 _ERNIE_FUSED_NORM_VERIFIED = False
@@ -197,7 +163,7 @@ def _ernie_gated_norm_scale_shift(
             _ERNIE_FUSED_GATED_NORM_DISABLED = True
             return ref, res_ref
 
-    res = _ernie_residual_gate_add(residual, update, gate)
+    res = residual_gate_add(residual, update, gate)
     return _eager_norm_scale_shift(norm, res, scale, shift), res
 
 
@@ -421,7 +387,7 @@ class ErnieImageSharedAdaLNBlock(nn.Module):
         x, residual = _ernie_gated_norm_scale_shift(
             self.adaLN_mlp_ln, residual, attn_out, gate_msa, scale_mlp, shift_mlp
         )
-        x = _ernie_residual_gate_add(residual, self.mlp(x), gate_mlp)
+        x = residual_gate_add(residual, self.mlp(x), gate_mlp)
 
         return x
 

python/sglang/multimodal_gen/runtime/models/dits/flux.py

@@ -30,6 +30,7 @@ from torch.nn import LayerNorm as LayerNorm
 
 from sglang.kernels.ops.diffusion.fused_linear_gelu import (
     can_fuse_linear_gelu,
+    fused_gelu_active,
     fused_linear_gelu_tanh,
     mark_fused_gelu_site,
 )
@@ -39,14 +40,8 @@ from sglang.kernels.ops.diffusion.fused_ln_modulate import (
     fused_ln_modulate_active,
     mark_fused_ln_modulate_site,
 )
-from sglang.kernels.ops.diffusion.modulate_scale_shift import (
+from sglang.kernels.ops.diffusion.modulate_scale_shift import modulate_scale_shift
-    can_use_modulate_scale_shift_cuda,
+from sglang.kernels.ops.diffusion.residual_gate_add import residual_gate_add
-    modulate_scale_shift_cuda,
-)
-from sglang.kernels.ops.diffusion.residual_gate_add import (
-    can_use_residual_gate_add_cuda,
-    residual_gate_add_cuda,
-)
 from sglang.multimodal_gen.configs.models.dits.flux import FluxConfig
 from sglang.multimodal_gen.runtime.distributed import (
     divide,
@@ -96,68 +91,6 @@ from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 
 logger = init_logger(__name__)  # pylint: disable=invalid-name
 
-_FLUX_RESIDUAL_GATE_CUDA_DISABLED = False
-
-
-def _flux_residual_gate_add(
-    residual: torch.Tensor,
-    update: torch.Tensor,
-    gate: torch.Tensor,
-) -> torch.Tensor:
-    """Single-kernel ``residual + gate * update``, bit-exact vs the eager pair.
-
-    Restricted to half dtypes: there the kernel reproduces the eager pair's
-    two-step rounding exactly (verified by ``torch.equal``), while for fp32 it
-    would contract to an fma (one rounding) and stop being bit-exact. The
-    kernel's row-broadcast gate only covers ``[1, ..., 1, D]``; batched
-    ``[B>1, 1, D]`` gates fail ``can_use_residual_gate_add_cuda`` and take the
-    eager fallback below.
-    """
-    global _FLUX_RESIDUAL_GATE_CUDA_DISABLED
-
-    if (
-        not _FLUX_RESIDUAL_GATE_CUDA_DISABLED
-        and residual.dtype in (torch.float16, torch.bfloat16)
-        and can_use_residual_gate_add_cuda(residual, update, gate)
-    ):
-        try:
-            return residual_gate_add_cuda(residual, update, gate)
-        except Exception as exc:
-            if torch.compiler.is_compiling():
-                raise
-            logger.warning_once(f"Disabling FLUX residual-gate CUDA fast path: {exc}")
-            _FLUX_RESIDUAL_GATE_CUDA_DISABLED = True
-
-    return residual + gate * update
-
-
-_FLUX_MODULATE_CUDA_DISABLED = False
-
-
-def _flux_modulate(
-    x: torch.Tensor, scale: torch.Tensor, shift: torch.Tensor
-) -> torch.Tensor:
-    """``x * (1 + scale[:, None]) + shift[:, None]`` in one CUDA kernel.
-
-    The kernel keeps the eager chain's per-op fp32-opmath/round-to-storage
-    boundaries, so it is bit-exact vs eager and needs no quality gate.
-    Guarded inputs fall back to the eager expression.
-    """
-    global _FLUX_MODULATE_CUDA_DISABLED
-
-    if not _FLUX_MODULATE_CUDA_DISABLED and can_use_modulate_scale_shift_cuda(
-        x, scale, shift
-    ):
-        try:
-            return modulate_scale_shift_cuda(x, scale, shift)
-        except Exception as exc:
-            if torch.compiler.is_compiling():
-                raise
-            logger.warning_once(f"Disabling FLUX modulate CUDA fast path: {exc}")
-            _FLUX_MODULATE_CUDA_DISABLED = True
-
-    return x * (1 + scale[:, None]) + shift[:, None]
-
 
 def _flux_norm_modulate(
     site: nn.Module,
@@ -174,7 +107,7 @@ def _flux_norm_modulate(
     """
     if fused_ln_modulate_active(site) and can_fuse_ln_modulate(x, scale, shift):
         return fused_ln_modulate(x, scale, shift, norm.eps)
-    return _flux_modulate(norm(x), scale, shift)
+    return modulate_scale_shift(norm(x), scale, shift)
 
 
 class FluxAdaLayerNormZero(AdaLayerNormZero):
@@ -412,9 +345,7 @@ class FluxGELU(nn.Module):
         mark_fused_gelu_site(self, "proj")
 
     def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
-        if self._sgl_fused_gelu_enabled and can_fuse_linear_gelu(
+        if fused_gelu_active(self) and can_fuse_linear_gelu(self.proj, hidden_states):
-            self.proj, hidden_states
-        ):
             return fused_linear_gelu_tanh(
                 hidden_states, self.proj.weight, self.proj.bias
             )
@@ -438,9 +369,7 @@ class FluxFusedGELUProj(nn.Module):
         mark_fused_gelu_site(self, "proj")
 
     def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
-        if self._sgl_fused_gelu_enabled and can_fuse_linear_gelu(
+        if fused_gelu_active(self) and can_fuse_linear_gelu(self.proj, hidden_states):
-            self.proj, hidden_states
-        ):
             return fused_linear_gelu_tanh(
                 hidden_states, self.proj.weight, self.proj.bias
             )
@@ -916,7 +845,7 @@ class FluxSingleTransformerBlock(nn.Module):
             hidden_states = gate * hidden_states
             hidden_states = residual + hidden_states
         else:
-            if self._sgl_fused_gelu_enabled and can_fuse_linear_gelu(
+            if fused_gelu_active(self) and can_fuse_linear_gelu(
                 self.proj_mlp, norm_hidden_states
             ):
                 mlp_hidden_states = fused_linear_gelu_tanh(
@@ -936,7 +865,7 @@ class FluxSingleTransformerBlock(nn.Module):
             hidden_states = torch.cat([attn_output, mlp_hidden_states], dim=2)
             gate = gate.unsqueeze(1)
             proj_out, _ = self.proj_out(hidden_states)
-            hidden_states = _flux_residual_gate_add(residual, proj_out, gate)
+            hidden_states = residual_gate_add(residual, proj_out, gate)
 
         if hidden_states.dtype == torch.float16:
             hidden_states = hidden_states.clip(-65504, 65504)
@@ -1074,7 +1003,7 @@ class FluxTransformerBlock(nn.Module):
             attn_output, context_attn_output, ip_attn_output = attention_outputs
 
         # Process attention outputs for the `hidden_states`.
-        hidden_states = _flux_residual_gate_add(
+        hidden_states = residual_gate_add(
             hidden_states, attn_output, gate_msa.unsqueeze(1)
         )
         if self.use_nunchaku_structure:
@@ -1088,14 +1017,14 @@ class FluxTransformerBlock(nn.Module):
             )
 
         ff_output = self.ff(norm_hidden_states)
-        hidden_states = _flux_residual_gate_add(
+        hidden_states = residual_gate_add(
             hidden_states, ff_output, gate_mlp.unsqueeze(1)
         )
 
         if len(attention_outputs) == 3:
             hidden_states = hidden_states + ip_attn_output
         # Process attention outputs for the `encoder_hidden_states`.
-        encoder_hidden_states = _flux_residual_gate_add(
+        encoder_hidden_states = residual_gate_add(
             encoder_hidden_states, context_attn_output, c_gate_msa.unsqueeze(1)
         )
 
@@ -1114,7 +1043,7 @@ class FluxTransformerBlock(nn.Module):
             )
 
         context_ff_output = self.ff_context(norm_encoder_hidden_states)
-        encoder_hidden_states = _flux_residual_gate_add(
+        encoder_hidden_states = residual_gate_add(
             encoder_hidden_states, context_ff_output, c_gate_mlp.unsqueeze(1)
         )
         if encoder_hidden_states.dtype == torch.float16:

python/sglang/multimodal_gen/runtime/models/dits/flux_2.py

@@ -20,10 +20,7 @@ from diffusers.models.attention import AttentionModuleMixin
 from diffusers.models.embeddings import TimestepEmbedding, Timesteps
 from diffusers.models.normalization import AdaLayerNormContinuous
 
-from sglang.kernels.ops.diffusion.residual_gate_add import (
+from sglang.kernels.ops.diffusion.residual_gate_add import residual_gate_add
-    can_use_residual_gate_add_cuda,
-    residual_gate_add_cuda,
-)
 from sglang.multimodal_gen.configs.models.dits.flux import FluxConfig
 from sglang.multimodal_gen.runtime.distributed import (
     divide,
@@ -69,40 +66,6 @@ from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 
 logger = init_logger(__name__)  # pylint: disable=invalid-name
 
-_FLUX2_RESIDUAL_GATE_CUDA_DISABLED = False
-
-
-def _flux2_residual_gate_add(
-    residual: torch.Tensor,
-    update: torch.Tensor,
-    gate: torch.Tensor,
-) -> torch.Tensor:
-    """Single-kernel ``residual + gate * update``, bit-exact vs the eager pair.
-
-    Restricted to half dtypes: there the kernel reproduces the eager pair's
-    two-step rounding exactly (verified by ``torch.equal``), while for fp32 it
-    would contract to an fma (one rounding) and stop being bit-exact. The
-    kernel's row-broadcast gate only covers ``[1, ..., 1, D]``; batched
-    ``[B>1, 1, D]`` gates fail ``can_use_residual_gate_add_cuda`` and take the
-    eager fallback below.
-    """
-    global _FLUX2_RESIDUAL_GATE_CUDA_DISABLED
-
-    if (
-        not _FLUX2_RESIDUAL_GATE_CUDA_DISABLED
-        and residual.dtype in (torch.float16, torch.bfloat16)
-        and can_use_residual_gate_add_cuda(residual, update, gate)
-    ):
-        try:
-            return residual_gate_add_cuda(residual, update, gate)
-        except Exception as exc:
-            if torch.compiler.is_compiling():
-                raise
-            logger.warning_once(f"Disabling FLUX.2 residual-gate CUDA fast path: {exc}")
-            _FLUX2_RESIDUAL_GATE_CUDA_DISABLED = True
-
-    return residual + gate * update
-
 
 def _get_qkv_projections(
     attn: "Flux2Attention", hidden_states, encoder_hidden_states=None
@@ -694,7 +657,7 @@ class Flux2SingleTransformerBlock(nn.Module):
             **joint_attention_kwargs,
         )
 
-        hidden_states = _flux2_residual_gate_add(hidden_states, attn_output, mod_gate)
+        hidden_states = residual_gate_add(hidden_states, attn_output, mod_gate)
         if hidden_states.dtype == torch.float16:
             hidden_states = hidden_states.clip(-65504, 65504)
 
@@ -779,15 +742,21 @@ class Flux2TransformerBlock(nn.Module):
         joint_attention_kwargs = joint_attention_kwargs or {}
 
         # Modulation parameters shape: [1, 1, self.dim]
-        (shift_msa, scale_msa, gate_msa), (
+        (
+            (shift_msa, scale_msa, gate_msa),
+            (
                 shift_mlp,
                 scale_mlp,
                 gate_mlp,
+            ),
         ) = temb_mod_params_img
-        (c_shift_msa, c_scale_msa, c_gate_msa), (
+        (
+            (c_shift_msa, c_scale_msa, c_gate_msa),
+            (
                 c_shift_mlp,
                 c_scale_mlp,
                 c_gate_mlp,
+            ),
         ) = temb_mod_params_txt
 
         # Img stream
@@ -812,16 +781,16 @@ class Flux2TransformerBlock(nn.Module):
         attn_output, context_attn_output = attention_outputs
 
         # Process attention outputs for the image stream (`hidden_states`).
-        hidden_states = _flux2_residual_gate_add(hidden_states, attn_output, gate_msa)
+        hidden_states = residual_gate_add(hidden_states, attn_output, gate_msa)
 
         norm_hidden_states = self.norm2(hidden_states)
         norm_hidden_states = norm_hidden_states * (1 + scale_mlp) + shift_mlp
 
         ff_output = self.ff(norm_hidden_states)
-        hidden_states = _flux2_residual_gate_add(hidden_states, ff_output, gate_mlp)
+        hidden_states = residual_gate_add(hidden_states, ff_output, gate_mlp)
 
         # Process attention outputs for the text stream (`encoder_hidden_states`).
-        encoder_hidden_states = _flux2_residual_gate_add(
+        encoder_hidden_states = residual_gate_add(
             encoder_hidden_states, context_attn_output, c_gate_msa
         )
 
@@ -831,7 +800,7 @@ class Flux2TransformerBlock(nn.Module):
         )
 
         context_ff_output = self.ff_context(norm_encoder_hidden_states)
-        encoder_hidden_states = _flux2_residual_gate_add(
+        encoder_hidden_states = residual_gate_add(
             encoder_hidden_states, context_ff_output, c_gate_mlp
         )
         if encoder_hidden_states.dtype == torch.float16:

python/sglang/multimodal_gen/runtime/models/dits/glm_image.py

@@ -20,13 +20,11 @@ import torch.nn.functional as F
 
 from sglang.kernels.ops.diffusion.fused_linear_gelu import (
     can_fuse_linear_gelu,
+    fused_gelu_active,
     fused_linear_gelu_tanh,
     mark_fused_gelu_site,
 )
-from sglang.kernels.ops.diffusion.residual_gate_add import (
+from sglang.kernels.ops.diffusion.residual_gate_add import residual_gate_add
-    can_use_residual_gate_add_cuda,
-    residual_gate_add_cuda,
-)
 from sglang.kernels.ops.diffusion.triton.layernorm_modulate import (
     can_use_fused_layernorm_modulate,
     can_use_fused_qk_head_layernorm,
@@ -79,7 +77,6 @@ _GLM_FUSED_LN_MOD_DISABLED = False
 _GLM_FUSED_LN_MOD_VERIFIED = False
 _GLM_FUSED_QK_LN_DISABLED = False
 _GLM_FUSED_QK_LN_VERIFIED = False
-_GLM_RESIDUAL_GATE_CUDA_DISABLED = False
 
 
 def _eager_ln_modulate(
@@ -189,34 +186,6 @@ def _glm_qk_layernorm(
     return norm_q(query).to(dtype=dtype), norm_k(key).to(dtype=dtype)
 
 
-def _glm_residual_gate_add(
-    residual: torch.Tensor,
-    update: torch.Tensor,
-    gate: torch.Tensor,
-) -> torch.Tensor:
-    """Single-kernel ``residual + gate * update``, bit-exact vs the eager pair.
-
-    Half dtypes only: for fp32 the kernel would contract to an fma (one
-    rounding) and stop being bit-exact.
-    """
-    global _GLM_RESIDUAL_GATE_CUDA_DISABLED
-
-    if (
-        not _GLM_RESIDUAL_GATE_CUDA_DISABLED
-        and residual.dtype in (torch.float16, torch.bfloat16)
-        and can_use_residual_gate_add_cuda(residual, update, gate)
-    ):
-        try:
-            return residual_gate_add_cuda(residual, update, gate)
-        except Exception as exc:
-            if torch.compiler.is_compiling():
-                raise
-            logger.warning_once(f"Disabling GLM residual-gate CUDA fast path: {exc}")
-            _GLM_RESIDUAL_GATE_CUDA_DISABLED = True
-
-    return residual + gate * update
-
-
 class GlmImageLayerKVCache:
     """KV cache for GlmImage model."""
 
@@ -491,9 +460,7 @@ class GlmImageGELU(nn.Module):
         mark_fused_gelu_site(self, "proj")
 
     def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
-        if self._sgl_fused_gelu_enabled and can_fuse_linear_gelu(
+        if fused_gelu_active(self) and can_fuse_linear_gelu(self.proj, hidden_states):
-            self.proj, hidden_states
-        ):
             return fused_linear_gelu_tanh(
                 hidden_states, self.proj.weight, self.proj.bias
             )
@@ -837,10 +804,10 @@ class GlmImageTransformerBlock(nn.Module):
 
         ff_output = self.ff(norm_hidden_states)
         ff_output_context = self.ff(norm_encoder_hidden_states)
-        hidden_states = _glm_residual_gate_add(
+        hidden_states = residual_gate_add(
             hidden_states, ff_output, gate_mlp.unsqueeze(1)
         )
-        encoder_hidden_states = _glm_residual_gate_add(
+        encoder_hidden_states = residual_gate_add(
             encoder_hidden_states, ff_output_context, c_gate_mlp.unsqueeze(1)
         )
 

python/sglang/multimodal_gen/runtime/models/dits/ideogram.py

@@ -8,6 +8,7 @@ import torch.nn as nn
 import torch.nn.functional as F
 
 from sglang.kernels.ops.diffusion.fused_gate_rmsnorm import (
+    fused_gate_rmsnorm_active,
     fused_rmsnorm_scale,
     fused_rmsnorm_tanh_residual,
     mark_fused_gate_rmsnorm_site,
@@ -394,7 +395,7 @@ class Ideogram4TransformerBlock(nn.Module):
             adaln_input
         ).chunk(4, dim=-1)
         enable_fused = (
-            self._sgl_fused_gate_rmsnorm_enabled and not torch.compiler.is_compiling()
+            fused_gate_rmsnorm_active(self) and not torch.compiler.is_compiling()
         )
         attn_out = self.attention(
             _norm_scale(x, scale_msa, self.attention_norm1, enable_fused),

python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py

@@ -14,10 +14,7 @@ from sglang.kernels.ops.diffusion.ltx2_qknorm_split_rope import (
     can_use_ltx2_qknorm_split_rope_cuda,
     ltx2_qknorm_split_rope_cuda,
 )
-from sglang.kernels.ops.diffusion.residual_gate_add import (
+from sglang.kernels.ops.diffusion.residual_gate_add import residual_gate_add
-    can_use_residual_gate_add_cuda,
-    residual_gate_add_cuda,
-)
 from sglang.multimodal_gen.configs.models.dits.ltx_2 import LTX2ArchConfig, LTX2Config
 from sglang.multimodal_gen.runtime.distributed import (
     get_sp_parallel_rank,
@@ -56,31 +53,9 @@ logger = init_logger(__name__)
 
 ADALN_NUM_BASE_PARAMS = 6
 ADALN_NUM_CROSS_ATTN_PARAMS = 3
-_LTX2_RESIDUAL_GATE_CUDA_DISABLED = False
 _LTX2_QKNORM_SPLIT_ROPE_CUDA_DISABLED = False
 
 
-def _ltx2_residual_gate_add(
-    residual: torch.Tensor,
-    update: torch.Tensor,
-    gate: torch.Tensor,
-) -> torch.Tensor:
-    global _LTX2_RESIDUAL_GATE_CUDA_DISABLED
-
-    if not _LTX2_RESIDUAL_GATE_CUDA_DISABLED and can_use_residual_gate_add_cuda(
-        residual, update, gate
-    ):
-        try:
-            return residual_gate_add_cuda(residual, update, gate)
-        except Exception as exc:
-            if torch.compiler.is_compiling():
-                raise
-            logger.warning_once(f"Disabling LTX2 residual-gate CUDA fast path: {exc}")
-            _LTX2_RESIDUAL_GATE_CUDA_DISABLED = True
-
-    return residual + update * gate
-
-
 def _ltx2_try_fused_qknorm_split_rope(
     q: torch.Tensor,
     k: torch.Tensor,
@@ -1237,9 +1212,7 @@ class LTX2TransformerBlock(nn.Module):
             gather_context_kv_for_sp=audio_replicated_for_sp,
             context_replicated_prefix_len=video_memory_prefix_len,
         )
-        hidden_states = _ltx2_residual_gate_add(
+        hidden_states = residual_gate_add(hidden_states, attn_hidden_states, vgate_msa)
-            hidden_states, attn_hidden_states, vgate_msa
-        )
 
         if audio_ada_values is None:
             ashift_msa, ascale_msa, agate_msa = self.get_ada_values(
@@ -1259,7 +1232,7 @@ class LTX2TransformerBlock(nn.Module):
             all_perturbed=skip_audio_self_attn,
             skip_sequence_parallel_override=audio_replicated_for_sp,
         )
-        audio_hidden_states = _ltx2_residual_gate_add(
+        audio_hidden_states = residual_gate_add(
             audio_hidden_states, attn_audio_hidden_states, agate_msa
         )
         # 2. Prompt Cross-Attention
@@ -1289,7 +1262,7 @@ class LTX2TransformerBlock(nn.Module):
                 context=mod_encoder_hidden_states,
                 mask=encoder_attention_mask,
             )
-            hidden_states = _ltx2_residual_gate_add(
+            hidden_states = residual_gate_add(
                 hidden_states, attn_hidden_states, vgate_q
             )
 
@@ -1317,7 +1290,7 @@ class LTX2TransformerBlock(nn.Module):
                 context=mod_audio_encoder_hidden_states,
                 mask=audio_encoder_attention_mask,
             )
-            audio_hidden_states = _ltx2_residual_gate_add(
+            audio_hidden_states = residual_gate_add(
                 audio_hidden_states, attn_audio_hidden_states, agate_q
             )
         else:
@@ -1419,7 +1392,7 @@ class LTX2TransformerBlock(nn.Module):
                 a2v_attn_hidden_states = (
                     a2v_attn_hidden_states * a2v_cross_attn_perturbation_mask
                 )
-            hidden_states = _ltx2_residual_gate_add(
+            hidden_states = residual_gate_add(
                 hidden_states, a2v_attn_hidden_states, a2v_gate
             )
 
@@ -1445,7 +1418,7 @@ class LTX2TransformerBlock(nn.Module):
                 v2a_attn_hidden_states = (
                     v2a_attn_hidden_states * v2a_cross_attn_perturbation_mask
                 )
-            audio_hidden_states = _ltx2_residual_gate_add(
+            audio_hidden_states = residual_gate_add(
                 audio_hidden_states, v2a_attn_hidden_states, v2a_gate
             )
         # 4. Feedforward
@@ -1459,7 +1432,7 @@ class LTX2TransformerBlock(nn.Module):
             self.rms_norm(hidden_states, self.norm_eps) * (1 + vscale_mlp) + vshift_mlp
         )
         ff_output = self.ff(norm_hidden_states)
-        hidden_states = _ltx2_residual_gate_add(hidden_states, ff_output, vgate_mlp)
+        hidden_states = residual_gate_add(hidden_states, ff_output, vgate_mlp)
 
         if audio_ada_values is None:
             ashift_mlp, ascale_mlp, agate_mlp = self.get_ada_values(
@@ -1472,7 +1445,7 @@ class LTX2TransformerBlock(nn.Module):
             + ashift_mlp
         )
         audio_ff_output = self.audio_ff(norm_audio_hidden_states)
-        audio_hidden_states = _ltx2_residual_gate_add(
+        audio_hidden_states = residual_gate_add(
             audio_hidden_states, audio_ff_output, agate_mlp
         )
         return hidden_states, audio_hidden_states

python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py

@@ -16,6 +16,7 @@ from diffusers.models.normalization import AdaLayerNormContinuous
 
 from sglang.kernels.ops.diffusion.fused_linear_gelu import (
     can_fuse_linear_gelu,
+    fused_gelu_active,
     fused_linear_gelu_tanh,
     mark_fused_gelu_site,
 )
@@ -864,9 +865,7 @@ class QwenImageGELU(nn.Module):
         mark_fused_gelu_site(self, "proj")
 
     def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
-        if self._sgl_fused_gelu_enabled and can_fuse_linear_gelu(
+        if fused_gelu_active(self) and can_fuse_linear_gelu(self.proj, hidden_states):
-            self.proj, hidden_states
-        ):
             return fused_linear_gelu_tanh(
                 hidden_states, self.proj.weight, self.proj.bias
             )

python/sglang/multimodal_gen/runtime/models/dits/zimage.py

@@ -83,11 +83,11 @@ def zimage_rmsnorm_tanh_mul_add(
     enable_fused: bool = True,
 ) -> torch.Tensor:
     if enable_fused:
-        from sglang.kernels.ops.diffusion.triton.zimage_native_norm import (
+        from sglang.kernels.ops.diffusion.triton.native_bf16_rmsnorm import (
-            zimage_rmsnorm_tanh_residual,
+            rmsnorm_tanh_residual,
         )
 
-        y = zimage_rmsnorm_tanh_residual(
+        y = rmsnorm_tanh_residual(
             x,
             gate,
             residual,
@@ -106,11 +106,11 @@ def zimage_rmsnorm_scale(
     enable_fused: bool = True,
 ) -> torch.Tensor:
     if enable_fused:
-        from sglang.kernels.ops.diffusion.triton.zimage_native_norm import (
+        from sglang.kernels.ops.diffusion.triton.native_bf16_rmsnorm import (
-            zimage_rmsnorm_scale as fused_zimage_rmsnorm_scale,
+            rmsnorm_scale,
         )
 
-        y = fused_zimage_rmsnorm_scale(
+        y = rmsnorm_scale(
             x,
             norm.weight.data.to(device=x.device, dtype=x.dtype).contiguous(),
             scale,

python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py

@@ -143,6 +143,26 @@ from sglang.multimodal_gen.runtime.utils.torch_compile import (
 
 logger = init_logger(__name__)
 
+_QUALITY_FUSION_HANDLERS: tuple[
+    tuple[str, Callable[[nn.Module], bool], Callable[[nn.Module], None]], ...
+] = (
+    (
+        "fused linear+GELU (cublasLt epilogue)",
+        mount_fused_linear_gelu,
+        unmount_fused_linear_gelu,
+    ),
+    (
+        "fused LN+modulate (affine folding)",
+        mount_fused_ln_modulate,
+        unmount_fused_ln_modulate,
+    ),
+    (
+        "fused gate RMSNorm (BF16-native Triton)",
+        mount_fused_gate_rmsnorm,
+        unmount_fused_gate_rmsnorm,
+    ),
+)
+
 
 def _ensure_tensor_model_output(model_output):
     sample = getattr(model_output, "sample", None)
@@ -233,8 +253,7 @@ class DenoisingStage(PipelineStage, RolloutDenoisingMixin):
         # cache-dit state (for delayed mounting and idempotent control)
         self._cache_dit_enabled = False
         self._cached_num_steps = None
-        # quality="high" fusion state: whether the cublasLt linear+GELU and
+        # Whether request-scoped quality="high" fusions are currently mounted.
-        # fused gate-RMSNorm sites are currently mounted on the transformers.
         self._quality_fusions_mounted = False
         self._torch_compile_registry = CompiledModuleRegistry()
         # Breakable CUDA graph runners, one per transformer module (lazy).
@@ -465,44 +484,28 @@ class DenoisingStage(PipelineStage, RolloutDenoisingMixin):
     def _maybe_toggle_quality_fusions(self, batch: Req) -> None:
         """Mount/unmount the ``quality="high"`` fusions for this batch.
 
-        The cublasLt linear+GELU epilogue and the fused gate-RMSNorm Triton
+        These fusions are numerically equivalent only at half-precision
-        kernels are numerically equivalent only at half-precision rounding
+        rounding level (not bit-exact), so they are mounted for
-        level (not bit-exact), so they are mounted for ``quality="high"``
+        ``quality="high"`` requests and unmounted otherwise. The
-        requests and unmounted otherwise -- the ``"lossless"`` default runs
+        ``"lossless"`` default runs the reference path bit-for-bit. ``quality``
-        the unmodified reference path bit-for-bit. ``quality`` participates
+        participates in the dynamic-batch signature, making this transition
-        in the dynamic-batch signature, so a worker batch is uniform in
+        safe at the batch boundary. Mounting is all-or-nothing per transformer
-        ``quality`` and this process-wide transition is safe at the batch
+        and fusion family; models without marked sites are no-ops.
-        boundary. Mounting is all-or-nothing per transformer and per fusion
-        family (any ineligible marked site keeps the whole transformer on
-        that family's reference path); models without marked sites are
-        no-ops.
         """
         want = getattr(batch.sampling_params, "quality", "lossless") == "high"
         if want == self._quality_fusions_mounted:
             return
-        mounted_gelu = False
+        mounted_fusions: set[str] = set()
-        mounted_gate_norm = False
-        mounted_ln_modulate = False
         for transformer in filter(None, [self.transformer, self.transformer_2]):
+            for description, mount, unmount in _QUALITY_FUSION_HANDLERS:
                 if want:
-                mounted_gelu |= mount_fused_linear_gelu(transformer)
+                    if mount(transformer):
-                mounted_gate_norm |= mount_fused_gate_rmsnorm(transformer)
+                        mounted_fusions.add(description)
-                mounted_ln_modulate |= mount_fused_ln_modulate(transformer)
                 else:
-                unmount_fused_linear_gelu(transformer)
+                    unmount(transformer)
-                unmount_fused_gate_rmsnorm(transformer)
-                unmount_fused_ln_modulate(transformer)
         self._quality_fusions_mounted = want
-        if want and mounted_gelu:
+        for description in sorted(mounted_fusions):
-            logger.info(
+            logger.info("Mounted %s for quality=high", description)
-                "Mounted fused linear+GELU (cublasLt epilogue) for quality=high"
-            )
-        if want and mounted_ln_modulate:
-            logger.info("Mounted fused LN+modulate (affine folding) for quality=high")
-        if want and mounted_gate_norm:
-            logger.info(
-                "Mounted fused gate RMSNorm (Z-Image Triton suite) for quality=high"
-            )
 
     def _cache_dit_dual_model_name(self) -> str:
         return "wan2.2"

test/registered/kernels/ops/diffusion/test_ernie_residual_gate_add.py

@@ -1,35 +0,0 @@
-"""ERNIE residual-gate fast path must stay bit-exact vs the eager pair."""
-
-import sys
-
-import pytest
-import torch
-
-from sglang.multimodal_gen.runtime.models.dits.ernie_image import (
-    _ernie_residual_gate_add,
-)
-from sglang.test.ci.ci_register import register_cuda_ci
-
-register_cuda_ci(est_time=4, stage="base-b-kernel-unit", runner_config="1-gpu-large")
-pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
-
-
-@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16, torch.float32])
-def test_residual_gate_add_is_bit_exact(dtype):
-    # Real ERNIE-Image shapes: hidden 4096, 1024^2 image tokens + text tokens.
-    # fp32 exercises the eager fallback (fast path is half-dtype only).
-    torch.manual_seed(0)
-    residual = torch.randn(1, 4216, 4096, device="cuda", dtype=dtype)
-    update = torch.randn_like(residual)
-    gate = torch.randn(1, 1, 4096, device="cuda", dtype=dtype)
-    out = _ernie_residual_gate_add(residual, update, gate)
-    assert torch.equal(out, residual + gate * update)
-
-    # Full-shape gate takes the same kernel path and must stay exact too.
-    gate_full = gate.expand_as(residual).contiguous()
-    out_full = _ernie_residual_gate_add(residual, update, gate_full)
-    assert torch.equal(out_full, residual + gate_full * update)
-
-
-if __name__ == "__main__":
-    sys.exit(pytest.main([__file__]))

test/registered/kernels/ops/diffusion/test_flux2_vae_fastpath.py

@@ -32,6 +32,11 @@ def test_flux2_vae_fastpath():
         gn_kernel.group_norm_silu_4d(x.contiguous(), gn.weight, gn.bias, 32, 1e-6)
         is None
     )
+    assert gn_kernel.group_norm_silu_4d(x, gn.weight.cpu(), gn.bias, 32, 1e-6) is None
+    assert (
+        gn_kernel.group_norm_silu_4d(x[..., :0, :], gn.weight, gn.bias, 32, 1e-6)
+        is None
+    )
 
     gate.enabled = True
     fast = fused_gn(x)

test/registered/kernels/ops/diffusion/test_fused_gate_rmsnorm.py

@@ -1,4 +1,4 @@
-"""Core checks for the quality-gated fused gate-RMSNorm (Z-Image suite reuse)."""
+"""Core checks for the quality-gated fused gate-RMSNorm path."""
 
 import sys
 
@@ -44,10 +44,10 @@ def test_fused_matches_ideogram_reference():
 def test_mount_guards_all_or_nothing():
     good, bad = _Site(), _Site(torch.float32)
     assert not fgn.mount_fused_gate_rmsnorm(nn.ModuleList([good, bad]))
-    assert not good._sgl_fused_gate_rmsnorm_enabled
+    assert not fgn.fused_gate_rmsnorm_active(good)
     assert fgn.mount_fused_gate_rmsnorm(good)
     fgn.unmount_fused_gate_rmsnorm(good)
-    assert not good._sgl_fused_gate_rmsnorm_enabled
+    assert not fgn.fused_gate_rmsnorm_active(good)
 
 
 if __name__ == "__main__":

test/registered/kernels/ops/diffusion/test_fused_linear_gelu.py

@@ -21,7 +21,7 @@ class _Site(nn.Module):
         gelu.mark_fused_gelu_site(self, "proj")
 
     def forward(self, x):
-        if self._sgl_fused_gelu_enabled and gelu.can_fuse_linear_gelu(self.proj, x):
+        if gelu.fused_gelu_active(self) and gelu.can_fuse_linear_gelu(self.proj, x):
             return gelu.fused_linear_gelu_tanh(x, self.proj.weight, self.proj.bias)
         return F.gelu(self.proj(x), approximate="tanh")
 
@@ -59,7 +59,7 @@ def test_mount_guards_and_lossless_path():
     good, bad = _Site(), _Site(torch.float32)
     model = nn.ModuleList([good, bad])
     assert not gelu.mount_fused_linear_gelu(model)
-    assert not good._sgl_fused_gelu_enabled
+    assert not gelu.fused_gelu_active(good)
 
     x = torch.randn(16, 64, device="cuda", dtype=torch.bfloat16)
     ref = good(x)
@@ -72,5 +72,15 @@ def test_mount_guards_and_lossless_path():
     assert not gelu.can_fuse_linear_gelu(good.proj, x.float())
 
 
+@torch.no_grad()
+def test_mounted_site_torch_compile_fullgraph():
+    site = _Site()
+    x = torch.randn(16, 64, device="cuda", dtype=torch.bfloat16)
+    assert gelu.mount_fused_linear_gelu(site)
+    expected = site(x)
+    actual = torch.compile(site, fullgraph=True)(x)
+    torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2)
+
+
 if __name__ == "__main__":
     sys.exit(pytest.main([__file__]))

test/registered/kernels/ops/diffusion/test_fused_ln_modulate.py

@@ -50,6 +50,32 @@ def test_fused_ln_modulate_guards_and_mount_protocol():
     assert not mount_fused_ln_modulate(nn.Module())  # no marked sites
 
 
+@torch.no_grad()
+def test_mounted_ln_modulate_site_torch_compile_fullgraph():
+    class Site(nn.Module):
+        def __init__(self):
+            super().__init__()
+            mark_fused_ln_modulate_site(self)
+
+        def forward(self, x, scale, shift):
+            if fused_ln_modulate_active(self) and can_fuse_ln_modulate(x, scale, shift):
+                return fused_ln_modulate(x, scale, shift, eps=1e-6)
+            return (
+                nn.functional.layer_norm(x, (x.shape[-1],), eps=1e-6)
+                * (1 + scale[:, None])
+                + shift[:, None]
+            )
+
+    site = Site()
+    assert mount_fused_ln_modulate(site)
+    x = torch.randn(1, 64, 128, device="cuda", dtype=torch.bfloat16)
+    scale = torch.randn(1, 128, device="cuda", dtype=torch.bfloat16)
+    shift = torch.randn_like(scale)
+    expected = site(x, scale, shift)
+    actual = torch.compile(site, fullgraph=True)(x, scale, shift)
+    torch.testing.assert_close(actual, expected, atol=0.0625, rtol=0.05)
+
+
 if __name__ == "__main__":
     import sys
 

test/registered/kernels/ops/diffusion/test_modulate_scale_shift.py

@@ -3,6 +3,7 @@ import torch
 
 from sglang.kernels.ops.diffusion.modulate_scale_shift import (
     can_use_modulate_scale_shift_cuda,
+    modulate_scale_shift,
     modulate_scale_shift_cuda,
 )
 from sglang.test.ci.ci_register import register_cuda_ci
@@ -48,6 +49,7 @@ def test_modulate_scale_shift_guards_reject_fp32():
     x = torch.randn((1, 64, 64), device="cuda", dtype=torch.float32)
     row = torch.randn((1, 64), device="cuda", dtype=torch.float32)
     assert not can_use_modulate_scale_shift_cuda(x, row, row)
+    assert torch.equal(modulate_scale_shift(x, row, row), _eager(x, row, row))
 
 
 if __name__ == "__main__":

test/registered/kernels/ops/diffusion/test_native_bf16_rmsnorm.py

@@ -1,9 +1,9 @@
 import pytest
 import torch
 
-from sglang.kernels.ops.diffusion.triton.zimage_native_norm import (
+from sglang.kernels.ops.diffusion.triton.native_bf16_rmsnorm import (
-    zimage_rmsnorm_scale,
+    rmsnorm_scale,
-    zimage_rmsnorm_tanh_residual,
+    rmsnorm_tanh_residual,
 )
 from sglang.test.ci.ci_register import register_cuda_ci
 
@@ -21,26 +21,28 @@ def _native_bf16_rmsnorm(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor:
     return ((x * rstd).to(torch.bfloat16) * weight).to(torch.bfloat16)
 
 
-def test_zimage_native_norm_rejects_cpu_inputs():
+def test_native_bf16_rmsnorm_rejects_unsupported_inputs():
     x = torch.randn(2, 3, 16, dtype=torch.bfloat16)
     weight = torch.randn(16, dtype=torch.bfloat16)
     modulation = torch.randn(2, 1, 16, dtype=torch.bfloat16)
     residual = torch.randn_like(x)
 
-    assert zimage_rmsnorm_scale(x, weight, modulation, EPS) is None
+    assert rmsnorm_scale(x, weight, modulation, EPS) is None
-    assert zimage_rmsnorm_tanh_residual(x, modulation, residual, weight, EPS) is None
+    assert rmsnorm_tanh_residual(x, modulation, residual, weight, EPS) is None
+    assert rmsnorm_scale(x, weight[:-1], modulation, EPS) is None
+    assert rmsnorm_tanh_residual(x, modulation, residual[..., :-1], weight, EPS) is None
 
 
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
 @pytest.mark.parametrize("shape", [(1, 32, 2560), (2, 17, 256)])
-def test_zimage_rmsnorm_scale_matches_native_bf16(shape):
+def test_rmsnorm_scale_matches_native_bf16(shape):
     torch.manual_seed(0)
     batch, _, dim = shape
     x = torch.randn(shape, device="cuda", dtype=torch.bfloat16)
     weight = torch.randn(dim, device="cuda", dtype=torch.bfloat16)
     scale = torch.randn(batch, 1, dim, device="cuda", dtype=torch.bfloat16)
 
-    actual = zimage_rmsnorm_scale(x, weight, scale, EPS)
+    actual = rmsnorm_scale(x, weight, scale, EPS)
     expected = (_native_bf16_rmsnorm(x, weight) * scale).to(torch.bfloat16)
 
     assert actual is not None
@@ -49,7 +51,7 @@ def test_zimage_rmsnorm_scale_matches_native_bf16(shape):
 
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
 @pytest.mark.parametrize("shape", [(1, 32, 2560), (2, 17, 256)])
-def test_zimage_rmsnorm_tanh_residual_matches_native_bf16(shape):
+def test_rmsnorm_tanh_residual_matches_native_bf16(shape):
     torch.manual_seed(0)
     batch, _, dim = shape
     x = torch.randn(shape, device="cuda", dtype=torch.bfloat16)
@@ -57,7 +59,7 @@ def test_zimage_rmsnorm_tanh_residual_matches_native_bf16(shape):
     residual = torch.randn(shape, device="cuda", dtype=torch.bfloat16)
     weight = torch.randn(dim, device="cuda", dtype=torch.bfloat16)
 
-    actual = zimage_rmsnorm_tanh_residual(x, gate, residual, weight, EPS)
+    actual = rmsnorm_tanh_residual(x, gate, residual, weight, EPS)
     norm = _native_bf16_rmsnorm(x, weight)
     gated = (torch.tanh(gate.float()).to(torch.bfloat16) * norm).to(torch.bfloat16)
     expected = (residual + gated).to(torch.bfloat16)
@@ -68,15 +70,15 @@ def test_zimage_rmsnorm_tanh_residual_matches_native_bf16(shape):
 
 
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
-def test_zimage_native_norm_rejects_hidden_size_above_limit():
+def test_native_bf16_rmsnorm_rejects_hidden_size_above_limit():
     dim = 8448
     x = torch.empty(1, 1, dim, device="cuda", dtype=torch.bfloat16)
     weight = torch.empty(dim, device="cuda", dtype=torch.bfloat16)
     modulation = torch.empty(1, 1, dim, device="cuda", dtype=torch.bfloat16)
     residual = torch.empty_like(x)
 
-    assert zimage_rmsnorm_scale(x, weight, modulation, EPS) is None
+    assert rmsnorm_scale(x, weight, modulation, EPS) is None
-    assert zimage_rmsnorm_tanh_residual(x, modulation, residual, weight, EPS) is None
+    assert rmsnorm_tanh_residual(x, modulation, residual, weight, EPS) is None
 
 
 if __name__ == "__main__":

test/registered/kernels/ops/diffusion/test_qknorm_rope.py

@@ -84,6 +84,17 @@ def fused_qknorm_rope(
     )
 
 
+def test_qknorm_rope_rejects_unsupported_dtypes() -> None:
+    from sglang.kernels.ops.diffusion.qknorm_rope import (
+        can_use_fused_inplace_qknorm_rope,
+    )
+
+    assert not can_use_fused_inplace_qknorm_rope(128, 128, False, torch.float32)
+    assert not can_use_fused_inplace_qknorm_rope(
+        128, 128, False, torch.bfloat16, torch.float64
+    )
+
+
 BS_LIST = [2**n for n in range(13)]
 BS_LIST += [x + 1 for x in BS_LIST]
 BS_LIST = get_ci_test_range(BS_LIST, [1, 9, 129, 257, 2049, 4097])
@@ -205,5 +216,28 @@ def test_qknorm_rope_preserves_split_bf16_rounding() -> None:
     assert torch.equal(k_ref, k_fused)
 
 
+def test_qknorm_rope_accepts_empty_token_dimension() -> None:
+    from sglang.kernels.ops.diffusion.qknorm_rope import fused_inplace_qknorm_rope
+
+    num_heads, head_dim = 8, 128
+    q = torch.empty(0, num_heads, head_dim, device=DEVICE, dtype=DTYPE)
+    k = torch.empty_like(q)
+    weight = torch.ones(head_dim, device=DEVICE, dtype=DTYPE)
+    cache = create_cos_sin_cache(head_dim, 1)
+    positions = torch.empty(0, device=DEVICE, dtype=torch.int64)
+
+    fused_inplace_qknorm_rope(
+        q,
+        k,
+        weight,
+        weight,
+        cache,
+        positions,
+        is_neox=False,
+        rope_dim=head_dim,
+    )
+    assert q.numel() == k.numel() == 0
+
+
 if __name__ == "__main__":
     sys.exit(pytest.main([__file__, "-v", "-s"]))

test/registered/kernels/ops/diffusion/test_quality_gate.py

@@ -0,0 +1,47 @@
+import sys
+
+import pytest
+import torch.nn as nn
+
+from sglang.kernels.ops.diffusion.quality_gate import QualityGatedFusion
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=2, suite="base-a-test-cpu")
+
+
+def test_quality_gate_mounts_and_unmounts_all_sites():
+    fusion = QualityGatedFusion(
+        name="test fusion",
+        marker_attr="_test_fusion_site",
+        enabled_attr="_test_fusion_enabled",
+    )
+    root = nn.ModuleList([nn.Module(), nn.Module()])
+    for index, site in enumerate(root):
+        fusion.mark(site, index)
+
+    assert [fusion.metadata(site) for site in root] == [0, 1]
+    assert fusion.mount(root)
+    assert all(fusion.is_enabled(site) for site in root)
+    fusion.unmount(root)
+    assert not any(fusion.is_enabled(site) for site in root)
+
+
+def test_quality_gate_rejection_is_all_or_nothing():
+    fusion = QualityGatedFusion(
+        name="test fusion",
+        marker_attr="_test_fusion_site",
+        enabled_attr="_test_fusion_enabled",
+    )
+    root = nn.ModuleList([nn.Module(), nn.Module()])
+    for index, site in enumerate(root):
+        fusion.mark(site, index)
+
+    assert not fusion.mount(
+        root, reject_reason=lambda site: "rejected" if fusion.metadata(site) else None
+    )
+    assert not any(fusion.is_enabled(site) for site in root)
+    assert not fusion.mount(nn.Module())
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__, "-v"]))

test/registered/kernels/ops/diffusion/test_residual_gate_add.py

@@ -5,6 +5,7 @@ import torch
 
 from sglang.kernels.ops.diffusion.residual_gate_add import (
     can_use_residual_gate_add_cuda,
+    residual_gate_add,
     residual_gate_add_cuda,
 )
 from sglang.test.ci.ci_register import register_cuda_ci
@@ -25,6 +26,8 @@ CASES = [
     ((1, 4608, 3072), (1, 1, 3072)),
     # FLUX.2-dev (D=6144) joint sequence.
     ((1, 4608, 6144), (1, 1, 6144)),
+    # ERNIE-4.5-VL 1024^2 image tokens plus text tokens.
+    ((1, 4216, 4096), (1, 1, 4096)),
 ]
 
 
@@ -55,6 +58,7 @@ def test_residual_gate_add_matches_torch(residual_shape, gate_shape):
     out = residual_gate_add_cuda(residual, update, gate)
     ref = residual + update * gate
     _assert_matches_torch(out, ref)
+    assert torch.equal(residual_gate_add(residual, update, gate), ref)
 
 
 @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32])
@@ -79,6 +83,13 @@ def test_can_use_residual_gate_add_cuda_rejects_unsupported_inputs():
     assert not can_use_residual_gate_add_cuda(residual, update.float(), gate)
     assert not can_use_residual_gate_add_cuda(residual, update[:, ::2], gate)
     assert not can_use_residual_gate_add_cuda(residual, update, gate[:, :, ::2])
+    empty_residual = residual[:, :0]
+    empty_update = update[:, :0]
+    assert not can_use_residual_gate_add_cuda(empty_residual, empty_update, gate)
+    assert torch.equal(
+        residual_gate_add(empty_residual, empty_update, gate),
+        empty_residual + empty_update * gate,
+    )
 
     # Only [1, ..., 1, D] row-broadcast gates are supported; a batched
     # [B>1, 1, D] gate is not row-broadcast here and must fall back.
@@ -88,6 +99,10 @@ def test_can_use_residual_gate_add_cuda_rejects_unsupported_inputs():
     assert not can_use_residual_gate_add_cuda(
         batched_residual, batched_update, batched_gate
     )
+    assert torch.equal(
+        residual_gate_add(batched_residual, batched_update, batched_gate),
+        batched_residual + batched_update * batched_gate,
+    )
 
 
 def test_residual_gate_add_custom_op_torch_compile_fullgraph():
@@ -96,7 +111,7 @@ def test_residual_gate_add_custom_op_torch_compile_fullgraph():
     gate = torch.randn((1, 1, 128), device="cuda", dtype=torch.bfloat16)
 
     def fn(residual, update, gate):
-        return residual_gate_add_cuda(residual, update, gate)
+        return residual_gate_add(residual, update, gate)
 
     compiled = torch.compile(fn, fullgraph=True)
     out = compiled(residual, update, gate)

test/registered/kernels/ops/diffusion/test_scale_shift.py

@@ -0,0 +1,40 @@
+import sys
+
+import pytest
+import torch
+
+from sglang.kernels.ops.diffusion.triton.scale_shift import (
+    try_fused_scaled_residual_add_exact,
+)
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=5, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+
+
+@torch.no_grad()
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
+def test_scaled_residual_add_is_bit_exact(dtype):
+    torch.manual_seed(0)
+    residual = torch.randn(2, 17, 64, device="cuda", dtype=torch.float32)
+    x = torch.randn(2, 17, 64, device="cuda", dtype=dtype)
+    scale = torch.randn(64, device="cuda", dtype=torch.float32)
+
+    actual = try_fused_scaled_residual_add_exact(residual, x, scale)
+    expected = residual + x * scale
+    assert actual is not None
+    assert torch.equal(actual, expected)
+
+
+@torch.no_grad()
+def test_scaled_residual_add_rejects_unsupported_inputs():
+    residual = torch.empty(2, 3, 8, device="cuda", dtype=torch.float32)
+    x = torch.empty_like(residual)
+    scale = torch.empty(8, device="cuda", dtype=torch.float32)
+
+    assert try_fused_scaled_residual_add_exact(residual, x, scale) is None
+    assert try_fused_scaled_residual_add_exact(residual, x.half(), scale[:-1]) is None
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__, "-v", "-s"]))

test/registered/kernels/ops/diffusion/test_timestep_embedding.py

@@ -134,12 +134,12 @@ def test_timestep_embedding_perf():
         end = torch.cuda.Event(enable_timing=True)
 
         for _ in range(warmup_times):
-            output_fn = kernel_fn(*args, **kwargs)
+            kernel_fn(*args, **kwargs)
         torch.cuda.synchronize()
 
         start.record()
         for _ in range(repeat_times):
-            output_fn = kernel_fn(*args, **kwargs)
+            kernel_fn(*args, **kwargs)
         end.record()
         end.synchronize()
         return start.elapsed_time(end) / repeat_times

test/registered/kernels/ops/diffusion/test_ulysses_qkv.py

@@ -0,0 +1,54 @@
+import sys
+
+import pytest
+import torch
+
+from sglang.kernels.ops.diffusion.triton.ulysses_qkv import (
+    pack_qkv_destination_major,
+)
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=8, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+
+
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
+def test_pack_qkv_destination_major_is_bit_exact(dtype):
+    torch.manual_seed(0)
+    rows, world_size, global_heads, head_size = 17, 4, 12, 64
+    q, k, v = (
+        torch.randn(rows, global_heads, head_size, device="cuda", dtype=dtype)
+        for _ in range(3)
+    )
+
+    local_heads = global_heads // world_size
+    expected = torch.empty(
+        world_size,
+        rows,
+        local_heads,
+        3 * head_size,
+        device="cuda",
+        dtype=dtype,
+    )
+    for index, tensor in enumerate((q, k, v)):
+        shards = tensor.view(rows, world_size, local_heads, head_size).permute(
+            1, 0, 2, 3
+        )
+        expected[..., index * head_size : (index + 1) * head_size].copy_(shards)
+
+    actual = pack_qkv_destination_major(q, k, v, world_size)
+    assert torch.equal(actual, expected)
+
+
+def test_pack_qkv_destination_major_validates_inputs():
+    q = torch.empty(2, 4, 8, device="cuda", dtype=torch.bfloat16)
+    with pytest.raises(ValueError, match="same 3D shape"):
+        pack_qkv_destination_major(q, q[:, :-1], q, 2)
+    with pytest.raises(ValueError, match="divide global_heads"):
+        pack_qkv_destination_major(q, q, q, 3)
+    with pytest.raises(ValueError, match="expected shape"):
+        pack_qkv_destination_major(q, q, q, 2, out=torch.empty_like(q))
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__, "-v", "-s"]))

test/registered/kernels/ops/diffusion/test_varlen_pack_pad.py

@@ -55,7 +55,10 @@ def _build_mask(bs, s_txt, s_img, valid_txt_lens):
 
 def _ref_pack(q, k, v, indices):
     bs, seq = q.shape[:2]
-    flat = lambda t: t.reshape(bs * seq, *t.shape[2:])
+
+    def flat(t):
+        return t.reshape(bs * seq, *t.shape[2:])
+
     return (
         flat(q).index_select(0, indices),
         flat(k).index_select(0, indices),
@@ -64,7 +67,6 @@ def _ref_pack(q, k, v, indices):
 
 
 def _ref_scatter(out_unpad, indices, bs, seq):
-    n_valid = indices.shape[0]
     _, num_heads, head_dim = out_unpad.shape
     flat = torch.zeros(
         bs * seq, num_heads, head_dim, dtype=out_unpad.dtype, device=DEVICE

test/registered/kernels/ops/diffusion/test_wan_vae_fastpath.py

@@ -66,5 +66,14 @@ def test_fused_module_gate_dispatch() -> None:
     assert torch.equal(fused(x), expected)
 
 
+@torch.no_grad()
+def test_kernel_rejects_empty_input() -> None:
+    x = torch.empty(1, 96, 0, 2, 2, device="cuda", dtype=torch.bfloat16).to(
+        memory_format=torch.channels_last_3d
+    )
+    gamma = torch.ones(96, 1, 1, 1, device="cuda", dtype=torch.bfloat16)
+    assert wan_rmsnorm_silu(x, gamma) is None
+
+
 if __name__ == "__main__":
     sys.exit(pytest.main([__file__, "-v", "-s"]))