571 lines
18 KiB
Python
571 lines
18 KiB
Python
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
from typing import TYPE_CHECKING, List, Optional
|
|
|
|
import torch
|
|
|
|
from sglang.srt.layers import deep_gemm_wrapper
|
|
from sglang.srt.layers.moe.moe_runner.base import (
|
|
MoeQuantInfo,
|
|
MoeRunnerConfig,
|
|
MoeRunnerCore,
|
|
RunnerInput,
|
|
RunnerOutput,
|
|
register_post_permute,
|
|
register_pre_permute,
|
|
)
|
|
from sglang.srt.layers.moe.utils import MoeRunnerBackend
|
|
from sglang.srt.utils import ceil_div, dispose_tensor, get_bool_env_var, is_hip, is_npu
|
|
from sglang.srt.utils.offloader import get_offloader
|
|
|
|
if TYPE_CHECKING:
|
|
from sglang.srt.layers.moe.token_dispatcher.deepep import (
|
|
DeepEPLLCombineInput,
|
|
DeepEPLLDispatchOutput,
|
|
DeepEPNormalCombineInput,
|
|
DeepEPNormalDispatchOutput,
|
|
)
|
|
from sglang.srt.layers.moe.token_dispatcher.standard import (
|
|
StandardCombineInput,
|
|
StandardDispatchOutput,
|
|
)
|
|
|
|
_is_hip = is_hip()
|
|
_is_npu = is_npu()
|
|
_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip
|
|
|
|
if not (_is_npu or _is_hip):
|
|
from sgl_kernel import silu_and_mul
|
|
|
|
|
|
# TODO(kaixih@nvidia): ideally we should merge this logic into
|
|
# `fill_gateup_input_triton_kernel` to directly generate e8m0 scale.
|
|
@torch.compile
|
|
def _cast_to_e8m0_with_rounding_up(x: torch.Tensor) -> torch.Tensor:
|
|
temp = x.to(torch.float32).view(torch.int32)
|
|
exp = torch.bitwise_right_shift(temp, 23)
|
|
mant = torch.bitwise_and(temp, 0x7FFFFF)
|
|
is_ru = torch.logical_and(
|
|
torch.logical_and((mant > 0), (exp != 0xFE)),
|
|
~torch.logical_and((exp == 0), (mant <= 0x400000)),
|
|
)
|
|
exp = torch.where(is_ru, exp + 1, exp)
|
|
new_x = exp.to(torch.uint8).view(torch.int)
|
|
return new_x.transpose(1, 2).contiguous().transpose(1, 2)
|
|
|
|
|
|
def copy_list_to_gpu_no_ce(arr: List[int]):
|
|
from sgl_kernel.elementwise import copy_to_gpu_no_ce
|
|
|
|
tensor_cpu = torch.tensor(arr, dtype=torch.int32, device="cpu")
|
|
tensor_gpu = torch.empty_like(tensor_cpu, device="cuda")
|
|
copy_to_gpu_no_ce(tensor_cpu, tensor_gpu)
|
|
return tensor_gpu
|
|
|
|
|
|
@dataclass
|
|
class DeepGemmRunnerInput(RunnerInput):
|
|
hidden_states: torch.Tensor
|
|
hidden_states_scale: torch.Tensor
|
|
use_masked_gemm: bool
|
|
masked_m: Optional[torch.Tensor] = None
|
|
expected_m: Optional[int] = None
|
|
m_indices: Optional[torch.Tensor] = None
|
|
|
|
@property
|
|
def runner_backend(self) -> MoeRunnerBackend:
|
|
return MoeRunnerBackend.DEEP_GEMM
|
|
|
|
|
|
@dataclass
|
|
class DeepGemmRunnerOutput(RunnerOutput):
|
|
hidden_states: torch.Tensor
|
|
|
|
@property
|
|
def runner_backend(self) -> MoeRunnerBackend:
|
|
return MoeRunnerBackend.DEEP_GEMM
|
|
|
|
|
|
@dataclass
|
|
class DeepGemmMoeQuantInfo(MoeQuantInfo):
|
|
w13_weight: torch.Tensor
|
|
w2_weight: torch.Tensor
|
|
use_fp8: bool
|
|
w13_scale: Optional[torch.Tensor] = None
|
|
w2_scale: Optional[torch.Tensor] = None
|
|
block_shape: Optional[List[int]] = None
|
|
|
|
|
|
class DeepGemmRunnerCore(MoeRunnerCore):
|
|
def __init__(self, config: MoeRunnerConfig):
|
|
super().__init__(config)
|
|
assert self.config.activation == "silu"
|
|
|
|
def run(
|
|
self,
|
|
runner_input: DeepGemmRunnerInput,
|
|
quant_info: DeepGemmMoeQuantInfo,
|
|
running_state: dict,
|
|
) -> DeepGemmRunnerOutput:
|
|
|
|
if not runner_input.use_masked_gemm:
|
|
hidden_states = self._run_contiguous_gemm(
|
|
runner_input, quant_info, running_state
|
|
)
|
|
else:
|
|
hidden_states = self._run_masked_gemm(
|
|
runner_input, quant_info, running_state
|
|
)
|
|
return DeepGemmRunnerOutput(hidden_states=hidden_states)
|
|
|
|
def _run_contiguous_gemm(
|
|
self,
|
|
runner_input: DeepGemmRunnerInput,
|
|
quant_info: DeepGemmMoeQuantInfo,
|
|
running_state: dict,
|
|
) -> torch.Tensor:
|
|
|
|
from sglang.srt.layers.moe.ep_moe.kernels import tma_align_input_scale
|
|
from sglang.srt.layers.quantization.fp8_kernel import (
|
|
sglang_per_token_group_quant_fp8,
|
|
)
|
|
|
|
hidden_states = runner_input.hidden_states
|
|
hidden_states_scale = runner_input.hidden_states_scale
|
|
all_tokens = running_state["all_tokens"]
|
|
hidden_states_device = running_state["hidden_states_device"]
|
|
hidden_states_dtype = running_state["hidden_states_dtype"]
|
|
hidden_states_shape = running_state["hidden_states_shape"]
|
|
m_indices = runner_input.m_indices
|
|
|
|
N = quant_info.w13_weight.size(1)
|
|
K = hidden_states_shape[1]
|
|
scale_block_size = 128
|
|
|
|
w13_weight_fp8 = (
|
|
quant_info.w13_weight,
|
|
quant_info.w13_scale,
|
|
)
|
|
w2_weight_fp8 = (quant_info.w2_weight, quant_info.w2_scale)
|
|
|
|
gateup_output = torch.empty(
|
|
(all_tokens, N),
|
|
device=hidden_states_device,
|
|
dtype=torch.bfloat16,
|
|
)
|
|
if not deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0:
|
|
hidden_states_scale = tma_align_input_scale(hidden_states_scale)
|
|
deep_gemm_wrapper.grouped_gemm_nt_f8f8bf16_contig(
|
|
(hidden_states, hidden_states_scale),
|
|
w13_weight_fp8,
|
|
gateup_output,
|
|
m_indices,
|
|
)
|
|
|
|
dispose_tensor(hidden_states)
|
|
dispose_tensor(hidden_states_scale)
|
|
|
|
down_input = torch.empty(
|
|
(
|
|
all_tokens,
|
|
N // 2,
|
|
),
|
|
device=gateup_output.device,
|
|
dtype=torch.bfloat16,
|
|
)
|
|
silu_and_mul(gateup_output.view(-1, N), down_input)
|
|
del gateup_output
|
|
|
|
down_input_fp8, down_input_scale = sglang_per_token_group_quant_fp8(
|
|
down_input,
|
|
scale_block_size,
|
|
column_major_scales=deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0,
|
|
scale_tma_aligned=deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0,
|
|
scale_ue8m0=deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0,
|
|
)
|
|
del down_input
|
|
|
|
down_output = torch.empty(
|
|
(all_tokens, K),
|
|
device=hidden_states_device,
|
|
dtype=torch.bfloat16,
|
|
)
|
|
if not deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0:
|
|
down_input_scale = tma_align_input_scale(down_input_scale)
|
|
|
|
deep_gemm_wrapper.grouped_gemm_nt_f8f8bf16_contig(
|
|
(down_input_fp8, down_input_scale),
|
|
w2_weight_fp8,
|
|
down_output,
|
|
m_indices,
|
|
)
|
|
|
|
return down_output
|
|
|
|
def _run_masked_gemm(
|
|
self,
|
|
runner_input: DeepGemmRunnerInput,
|
|
quant_info: DeepGemmMoeQuantInfo,
|
|
running_state: dict,
|
|
) -> torch.Tensor:
|
|
|
|
from sglang.srt.layers import deep_gemm_wrapper
|
|
from sglang.srt.layers.moe.ep_moe.kernels import (
|
|
silu_and_mul_masked_post_quant_fwd,
|
|
)
|
|
|
|
hidden_states = runner_input.hidden_states
|
|
hidden_states_scale = runner_input.hidden_states_scale
|
|
masked_m = runner_input.masked_m
|
|
expected_m = runner_input.expected_m
|
|
|
|
w13_weight = quant_info.w13_weight
|
|
w2_weight = quant_info.w2_weight
|
|
w13_scale = quant_info.w13_scale
|
|
w2_scale = quant_info.w2_scale
|
|
|
|
hidden_states_device = running_state["hidden_states_device"]
|
|
|
|
# GroupGemm-0
|
|
if deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0:
|
|
if hidden_states_scale.dtype != torch.int:
|
|
b, s_mn, s_k = hidden_states_scale.shape
|
|
assert (
|
|
s_mn % 4 == 0 and s_k % 4 == 0
|
|
), f"scales must be aligned to 4, but got ({b}, {s_mn}, {s_k})"
|
|
hidden_states_scale = _cast_to_e8m0_with_rounding_up(
|
|
hidden_states_scale
|
|
)
|
|
else:
|
|
hidden_states_scale = deep_gemm_wrapper.get_mn_major_tma_aligned_tensor(
|
|
hidden_states_scale
|
|
)
|
|
|
|
num_groups, m, k = hidden_states.shape
|
|
n = w13_weight.size(1)
|
|
gateup_output = torch.empty(
|
|
(num_groups, m, n), device=hidden_states_device, dtype=torch.bfloat16
|
|
)
|
|
deep_gemm_wrapper.grouped_gemm_nt_f8f8bf16_masked(
|
|
(hidden_states, hidden_states_scale),
|
|
(w13_weight, w13_scale),
|
|
gateup_output,
|
|
masked_m,
|
|
expected_m,
|
|
)
|
|
dispose_tensor(hidden_states)
|
|
dispose_tensor(hidden_states_scale)
|
|
|
|
# Act
|
|
down_input = torch.empty(
|
|
(
|
|
gateup_output.shape[0],
|
|
gateup_output.shape[1],
|
|
gateup_output.shape[2] // 2,
|
|
),
|
|
device=hidden_states_device,
|
|
dtype=torch.float8_e4m3fn,
|
|
)
|
|
scale_block_size = 128
|
|
down_input_scale = torch.empty(
|
|
(
|
|
gateup_output.shape[0],
|
|
gateup_output.shape[1],
|
|
gateup_output.shape[2] // 2 // scale_block_size,
|
|
),
|
|
device=hidden_states_device,
|
|
dtype=torch.float32,
|
|
)
|
|
silu_and_mul_masked_post_quant_fwd(
|
|
gateup_output,
|
|
down_input,
|
|
down_input_scale,
|
|
scale_block_size,
|
|
masked_m,
|
|
scale_ue8m0=deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0,
|
|
)
|
|
del gateup_output
|
|
|
|
# GroupGemm-1
|
|
n = w2_weight.shape[1]
|
|
|
|
if not deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0:
|
|
down_input_scale = deep_gemm_wrapper.get_mn_major_tma_aligned_tensor(
|
|
down_input_scale
|
|
)
|
|
|
|
down_output = torch.empty(
|
|
(num_groups, m, n), device=hidden_states_device, dtype=torch.bfloat16
|
|
)
|
|
deep_gemm_wrapper.grouped_gemm_nt_f8f8bf16_masked(
|
|
(down_input, down_input_scale),
|
|
(w2_weight, w2_scale),
|
|
down_output,
|
|
masked_m,
|
|
expected_m,
|
|
)
|
|
|
|
return down_output
|
|
|
|
@property
|
|
def runner_backend(self) -> MoeRunnerBackend:
|
|
return MoeRunnerBackend.DEEP_GEMM
|
|
|
|
|
|
@register_pre_permute("standard", "deep_gemm")
|
|
def pre_permute_standard_to_deep_gemm(
|
|
dispatch_output: StandardDispatchOutput,
|
|
quant_info: DeepGemmMoeQuantInfo,
|
|
runner_config: MoeRunnerConfig,
|
|
running_state: dict,
|
|
) -> DeepGemmRunnerInput:
|
|
|
|
from sglang.srt.layers.moe.ep_moe.kernels import moe_ep_deepgemm_preprocess
|
|
|
|
hidden_states, topk_output = dispatch_output
|
|
topk_weights, topk_ids, _ = topk_output
|
|
|
|
hidden_states_shape = hidden_states.shape
|
|
hidden_states_dtype = hidden_states.dtype
|
|
hidden_states_device = hidden_states.device
|
|
hidden_states_ref = hidden_states
|
|
|
|
topk_weights, topk_ids = topk_weights, topk_ids
|
|
|
|
# PreReorder
|
|
masked_m, expected_m, src2dst, hidden_states, hidden_states_scale = (
|
|
moe_ep_deepgemm_preprocess(
|
|
topk_ids,
|
|
runner_config.num_local_experts,
|
|
hidden_states,
|
|
runner_config.top_k,
|
|
quant_info.block_shape,
|
|
)
|
|
)
|
|
|
|
dispose_tensor(hidden_states_ref)
|
|
|
|
running_state["topk_ids"] = topk_ids
|
|
running_state["topk_weights"] = topk_weights
|
|
running_state["hidden_states_shape"] = hidden_states_shape
|
|
running_state["hidden_states_dtype"] = hidden_states_dtype
|
|
running_state["hidden_states_device"] = hidden_states_device
|
|
running_state["src2dst"] = src2dst
|
|
|
|
return DeepGemmRunnerInput(
|
|
hidden_states=hidden_states,
|
|
hidden_states_scale=hidden_states_scale,
|
|
use_masked_gemm=True,
|
|
masked_m=masked_m,
|
|
expected_m=expected_m,
|
|
)
|
|
|
|
|
|
@register_post_permute("deep_gemm", "standard")
|
|
def post_permute_deep_gemm_to_standard(
|
|
runner_output: DeepGemmRunnerOutput,
|
|
quant_info: DeepGemmMoeQuantInfo,
|
|
runner_config: MoeRunnerConfig,
|
|
running_state: dict,
|
|
) -> StandardCombineInput:
|
|
from sglang.srt.layers.moe.ep_moe.kernels import post_reorder_triton_kernel
|
|
from sglang.srt.layers.moe.token_dispatcher.standard import StandardCombineInput
|
|
|
|
hidden_states_shape = running_state["hidden_states_shape"]
|
|
hidden_states_dtype = running_state["hidden_states_dtype"]
|
|
hidden_states_device = running_state["hidden_states_device"]
|
|
src2dst = running_state["src2dst"]
|
|
topk_ids = running_state["topk_ids"]
|
|
topk_weights = running_state["topk_weights"]
|
|
|
|
output = torch.empty(
|
|
hidden_states_shape, dtype=hidden_states_dtype, device=hidden_states_device
|
|
)
|
|
post_reorder_triton_kernel[(hidden_states_shape[0],)](
|
|
runner_output.hidden_states,
|
|
output,
|
|
src2dst,
|
|
topk_ids,
|
|
topk_weights,
|
|
runner_config.top_k,
|
|
hidden_states_shape[1],
|
|
BLOCK_SIZE=512,
|
|
)
|
|
|
|
dispose_tensor(runner_output.hidden_states)
|
|
|
|
if runner_config.routed_scaling_factor is not None:
|
|
output *= runner_config.routed_scaling_factor
|
|
|
|
return StandardCombineInput(
|
|
hidden_states=output,
|
|
)
|
|
|
|
|
|
@register_pre_permute("deepep_ll", "deep_gemm")
|
|
def pre_permute_deepep_ll_to_deep_gemm(
|
|
dispatch_output: DeepEPLLDispatchOutput,
|
|
quant_info: DeepGemmMoeQuantInfo,
|
|
runner_config: MoeRunnerConfig,
|
|
running_state: dict,
|
|
) -> DeepGemmRunnerInput:
|
|
|
|
hidden_states, hidden_states_scale, topk_ids, topk_weights, masked_m, expected_m = (
|
|
dispatch_output
|
|
)
|
|
|
|
running_state["topk_ids"] = topk_ids
|
|
running_state["topk_weights"] = topk_weights
|
|
running_state["hidden_states_shape"] = hidden_states.shape
|
|
running_state["hidden_states_dtype"] = hidden_states.dtype
|
|
running_state["hidden_states_device"] = hidden_states.device
|
|
|
|
return DeepGemmRunnerInput(
|
|
hidden_states=hidden_states,
|
|
hidden_states_scale=hidden_states_scale,
|
|
use_masked_gemm=True,
|
|
masked_m=masked_m,
|
|
expected_m=expected_m,
|
|
)
|
|
|
|
|
|
@register_post_permute("deep_gemm", "deepep_ll")
|
|
def post_permute_deep_gemm_to_deepep_ll(
|
|
runner_output: DeepGemmRunnerOutput,
|
|
quant_info: DeepGemmMoeQuantInfo,
|
|
runner_config: MoeRunnerConfig,
|
|
running_state: dict,
|
|
) -> DeepEPLLCombineInput:
|
|
|
|
from sglang.srt.layers.moe.token_dispatcher.deepep import DeepEPLLCombineInput
|
|
|
|
return DeepEPLLCombineInput(
|
|
hidden_states=runner_output.hidden_states,
|
|
topk_ids=running_state["topk_ids"],
|
|
topk_weights=running_state["topk_weights"],
|
|
)
|
|
|
|
|
|
@register_pre_permute("deepep_normal", "deep_gemm")
|
|
def pre_permute_deepep_normal_to_deep_gemm(
|
|
dispatch_output: DeepEPNormalDispatchOutput,
|
|
quant_info: DeepGemmMoeQuantInfo,
|
|
runner_config: MoeRunnerConfig,
|
|
running_state: dict,
|
|
) -> DeepGemmRunnerInput:
|
|
|
|
from sglang.srt.layers.moe.ep_moe.kernels import ep_scatter
|
|
|
|
(
|
|
hidden_states,
|
|
hidden_states_scale,
|
|
topk_ids,
|
|
topk_weights,
|
|
num_recv_tokens_per_expert,
|
|
) = dispatch_output
|
|
assert runner_config.activation == "silu"
|
|
|
|
all_tokens = sum(num_recv_tokens_per_expert)
|
|
running_state["all_tokens"] = all_tokens
|
|
|
|
K = hidden_states.shape[1]
|
|
|
|
hidden_states_shape = hidden_states.shape
|
|
hidden_states_device = hidden_states.device
|
|
hidden_states_dtype = hidden_states.dtype
|
|
|
|
running_state["hidden_states_shape"] = hidden_states_shape
|
|
running_state["hidden_states_device"] = hidden_states_device
|
|
running_state["hidden_states_dtype"] = hidden_states_dtype
|
|
running_state["topk_ids"] = topk_ids
|
|
running_state["topk_weights"] = topk_weights
|
|
|
|
input_tensor = torch.empty(
|
|
(all_tokens, K),
|
|
device=hidden_states.device,
|
|
dtype=hidden_states.dtype,
|
|
)
|
|
if deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0:
|
|
# TODO check whether need `zeros`
|
|
input_tensor_scale = torch.zeros(
|
|
(ceil_div(K // 128, 4), all_tokens),
|
|
device=hidden_states.device,
|
|
dtype=torch.int,
|
|
).transpose(0, 1)
|
|
else:
|
|
input_tensor_scale = torch.empty(
|
|
(all_tokens, K // 128),
|
|
device=hidden_states.device,
|
|
dtype=torch.float32,
|
|
)
|
|
m_indices = torch.empty(all_tokens, device=hidden_states.device, dtype=torch.int32)
|
|
output_index = torch.empty_like(topk_ids)
|
|
|
|
if get_offloader().forbid_copy_engine_usage:
|
|
num_recv_tokens_per_expert_gpu = copy_list_to_gpu_no_ce(
|
|
num_recv_tokens_per_expert
|
|
)
|
|
else:
|
|
num_recv_tokens_per_expert_gpu = torch.tensor(
|
|
num_recv_tokens_per_expert,
|
|
dtype=torch.int32,
|
|
pin_memory=True,
|
|
device="cpu",
|
|
).cuda(non_blocking=True)
|
|
expert_start_loc = torch.empty_like(num_recv_tokens_per_expert_gpu)
|
|
|
|
ep_scatter(
|
|
hidden_states,
|
|
hidden_states_scale,
|
|
topk_ids,
|
|
num_recv_tokens_per_expert_gpu,
|
|
expert_start_loc,
|
|
input_tensor,
|
|
input_tensor_scale,
|
|
m_indices,
|
|
output_index,
|
|
scale_ue8m0=deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0,
|
|
)
|
|
dispose_tensor(hidden_states)
|
|
dispose_tensor(hidden_states_scale)
|
|
|
|
running_state["output_index"] = output_index
|
|
|
|
return DeepGemmRunnerInput(
|
|
hidden_states=input_tensor,
|
|
hidden_states_scale=input_tensor_scale,
|
|
use_masked_gemm=False,
|
|
m_indices=m_indices,
|
|
)
|
|
|
|
|
|
@register_post_permute("deep_gemm", "deepep_normal")
|
|
def post_permute_deep_gemm_to_deepep_normal(
|
|
runner_output: DeepGemmRunnerOutput,
|
|
quant_info: DeepGemmMoeQuantInfo,
|
|
runner_config: MoeRunnerConfig,
|
|
running_state: dict,
|
|
) -> DeepEPNormalCombineInput:
|
|
|
|
from sglang.srt.layers.moe.ep_moe.kernels import ep_gather
|
|
from sglang.srt.layers.moe.token_dispatcher.deepep import DeepEPNormalCombineInput
|
|
|
|
hidden_states = runner_output.hidden_states
|
|
topk_ids = running_state["topk_ids"]
|
|
topk_weights = running_state["topk_weights"]
|
|
output_index = running_state["output_index"]
|
|
|
|
gather_out = torch.empty(
|
|
running_state["hidden_states_shape"],
|
|
device=running_state["hidden_states_device"],
|
|
dtype=torch.bfloat16,
|
|
)
|
|
ep_gather(hidden_states, topk_ids, topk_weights, output_index, gather_out)
|
|
|
|
return DeepEPNormalCombineInput(
|
|
hidden_states=gather_out,
|
|
topk_ids=running_state["topk_ids"],
|
|
topk_weights=running_state["topk_weights"],
|
|
)
|