[diffusion] refactor: refactor diffusion triton kernels (#18966)
This commit is contained in:
Xiaoyu Zhang
@@ -1,6 +1,3 @@
|
||||
# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
|
||||
|
||||
# TODO: for temporary usage, expecting a refactor
|
||||
from typing import Optional
|
||||
|
||||
import torch
|
||||
@@ -8,507 +5,6 @@ import triton # type: ignore
|
||||
import triton.language as tl # type: ignore
|
||||
from torch import Tensor
|
||||
|
||||
from sglang.multimodal_gen.runtime.platforms import current_platform
|
||||
|
||||
|
||||
@triton.autotune(
|
||||
configs=[
|
||||
triton.Config({"BLOCK_N": 64}, num_warps=2),
|
||||
triton.Config({"BLOCK_N": 128}, num_warps=4),
|
||||
triton.Config({"BLOCK_N": 256}, num_warps=4),
|
||||
triton.Config({"BLOCK_N": 512}, num_warps=4),
|
||||
triton.Config({"BLOCK_N": 1024}, num_warps=8),
|
||||
],
|
||||
key=["inner_dim"],
|
||||
)
|
||||
@triton.jit
|
||||
def _fused_scale_shift_4d_kernel(
|
||||
output_ptr,
|
||||
normalized_ptr,
|
||||
scale_ptr,
|
||||
shift_ptr,
|
||||
scale_constant: tl.constexpr, # scale_constant is either 0 or 1.
|
||||
rows,
|
||||
inner_dim,
|
||||
seq_len,
|
||||
num_frames,
|
||||
frame_seqlen,
|
||||
BLOCK_N: tl.constexpr,
|
||||
):
|
||||
pid_row = tl.program_id(0)
|
||||
pid_col = tl.program_id(1)
|
||||
|
||||
col_offsets = pid_col * BLOCK_N + tl.arange(0, BLOCK_N)
|
||||
mask = col_offsets < inner_dim
|
||||
|
||||
# Pointers for normalized and output
|
||||
row_base = pid_row * inner_dim
|
||||
norm_ptrs = normalized_ptr + row_base + col_offsets
|
||||
out_ptrs = output_ptr + row_base + col_offsets
|
||||
|
||||
# Pointers for scale and shift for 4D
|
||||
b_idx = pid_row // seq_len
|
||||
t_idx = pid_row % seq_len
|
||||
frame_idx_in_batch = t_idx // frame_seqlen
|
||||
|
||||
scale_row_idx = b_idx * num_frames + frame_idx_in_batch
|
||||
scale_ptrs = scale_ptr + scale_row_idx * inner_dim + col_offsets
|
||||
shift_ptrs = shift_ptr + scale_row_idx * inner_dim + col_offsets
|
||||
|
||||
normalized = tl.load(norm_ptrs, mask=mask, other=0.0)
|
||||
scale = tl.load(scale_ptrs, mask=mask, other=0.0)
|
||||
shift = tl.load(shift_ptrs, mask=mask, other=0.0)
|
||||
|
||||
scale_const_tensor = tl.full([BLOCK_N], scale_constant, dtype=scale.dtype)
|
||||
output = normalized * (scale_const_tensor + scale) + shift
|
||||
|
||||
tl.store(out_ptrs, output, mask=mask)
|
||||
|
||||
|
||||
@triton.jit
|
||||
def fuse_scale_shift_kernel_blc_opt(
|
||||
x_ptr,
|
||||
shift_ptr,
|
||||
scale_ptr,
|
||||
scale_constant: tl.constexpr, # scale_constant is either 0 or 1.,
|
||||
y_ptr,
|
||||
B,
|
||||
L,
|
||||
C,
|
||||
stride_x_b,
|
||||
stride_x_l,
|
||||
stride_x_c,
|
||||
stride_s_b,
|
||||
stride_s_l,
|
||||
stride_s_c,
|
||||
stride_sc_b,
|
||||
stride_sc_l,
|
||||
stride_sc_c,
|
||||
SCALE_IS_SCALAR: tl.constexpr,
|
||||
SHIFT_IS_SCALAR: tl.constexpr,
|
||||
BLOCK_L: tl.constexpr,
|
||||
BLOCK_C: tl.constexpr,
|
||||
):
|
||||
pid_l = tl.program_id(0)
|
||||
pid_c = tl.program_id(1)
|
||||
pid_b = tl.program_id(2)
|
||||
|
||||
l_offsets = pid_l * BLOCK_L + tl.arange(0, BLOCK_L)
|
||||
c_offsets = pid_c * BLOCK_C + tl.arange(0, BLOCK_C)
|
||||
|
||||
mask_l = l_offsets < L
|
||||
mask_c = c_offsets < C
|
||||
mask = mask_l[:, None] & mask_c[None, :]
|
||||
|
||||
x_off = (
|
||||
pid_b * stride_x_b
|
||||
+ l_offsets[:, None] * stride_x_l
|
||||
+ c_offsets[None, :] * stride_x_c
|
||||
)
|
||||
x = tl.load(x_ptr + x_off, mask=mask, other=0)
|
||||
|
||||
if SHIFT_IS_SCALAR:
|
||||
shift_val = tl.load(shift_ptr)
|
||||
shift = tl.full((BLOCK_L, BLOCK_C), shift_val, dtype=shift_val.dtype)
|
||||
else:
|
||||
s_off = (
|
||||
pid_b * stride_s_b
|
||||
+ l_offsets[:, None] * stride_s_l
|
||||
+ c_offsets[None, :] * stride_s_c
|
||||
)
|
||||
shift = tl.load(shift_ptr + s_off, mask=mask, other=0)
|
||||
|
||||
if SCALE_IS_SCALAR:
|
||||
scale_val = tl.load(scale_ptr)
|
||||
scale = tl.full((BLOCK_L, BLOCK_C), scale_val, dtype=scale_val.dtype)
|
||||
else:
|
||||
sc_off = (
|
||||
pid_b * stride_sc_b
|
||||
+ l_offsets[:, None] * stride_sc_l
|
||||
+ c_offsets[None, :] * stride_sc_c
|
||||
)
|
||||
scale = tl.load(scale_ptr + sc_off, mask=mask, other=0)
|
||||
|
||||
y = x * (scale_constant + scale) + shift
|
||||
tl.store(y_ptr + x_off, y, mask=mask)
|
||||
|
||||
|
||||
@triton.jit
|
||||
def fuse_scale_shift_gate_select01_kernel_blc_opt(
|
||||
x_ptr,
|
||||
shift0_ptr,
|
||||
scale0_ptr,
|
||||
gate0_ptr,
|
||||
shift1_ptr,
|
||||
scale1_ptr,
|
||||
gate1_ptr,
|
||||
index_ptr,
|
||||
y_ptr,
|
||||
gate_out_ptr,
|
||||
B,
|
||||
L,
|
||||
C,
|
||||
stride_x_b,
|
||||
stride_x_l,
|
||||
stride_x_c,
|
||||
stride_s0_b,
|
||||
stride_s0_c,
|
||||
stride_sc0_b,
|
||||
stride_sc0_c,
|
||||
stride_g0_b,
|
||||
stride_g0_c,
|
||||
stride_s1_b,
|
||||
stride_s1_c,
|
||||
stride_sc1_b,
|
||||
stride_sc1_c,
|
||||
stride_g1_b,
|
||||
stride_g1_c,
|
||||
stride_i_b,
|
||||
stride_i_l,
|
||||
stride_go_b,
|
||||
stride_go_l,
|
||||
stride_go_c,
|
||||
BLOCK_L: tl.constexpr,
|
||||
BLOCK_C: tl.constexpr,
|
||||
):
|
||||
pid_l = tl.program_id(0)
|
||||
pid_c = tl.program_id(1)
|
||||
pid_b = tl.program_id(2)
|
||||
|
||||
l_offsets = pid_l * BLOCK_L + tl.arange(0, BLOCK_L)
|
||||
c_offsets = pid_c * BLOCK_C + tl.arange(0, BLOCK_C)
|
||||
|
||||
mask_l = l_offsets < L
|
||||
mask_c = c_offsets < C
|
||||
mask = mask_l[:, None] & mask_c[None, :]
|
||||
|
||||
x_off = (
|
||||
pid_b * stride_x_b
|
||||
+ l_offsets[:, None] * stride_x_l
|
||||
+ c_offsets[None, :] * stride_x_c
|
||||
)
|
||||
x = tl.load(x_ptr + x_off, mask=mask, other=0)
|
||||
|
||||
idx_off = pid_b * stride_i_b + l_offsets * stride_i_l
|
||||
idx = tl.load(index_ptr + idx_off, mask=mask_l, other=0).to(tl.int1)[:, None]
|
||||
|
||||
s0_off = pid_b * stride_s0_b + c_offsets[None, :] * stride_s0_c
|
||||
sc0_off = pid_b * stride_sc0_b + c_offsets[None, :] * stride_sc0_c
|
||||
g0_off = pid_b * stride_g0_b + c_offsets[None, :] * stride_g0_c
|
||||
s1_off = pid_b * stride_s1_b + c_offsets[None, :] * stride_s1_c
|
||||
sc1_off = pid_b * stride_sc1_b + c_offsets[None, :] * stride_sc1_c
|
||||
g1_off = pid_b * stride_g1_b + c_offsets[None, :] * stride_g1_c
|
||||
|
||||
shift0 = tl.load(shift0_ptr + s0_off, mask=mask_c[None, :], other=0)
|
||||
scale0 = tl.load(scale0_ptr + sc0_off, mask=mask_c[None, :], other=0)
|
||||
gate0 = tl.load(gate0_ptr + g0_off, mask=mask_c[None, :], other=0)
|
||||
shift1 = tl.load(shift1_ptr + s1_off, mask=mask_c[None, :], other=0)
|
||||
scale1 = tl.load(scale1_ptr + sc1_off, mask=mask_c[None, :], other=0)
|
||||
gate1 = tl.load(gate1_ptr + g1_off, mask=mask_c[None, :], other=0)
|
||||
|
||||
shift = tl.where(idx, shift1, shift0)
|
||||
scale = tl.where(idx, scale1, scale0)
|
||||
gate = tl.where(idx, gate1, gate0)
|
||||
|
||||
y = x * (1 + scale) + shift
|
||||
tl.store(y_ptr + x_off, y, mask=mask)
|
||||
|
||||
go_off = (
|
||||
pid_b * stride_go_b
|
||||
+ l_offsets[:, None] * stride_go_l
|
||||
+ c_offsets[None, :] * stride_go_c
|
||||
)
|
||||
tl.store(gate_out_ptr + go_off, gate, mask=mask)
|
||||
|
||||
|
||||
def fuse_scale_shift_kernel(
|
||||
x: torch.Tensor,
|
||||
scale: torch.Tensor,
|
||||
shift: torch.Tensor,
|
||||
scale_constant: float = 1.0,
|
||||
block_l: int = 128,
|
||||
block_c: int = 128,
|
||||
):
|
||||
assert x.is_cuda and scale.is_cuda
|
||||
assert x.is_contiguous()
|
||||
|
||||
B, L, C = x.shape
|
||||
output = torch.empty_like(x)
|
||||
|
||||
if scale.dim() == 4:
|
||||
# scale/shift: [B, F, 1, C]
|
||||
rows = B * L
|
||||
x_2d = x.view(rows, C)
|
||||
output_2d = output.view(rows, C)
|
||||
grid = lambda META: (rows, triton.cdiv(C, META["BLOCK_N"]))
|
||||
num_frames = scale.shape[1]
|
||||
assert (
|
||||
L % num_frames == 0
|
||||
), "seq_len must be divisible by num_frames for 4D scale/shift"
|
||||
frame_seqlen = L // num_frames
|
||||
|
||||
# Compact [B, F, C] without the singleton dim into [B*F, C]
|
||||
scale_reshaped = scale.squeeze(2).reshape(-1, C).contiguous()
|
||||
shift_reshaped = shift.squeeze(2).reshape(-1, C).contiguous()
|
||||
|
||||
_fused_scale_shift_4d_kernel[grid](
|
||||
output_2d,
|
||||
x_2d,
|
||||
scale_reshaped,
|
||||
shift_reshaped,
|
||||
scale_constant,
|
||||
rows,
|
||||
C,
|
||||
L,
|
||||
num_frames,
|
||||
frame_seqlen,
|
||||
)
|
||||
else:
|
||||
# 2D: [B, C] or [1, C] -> treat as [B, 1, C] and broadcast over L
|
||||
# 3D: [B, L, C] (or broadcastable variants like [B, 1, C], [1, L, C], [1, 1, C])
|
||||
# Also support scalar (0D or 1-element)
|
||||
if scale.dim() == 0 or (scale.dim() == 1 and scale.numel() == 1):
|
||||
scale_blc = scale.reshape(1)
|
||||
elif scale.dim() == 2:
|
||||
scale_blc = scale[:, None, :]
|
||||
elif scale.dim() == 3:
|
||||
scale_blc = scale
|
||||
else:
|
||||
raise ValueError("scale must be 0D/1D(1)/2D/3D or 4D")
|
||||
|
||||
if shift.dim() == 0 or (shift.dim() == 1 and shift.numel() == 1):
|
||||
shift_blc = shift.reshape(1)
|
||||
elif shift.dim() == 2:
|
||||
shift_blc = shift[:, None, :]
|
||||
elif shift.dim() == 3:
|
||||
shift_blc = shift
|
||||
else:
|
||||
# broadcast later via expand if possible
|
||||
shift_blc = shift
|
||||
|
||||
need_scale_scalar = scale_blc.dim() == 1 and scale_blc.numel() == 1
|
||||
need_shift_scalar = shift_blc.dim() == 1 and shift_blc.numel() == 1
|
||||
|
||||
if not need_scale_scalar:
|
||||
scale_exp = scale_blc.expand(B, L, C)
|
||||
s_sb, s_sl, s_sc = scale_exp.stride()
|
||||
else:
|
||||
s_sb = s_sl = s_sc = 0
|
||||
|
||||
if not need_shift_scalar:
|
||||
shift_exp = shift_blc.expand(B, L, C)
|
||||
sh_sb, sh_sl, sh_sc = shift_exp.stride()
|
||||
else:
|
||||
sh_sb = sh_sl = sh_sc = 0
|
||||
|
||||
# If both scalars and both zero, copy fast-path
|
||||
if need_scale_scalar and need_shift_scalar:
|
||||
if (scale_blc.abs().max() == 0) and (shift_blc.abs().max() == 0):
|
||||
output.copy_(x)
|
||||
return output
|
||||
|
||||
grid = (triton.cdiv(L, block_l), triton.cdiv(C, block_c), B)
|
||||
fuse_scale_shift_kernel_blc_opt[grid](
|
||||
x,
|
||||
shift_blc if need_shift_scalar else shift_exp,
|
||||
scale_blc if need_scale_scalar else scale_exp,
|
||||
scale_constant,
|
||||
output,
|
||||
B,
|
||||
L,
|
||||
C,
|
||||
x.stride(0),
|
||||
x.stride(1),
|
||||
x.stride(2),
|
||||
sh_sb,
|
||||
sh_sl,
|
||||
sh_sc,
|
||||
s_sb,
|
||||
s_sl,
|
||||
s_sc,
|
||||
SCALE_IS_SCALAR=need_scale_scalar,
|
||||
SHIFT_IS_SCALAR=need_shift_scalar,
|
||||
BLOCK_L=block_l,
|
||||
BLOCK_C=block_c,
|
||||
num_warps=4,
|
||||
num_stages=2,
|
||||
)
|
||||
return output
|
||||
|
||||
|
||||
def fuse_scale_shift_gate_select01_kernel(
|
||||
x: torch.Tensor,
|
||||
scale0: torch.Tensor,
|
||||
shift0: torch.Tensor,
|
||||
gate0: torch.Tensor,
|
||||
scale1: torch.Tensor,
|
||||
shift1: torch.Tensor,
|
||||
gate1: torch.Tensor,
|
||||
index: torch.Tensor,
|
||||
block_l: int = 128,
|
||||
block_c: int = 128,
|
||||
):
|
||||
assert x.is_contiguous()
|
||||
B, L, C = x.shape
|
||||
output = torch.empty_like(x)
|
||||
gate_out = torch.empty_like(x)
|
||||
|
||||
if (
|
||||
scale0.dim() != 2
|
||||
or shift0.dim() != 2
|
||||
or gate0.dim() != 2
|
||||
or scale1.dim() != 2
|
||||
or shift1.dim() != 2
|
||||
or gate1.dim() != 2
|
||||
):
|
||||
raise ValueError("scale0/shift0/gate0/scale1/shift1/gate1 must be 2D [B, C]")
|
||||
if index.dim() != 2:
|
||||
raise ValueError("index must be 2D [B, L]")
|
||||
|
||||
grid = (triton.cdiv(L, block_l), triton.cdiv(C, block_c), B)
|
||||
fuse_scale_shift_gate_select01_kernel_blc_opt[grid](
|
||||
x,
|
||||
shift0,
|
||||
scale0,
|
||||
gate0,
|
||||
shift1,
|
||||
scale1,
|
||||
gate1,
|
||||
index,
|
||||
output,
|
||||
gate_out,
|
||||
B,
|
||||
L,
|
||||
C,
|
||||
x.stride(0),
|
||||
x.stride(1),
|
||||
x.stride(2),
|
||||
shift0.stride(0),
|
||||
shift0.stride(1),
|
||||
scale0.stride(0),
|
||||
scale0.stride(1),
|
||||
gate0.stride(0),
|
||||
gate0.stride(1),
|
||||
shift1.stride(0),
|
||||
shift1.stride(1),
|
||||
scale1.stride(0),
|
||||
scale1.stride(1),
|
||||
gate1.stride(0),
|
||||
gate1.stride(1),
|
||||
index.stride(0),
|
||||
index.stride(1),
|
||||
gate_out.stride(0),
|
||||
gate_out.stride(1),
|
||||
gate_out.stride(2),
|
||||
BLOCK_L=block_l,
|
||||
BLOCK_C=block_c,
|
||||
num_warps=4,
|
||||
num_stages=2,
|
||||
)
|
||||
return output, gate_out
|
||||
|
||||
|
||||
@triton.autotune(
|
||||
configs=[
|
||||
triton.Config({"BLOCK_HS_HALF": 32}, num_warps=2),
|
||||
triton.Config({"BLOCK_HS_HALF": 64}, num_warps=4),
|
||||
triton.Config({"BLOCK_HS_HALF": 128}, num_warps=4),
|
||||
triton.Config({"BLOCK_HS_HALF": 256}, num_warps=8),
|
||||
],
|
||||
key=["head_size", "interleaved"],
|
||||
)
|
||||
@triton.jit
|
||||
def _rotary_embedding_kernel(
|
||||
output_ptr,
|
||||
x_ptr,
|
||||
cos_ptr,
|
||||
sin_ptr,
|
||||
num_heads,
|
||||
head_size,
|
||||
num_tokens,
|
||||
stride_x_row,
|
||||
stride_cos_row,
|
||||
stride_sin_row,
|
||||
interleaved: tl.constexpr,
|
||||
BLOCK_HS_HALF: tl.constexpr,
|
||||
):
|
||||
row_idx = tl.program_id(0)
|
||||
token_idx = (row_idx // num_heads) % num_tokens
|
||||
|
||||
x_row_ptr = x_ptr + row_idx * stride_x_row
|
||||
cos_row_ptr = cos_ptr + token_idx * stride_cos_row
|
||||
sin_row_ptr = sin_ptr + token_idx * stride_sin_row
|
||||
output_row_ptr = output_ptr + row_idx * stride_x_row
|
||||
|
||||
# half size for x1 and x2
|
||||
head_size_half = head_size // 2
|
||||
|
||||
for block_start in range(0, head_size_half, BLOCK_HS_HALF):
|
||||
offsets_half = block_start + tl.arange(0, BLOCK_HS_HALF)
|
||||
mask = offsets_half < head_size_half
|
||||
|
||||
cos_vals = tl.load(cos_row_ptr + offsets_half, mask=mask, other=0.0)
|
||||
sin_vals = tl.load(sin_row_ptr + offsets_half, mask=mask, other=0.0)
|
||||
|
||||
offsets_x1 = 2 * offsets_half
|
||||
offsets_x2 = 2 * offsets_half + 1
|
||||
|
||||
x1_vals = tl.load(x_row_ptr + offsets_x1, mask=mask, other=0.0)
|
||||
x2_vals = tl.load(x_row_ptr + offsets_x2, mask=mask, other=0.0)
|
||||
|
||||
x1_fp32 = x1_vals.to(tl.float32)
|
||||
x2_fp32 = x2_vals.to(tl.float32)
|
||||
cos_fp32 = cos_vals.to(tl.float32)
|
||||
sin_fp32 = sin_vals.to(tl.float32)
|
||||
o1_vals = tl.fma(-x2_fp32, sin_fp32, x1_fp32 * cos_fp32)
|
||||
o2_vals = tl.fma(x1_fp32, sin_fp32, x2_fp32 * cos_fp32)
|
||||
|
||||
tl.store(output_row_ptr + offsets_x1, o1_vals.to(x1_vals.dtype), mask=mask)
|
||||
tl.store(output_row_ptr + offsets_x2, o2_vals.to(x2_vals.dtype), mask=mask)
|
||||
|
||||
|
||||
def apply_rotary_embedding(
|
||||
x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, interleaved: bool = False
|
||||
) -> torch.Tensor:
|
||||
output = torch.empty_like(x)
|
||||
|
||||
if x.dim() > 3:
|
||||
bsz, num_tokens, num_heads, head_size = x.shape
|
||||
else:
|
||||
num_tokens, num_heads, head_size = x.shape
|
||||
bsz = 1
|
||||
|
||||
assert head_size % 2 == 0, "head_size must be divisible by 2"
|
||||
|
||||
x_reshaped = x.view(-1, head_size)
|
||||
output_reshaped = output.view(-1, head_size)
|
||||
|
||||
# num_tokens per head, 1 token per block
|
||||
grid = (bsz * num_tokens * num_heads,)
|
||||
|
||||
if interleaved and cos.shape[-1] == head_size:
|
||||
cos = cos[..., ::2].contiguous()
|
||||
sin = sin[..., ::2].contiguous()
|
||||
else:
|
||||
cos = cos.contiguous()
|
||||
sin = sin.contiguous()
|
||||
|
||||
_rotary_embedding_kernel[grid](
|
||||
output_reshaped,
|
||||
x_reshaped,
|
||||
cos,
|
||||
sin,
|
||||
num_heads,
|
||||
head_size,
|
||||
num_tokens,
|
||||
x_reshaped.stride(0),
|
||||
cos.stride(0),
|
||||
sin.stride(0),
|
||||
interleaved,
|
||||
)
|
||||
|
||||
return output
|
||||
|
||||
|
||||
# RMSNorm-fp32
|
||||
def maybe_contiguous_lastdim(x):
|
||||
@@ -1122,86 +618,3 @@ def rms_norm_fn(
|
||||
out,
|
||||
residual_out,
|
||||
)
|
||||
|
||||
|
||||
# Adapted from https://github.com/ModelTC/LightX2V/blob/main/lightx2v/common/ops/norm/triton_ops.py#L905-L956
|
||||
@triton.jit
|
||||
def _rms_norm_tiled_onepass(
|
||||
y_ptr,
|
||||
x_ptr,
|
||||
w_ptr,
|
||||
SEQ: tl.constexpr,
|
||||
DIM: tl.constexpr,
|
||||
EPS: tl.constexpr,
|
||||
BLOCK_SIZE_SEQ: tl.constexpr,
|
||||
BLOCK_SIZE_DIM: tl.constexpr,
|
||||
):
|
||||
seq_blk_id = tl.program_id(0)
|
||||
seq_id = seq_blk_id * BLOCK_SIZE_SEQ
|
||||
|
||||
seq_offset = seq_id + tl.arange(0, BLOCK_SIZE_SEQ)[:, None]
|
||||
s_mask = seq_offset < SEQ
|
||||
d_offset = tl.arange(0, BLOCK_SIZE_DIM)[None, :]
|
||||
d_mask = d_offset < DIM
|
||||
y_blk = y_ptr + seq_offset * DIM + d_offset
|
||||
x_blk = x_ptr + seq_offset * DIM + d_offset
|
||||
mask = s_mask & d_mask
|
||||
|
||||
x = tl.load(x_blk, mask=mask, other=0.0).to(tl.float32)
|
||||
mean_square = tl.sum(x * x, axis=1, keep_dims=True) / DIM
|
||||
rstd = tl.math.rsqrt(mean_square + EPS)
|
||||
w = tl.load(w_ptr + d_offset, mask=d_mask)
|
||||
tl.store(y_blk, x * rstd * w, mask=mask)
|
||||
|
||||
|
||||
def triton_one_pass_rms_norm(x: torch.Tensor, w: torch.Tensor, eps: float = 1e-6):
|
||||
shape = x.shape
|
||||
x = x.contiguous()
|
||||
y = torch.empty_like(x)
|
||||
x_view = x.reshape(-1, shape[-1])
|
||||
y_view = y.reshape(-1, shape[-1])
|
||||
S, D = x_view.shape
|
||||
|
||||
BLOCK_SIZE_SEQ = min(16, triton.next_power_of_2(max(1, S // 512)))
|
||||
grid = (triton.cdiv(S, BLOCK_SIZE_SEQ),)
|
||||
|
||||
with torch.cuda.device(x.device):
|
||||
torch.library.wrap_triton(_rms_norm_tiled_onepass)[grid](
|
||||
y_view,
|
||||
x_view,
|
||||
w,
|
||||
S,
|
||||
D,
|
||||
eps,
|
||||
BLOCK_SIZE_DIM=triton.next_power_of_2(D),
|
||||
BLOCK_SIZE_SEQ=BLOCK_SIZE_SEQ,
|
||||
)
|
||||
return y
|
||||
|
||||
|
||||
if current_platform.is_npu():
|
||||
# TODO: remove this when triton ascend bug is fixed
|
||||
def fuse_scale_shift_native(
|
||||
x: torch.Tensor,
|
||||
scale: torch.Tensor,
|
||||
shift: torch.Tensor,
|
||||
block_l: int = 128,
|
||||
block_c: int = 128,
|
||||
):
|
||||
return x * (1 + scale) + shift
|
||||
|
||||
fuse_scale_shift_kernel = fuse_scale_shift_native
|
||||
|
||||
# TODO: remove this when triton ascend bug is fixed
|
||||
def apply_rotary_embedding_native(
|
||||
x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, interleaved: bool = False
|
||||
) -> torch.Tensor:
|
||||
cos = cos.unsqueeze(-2).to(x.dtype)
|
||||
sin = sin.unsqueeze(-2).to(x.dtype)
|
||||
x1 = x[..., ::2]
|
||||
x2 = x[..., 1::2]
|
||||
o1 = x1 * cos - x2 * sin
|
||||
o2 = x2 * cos + x1 * sin
|
||||
return torch.stack((o1, o2), dim=-1).flatten(-2)
|
||||
|
||||
apply_rotary_embedding = apply_rotary_embedding_native
|
||||
25
python/sglang/jit_kernel/diffusion/triton/npu_fallback.py
Normal file
25
python/sglang/jit_kernel/diffusion/triton/npu_fallback.py
Normal file
@@ -0,0 +1,25 @@
|
||||
import torch
|
||||
|
||||
|
||||
# TODO: remove this when triton ascend bug is fixed
|
||||
def fuse_scale_shift_native(
|
||||
x: torch.Tensor,
|
||||
scale: torch.Tensor,
|
||||
shift: torch.Tensor,
|
||||
block_l: int = 128,
|
||||
block_c: int = 128,
|
||||
):
|
||||
return x * (1 + scale) + shift
|
||||
|
||||
|
||||
# TODO: remove this when triton ascend bug is fixed
|
||||
def apply_rotary_embedding_native(
|
||||
x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, interleaved: bool = False
|
||||
) -> torch.Tensor:
|
||||
cos = cos.unsqueeze(-2).to(x.dtype)
|
||||
sin = sin.unsqueeze(-2).to(x.dtype)
|
||||
x1 = x[..., ::2]
|
||||
x2 = x[..., 1::2]
|
||||
o1 = x1 * cos - x2 * sin
|
||||
o2 = x2 * cos + x1 * sin
|
||||
return torch.stack((o1, o2), dim=-1).flatten(-2)
|
||||
58
python/sglang/jit_kernel/diffusion/triton/rmsnorm_onepass.py
Normal file
58
python/sglang/jit_kernel/diffusion/triton/rmsnorm_onepass.py
Normal file
@@ -0,0 +1,58 @@
|
||||
import torch
|
||||
import triton # type: ignore
|
||||
import triton.language as tl # type: ignore
|
||||
|
||||
|
||||
# Adapted from https://github.com/ModelTC/LightX2V/blob/main/lightx2v/common/ops/norm/triton_ops.py#L905-L956
|
||||
@triton.jit
|
||||
def _rms_norm_tiled_onepass(
|
||||
y_ptr,
|
||||
x_ptr,
|
||||
w_ptr,
|
||||
SEQ: tl.constexpr,
|
||||
DIM: tl.constexpr,
|
||||
EPS: tl.constexpr,
|
||||
BLOCK_SIZE_SEQ: tl.constexpr,
|
||||
BLOCK_SIZE_DIM: tl.constexpr,
|
||||
):
|
||||
seq_blk_id = tl.program_id(0)
|
||||
seq_id = seq_blk_id * BLOCK_SIZE_SEQ
|
||||
|
||||
seq_offset = seq_id + tl.arange(0, BLOCK_SIZE_SEQ)[:, None]
|
||||
s_mask = seq_offset < SEQ
|
||||
d_offset = tl.arange(0, BLOCK_SIZE_DIM)[None, :]
|
||||
d_mask = d_offset < DIM
|
||||
y_blk = y_ptr + seq_offset * DIM + d_offset
|
||||
x_blk = x_ptr + seq_offset * DIM + d_offset
|
||||
mask = s_mask & d_mask
|
||||
|
||||
x = tl.load(x_blk, mask=mask, other=0.0).to(tl.float32)
|
||||
mean_square = tl.sum(x * x, axis=1, keep_dims=True) / DIM
|
||||
rstd = tl.math.rsqrt(mean_square + EPS)
|
||||
w = tl.load(w_ptr + d_offset, mask=d_mask)
|
||||
tl.store(y_blk, x * rstd * w, mask=mask)
|
||||
|
||||
|
||||
def triton_one_pass_rms_norm(x: torch.Tensor, w: torch.Tensor, eps: float = 1e-6):
|
||||
shape = x.shape
|
||||
x = x.contiguous()
|
||||
y = torch.empty_like(x)
|
||||
x_view = x.reshape(-1, shape[-1])
|
||||
y_view = y.reshape(-1, shape[-1])
|
||||
S, D = x_view.shape
|
||||
|
||||
BLOCK_SIZE_SEQ = min(16, triton.next_power_of_2(max(1, S // 512)))
|
||||
grid = (triton.cdiv(S, BLOCK_SIZE_SEQ),)
|
||||
|
||||
with torch.get_device_module().device(x.device):
|
||||
torch.library.wrap_triton(_rms_norm_tiled_onepass)[grid](
|
||||
y_view,
|
||||
x_view,
|
||||
w,
|
||||
S,
|
||||
D,
|
||||
eps,
|
||||
BLOCK_SIZE_DIM=triton.next_power_of_2(D),
|
||||
BLOCK_SIZE_SEQ=BLOCK_SIZE_SEQ,
|
||||
)
|
||||
return y
|
||||
113
python/sglang/jit_kernel/diffusion/triton/rotary.py
Normal file
113
python/sglang/jit_kernel/diffusion/triton/rotary.py
Normal file
@@ -0,0 +1,113 @@
|
||||
import torch
|
||||
import triton # type: ignore
|
||||
import triton.language as tl # type: ignore
|
||||
|
||||
from sglang.multimodal_gen.runtime.platforms import current_platform
|
||||
|
||||
|
||||
@triton.autotune(
|
||||
configs=[
|
||||
triton.Config({"BLOCK_HS_HALF": 32}, num_warps=2),
|
||||
triton.Config({"BLOCK_HS_HALF": 64}, num_warps=4),
|
||||
triton.Config({"BLOCK_HS_HALF": 128}, num_warps=4),
|
||||
triton.Config({"BLOCK_HS_HALF": 256}, num_warps=8),
|
||||
],
|
||||
key=["head_size", "interleaved"],
|
||||
)
|
||||
@triton.jit
|
||||
def _rotary_embedding_kernel(
|
||||
output_ptr,
|
||||
x_ptr,
|
||||
cos_ptr,
|
||||
sin_ptr,
|
||||
num_heads,
|
||||
head_size,
|
||||
num_tokens,
|
||||
stride_x_row,
|
||||
stride_cos_row,
|
||||
stride_sin_row,
|
||||
interleaved: tl.constexpr,
|
||||
BLOCK_HS_HALF: tl.constexpr,
|
||||
):
|
||||
row_idx = tl.program_id(0)
|
||||
token_idx = (row_idx // num_heads) % num_tokens
|
||||
|
||||
x_row_ptr = x_ptr + row_idx * stride_x_row
|
||||
cos_row_ptr = cos_ptr + token_idx * stride_cos_row
|
||||
sin_row_ptr = sin_ptr + token_idx * stride_sin_row
|
||||
output_row_ptr = output_ptr + row_idx * stride_x_row
|
||||
|
||||
# half size for x1 and x2
|
||||
head_size_half = head_size // 2
|
||||
|
||||
for block_start in range(0, head_size_half, BLOCK_HS_HALF):
|
||||
offsets_half = block_start + tl.arange(0, BLOCK_HS_HALF)
|
||||
mask = offsets_half < head_size_half
|
||||
|
||||
cos_vals = tl.load(cos_row_ptr + offsets_half, mask=mask, other=0.0)
|
||||
sin_vals = tl.load(sin_row_ptr + offsets_half, mask=mask, other=0.0)
|
||||
|
||||
offsets_x1 = 2 * offsets_half
|
||||
offsets_x2 = 2 * offsets_half + 1
|
||||
|
||||
x1_vals = tl.load(x_row_ptr + offsets_x1, mask=mask, other=0.0)
|
||||
x2_vals = tl.load(x_row_ptr + offsets_x2, mask=mask, other=0.0)
|
||||
|
||||
x1_fp32 = x1_vals.to(tl.float32)
|
||||
x2_fp32 = x2_vals.to(tl.float32)
|
||||
cos_fp32 = cos_vals.to(tl.float32)
|
||||
sin_fp32 = sin_vals.to(tl.float32)
|
||||
o1_vals = tl.fma(-x2_fp32, sin_fp32, x1_fp32 * cos_fp32)
|
||||
o2_vals = tl.fma(x1_fp32, sin_fp32, x2_fp32 * cos_fp32)
|
||||
|
||||
tl.store(output_row_ptr + offsets_x1, o1_vals.to(x1_vals.dtype), mask=mask)
|
||||
tl.store(output_row_ptr + offsets_x2, o2_vals.to(x2_vals.dtype), mask=mask)
|
||||
|
||||
|
||||
def apply_rotary_embedding(
|
||||
x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, interleaved: bool = False
|
||||
) -> torch.Tensor:
|
||||
output = torch.empty_like(x)
|
||||
|
||||
if x.dim() > 3:
|
||||
bsz, num_tokens, num_heads, head_size = x.shape
|
||||
else:
|
||||
num_tokens, num_heads, head_size = x.shape
|
||||
bsz = 1
|
||||
|
||||
assert head_size % 2 == 0, "head_size must be divisible by 2"
|
||||
|
||||
x_reshaped = x.view(-1, head_size)
|
||||
output_reshaped = output.view(-1, head_size)
|
||||
|
||||
# num_tokens per head, 1 token per block
|
||||
grid = (bsz * num_tokens * num_heads,)
|
||||
|
||||
if interleaved and cos.shape[-1] == head_size:
|
||||
cos = cos[..., ::2].contiguous()
|
||||
sin = sin[..., ::2].contiguous()
|
||||
else:
|
||||
cos = cos.contiguous()
|
||||
sin = sin.contiguous()
|
||||
|
||||
_rotary_embedding_kernel[grid](
|
||||
output_reshaped,
|
||||
x_reshaped,
|
||||
cos,
|
||||
sin,
|
||||
num_heads,
|
||||
head_size,
|
||||
num_tokens,
|
||||
x_reshaped.stride(0),
|
||||
cos.stride(0),
|
||||
sin.stride(0),
|
||||
interleaved,
|
||||
)
|
||||
|
||||
return output
|
||||
|
||||
|
||||
if current_platform.is_npu():
|
||||
from .npu_fallback import apply_rotary_embedding_native
|
||||
|
||||
apply_rotary_embedding = apply_rotary_embedding_native
|
||||
408
python/sglang/jit_kernel/diffusion/triton/scale_shift.py
Normal file
408
python/sglang/jit_kernel/diffusion/triton/scale_shift.py
Normal file
@@ -0,0 +1,408 @@
|
||||
import torch
|
||||
import triton # type: ignore
|
||||
import triton.language as tl # type: ignore
|
||||
|
||||
from sglang.multimodal_gen.runtime.platforms import current_platform
|
||||
|
||||
|
||||
@triton.autotune(
|
||||
configs=[
|
||||
triton.Config({"BLOCK_N": 64}, num_warps=2),
|
||||
triton.Config({"BLOCK_N": 128}, num_warps=4),
|
||||
triton.Config({"BLOCK_N": 256}, num_warps=4),
|
||||
triton.Config({"BLOCK_N": 512}, num_warps=4),
|
||||
triton.Config({"BLOCK_N": 1024}, num_warps=8),
|
||||
],
|
||||
key=["inner_dim"],
|
||||
)
|
||||
@triton.jit
|
||||
def _fused_scale_shift_4d_kernel(
|
||||
output_ptr,
|
||||
normalized_ptr,
|
||||
scale_ptr,
|
||||
shift_ptr,
|
||||
scale_constant: tl.constexpr, # scale_constant is either 0 or 1.
|
||||
rows,
|
||||
inner_dim,
|
||||
seq_len,
|
||||
num_frames,
|
||||
frame_seqlen,
|
||||
BLOCK_N: tl.constexpr,
|
||||
):
|
||||
pid_row = tl.program_id(0)
|
||||
pid_col = tl.program_id(1)
|
||||
|
||||
col_offsets = pid_col * BLOCK_N + tl.arange(0, BLOCK_N)
|
||||
mask = col_offsets < inner_dim
|
||||
|
||||
# Pointers for normalized and output
|
||||
row_base = pid_row * inner_dim
|
||||
norm_ptrs = normalized_ptr + row_base + col_offsets
|
||||
out_ptrs = output_ptr + row_base + col_offsets
|
||||
|
||||
# Pointers for scale and shift for 4D
|
||||
b_idx = pid_row // seq_len
|
||||
t_idx = pid_row % seq_len
|
||||
frame_idx_in_batch = t_idx // frame_seqlen
|
||||
|
||||
scale_row_idx = b_idx * num_frames + frame_idx_in_batch
|
||||
scale_ptrs = scale_ptr + scale_row_idx * inner_dim + col_offsets
|
||||
shift_ptrs = shift_ptr + scale_row_idx * inner_dim + col_offsets
|
||||
|
||||
normalized = tl.load(norm_ptrs, mask=mask, other=0.0)
|
||||
scale = tl.load(scale_ptrs, mask=mask, other=0.0)
|
||||
shift = tl.load(shift_ptrs, mask=mask, other=0.0)
|
||||
|
||||
scale_const_tensor = tl.full([BLOCK_N], scale_constant, dtype=scale.dtype)
|
||||
output = normalized * (scale_const_tensor + scale) + shift
|
||||
|
||||
tl.store(out_ptrs, output, mask=mask)
|
||||
|
||||
|
||||
@triton.jit
|
||||
def fuse_scale_shift_kernel_blc_opt(
|
||||
x_ptr,
|
||||
shift_ptr,
|
||||
scale_ptr,
|
||||
scale_constant: tl.constexpr, # scale_constant is either 0 or 1.,
|
||||
y_ptr,
|
||||
B,
|
||||
L,
|
||||
C,
|
||||
stride_x_b,
|
||||
stride_x_l,
|
||||
stride_x_c,
|
||||
stride_s_b,
|
||||
stride_s_l,
|
||||
stride_s_c,
|
||||
stride_sc_b,
|
||||
stride_sc_l,
|
||||
stride_sc_c,
|
||||
SCALE_IS_SCALAR: tl.constexpr,
|
||||
SHIFT_IS_SCALAR: tl.constexpr,
|
||||
BLOCK_L: tl.constexpr,
|
||||
BLOCK_C: tl.constexpr,
|
||||
):
|
||||
pid_l = tl.program_id(0)
|
||||
pid_c = tl.program_id(1)
|
||||
pid_b = tl.program_id(2)
|
||||
|
||||
l_offsets = pid_l * BLOCK_L + tl.arange(0, BLOCK_L)
|
||||
c_offsets = pid_c * BLOCK_C + tl.arange(0, BLOCK_C)
|
||||
|
||||
mask_l = l_offsets < L
|
||||
mask_c = c_offsets < C
|
||||
mask = mask_l[:, None] & mask_c[None, :]
|
||||
|
||||
x_off = (
|
||||
pid_b * stride_x_b
|
||||
+ l_offsets[:, None] * stride_x_l
|
||||
+ c_offsets[None, :] * stride_x_c
|
||||
)
|
||||
x = tl.load(x_ptr + x_off, mask=mask, other=0)
|
||||
|
||||
if SHIFT_IS_SCALAR:
|
||||
shift_val = tl.load(shift_ptr)
|
||||
shift = tl.full((BLOCK_L, BLOCK_C), shift_val, dtype=shift_val.dtype)
|
||||
else:
|
||||
s_off = (
|
||||
pid_b * stride_s_b
|
||||
+ l_offsets[:, None] * stride_s_l
|
||||
+ c_offsets[None, :] * stride_s_c
|
||||
)
|
||||
shift = tl.load(shift_ptr + s_off, mask=mask, other=0)
|
||||
|
||||
if SCALE_IS_SCALAR:
|
||||
scale_val = tl.load(scale_ptr)
|
||||
scale = tl.full((BLOCK_L, BLOCK_C), scale_val, dtype=scale_val.dtype)
|
||||
else:
|
||||
sc_off = (
|
||||
pid_b * stride_sc_b
|
||||
+ l_offsets[:, None] * stride_sc_l
|
||||
+ c_offsets[None, :] * stride_sc_c
|
||||
)
|
||||
scale = tl.load(scale_ptr + sc_off, mask=mask, other=0)
|
||||
|
||||
y = x * (scale_constant + scale) + shift
|
||||
tl.store(y_ptr + x_off, y, mask=mask)
|
||||
|
||||
|
||||
@triton.jit
|
||||
def fuse_scale_shift_gate_select01_kernel_blc_opt(
|
||||
x_ptr,
|
||||
shift0_ptr,
|
||||
scale0_ptr,
|
||||
gate0_ptr,
|
||||
shift1_ptr,
|
||||
scale1_ptr,
|
||||
gate1_ptr,
|
||||
index_ptr,
|
||||
y_ptr,
|
||||
gate_out_ptr,
|
||||
B,
|
||||
L,
|
||||
C,
|
||||
stride_x_b,
|
||||
stride_x_l,
|
||||
stride_x_c,
|
||||
stride_s0_b,
|
||||
stride_s0_c,
|
||||
stride_sc0_b,
|
||||
stride_sc0_c,
|
||||
stride_g0_b,
|
||||
stride_g0_c,
|
||||
stride_s1_b,
|
||||
stride_s1_c,
|
||||
stride_sc1_b,
|
||||
stride_sc1_c,
|
||||
stride_g1_b,
|
||||
stride_g1_c,
|
||||
stride_i_b,
|
||||
stride_i_l,
|
||||
stride_go_b,
|
||||
stride_go_l,
|
||||
stride_go_c,
|
||||
BLOCK_L: tl.constexpr,
|
||||
BLOCK_C: tl.constexpr,
|
||||
):
|
||||
pid_l = tl.program_id(0)
|
||||
pid_c = tl.program_id(1)
|
||||
pid_b = tl.program_id(2)
|
||||
|
||||
l_offsets = pid_l * BLOCK_L + tl.arange(0, BLOCK_L)
|
||||
c_offsets = pid_c * BLOCK_C + tl.arange(0, BLOCK_C)
|
||||
|
||||
mask_l = l_offsets < L
|
||||
mask_c = c_offsets < C
|
||||
mask = mask_l[:, None] & mask_c[None, :]
|
||||
|
||||
x_off = (
|
||||
pid_b * stride_x_b
|
||||
+ l_offsets[:, None] * stride_x_l
|
||||
+ c_offsets[None, :] * stride_x_c
|
||||
)
|
||||
x = tl.load(x_ptr + x_off, mask=mask, other=0)
|
||||
|
||||
idx_off = pid_b * stride_i_b + l_offsets * stride_i_l
|
||||
idx = tl.load(index_ptr + idx_off, mask=mask_l, other=0).to(tl.int1)[:, None]
|
||||
|
||||
s0_off = pid_b * stride_s0_b + c_offsets[None, :] * stride_s0_c
|
||||
sc0_off = pid_b * stride_sc0_b + c_offsets[None, :] * stride_sc0_c
|
||||
g0_off = pid_b * stride_g0_b + c_offsets[None, :] * stride_g0_c
|
||||
s1_off = pid_b * stride_s1_b + c_offsets[None, :] * stride_s1_c
|
||||
sc1_off = pid_b * stride_sc1_b + c_offsets[None, :] * stride_sc1_c
|
||||
g1_off = pid_b * stride_g1_b + c_offsets[None, :] * stride_g1_c
|
||||
|
||||
shift0 = tl.load(shift0_ptr + s0_off, mask=mask_c[None, :], other=0)
|
||||
scale0 = tl.load(scale0_ptr + sc0_off, mask=mask_c[None, :], other=0)
|
||||
gate0 = tl.load(gate0_ptr + g0_off, mask=mask_c[None, :], other=0)
|
||||
shift1 = tl.load(shift1_ptr + s1_off, mask=mask_c[None, :], other=0)
|
||||
scale1 = tl.load(scale1_ptr + sc1_off, mask=mask_c[None, :], other=0)
|
||||
gate1 = tl.load(gate1_ptr + g1_off, mask=mask_c[None, :], other=0)
|
||||
|
||||
shift = tl.where(idx, shift1, shift0)
|
||||
scale = tl.where(idx, scale1, scale0)
|
||||
gate = tl.where(idx, gate1, gate0)
|
||||
|
||||
y = x * (1 + scale) + shift
|
||||
tl.store(y_ptr + x_off, y, mask=mask)
|
||||
|
||||
go_off = (
|
||||
pid_b * stride_go_b
|
||||
+ l_offsets[:, None] * stride_go_l
|
||||
+ c_offsets[None, :] * stride_go_c
|
||||
)
|
||||
tl.store(gate_out_ptr + go_off, gate, mask=mask)
|
||||
|
||||
|
||||
def fuse_scale_shift_kernel(
|
||||
x: torch.Tensor,
|
||||
scale: torch.Tensor,
|
||||
shift: torch.Tensor,
|
||||
scale_constant: float = 1.0,
|
||||
block_l: int = 128,
|
||||
block_c: int = 128,
|
||||
):
|
||||
assert x.is_cuda and scale.is_cuda
|
||||
assert x.is_contiguous()
|
||||
|
||||
B, L, C = x.shape
|
||||
output = torch.empty_like(x)
|
||||
|
||||
if scale.dim() == 4:
|
||||
# scale/shift: [B, F, 1, C]
|
||||
rows = B * L
|
||||
x_2d = x.view(rows, C)
|
||||
output_2d = output.view(rows, C)
|
||||
grid = lambda META: (rows, triton.cdiv(C, META["BLOCK_N"]))
|
||||
num_frames = scale.shape[1]
|
||||
assert (
|
||||
L % num_frames == 0
|
||||
), "seq_len must be divisible by num_frames for 4D scale/shift"
|
||||
frame_seqlen = L // num_frames
|
||||
|
||||
# Compact [B, F, C] without the singleton dim into [B*F, C]
|
||||
scale_reshaped = scale.squeeze(2).reshape(-1, C).contiguous()
|
||||
shift_reshaped = shift.squeeze(2).reshape(-1, C).contiguous()
|
||||
|
||||
_fused_scale_shift_4d_kernel[grid](
|
||||
output_2d,
|
||||
x_2d,
|
||||
scale_reshaped,
|
||||
shift_reshaped,
|
||||
scale_constant,
|
||||
rows,
|
||||
C,
|
||||
L,
|
||||
num_frames,
|
||||
frame_seqlen,
|
||||
)
|
||||
else:
|
||||
# 2D: [B, C] or [1, C] -> treat as [B, 1, C] and broadcast over L
|
||||
# 3D: [B, L, C] (or broadcastable variants like [B, 1, C], [1, L, C], [1, 1, C])
|
||||
# Also support scalar (0D or 1-element)
|
||||
if scale.dim() == 0 or (scale.dim() == 1 and scale.numel() == 1):
|
||||
scale_blc = scale.reshape(1)
|
||||
elif scale.dim() == 2:
|
||||
scale_blc = scale[:, None, :]
|
||||
elif scale.dim() == 3:
|
||||
scale_blc = scale
|
||||
else:
|
||||
raise ValueError("scale must be 0D/1D(1)/2D/3D or 4D")
|
||||
|
||||
if shift.dim() == 0 or (shift.dim() == 1 and shift.numel() == 1):
|
||||
shift_blc = shift.reshape(1)
|
||||
elif shift.dim() == 2:
|
||||
shift_blc = shift[:, None, :]
|
||||
elif shift.dim() == 3:
|
||||
shift_blc = shift
|
||||
else:
|
||||
# broadcast later via expand if possible
|
||||
shift_blc = shift
|
||||
|
||||
need_scale_scalar = scale_blc.dim() == 1 and scale_blc.numel() == 1
|
||||
need_shift_scalar = shift_blc.dim() == 1 and shift_blc.numel() == 1
|
||||
|
||||
if not need_scale_scalar:
|
||||
scale_exp = scale_blc.expand(B, L, C)
|
||||
s_sb, s_sl, s_sc = scale_exp.stride()
|
||||
else:
|
||||
s_sb = s_sl = s_sc = 0
|
||||
|
||||
if not need_shift_scalar:
|
||||
shift_exp = shift_blc.expand(B, L, C)
|
||||
sh_sb, sh_sl, sh_sc = shift_exp.stride()
|
||||
else:
|
||||
sh_sb = sh_sl = sh_sc = 0
|
||||
|
||||
# If both scalars and both zero, copy fast-path
|
||||
if need_scale_scalar and need_shift_scalar:
|
||||
if (scale_blc.abs().max() == 0) and (shift_blc.abs().max() == 0):
|
||||
output.copy_(x)
|
||||
return output
|
||||
|
||||
grid = (triton.cdiv(L, block_l), triton.cdiv(C, block_c), B)
|
||||
fuse_scale_shift_kernel_blc_opt[grid](
|
||||
x,
|
||||
shift_blc if need_shift_scalar else shift_exp,
|
||||
scale_blc if need_scale_scalar else scale_exp,
|
||||
scale_constant,
|
||||
output,
|
||||
B,
|
||||
L,
|
||||
C,
|
||||
x.stride(0),
|
||||
x.stride(1),
|
||||
x.stride(2),
|
||||
sh_sb,
|
||||
sh_sl,
|
||||
sh_sc,
|
||||
s_sb,
|
||||
s_sl,
|
||||
s_sc,
|
||||
SCALE_IS_SCALAR=need_scale_scalar,
|
||||
SHIFT_IS_SCALAR=need_shift_scalar,
|
||||
BLOCK_L=block_l,
|
||||
BLOCK_C=block_c,
|
||||
num_warps=4,
|
||||
num_stages=2,
|
||||
)
|
||||
return output
|
||||
|
||||
|
||||
def fuse_scale_shift_gate_select01_kernel(
|
||||
x: torch.Tensor,
|
||||
scale0: torch.Tensor,
|
||||
shift0: torch.Tensor,
|
||||
gate0: torch.Tensor,
|
||||
scale1: torch.Tensor,
|
||||
shift1: torch.Tensor,
|
||||
gate1: torch.Tensor,
|
||||
index: torch.Tensor,
|
||||
block_l: int = 128,
|
||||
block_c: int = 128,
|
||||
):
|
||||
assert x.is_contiguous()
|
||||
B, L, C = x.shape
|
||||
output = torch.empty_like(x)
|
||||
gate_out = torch.empty_like(x)
|
||||
|
||||
if (
|
||||
scale0.dim() != 2
|
||||
or shift0.dim() != 2
|
||||
or gate0.dim() != 2
|
||||
or scale1.dim() != 2
|
||||
or shift1.dim() != 2
|
||||
or gate1.dim() != 2
|
||||
):
|
||||
raise ValueError("scale0/shift0/gate0/scale1/shift1/gate1 must be 2D [B, C]")
|
||||
if index.dim() != 2:
|
||||
raise ValueError("index must be 2D [B, L]")
|
||||
|
||||
grid = (triton.cdiv(L, block_l), triton.cdiv(C, block_c), B)
|
||||
fuse_scale_shift_gate_select01_kernel_blc_opt[grid](
|
||||
x,
|
||||
shift0,
|
||||
scale0,
|
||||
gate0,
|
||||
shift1,
|
||||
scale1,
|
||||
gate1,
|
||||
index,
|
||||
output,
|
||||
gate_out,
|
||||
B,
|
||||
L,
|
||||
C,
|
||||
x.stride(0),
|
||||
x.stride(1),
|
||||
x.stride(2),
|
||||
shift0.stride(0),
|
||||
shift0.stride(1),
|
||||
scale0.stride(0),
|
||||
scale0.stride(1),
|
||||
gate0.stride(0),
|
||||
gate0.stride(1),
|
||||
shift1.stride(0),
|
||||
shift1.stride(1),
|
||||
scale1.stride(0),
|
||||
scale1.stride(1),
|
||||
gate1.stride(0),
|
||||
gate1.stride(1),
|
||||
index.stride(0),
|
||||
index.stride(1),
|
||||
gate_out.stride(0),
|
||||
gate_out.stride(1),
|
||||
gate_out.stride(2),
|
||||
BLOCK_L=block_l,
|
||||
BLOCK_C=block_c,
|
||||
num_warps=4,
|
||||
num_stages=2,
|
||||
)
|
||||
return output, gate_out
|
||||
|
||||
|
||||
if current_platform.is_npu():
|
||||
from .npu_fallback import fuse_scale_shift_native
|
||||
|
||||
fuse_scale_shift_kernel = fuse_scale_shift_native
|
||||
@@ -1,7 +1,7 @@
|
||||
import torch
|
||||
|
||||
from sglang.jit_kernel.diffusion.triton.scale_shift import fuse_scale_shift_kernel
|
||||
from sglang.multimodal_gen.runtime.layers.custom_op import CustomOp
|
||||
from sglang.multimodal_gen.runtime.layers.triton_ops import fuse_scale_shift_kernel
|
||||
|
||||
|
||||
class MulAdd(CustomOp):
|
||||
|
||||
@@ -20,6 +20,9 @@ if _is_cuda:
|
||||
if _is_npu:
|
||||
import torch_npu
|
||||
|
||||
from sglang.jit_kernel.diffusion.triton.norm import norm_infer, rms_norm_fn
|
||||
from sglang.jit_kernel.diffusion.triton.rmsnorm_onepass import triton_one_pass_rms_norm
|
||||
from sglang.jit_kernel.diffusion.triton.scale_shift import fuse_scale_shift_kernel
|
||||
from sglang.jit_kernel.norm import can_use_fused_inplace_qknorm, fused_inplace_qknorm
|
||||
from sglang.multimodal_gen.runtime.distributed.parallel_state import (
|
||||
get_tensor_model_parallel_rank,
|
||||
@@ -27,12 +30,6 @@ from sglang.multimodal_gen.runtime.distributed.parallel_state import (
|
||||
get_tp_group,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.layers.custom_op import CustomOp
|
||||
from sglang.multimodal_gen.runtime.layers.triton_ops import (
|
||||
fuse_scale_shift_kernel,
|
||||
norm_infer,
|
||||
rms_norm_fn,
|
||||
triton_one_pass_rms_norm,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.utils.common import get_bool_env_var
|
||||
|
||||
|
||||
|
||||
@@ -32,9 +32,9 @@ from typing import Any, Optional, Tuple
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.jit_kernel.diffusion.triton.rotary import apply_rotary_embedding
|
||||
from sglang.multimodal_gen.runtime.distributed.parallel_state import get_sp_group
|
||||
from sglang.multimodal_gen.runtime.layers.custom_op import CustomOp
|
||||
from sglang.multimodal_gen.runtime.layers.triton_ops import apply_rotary_embedding
|
||||
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
@@ -15,6 +15,9 @@ from diffusers.models.embeddings import TimestepEmbedding, Timesteps
|
||||
from diffusers.models.modeling_outputs import Transformer2DModelOutput
|
||||
from diffusers.models.normalization import AdaLayerNormContinuous
|
||||
|
||||
from sglang.jit_kernel.diffusion.triton.scale_shift import (
|
||||
fuse_scale_shift_gate_select01_kernel,
|
||||
)
|
||||
from sglang.multimodal_gen.configs.models.dits.qwenimage import QwenImageDitConfig
|
||||
from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
|
||||
from sglang.multimodal_gen.runtime.layers.attention import USPAttention
|
||||
@@ -41,9 +44,6 @@ from sglang.multimodal_gen.runtime.layers.quantization.configs.nunchaku_config i
|
||||
from sglang.multimodal_gen.runtime.layers.rotary_embedding import (
|
||||
apply_flashinfer_rope_qk_inplace,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.layers.triton_ops import (
|
||||
fuse_scale_shift_gate_select01_kernel,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT
|
||||
from sglang.multimodal_gen.runtime.platforms import (
|
||||
AttentionBackendEnum,
|
||||
|
||||
Reference in New Issue
Block a user