From 3add35e26dc0623d6647e226de7d17754bb61804 Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang <1182563586@qq.com> Date: Tue, 30 Jun 2026 11:38:22 +0800 Subject: [PATCH] [Diffusion] Reuse shared AlignedVector and tidy jit_kernel/diffusion (#29664) Co-authored-by: Claude Opus 4.8 --- .../csrc/diffusion/causal_conv3d_cat_pad.cuh | 10 ++-- .../csrc/diffusion/residual_gate_add.cuh | 37 +++++------- .../csrc/diffusion/timestep_embedding.cuh | 58 ++++++++++--------- .../cutedsl/norm_tanh_mul_add_norm_scale.py | 47 ++------------- .../scale_residual_norm_scale_shift.py | 55 +++--------------- .../jit_kernel/diffusion/cutedsl/utils.py | 42 ++++++++++++++ .../jit_kernel/diffusion/triton/norm.py | 1 - .../diffusion/triton/sana_wm_gdn.py | 16 +---- .../diffusion/triton/scale_shift.py | 2 - .../sglang/jit_kernel/timestep_embedding.py | 7 ++- 10 files changed, 116 insertions(+), 159 deletions(-) diff --git a/python/sglang/jit_kernel/csrc/diffusion/causal_conv3d_cat_pad.cuh b/python/sglang/jit_kernel/csrc/diffusion/causal_conv3d_cat_pad.cuh index e5b447318..a0adf7cd3 100644 --- a/python/sglang/jit_kernel/csrc/diffusion/causal_conv3d_cat_pad.cuh +++ b/python/sglang/jit_kernel/csrc/diffusion/causal_conv3d_cat_pad.cuh @@ -15,6 +15,7 @@ #include // For RuntimeCheck, div_ceil #include // For LaunchKernel +#include // For device::AlignedVector #include @@ -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; const int64_t nthreads = static_cast(gridDim.x) * blockDim.x; for (int64_t vid = static_cast(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(out)[vid] = pack.raw; + pack.store(out, vid); } } diff --git a/python/sglang/jit_kernel/csrc/diffusion/residual_gate_add.cuh b/python/sglang/jit_kernel/csrc/diffusion/residual_gate_add.cuh index 693641880..25a68e99e 100644 --- a/python/sglang/jit_kernel/csrc/diffusion/residual_gate_add.cuh +++ b/python/sglang/jit_kernel/csrc/diffusion/residual_gate_add.cuh @@ -17,6 +17,7 @@ #include // For dtype_trait conversions #include // For LaunchKernel and CUDA dtype aliases +#include // For device::AlignedVector #include @@ -102,13 +103,6 @@ __device__ __forceinline__ T residual_gate_value(T residual, T update, T gate) { return dtype_trait::from(to_float(residual) + to_float(product)); } -template -union Vec16 { - static constexpr int kElems = 16 / sizeof(T); - uint4 raw; - T elems[kElems]; -}; - template __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; const int64_t stride = static_cast(gridDim.x) * blockDim.x; for (int64_t v = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; v < n_vec; v += stride) { - const Vec16 r{.raw = reinterpret_cast(residual)[v]}; - const Vec16 u{.raw = reinterpret_cast(update)[v]}; - const Vec16 g{.raw = reinterpret_cast(gate)[v]}; - - Vec16 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(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; const int64_t col_vec = static_cast(blockIdx.x) * kBcastColsVecPerBlock + threadIdx.x; if (col_vec >= row_vec) { return; } - const Vec16 g{.raw = SGLANG_LDG(reinterpret_cast(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 r{.raw = reinterpret_cast(residual)[v]}; - const Vec16 u{.raw = reinterpret_cast(update)[v]}; - - Vec16 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(out)[v] = o.raw; + o.store(out, v); } } } diff --git a/python/sglang/jit_kernel/csrc/diffusion/timestep_embedding.cuh b/python/sglang/jit_kernel/csrc/diffusion/timestep_embedding.cuh index 2d29da50b..28b59798d 100644 --- a/python/sglang/jit_kernel/csrc/diffusion/timestep_embedding.cuh +++ b/python/sglang/jit_kernel/csrc/diffusion/timestep_embedding.cuh @@ -1,9 +1,12 @@ +#pragma once + #include #include #include #include #include +#include // For device::AlignedVector #include #include @@ -14,8 +17,12 @@ #include #include +namespace sglang_timestep_embedding { + namespace { +constexpr int kVec = 4; // 16B float vector store + template __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; + int row_idx = static_cast(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(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(output_batch_base_ptr + thread_offset * 4); - top_half = reinterpret_cast(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(output_batch_base_ptr + thread_offset * 4); - bottom_half = reinterpret_cast(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(blockDim.x); } @@ -118,6 +120,8 @@ inline void launch_timestep_embedding( } } +} // namespace + template void timestep_embedding( tvm::ffi::TensorView input, @@ -147,4 +151,4 @@ void timestep_embedding( launch_timestep_embedding(input, output, dim, flip_sin_to_cos, downscale_freq_shift, scale, max_period); } -} // namespace +} // namespace sglang_timestep_embedding diff --git a/python/sglang/jit_kernel/diffusion/cutedsl/norm_tanh_mul_add_norm_scale.py b/python/sglang/jit_kernel/diffusion/cutedsl/norm_tanh_mul_add_norm_scale.py index c5ae1c619..87baffe75 100644 --- a/python/sglang/jit_kernel/diffusion/cutedsl/norm_tanh_mul_add_norm_scale.py +++ b/python/sglang/jit_kernel/diffusion/cutedsl/norm_tanh_mul_add_norm_scale.py @@ -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 diff --git a/python/sglang/jit_kernel/diffusion/cutedsl/scale_residual_norm_scale_shift.py b/python/sglang/jit_kernel/diffusion/cutedsl/scale_residual_norm_scale_shift.py index c835fea63..f560a426f 100644 --- a/python/sglang/jit_kernel/diffusion/cutedsl/scale_residual_norm_scale_shift.py +++ b/python/sglang/jit_kernel/diffusion/cutedsl/scale_residual_norm_scale_shift.py @@ -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( diff --git a/python/sglang/jit_kernel/diffusion/cutedsl/utils.py b/python/sglang/jit_kernel/diffusion/cutedsl/utils.py index d23c2342b..eb41bd6a8 100644 --- a/python/sglang/jit_kernel/diffusion/cutedsl/utils.py +++ b/python/sglang/jit_kernel/diffusion/cutedsl/utils.py @@ -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) diff --git a/python/sglang/jit_kernel/diffusion/triton/norm.py b/python/sglang/jit_kernel/diffusion/triton/norm.py index 31ee451a4..bc44a2be1 100644 --- a/python/sglang/jit_kernel/diffusion/triton/norm.py +++ b/python/sglang/jit_kernel/diffusion/triton/norm.py @@ -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 diff --git a/python/sglang/jit_kernel/diffusion/triton/sana_wm_gdn.py b/python/sglang/jit_kernel/diffusion/triton/sana_wm_gdn.py index 54ee1f936..c9d453b05 100644 --- a/python/sglang/jit_kernel/diffusion/triton/sana_wm_gdn.py +++ b/python/sglang/jit_kernel/diffusion/triton/sana_wm_gdn.py @@ -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, diff --git a/python/sglang/jit_kernel/diffusion/triton/scale_shift.py b/python/sglang/jit_kernel/diffusion/triton/scale_shift.py index 00b5e77fa..dff528e87 100644 --- a/python/sglang/jit_kernel/diffusion/triton/scale_shift.py +++ b/python/sglang/jit_kernel/diffusion/triton/scale_shift.py @@ -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, diff --git a/python/sglang/jit_kernel/timestep_embedding.py b/python/sglang/jit_kernel/timestep_embedding.py index c65c145fa..08a941a6e 100644 --- a/python/sglang/jit_kernel/timestep_embedding.py +++ b/python/sglang/jit_kernel/timestep_embedding.py @@ -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}>", + ) + ], )