[Kernel] Remove unused implementations and stale registry entries (#32636)
This commit is contained in:
Xiaoyu Zhang
@@ -1,38 +1,13 @@
|
||||
from typing import Optional, Tuple
|
||||
|
||||
import torch
|
||||
import triton
|
||||
import triton.language as tl
|
||||
|
||||
from sglang.kernels.jit.utils import is_arch_support_pdl
|
||||
from sglang.kernels.ops.activation.softcap import softcap_out as fused_softcap
|
||||
from sglang.srt.utils import is_hip
|
||||
from sglang.srt.utils.custom_op import register_custom_op
|
||||
|
||||
_is_hip = is_hip()
|
||||
|
||||
|
||||
# cast to float + softcap
|
||||
class Softcap:
|
||||
def __init__(self, softcap_const: float):
|
||||
self.softcap_const = softcap_const
|
||||
|
||||
def __call__(self, *args, **kwargs):
|
||||
return self.forward(*args, **kwargs)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
if x.is_cuda:
|
||||
return self.forward_cuda(x)
|
||||
else:
|
||||
return self.forward_native(x)
|
||||
|
||||
def forward_native(self, x: torch.Tensor) -> torch.Tensor:
|
||||
return torch.tanh(x.float() / self.softcap_const) * self.softcap_const
|
||||
|
||||
def forward_cuda(self, x: torch.Tensor, autotune=False) -> torch.Tensor:
|
||||
return fused_softcap(x, self.softcap_const, autotune=autotune)
|
||||
|
||||
|
||||
rmsnorm_autotune = triton.autotune(
|
||||
configs=[
|
||||
triton.Config(kwargs={"BLOCK_SIZE": 1024}, num_warps=4, num_stages=1),
|
||||
@@ -213,136 +188,6 @@ def fused_rmsnorm(x, weight, eps, autotune=False, inplace=False):
|
||||
return output
|
||||
|
||||
|
||||
class FusedDualResidualRMSNorm:
|
||||
"""
|
||||
Fused implementation of
|
||||
y = RMSNorm2(RMSNorm1(x) + residual))
|
||||
"""
|
||||
|
||||
def __init__(self, rmsnorm1, rmsnorm2) -> None: # the one after rmsnorm1
|
||||
self.rmsnorm1 = rmsnorm1
|
||||
self.rmsnorm2 = rmsnorm2
|
||||
self.variance_epsilon = self.rmsnorm1.variance_epsilon
|
||||
assert self.rmsnorm1.variance_epsilon == self.rmsnorm2.variance_epsilon
|
||||
assert self.rmsnorm1.weight.shape == self.rmsnorm2.weight.shape
|
||||
|
||||
def __call__(self, *args, **kwargs):
|
||||
return self.forward(*args, **kwargs)
|
||||
|
||||
def forward(
|
||||
self, x: torch.Tensor, residual: torch.Tensor
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
if x.is_cuda:
|
||||
return self.forward_cuda(x, residual)
|
||||
else:
|
||||
return self.forward_flashinfer(x, residual)
|
||||
|
||||
def forward_cuda(
|
||||
self, x: torch.Tensor, residual: torch.Tensor, autotune=False
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
return fused_dual_residual_rmsnorm(
|
||||
x,
|
||||
residual,
|
||||
self.rmsnorm1.weight,
|
||||
self.rmsnorm2.weight,
|
||||
self.variance_epsilon,
|
||||
autotune=autotune,
|
||||
)
|
||||
|
||||
def forward_flashinfer(
|
||||
self,
|
||||
x: torch.Tensor,
|
||||
residual: torch.Tensor,
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
normed1 = self.rmsnorm1(x)
|
||||
residual = normed1 + residual
|
||||
return self.rmsnorm2(residual), residual
|
||||
|
||||
def forward_native(
|
||||
self,
|
||||
x: torch.Tensor,
|
||||
residual: torch.Tensor,
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
normed1 = self.rmsnorm1.forward_native(x)
|
||||
residual = normed1 + residual
|
||||
return self.rmsnorm2.forward_native(residual), residual
|
||||
|
||||
|
||||
@triton.jit
|
||||
def experts_combine_kernel(
|
||||
out_hidden_states,
|
||||
moe_hidden_states,
|
||||
mlp_hidden_states,
|
||||
combine_k: tl.constexpr,
|
||||
hidden_dim: tl.constexpr,
|
||||
BLOCK_SIZE: tl.constexpr,
|
||||
):
|
||||
pid = tl.program_id(0)
|
||||
start_index_mlp = pid * hidden_dim
|
||||
start_index_rmoe = pid * hidden_dim * combine_k
|
||||
offsets = tl.arange(0, BLOCK_SIZE)
|
||||
mask = offsets < hidden_dim
|
||||
combine_k_offsets = tl.arange(0, combine_k)
|
||||
|
||||
moe_x = tl.load(
|
||||
moe_hidden_states
|
||||
+ start_index_rmoe
|
||||
+ combine_k_offsets[:, None] * hidden_dim
|
||||
+ offsets[None, :],
|
||||
mask=mask[None, :],
|
||||
other=0.0,
|
||||
)
|
||||
moe_x = tl.sum(moe_x, axis=0)
|
||||
mlp_x = tl.load(mlp_hidden_states + start_index_mlp + offsets, mask=mask, other=0.0)
|
||||
combined_x = (moe_x + mlp_x) / 1.4142135623730951
|
||||
|
||||
tl.store(out_hidden_states + start_index_mlp + offsets, combined_x, mask=mask)
|
||||
|
||||
|
||||
@register_custom_op(out_shape="mlp_hidden_states")
|
||||
def experts_combine_triton(
|
||||
moe_hidden_states: torch.Tensor,
|
||||
mlp_hidden_states: torch.Tensor,
|
||||
output_buffer: Optional[torch.Tensor] = None,
|
||||
) -> torch.Tensor:
|
||||
assert moe_hidden_states.is_contiguous()
|
||||
assert mlp_hidden_states.is_contiguous()
|
||||
|
||||
if len(moe_hidden_states.shape) == 2:
|
||||
combine_k = 1 # pre-combined
|
||||
else:
|
||||
combine_k = moe_hidden_states.shape[1]
|
||||
|
||||
if output_buffer is None:
|
||||
out_hidden_states = torch.empty_like(mlp_hidden_states)
|
||||
else:
|
||||
flat_output_buffer = output_buffer.view(mlp_hidden_states.dtype).reshape(-1)
|
||||
assert flat_output_buffer.numel() >= mlp_hidden_states.numel()
|
||||
out_hidden_states = flat_output_buffer[: mlp_hidden_states.numel()].reshape(
|
||||
mlp_hidden_states.shape
|
||||
)
|
||||
|
||||
bs, hidden_dim = mlp_hidden_states.shape
|
||||
|
||||
config = {
|
||||
"BLOCK_SIZE": triton.next_power_of_2(hidden_dim),
|
||||
"num_warps": max(
|
||||
min(triton.next_power_of_2(triton.cdiv(hidden_dim, 1024)), 8), 4
|
||||
),
|
||||
}
|
||||
|
||||
experts_combine_kernel[(bs,)](
|
||||
out_hidden_states,
|
||||
moe_hidden_states,
|
||||
mlp_hidden_states,
|
||||
combine_k,
|
||||
hidden_dim,
|
||||
**config,
|
||||
)
|
||||
|
||||
return out_hidden_states
|
||||
|
||||
|
||||
# gelu on first half of vector
|
||||
@triton.jit
|
||||
def gelu_and_mul_kernel(
|
||||
|
||||
Reference in New Issue
Block a user