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

[Diffusion] Reuse shared AlignedVector and tidy jit_kernel/diffusion (#29664)

Browse Source 
Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
Xiaoyu Zhang
2026-06-30 11:38:22 +08:00
committed by GitHub
co-authored by Claude Opus 4.8
parent 8dbf04fc56
commit 3add35e26d
10 changed files with 116 additions and 159 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/jit_kernel/csrc/diffusion/causal_conv3d_cat_pad.cuh
+4 -6
View File
@@ -15,6 +15,7 @@
#include <sgl_kernel/utils.h> // For RuntimeCheck, div_ceil
#include <sgl_kernel/utils.cuh> // For LaunchKernel
#include <sgl_kernel/vec.cuh> // For device::AlignedVector
#include <cstdint>
@@ -41,10 +42,7 @@ __global__ void __launch_bounds__(kBlockSize) cat_pad_flat_kernel(
int64_t pad_d_left,
int64_t pad_h_top,
int64_t pad_w_left) {
union Pack {
ET elem[kVec];
uint4 raw;
};
using Pack = device::AlignedVector<ET, kVec>;
const int64_t nthreads = static_cast<int64_t>(gridDim.x) * blockDim.x;
for (int64_t vid = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x; vid < total_vecs; vid += nthreads) {
@@ -81,7 +79,7 @@ __global__ void __launch_bounds__(kBlockSize) cat_pad_flat_kernel(
value = SGLANG_LDG(src + iw);
}
}
pack.elem[i] = value;
pack[i] = value;
if (++ow == out_w) {
ow = 0;
@@ -110,7 +108,7 @@ __global__ void __launch_bounds__(kBlockSize) cat_pad_flat_kernel(
}
}
reinterpret_cast<uint4*>(out)[vid] = pack.raw;
pack.store(out, vid);
}
}
python/sglang/jit_kernel/csrc/diffusion/residual_gate_add.cuh
+16 -21
View File
@@ -17,6 +17,7 @@
#include <sgl_kernel/type.cuh> // For dtype_trait conversions
#include <sgl_kernel/utils.cuh> // For LaunchKernel and CUDA dtype aliases
#include <sgl_kernel/vec.cuh> // For device::AlignedVector
#include <cstdint>
@@ -102,13 +103,6 @@ __device__ __forceinline__ T residual_gate_value(T residual, T update, T gate) {
return dtype_trait<T>::from(to_float(residual) + to_float(product));
}
template <typename T>
union Vec16 {
static constexpr int kElems = 16 / sizeof(T);
uint4 raw;
T elems[kElems];
};
template <typename T, int kVec>
__global__ void residual_gate_add_vec_kernel(
const T* __restrict__ residual,
@@ -116,18 +110,18 @@ __global__ void residual_gate_add_vec_kernel(
const T* __restrict__ gate,
T* __restrict__ out,
int64_t n_vec) {
using Vec = device::AlignedVector<T, kVec>;
const int64_t stride = static_cast<int64_t>(gridDim.x) * blockDim.x;
for (int64_t v = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x; v < n_vec; v += stride) {
const Vec16<T> r{.raw = reinterpret_cast<const uint4*>(residual)[v]};
const Vec16<T> u{.raw = reinterpret_cast<const uint4*>(update)[v]};
const Vec16<T> g{.raw = reinterpret_cast<const uint4*>(gate)[v]};
Vec16<T> o;
Vec r, u, g, o;
r.load(residual, v);
u.load(update, v);
g.load(gate, v);
#pragma unroll
for (int i = 0; i < kVec; ++i) {
o.elems[i] = residual_gate_value(r.elems[i], u.elems[i], g.elems[i]);
o[i] = residual_gate_value(r[i], u[i], g[i]);
}
reinterpret_cast<uint4*>(out)[v] = o.raw;
o.store(out, v);
}
}
@@ -139,12 +133,14 @@ __global__ void residual_gate_add_bcast_row_tile_kernel(
T* __restrict__ out,
int64_t rows,
int64_t row_vec) {
using Vec = device::AlignedVector<T, kVec>;
const int64_t col_vec = static_cast<int64_t>(blockIdx.x) * kBcastColsVecPerBlock + threadIdx.x;
if (col_vec >= row_vec) {
return;
}
const Vec16<T> g{.raw = SGLANG_LDG(reinterpret_cast<const uint4*>(gate) + col_vec)};
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.
@@ -156,15 +152,14 @@ __global__ void residual_gate_add_bcast_row_tile_kernel(
const int64_t row = row_base + row_offset;
if (row < rows) {
const int64_t v = row * row_vec + col_vec;
const Vec16<T> r{.raw = reinterpret_cast<const uint4*>(residual)[v]};
const Vec16<T> u{.raw = reinterpret_cast<const uint4*>(update)[v]};
Vec16<T> o;
Vec r, u, o;
r.load(residual, v);
u.load(update, v);
#pragma unroll
for (int i = 0; i < kVec; ++i) {
o.elems[i] = residual_gate_value(r.elems[i], u.elems[i], g.elems[i]);
o[i] = residual_gate_value(r[i], u[i], g[i]);
}
reinterpret_cast<uint4*>(out)[v] = o.raw;
o.store(out, v);
}
}
}
python/sglang/jit_kernel/csrc/diffusion/timestep_embedding.cuh
+31 -27
View File
@@ -1,9 +1,12 @@
#pragma once
#include <sgl_kernel/tensor.h>
#include <sgl_kernel/utils.h>
#include <sgl_kernel/math.cuh>
#include <sgl_kernel/type.cuh>
#include <sgl_kernel/utils.cuh>
#include <sgl_kernel/vec.cuh> // For device::AlignedVector
#include <dlpack/dlpack.h>
#include <tvm/ffi/container/tensor.h>
@@ -14,8 +17,12 @@
#include <cuda_runtime.h>
#include <type_traits>
namespace sglang_timestep_embedding {
namespace {
constexpr int kVec = 4; // 16B float vector store
template <bool kFlipSinToCos, typename TIn>
__global__ void timestep_embedding_kernel(
const TIn* __restrict__ t_ptr,
@@ -24,6 +31,8 @@ __global__ void timestep_embedding_kernel(
float neg_log_max_period,
float scale,
int batch_size) {
using Vec = device::AlignedVector<float, kVec>;
int row_idx = static_cast<int>(blockIdx.x * blockDim.y + threadIdx.y);
if (row_idx >= batch_size) {
return;
@@ -34,36 +43,29 @@ __global__ void timestep_embedding_kernel(
int half_dim = dim / 2;
int thread_offset = static_cast<int>(threadIdx.x);
while (thread_offset * 4 < half_dim) {
float4* top_half;
float4* bottom_half;
while (thread_offset * kVec < half_dim) {
// !flip: output is [sin | cos]; flip: output is [cos | sin].
float* cos_dst;
float* sin_dst;
if constexpr (!kFlipSinToCos) {
bottom_half = reinterpret_cast<float4*>(output_batch_base_ptr + thread_offset * 4);
top_half = reinterpret_cast<float4*>(output_batch_base_ptr + half_dim + thread_offset * 4);
sin_dst = output_batch_base_ptr + thread_offset * kVec;
cos_dst = output_batch_base_ptr + half_dim + thread_offset * kVec;
} else {
top_half = reinterpret_cast<float4*>(output_batch_base_ptr + thread_offset * 4);
bottom_half = reinterpret_cast<float4*>(output_batch_base_ptr + half_dim + thread_offset * 4);
cos_dst = output_batch_base_ptr + thread_offset * kVec;
sin_dst = output_batch_base_ptr + half_dim + thread_offset * kVec;
}
float4 vals;
vals.x = scale * t_val * device::math::exp(neg_log_max_period * __int2float_rn(thread_offset * 4 + 0));
vals.y = scale * t_val * device::math::exp(neg_log_max_period * __int2float_rn(thread_offset * 4 + 1));
vals.z = scale * t_val * device::math::exp(neg_log_max_period * __int2float_rn(thread_offset * 4 + 2));
vals.w = scale * t_val * device::math::exp(neg_log_max_period * __int2float_rn(thread_offset * 4 + 3));
float4 cos_vals;
cos_vals.x = device::math::cos(vals.x);
cos_vals.y = device::math::cos(vals.y);
cos_vals.z = device::math::cos(vals.z);
cos_vals.w = device::math::cos(vals.w);
*top_half = cos_vals;
float4 sin_vals;
sin_vals.x = device::math::sin(vals.x);
sin_vals.y = device::math::sin(vals.y);
sin_vals.z = device::math::sin(vals.z);
sin_vals.w = device::math::sin(vals.w);
*bottom_half = sin_vals;
Vec cos_vec;
Vec sin_vec;
#pragma unroll
for (int i = 0; i < kVec; ++i) {
const float angle =
scale * t_val * device::math::exp(neg_log_max_period * __int2float_rn(thread_offset * kVec + i));
cos_vec[i] = device::math::cos(angle);
sin_vec[i] = device::math::sin(angle);
}
cos_vec.store(cos_dst);
sin_vec.store(sin_dst);
thread_offset += static_cast<int>(blockDim.x);
}
@@ -118,6 +120,8 @@ inline void launch_timestep_embedding(
}
}
} // namespace
template <typename TIn>
void timestep_embedding(
tvm::ffi::TensorView input,
@@ -147,4 +151,4 @@ void timestep_embedding(
launch_timestep_embedding<TIn>(input, output, dim, flip_sin_to_cos, downscale_freq_shift, scale, max_period);
}
} // namespace
} // namespace sglang_timestep_embedding
python/sglang/jit_kernel/diffusion/cutedsl/norm_tanh_mul_add_norm_scale.py
+6 -41
View File
@@ -10,50 +10,15 @@ from sglang.jit_kernel.diffusion.cutedsl.common.norm_fusion import (
broadcast_tensor_for_bsfd,
tensor_slice_for_bsfd,
)
from sglang.jit_kernel.diffusion.cutedsl.utils import TORCH_TO_CUTE_DTYPE, WARP_SIZE
from sglang.jit_kernel.diffusion.cutedsl.utils import (
WARP_SIZE,
to_cute_arg,
to_fake_cute_args,
)
_COMPILE_CACHE = {}
def to_cute_arg(
t,
*,
assume_aligned: Optional[int] = 32,
use_32bit_stride: bool = False,
enable_tvm_ffi: bool = True,
):
"""
Convert a Python value into a CuTeDSL value.
"""
if isinstance(t, torch.Tensor):
return cute.runtime.from_dlpack(
t,
assumed_align=assume_aligned,
use_32bit_stride=use_32bit_stride,
enable_tvm_ffi=enable_tvm_ffi,
)
if isinstance(t, int):
return cutlass.Int32(t)
if isinstance(t, float):
return cutlass.Float32(t)
return t
def to_fake_cute_args(t: torch.Tensor):
if isinstance(t, torch.Tensor):
# Only keep the last dim as compile-time value to maximum compiled kernel reuse
# e.g. (1,2,1536):(3027,1536,1) -> (?,?,1536):(?,?,1)
D = t.shape[-1]
dtype = TORCH_TO_CUTE_DTYPE[t.dtype]
shape = (*(cute.sym_int() for _ in range(t.ndim - 1)), D)
stride = (*(cute.sym_int(divisibility=D) for _ in range(t.ndim - 1)), 1)
fake_t = cute.runtime.make_fake_tensor(
dtype, shape, stride, memspace=cute.AddressSpace.gmem, assumed_align=32
)
return fake_t
return to_cute_arg(t)
class NormTanhMulAddNormScale:
@classmethod
def make_hash_key(cls, *inputs):
@@ -166,7 +131,7 @@ class NormTanhMulAddNormScale:
@cute.jit
def copy_if(src, dst):
if cutlass.const_expr(
isinstance(src, cute.Tensor) and isinstance(src, cute.Tensor)
isinstance(src, cute.Tensor) and isinstance(dst, cute.Tensor)
):
cute.autovec_copy(src, dst) # LDG.128
python/sglang/jit_kernel/diffusion/cutedsl/scale_residual_norm_scale_shift.py
+9 -46
View File
@@ -10,50 +10,14 @@ from sglang.jit_kernel.diffusion.cutedsl.common.norm_fusion import (
broadcast_tensor_for_bsfd,
tensor_slice_for_bsfd,
)
from sglang.jit_kernel.diffusion.cutedsl.utils import TORCH_TO_CUTE_DTYPE, WARP_SIZE
from sglang.jit_kernel.diffusion.cutedsl.utils import (
WARP_SIZE,
to_fake_cute_args,
)
_COMPILE_CACHE = {}
def to_cute_arg(
t,
*,
assume_aligned: Optional[int] = 32,
use_32bit_stride: bool = False,
enable_tvm_ffi: bool = True,
):
"""
Convert a Python value into a CuTeDSL value.
"""
if isinstance(t, torch.Tensor):
return cute.runtime.from_dlpack(
t,
assumed_align=assume_aligned,
use_32bit_stride=use_32bit_stride,
enable_tvm_ffi=enable_tvm_ffi,
)
if isinstance(t, int):
return cutlass.Int32(t)
if isinstance(t, float):
return cutlass.Float32(t)
return t
def to_fake_cute_args(t: torch.Tensor):
if isinstance(t, torch.Tensor):
# Only keep the last dim as compile-time value to maximum compiled kernel reuse
# e.g. (1,2,1536):(3027,1536,1) -> (?,?,1536):(?,?,1)
D = t.shape[-1]
dtype = TORCH_TO_CUTE_DTYPE[t.dtype]
shape = (*(cute.sym_int() for _ in range(t.ndim - 1)), D)
stride = (*(cute.sym_int(divisibility=D) for _ in range(t.ndim - 1)), 1)
fake_t = cute.runtime.make_fake_tensor(
dtype, shape, stride, memspace=cute.AddressSpace.gmem, assumed_align=32
)
return fake_t
return to_cute_arg(t)
class ScaleResidualNormScaleShift:
@classmethod
def make_hash_key(cls, *inputs):
@@ -234,7 +198,7 @@ def validate_x(t: torch.Tensor, B: int, S: int, D: int):
raise ValueError(f"Validate failed: not contiguous on dim D.")
def validate_weight_bias(t: Optional[torch.Tensor], B: int, S: int, D: int):
def validate_weight_bias(t: Optional[torch.Tensor], D: int):
if t is None:
return
if t.dtype not in (torch.float16, torch.bfloat16, torch.float32):
@@ -312,8 +276,8 @@ def fused_norm_scale_shift(
# Tensor Validation
BSD = x.shape
validate_x(x, *BSD)
validate_weight_bias(weight, *BSD)
validate_weight_bias(bias, *BSD)
validate_weight_bias(weight, BSD[-1])
validate_weight_bias(bias, BSD[-1])
validate_scale_shift(scale, *BSD)
validate_scale_shift(shift, *BSD)
@@ -399,14 +363,13 @@ def fused_scale_residual_norm_scale_shift(
validate_x(x, *BSD)
validate_x(residual, *BSD)
validate_gate(gate, *BSD)
validate_weight_bias(weight, *BSD)
validate_weight_bias(bias, *BSD)
validate_weight_bias(weight, BSD[-1])
validate_weight_bias(bias, BSD[-1])
validate_scale_shift(scale, *BSD)
validate_scale_shift(shift, *BSD)
if norm_type == "layer" or norm_type == "rms":
stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream)
# if norm_type == "layer" or norm_type == "rms":
D = x.shape[-1]
if D % 256 != 0 or D > 8192:
raise ValueError(
python/sglang/jit_kernel/diffusion/cutedsl/utils.py
+42
View File
@@ -1,4 +1,7 @@
from typing import Optional
import cutlass
import cutlass.cute as cute
import torch
WARP_SIZE = 32
@@ -8,3 +11,42 @@ TORCH_TO_CUTE_DTYPE = {
torch.bfloat16: cutlass.BFloat16,
torch.float32: cutlass.Float32,
}
def to_cute_arg(
t,
*,
assume_aligned: Optional[int] = 32,
use_32bit_stride: bool = False,
enable_tvm_ffi: bool = True,
):
"""
Convert a Python value into a CuTeDSL value.
"""
if isinstance(t, torch.Tensor):
return cute.runtime.from_dlpack(
t,
assumed_align=assume_aligned,
use_32bit_stride=use_32bit_stride,
enable_tvm_ffi=enable_tvm_ffi,
)
if isinstance(t, int):
return cutlass.Int32(t)
if isinstance(t, float):
return cutlass.Float32(t)
return t
def to_fake_cute_args(t: torch.Tensor):
if isinstance(t, torch.Tensor):
# Only keep the last dim as compile-time value to maximum compiled kernel reuse
# e.g. (1,2,1536):(3027,1536,1) -> (?,?,1536):(?,?,1)
D = t.shape[-1]
dtype = TORCH_TO_CUTE_DTYPE[t.dtype]
shape = (*(cute.sym_int() for _ in range(t.ndim - 1)), D)
stride = (*(cute.sym_int(divisibility=D) for _ in range(t.ndim - 1)), 1)
fake_t = cute.runtime.make_fake_tensor(
dtype, shape, stride, memspace=cute.AddressSpace.gmem, assumed_align=32
)
return fake_t
return to_cute_arg(t)
python/sglang/jit_kernel/diffusion/triton/norm.py
-1
View File
@@ -38,7 +38,6 @@ def triton_autotune_configs():
for warp_count in [1, 2, 4, 8, 16, 32]
if warp_count * warp_size <= max_threads_per_block
]
# return [triton.Config({}, num_warps=8)]
# Copied from flash-attn
python/sglang/jit_kernel/diffusion/triton/sana_wm_gdn.py
+2 -14
View File
@@ -120,18 +120,6 @@ def prepare_rope_tables(
return rope_cos.contiguous(), rope_sin.contiguous()
def _precompute_inv_rms(
qkv: torch.Tensor, idx: int, C: int, eps: float = 1e-5
) -> torch.Tensor:
"""Compute 1/RMS for one component of QKV over the full C = H*D channel dim.
qkv: (B, N, 3, H, D); idx: 0=Q, 1=K, 2=V; C: H*D. Returns (B, N) float32.
"""
raw = qkv[:, :, idx].float() # (B, N, H, D)
sq_sum = (raw * raw).sum(dim=(-2, -1)) # (B, N)
return torch.rsqrt(sq_sum / C + eps)
# =====================================================================
# Fused single-pass Q+K inverse-RMS Triton kernel
# =====================================================================
@@ -182,7 +170,7 @@ def fused_qk_inv_rms(
) -> tuple[torch.Tensor, torch.Tensor]:
"""Single-pass Triton fused Q+K inverse-RMS.
Replaces two ``_precompute_inv_rms`` scans with one launch that reads each
Replaces two separate PyTorch RMS scans with one launch that reads each
``(b, n)`` row of ``qkv`` exactly once.
qkv: (B, N, 3, H, D) contiguous. Returns (q_inv_rms, k_inv_rms), each (B, N) float32.
"""
@@ -230,7 +218,7 @@ def fused_bigdn_func(
"""Bidirectional fused GDN. Returns ``(B, N, H, D)``.
Thin entry point kept for call-site stability; delegates to
:func:`fused_bigdn_bidi_chunkwise` from ``fused_gdn_chunkwise``.
:func:`fused_bigdn_bidi_chunkwise` from ``sana_wm_gdn_chunkwise``.
"""
from sglang.jit_kernel.diffusion.triton.sana_wm_gdn_chunkwise import (
fused_bigdn_bidi_chunkwise,
python/sglang/jit_kernel/diffusion/triton/scale_shift.py
-2
View File
@@ -224,7 +224,6 @@ def _fused_scale_shift_4d_kernel(
scale_ptr,
shift_ptr,
scale_constant: tl.constexpr, # scale_constant is either 0 or 1.
rows,
inner_dim,
seq_len,
num_frames,
@@ -386,7 +385,6 @@ def fuse_scale_shift_kernel(
scale_reshaped,
shift_reshaped,
scale_constant,
rows,
C,
L,
num_frames,
python/sglang/jit_kernel/timestep_embedding.py
+6 -1
View File
@@ -18,7 +18,12 @@ def _jit_timestep_embedding_module(dtype: torch.dtype) -> Module:
"timestep_embedding",
*args,
cuda_files=["diffusion/timestep_embedding.cuh"],
cuda_wrappers=[("timestep_embedding", f"timestep_embedding<{args}>")],
cuda_wrappers=[
(
"timestep_embedding",
f"sglang_timestep_embedding::timestep_embedding<{args}>",
)
],
)