[Qwen3-Next] Support gdn fused_rms_norm_gated (#19434)
Co-authored-by: luoyuan.luo <luoyuan.luo@antgroup.com>
This commit is contained in:
Yuan Luo
388
python/sglang/srt/layers/attention/fla/fused_norm_gate.py
Normal file
388
python/sglang/srt/layers/attention/fla/fused_norm_gate.py
Normal file
@@ -0,0 +1,388 @@
|
||||
# Adapt from https://github.com/fla-org/flash-linear-attention/blob/main/fla/modules/fused_norm_gate.py
|
||||
# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
|
||||
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import triton
|
||||
import triton.language as tl
|
||||
|
||||
from sglang.srt.utils import (
|
||||
cdiv,
|
||||
cpu_has_amx_support,
|
||||
is_cpu,
|
||||
is_npu,
|
||||
next_power_of_2,
|
||||
)
|
||||
|
||||
_is_npu = is_npu()
|
||||
_use_cpu = is_cpu() and cpu_has_amx_support()
|
||||
|
||||
# Maximum rows per Triton block for layernorm gated kernel
|
||||
MAX_ROWS_PER_BLOCK = 4
|
||||
|
||||
|
||||
@triton.jit
|
||||
def layer_norm_gated_fwd_kernel(
|
||||
x, # pointer to the input
|
||||
g, # pointer to the gate
|
||||
y, # pointer to the output
|
||||
w, # pointer to the weights
|
||||
b, # pointer to the biases
|
||||
residual, # pointer to the residual
|
||||
residual_out, # pointer to the residual
|
||||
mean, # pointer to the mean
|
||||
rstd, # pointer to the 1/std
|
||||
eps, # epsilon to avoid division by zero
|
||||
T, # number of rows in x
|
||||
D: tl.constexpr, # number of columns in x
|
||||
BT: tl.constexpr,
|
||||
BD: tl.constexpr,
|
||||
ACTIVATION: tl.constexpr,
|
||||
IS_RMS_NORM: tl.constexpr,
|
||||
STORE_RESIDUAL_OUT: tl.constexpr,
|
||||
HAS_RESIDUAL: tl.constexpr,
|
||||
HAS_WEIGHT: tl.constexpr,
|
||||
HAS_BIAS: tl.constexpr,
|
||||
):
|
||||
i_t = tl.program_id(0)
|
||||
|
||||
o_d = tl.arange(0, BD)
|
||||
m_d = o_d < D
|
||||
|
||||
p_x = tl.make_block_ptr(x, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0))
|
||||
b_x = tl.load(p_x, boundary_check=(0, 1)).to(tl.float32)
|
||||
if HAS_RESIDUAL:
|
||||
p_res = tl.make_block_ptr(
|
||||
residual, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0)
|
||||
)
|
||||
b_x += tl.load(p_res, boundary_check=(0, 1)).to(tl.float32)
|
||||
if STORE_RESIDUAL_OUT:
|
||||
p_res_out = tl.make_block_ptr(
|
||||
residual_out, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0)
|
||||
)
|
||||
tl.store(p_res_out, b_x.to(p_res_out.dtype.element_ty), boundary_check=(0, 1))
|
||||
if not IS_RMS_NORM:
|
||||
b_mean = tl.sum(b_x, axis=1) / D
|
||||
p_mean = tl.make_block_ptr(mean, (T,), (1,), (i_t * BT,), (BT,), (0,))
|
||||
tl.store(p_mean, b_mean.to(p_mean.dtype.element_ty), boundary_check=(0,))
|
||||
b_xbar = tl.where(m_d[None, :], b_x - b_mean[:, None], 0.0)
|
||||
b_var = tl.sum(b_xbar * b_xbar, axis=1) / D
|
||||
else:
|
||||
b_xbar = tl.where(m_d[None, :], b_x, 0.0)
|
||||
b_var = tl.sum(b_xbar * b_xbar, axis=1) / D
|
||||
b_rstd = 1 / tl.sqrt(b_var + eps)
|
||||
|
||||
p_rstd = tl.make_block_ptr(rstd, (T,), (1,), (i_t * BT,), (BT,), (0,))
|
||||
tl.store(p_rstd, b_rstd.to(p_rstd.dtype.element_ty), boundary_check=(0,))
|
||||
|
||||
if HAS_WEIGHT:
|
||||
b_w = tl.load(w + o_d, mask=m_d).to(tl.float32)
|
||||
if HAS_BIAS:
|
||||
b_b = tl.load(b + o_d, mask=m_d).to(tl.float32)
|
||||
b_x_hat = (
|
||||
(b_x - b_mean[:, None]) * b_rstd[:, None]
|
||||
if not IS_RMS_NORM
|
||||
else b_x * b_rstd[:, None]
|
||||
)
|
||||
b_y = b_x_hat * b_w[None, :] if HAS_WEIGHT else b_x_hat
|
||||
if HAS_BIAS:
|
||||
b_y = b_y + b_b[None, :]
|
||||
|
||||
# swish/sigmoid output gate
|
||||
p_g = tl.make_block_ptr(g, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0))
|
||||
b_g = tl.load(p_g, boundary_check=(0, 1)).to(tl.float32)
|
||||
if ACTIVATION == "swish" or ACTIVATION == "silu":
|
||||
b_y = b_y * b_g * tl.sigmoid(b_g)
|
||||
elif ACTIVATION == "sigmoid":
|
||||
b_y = b_y * tl.sigmoid(b_g)
|
||||
|
||||
# Write output
|
||||
p_y = tl.make_block_ptr(y, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0))
|
||||
tl.store(p_y, b_y.to(p_y.dtype.element_ty), boundary_check=(0, 1))
|
||||
|
||||
|
||||
@triton.jit
|
||||
def layer_norm_gated_fwd_kernel1(
|
||||
x, # pointer to the input
|
||||
g, # pointer to the gate
|
||||
y, # pointer to the output
|
||||
w, # pointer to the weights
|
||||
b, # pointer to the biases
|
||||
residual, # pointer to the residual
|
||||
residual_out, # pointer to the residual
|
||||
mean, # pointer to the mean
|
||||
rstd, # pointer to the 1/std
|
||||
eps, # epsilon to avoid division by zero
|
||||
D: tl.constexpr, # number of columns in x
|
||||
BD: tl.constexpr,
|
||||
ACTIVATION: tl.constexpr,
|
||||
IS_RMS_NORM: tl.constexpr,
|
||||
STORE_RESIDUAL_OUT: tl.constexpr,
|
||||
HAS_RESIDUAL: tl.constexpr,
|
||||
HAS_WEIGHT: tl.constexpr,
|
||||
HAS_BIAS: tl.constexpr,
|
||||
):
|
||||
i_t = tl.program_id(0)
|
||||
x += i_t * D
|
||||
y += i_t * D
|
||||
g += i_t * D
|
||||
if HAS_RESIDUAL:
|
||||
residual += i_t * D
|
||||
if STORE_RESIDUAL_OUT:
|
||||
residual_out += i_t * D
|
||||
|
||||
o_d = tl.arange(0, BD)
|
||||
m_d = o_d < D
|
||||
b_x = tl.load(x + o_d, mask=m_d, other=0.0).to(tl.float32)
|
||||
if HAS_RESIDUAL:
|
||||
b_x += tl.load(residual + o_d, mask=m_d, other=0.0).to(tl.float32)
|
||||
if STORE_RESIDUAL_OUT:
|
||||
tl.store(residual_out + o_d, b_x, mask=m_d)
|
||||
if not IS_RMS_NORM:
|
||||
b_mean = tl.sum(b_x, axis=0) / D
|
||||
tl.store(mean + i_t, b_mean)
|
||||
b_xbar = tl.where(m_d, b_x - b_mean, 0.0)
|
||||
b_var = tl.sum(b_xbar * b_xbar, axis=0) / D
|
||||
else:
|
||||
b_xbar = tl.where(m_d, b_x, 0.0)
|
||||
b_var = tl.sum(b_xbar * b_xbar, axis=0) / D
|
||||
b_rstd = 1 / tl.sqrt(b_var + eps)
|
||||
tl.store(rstd + i_t, b_rstd)
|
||||
|
||||
if HAS_WEIGHT:
|
||||
b_w = tl.load(w + o_d, mask=m_d).to(tl.float32)
|
||||
if HAS_BIAS:
|
||||
b_b = tl.load(b + o_d, mask=m_d).to(tl.float32)
|
||||
b_x_hat = (b_x - b_mean) * b_rstd if not IS_RMS_NORM else b_x * b_rstd
|
||||
b_y = b_x_hat * b_w if HAS_WEIGHT else b_x_hat
|
||||
if HAS_BIAS:
|
||||
b_y = b_y + b_b
|
||||
|
||||
# swish/sigmoid output gate
|
||||
b_g = tl.load(g + o_d, mask=m_d, other=0.0).to(tl.float32)
|
||||
if ACTIVATION == "swish" or ACTIVATION == "silu":
|
||||
b_y = b_y * b_g * tl.sigmoid(b_g)
|
||||
elif ACTIVATION == "sigmoid":
|
||||
b_y = b_y * tl.sigmoid(b_g)
|
||||
|
||||
# Write output
|
||||
tl.store(y + o_d, b_y, mask=m_d)
|
||||
|
||||
|
||||
def layer_norm_gated_fwd(
|
||||
x: torch.Tensor,
|
||||
g: torch.Tensor,
|
||||
weight: torch.Tensor,
|
||||
bias: torch.Tensor,
|
||||
activation: str = "swish",
|
||||
eps: float = 1e-5,
|
||||
residual: torch.Tensor = None,
|
||||
out_dtype: torch.dtype = None,
|
||||
residual_dtype: torch.dtype = None,
|
||||
is_rms_norm: bool = False,
|
||||
):
|
||||
if residual is not None:
|
||||
residual_dtype = residual.dtype
|
||||
T, D = x.shape
|
||||
if residual is not None:
|
||||
assert residual.shape == (T, D)
|
||||
if weight is not None:
|
||||
assert weight.shape == (D,)
|
||||
if bias is not None:
|
||||
assert bias.shape == (D,)
|
||||
# allocate output
|
||||
y = x if out_dtype is None else torch.empty_like(x, dtype=out_dtype)
|
||||
if residual is not None or (
|
||||
residual_dtype is not None and residual_dtype != x.dtype
|
||||
):
|
||||
residual_out = torch.empty(T, D, device=x.device, dtype=residual_dtype)
|
||||
else:
|
||||
residual_out = None
|
||||
mean = (
|
||||
torch.empty((T,), dtype=torch.float, device=x.device)
|
||||
if not is_rms_norm
|
||||
else None
|
||||
)
|
||||
rstd = torch.empty((T,), dtype=torch.float, device=x.device)
|
||||
# Less than 64KB per feature: enqueue fused kernel
|
||||
MAX_FUSED_SIZE = 65536 // x.element_size()
|
||||
BD = min(MAX_FUSED_SIZE, next_power_of_2(D))
|
||||
if D > BD:
|
||||
raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.")
|
||||
# heuristics for number of warps
|
||||
|
||||
if D <= 512:
|
||||
BT = 32
|
||||
layer_norm_gated_fwd_kernel[(cdiv(T, BT),)](
|
||||
x=x,
|
||||
g=g,
|
||||
y=y,
|
||||
w=weight,
|
||||
b=bias,
|
||||
residual=residual,
|
||||
residual_out=residual_out,
|
||||
mean=mean,
|
||||
rstd=rstd,
|
||||
eps=eps,
|
||||
T=T,
|
||||
D=D,
|
||||
BD=BD,
|
||||
BT=BT,
|
||||
ACTIVATION=activation,
|
||||
IS_RMS_NORM=is_rms_norm,
|
||||
STORE_RESIDUAL_OUT=residual_out is not None,
|
||||
HAS_RESIDUAL=residual is not None,
|
||||
HAS_WEIGHT=weight is not None,
|
||||
HAS_BIAS=bias is not None,
|
||||
num_warps=4,
|
||||
)
|
||||
else:
|
||||
layer_norm_gated_fwd_kernel1[(T,)](
|
||||
x=x,
|
||||
g=g,
|
||||
y=y,
|
||||
w=weight,
|
||||
b=bias,
|
||||
residual=residual,
|
||||
residual_out=residual_out,
|
||||
mean=mean,
|
||||
rstd=rstd,
|
||||
eps=eps,
|
||||
D=D,
|
||||
BD=BD,
|
||||
ACTIVATION=activation,
|
||||
IS_RMS_NORM=is_rms_norm,
|
||||
STORE_RESIDUAL_OUT=residual_out is not None,
|
||||
HAS_RESIDUAL=residual is not None,
|
||||
HAS_WEIGHT=weight is not None,
|
||||
HAS_BIAS=bias is not None,
|
||||
num_warps=4,
|
||||
)
|
||||
# residual_out is None if residual is None and residual_dtype == input_dtype
|
||||
return y, mean, rstd, residual_out if residual_out is not None else x
|
||||
|
||||
|
||||
class LayerNormGatedFunction(torch.autograd.Function):
|
||||
@staticmethod
|
||||
def forward(
|
||||
ctx,
|
||||
x: torch.Tensor,
|
||||
g: torch.Tensor,
|
||||
weight: torch.Tensor,
|
||||
bias: torch.Tensor,
|
||||
activation: str,
|
||||
residual: torch.Tensor | None = None,
|
||||
eps: float = 1e-6,
|
||||
prenorm: bool = False,
|
||||
residual_in_fp32: bool = False,
|
||||
is_rms_norm: bool = False,
|
||||
):
|
||||
x_shape_og = x.shape
|
||||
g_shape_og = g.shape
|
||||
# reshape input data into 2D tensor
|
||||
x = x.reshape(-1, x.shape[-1])
|
||||
g = g.reshape(-1, g.shape[-1])
|
||||
if residual is not None:
|
||||
assert residual.shape == x_shape_og
|
||||
residual = residual.reshape(-1, residual.shape[-1])
|
||||
residual_dtype = (
|
||||
residual.dtype
|
||||
if residual is not None
|
||||
else (torch.float if residual_in_fp32 else None)
|
||||
)
|
||||
y, mean, rstd, residual_out = layer_norm_gated_fwd(
|
||||
x=x,
|
||||
g=g,
|
||||
weight=weight,
|
||||
bias=bias,
|
||||
activation=activation,
|
||||
eps=eps,
|
||||
residual=residual,
|
||||
residual_dtype=residual_dtype,
|
||||
is_rms_norm=is_rms_norm,
|
||||
)
|
||||
ctx.save_for_backward(residual_out, g, weight, bias, mean, rstd)
|
||||
ctx.x_shape_og = x_shape_og
|
||||
ctx.g_shape_og = g_shape_og
|
||||
ctx.activation = activation
|
||||
ctx.eps = eps
|
||||
ctx.is_rms_norm = is_rms_norm
|
||||
ctx.has_residual = residual is not None
|
||||
ctx.prenorm = prenorm
|
||||
ctx.x_dtype = x.dtype
|
||||
y = y.reshape(x_shape_og)
|
||||
return y if not prenorm else (y, residual_out.reshape(x_shape_og))
|
||||
|
||||
|
||||
def rms_norm_gated(
|
||||
x: torch.Tensor,
|
||||
g: torch.Tensor,
|
||||
weight: torch.Tensor,
|
||||
bias: torch.Tensor,
|
||||
activation: str = "swish",
|
||||
residual: torch.Tensor | None = None,
|
||||
prenorm: bool = False,
|
||||
residual_in_fp32: bool = False,
|
||||
eps: float = 1e-6,
|
||||
):
|
||||
return LayerNormGatedFunction.apply(
|
||||
x,
|
||||
g,
|
||||
weight,
|
||||
bias,
|
||||
activation,
|
||||
residual,
|
||||
eps,
|
||||
prenorm,
|
||||
residual_in_fp32,
|
||||
True,
|
||||
)
|
||||
|
||||
|
||||
class FusedRMSNormGated(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
hidden_size: int,
|
||||
elementwise_affine: bool = True,
|
||||
eps: float = 1e-5,
|
||||
activation: str = "swish",
|
||||
device: torch.device | None = None,
|
||||
dtype: torch.dtype | None = None,
|
||||
) -> None:
|
||||
factory_kwargs = {"device": device, "dtype": dtype}
|
||||
super().__init__()
|
||||
|
||||
self.hidden_size = hidden_size
|
||||
self.elementwise_affine = elementwise_affine
|
||||
self.eps = eps
|
||||
self.activation = activation
|
||||
|
||||
if self.activation not in ["swish", "silu", "sigmoid"]:
|
||||
raise ValueError(f"Unsupported activation: {self.activation}")
|
||||
|
||||
if elementwise_affine:
|
||||
self.weight = nn.Parameter(torch.empty(hidden_size, **factory_kwargs))
|
||||
else:
|
||||
self.register_parameter("weight", None)
|
||||
self.register_parameter("bias", None)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x: torch.Tensor,
|
||||
g: torch.Tensor,
|
||||
residual: torch.Tensor | None = None,
|
||||
prenorm: bool = False,
|
||||
residual_in_fp32: bool = False,
|
||||
) -> torch.Tensor:
|
||||
return rms_norm_gated(
|
||||
x,
|
||||
g,
|
||||
self.weight,
|
||||
self.bias,
|
||||
self.activation,
|
||||
residual=residual,
|
||||
eps=self.eps,
|
||||
prenorm=prenorm,
|
||||
residual_in_fp32=residual_in_fp32,
|
||||
)
|
||||
@@ -5,12 +5,12 @@
|
||||
# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import triton
|
||||
import triton.language as tl
|
||||
|
||||
from sglang.srt.layers.attention.fla.chunk_delta_h import chunk_gated_delta_rule_fwd_h
|
||||
from sglang.srt.layers.attention.fla.cumsum import chunk_local_cumsum
|
||||
from sglang.srt.layers.attention.fla.fused_norm_gate import layer_norm_gated_fwd
|
||||
from sglang.srt.layers.attention.fla.fused_recurrent import (
|
||||
fused_recurrent_gated_delta_rule_fwd_kernel,
|
||||
)
|
||||
@@ -155,247 +155,6 @@ def fused_recurrent_kda(
|
||||
return o, final_state
|
||||
|
||||
|
||||
@triton.jit
|
||||
def layer_norm_gated_fwd_kernel(
|
||||
x, # pointer to the input
|
||||
g, # pointer to the gate
|
||||
y, # pointer to the output
|
||||
w, # pointer to the weights
|
||||
b, # pointer to the biases
|
||||
residual, # pointer to the residual
|
||||
residual_out, # pointer to the residual
|
||||
mean, # pointer to the mean
|
||||
rstd, # pointer to the 1/std
|
||||
eps, # epsilon to avoid division by zero
|
||||
T, # number of rows in x
|
||||
D: tl.constexpr, # number of columns in x
|
||||
BT: tl.constexpr,
|
||||
BD: tl.constexpr,
|
||||
ACTIVATION: tl.constexpr,
|
||||
IS_RMS_NORM: tl.constexpr,
|
||||
STORE_RESIDUAL_OUT: tl.constexpr,
|
||||
HAS_RESIDUAL: tl.constexpr,
|
||||
HAS_WEIGHT: tl.constexpr,
|
||||
HAS_BIAS: tl.constexpr,
|
||||
):
|
||||
i_t = tl.program_id(0)
|
||||
|
||||
o_d = tl.arange(0, BD)
|
||||
m_d = o_d < D
|
||||
|
||||
p_x = tl.make_block_ptr(x, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0))
|
||||
b_x = tl.load(p_x, boundary_check=(0, 1)).to(tl.float32)
|
||||
if HAS_RESIDUAL:
|
||||
p_res = tl.make_block_ptr(
|
||||
residual, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0)
|
||||
)
|
||||
b_x += tl.load(p_res, boundary_check=(0, 1)).to(tl.float32)
|
||||
if STORE_RESIDUAL_OUT:
|
||||
p_res_out = tl.make_block_ptr(
|
||||
residual_out, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0)
|
||||
)
|
||||
tl.store(p_res_out, b_x.to(p_res_out.dtype.element_ty), boundary_check=(0, 1))
|
||||
if not IS_RMS_NORM:
|
||||
b_mean = tl.sum(b_x, axis=1) / D
|
||||
p_mean = tl.make_block_ptr(mean, (T,), (1,), (i_t * BT,), (BT,), (0,))
|
||||
tl.store(p_mean, b_mean.to(p_mean.dtype.element_ty), boundary_check=(0,))
|
||||
b_xbar = tl.where(m_d[None, :], b_x - b_mean[:, None], 0.0)
|
||||
b_var = tl.sum(b_xbar * b_xbar, axis=1) / D
|
||||
else:
|
||||
b_xbar = tl.where(m_d[None, :], b_x, 0.0)
|
||||
b_var = tl.sum(b_xbar * b_xbar, axis=1) / D
|
||||
b_rstd = 1 / tl.sqrt(b_var + eps)
|
||||
|
||||
p_rstd = tl.make_block_ptr(rstd, (T,), (1,), (i_t * BT,), (BT,), (0,))
|
||||
tl.store(p_rstd, b_rstd.to(p_rstd.dtype.element_ty), boundary_check=(0,))
|
||||
|
||||
if HAS_WEIGHT:
|
||||
b_w = tl.load(w + o_d, mask=m_d).to(tl.float32)
|
||||
if HAS_BIAS:
|
||||
b_b = tl.load(b + o_d, mask=m_d).to(tl.float32)
|
||||
b_x_hat = (
|
||||
(b_x - b_mean[:, None]) * b_rstd[:, None]
|
||||
if not IS_RMS_NORM
|
||||
else b_x * b_rstd[:, None]
|
||||
)
|
||||
b_y = b_x_hat * b_w[None, :] if HAS_WEIGHT else b_x_hat
|
||||
if HAS_BIAS:
|
||||
b_y = b_y + b_b[None, :]
|
||||
|
||||
# swish/sigmoid output gate
|
||||
p_g = tl.make_block_ptr(g, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0))
|
||||
b_g = tl.load(p_g, boundary_check=(0, 1)).to(tl.float32)
|
||||
if ACTIVATION == "swish" or ACTIVATION == "silu":
|
||||
b_y = b_y * b_g * tl.sigmoid(b_g)
|
||||
elif ACTIVATION == "sigmoid":
|
||||
b_y = b_y * tl.sigmoid(b_g)
|
||||
|
||||
# Write output
|
||||
p_y = tl.make_block_ptr(y, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0))
|
||||
tl.store(p_y, b_y.to(p_y.dtype.element_ty), boundary_check=(0, 1))
|
||||
|
||||
|
||||
@triton.jit
|
||||
def layer_norm_gated_fwd_kernel1(
|
||||
x, # pointer to the input
|
||||
g, # pointer to the gate
|
||||
y, # pointer to the output
|
||||
w, # pointer to the weights
|
||||
b, # pointer to the biases
|
||||
residual, # pointer to the residual
|
||||
residual_out, # pointer to the residual
|
||||
mean, # pointer to the mean
|
||||
rstd, # pointer to the 1/std
|
||||
eps, # epsilon to avoid division by zero
|
||||
D: tl.constexpr, # number of columns in x
|
||||
BD: tl.constexpr,
|
||||
ACTIVATION: tl.constexpr,
|
||||
IS_RMS_NORM: tl.constexpr,
|
||||
STORE_RESIDUAL_OUT: tl.constexpr,
|
||||
HAS_RESIDUAL: tl.constexpr,
|
||||
HAS_WEIGHT: tl.constexpr,
|
||||
HAS_BIAS: tl.constexpr,
|
||||
):
|
||||
i_t = tl.program_id(0)
|
||||
x += i_t * D
|
||||
y += i_t * D
|
||||
g += i_t * D
|
||||
if HAS_RESIDUAL:
|
||||
residual += i_t * D
|
||||
if STORE_RESIDUAL_OUT:
|
||||
residual_out += i_t * D
|
||||
|
||||
o_d = tl.arange(0, BD)
|
||||
m_d = o_d < D
|
||||
b_x = tl.load(x + o_d, mask=m_d, other=0.0).to(tl.float32)
|
||||
if HAS_RESIDUAL:
|
||||
b_x += tl.load(residual + o_d, mask=m_d, other=0.0).to(tl.float32)
|
||||
if STORE_RESIDUAL_OUT:
|
||||
tl.store(residual_out + o_d, b_x, mask=m_d)
|
||||
if not IS_RMS_NORM:
|
||||
b_mean = tl.sum(b_x, axis=0) / D
|
||||
tl.store(mean + i_t, b_mean)
|
||||
b_xbar = tl.where(m_d, b_x - b_mean, 0.0)
|
||||
b_var = tl.sum(b_xbar * b_xbar, axis=0) / D
|
||||
else:
|
||||
b_xbar = tl.where(m_d, b_x, 0.0)
|
||||
b_var = tl.sum(b_xbar * b_xbar, axis=0) / D
|
||||
b_rstd = 1 / tl.sqrt(b_var + eps)
|
||||
tl.store(rstd + i_t, b_rstd)
|
||||
|
||||
if HAS_WEIGHT:
|
||||
b_w = tl.load(w + o_d, mask=m_d).to(tl.float32)
|
||||
if HAS_BIAS:
|
||||
b_b = tl.load(b + o_d, mask=m_d).to(tl.float32)
|
||||
b_x_hat = (b_x - b_mean) * b_rstd if not IS_RMS_NORM else b_x * b_rstd
|
||||
b_y = b_x_hat * b_w if HAS_WEIGHT else b_x_hat
|
||||
if HAS_BIAS:
|
||||
b_y = b_y + b_b
|
||||
|
||||
# swish/sigmoid output gate
|
||||
b_g = tl.load(g + o_d, mask=m_d, other=0.0).to(tl.float32)
|
||||
if ACTIVATION == "swish" or ACTIVATION == "silu":
|
||||
b_y = b_y * b_g * tl.sigmoid(b_g)
|
||||
elif ACTIVATION == "sigmoid":
|
||||
b_y = b_y * tl.sigmoid(b_g)
|
||||
|
||||
# Write output
|
||||
tl.store(y + o_d, b_y, mask=m_d)
|
||||
|
||||
|
||||
def layer_norm_gated_fwd(
|
||||
x: torch.Tensor,
|
||||
g: torch.Tensor,
|
||||
weight: torch.Tensor,
|
||||
bias: torch.Tensor,
|
||||
activation: str = "swish",
|
||||
eps: float = 1e-5,
|
||||
residual: torch.Tensor = None,
|
||||
out_dtype: torch.dtype = None,
|
||||
residual_dtype: torch.dtype = None,
|
||||
is_rms_norm: bool = False,
|
||||
):
|
||||
if residual is not None:
|
||||
residual_dtype = residual.dtype
|
||||
T, D = x.shape
|
||||
if residual is not None:
|
||||
assert residual.shape == (T, D)
|
||||
if weight is not None:
|
||||
assert weight.shape == (D,)
|
||||
if bias is not None:
|
||||
assert bias.shape == (D,)
|
||||
# allocate output
|
||||
y = x if out_dtype is None else torch.empty_like(x, dtype=out_dtype)
|
||||
if residual is not None or (
|
||||
residual_dtype is not None and residual_dtype != x.dtype
|
||||
):
|
||||
residual_out = torch.empty(T, D, device=x.device, dtype=residual_dtype)
|
||||
else:
|
||||
residual_out = None
|
||||
mean = (
|
||||
torch.empty((T,), dtype=torch.float, device=x.device)
|
||||
if not is_rms_norm
|
||||
else None
|
||||
)
|
||||
rstd = torch.empty((T,), dtype=torch.float, device=x.device)
|
||||
# Less than 64KB per feature: enqueue fused kernel
|
||||
MAX_FUSED_SIZE = 65536 // x.element_size()
|
||||
BD = min(MAX_FUSED_SIZE, next_power_of_2(D))
|
||||
if D > BD:
|
||||
raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.")
|
||||
# heuristics for number of warps
|
||||
|
||||
if D <= 512:
|
||||
BT = 32
|
||||
layer_norm_gated_fwd_kernel[(cdiv(T, BT),)](
|
||||
x=x,
|
||||
g=g,
|
||||
y=y,
|
||||
w=weight,
|
||||
b=bias,
|
||||
residual=residual,
|
||||
residual_out=residual_out,
|
||||
mean=mean,
|
||||
rstd=rstd,
|
||||
eps=eps,
|
||||
T=T,
|
||||
D=D,
|
||||
BD=BD,
|
||||
BT=BT,
|
||||
ACTIVATION=activation,
|
||||
IS_RMS_NORM=is_rms_norm,
|
||||
STORE_RESIDUAL_OUT=residual_out is not None,
|
||||
HAS_RESIDUAL=residual is not None,
|
||||
HAS_WEIGHT=weight is not None,
|
||||
HAS_BIAS=bias is not None,
|
||||
num_warps=4,
|
||||
)
|
||||
else:
|
||||
layer_norm_gated_fwd_kernel1[(T,)](
|
||||
x=x,
|
||||
g=g,
|
||||
y=y,
|
||||
w=weight,
|
||||
b=bias,
|
||||
residual=residual,
|
||||
residual_out=residual_out,
|
||||
mean=mean,
|
||||
rstd=rstd,
|
||||
eps=eps,
|
||||
D=D,
|
||||
BD=BD,
|
||||
ACTIVATION=activation,
|
||||
IS_RMS_NORM=is_rms_norm,
|
||||
STORE_RESIDUAL_OUT=residual_out is not None,
|
||||
HAS_RESIDUAL=residual is not None,
|
||||
HAS_WEIGHT=weight is not None,
|
||||
HAS_BIAS=bias is not None,
|
||||
num_warps=4,
|
||||
)
|
||||
# residual_out is None if residual is None and residual_dtype == input_dtype
|
||||
return y, mean, rstd, residual_out if residual_out is not None else x
|
||||
|
||||
|
||||
def rms_norm_gated(
|
||||
x: torch.Tensor,
|
||||
g: torch.Tensor,
|
||||
@@ -434,54 +193,6 @@ def rms_norm_gated(
|
||||
return y if not prenorm else (y, residual_out.reshape(x_shape_og))
|
||||
|
||||
|
||||
class FusedRMSNormGated(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
hidden_size: int,
|
||||
elementwise_affine: bool = True,
|
||||
eps: float = 1e-5,
|
||||
activation: str = "swish",
|
||||
device: torch.device | None = None,
|
||||
dtype: torch.dtype | None = None,
|
||||
) -> None:
|
||||
factory_kwargs = {"device": device, "dtype": dtype}
|
||||
super().__init__()
|
||||
|
||||
self.hidden_size = hidden_size
|
||||
self.elementwise_affine = elementwise_affine
|
||||
self.eps = eps
|
||||
self.activation = activation
|
||||
|
||||
if self.activation not in ["swish", "silu", "sigmoid"]:
|
||||
raise ValueError(f"Unsupported activation: {self.activation}")
|
||||
|
||||
if elementwise_affine:
|
||||
self.weight = nn.Parameter(torch.empty(hidden_size, **factory_kwargs))
|
||||
else:
|
||||
self.register_parameter("weight", None)
|
||||
self.register_parameter("bias", None)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x: torch.Tensor,
|
||||
g: torch.Tensor,
|
||||
residual: torch.Tensor | None = None,
|
||||
prenorm: bool = False,
|
||||
residual_in_fp32: bool = False,
|
||||
) -> torch.Tensor:
|
||||
return rms_norm_gated(
|
||||
x,
|
||||
g,
|
||||
self.weight,
|
||||
self.bias,
|
||||
self.activation,
|
||||
residual=residual,
|
||||
eps=self.eps,
|
||||
prenorm=prenorm,
|
||||
residual_in_fp32=residual_in_fp32,
|
||||
)
|
||||
|
||||
|
||||
@triton.autotune(
|
||||
configs=[
|
||||
triton.Config({"BK": BK}, num_warps=num_warps, num_stages=num_stages)
|
||||
|
||||
@@ -15,7 +15,8 @@ from sglang.srt.distributed import (
|
||||
tensor_model_parallel_all_reduce,
|
||||
)
|
||||
from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder
|
||||
from sglang.srt.layers.attention.fla.kda import FusedRMSNormGated, fused_kda_gate
|
||||
from sglang.srt.layers.attention.fla.fused_norm_gate import FusedRMSNormGated
|
||||
from sglang.srt.layers.attention.fla.kda import fused_kda_gate
|
||||
from sglang.srt.layers.dp_attention import get_attention_tp_rank, get_attention_tp_size
|
||||
from sglang.srt.layers.layernorm import RMSNorm
|
||||
from sglang.srt.layers.linear import (
|
||||
|
||||
@@ -53,6 +53,9 @@ from sglang.srt.utils import (
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
from sglang.srt.layers.attention.fla.fused_norm_gate import FusedRMSNormGated
|
||||
|
||||
_is_cuda = is_cuda()
|
||||
_is_npu = is_npu()
|
||||
_is_cpu = is_cpu()
|
||||
@@ -291,14 +294,23 @@ class Qwen3GatedDeltaNet(nn.Module):
|
||||
|
||||
set_weight_attrs(self.A_log, {"weight_loader": sharded_weight_loader(0)})
|
||||
set_weight_attrs(self.dt_bias, {"weight_loader": sharded_weight_loader(0)})
|
||||
|
||||
self.norm = RMSNormGated(
|
||||
self.head_v_dim,
|
||||
eps=self.layer_norm_epsilon,
|
||||
group_size=None,
|
||||
norm_before_gate=True,
|
||||
device=torch.get_device_module().current_device(),
|
||||
dtype=config.torch_dtype,
|
||||
self.norm = (
|
||||
RMSNormGated(
|
||||
self.head_v_dim,
|
||||
eps=self.layer_norm_epsilon,
|
||||
group_size=None,
|
||||
norm_before_gate=True,
|
||||
device=torch.get_device_module().current_device(),
|
||||
dtype=config.torch_dtype,
|
||||
)
|
||||
if get_global_server_args().enable_piecewise_cuda_graph
|
||||
else FusedRMSNormGated(
|
||||
self.head_v_dim,
|
||||
eps=self.layer_norm_epsilon,
|
||||
activation=self.activation,
|
||||
device=torch.get_device_module().current_device(),
|
||||
dtype=config.torch_dtype,
|
||||
)
|
||||
)
|
||||
|
||||
self.out_proj = RowParallelLinear(
|
||||
|
||||
Reference in New Issue
Block a user