minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 14 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Xiaoyu Zhang
2026-08-09 00:37:00 +08:00
committed by GitHub
parent ec5199b906
commit dc9624deb2
50 changed files with 1358 additions and 1376 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/kernels/jit/csrc/diffusion/modulate_scale_shift.cuh
+93 -177
View File
@@ -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 <sgl_kernel/tensor.h> // For host dtype helpers and TensorView metadata
#include <sgl_kernel/utils.h> // For RuntimeCheck and div_ceil
#include <sgl_kernel/tensor.h>
#include <sgl_kernel/utils.h>
#include <sgl_kernel/type.cuh> // For DTypeTrait conversions
#include <sgl_kernel/utils.cuh> // For LaunchKernel and CUDA dtype aliases
#include <sgl_kernel/vec.cuh> // For device::AlignedVector
#include <sgl_kernel/type.cuh>
#include <sgl_kernel/utils.cuh>
#include <sgl_kernel/vec.cuh>
#include <algorithm>
#include <cstdint>
#include <type_traits>
namespace sglang {
namespace sglang_modulate_scale_shift {
namespace modulate_scale_shift {
namespace {
constexpr int kRowsPerBlock = 4;
constexpr int kColsVecPerBlock = 256;
constexpr int64_t kMaxGrid = 65535;
inline const char* data_ptr(const tvm::ffi::TensorView& t) {
return static_cast<const char*>(t.data_ptr()) + t.byte_offset();
}
inline char* mutable_data_ptr(const tvm::ffi::TensorView& t) {
return static_cast<char*>(t.data_ptr()) + t.byte_offset();
}
inline bool aligned16(const void* p) {
return (reinterpret_cast<uintptr_t>(p) & 0xF) == 0;
}
inline int64_t numel(const tvm::ffi::TensorView& t) {
int64_t n = 1;
for (int i = 0; i < t.ndim(); ++i) {
n *= t.size(i);
}
return n;
}
inline bool is_dense_contiguous(const tvm::ffi::TensorView& t) {
int64_t expected = 1;
for (int i = t.ndim() - 1; i >= 0; --i) {
if (t.size(i) == 1) {
continue;
}
if (t.stride(i) != expected) {
return false;
}
expected *= t.size(i);
}
return true;
}
constexpr uint32_t kRowsPerBlock = 4;
constexpr uint32_t kColsVecPerBlock = 256;
constexpr uint32_t kMaxGridY = 65535;
constexpr uintptr_t kAlignment = 16;
template <typename T>
inline void check_dtype(const tvm::ffi::TensorView& t) {
host::RuntimeCheck(host::is_type<T>(t.dtype()), "unexpected dtype for modulate_scale_shift");
}
template <typename T>
__device__ __forceinline__ float to_float(T v) {
return static_cast<float>(v);
}
template <>
__device__ __forceinline__ float to_float<fp16_t>(fp16_t v) {
return __half2float(v);
}
template <>
__device__ __forceinline__ float to_float<bf16_t>(bf16_t v) {
return __bfloat162float(v);
}
template <typename T>
__device__ __forceinline__ T modulate_value(T x, T scale, T shift) {
// Round each intermediate back to T (the eager chain's kernel boundaries;
// also blocks fmul+fadd FMA contraction).
const T one_plus_scale = DTypeTrait<T>::from(1.0f + to_float(scale));
const T product = DTypeTrait<T>::from(to_float(x) * to_float(one_plus_scale));
return DTypeTrait<T>::from(to_float(product) + to_float(shift));
SGL_DEVICE T modulate_value(T x, T scale, T shift) {
const T one_plus_scale = device::cast<T>(1.0f + device::cast<fp32_t>(scale));
const T product = device::cast<T>(device::cast<fp32_t>(x) * device::cast<fp32_t>(one_plus_scale));
return device::cast<T>(device::cast<fp32_t>(product) + device::cast<fp32_t>(shift));
}
template <typename T, int kVec>
__global__ void modulate_scale_shift_vec_kernel(
__global__ void modulate_scale_shift_kernel(
T* __restrict__ out,
const T* __restrict__ x,
const T* __restrict__ scale,
const T* __restrict__ shift,
T* __restrict__ out,
int64_t rows,
int64_t rows_per_batch,
int64_t row_vec) {
@@ -112,114 +49,93 @@ __global__ void modulate_scale_shift_vec_kernel(
return;
}
// Grid-stride: the row-tile count can exceed the gridDim.y hardware limit.
const int64_t row_tile_stride = static_cast<int64_t>(gridDim.y) * kRowsPerBlock;
for (int64_t row_base = static_cast<int64_t>(blockIdx.y) * kRowsPerBlock; row_base < rows;
row_base += row_tile_stride) {
const int64_t row_stride = static_cast<int64_t>(gridDim.y) * kRowsPerBlock;
for (int64_t row_base = static_cast<int64_t>(blockIdx.y) * kRowsPerBlock; row_base < rows; row_base += row_stride) {
#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 <typename T>
inline void launch_modulate_scale_shift(
const tvm::ffi::TensorView& out,
const tvm::ffi::TensorView& x,
const tvm::ffi::TensorView& scale,
const tvm::ffi::TensorView& shift) {
const int64_t total = numel(x);
if (total == 0) {
return;
}
const int64_t D = x.size(x.ndim() - 1);
const int64_t rows = total / D;
const int64_t batches = scale.size(0);
const int64_t rows_per_batch = rows / batches;
const T* x_ptr = reinterpret_cast<const T*>(data_ptr(x));
const T* scale_ptr = reinterpret_cast<const T*>(data_ptr(scale));
const T* shift_ptr = reinterpret_cast<const T*>(data_ptr(shift));
T* out_ptr = reinterpret_cast<T*>(mutable_data_ptr(out));
constexpr int kVec = 16 / sizeof(T);
host::RuntimeCheck(
aligned16(x_ptr) && aligned16(scale_ptr) && aligned16(shift_ptr) && aligned16(out_ptr),
"modulate_scale_shift requires 16-byte aligned tensors");
host::RuntimeCheck(D % kVec == 0, "modulate_scale_shift requires D to be a multiple of the vector width");
const int64_t row_vec = D / kVec;
const int64_t col_blocks = host::div_ceil(row_vec, static_cast<int64_t>(kColsVecPerBlock));
const int64_t row_tiles = host::div_ceil(rows, static_cast<int64_t>(kRowsPerBlock));
const int64_t row_blocks = row_tiles > kMaxGrid ? kMaxGrid : row_tiles;
host::LaunchKernel(
dim3(static_cast<uint32_t>(col_blocks), static_cast<uint32_t>(row_blocks)), dim3(kColsVecPerBlock), out.device())(
modulate_scale_shift_vec_kernel<T, kVec>, x_ptr, scale_ptr, shift_ptr, out_ptr, rows, rows_per_batch, row_vec);
}
template <typename T>
inline void validate_modulate_scale_shift(
const tvm::ffi::TensorView& out,
const tvm::ffi::TensorView& x,
const tvm::ffi::TensorView& scale,
const tvm::ffi::TensorView& shift) {
check_dtype<T>(out);
check_dtype<T>(x);
check_dtype<T>(scale);
check_dtype<T>(shift);
host::RuntimeCheck(x.device().device_type == kDLCUDA, "x must be CUDA");
host::RuntimeCheck(scale.device().device_type == kDLCUDA, "scale must be CUDA");
host::RuntimeCheck(shift.device().device_type == kDLCUDA, "shift must be CUDA");
host::RuntimeCheck(out.device().device_type == kDLCUDA, "out must be CUDA");
host::RuntimeCheck(
x.device().device_id == scale.device().device_id && x.device().device_id == shift.device().device_id &&
x.device().device_id == out.device().device_id,
"x/scale/shift/out must be on the same CUDA device");
host::RuntimeCheck(x.ndim() == 3, "x must be [B, L, D]");
host::RuntimeCheck(scale.ndim() == 2, "scale must be [B, D]");
host::RuntimeCheck(shift.ndim() == 2, "shift must be [B, D]");
host::RuntimeCheck(out.ndim() == x.ndim(), "out rank must match x");
for (int i = 0; i < x.ndim(); ++i) {
host::RuntimeCheck(out.size(i) == x.size(i), "out shape must match x");
}
host::RuntimeCheck(scale.size(0) == x.size(0), "scale batch dim must match x");
host::RuntimeCheck(scale.size(1) == x.size(2), "scale last dim must match x");
host::RuntimeCheck(shift.size(0) == scale.size(0) && shift.size(1) == scale.size(1), "shift shape must match scale");
host::RuntimeCheck(is_dense_contiguous(x), "x must be contiguous");
host::RuntimeCheck(is_dense_contiguous(scale), "scale must be contiguous");
host::RuntimeCheck(is_dense_contiguous(shift), "shift must be contiguous");
host::RuntimeCheck(is_dense_contiguous(out), "out must be contiguous");
host::RuntimeCheck(data_ptr(out) != data_ptr(x), "out must not alias x");
host::RuntimeCheck(data_ptr(out) != data_ptr(scale), "out must not alias scale");
host::RuntimeCheck(data_ptr(out) != data_ptr(shift), "out must not alias shift");
}
} // namespace
/**
* \brief Validate and launch bit-exact diffusion adaLN modulation.
*
* \tparam T Activation type: fp16_t or bf16_t.
*/
template <typename T>
struct ModulateScaleShiftKernel {
static_assert(std::is_same_v<T, fp16_t> || std::is_same_v<T, bf16_t>);
/**
* \param out Output tensor with shape [B, L, D].
* \param x Input tensor with shape [B, L, D].
* \param scale Scale tensor with shape [B, D].
* \param shift Shift tensor with shape [B, D].
*/
static void
run(tvm::ffi::TensorView out, tvm::ffi::TensorView x, tvm::ffi::TensorView scale, tvm::ffi::TensorView shift) {
validate_modulate_scale_shift<T>(out, x, scale, shift);
launch_modulate_scale_shift<T>(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<kDLCUDA>();
TensorMatcher({B, L, D}).with_dtype<T>().with_device(device).verify(out).verify(x);
TensorMatcher({B, D}).with_dtype<T>().with_device(device).verify(scale).verify(shift);
const int64_t batch = B.unwrap();
const int64_t sequence_length = L.unwrap();
const int64_t hidden_size = D.unwrap();
const int64_t rows = batch * sequence_length;
if (rows == 0 || hidden_size == 0) {
return;
}
constexpr int kVec = kAlignment / sizeof(T);
CHECK_HOST(hidden_size % kVec == 0) << "hidden size must be a multiple of " << kVec;
auto* out_ptr = static_cast<T*>(out.data_ptr());
const auto* x_ptr = static_cast<const T*>(x.data_ptr());
const auto* scale_ptr = static_cast<const T*>(scale.data_ptr());
const auto* shift_ptr = static_cast<const T*>(shift.data_ptr());
CHECK_HOST(
reinterpret_cast<uintptr_t>(out_ptr) % kAlignment == 0 &&
reinterpret_cast<uintptr_t>(x_ptr) % kAlignment == 0 &&
reinterpret_cast<uintptr_t>(scale_ptr) % kAlignment == 0 &&
reinterpret_cast<uintptr_t>(shift_ptr) % kAlignment == 0)
<< "modulate_scale_shift requires 16-byte aligned tensors";
CHECK_HOST(out_ptr != x_ptr && out_ptr != scale_ptr && out_ptr != shift_ptr) << "output must not alias an input";
const int64_t row_vec = hidden_size / kVec;
const auto col_blocks = static_cast<uint32_t>(div_ceil(row_vec, static_cast<int64_t>(kColsVecPerBlock)));
const int64_t row_tiles = div_ceil(rows, static_cast<int64_t>(kRowsPerBlock));
const auto row_blocks = static_cast<uint32_t>(std::min<int64_t>(row_tiles, kMaxGridY));
LaunchKernel(dim3(col_blocks, row_blocks), kColsVecPerBlock, device.unwrap())(
modulate_scale_shift_kernel<T, kVec>, out_ptr, x_ptr, scale_ptr, shift_ptr, rows, sequence_length, row_vec);
}
};
} // namespace sglang_modulate_scale_shift
} // namespace modulate_scale_shift
} // namespace sglang
python/sglang/kernels/jit/csrc/diffusion/qknorm_rope.cuh
+1
View File
@@ -277,6 +277,7 @@ struct QKNormRopeKernel {
const auto num_tokens = static_cast<uint32_t>(N.unwrap());
const auto num_qo_heads = static_cast<uint32_t>(Q.unwrap());
const auto num_kv_heads = static_cast<uint32_t>(K.unwrap());
if (num_tokens == 0 || (num_qo_heads == 0 && num_kv_heads == 0)) return;
const auto q_stride_bytes = static_cast<int64_t>(Dq.unwrap() * sizeof(DType));
const auto k_stride_bytes = static_cast<int64_t>(Dk.unwrap() * sizeof(DType));
const auto head_stride_bytes = static_cast<int64_t>(Dd.unwrap() * sizeof(DType));
python/sglang/kernels/jit/csrc/diffusion/residual_gate_add.cuh
+143 -252
View File
@@ -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 <sgl_kernel/tensor.h> // For host dtype helpers and TensorView metadata
#include <sgl_kernel/utils.h> // For RuntimeCheck and div_ceil
#include <sgl_kernel/tensor.h>
#include <sgl_kernel/utils.h>
#include <sgl_kernel/type.cuh> // For DTypeTrait conversions
#include <sgl_kernel/utils.cuh> // For LaunchKernel and CUDA dtype aliases
#include <sgl_kernel/vec.cuh> // For device::AlignedVector
#include <sgl_kernel/type.cuh>
#include <sgl_kernel/utils.cuh>
#include <sgl_kernel/vec.cuh>
#include <algorithm>
#include <cstdint>
#include <type_traits>
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<const char*>(t.data_ptr()) + t.byte_offset();
}
inline char* mutable_data_ptr(const tvm::ffi::TensorView& t) {
return static_cast<char*>(t.data_ptr()) + t.byte_offset();
}
inline bool aligned16(const void* p) {
return (reinterpret_cast<uintptr_t>(p) & 0xF) == 0;
}
inline int64_t numel(const tvm::ffi::TensorView& t) {
int64_t n = 1;
for (int i = 0; i < t.ndim(); ++i) {
n *= t.size(i);
}
return n;
}
inline int64_t grid_for(int64_t total) {
int64_t grid = host::div_ceil(total, static_cast<int64_t>(kBlockSize));
if (grid < 1) {
grid = 1;
}
if (grid > kMaxGrid) {
grid = kMaxGrid;
}
return grid;
}
inline bool is_dense_contiguous(const tvm::ffi::TensorView& t) {
int64_t expected = 1;
for (int i = t.ndim() - 1; i >= 0; --i) {
if (t.size(i) == 1) {
continue;
}
if (t.stride(i) != expected) {
return false;
}
expected *= t.size(i);
}
return true;
}
enum class GateMode : int { kFull, kBroadcastRow };
template <typename T>
inline void check_dtype(const tvm::ffi::TensorView& t) {
host::RuntimeCheck(host::is_type<T>(t.dtype()), "unexpected dtype for residual_gate_add");
}
template <typename T>
__device__ __forceinline__ float to_float(T v) {
return static_cast<float>(v);
}
template <>
__device__ __forceinline__ float to_float<fp16_t>(fp16_t v) {
return __half2float(v);
}
template <>
__device__ __forceinline__ float to_float<bf16_t>(bf16_t v) {
return __bfloat162float(v);
}
template <typename T>
__device__ __forceinline__ T residual_gate_value(T residual, T update, T gate) {
const T product = DTypeTrait<T>::from(to_float(update) * to_float(gate));
return DTypeTrait<T>::from(to_float(residual) + to_float(product));
SGL_DEVICE T residual_gate_value(T residual, T update, T gate) {
const T product = device::cast<T>(device::cast<fp32_t>(update) * device::cast<fp32_t>(gate));
return device::cast<T>(device::cast<fp32_t>(residual) + device::cast<fp32_t>(product));
}
template <typename T, int kVec>
__global__ void residual_gate_add_vec_kernel(
T* __restrict__ out,
const T* __restrict__ residual,
const T* __restrict__ update,
const T* __restrict__ gate,
T* __restrict__ out,
int64_t n_vec) {
int64_t num_vectors) {
using Vec = device::AlignedVector<T, kVec>;
const int64_t stride = static_cast<int64_t>(gridDim.x) * blockDim.x;
for (int64_t v = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x; v < n_vec; v += stride) {
Vec r, u, g, o;
r.load(residual, v);
u.load(update, v);
g.load(gate, v);
for (int64_t vector = static_cast<int64_t>(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 <typename T, int kVec>
__global__ void residual_gate_add_bcast_row_tile_kernel(
__global__ void residual_gate_add_broadcast_kernel(
T* __restrict__ out,
const T* __restrict__ residual,
const T* __restrict__ update,
const T* __restrict__ gate,
T* __restrict__ out,
int64_t rows,
int64_t row_vec) {
int64_t row_vectors) {
using Vec = device::AlignedVector<T, kVec>;
const int64_t col_vec = static_cast<int64_t>(blockIdx.x) * kBcastColsVecPerBlock + threadIdx.x;
if (col_vec >= row_vec) {
const int64_t column = static_cast<int64_t>(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<int64_t>(gridDim.y) * kBcastRowsPerBlock;
for (int64_t row_base = static_cast<int64_t>(blockIdx.y) * kBcastRowsPerBlock; row_base < rows;
row_base += row_tile_stride) {
Vec gate_vec;
gate_vec.load(gate, column);
const int64_t row_stride = static_cast<int64_t>(gridDim.y) * kBroadcastRowsPerBlock;
for (int64_t row_base = static_cast<int64_t>(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 <typename T, GateMode kGate>
template <typename T, GateMode kGateMode>
__global__ void residual_gate_add_scalar_kernel(
T* __restrict__ out,
const T* __restrict__ residual,
const T* __restrict__ update,
const T* __restrict__ gate,
T* __restrict__ out,
int64_t begin,
int64_t total,
int64_t D) {
int64_t numel,
int64_t hidden_size) {
const int64_t stride = static_cast<int64_t>(gridDim.x) * blockDim.x;
for (int64_t i = begin + static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x; i < total; i += stride) {
const T gate_value = kGate == GateMode::kFull ? gate[i] : SGLANG_LDG(gate + (i % D));
out[i] = residual_gate_value(residual[i], update[i], gate_value);
for (int64_t index = static_cast<int64_t>(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 <typename T>
inline void launch_residual_gate_add(
const tvm::ffi::TensorView& out,
const tvm::ffi::TensorView& residual,
const tvm::ffi::TensorView& update,
const tvm::ffi::TensorView& gate,
GateMode mode) {
const int64_t total = numel(residual);
if (total == 0) {
return;
}
const int64_t D = residual.size(residual.ndim() - 1);
const T* residual_ptr = reinterpret_cast<const T*>(data_ptr(residual));
const T* update_ptr = reinterpret_cast<const T*>(data_ptr(update));
const T* gate_ptr = reinterpret_cast<const T*>(data_ptr(gate));
T* out_ptr = reinterpret_cast<T*>(mutable_data_ptr(out));
constexpr int kVec = 16 / sizeof(T);
const bool vec_ok = aligned16(residual_ptr) && aligned16(update_ptr) && aligned16(gate_ptr) && aligned16(out_ptr) &&
(D % kVec == 0) && (mode == GateMode::kBcastRow || total % kVec == 0);
int64_t done = 0;
if (vec_ok) {
const int64_t n_vec = total / kVec;
const int64_t row_vec = D / kVec;
if (mode == GateMode::kFull) {
host::LaunchKernel(static_cast<uint32_t>(grid_for(n_vec)), kBlockSize, out.device())(
residual_gate_add_vec_kernel<T, kVec>, residual_ptr, update_ptr, gate_ptr, out_ptr, n_vec);
} else {
const int64_t rows = total / D;
const int64_t col_blocks = host::div_ceil(row_vec, static_cast<int64_t>(kBcastColsVecPerBlock));
const int64_t row_tiles = host::div_ceil(rows, static_cast<int64_t>(kBcastRowsPerBlock));
const int64_t row_blocks = row_tiles > kMaxGrid ? kMaxGrid : row_tiles;
host::LaunchKernel(
dim3(static_cast<uint32_t>(col_blocks), static_cast<uint32_t>(row_blocks)),
dim3(kBcastColsVecPerBlock),
out.device())(
residual_gate_add_bcast_row_tile_kernel<T, kVec>, residual_ptr, update_ptr, gate_ptr, out_ptr, rows, row_vec);
}
done = n_vec * kVec;
}
if (done < total) {
if (mode == GateMode::kFull) {
host::LaunchKernel(static_cast<uint32_t>(grid_for(total - done)), kBlockSize, out.device())(
residual_gate_add_scalar_kernel<T, GateMode::kFull>,
residual_ptr,
update_ptr,
gate_ptr,
out_ptr,
done,
total,
D);
} else {
host::LaunchKernel(static_cast<uint32_t>(grid_for(total - done)), kBlockSize, out.device())(
residual_gate_add_scalar_kernel<T, GateMode::kBcastRow>,
residual_ptr,
update_ptr,
gate_ptr,
out_ptr,
done,
total,
D);
}
}
}
template <typename T>
inline GateMode validate_residual_gate_add(
const tvm::ffi::TensorView& out,
const tvm::ffi::TensorView& residual,
const tvm::ffi::TensorView& update,
const tvm::ffi::TensorView& gate) {
check_dtype<T>(out);
check_dtype<T>(residual);
check_dtype<T>(update);
check_dtype<T>(gate);
host::RuntimeCheck(residual.device().device_type == kDLCUDA, "residual must be CUDA");
host::RuntimeCheck(update.device().device_type == kDLCUDA, "update must be CUDA");
host::RuntimeCheck(gate.device().device_type == kDLCUDA, "gate must be CUDA");
host::RuntimeCheck(out.device().device_type == kDLCUDA, "out must be CUDA");
host::RuntimeCheck(
residual.device().device_id == update.device().device_id &&
residual.device().device_id == gate.device().device_id &&
residual.device().device_id == out.device().device_id,
"residual/update/gate/out must be on the same CUDA device");
host::RuntimeCheck(residual.ndim() >= 2, "residual must be at least 2D");
host::RuntimeCheck(update.ndim() == residual.ndim(), "update rank must match residual");
host::RuntimeCheck(out.ndim() == residual.ndim(), "out rank must match residual");
for (int i = 0; i < residual.ndim(); ++i) {
host::RuntimeCheck(update.size(i) == residual.size(i), "update shape must match residual");
host::RuntimeCheck(out.size(i) == residual.size(i), "out shape must match residual");
}
host::RuntimeCheck(is_dense_contiguous(residual), "residual must be contiguous");
host::RuntimeCheck(is_dense_contiguous(update), "update must be contiguous");
host::RuntimeCheck(is_dense_contiguous(out), "out must be contiguous");
host::RuntimeCheck(is_dense_contiguous(gate), "gate must be contiguous");
host::RuntimeCheck(data_ptr(out) != data_ptr(residual), "out must not alias residual");
host::RuntimeCheck(data_ptr(out) != data_ptr(update), "out must not alias update");
host::RuntimeCheck(data_ptr(out) != data_ptr(gate), "out must not alias gate");
const int D_dim = residual.ndim() - 1;
const int row_dim = residual.ndim() - 2;
host::RuntimeCheck(gate.ndim() == residual.ndim(), "gate rank must match residual");
host::RuntimeCheck(gate.size(D_dim) == residual.size(D_dim), "gate last dim must match residual");
bool full_gate = true;
for (int i = 0; i < residual.ndim(); ++i) {
full_gate = full_gate && gate.size(i) == residual.size(i);
}
if (full_gate) {
return GateMode::kFull;
}
host::RuntimeCheck(gate.size(row_dim) == 1, "broadcast gate row dim must be 1");
for (int i = 0; i < D_dim; ++i) {
host::RuntimeCheck(gate.size(i) == 1, "broadcast gate leading dims must be 1");
}
return GateMode::kBcastRow;
}
} // namespace
/**
* \brief Validate and launch a bit-exact residual-gate update.
*
* Python flattens the tensors before dispatch. A broadcast gate contains one
* hidden-size row; a full gate has the same number of elements as the inputs.
*/
template <typename T>
struct ResidualGateAddKernel {
static_assert(std::is_same_v<T, fp16_t> || std::is_same_v<T, bf16_t> || std::is_same_v<T, fp32_t>);
static void
run(tvm::ffi::TensorView out, tvm::ffi::TensorView residual, tvm::ffi::TensorView update, tvm::ffi::TensorView gate) {
const GateMode mode = validate_residual_gate_add<T>(out, residual, update, gate);
launch_residual_gate_add<T>(out, residual, update, gate, mode);
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<kDLCUDA>();
TensorMatcher({N}).with_dtype<T>().with_device(device).verify(out).verify(residual).verify(update);
TensorMatcher({G}).with_dtype<T>().with_device(device).verify(gate);
const int64_t numel = N.unwrap();
CHECK_HOST(hidden_size > 0 && numel % hidden_size == 0) << "hidden size must be positive and divide the input size";
CHECK_HOST(G.unwrap() == (broadcast_gate ? hidden_size : numel)) << "gate size does not match its mode";
if (numel == 0) {
return;
}
auto* out_ptr = static_cast<T*>(out.data_ptr());
const auto* residual_ptr = static_cast<const T*>(residual.data_ptr());
const auto* update_ptr = static_cast<const T*>(update.data_ptr());
const auto* gate_ptr = static_cast<const T*>(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<uintptr_t>(out_ptr) % kAlignment == 0 &&
reinterpret_cast<uintptr_t>(residual_ptr) % kAlignment == 0 &&
reinterpret_cast<uintptr_t>(update_ptr) % kAlignment == 0 &&
reinterpret_cast<uintptr_t>(gate_ptr) % kAlignment == 0;
const bool vectorized = aligned && hidden_size % kVec == 0;
if (vectorized) {
const int64_t num_vectors = numel / kVec;
if (!broadcast_gate) {
const auto blocks =
static_cast<uint32_t>(std::min<int64_t>(div_ceil(num_vectors, static_cast<int64_t>(kBlockSize)), kMaxGrid));
LaunchKernel(blocks, kBlockSize, device.unwrap())(
residual_gate_add_vec_kernel<T, kVec>, out_ptr, residual_ptr, update_ptr, gate_ptr, num_vectors);
return;
}
const int64_t rows = numel / hidden_size;
const int64_t row_vectors = hidden_size / kVec;
const auto column_blocks =
static_cast<uint32_t>(div_ceil(row_vectors, static_cast<int64_t>(kBroadcastColsPerBlock)));
const int64_t row_tiles = div_ceil(rows, static_cast<int64_t>(kBroadcastRowsPerBlock));
const auto row_blocks = static_cast<uint32_t>(std::min<int64_t>(row_tiles, kMaxGrid));
LaunchKernel(dim3(column_blocks, row_blocks), kBroadcastColsPerBlock, device.unwrap())(
residual_gate_add_broadcast_kernel<T, kVec>, out_ptr, residual_ptr, update_ptr, gate_ptr, rows, row_vectors);
return;
}
const auto blocks =
static_cast<uint32_t>(std::min<int64_t>(div_ceil(numel, static_cast<int64_t>(kBlockSize)), kMaxGrid));
if (broadcast_gate) {
LaunchKernel(blocks, kBlockSize, device.unwrap())(
residual_gate_add_scalar_kernel<T, GateMode::kBroadcastRow>,
out_ptr,
residual_ptr,
update_ptr,
gate_ptr,
numel,
hidden_size);
} else {
LaunchKernel(blocks, kBlockSize, device.unwrap())(
residual_gate_add_scalar_kernel<T, GateMode::kFull>,
out_ptr,
residual_ptr,
update_ptr,
gate_ptr,
numel,
hidden_size);
}
}
};
python/sglang/kernels/jit/csrc/diffusion/usp_relayout.cuh
+97 -132
View File
@@ -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 <sgl_kernel/tensor.h> // For host dtype helpers and TensorView metadata
#include <sgl_kernel/utils.h> // For RuntimeCheck and div_ceil
#include <sgl_kernel/tensor.h>
#include <sgl_kernel/utils.h>
#include <sgl_kernel/type.cuh> // For CUDA dtype aliases
#include <sgl_kernel/utils.cuh> // For LaunchKernel
#include <sgl_kernel/vec.cuh> // For device::AlignedVector
#include <sgl_kernel/type.cuh>
#include <sgl_kernel/utils.cuh>
#include <sgl_kernel/vec.cuh>
#include <algorithm>
#include <cstdint>
#include <type_traits>
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<const char*>(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<char*>(t.data_ptr()) + t.byte_offset();
}
inline bool aligned16(const void* p) {
return (reinterpret_cast<uintptr_t>(p) & 0xF) == 0;
}
inline int64_t numel(const tvm::ffi::TensorView& t) {
int64_t n = 1;
for (int i = 0; i < t.ndim(); ++i) {
n *= t.size(i);
}
return n;
}
inline int64_t grid_for(int64_t total) {
int64_t grid = host::div_ceil(total, static_cast<int64_t>(kBlockSize));
if (grid < 1) {
grid = 1;
}
if (grid > kMaxGrid) {
grid = kMaxGrid;
}
return grid;
}
inline bool is_dense_contiguous(const tvm::ffi::TensorView& t) {
int64_t expected = 1;
for (int i = t.ndim() - 1; i >= 0; --i) {
if (t.size(i) == 1) {
continue;
}
if (t.stride(i) != expected) {
return false;
}
expected *= t.size(i);
}
return true;
}
template <typename T>
inline void check_dtype(const tvm::ffi::TensorView& t) {
host::RuntimeCheck(host::is_type<T>(t.dtype()), "unexpected dtype for usp_merge_heads tensor");
}
// out[b, s, w, h, c] = x[w, s, b, h, c]
// out[b, s, w, h, d] = x[w, s, b, h, d]
template <typename T, int kVec>
__global__ void usp_merge_heads_vec_kernel(
T* __restrict__ out,
const T* __restrict__ x,
int64_t n_vec,
int64_t 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<int64_t>(gridDim.x) * blockDim.x;
for (int64_t i = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x; i < n_vec; i += stride) {
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<int64_t>(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<T, kVec> 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<T, kVec> value;
value.load(x, source);
value.store(out, index);
}
}
@@ -113,66 +61,83 @@ template <typename T>
__global__ void usp_merge_heads_scalar_kernel(
T* __restrict__ out,
const T* __restrict__ x,
int64_t total,
int64_t numel,
int64_t head_dim,
int64_t h_local,
int64_t local_heads,
int64_t batch,
int64_t seq,
int64_t world) {
int64_t sequence_length,
int64_t world_size) {
const int64_t stride = static_cast<int64_t>(gridDim.x) * blockDim.x;
for (int64_t i = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x; i < total; i += stride) {
int64_t rest = i;
const int64_t c = rest % head_dim;
for (int64_t index = static_cast<int64_t>(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 <typename T>
struct UspMergeHeadsKernel {
static void run(tvm::ffi::TensorView out, tvm::ffi::TensorView x) {
check_dtype<T>(out);
check_dtype<T>(x);
host::RuntimeCheck(x.ndim() == 5, "x must be [W, S, B, h_local, D]");
host::RuntimeCheck(out.ndim() == 5, "out must be [B, S, W, h_local, D]");
for (auto* t : {&x, &out}) {
host::RuntimeCheck(t->device().device_type == kDLCUDA, "usp_merge_heads tensors must be CUDA");
host::RuntimeCheck(is_dense_contiguous(*t), "usp_merge_heads tensors must be contiguous");
}
const int64_t world = x.size(0);
const int64_t seq = x.size(1);
const int64_t batch = x.size(2);
const int64_t h_local = x.size(3);
const int64_t head_dim = x.size(4);
host::RuntimeCheck(
out.size(0) == batch && out.size(1) == seq && out.size(2) == world && out.size(3) == h_local &&
out.size(4) == head_dim,
"out must be the [B, S, W, h_local, D] permutation of x");
static_assert(std::is_same_v<T, fp16_t> || std::is_same_v<T, bf16_t> || std::is_same_v<T, fp32_t>);
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<kDLCUDA>();
TensorMatcher({W, S, B, H, D}).with_dtype<T>().with_device(device).verify(x);
TensorMatcher({B, S, W, H, D}).with_dtype<T>().with_device(device).verify(out);
const int64_t world_size = W.unwrap();
const int64_t sequence_length = S.unwrap();
const int64_t batch = B.unwrap();
const int64_t local_heads = H.unwrap();
const int64_t head_dim = D.unwrap();
const int64_t numel = world_size * sequence_length * batch * local_heads * head_dim;
if (numel == 0) {
return;
}
T* out_ptr = reinterpret_cast<T*>(mutable_data_ptr(out));
const T* x_ptr = reinterpret_cast<const T*>(data_ptr(x));
auto* out_ptr = static_cast<T*>(out.data_ptr());
const auto* x_ptr = static_cast<const T*>(x.data_ptr());
CHECK_HOST(out_ptr != x_ptr) << "output must not alias the input";
const auto launch = [&](auto kernel, int64_t work_items, auto... args) {
const auto blocks =
static_cast<uint32_t>(std::min<int64_t>(div_ceil(work_items, static_cast<int64_t>(kBlockSize)), kMaxGrid));
LaunchKernel(blocks, kBlockSize, device.unwrap())(kernel, out_ptr, x_ptr, work_items, args...);
};
constexpr int kVec = 16 / sizeof(T);
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<uint32_t>(grid_for(n_vec)), kBlockSize, out.device())(
usp_merge_heads_vec_kernel<T, kVec>, 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<uintptr_t>(out_ptr) % kAlignment == 0 &&
reinterpret_cast<uintptr_t>(x_ptr) % kAlignment == 0;
if (vectorized) {
launch(
usp_merge_heads_vec_kernel<T, kVec>,
numel / kVec,
head_dim / kVec,
local_heads,
batch,
sequence_length,
world_size);
} else {
host::LaunchKernel(static_cast<uint32_t>(grid_for(total)), kBlockSize, out.device())(
usp_merge_heads_scalar_kernel<T>, out_ptr, x_ptr, total, head_dim, h_local, batch, seq, world);
launch(usp_merge_heads_scalar_kernel<T>, numel, head_dim, local_heads, batch, sequence_length, world_size);
}
}
};
python/sglang/kernels/ops/diffusion/__init__.py
+1 -13
View File
@@ -36,24 +36,12 @@ register_kernel(
KernelSpec(
op="diffusion.residual_gate_add",
backend=KernelBackend.JIT,
target="sglang.kernels.ops.diffusion.residual_gate_add:residual_gate_add_cuda",
target="sglang.kernels.ops.diffusion.residual_gate_add:residual_gate_add",
capabilities=_CUDA,
format_signature=FormatSignature(description="residual + gate * update"),
description="Fused residual gate-add (sglang.kernels.jit).",
)
)
register_kernel(
KernelSpec(
op="diffusion.fused_linear_gelu_tanh",
backend=KernelBackend.TORCH,
target="sglang.kernels.ops.diffusion.fused_linear_gelu:fused_linear_gelu_tanh",
capabilities=_CUDA,
format_signature=FormatSignature(
description="linear + tanh-GELU via the cublasLt GELU epilogue"
),
description="Fused up-proj GEMM + tanh-GELU (torch._addmm_activation).",
)
)
register_kernel(
KernelSpec(
op="diffusion.fused_inplace_qknorm_rope",
python/sglang/kernels/ops/diffusion/cutedsl/scale_residual_norm_scale_shift.py
+5 -5
View File
@@ -195,7 +195,7 @@ def validate_x(t: torch.Tensor, B: int, S: int, D: int):
if t.shape != (B, S, D):
raise ValueError(f"Validate failed: unsupported tensor shape: {t.shape}.")
if t.stride()[-1] != 1:
raise ValueError(f"Validate failed: not contiguous on dim D.")
raise ValueError("Validate failed: not contiguous on dim D.")
def validate_weight_bias(t: Optional[torch.Tensor], D: int):
@@ -206,7 +206,7 @@ def validate_weight_bias(t: Optional[torch.Tensor], D: int):
if t.shape != (D,):
raise ValueError(f"Validate failed: unsupported tensor shape: {t.shape}.")
if t.stride()[-1] != 1:
raise ValueError(f"Validate failed: not contiguous on dim D.")
raise ValueError("Validate failed: not contiguous on dim D.")
def validate_scale_shift(t: torch.Tensor, B: int, S: int, D: int):
@@ -230,7 +230,7 @@ def validate_scale_shift(t: torch.Tensor, B: int, S: int, D: int):
if failed:
raise ValueError(f"Validate failed: unsupported tensor shape: {t.shape}.")
if t.stride()[-1] != 1:
raise ValueError(f"Validate failed: not contiguous on dim D.")
raise ValueError("Validate failed: not contiguous on dim D.")
def validate_gate(t: Union[torch.Tensor, int], B: int, S: int, D: int):
@@ -311,7 +311,7 @@ def fused_norm_scale_shift(
compiled_fn(*torch_tensors, eps, stream)
return y
else:
raise ValueError(f'norm_type must be one of "layer" and "rms"')
raise ValueError('norm_type must be one of "layer" and "rms"')
@fused_norm_scale_shift.register_fake
@@ -401,7 +401,7 @@ def fused_scale_residual_norm_scale_shift(
compiled_fn(*torch_tensors, eps, stream)
return y, resi_out
else:
raise ValueError(f'norm_type must be one of "layer" and "rms"')
raise ValueError('norm_type must be one of "layer" and "rms"')
@fused_scale_residual_norm_scale_shift.register_fake
python/sglang/kernels/ops/diffusion/fused_gate_rmsnorm.py
+26 -39
View File
@@ -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)
python/sglang/kernels/ops/diffusion/fused_linear_gelu.py
+18 -28
View File
@@ -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)
python/sglang/kernels/ops/diffusion/fused_ln_modulate.py
+12 -16
View File
@@ -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(
python/sglang/kernels/ops/diffusion/modulate_scale_shift.py
+37 -2
View File
@@ -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",
]
python/sglang/kernels/ops/diffusion/qknorm_rope.py
+10
View File
@@ -19,6 +19,9 @@ if TYPE_CHECKING:
logger = logging.getLogger(__name__)
_SUPPORTED_DTYPES = (torch.float16, torch.bfloat16)
_SUPPORTED_CACHE_DTYPES = (*_SUPPORTED_DTYPES, torch.float32)
@cache_once
def _jit_qknorm_rope_module(
@@ -56,6 +59,13 @@ def can_use_fused_inplace_qknorm_rope(
cache_dtype: torch.dtype = torch.float32,
round_norm_before_rope: bool = False,
) -> bool:
if dtype not in _SUPPORTED_DTYPES or cache_dtype not in _SUPPORTED_CACHE_DTYPES:
logger.warning(
"Unsupported dtype pair (%s, %s) for JIT fused QKNorm+RoPE",
dtype,
cache_dtype,
)
return False
if head_dim not in (64, 128, 256):
logger.warning(f"Unsupported head_dim={head_dim} for JIT fused QKNorm+RoPE")
return False
python/sglang/kernels/ops/diffusion/quality_gate.py
+82
View File
@@ -0,0 +1,82 @@
"""Shared module-site protocol for request-scoped diffusion fast paths."""
from __future__ import annotations
import logging
from collections.abc import Callable, Iterator
from typing import Any
from torch import nn
RejectReason = Callable[[nn.Module], str | None]
class QualityGatedFusion:
"""Track and toggle one family of opt-in fusion sites.
The marker metadata describes the model attribute(s) owned by a site. The
enabled flag remains a plain module attribute so compiled model forwards
can read it without depending on this helper object.
"""
__slots__ = ("enabled_attr", "marker_attr", "name")
def __init__(self, *, name: str, marker_attr: str, enabled_attr: str) -> None:
self.name = name
self.marker_attr = marker_attr
self.enabled_attr = enabled_attr
def mark(self, module: nn.Module, metadata: Any = True) -> None:
setattr(module, self.marker_attr, metadata)
setattr(module, self.enabled_attr, False)
def metadata(self, module: nn.Module, default: Any = None) -> Any:
return getattr(module, self.marker_attr, default)
def is_enabled(self, module: nn.Module) -> bool:
return bool(getattr(module, self.enabled_attr, False))
def iter_sites(self, root: nn.Module) -> Iterator[nn.Module]:
for module in root.modules():
if hasattr(module, self.marker_attr):
yield module
def mount(
self,
root: nn.Module,
*,
reject_reason: RejectReason | None = None,
logger: logging.Logger | None = None,
) -> bool:
"""Enable every eligible site, or leave the whole family disabled."""
sites = list(self.iter_sites(root))
if not sites:
return False
if reject_reason is not None:
for site in sites:
reason = reject_reason(site)
if reason is None:
continue
self._set_enabled(sites, False)
if logger is not None:
logger.info(
"%s: %s site failed static guards (%s); keeping the "
"whole model on the reference path",
self.name,
type(site).__name__,
reason,
)
return False
self._set_enabled(sites, True)
return True
def unmount(self, root: nn.Module) -> None:
self._set_enabled(self.iter_sites(root), False)
def _set_enabled(
self, sites: Iterator[nn.Module] | list[nn.Module], enabled: bool
) -> None:
for site in sites:
setattr(site, self.enabled_attr, enabled)
python/sglang/kernels/ops/diffusion/residual_gate_add.py
+54 -3
View File
@@ -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",
]
python/sglang/kernels/ops/diffusion/triton/group_norm_silu_twopass.py
+32 -24
View File
@@ -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
python/sglang/kernels/ops/diffusion/triton/indexed_modulation.py
+4 -11
View File
@@ -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,
python/sglang/kernels/ops/diffusion/triton/layernorm_modulate.py
+10 -51
View File
@@ -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, :]
python/sglang/kernels/ops/diffusion/triton/mps_fallback.py
+6 -9
View File
@@ -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.
python/sglang/kernels/ops/diffusion/triton/native_bf16_rmsnorm.py
+208
View File
@@ -0,0 +1,208 @@
# SPDX-License-Identifier: Apache-2.0
"""BF16-native RMSNorm fusions shared by diffusion transformer models."""
from __future__ import annotations
import torch
import triton # type: ignore
import triton.language as tl # type: ignore
MAX_HIDDEN_SIZE = 8192
@triton.jit
def _tanh(x):
return 2.0 / (1.0 + tl.exp(-2.0 * x)) - 1.0
@triton.jit
def _rmsnorm_scale_kernel(
y_ptr,
x_ptr,
weight_ptr,
scale_ptr,
x_row_stride,
scale_row_stride,
seq_len,
dim: tl.constexpr,
eps: tl.constexpr,
block_dim: tl.constexpr,
):
row = tl.program_id(0)
offsets = tl.arange(0, block_dim)
mask = offsets < dim
x = tl.load(x_ptr + row * x_row_stride + offsets, mask=mask, other=0.0)
square = (x * x).to(tl.bfloat16)
mean_square = (tl.sum(square, axis=0) / dim).to(tl.bfloat16)
rstd = tl.rsqrt((mean_square + eps).to(tl.bfloat16).to(tl.float32)).to(tl.bfloat16)
batch = row // seq_len
weight = tl.load(weight_ptr + offsets, mask=mask, other=0.0)
scale = tl.load(
scale_ptr + batch * scale_row_stride + offsets, mask=mask, other=0.0
)
y = (((x * rstd).to(tl.bfloat16) * weight).to(tl.bfloat16) * scale).to(tl.bfloat16)
tl.store(y_ptr + row * dim + offsets, y, mask=mask)
@triton.jit
def _rmsnorm_tanh_residual_kernel(
y_ptr,
x_ptr,
gate_ptr,
residual_ptr,
weight_ptr,
x_row_stride,
gate_row_stride,
residual_row_stride,
seq_len,
dim: tl.constexpr,
eps: tl.constexpr,
block_dim: tl.constexpr,
):
row = tl.program_id(0)
offsets = tl.arange(0, block_dim)
mask = offsets < dim
x = tl.load(x_ptr + row * x_row_stride + offsets, mask=mask, other=0.0)
square = (x * x).to(tl.bfloat16)
mean_square = (tl.sum(square, axis=0) / dim).to(tl.bfloat16)
rstd = tl.rsqrt((mean_square + eps).to(tl.bfloat16).to(tl.float32)).to(tl.bfloat16)
batch = row // seq_len
gate = tl.load(gate_ptr + batch * gate_row_stride + offsets, mask=mask, other=0.0)
residual = tl.load(
residual_ptr + row * residual_row_stride + offsets, mask=mask, other=0.0
)
weight = tl.load(weight_ptr + offsets, mask=mask, other=0.0)
norm = ((x * rstd).to(tl.bfloat16) * weight).to(tl.bfloat16)
gated = (_tanh(gate.to(tl.float32)).to(tl.bfloat16) * norm).to(tl.bfloat16)
y = (residual + gated).to(tl.bfloat16)
tl.store(y_ptr + row * dim + offsets, y, mask=mask)
def _flat_row_stride(x: torch.Tensor) -> int | None:
if x.dim() < 2 or x.stride(-1) != 1:
return None
row_stride = x.stride(-2)
expected_stride = row_stride * x.shape[-2]
for dim in range(x.dim() - 3, -1, -1):
if x.stride(dim) != expected_stride:
return None
expected_stride *= x.shape[dim]
return row_stride
def _can_use_operand(
x: torch.Tensor, weight: torch.Tensor, other: torch.Tensor
) -> bool:
return (
x.is_cuda
and weight.is_cuda
and other.is_cuda
and x.device == weight.device == other.device
and x.dtype == weight.dtype == other.dtype == torch.bfloat16
and x.dim() >= 2
and 0 < x.shape[-1] <= MAX_HIDDEN_SIZE
and x.numel() > 0
and weight.shape == (x.shape[-1],)
and weight.is_contiguous()
and other.dim() >= 2
and other.shape[-1] == x.shape[-1]
and other.numel() > 0
and _flat_row_stride(x) is not None
and _flat_row_stride(other) is not None
)
def rmsnorm_scale(
x: torch.Tensor,
weight: torch.Tensor,
scale: torch.Tensor,
eps: float,
) -> torch.Tensor | None:
"""Apply BF16-native ``RMSNorm(x) * scale`` or return ``None``."""
if not _can_use_operand(x, weight, scale):
return None
dim = x.shape[-1]
x_rows = x.numel() // dim
scale_rows = scale.numel() // dim
if x_rows % scale_rows != 0:
return None
x_row_stride = _flat_row_stride(x)
scale_row_stride = _flat_row_stride(scale)
if x_row_stride is None or scale_row_stride is None:
return None
out = torch.empty_like(x, memory_format=torch.contiguous_format)
with torch.get_device_module().device(x.device):
_rmsnorm_scale_kernel[(x_rows,)](
out.reshape(-1, dim),
x,
weight,
scale,
x_row_stride,
scale_row_stride,
x_rows // scale_rows,
dim,
eps,
block_dim=triton.next_power_of_2(dim),
num_warps=8,
)
return out
def rmsnorm_tanh_residual(
x: torch.Tensor,
gate: torch.Tensor,
residual: torch.Tensor,
weight: torch.Tensor,
eps: float,
) -> torch.Tensor | None:
"""Apply BF16-native gated RMSNorm residual fusion or return ``None``."""
if not _can_use_operand(x, weight, gate):
return None
if (
residual.device != x.device
or residual.dtype != x.dtype
or residual.shape != x.shape
or _flat_row_stride(residual) is None
):
return None
dim = x.shape[-1]
x_rows = x.numel() // dim
gate_rows = gate.numel() // dim
if x_rows % gate_rows != 0:
return None
x_row_stride = _flat_row_stride(x)
gate_row_stride = _flat_row_stride(gate)
residual_row_stride = _flat_row_stride(residual)
if x_row_stride is None or gate_row_stride is None or residual_row_stride is None:
return None
out = torch.empty_like(x, memory_format=torch.contiguous_format)
with torch.get_device_module().device(x.device):
_rmsnorm_tanh_residual_kernel[(x_rows,)](
out.reshape(-1, dim),
x,
gate,
residual,
weight,
x_row_stride,
gate_row_stride,
residual_row_stride,
x_rows // gate_rows,
dim,
eps,
block_dim=triton.next_power_of_2(dim),
num_warps=8,
)
return out
__all__ = ["MAX_HIDDEN_SIZE", "rmsnorm_scale", "rmsnorm_tanh_residual"]
python/sglang/kernels/ops/diffusion/triton/numerics.py
+72
View File
@@ -0,0 +1,72 @@
# SPDX-License-Identifier: Apache-2.0
"""Numerical primitives shared by bit-exact diffusion Triton kernels."""
import triton # type: ignore
import triton.language as tl # type: ignore
_FLT_MIN = tl.constexpr(1.1754943508222875e-38)
@triton.jit
def round_bf16_to_fp32(value):
"""RNE-round fp32 to bf16 precision while keeping an fp32 register."""
bits = value.to(tl.int32, bitcast=True)
rounding_bias = 0x7FFF + ((bits >> 16) & 1)
rounded_bits = (bits + rounding_bias) & -65536
return rounded_bits.to(tl.float32, bitcast=True)
@triton.jit
def mul_rn_f32(x, y):
"""Opaque correctly-rounded fp32 multiply that blocks FMA contraction."""
return tl.inline_asm_elementwise(
asm="mul.rn.f32 $0, $1, $2;",
constraints="=f,f,f",
args=[x, y],
dtype=tl.float32,
is_pure=True,
pack=1,
)
@triton.jit
def div_rn_f32(x, y):
"""IEEE correctly-rounded fp32 division."""
return tl.inline_asm_elementwise(
asm="div.rn.f32 $0, $1, $2;",
constraints="=f,f,f",
args=[x, y],
dtype=tl.float32,
is_pure=True,
pack=1,
)
@triton.jit
def rsqrt_approx_f32(x):
return tl.inline_asm_elementwise(
asm="rsqrt.approx.f32 $0, $1;",
constraints="=f,f",
args=[x],
dtype=tl.float32,
is_pure=True,
pack=1,
)
@triton.jit
def cuda_rsqrtf(x):
"""Match CUDA ``rsqrtf``, including its subnormal rescaling path."""
is_subnormal = tl.abs(x) < _FLT_MIN
scaled = tl.where(is_subnormal, x * 16777216.0, x)
result = rsqrt_approx_f32(scaled)
return tl.where(is_subnormal, result * 4096.0, result)
__all__ = [
"cuda_rsqrtf",
"div_rn_f32",
"mul_rn_f32",
"round_bf16_to_fp32",
"rsqrt_approx_f32",
]
python/sglang/kernels/ops/diffusion/triton/rmsnorm_scale_shift_bitexact.py
+12 -43
View File
@@ -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
python/sglang/kernels/ops/diffusion/triton/sana_wm_gdn_chunkwise.py
+7 -2
View File
@@ -885,8 +885,13 @@ def phase_b_triton(
load_init = init_state_kv is not None
dummy = torch.empty(1, device=device, dtype=fdtype)
full_M = lambda: torch.empty(BH, F, BLOCK_D, BLOCK_D, device=device, dtype=fdtype)
full_z = lambda: torch.empty(BH, F, BLOCK_D, device=device, dtype=fdtype)
def full_M():
return torch.empty(BH, F, BLOCK_D, BLOCK_D, device=device, dtype=fdtype)
def full_z():
return torch.empty(BH, F, BLOCK_D, device=device, dtype=fdtype)
M_fwd = dummy if direction == 2 else full_M()
z_fwd = dummy if (direction == 2 or skip_z) else full_z()
# Combined-history reuses M_fwd/z_fwd as M_hist/z_hist; rev outputs are
python/sglang/kernels/ops/diffusion/triton/scale_shift.py
+3 -18
View File
@@ -2,26 +2,10 @@ import torch
import triton # type: ignore
import triton.language as tl # type: ignore
from sglang.kernels.ops.diffusion.triton.numerics import mul_rn_f32
from sglang.multimodal_gen.runtime.platforms import current_platform
@triton.jit
def _fp32_mul_add_rn(x, scale, residual):
"""Match separate CUDA FP32 multiply and add rounding (no FMA)."""
return tl.inline_asm_elementwise(
asm="""{
.reg .f32 product;
mul.rn.f32 product, $1, $2;
add.rn.f32 $0, $3, product;
}""",
constraints="=f,f,f,f",
args=(x, scale, residual),
dtype=tl.float32,
is_pure=True,
pack=1,
)
@triton.jit
def _fused_scaled_residual_add_exact_kernel(
output_ptr,
@@ -37,7 +21,8 @@ def _fused_scaled_residual_add_exact_kernel(
x = tl.load(x_ptr + offsets, mask=mask).to(tl.float32)
scale = tl.load(scale_ptr + offsets % width, mask=mask)
residual = tl.load(residual_ptr + offsets, mask=mask)
output = _fp32_mul_add_rn(x, scale, residual)
# The opaque multiply keeps Triton from contracting this into an FMA.
output = residual + mul_rn_f32(x, scale)
tl.store(output_ptr + offsets, output, mask=mask)
python/sglang/kernels/ops/diffusion/triton/ulysses_qkv.py
+44 -20
View File
@@ -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"]
python/sglang/kernels/ops/diffusion/triton/wan_rmsnorm_silu.py
+2 -1
View File
@@ -119,7 +119,7 @@ def _triton_wan_rmsnorm_silu_cuda(
block_c = triton.next_power_of_2(channels)
num_warps = 1 if block_c <= 64 else 4 if block_c <= 512 else 8
with torch.cuda.device(x.device):
with torch.get_device_module().device(x.device):
_wan_rmsnorm_silu_kernel[(bsz * t_size * h_size * w_size,)](
x,
gamma,
@@ -161,6 +161,7 @@ def can_use_wan_rmsnorm_silu(
and not x.requires_grad
and x.dtype in _SUPPORTED_DTYPES
and x.ndim == 5
and x.numel() > 0
and 0 < x.shape[1] <= _MAX_CHANNELS
and x.is_contiguous(memory_format=torch.channels_last_3d)
and _affine_supported(x, gamma)
python/sglang/kernels/ops/diffusion/triton/zimage_native_norm.py
+12 -183
View File
@@ -1,187 +1,13 @@
# SPDX-License-Identifier: Apache-2.0
"""Z-Image-specific bit-exact per-head RMSNorm kernel."""
from __future__ import annotations
import torch
import triton # type: ignore
import triton.language as tl # type: ignore
@triton.jit
def _tanh(x):
return 2.0 / (1.0 + tl.exp(-2.0 * x)) - 1.0
@triton.jit
def _rmsnorm_scale_kernel(
y_ptr,
x_ptr,
weight_ptr,
scale_ptr,
x_row_stride,
scale_row_stride,
seq_len,
dim: tl.constexpr,
eps: tl.constexpr,
block_dim: tl.constexpr,
):
row = tl.program_id(0)
offsets = tl.arange(0, block_dim)
mask = offsets < dim
x = tl.load(x_ptr + row * x_row_stride + offsets, mask=mask, other=0.0)
square = (x * x).to(tl.bfloat16)
mean_square = (tl.sum(square, axis=0) / dim).to(tl.bfloat16)
rstd = tl.rsqrt((mean_square + eps).to(tl.bfloat16).to(tl.float32)).to(tl.bfloat16)
batch = row // seq_len
weight = tl.load(weight_ptr + offsets, mask=mask, other=0.0)
scale = tl.load(
scale_ptr + batch * scale_row_stride + offsets, mask=mask, other=0.0
)
y = (((x * rstd).to(tl.bfloat16) * weight).to(tl.bfloat16) * scale).to(tl.bfloat16)
tl.store(y_ptr + row * dim + offsets, y, mask=mask)
@triton.jit
def _rmsnorm_tanh_residual_kernel(
y_ptr,
x_ptr,
gate_ptr,
residual_ptr,
weight_ptr,
x_row_stride,
gate_row_stride,
residual_row_stride,
seq_len,
dim: tl.constexpr,
eps: tl.constexpr,
block_dim: tl.constexpr,
):
row = tl.program_id(0)
offsets = tl.arange(0, block_dim)
mask = offsets < dim
x = tl.load(x_ptr + row * x_row_stride + offsets, mask=mask, other=0.0)
square = (x * x).to(tl.bfloat16)
mean_square = (tl.sum(square, axis=0) / dim).to(tl.bfloat16)
rstd = tl.rsqrt((mean_square + eps).to(tl.bfloat16).to(tl.float32)).to(tl.bfloat16)
batch = row // seq_len
gate = tl.load(gate_ptr + batch * gate_row_stride + offsets, mask=mask, other=0.0)
residual = tl.load(
residual_ptr + row * residual_row_stride + offsets, mask=mask, other=0.0
)
weight = tl.load(weight_ptr + offsets, mask=mask, other=0.0)
norm = ((x * rstd).to(tl.bfloat16) * weight).to(tl.bfloat16)
gated = (_tanh(gate.to(tl.float32)).to(tl.bfloat16) * norm).to(tl.bfloat16)
y = (residual + gated).to(tl.bfloat16)
tl.store(y_ptr + row * dim + offsets, y, mask=mask)
def _flat_row_stride(x: torch.Tensor) -> int | None:
if x.dim() < 2 or x.stride(-1) != 1:
return None
row_stride = x.stride(-2)
expected_stride = row_stride * x.shape[-2]
for dim in range(x.dim() - 3, -1, -1):
if x.stride(dim) != expected_stride:
return None
expected_stride *= x.shape[dim]
return row_stride
def _can_use(x: torch.Tensor, weight: torch.Tensor, other: torch.Tensor) -> bool:
return (
x.is_cuda
and weight.is_cuda
and other.is_cuda
and x.dtype == torch.bfloat16
and weight.dtype == torch.bfloat16
and other.dtype == torch.bfloat16
and weight.is_contiguous()
and x.shape[-1] <= 8192
and _flat_row_stride(x) is not None
and _flat_row_stride(other) is not None
)
def zimage_rmsnorm_scale(
x: torch.Tensor,
weight: torch.Tensor,
scale: torch.Tensor,
eps: float,
) -> torch.Tensor | None:
if not _can_use(x, weight, scale):
return None
shape = x.shape
dim = shape[-1]
x_rows = x.numel() // dim
scale_rows = scale.numel() // dim
if x_rows % scale_rows != 0:
return None
seq_len = x_rows // scale_rows
x_row_stride = _flat_row_stride(x)
scale_row_stride = _flat_row_stride(scale)
if x_row_stride is None or scale_row_stride is None:
return None
y = torch.empty_like(x, memory_format=torch.contiguous_format)
with torch.get_device_module().device(x.device):
_rmsnorm_scale_kernel[(x_rows,)](
y.reshape(-1, dim),
x,
weight,
scale,
x_row_stride,
scale_row_stride,
seq_len,
dim,
eps,
block_dim=triton.next_power_of_2(dim),
num_warps=8,
)
return y
def zimage_rmsnorm_tanh_residual(
x: torch.Tensor,
gate: torch.Tensor,
residual: torch.Tensor,
weight: torch.Tensor,
eps: float,
) -> torch.Tensor | None:
if not (_can_use(x, weight, gate) and residual.is_cuda):
return None
if residual.dtype != x.dtype or _flat_row_stride(residual) is None:
return None
shape = x.shape
dim = shape[-1]
x_rows = x.numel() // dim
gate_rows = gate.numel() // dim
if x_rows % gate_rows != 0:
return None
seq_len = x_rows // gate_rows
x_row_stride = _flat_row_stride(x)
gate_row_stride = _flat_row_stride(gate)
residual_row_stride = _flat_row_stride(residual)
if x_row_stride is None or gate_row_stride is None or residual_row_stride is None:
return None
y = torch.empty_like(x, memory_format=torch.contiguous_format)
with torch.get_device_module().device(x.device):
_rmsnorm_tanh_residual_kernel[(x_rows,)](
y.reshape(-1, dim),
x,
gate,
residual,
weight,
x_row_stride,
gate_row_stride,
residual_row_stride,
seq_len,
dim,
eps,
block_dim=triton.next_power_of_2(dim),
num_warps=8,
)
return y
@triton.jit
def _qk_rmsnorm_native_kernel(
y_ptr,
@@ -294,9 +120,11 @@ def can_use_qk_rmsnorm_native(
return (
x.is_cuda
and weight.is_cuda
and x.dtype == torch.bfloat16
and x.device == weight.device
and x.dtype == weight.dtype == torch.bfloat16
and x.numel() > 0
and head_dim == 128
and weight.numel() == head_dim
and weight.shape == (head_dim,)
and weight.is_contiguous()
and _qk_head_token_stride(x, head_dim) is not None
)
@@ -319,8 +147,6 @@ def zimage_qk_rmsnorm_native(
token_stride = _qk_head_token_stride(x, head_dim)
if token_stride is None:
return None
if weight.dtype != x.dtype:
weight = weight.to(dtype=x.dtype)
nheads = x.shape[2]
n_rows = x.shape[0] * x.shape[1] * nheads
rows_per_prog = 8
@@ -340,3 +166,6 @@ def zimage_qk_rmsnorm_native(
num_warps=8,
)
return y
__all__ = ["can_use_qk_rmsnorm_native", "zimage_qk_rmsnorm_native"]
python/sglang/kernels/ops/diffusion/usp_relayout.py
+2
View File
@@ -16,6 +16,8 @@ _SUPPORTED_DTYPES = (torch.float16, torch.bfloat16, torch.float32)
@cache_once
def _jit_usp_relayout_module(dtype: torch.dtype) -> Module:
if dtype not in _SUPPORTED_DTYPES:
raise RuntimeError(f"Unsupported usp_merge_heads dtype: {dtype}")
args = make_cpp_args(dtype)
return load_jit(
"diffusion_usp_relayout",
python/sglang/multimodal_gen/runtime/models/dits/ernie_image.py
+3 -37
View File
@@ -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
python/sglang/multimodal_gen/runtime/models/dits/flux.py
+12 -83
View File
@@ -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:
python/sglang/multimodal_gen/runtime/models/dits/flux_2.py
+20 -51
View File
@@ -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:
python/sglang/multimodal_gen/runtime/models/dits/glm_image.py
+5 -38
View File
@@ -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)
)
python/sglang/multimodal_gen/runtime/models/dits/ideogram.py
+2 -1
View File
@@ -8,6 +8,7 @@ import torch.nn as nn
import torch.nn.functional as F
from sglang.kernels.ops.diffusion.fused_gate_rmsnorm import (
fused_gate_rmsnorm_active,
fused_rmsnorm_scale,
fused_rmsnorm_tanh_residual,
mark_fused_gate_rmsnorm_site,
@@ -394,7 +395,7 @@ class Ideogram4TransformerBlock(nn.Module):
adaln_input
).chunk(4, dim=-1)
enable_fused = (
self._sgl_fused_gate_rmsnorm_enabled and not torch.compiler.is_compiling()
fused_gate_rmsnorm_active(self) and not torch.compiler.is_compiling()
)
attn_out = self.attention(
_norm_scale(x, scale_msa, self.attention_norm1, enable_fused),
python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py
+9 -36
View File
@@ -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
python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py
+2 -3
View File
@@ -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
)
python/sglang/multimodal_gen/runtime/models/dits/zimage.py
+6 -6
View File
@@ -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,
python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py
+37 -34
View File
@@ -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"