diff --git a/python/sglang/kernels/jit/csrc/diffusion/modulate_scale_shift.cuh b/python/sglang/kernels/jit/csrc/diffusion/modulate_scale_shift.cuh index 8e5ed6b1c..96ee5b378 100644 --- a/python/sglang/kernels/jit/csrc/diffusion/modulate_scale_shift.cuh +++ b/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 -// storage dtype (the per-op kernel boundaries of the eager aten chain): -// -// 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). +// Reproduces the eager storage-dtype rounding boundaries: +// out = round(round(x * round(1 + scale)) + shift) #pragma once -#include // For host dtype helpers and TensorView metadata -#include // For RuntimeCheck and div_ceil +#include +#include -#include // For DTypeTrait conversions -#include // For LaunchKernel and CUDA dtype aliases -#include // For device::AlignedVector +#include +#include +#include +#include #include +#include namespace sglang { -namespace sglang_modulate_scale_shift { +namespace modulate_scale_shift { namespace { -constexpr int kRowsPerBlock = 4; -constexpr int kColsVecPerBlock = 256; -constexpr int64_t kMaxGrid = 65535; - -inline const char* data_ptr(const tvm::ffi::TensorView& t) { - return static_cast(t.data_ptr()) + t.byte_offset(); -} - -inline char* mutable_data_ptr(const tvm::ffi::TensorView& t) { - return static_cast(t.data_ptr()) + t.byte_offset(); -} - -inline bool aligned16(const void* p) { - return (reinterpret_cast(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; -} +constexpr uint32_t kRowsPerBlock = 4; +constexpr uint32_t kColsVecPerBlock = 256; +constexpr uint32_t kMaxGridY = 65535; +constexpr uintptr_t kAlignment = 16; template -inline void check_dtype(const tvm::ffi::TensorView& t) { - host::RuntimeCheck(host::is_type(t.dtype()), "unexpected dtype for modulate_scale_shift"); -} - -template -__device__ __forceinline__ float to_float(T v) { - return static_cast(v); -} - -template <> -__device__ __forceinline__ float to_float(fp16_t v) { - return __half2float(v); -} - -template <> -__device__ __forceinline__ float to_float(bf16_t v) { - return __bfloat162float(v); -} - -template -__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::from(1.0f + to_float(scale)); - const T product = DTypeTrait::from(to_float(x) * to_float(one_plus_scale)); - return DTypeTrait::from(to_float(product) + to_float(shift)); +SGL_DEVICE T modulate_value(T x, T scale, T shift) { + const T one_plus_scale = device::cast(1.0f + device::cast(scale)); + const T product = device::cast(device::cast(x) * device::cast(one_plus_scale)); + return device::cast(device::cast(product) + device::cast(shift)); } template -__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_tile_stride = static_cast(gridDim.y) * kRowsPerBlock; - for (int64_t row_base = static_cast(blockIdx.y) * kRowsPerBlock; row_base < rows; - row_base += row_tile_stride) { + const int64_t row_stride = static_cast(gridDim.y) * kRowsPerBlock; + for (int64_t row_base = static_cast(blockIdx.y) * kRowsPerBlock; row_base < rows; row_base += row_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) { - const int64_t batch = row / rows_per_batch; - const int64_t mod_v = batch * row_vec + col_vec; - const int64_t v = row * row_vec + col_vec; - Vec xv, s, b, o; - s.load(scale, mod_v); - b.load(shift, mod_v); - xv.load(x, v); -#pragma unroll - for (int i = 0; i < kVec; ++i) { - o[i] = modulate_value(xv[i], s[i], b[i]); - } - o.store(out, v); + if (row >= rows) { + continue; } + + const int64_t batch = row / rows_per_batch; + const int64_t modulation_offset = batch * row_vec + col_vec; + const int64_t activation_offset = row * row_vec + col_vec; + Vec x_vec, scale_vec, shift_vec, out_vec; + x_vec.load(x, activation_offset); + scale_vec.load(scale, modulation_offset); + shift_vec.load(shift, modulation_offset); +#pragma unroll + for (int i = 0; i < kVec; ++i) { + out_vec[i] = modulate_value(x_vec[i], scale_vec[i], shift_vec[i]); + } + out_vec.store(out, activation_offset); } } } -template -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(data_ptr(x)); - const T* scale_ptr = reinterpret_cast(data_ptr(scale)); - const T* shift_ptr = reinterpret_cast(data_ptr(shift)); - T* out_ptr = reinterpret_cast(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(kColsVecPerBlock)); - const int64_t row_tiles = host::div_ceil(rows, static_cast(kRowsPerBlock)); - const int64_t row_blocks = row_tiles > kMaxGrid ? kMaxGrid : row_tiles; - host::LaunchKernel( - dim3(static_cast(col_blocks), static_cast(row_blocks)), dim3(kColsVecPerBlock), out.device())( - modulate_scale_shift_vec_kernel, x_ptr, scale_ptr, shift_ptr, out_ptr, rows, rows_per_batch, row_vec); -} - -template -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(out); - check_dtype(x); - check_dtype(scale); - check_dtype(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 struct ModulateScaleShiftKernel { + static_assert(std::is_same_v || std::is_same_v); + + /** + * \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(out, x, scale, shift); - launch_modulate_scale_shift(out, x, scale, shift); + using namespace host; + + auto B = SymbolicSize{"batch"}; + auto L = SymbolicSize{"sequence_length"}; + auto D = SymbolicSize{"hidden_size"}; + auto device = SymbolicDevice{}; + device.set_options(); + + TensorMatcher({B, L, D}).with_dtype().with_device(device).verify(out).verify(x); + TensorMatcher({B, D}).with_dtype().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(out.data_ptr()); + const auto* x_ptr = static_cast(x.data_ptr()); + const auto* scale_ptr = static_cast(scale.data_ptr()); + const auto* shift_ptr = static_cast(shift.data_ptr()); + CHECK_HOST( + reinterpret_cast(out_ptr) % kAlignment == 0 && + reinterpret_cast(x_ptr) % kAlignment == 0 && + reinterpret_cast(scale_ptr) % kAlignment == 0 && + reinterpret_cast(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(div_ceil(row_vec, static_cast(kColsVecPerBlock))); + const int64_t row_tiles = div_ceil(rows, static_cast(kRowsPerBlock)); + const auto row_blocks = static_cast(std::min(row_tiles, kMaxGridY)); + LaunchKernel(dim3(col_blocks, row_blocks), kColsVecPerBlock, device.unwrap())( + modulate_scale_shift_kernel, out_ptr, x_ptr, scale_ptr, shift_ptr, rows, sequence_length, row_vec); } }; -} // namespace sglang_modulate_scale_shift +} // namespace modulate_scale_shift } // namespace sglang diff --git a/python/sglang/kernels/jit/csrc/diffusion/qknorm_rope.cuh b/python/sglang/kernels/jit/csrc/diffusion/qknorm_rope.cuh index 2494e2d85..0049f5258 100644 --- a/python/sglang/kernels/jit/csrc/diffusion/qknorm_rope.cuh +++ b/python/sglang/kernels/jit/csrc/diffusion/qknorm_rope.cuh @@ -277,6 +277,7 @@ struct QKNormRopeKernel { const auto num_tokens = static_cast(N.unwrap()); const auto num_qo_heads = static_cast(Q.unwrap()); const auto num_kv_heads = static_cast(K.unwrap()); + if (num_tokens == 0 || (num_qo_heads == 0 && num_kv_heads == 0)) return; const auto q_stride_bytes = static_cast(Dq.unwrap() * sizeof(DType)); const auto k_stride_bytes = static_cast(Dk.unwrap() * sizeof(DType)); const auto head_stride_bytes = static_cast(Dd.unwrap() * sizeof(DType)); diff --git a/python/sglang/kernels/jit/csrc/diffusion/residual_gate_add.cuh b/python/sglang/kernels/jit/csrc/diffusion/residual_gate_add.cuh index c7a3cf9e6..75670e5fd 100644 --- a/python/sglang/kernels/jit/csrc/diffusion/residual_gate_add.cuh +++ b/python/sglang/kernels/jit/csrc/diffusion/residual_gate_add.cuh @@ -1,314 +1,205 @@ -// CUDA fast path for diffusion residual-gate elementwise updates. -// -// Implements: +// CUDA fast path for bit-exact diffusion residual-gate updates: // 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 // For host dtype helpers and TensorView metadata -#include // For RuntimeCheck and div_ceil +#include +#include -#include // For DTypeTrait conversions -#include // For LaunchKernel and CUDA dtype aliases -#include // For device::AlignedVector +#include +#include +#include +#include #include +#include namespace sglang { namespace residual_gate_add { -constexpr int kBlockSize = 256; -constexpr int kBcastRowsPerBlock = 4; -constexpr int kBcastColsVecPerBlock = 256; -constexpr int64_t kMaxGrid = 65535; +namespace { -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) { - return static_cast(t.data_ptr()) + t.byte_offset(); -} - -inline char* mutable_data_ptr(const tvm::ffi::TensorView& t) { - return static_cast(t.data_ptr()) + t.byte_offset(); -} - -inline bool aligned16(const void* p) { - return (reinterpret_cast(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(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; -} +enum class GateMode : int { kFull, kBroadcastRow }; template -inline void check_dtype(const tvm::ffi::TensorView& t) { - host::RuntimeCheck(host::is_type(t.dtype()), "unexpected dtype for residual_gate_add"); -} - -template -__device__ __forceinline__ float to_float(T v) { - return static_cast(v); -} - -template <> -__device__ __forceinline__ float to_float(fp16_t v) { - return __half2float(v); -} - -template <> -__device__ __forceinline__ float to_float(bf16_t v) { - return __bfloat162float(v); -} - -template -__device__ __forceinline__ T residual_gate_value(T residual, T update, T gate) { - const T product = DTypeTrait::from(to_float(update) * to_float(gate)); - return DTypeTrait::from(to_float(residual) + to_float(product)); +SGL_DEVICE T residual_gate_value(T residual, T update, T gate) { + const T product = device::cast(device::cast(update) * device::cast(gate)); + return device::cast(device::cast(residual) + device::cast(product)); } template __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 n_vec) { + int64_t num_vectors) { using Vec = device::AlignedVector; const int64_t stride = static_cast(gridDim.x) * blockDim.x; - for (int64_t v = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; v < n_vec; v += stride) { - Vec r, u, g, o; - r.load(residual, v); - u.load(update, v); - g.load(gate, v); + for (int64_t vector = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; vector < num_vectors; + vector += stride) { + Vec residual_vec, update_vec, gate_vec, out_vec; + residual_vec.load(residual, vector); + 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 -__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; - const int64_t col_vec = static_cast(blockIdx.x) * kBcastColsVecPerBlock + threadIdx.x; - if (col_vec >= row_vec) { + const int64_t column = static_cast(blockIdx.x) * kBroadcastColsPerBlock + threadIdx.x; + if (column >= row_vectors) { return; } - Vec g; - g.load(gate, col_vec); - - // Grid-stride over row tiles so the launch stays valid even when the number - // of row tiles exceeds the gridDim.y hardware limit. - const int64_t row_tile_stride = static_cast(gridDim.y) * kBcastRowsPerBlock; - for (int64_t row_base = static_cast(blockIdx.y) * kBcastRowsPerBlock; row_base < rows; - row_base += row_tile_stride) { + Vec gate_vec; + gate_vec.load(gate, column); + const int64_t row_stride = static_cast(gridDim.y) * kBroadcastRowsPerBlock; + for (int64_t row_base = static_cast(blockIdx.y) * kBroadcastRowsPerBlock; row_base < rows; + row_base += row_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) { - const int64_t v = row * row_vec + col_vec; - Vec r, u, o; - r.load(residual, v); - u.load(update, v); -#pragma unroll - for (int i = 0; i < kVec; ++i) { - o[i] = residual_gate_value(r[i], u[i], g[i]); - } - o.store(out, v); + if (row >= rows) { + continue; } + const int64_t vector = row * row_vectors + column; + 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) { + out_vec[i] = residual_gate_value(residual_vec[i], update_vec[i], gate_vec[i]); + } + out_vec.store(out, vector); } } } -template +template __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 begin, - int64_t total, - int64_t D) { + int64_t numel, + int64_t hidden_size) { const int64_t stride = static_cast(gridDim.x) * blockDim.x; - for (int64_t i = begin + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; i < total; i += stride) { - const T gate_value = kGate == GateMode::kFull ? gate[i] : SGLANG_LDG(gate + (i % D)); - out[i] = residual_gate_value(residual[i], update[i], gate_value); + for (int64_t index = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; index < numel; index += stride) { + const T gate_value = kGateMode == GateMode::kFull ? gate[index] : SGLANG_LDG(gate + index % hidden_size); + out[index] = residual_gate_value(residual[index], update[index], gate_value); } } -template -inline void launch_residual_gate_add( - const tvm::ffi::TensorView& out, - const tvm::ffi::TensorView& residual, - const tvm::ffi::TensorView& update, - const tvm::ffi::TensorView& gate, - GateMode mode) { - const int64_t total = numel(residual); - if (total == 0) { - return; - } - - const int64_t D = residual.size(residual.ndim() - 1); - const T* residual_ptr = reinterpret_cast(data_ptr(residual)); - const T* update_ptr = reinterpret_cast(data_ptr(update)); - const T* gate_ptr = reinterpret_cast(data_ptr(gate)); - T* out_ptr = reinterpret_cast(mutable_data_ptr(out)); - constexpr int kVec = 16 / sizeof(T); - - const bool vec_ok = aligned16(residual_ptr) && aligned16(update_ptr) && aligned16(gate_ptr) && aligned16(out_ptr) && - (D % kVec == 0) && (mode == GateMode::kBcastRow || total % kVec == 0); - - int64_t done = 0; - if (vec_ok) { - const int64_t n_vec = total / kVec; - const int64_t row_vec = D / kVec; - if (mode == GateMode::kFull) { - host::LaunchKernel(static_cast(grid_for(n_vec)), kBlockSize, out.device())( - residual_gate_add_vec_kernel, residual_ptr, update_ptr, gate_ptr, out_ptr, n_vec); - } else { - const int64_t rows = total / D; - const int64_t col_blocks = host::div_ceil(row_vec, static_cast(kBcastColsVecPerBlock)); - const int64_t row_tiles = host::div_ceil(rows, static_cast(kBcastRowsPerBlock)); - const int64_t row_blocks = row_tiles > kMaxGrid ? kMaxGrid : row_tiles; - host::LaunchKernel( - dim3(static_cast(col_blocks), static_cast(row_blocks)), - dim3(kBcastColsVecPerBlock), - out.device())( - residual_gate_add_bcast_row_tile_kernel, 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(grid_for(total - done)), kBlockSize, out.device())( - residual_gate_add_scalar_kernel, - residual_ptr, - update_ptr, - gate_ptr, - out_ptr, - done, - total, - D); - } else { - host::LaunchKernel(static_cast(grid_for(total - done)), kBlockSize, out.device())( - residual_gate_add_scalar_kernel, - residual_ptr, - update_ptr, - gate_ptr, - out_ptr, - done, - total, - D); - } - } -} - -template -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(out); - check_dtype(residual); - check_dtype(update); - check_dtype(gate); - host::RuntimeCheck(residual.device().device_type == kDLCUDA, "residual must be CUDA"); - host::RuntimeCheck(update.device().device_type == kDLCUDA, "update must be CUDA"); - host::RuntimeCheck(gate.device().device_type == kDLCUDA, "gate must be CUDA"); - host::RuntimeCheck(out.device().device_type == kDLCUDA, "out must be CUDA"); - host::RuntimeCheck( - residual.device().device_id == update.device().device_id && - residual.device().device_id == gate.device().device_id && - residual.device().device_id == out.device().device_id, - "residual/update/gate/out must be on the same CUDA device"); - host::RuntimeCheck(residual.ndim() >= 2, "residual must be at least 2D"); - host::RuntimeCheck(update.ndim() == residual.ndim(), "update rank must match residual"); - host::RuntimeCheck(out.ndim() == residual.ndim(), "out rank must match residual"); - for (int i = 0; i < residual.ndim(); ++i) { - host::RuntimeCheck(update.size(i) == residual.size(i), "update shape must match residual"); - host::RuntimeCheck(out.size(i) == residual.size(i), "out shape must match residual"); - } - host::RuntimeCheck(is_dense_contiguous(residual), "residual must be contiguous"); - host::RuntimeCheck(is_dense_contiguous(update), "update must be contiguous"); - host::RuntimeCheck(is_dense_contiguous(out), "out must be contiguous"); - host::RuntimeCheck(is_dense_contiguous(gate), "gate must be contiguous"); - host::RuntimeCheck(data_ptr(out) != data_ptr(residual), "out must not alias residual"); - host::RuntimeCheck(data_ptr(out) != data_ptr(update), "out must not alias update"); - host::RuntimeCheck(data_ptr(out) != data_ptr(gate), "out must not alias gate"); - - const int D_dim = residual.ndim() - 1; - const int row_dim = residual.ndim() - 2; - host::RuntimeCheck(gate.ndim() == residual.ndim(), "gate rank must match residual"); - host::RuntimeCheck(gate.size(D_dim) == residual.size(D_dim), "gate last dim must match residual"); - - bool full_gate = true; - for (int i = 0; i < residual.ndim(); ++i) { - full_gate = full_gate && gate.size(i) == residual.size(i); - } - if (full_gate) { - return GateMode::kFull; - } - - host::RuntimeCheck(gate.size(row_dim) == 1, "broadcast gate row dim must be 1"); - for (int i = 0; i < D_dim; ++i) { - host::RuntimeCheck(gate.size(i) == 1, "broadcast gate leading dims must be 1"); - } - return GateMode::kBcastRow; -} +} // namespace +/** + * \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 struct ResidualGateAddKernel { + static_assert(std::is_same_v || std::is_same_v || std::is_same_v); + 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(out, residual, update, gate); - launch_residual_gate_add(out, residual, update, gate, mode); + run(tvm::ffi::TensorView out, + tvm::ffi::TensorView residual, + tvm::ffi::TensorView update, + 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(); + TensorMatcher({N}).with_dtype().with_device(device).verify(out).verify(residual).verify(update); + TensorMatcher({G}).with_dtype().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; + } + + auto* out_ptr = static_cast(out.data_ptr()); + const auto* residual_ptr = static_cast(residual.data_ptr()); + const auto* update_ptr = static_cast(update.data_ptr()); + const auto* gate_ptr = static_cast(gate.data_ptr()); + CHECK_HOST(out_ptr != residual_ptr && out_ptr != update_ptr && out_ptr != gate_ptr) + << "output must not alias an input"; + + constexpr int kVec = kAlignment / sizeof(T); + const bool aligned = reinterpret_cast(out_ptr) % kAlignment == 0 && + reinterpret_cast(residual_ptr) % kAlignment == 0 && + reinterpret_cast(update_ptr) % kAlignment == 0 && + reinterpret_cast(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(std::min(div_ceil(num_vectors, static_cast(kBlockSize)), kMaxGrid)); + LaunchKernel(blocks, kBlockSize, device.unwrap())( + residual_gate_add_vec_kernel, out_ptr, residual_ptr, update_ptr, gate_ptr, num_vectors); + return; + } + + const int64_t rows = numel / hidden_size; + const int64_t row_vectors = hidden_size / kVec; + const auto column_blocks = + static_cast(div_ceil(row_vectors, static_cast(kBroadcastColsPerBlock))); + const int64_t row_tiles = div_ceil(rows, static_cast(kBroadcastRowsPerBlock)); + const auto row_blocks = static_cast(std::min(row_tiles, kMaxGrid)); + LaunchKernel(dim3(column_blocks, row_blocks), kBroadcastColsPerBlock, device.unwrap())( + residual_gate_add_broadcast_kernel, out_ptr, residual_ptr, update_ptr, gate_ptr, rows, row_vectors); + return; + } + + const auto blocks = + static_cast(std::min(div_ceil(numel, static_cast(kBlockSize)), kMaxGrid)); + if (broadcast_gate) { + LaunchKernel(blocks, kBlockSize, device.unwrap())( + residual_gate_add_scalar_kernel, + out_ptr, + residual_ptr, + update_ptr, + gate_ptr, + numel, + hidden_size); + } else { + LaunchKernel(blocks, kBlockSize, device.unwrap())( + residual_gate_add_scalar_kernel, + out_ptr, + residual_ptr, + update_ptr, + gate_ptr, + numel, + hidden_size); + } } }; diff --git a/python/sglang/kernels/jit/csrc/diffusion/usp_relayout.cuh b/python/sglang/kernels/jit/csrc/diffusion/usp_relayout.cuh index 6ec5b4968..92f670af9 100644 --- a/python/sglang/kernels/jit/csrc/diffusion/usp_relayout.cuh +++ b/python/sglang/kernels/jit/csrc/diffusion/usp_relayout.cuh @@ -1,111 +1,59 @@ -// CUDA fast path for the Ulysses sequence-parallel output head merge. -// -// 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. +// CUDA fast path for the Ulysses sequence-parallel output-head merge: +// [W, S, B, H, D] -> [B, S, W, H, D] #pragma once -#include // For host dtype helpers and TensorView metadata -#include // For RuntimeCheck and div_ceil +#include +#include -#include // For CUDA dtype aliases -#include // For LaunchKernel -#include // For device::AlignedVector +#include +#include +#include +#include #include +#include namespace sglang { namespace usp_relayout { -constexpr int kBlockSize = 256; -constexpr int64_t kMaxGrid = 65535; +namespace { -inline const char* data_ptr(const tvm::ffi::TensorView& t) { - return static_cast(t.data_ptr()) + t.byte_offset(); -} +constexpr uint32_t kBlockSize = 256; +constexpr uint32_t kMaxGrid = 65535; +constexpr uintptr_t kAlignment = 16; -inline char* mutable_data_ptr(const tvm::ffi::TensorView& t) { - return static_cast(t.data_ptr()) + t.byte_offset(); -} - -inline bool aligned16(const void* p) { - return (reinterpret_cast(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(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 -inline void check_dtype(const tvm::ffi::TensorView& t) { - host::RuntimeCheck(host::is_type(t.dtype()), "unexpected dtype for usp_merge_heads tensor"); -} - -// out[b, s, w, h, c] = x[w, s, b, h, c] +// out[b, s, w, h, d] = x[w, s, b, h, d] template __global__ void usp_merge_heads_vec_kernel( T* __restrict__ out, const T* __restrict__ x, - int64_t n_vec, - int64_t d_vec, // D / kVec - int64_t h_local, + int64_t num_vectors, + int64_t head_vectors, + int64_t local_heads, int64_t batch, - int64_t seq, - int64_t world) { + int64_t sequence_length, + int64_t world_size) { const int64_t stride = static_cast(gridDim.x) * blockDim.x; - for (int64_t i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; i < n_vec; i += stride) { - int64_t rest = i; - const int64_t c_vec = rest % d_vec; - rest /= d_vec; - const int64_t h = rest % h_local; - rest /= h_local; - const int64_t w = rest % world; - rest /= world; - const int64_t s = rest % seq; - const int64_t b = rest / seq; + for (int64_t index = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; index < num_vectors; + index += stride) { + int64_t rest = index; + const int64_t head_offset = rest % head_vectors; + rest /= head_vectors; + const int64_t head = rest % local_heads; + rest /= local_heads; + const int64_t rank = rest % world_size; + rest /= world_size; + 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; - device::AlignedVector val; - val.load(x, src_vec); - val.store(out, i); + const int64_t source = + ((((rank * sequence_length + sequence) * batch + batch_index) * local_heads) + head) * head_vectors + + head_offset; + device::AlignedVector value; + value.load(x, source); + value.store(out, index); } } @@ -113,66 +61,83 @@ template __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 world) { + int64_t sequence_length, + int64_t world_size) { const int64_t stride = static_cast(gridDim.x) * blockDim.x; - for (int64_t i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; i < total; i += stride) { - int64_t rest = i; - const int64_t c = rest % head_dim; + for (int64_t index = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; index < numel; index += stride) { + int64_t rest = index; + const int64_t head_offset = rest % head_dim; rest /= head_dim; - const int64_t h = rest % h_local; - rest /= h_local; - const int64_t w = rest % world; - rest /= world; - const int64_t s = rest % seq; - const int64_t b = rest / seq; + const int64_t head = rest % local_heads; + rest /= local_heads; + const int64_t rank = rest % world_size; + rest /= world_size; + const int64_t sequence = rest % sequence_length; + 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 struct UspMergeHeadsKernel { - static void run(tvm::ffi::TensorView out, tvm::ffi::TensorView x) { - check_dtype(out); - check_dtype(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"); + static_assert(std::is_same_v || std::is_same_v || std::is_same_v); - const int64_t total = numel(x); - if (total == 0) { + static void run(tvm::ffi::TensorView out, tvm::ffi::TensorView x) { + 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(); + TensorMatcher({W, S, B, H, D}).with_dtype().with_device(device).verify(x); + TensorMatcher({B, S, W, H, D}).with_dtype().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(mutable_data_ptr(out)); - const T* x_ptr = reinterpret_cast(data_ptr(x)); + auto* out_ptr = static_cast(out.data_ptr()); + const auto* x_ptr = static_cast(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(std::min(div_ceil(work_items, static_cast(kBlockSize)), kMaxGrid)); + LaunchKernel(blocks, kBlockSize, device.unwrap())(kernel, out_ptr, x_ptr, work_items, args...); + }; - constexpr int kVec = 16 / sizeof(T); - const bool vec_ok = (head_dim % kVec == 0) && aligned16(out_ptr) && aligned16(x_ptr); - if (vec_ok) { - const int64_t n_vec = total / kVec; - host::LaunchKernel(static_cast(grid_for(n_vec)), kBlockSize, out.device())( - usp_merge_heads_vec_kernel, out_ptr, x_ptr, n_vec, head_dim / kVec, h_local, batch, seq, world); + constexpr int kVec = kAlignment / sizeof(T); + const bool vectorized = head_dim % kVec == 0 && reinterpret_cast(out_ptr) % kAlignment == 0 && + reinterpret_cast(x_ptr) % kAlignment == 0; + if (vectorized) { + launch( + usp_merge_heads_vec_kernel, + numel / kVec, + head_dim / kVec, + local_heads, + batch, + sequence_length, + world_size); } else { - host::LaunchKernel(static_cast(grid_for(total)), kBlockSize, out.device())( - usp_merge_heads_scalar_kernel, out_ptr, x_ptr, total, head_dim, h_local, batch, seq, world); + launch(usp_merge_heads_scalar_kernel, numel, head_dim, local_heads, batch, sequence_length, world_size); } } }; diff --git a/python/sglang/kernels/ops/diffusion/__init__.py b/python/sglang/kernels/ops/diffusion/__init__.py index 556549b26..6c97dd8aa 100644 --- a/python/sglang/kernels/ops/diffusion/__init__.py +++ b/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", diff --git a/python/sglang/kernels/ops/diffusion/cutedsl/scale_residual_norm_scale_shift.py b/python/sglang/kernels/ops/diffusion/cutedsl/scale_residual_norm_scale_shift.py index 2905cda24..e0c3c4ce3 100644 --- a/python/sglang/kernels/ops/diffusion/cutedsl/scale_residual_norm_scale_shift.py +++ b/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 diff --git a/python/sglang/kernels/ops/diffusion/fused_gate_rmsnorm.py b/python/sglang/kernels/ops/diffusion/fused_gate_rmsnorm.py index c8313cbf1..8d210b85c 100644 --- a/python/sglang/kernels/ops/diffusion/fused_gate_rmsnorm.py +++ b/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 -bf16-native Triton kernels -(:mod:`sglang.kernels.ops.diffusion.triton.zimage_native_norm`) fuse each +attention/FFN and ``x + tanh(gate) * RMSNorm(out)`` after. Shared BF16-native +Triton kernels +(: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 ( - zimage_rmsnorm_scale, + from sglang.kernels.ops.diffusion.triton.native_bf16_rmsnorm import ( + 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 ( - zimage_rmsnorm_tanh_residual, + from sglang.kernels.ops.diffusion.triton.native_bf16_rmsnorm import ( + 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)) - setattr(module, _SITE_ENABLED_ATTR, False) + _FUSION.mark(module, tuple(norm_attrs)) -def iter_fused_gate_rmsnorm_sites(root: nn.Module) -> Iterator[nn.Module]: - """Yield every marked site under ``root`` (including ``root``).""" - for module in root.modules(): - if getattr(module, _SITE_NORM_ATTRS, None) is not None: - yield module +def fused_gate_rmsnorm_active(module: nn.Module) -> bool: + """Whether the quality-gated fused path is mounted on ``module``.""" + return _FUSION.is_enabled(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)) - 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 + return _FUSION.mount(root, reject_reason=_static_reject_reason, logger=logger) 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): - setattr(site, _SITE_ENABLED_ATTR, False) + _FUSION.unmount(root) diff --git a/python/sglang/kernels/ops/diffusion/fused_linear_gelu.py b/python/sglang/kernels/ops/diffusion/fused_linear_gelu.py index 11774ca02..64b2f33dc 100644 --- a/python/sglang/kernels/ops/diffusion/fused_linear_gelu.py +++ b/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) - setattr(module, _SITE_ENABLED_ATTR, False) + _FUSION.mark(module, linear_attr) -def iter_fused_gelu_sites(root: nn.Module) -> Iterator[nn.Module]: - """Yield every marked fusion site under ``root`` (including ``root``).""" - for module in root.modules(): - if getattr(module, _SITE_LINEAR_ATTR, None) is not None: - yield module +def fused_gelu_active(module: nn.Module) -> bool: + """Whether the quality-gated fused path is mounted on ``module``.""" + return _FUSION.is_enabled(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)) - 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 + return _FUSION.mount(root, reject_reason=_site_reject_reason, logger=logger) 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): - setattr(site, _SITE_ENABLED_ATTR, False) + _FUSION.unmount(root) diff --git a/python/sglang/kernels/ops/diffusion/fused_ln_modulate.py b/python/sglang/kernels/ops/diffusion/fused_ln_modulate.py index 7383596cd..29a75699f 100644 --- a/python/sglang/kernels/ops/diffusion/fused_ln_modulate.py +++ b/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) - - -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 + return _FUSION.is_enabled(module) def mount_fused_ln_modulate(root: nn.Module) -> bool: - sites = list(iter_fused_ln_modulate_sites(root)) - for site in sites: - setattr(site, _SITE_ENABLED_ATTR, True) - return bool(sites) + return _FUSION.mount(root) def unmount_fused_ln_modulate(root: nn.Module) -> None: - for site in iter_fused_ln_modulate_sites(root): - setattr(site, _SITE_ENABLED_ATTR, False) + _FUSION.unmount(root) def can_fuse_ln_modulate( diff --git a/python/sglang/kernels/ops/diffusion/modulate_scale_shift.py b/python/sglang/kernels/ops/diffusion/modulate_scale_shift.py index 8a635e79c..7c4d59b99 100644 --- a/python/sglang/kernels/ops/diffusion/modulate_scale_shift.py +++ b/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"ModulateScaleShiftKernel<{args}>::run", + f"modulate_scale_shift::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", +] diff --git a/python/sglang/kernels/ops/diffusion/qknorm_rope.py b/python/sglang/kernels/ops/diffusion/qknorm_rope.py index 80149d5f3..cbf67640d 100644 --- a/python/sglang/kernels/ops/diffusion/qknorm_rope.py +++ b/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 diff --git a/python/sglang/kernels/ops/diffusion/quality_gate.py b/python/sglang/kernels/ops/diffusion/quality_gate.py new file mode 100644 index 000000000..1893c45d1 --- /dev/null +++ b/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) diff --git a/python/sglang/kernels/ops/diffusion/residual_gate_add.py b/python/sglang/kernels/ops/diffusion/residual_gate_add.py index 4e896e868..6bd81aa62 100644 --- a/python/sglang/kernels/ops/diffusion/residual_gate_add.py +++ b/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 gate.shape[row_dim] == 1 and all(size == 1 for size in gate.shape[:-1]) + return 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", +] diff --git a/python/sglang/kernels/ops/diffusion/triton/group_norm_silu_twopass.py b/python/sglang/kernels/ops/diffusion/triton/group_norm_silu_twopass.py index 7af90e106..fa4bad151 100644 --- a/python/sglang/kernels/ops/diffusion/triton/group_norm_silu_twopass.py +++ b/python/sglang/kernels/ops/diffusion/triton/group_norm_silu_twopass.py @@ -133,35 +133,36 @@ 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) - _gn_partial_rows_kernel[(nchunks, n_batch)]( - x3, psum, psq, rows, rows_per_prog, C=c, BLOCK_R=block_r, num_warps=4 - ) - ss = torch.empty((n_batch, 2, c), device=x3.device, dtype=torch.float32) - block_k = max(1, min(4096 // cpg, triton.next_power_of_2(nchunks))) - _gn_finalize_kernel[(num_groups, n_batch)]( - psum, - psq, - weight, - bias, - ss, - nchunks, - rows * cpg, - eps, - c, - CPG=cpg, - BLOCK_K=block_k, - num_warps=4, - ) - y3 = torch.empty_like(x3) - _gn_apply_rows_kernel[(triton.cdiv(rows, block_r), n_batch)]( - x3, ss, y3, rows, C=c, BLOCK_R=block_r, SILU=apply_silu, num_warps=4 - ) + 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 + ) + ss = torch.empty((n_batch, 2, c), device=x3.device, dtype=torch.float32) + block_k = max(1, min(4096 // cpg, triton.next_power_of_2(nchunks))) + _gn_finalize_kernel[(num_groups, n_batch)]( + psum, + psq, + weight, + bias, + ss, + nchunks, + rows * cpg, + eps, + c, + CPG=cpg, + BLOCK_K=block_k, + num_warps=4, + ) + y3 = torch.empty_like(x3) + _gn_apply_rows_kernel[(triton.cdiv(rows, block_r), n_batch)]( + x3, ss, y3, rows, C=c, BLOCK_R=block_r, SILU=apply_silu, num_warps=4 + ) return y3 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 diff --git a/python/sglang/kernels/ops/diffusion/triton/indexed_modulation.py b/python/sglang/kernels/ops/diffusion/triton/indexed_modulation.py index 58b23114c..986138384 100644 --- a/python/sglang/kernels/ops/diffusion/triton/indexed_modulation.py +++ b/python/sglang/kernels/ops/diffusion/triton/indexed_modulation.py @@ -4,14 +4,7 @@ import torch import triton import triton.language as tl - -@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) +from sglang.kernels.ops.diffusion.triton.numerics import round_bf16_to_fp32 @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) - scaled = _round_bf16_to_fp32(x * one_plus_scale) + one_plus_scale = round_bf16_to_fp32(1.0 + 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, diff --git a/python/sglang/kernels/ops/diffusion/triton/layernorm_modulate.py b/python/sglang/kernels/ops/diffusion/triton/layernorm_modulate.py index 432b2808d..976925960 100644 --- a/python/sglang/kernels/ops/diffusion/triton/layernorm_modulate.py +++ b/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) - y = _round_bf16_to_fp32(y * one_plus) + sh + one_plus = round_bf16_to_fp32(1.0 + sc) + 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, :] diff --git a/python/sglang/kernels/ops/diffusion/triton/mps_fallback.py b/python/sglang/kernels/ops/diffusion/triton/mps_fallback.py index 792d99580..6c1770099 100644 --- a/python/sglang/kernels/ops/diffusion/triton/mps_fallback.py +++ b/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, - fuse_scale_shift_kernel_native, + apply_rotary_embedding_native as apply_rotary_embedding_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. diff --git a/python/sglang/kernels/ops/diffusion/triton/native_bf16_rmsnorm.py b/python/sglang/kernels/ops/diffusion/triton/native_bf16_rmsnorm.py new file mode 100644 index 000000000..3b61a1caa --- /dev/null +++ b/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"] diff --git a/python/sglang/kernels/ops/diffusion/triton/numerics.py b/python/sglang/kernels/ops/diffusion/triton/numerics.py new file mode 100644 index 000000000..37c18a774 --- /dev/null +++ b/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", +] diff --git a/python/sglang/kernels/ops/diffusion/triton/rmsnorm_scale_shift_bitexact.py b/python/sglang/kernels/ops/diffusion/triton/rmsnorm_scale_shift_bitexact.py index 59002f493..521d1b0a9 100644 --- a/python/sglang/kernels/ops/diffusion/triton/rmsnorm_scale_shift_bitexact.py +++ b/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) - one_plus = _round_bf16_to_fp32(1.0 + sc) - prod = _round_bf16_to_fp32(y * one_plus) + y = round_bf16_to_fp32(xin * rcp * w) # (bf16)(x * rstd * w) + one_plus = round_bf16_to_fp32(1.0 + sc) + prod = round_bf16_to_fp32(y * one_plus) tl.store(out_ptr + row_base + cols, prod + sh) # store rounds to bf16 diff --git a/python/sglang/kernels/ops/diffusion/triton/sana_wm_gdn_chunkwise.py b/python/sglang/kernels/ops/diffusion/triton/sana_wm_gdn_chunkwise.py index 35ffd2d95..92fd1d104 100644 --- a/python/sglang/kernels/ops/diffusion/triton/sana_wm_gdn_chunkwise.py +++ b/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 diff --git a/python/sglang/kernels/ops/diffusion/triton/scale_shift.py b/python/sglang/kernels/ops/diffusion/triton/scale_shift.py index f5f416c2b..89572fd1a 100644 --- a/python/sglang/kernels/ops/diffusion/triton/scale_shift.py +++ b/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) diff --git a/python/sglang/kernels/ops/diffusion/triton/ulysses_qkv.py b/python/sglang/kernels/ops/diffusion/triton/ulysses_qkv.py index 9364373f9..291d776df 100644 --- a/python/sglang/kernels/ops/diffusion/triton/ulysses_qkv.py +++ b/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() - assert out.dtype == q.dtype and out.device == q.device + if not ( + 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,22 +97,26 @@ def pack_qkv_destination_major( return output block_size = 1024 - _pack_qkv_destination_major_kernel[(triton.cdiv(total_elements, block_size),)]( - output, - q, - k, - v, - total_elements, - rows, - local_heads, - head_size, - q.stride(0), - q.stride(1), - k.stride(0), - k.stride(1), - v.stride(0), - v.stride(1), - BLOCK_SIZE=block_size, - num_warps=8, - ) + with torch.get_device_module().device(q.device): + _pack_qkv_destination_major_kernel[(triton.cdiv(total_elements, block_size),)]( + output, + q, + k, + v, + total_elements, + rows, + local_heads, + head_size, + q.stride(0), + q.stride(1), + k.stride(0), + k.stride(1), + v.stride(0), + v.stride(1), + BLOCK_SIZE=block_size, + num_warps=8, + ) return output + + +__all__ = ["pack_qkv_destination_major"] diff --git a/python/sglang/kernels/ops/diffusion/triton/wan_rmsnorm_silu.py b/python/sglang/kernels/ops/diffusion/triton/wan_rmsnorm_silu.py index 275390cf7..f3d8c4a48 100644 --- a/python/sglang/kernels/ops/diffusion/triton/wan_rmsnorm_silu.py +++ b/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) diff --git a/python/sglang/kernels/ops/diffusion/triton/zimage_native_norm.py b/python/sglang/kernels/ops/diffusion/triton/zimage_native_norm.py index c605cee05..586cdf0d2 100644 --- a/python/sglang/kernels/ops/diffusion/triton/zimage_native_norm.py +++ b/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"] diff --git a/python/sglang/kernels/ops/diffusion/usp_relayout.py b/python/sglang/kernels/ops/diffusion/usp_relayout.py index 0d1483383..38321b4f1 100644 --- a/python/sglang/kernels/ops/diffusion/usp_relayout.py +++ b/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", diff --git a/python/sglang/multimodal_gen/runtime/models/dits/ernie_image.py b/python/sglang/multimodal_gen/runtime/models/dits/ernie_image.py index c7deb1c28..9e6acfe37 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/ernie_image.py +++ b/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 ( - can_use_residual_gate_add_cuda, - residual_gate_add_cuda, -) +from sglang.kernels.ops.diffusion.residual_gate_add import residual_gate_add 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 diff --git a/python/sglang/multimodal_gen/runtime/models/dits/flux.py b/python/sglang/multimodal_gen/runtime/models/dits/flux.py index 4ebf7ac40..3b6cdcafd 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/flux.py +++ b/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 ( - can_use_modulate_scale_shift_cuda, - 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.kernels.ops.diffusion.modulate_scale_shift import modulate_scale_shift +from sglang.kernels.ops.diffusion.residual_gate_add import residual_gate_add 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( - self.proj, hidden_states - ): + if fused_gelu_active(self) and can_fuse_linear_gelu(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( - self.proj, hidden_states - ): + if fused_gelu_active(self) and can_fuse_linear_gelu(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: diff --git a/python/sglang/multimodal_gen/runtime/models/dits/flux_2.py b/python/sglang/multimodal_gen/runtime/models/dits/flux_2.py index e875b0598..ddf9e5c44 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/flux_2.py +++ b/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 ( - can_use_residual_gate_add_cuda, - residual_gate_add_cuda, -) +from sglang.kernels.ops.diffusion.residual_gate_add import residual_gate_add 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_mlp, - scale_mlp, - gate_mlp, + ( + (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_mlp, - c_scale_mlp, - c_gate_mlp, + ( + (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: diff --git a/python/sglang/multimodal_gen/runtime/models/dits/glm_image.py b/python/sglang/multimodal_gen/runtime/models/dits/glm_image.py index ebc066b4a..ad2e10805 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/glm_image.py +++ b/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 ( - can_use_residual_gate_add_cuda, - residual_gate_add_cuda, -) +from sglang.kernels.ops.diffusion.residual_gate_add import residual_gate_add 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( - self.proj, hidden_states - ): + if fused_gelu_active(self) and can_fuse_linear_gelu(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) ) diff --git a/python/sglang/multimodal_gen/runtime/models/dits/ideogram.py b/python/sglang/multimodal_gen/runtime/models/dits/ideogram.py index 69666454d..4e48932a9 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/ideogram.py +++ b/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), diff --git a/python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py b/python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py index 64619287d..ce878e904 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py +++ b/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 ( - can_use_residual_gate_add_cuda, - residual_gate_add_cuda, -) +from sglang.kernels.ops.diffusion.residual_gate_add import residual_gate_add 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, attn_hidden_states, vgate_msa - ) + hidden_states = residual_gate_add(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 diff --git a/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py b/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py index 18107c1b6..70e0cd4b3 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py +++ b/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( - self.proj, hidden_states - ): + if fused_gelu_active(self) and can_fuse_linear_gelu(self.proj, hidden_states): return fused_linear_gelu_tanh( hidden_states, self.proj.weight, self.proj.bias ) diff --git a/python/sglang/multimodal_gen/runtime/models/dits/zimage.py b/python/sglang/multimodal_gen/runtime/models/dits/zimage.py index c413edc27..52541cb01 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/zimage.py +++ b/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 ( - zimage_rmsnorm_tanh_residual, + from sglang.kernels.ops.diffusion.triton.native_bf16_rmsnorm import ( + 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 ( - zimage_rmsnorm_scale as fused_zimage_rmsnorm_scale, + from sglang.kernels.ops.diffusion.triton.native_bf16_rmsnorm import ( + rmsnorm_scale, ) - y = fused_zimage_rmsnorm_scale( + y = rmsnorm_scale( x, norm.weight.data.to(device=x.device, dtype=x.dtype).contiguous(), scale, diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py index a53ed33e2..8f3e91aad 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py +++ b/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 - # fused gate-RMSNorm sites are currently mounted on the transformers. + # Whether request-scoped quality="high" fusions are currently mounted. 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 - kernels are numerically equivalent only at half-precision rounding - level (not bit-exact), so they are mounted for ``quality="high"`` - requests and unmounted otherwise -- the ``"lossless"`` default runs - the unmodified reference path bit-for-bit. ``quality`` participates - in the dynamic-batch signature, so a worker batch is uniform in - ``quality`` and this process-wide transition is safe at the batch - 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. + These fusions are numerically equivalent only at half-precision + rounding level (not bit-exact), so they are mounted for + ``quality="high"`` requests and unmounted otherwise. The + ``"lossless"`` default runs the reference path bit-for-bit. ``quality`` + participates in the dynamic-batch signature, making this transition + safe at the batch boundary. Mounting is all-or-nothing per transformer + and fusion family; 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_gate_norm = False - mounted_ln_modulate = False + mounted_fusions: set[str] = set() for transformer in filter(None, [self.transformer, self.transformer_2]): - if want: - mounted_gelu |= mount_fused_linear_gelu(transformer) - mounted_gate_norm |= mount_fused_gate_rmsnorm(transformer) - mounted_ln_modulate |= mount_fused_ln_modulate(transformer) - else: - unmount_fused_linear_gelu(transformer) - unmount_fused_gate_rmsnorm(transformer) - unmount_fused_ln_modulate(transformer) + for description, mount, unmount in _QUALITY_FUSION_HANDLERS: + if want: + if mount(transformer): + mounted_fusions.add(description) + else: + unmount(transformer) self._quality_fusions_mounted = want - if want and mounted_gelu: - logger.info( - "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" - ) + for description in sorted(mounted_fusions): + logger.info("Mounted %s for quality=high", description) def _cache_dit_dual_model_name(self) -> str: return "wan2.2" diff --git a/test/registered/kernels/ops/diffusion/test_ernie_residual_gate_add.py b/test/registered/kernels/ops/diffusion/test_ernie_residual_gate_add.py deleted file mode 100644 index 221a762ef..000000000 --- a/test/registered/kernels/ops/diffusion/test_ernie_residual_gate_add.py +++ /dev/null @@ -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__])) diff --git a/test/registered/kernels/ops/diffusion/test_flux2_vae_fastpath.py b/test/registered/kernels/ops/diffusion/test_flux2_vae_fastpath.py index ba796ce0e..9796f213e 100644 --- a/test/registered/kernels/ops/diffusion/test_flux2_vae_fastpath.py +++ b/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) diff --git a/test/registered/kernels/ops/diffusion/test_fused_gate_rmsnorm.py b/test/registered/kernels/ops/diffusion/test_fused_gate_rmsnorm.py index 5a300580b..57d72ecd1 100644 --- a/test/registered/kernels/ops/diffusion/test_fused_gate_rmsnorm.py +++ b/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__": diff --git a/test/registered/kernels/ops/diffusion/test_fused_linear_gelu.py b/test/registered/kernels/ops/diffusion/test_fused_linear_gelu.py index 0976a53d9..e30a9f8c9 100644 --- a/test/registered/kernels/ops/diffusion/test_fused_linear_gelu.py +++ b/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__])) diff --git a/test/registered/kernels/ops/diffusion/test_fused_ln_modulate.py b/test/registered/kernels/ops/diffusion/test_fused_ln_modulate.py index 049350ddf..ebc6ed72c 100644 --- a/test/registered/kernels/ops/diffusion/test_fused_ln_modulate.py +++ b/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 diff --git a/test/registered/kernels/ops/diffusion/test_modulate_scale_shift.py b/test/registered/kernels/ops/diffusion/test_modulate_scale_shift.py index a0aa6d8e8..68d841d8d 100644 --- a/test/registered/kernels/ops/diffusion/test_modulate_scale_shift.py +++ b/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__": diff --git a/test/registered/kernels/ops/diffusion/test_zimage_native_norm.py b/test/registered/kernels/ops/diffusion/test_native_bf16_rmsnorm.py similarity index 74% rename from test/registered/kernels/ops/diffusion/test_zimage_native_norm.py rename to test/registered/kernels/ops/diffusion/test_native_bf16_rmsnorm.py index 6c1b1d959..6c50d1c32 100644 --- a/test/registered/kernels/ops/diffusion/test_zimage_native_norm.py +++ b/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 ( - zimage_rmsnorm_scale, - zimage_rmsnorm_tanh_residual, +from sglang.kernels.ops.diffusion.triton.native_bf16_rmsnorm import ( + rmsnorm_scale, + 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 zimage_rmsnorm_tanh_residual(x, modulation, residual, weight, EPS) is None + assert rmsnorm_scale(x, weight, modulation, 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 zimage_rmsnorm_tanh_residual(x, modulation, residual, weight, EPS) is None + assert rmsnorm_scale(x, weight, modulation, EPS) is None + assert rmsnorm_tanh_residual(x, modulation, residual, weight, EPS) is None if __name__ == "__main__": diff --git a/test/registered/kernels/ops/diffusion/test_qknorm_rope.py b/test/registered/kernels/ops/diffusion/test_qknorm_rope.py index 2030e78ce..98c96d092 100644 --- a/test/registered/kernels/ops/diffusion/test_qknorm_rope.py +++ b/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"])) diff --git a/test/registered/kernels/ops/diffusion/test_quality_gate.py b/test/registered/kernels/ops/diffusion/test_quality_gate.py new file mode 100644 index 000000000..ba90c52eb --- /dev/null +++ b/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"])) diff --git a/test/registered/kernels/ops/diffusion/test_residual_gate_add.py b/test/registered/kernels/ops/diffusion/test_residual_gate_add.py index b45cddb0a..f281e7c89 100644 --- a/test/registered/kernels/ops/diffusion/test_residual_gate_add.py +++ b/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) diff --git a/test/registered/kernels/ops/diffusion/test_scale_shift.py b/test/registered/kernels/ops/diffusion/test_scale_shift.py new file mode 100644 index 000000000..aff10d181 --- /dev/null +++ b/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"])) diff --git a/test/registered/kernels/ops/diffusion/test_timestep_embedding.py b/test/registered/kernels/ops/diffusion/test_timestep_embedding.py index fb1a904e1..59e13a672 100644 --- a/test/registered/kernels/ops/diffusion/test_timestep_embedding.py +++ b/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 diff --git a/test/registered/kernels/ops/diffusion/test_ulysses_qkv.py b/test/registered/kernels/ops/diffusion/test_ulysses_qkv.py new file mode 100644 index 000000000..7dec40df2 --- /dev/null +++ b/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"])) diff --git a/test/registered/kernels/ops/diffusion/test_varlen_pack_pad.py b/test/registered/kernels/ops/diffusion/test_varlen_pack_pad.py index a99aecda5..222f1d1a1 100644 --- a/test/registered/kernels/ops/diffusion/test_varlen_pack_pad.py +++ b/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 diff --git a/test/registered/kernels/ops/diffusion/test_wan_vae_fastpath.py b/test/registered/kernels/ops/diffusion/test_wan_vae_fastpath.py index b45ced7a5..260f391c7 100644 --- a/test/registered/kernels/ops/diffusion/test_wan_vae_fastpath.py +++ b/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"]))