550 lines
17 KiB
Python
550 lines
17 KiB
Python
import torch
|
|
import triton
|
|
import triton.language as tl
|
|
|
|
from sglang.kernels.jit.utils import is_arch_support_pdl
|
|
from sglang.srt.utils import is_hip
|
|
|
|
_is_hip = is_hip()
|
|
|
|
|
|
rmsnorm_autotune = triton.autotune(
|
|
configs=[
|
|
triton.Config(kwargs={"BLOCK_SIZE": 1024}, num_warps=4, num_stages=1),
|
|
triton.Config(kwargs={"BLOCK_SIZE": 1024}, num_warps=8, num_stages=1),
|
|
triton.Config(kwargs={"BLOCK_SIZE": 1024}, num_warps=16, num_stages=1),
|
|
triton.Config(kwargs={"BLOCK_SIZE": 1024}, num_warps=4),
|
|
triton.Config(kwargs={"BLOCK_SIZE": 1024}, num_warps=8),
|
|
triton.Config(kwargs={"BLOCK_SIZE": 1024}, num_warps=16),
|
|
triton.Config(kwargs={"BLOCK_SIZE": 1024}, num_warps=4, num_stages=4),
|
|
triton.Config(kwargs={"BLOCK_SIZE": 1024}, num_warps=8, num_stages=4),
|
|
triton.Config(kwargs={"BLOCK_SIZE": 1024}, num_warps=16, num_stages=4),
|
|
triton.Config(kwargs={"BLOCK_SIZE": 1024}, num_warps=8, num_stages=8),
|
|
triton.Config(kwargs={"BLOCK_SIZE": 1024}, num_warps=16, num_stages=8),
|
|
triton.Config(kwargs={"BLOCK_SIZE": 2048}, num_warps=8),
|
|
triton.Config(kwargs={"BLOCK_SIZE": 2048}, num_warps=16),
|
|
triton.Config(kwargs={"BLOCK_SIZE": 2048}, num_warps=8, num_stages=4),
|
|
triton.Config(kwargs={"BLOCK_SIZE": 2048}, num_warps=16, num_stages=4),
|
|
triton.Config(kwargs={"BLOCK_SIZE": 4096}, num_warps=8),
|
|
triton.Config(kwargs={"BLOCK_SIZE": 4096}, num_warps=16),
|
|
triton.Config(kwargs={"BLOCK_SIZE": 8192}, num_warps=8),
|
|
triton.Config(kwargs={"BLOCK_SIZE": 8192}, num_warps=16),
|
|
triton.Config(kwargs={"BLOCK_SIZE": 8192}, num_warps=32),
|
|
triton.Config(kwargs={"BLOCK_SIZE": 8192}, num_warps=8, num_stages=1),
|
|
triton.Config(kwargs={"BLOCK_SIZE": 8192}, num_warps=16, num_stages=1),
|
|
triton.Config(kwargs={"BLOCK_SIZE": 8192}, num_warps=32, num_stages=1),
|
|
triton.Config(kwargs={"BLOCK_SIZE": 8192}, num_warps=8, num_stages=4),
|
|
triton.Config(kwargs={"BLOCK_SIZE": 8192}, num_warps=16, num_stages=4),
|
|
triton.Config(kwargs={"BLOCK_SIZE": 8192}, num_warps=32, num_stages=4),
|
|
triton.Config(kwargs={"BLOCK_SIZE": 16384}, num_warps=8),
|
|
triton.Config(kwargs={"BLOCK_SIZE": 16384}, num_warps=16),
|
|
triton.Config(kwargs={"BLOCK_SIZE": 16384}, num_warps=32),
|
|
triton.Config(kwargs={"BLOCK_SIZE": 16384}, num_warps=8, num_stages=1),
|
|
triton.Config(kwargs={"BLOCK_SIZE": 16384}, num_warps=16, num_stages=1),
|
|
triton.Config(kwargs={"BLOCK_SIZE": 16384}, num_warps=32, num_stages=1),
|
|
triton.Config(kwargs={"BLOCK_SIZE": 16384}, num_warps=8, num_stages=4),
|
|
triton.Config(kwargs={"BLOCK_SIZE": 16384}, num_warps=16, num_stages=4),
|
|
triton.Config(kwargs={"BLOCK_SIZE": 16384}, num_warps=32, num_stages=4),
|
|
],
|
|
key=["hidden_dim"],
|
|
)
|
|
|
|
|
|
@triton.jit
|
|
def fused_dual_residual_rmsnorm_kernel(
|
|
output_ptr,
|
|
mid_ptr,
|
|
activ_ptr,
|
|
residual_ptr,
|
|
weight1_ptr,
|
|
weight2_ptr,
|
|
eps: tl.constexpr,
|
|
hidden_dim: tl.constexpr,
|
|
BLOCK_SIZE: tl.constexpr,
|
|
):
|
|
pid = tl.program_id(axis=0)
|
|
input_start = pid * hidden_dim
|
|
|
|
offsets = tl.arange(0, BLOCK_SIZE)
|
|
mask = offsets < hidden_dim
|
|
|
|
a_ = tl.load(activ_ptr + input_start + offsets, mask=mask, other=0.0)
|
|
a = a_.to(tl.float32)
|
|
rms = tl.sqrt(tl.sum(a * a, axis=0) / hidden_dim + eps)
|
|
|
|
r = tl.load(residual_ptr + input_start + offsets, mask=mask, other=0.0)
|
|
w1_ = tl.load(weight1_ptr + offsets, mask=mask, other=0.0)
|
|
w1 = w1_.to(tl.float32)
|
|
|
|
a2r = r + (a / rms * w1).to(r.dtype)
|
|
tl.store(
|
|
mid_ptr + input_start + offsets,
|
|
a2r,
|
|
mask=mask,
|
|
)
|
|
|
|
a2r = a2r.to(tl.float32)
|
|
rms2 = tl.sqrt(tl.sum(a2r * a2r, axis=0) / hidden_dim + eps)
|
|
|
|
w2_ = tl.load(weight2_ptr + offsets, mask=mask, other=0.0)
|
|
w2 = w2_.to(tl.float32)
|
|
|
|
tl.store(
|
|
output_ptr + input_start + offsets,
|
|
a2r / rms2 * w2, # implicitly casts to output dtype here
|
|
mask=mask,
|
|
)
|
|
|
|
|
|
fused_dual_residual_rmsnorm_kernel_autotune = rmsnorm_autotune(
|
|
fused_dual_residual_rmsnorm_kernel
|
|
)
|
|
|
|
|
|
def fused_dual_residual_rmsnorm(x, residual, weight1, weight2, eps, autotune=False):
|
|
assert len(x.shape) == 2
|
|
assert (
|
|
x.shape == residual.shape and x.dtype == residual.dtype
|
|
), f"{x.shape=} {residual.shape=} {x.dtype=} {residual.dtype=}"
|
|
output, mid = torch.empty_like(x), torch.empty_like(x)
|
|
bs, hidden_dim = x.shape
|
|
if autotune:
|
|
fused_dual_residual_rmsnorm_kernel_autotune[(bs,)](
|
|
output, mid, x, residual, weight1, weight2, eps=eps, hidden_dim=hidden_dim
|
|
)
|
|
else:
|
|
max_warps = 16 if _is_hip else 32
|
|
config = {
|
|
"BLOCK_SIZE": triton.next_power_of_2(hidden_dim),
|
|
"num_warps": max(
|
|
min(triton.next_power_of_2(triton.cdiv(hidden_dim, 256)), max_warps), 4
|
|
),
|
|
}
|
|
|
|
fused_dual_residual_rmsnorm_kernel[(bs,)](
|
|
output,
|
|
mid,
|
|
x,
|
|
residual,
|
|
weight1,
|
|
weight2,
|
|
eps=eps,
|
|
hidden_dim=hidden_dim,
|
|
**config,
|
|
)
|
|
|
|
return output, mid
|
|
|
|
|
|
@triton.jit
|
|
def fused_rmsnorm_kernel(
|
|
output_ptr,
|
|
activ_ptr,
|
|
weight_ptr,
|
|
eps: tl.constexpr,
|
|
hidden_dim: tl.constexpr,
|
|
BLOCK_SIZE: tl.constexpr,
|
|
):
|
|
pid = tl.program_id(axis=0).to(tl.int64)
|
|
input_start = pid * hidden_dim
|
|
|
|
offsets = tl.arange(0, BLOCK_SIZE)
|
|
mask = offsets < hidden_dim
|
|
|
|
a_ = tl.load(activ_ptr + input_start + offsets, mask=mask, other=0.0)
|
|
a = a_.to(tl.float32)
|
|
rms = tl.sqrt(tl.sum(a * a, axis=0) / hidden_dim + eps)
|
|
|
|
w1_ = tl.load(weight_ptr + offsets, mask=mask, other=0.0)
|
|
w1 = w1_.to(tl.float32)
|
|
|
|
a_rms = a / rms * w1
|
|
|
|
tl.store(
|
|
output_ptr + input_start + offsets,
|
|
a_rms, # implicitly casts to output dtype here
|
|
mask=mask,
|
|
)
|
|
|
|
|
|
def fused_rmsnorm(x, weight, eps, autotune=False, inplace=False):
|
|
assert len(x.shape) == 2
|
|
if inplace:
|
|
output = x
|
|
else:
|
|
output = torch.empty_like(x)
|
|
bs, hidden_dim = x.shape
|
|
max_warps = 16 if _is_hip else 32
|
|
config = {
|
|
"BLOCK_SIZE": triton.next_power_of_2(hidden_dim),
|
|
"num_warps": max(
|
|
min(triton.next_power_of_2(triton.cdiv(hidden_dim, 256)), max_warps), 4
|
|
),
|
|
}
|
|
|
|
fused_rmsnorm_kernel[(bs,)](
|
|
output, x, weight, eps=eps, hidden_dim=hidden_dim, **config
|
|
)
|
|
return output
|
|
|
|
|
|
# gelu on first half of vector
|
|
@triton.jit
|
|
def gelu_and_mul_kernel(
|
|
out_hidden_states_ptr, # (bs, hidden_dim)
|
|
out_scales_ptr, # (bs,)
|
|
hidden_states_ptr, # (bs, hidden_dim * 2)
|
|
quant_max: tl.constexpr,
|
|
static_scale: tl.constexpr,
|
|
hidden_dim: tl.constexpr, # the output hidden_dim
|
|
BLOCK_SIZE: tl.constexpr,
|
|
):
|
|
pid = tl.program_id(axis=0)
|
|
|
|
input_start = pid * hidden_dim * 2
|
|
output_start = pid * hidden_dim
|
|
|
|
input1_offs = tl.arange(0, BLOCK_SIZE)
|
|
mask = tl.arange(0, BLOCK_SIZE) < hidden_dim # shared for input1, input3, output
|
|
input3_offs = hidden_dim + tl.arange(0, BLOCK_SIZE)
|
|
output_offs = tl.arange(0, BLOCK_SIZE)
|
|
|
|
x1 = tl.load(
|
|
hidden_states_ptr + input_start + input1_offs, mask=mask, other=0.0
|
|
).to(tl.float32)
|
|
x3 = tl.load(
|
|
hidden_states_ptr + input_start + input3_offs, mask=mask, other=0.0
|
|
).to(tl.float32)
|
|
|
|
# gelu
|
|
# cast down before mul to better match training?
|
|
gelu_x1 = 0.5 * (1.0 + tl.erf(x1 * 0.7071067811865475)) * x1
|
|
out = x3 * gelu_x1.to(hidden_states_ptr.dtype.element_ty)
|
|
|
|
if quant_max is not None:
|
|
raise NotImplementedError()
|
|
|
|
tl.store(out_hidden_states_ptr + output_start + output_offs, out, mask=mask)
|
|
|
|
|
|
def gelu_and_mul_triton(
|
|
hidden_states,
|
|
scales=None,
|
|
quantize=None, # dtype to quantize to
|
|
out=None,
|
|
):
|
|
bs, in_hidden_dim = hidden_states.shape
|
|
hidden_dim = in_hidden_dim // 2
|
|
|
|
if out is None:
|
|
out_hidden_states = torch.empty(
|
|
(bs, hidden_dim),
|
|
dtype=quantize or hidden_states.dtype,
|
|
device=hidden_states.device,
|
|
)
|
|
else:
|
|
assert out.shape == (bs, hidden_dim)
|
|
assert out.dtype == (quantize or hidden_states.dtype)
|
|
out_hidden_states = out
|
|
out_scales = None
|
|
static_scale = False
|
|
if quantize is not None:
|
|
if scales is None:
|
|
out_scales = torch.empty(
|
|
(bs,), dtype=torch.float32, device=hidden_states.device
|
|
)
|
|
else:
|
|
out_scales = scales
|
|
static_scale = True
|
|
|
|
max_warps = 16 if _is_hip else 32
|
|
config = {
|
|
# 8 ele per thread (not tuned)
|
|
"num_warps": max(
|
|
min(triton.next_power_of_2(triton.cdiv(hidden_dim, 8 * 32)), max_warps), 4
|
|
),
|
|
}
|
|
|
|
gelu_and_mul_kernel[(bs,)](
|
|
out_hidden_states,
|
|
out_scales,
|
|
hidden_states,
|
|
quant_max=torch.finfo(quantize).max if quantize is not None else None,
|
|
static_scale=static_scale,
|
|
hidden_dim=hidden_dim,
|
|
BLOCK_SIZE=triton.next_power_of_2(hidden_dim),
|
|
**config,
|
|
)
|
|
|
|
if quantize is not None:
|
|
return out_hidden_states, out_scales
|
|
else:
|
|
return out_hidden_states, None
|
|
|
|
|
|
# silu on first half of vector
|
|
@triton.jit
|
|
def silu_and_mul_kernel(
|
|
out_hidden_states_ptr, # (bs, hidden_dim)
|
|
out_scales_ptr, # (bs,)
|
|
hidden_states_ptr, # (bs, hidden_dim * 2)
|
|
quant_max: tl.constexpr,
|
|
static_scale: tl.constexpr,
|
|
hidden_dim: tl.constexpr, # the output hidden_dim
|
|
BLOCK_SIZE: tl.constexpr,
|
|
):
|
|
pid = tl.program_id(axis=0)
|
|
|
|
input_start = pid * hidden_dim * 2
|
|
output_start = pid * hidden_dim
|
|
|
|
input1_offs = tl.arange(0, BLOCK_SIZE)
|
|
mask = tl.arange(0, BLOCK_SIZE) < hidden_dim # shared for input1, input3, output
|
|
input3_offs = hidden_dim + tl.arange(0, BLOCK_SIZE)
|
|
output_offs = tl.arange(0, BLOCK_SIZE)
|
|
|
|
x1 = tl.load(
|
|
hidden_states_ptr + input_start + input1_offs, mask=mask, other=0.0
|
|
).to(tl.float32)
|
|
x3 = tl.load(
|
|
hidden_states_ptr + input_start + input3_offs, mask=mask, other=0.0
|
|
).to(tl.float32)
|
|
|
|
# silu
|
|
# cast down before mul to better match training?
|
|
silu_x1 = x1 * tl.sigmoid(x1)
|
|
out = x3 * silu_x1.to(hidden_states_ptr.dtype.element_ty)
|
|
|
|
if quant_max is not None:
|
|
raise NotImplementedError()
|
|
|
|
tl.store(out_hidden_states_ptr + output_start + output_offs, out, mask=mask)
|
|
|
|
|
|
def silu_and_mul_triton(
|
|
hidden_states,
|
|
scales=None,
|
|
quantize=None, # dtype to quantize to
|
|
out=None,
|
|
):
|
|
bs, in_hidden_dim = hidden_states.shape
|
|
hidden_dim = in_hidden_dim // 2
|
|
|
|
if out is None:
|
|
out_hidden_states = torch.empty(
|
|
(bs, hidden_dim),
|
|
dtype=quantize or hidden_states.dtype,
|
|
device=hidden_states.device,
|
|
)
|
|
else:
|
|
assert out.shape == (bs, hidden_dim)
|
|
assert out.dtype == (quantize or hidden_states.dtype)
|
|
out_hidden_states = out
|
|
out_scales = None
|
|
static_scale = False
|
|
if quantize is not None:
|
|
if scales is None:
|
|
out_scales = torch.empty(
|
|
(bs,), dtype=torch.float32, device=hidden_states.device
|
|
)
|
|
else:
|
|
out_scales = scales
|
|
static_scale = True
|
|
|
|
max_warps = 16 if _is_hip else 32
|
|
config = {
|
|
# 8 ele per thread (not tuned)
|
|
"num_warps": max(
|
|
min(triton.next_power_of_2(triton.cdiv(hidden_dim, 8 * 32)), max_warps), 4
|
|
),
|
|
}
|
|
|
|
silu_and_mul_kernel[(bs,)](
|
|
out_hidden_states,
|
|
out_scales,
|
|
hidden_states,
|
|
quant_max=torch.finfo(quantize).max if quantize is not None else None,
|
|
static_scale=static_scale,
|
|
hidden_dim=hidden_dim,
|
|
BLOCK_SIZE=triton.next_power_of_2(hidden_dim),
|
|
**config,
|
|
)
|
|
|
|
if quantize is not None:
|
|
return out_hidden_states, out_scales
|
|
else:
|
|
return out_hidden_states, None
|
|
|
|
|
|
@triton.jit
|
|
def _fused_sigmoid_mul_kernel(
|
|
output_ptr,
|
|
attn_output_ptr,
|
|
gate_ptr,
|
|
gate_stride_row,
|
|
gate_stride_head,
|
|
hidden_dim: tl.constexpr,
|
|
HEAD_DIM: tl.constexpr,
|
|
BLOCK_H: tl.constexpr,
|
|
):
|
|
"""Fuse sigmoid(gate) * attn_output into a single kernel."""
|
|
pid_row = tl.program_id(0).to(tl.int64)
|
|
pid_block = tl.program_id(1)
|
|
|
|
offsets = pid_block * BLOCK_H + tl.arange(0, BLOCK_H)
|
|
mask = offsets < hidden_dim
|
|
head = offsets // HEAD_DIM
|
|
d = offsets - head * HEAD_DIM
|
|
|
|
attn_off = pid_row * hidden_dim + offsets
|
|
attn = tl.load(attn_output_ptr + attn_off, mask=mask, other=0.0).to(tl.float32)
|
|
|
|
gate_off = pid_row * gate_stride_row + head * gate_stride_head + d
|
|
g = tl.load(gate_ptr + gate_off, mask=mask, other=0.0).to(tl.float32)
|
|
|
|
result = attn * tl.sigmoid(g)
|
|
tl.store(output_ptr + attn_off, result, mask=mask)
|
|
|
|
|
|
def fused_sigmoid_mul(
|
|
attn_output: torch.Tensor,
|
|
gate: torch.Tensor,
|
|
inplace: bool = False,
|
|
) -> torch.Tensor:
|
|
"""
|
|
Fused sigmoid-mul for attention output gating.
|
|
|
|
Equivalent to: attn_output * sigmoid(gate)
|
|
|
|
The production Qwen3.5 path passes a 3D strided gate. A single hidden-block
|
|
Triton kernel handles both that path and flat contiguous inputs.
|
|
|
|
When inplace=True, writes result back to attn_output and returns it.
|
|
|
|
Supports strided gate: if gate is 3D (num_tokens, num_heads, head_dim)
|
|
and attn_output is 2D (num_tokens, hidden_dim), the kernel reads gate
|
|
via explicit strides without requiring a contiguous copy.
|
|
"""
|
|
if gate.ndim == 3 and attn_output.ndim == 2:
|
|
# Strided gate path: gate is 3D (num_tokens, num_heads, head_dim)
|
|
num_tokens, num_heads, head_dim = gate.shape
|
|
hidden_dim = num_heads * head_dim
|
|
assert attn_output.shape == (num_tokens, hidden_dim)
|
|
gate_stride_row = gate.stride(0)
|
|
gate_stride_head = gate.stride(1)
|
|
else:
|
|
# Flat path: both tensors have the same shape
|
|
assert (
|
|
attn_output.shape == gate.shape
|
|
), "attn_output and gate must have the same shape"
|
|
hidden_dim = attn_output.shape[-1]
|
|
num_tokens = attn_output.numel() // hidden_dim
|
|
head_dim = hidden_dim
|
|
gate_stride_row = hidden_dim
|
|
gate_stride_head = hidden_dim
|
|
|
|
out = attn_output if inplace else torch.empty_like(attn_output)
|
|
block_h = 1024 if num_tokens < 1024 else 2048
|
|
grid = (num_tokens, triton.cdiv(hidden_dim, block_h))
|
|
_fused_sigmoid_mul_kernel[grid](
|
|
out,
|
|
attn_output,
|
|
gate,
|
|
gate_stride_row,
|
|
gate_stride_head,
|
|
hidden_dim,
|
|
HEAD_DIM=head_dim,
|
|
BLOCK_H=block_h,
|
|
num_warps=4,
|
|
)
|
|
return out
|
|
|
|
|
|
@triton.jit
|
|
def _fused_gate_sigmoid_mul_add_kernel(
|
|
hidden_states_ptr, # [num_tokens, hidden_dim]
|
|
gate_weight_ptr, # [hidden_dim]
|
|
shared_output_ptr, # [num_tokens, hidden_dim]
|
|
final_hidden_states_ptr, # [num_tokens, hidden_dim]
|
|
hidden_dim: tl.constexpr,
|
|
BLOCK_SIZE: tl.constexpr,
|
|
USE_PDL: tl.constexpr = False,
|
|
):
|
|
pid = tl.program_id(axis=0).to(tl.int64)
|
|
row_offset = pid * hidden_dim
|
|
|
|
offsets = tl.arange(0, BLOCK_SIZE)
|
|
mask = offsets < hidden_dim
|
|
|
|
w = tl.load(gate_weight_ptr + offsets, mask=mask, other=0.0).to(tl.float32)
|
|
|
|
if USE_PDL:
|
|
tl.extra.cuda.gdc_wait()
|
|
|
|
h = tl.load(hidden_states_ptr + row_offset + offsets, mask=mask, other=0.0).to(
|
|
tl.float32
|
|
)
|
|
s = tl.load(shared_output_ptr + row_offset + offsets, mask=mask, other=0.0).to(
|
|
tl.float32
|
|
)
|
|
f = tl.load(
|
|
final_hidden_states_ptr + row_offset + offsets, mask=mask, other=0.0
|
|
).to(tl.float32)
|
|
|
|
if USE_PDL:
|
|
tl.extra.cuda.gdc_launch_dependents()
|
|
|
|
gate_val = tl.sigmoid(tl.sum(h * w, axis=0))
|
|
result = f + gate_val * s
|
|
|
|
tl.store(final_hidden_states_ptr + row_offset + offsets, result, mask=mask)
|
|
|
|
|
|
def fused_gate_sigmoid_mul_add(
|
|
hidden_states: torch.Tensor,
|
|
gate_weight: torch.Tensor,
|
|
shared_output: torch.Tensor,
|
|
final_hidden_states: torch.Tensor,
|
|
) -> None:
|
|
"""
|
|
Fused gate-sigmoid-mul-add for MoE shared expert gating.
|
|
|
|
Equivalent to:
|
|
gate = hidden_states @ gate_weight
|
|
final_hidden_states += sigmoid(gate).unsqueeze(1) * shared_output
|
|
"""
|
|
assert hidden_states.is_contiguous(), "hidden_states must be contiguous"
|
|
assert gate_weight.is_contiguous(), "gate_weight must be contiguous"
|
|
assert shared_output.is_contiguous(), "shared_output must be contiguous"
|
|
assert final_hidden_states.is_contiguous(), "final_hidden_states must be contiguous"
|
|
|
|
num_tokens, hidden_dim = hidden_states.shape
|
|
assert gate_weight.shape == (hidden_dim,)
|
|
assert shared_output.shape == (num_tokens, hidden_dim)
|
|
assert final_hidden_states.shape == (num_tokens, hidden_dim)
|
|
|
|
max_warps = 16 if _is_hip else 32
|
|
config = {
|
|
"BLOCK_SIZE": triton.next_power_of_2(hidden_dim),
|
|
"num_warps": max(
|
|
min(triton.next_power_of_2(triton.cdiv(hidden_dim, 256)), max_warps), 4
|
|
),
|
|
}
|
|
|
|
if num_tokens >= 1024:
|
|
config["num_warps"] = min(config["num_warps"], 8)
|
|
|
|
use_pdl = is_arch_support_pdl()
|
|
pdl_kwargs = {"launch_pdl": True} if use_pdl else {}
|
|
|
|
_fused_gate_sigmoid_mul_add_kernel[(num_tokens,)](
|
|
hidden_states,
|
|
gate_weight,
|
|
shared_output,
|
|
final_hidden_states,
|
|
hidden_dim=hidden_dim,
|
|
USE_PDL=use_pdl,
|
|
**config,
|
|
**pdl_kwargs,
|
|
)
|