2469 lines
95 KiB
Python
2469 lines
95 KiB
Python
# Copyright 2024 SGLang Team
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
# ==============================================================================
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
import math
|
|
from dataclasses import dataclass
|
|
from enum import IntEnum, auto
|
|
from typing import (
|
|
TYPE_CHECKING,
|
|
Callable,
|
|
NamedTuple,
|
|
Optional,
|
|
Protocol,
|
|
Tuple,
|
|
TypeGuard,
|
|
runtime_checkable,
|
|
)
|
|
|
|
import torch
|
|
import torch.nn.functional as F
|
|
|
|
if TYPE_CHECKING:
|
|
from triton_kernels.tensor_details.ragged_tensor import RaggedTensorMetadata
|
|
|
|
from sglang.srt.runtime_context import get_exec, get_lora, get_parallel
|
|
|
|
try:
|
|
from triton_kernels.tensor import make_ragged_tensor_metadata
|
|
from triton_kernels.topk import topk as triton_kernels_topk
|
|
|
|
def routing(
|
|
logits,
|
|
n_expts_act,
|
|
sm_first=False,
|
|
expt_indx=None,
|
|
simulated_ep=1,
|
|
n_rows=None,
|
|
):
|
|
if simulated_ep != 1:
|
|
raise NotImplementedError(
|
|
"simulated_ep routing is not supported with triton_kernels 3.7.1"
|
|
)
|
|
|
|
if sm_first:
|
|
logits = torch.softmax(logits, dim=-1)
|
|
|
|
sparse_logits = triton_kernels_topk(
|
|
logits,
|
|
n_expts_act,
|
|
apply_softmax=not sm_first,
|
|
y_indx=expt_indx,
|
|
n_rows=n_rows,
|
|
)
|
|
dispatch_indx = sparse_logits.mask_metadata.row_sorted_indx
|
|
combine_indx = sparse_logits.mask_metadata.col_sorted_indx
|
|
gather_indx = torch.div(combine_indx, n_expts_act, rounding_mode="trunc")
|
|
scatter_indx = combine_indx
|
|
ragged_metadata = make_ragged_tensor_metadata(
|
|
sparse_logits.mask_metadata.col_sum, dispatch_indx.shape[0]
|
|
)
|
|
gate_scal = sparse_logits.vals.flatten()[combine_indx]
|
|
return ragged_metadata, gather_indx, scatter_indx, gate_scal, n_expts_act
|
|
|
|
except ImportError:
|
|
pass
|
|
|
|
from sglang.kernels.fused_op import BaseFusedOp
|
|
from sglang.kernels.ops.attention.dsv4 import mask_topk_ids
|
|
from sglang.srt.distributed import (
|
|
get_tp_group,
|
|
)
|
|
from sglang.srt.distributed.device_communicators.pynccl_allocator import (
|
|
use_symmetric_memory,
|
|
)
|
|
from sglang.srt.environ import envs
|
|
from sglang.srt.eplb import expert_location_dispatch
|
|
from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder
|
|
from sglang.srt.eplb.expert_location_dispatch import (
|
|
ExpertLocationDispatchInfo,
|
|
topk_ids_logical_to_physical,
|
|
)
|
|
from sglang.srt.layers.dp_attention import is_allocation_symmetric
|
|
from sglang.srt.layers.moe import get_moe_runner_backend
|
|
from sglang.srt.layers.moe.utils import (
|
|
has_per_rank_fused_shared_slots,
|
|
)
|
|
from sglang.srt.state_capturer.routed_experts import get_global_experts_capturer
|
|
from sglang.srt.utils import (
|
|
cpu_has_amx_support,
|
|
get_bool_env_var,
|
|
get_compiler_backend,
|
|
is_cpu,
|
|
is_cuda,
|
|
is_hip,
|
|
is_musa,
|
|
is_npu,
|
|
is_xpu,
|
|
)
|
|
|
|
_SGLANG_EXPERIMENTAL_LORA_OPTI = envs.SGLANG_EXPERIMENTAL_LORA_OPTI.get()
|
|
|
|
if TYPE_CHECKING:
|
|
from sglang.srt.layers.quantization import QuantizationConfig
|
|
|
|
|
|
logger = logging.getLogger(__name__)
|
|
_is_cuda = is_cuda()
|
|
_is_hip = is_hip()
|
|
_is_cpu = is_cpu()
|
|
_is_cpu_amx_available = cpu_has_amx_support()
|
|
_is_xpu = is_xpu()
|
|
_is_npu = is_npu()
|
|
_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip
|
|
_is_musa = is_musa()
|
|
|
|
# Epsilon added to the top-k weight sum before renormalization, matching the
|
|
# DeepSeek reference gate (modeling_deepseek.py: `topk_weight.sum(...) + 1e-20`)
|
|
# and flashinfer's trtllm routing kernels (mSumEpsilon). With sigmoid scoring
|
|
# plus a selection bias, a token whose selected experts all have deeply negative
|
|
# router logits can have every gathered sigmoid weight underflow to exactly
|
|
# zero; a bare division then yields 0/0 = NaN and poisons the token's output
|
|
# row. For healthy tokens the sum is >= sigmoid(logit_max) >> 1e-20, so results
|
|
# are unchanged. The renormalization is performed in float32 (the reference gate
|
|
# computes the whole gate in fp32); the epsilon underflows to zero in float16.
|
|
_RENORMALIZE_SUM_EPSILON = 1e-20
|
|
|
|
# Experimental: skip the HIP padded-token routing-weight masking entirely.
|
|
# Padded (CUDA-graph) rows are discarded downstream and the MoE combine is
|
|
# per-token, so zeroing their weights is in principle unnecessary. Gated off by
|
|
# default because it is a numerics-affecting change that must be validated with
|
|
# an accuracy run before becoming the default.
|
|
_skip_hip_pad_mask = get_bool_env_var("SGLANG_MORI_NO_PAD_MASK", "False")
|
|
|
|
|
|
if _is_cuda:
|
|
try:
|
|
from flashinfer.fused_moe import fused_topk_deepseek as _fused_topk_deepseek
|
|
|
|
from sglang.srt.utils.custom_op import register_custom_op
|
|
|
|
@register_custom_op(
|
|
op_name="fused_topk_deepseek",
|
|
mutates_args=["topk_weights", "topk_ids"],
|
|
)
|
|
def fused_topk_deepseek(
|
|
gating_output: torch.Tensor,
|
|
correction_bias: torch.Tensor,
|
|
num_expert_group: int,
|
|
topk_group: int,
|
|
topk: int,
|
|
scaling_factor: float,
|
|
topk_weights: torch.Tensor,
|
|
topk_ids: torch.Tensor,
|
|
renormalize: bool,
|
|
) -> None:
|
|
_fused_topk_deepseek(
|
|
gating_output,
|
|
correction_bias,
|
|
num_expert_group,
|
|
topk_group,
|
|
topk,
|
|
scaling_factor,
|
|
topk_weights,
|
|
topk_ids,
|
|
renormalize,
|
|
)
|
|
|
|
except ImportError:
|
|
fused_topk_deepseek = None
|
|
|
|
if _is_cuda or _is_hip or _is_xpu:
|
|
if _is_xpu:
|
|
# XPU has no tvm_ffi, so the CUDA JIT topk_sigmoid isn't reachable;
|
|
# use the AOT symbols from sgl_kernel directly. topk_sigmoid was aligned
|
|
# with the post-#28715 CUDA signature in sgl-kernel-xpu#285.
|
|
from sgl_kernel import topk_sigmoid, topk_softmax
|
|
else:
|
|
from sglang.kernels.ops.moe import topk_softmax
|
|
|
|
try:
|
|
from sglang.kernels.ops.moe.moe_topk_sigmoid import topk_sigmoid
|
|
except ImportError:
|
|
pass
|
|
if _use_aiter:
|
|
try:
|
|
from aiter import biased_grouped_topk as aiter_biased_grouped_topk
|
|
from aiter.fused_moe import fused_topk as aiter_fused_topk
|
|
except ImportError:
|
|
raise ImportError("aiter is required when SGLANG_USE_AITER is set to True")
|
|
if _is_musa:
|
|
try:
|
|
from mate import moe_fused_gate
|
|
except ImportError:
|
|
raise ImportError("mate is required for the biased grouped topk.")
|
|
|
|
from sglang.srt.hardware_backend.musa.kernels.topk import topk_sigmoid, topk_softmax
|
|
|
|
# -------------------------------- TopKConfig ---------------------------------------
|
|
|
|
|
|
@dataclass
|
|
class TopKConfig:
|
|
top_k: int
|
|
use_grouped_topk: bool = False
|
|
topk_group: Optional[int] = None
|
|
num_expert_group: Optional[int] = None
|
|
renormalize: bool = True
|
|
num_fused_shared_experts: int = 0
|
|
custom_routing_function: Optional[Callable] = None
|
|
correction_bias: Optional[torch.Tensor] = None
|
|
torch_native: bool = False
|
|
routed_scaling_factor: Optional[float] = None
|
|
apply_routed_scaling_factor_on_output: bool = False
|
|
fused_shared_experts_scaling_factor: Optional[float] = None
|
|
output_format: Optional[TopKOutputFormat] = None
|
|
scoring_func: str = "softmax"
|
|
# Draft-side MoE blocks set this False so they never write the target's
|
|
# process-global routed-experts capture buffer.
|
|
allow_routed_experts_capture: bool = True
|
|
|
|
|
|
# -------------------------------- TopKOutput ---------------------------------------
|
|
|
|
|
|
class TopKOutputChecker:
|
|
|
|
@staticmethod
|
|
def format_is_standard(topk_output: TopKOutput) -> TypeGuard[StandardTopKOutput]:
|
|
# ===== TO BE REFACTORED ====
|
|
# The experimental fused topk+pack carrier only exists under the master switch.
|
|
if _SGLANG_EXPERIMENTAL_LORA_OPTI:
|
|
return isinstance(
|
|
topk_output, (StandardTopKOutput, StandardTopKOutputPacked)
|
|
)
|
|
# ===== END TO BE REFACTORED ====
|
|
return isinstance(topk_output, StandardTopKOutput)
|
|
|
|
@staticmethod
|
|
def format_is_triton_kernels(
|
|
topk_output: TopKOutput,
|
|
) -> TypeGuard[TritonKernelTopKOutput]:
|
|
return isinstance(topk_output, TritonKernelTopKOutput)
|
|
|
|
@staticmethod
|
|
def format_is_bypassed(topk_output: TopKOutput) -> TypeGuard[BypassedTopKOutput]:
|
|
return isinstance(topk_output, BypassedTopKOutput)
|
|
|
|
@staticmethod
|
|
def format_is_packed(topk_output: TopKOutput) -> TypeGuard[PackedTopKOutput]:
|
|
return isinstance(topk_output, PackedTopKOutput)
|
|
|
|
|
|
class TopKOutputFormat(IntEnum):
|
|
STANDARD = auto()
|
|
TRITON_KERNEL = auto()
|
|
BYPASSED = auto()
|
|
PACKED = auto()
|
|
|
|
|
|
@runtime_checkable
|
|
class TopKOutput(Protocol):
|
|
"""Protocol for top-k outputs in different formats."""
|
|
|
|
@property
|
|
def format(self) -> TopKOutputFormat:
|
|
"""The format of the output."""
|
|
...
|
|
|
|
|
|
class StandardTopKOutput(NamedTuple):
|
|
"""Standard top-k output format."""
|
|
|
|
topk_weights: torch.Tensor
|
|
topk_ids: torch.Tensor
|
|
router_logits: torch.Tensor
|
|
|
|
@property
|
|
def format(self) -> TopKOutputFormat:
|
|
return TopKOutputFormat.STANDARD
|
|
|
|
|
|
# ===== TO BE REFACTORED ====
|
|
# Experimental fused topk+pack (SGLANG_OPT_LORA_FUSED_TOPK_PACK) carrier: the FlashInfer
|
|
# routed-MoE packed topk produced fused in the gating kernel. Kept a SEPARATE type rather
|
|
# than a 4th StandardTopKOutput field so the OSS `a, b, _ = topk_output` 3-tuple unpack
|
|
# stays valid; only the gated experimental MoE dispatch reads .packed_topk_ids (getattr).
|
|
class StandardTopKOutputPacked(NamedTuple):
|
|
topk_weights: torch.Tensor
|
|
topk_ids: torch.Tensor
|
|
router_logits: torch.Tensor
|
|
packed_topk_ids: torch.Tensor
|
|
|
|
@property
|
|
def format(self) -> TopKOutputFormat:
|
|
return TopKOutputFormat.STANDARD
|
|
|
|
|
|
# ===== END TO BE REFACTORED ====
|
|
|
|
|
|
class TritonKernelTopKOutput(NamedTuple):
|
|
"""Triton kernel top-k output format."""
|
|
|
|
a_ragged_metadata: RaggedTensorMetadata
|
|
gather_indx: torch.Tensor
|
|
scatter_indx: torch.Tensor
|
|
gate_scal: torch.Tensor
|
|
n_expts_act: int
|
|
|
|
@property
|
|
def format(self) -> TopKOutputFormat:
|
|
return TopKOutputFormat.TRITON_KERNEL
|
|
|
|
|
|
class BypassedTopKOutput(NamedTuple):
|
|
"""Bypassed top-k output format."""
|
|
|
|
hidden_states: torch.Tensor
|
|
router_logits: torch.Tensor
|
|
topk_config: TopKConfig
|
|
num_token_non_padded: Optional[torch.Tensor] = None
|
|
expert_location_dispatch_info: Optional[ExpertLocationDispatchInfo] = None
|
|
|
|
@property
|
|
def format(self) -> TopKOutputFormat:
|
|
return TopKOutputFormat.BYPASSED
|
|
|
|
def to_standard(self, layer_id: Optional[int] = None) -> StandardTopKOutput:
|
|
"""Materialize routing tensors. Used by MoE kernels that need explicit
|
|
topk_ids / topk_weights rather than doing routing internally."""
|
|
return select_experts(
|
|
hidden_states=self.hidden_states,
|
|
router_logits=self.router_logits,
|
|
topk_config=self.topk_config,
|
|
layer_id=layer_id,
|
|
num_token_non_padded=self.num_token_non_padded,
|
|
expert_location_dispatch_info=self.expert_location_dispatch_info,
|
|
)
|
|
|
|
|
|
class PackedTopKOutput(NamedTuple):
|
|
"""Packed top-k output format used by FlashInfer TRT-LLM routed MoE.
|
|
|
|
``packed_topk_ids`` is an int32 tensor of shape (num_tokens, top_k) where each
|
|
element encodes the expert id in the upper 16 bits and the bf16 routing
|
|
weight bits in the lower 16 bits, matching FlashInfer's packed layout.
|
|
"""
|
|
|
|
packed_topk_ids: torch.Tensor
|
|
router_logits: torch.Tensor
|
|
|
|
@property
|
|
def format(self) -> TopKOutputFormat:
|
|
return TopKOutputFormat.PACKED
|
|
|
|
|
|
def _make_round_robin_expert_ids(
|
|
num_tokens: int,
|
|
topk: int,
|
|
num_experts: int,
|
|
*,
|
|
device: torch.device,
|
|
dtype: torch.dtype,
|
|
layer_id: Optional[int] = None,
|
|
) -> torch.Tensor:
|
|
if topk == 0:
|
|
return torch.empty((num_tokens, 0), device=device, dtype=dtype)
|
|
|
|
step = max(num_experts // topk, 1)
|
|
layer_offset = 0 if layer_id is None else layer_id
|
|
offsets = torch.arange(num_tokens, device=device, dtype=dtype).unsqueeze(1)
|
|
steps = torch.arange(topk, device=device, dtype=dtype).unsqueeze(0) * step
|
|
return (offsets + layer_offset + steps) % num_experts
|
|
|
|
|
|
# -------------------------------- TopK ---------------------------------------
|
|
|
|
|
|
class TopK(BaseFusedOp):
|
|
"""
|
|
Parameters:
|
|
--top_k: The all number of top experts selected per token, including the fused shared expert(s).
|
|
--num_fused_shared_experts: num of shared experts, can be activate both in TP or EP mode.
|
|
--routed_scaling_factor: the scaling factor for routed experts in topk_weights.
|
|
--fused_shared_experts_scaling_factor: scaling factor applied to the fused shared experts'
|
|
topk weight (models pass 1/ep_size under standard EP, where the per-rank shared-expert
|
|
outputs are all-reduced).
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
top_k: int,
|
|
*,
|
|
layer_id: Optional[int] = None,
|
|
use_grouped_topk: bool = False,
|
|
topk_group: Optional[int] = None,
|
|
num_expert_group: Optional[int] = None,
|
|
renormalize: bool = True,
|
|
num_fused_shared_experts: int = 0,
|
|
custom_routing_function: Optional[Callable] = None,
|
|
scoring_func: str = "softmax",
|
|
correction_bias: Optional[torch.Tensor] = None,
|
|
quant_config: Optional[QuantizationConfig] = None,
|
|
routed_scaling_factor: Optional[float] = None,
|
|
apply_routed_scaling_factor_on_output: Optional[bool] = False,
|
|
output_format: Optional[TopKOutputFormat] = None,
|
|
fused_shared_experts_scaling_factor: Optional[float] = None,
|
|
is_fp4_experts: bool = False,
|
|
allow_routed_experts_capture: bool = True,
|
|
):
|
|
# NOTE: scoring_func is not used for now, but we keep it for future use
|
|
# see https://github.com/sgl-project/sglang/pull/4505 for more details
|
|
super().__init__()
|
|
|
|
if use_grouped_topk:
|
|
assert num_expert_group is not None and topk_group is not None
|
|
|
|
self.layer_id = layer_id
|
|
|
|
self.enable_waterfill = (
|
|
num_fused_shared_experts > 0 and get_exec().moe.enable_waterfill
|
|
)
|
|
|
|
self.waterfill_balancer = None
|
|
if self.enable_waterfill:
|
|
# TODO(ch-wan): Refactor shared-expert fusion and routed TopK fusion.
|
|
top_k -= num_fused_shared_experts
|
|
num_fused_shared_experts = 0
|
|
output_format = TopKOutputFormat.STANDARD
|
|
|
|
# Under the flashinfer_mxfp4 backend, fp4-expert ckpts take STANDARD
|
|
# (consumes topk_ids/weights); otherwise BYPASSED. No-op on other backends.
|
|
self.is_fp4_experts = is_fp4_experts
|
|
self.topk_config = TopKConfig(
|
|
top_k=top_k,
|
|
use_grouped_topk=use_grouped_topk,
|
|
renormalize=renormalize,
|
|
topk_group=topk_group,
|
|
num_expert_group=num_expert_group,
|
|
num_fused_shared_experts=num_fused_shared_experts,
|
|
custom_routing_function=custom_routing_function,
|
|
correction_bias=correction_bias,
|
|
routed_scaling_factor=routed_scaling_factor,
|
|
apply_routed_scaling_factor_on_output=apply_routed_scaling_factor_on_output,
|
|
fused_shared_experts_scaling_factor=fused_shared_experts_scaling_factor,
|
|
output_format=output_format,
|
|
scoring_func=scoring_func,
|
|
allow_routed_experts_capture=allow_routed_experts_capture,
|
|
)
|
|
|
|
def _apply_waterfill(self, topk_output: TopKOutput, num_tokens: int) -> TopKOutput:
|
|
if self.enable_waterfill and self.waterfill_balancer is None:
|
|
raise RuntimeError(
|
|
"Waterfill TopK must be prepared by ModelRunner before forward."
|
|
)
|
|
if self.waterfill_balancer is None:
|
|
return topk_output
|
|
assert TopKOutputChecker.format_is_standard(topk_output)
|
|
return self.waterfill_balancer.expand_topk(topk_output, num_tokens)
|
|
|
|
def forward_musa(self, *args, **kwargs) -> TopKOutput:
|
|
# MUSA follows the CUDA path explicitly: select_experts branches on
|
|
# _is_musa internally (hardware_backend.musa topk kernels), so the
|
|
# native path would bypass them.
|
|
return self.forward_cuda(*args, **kwargs)
|
|
|
|
def _torch_compile_forward(self, num_tokens: int) -> Optional[Callable]:
|
|
# torch.compile of the native TopK only pays off at bs=1; for larger
|
|
# batches keep the current optimized dispatch (see MultiPlatformOp
|
|
# history: the compiled path regressed bs > 1).
|
|
return self.forward_native if num_tokens == 1 else None
|
|
|
|
def forward_native(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
router_logits: torch.Tensor,
|
|
*,
|
|
num_token_non_padded: Optional[torch.Tensor] = None,
|
|
expert_location_dispatch_info: Optional[ExpertLocationDispatchInfo] = None,
|
|
) -> TopKOutput:
|
|
self.topk_config.torch_native = True
|
|
topk_output = select_experts(
|
|
hidden_states=hidden_states,
|
|
layer_id=self.layer_id,
|
|
router_logits=router_logits,
|
|
topk_config=self.topk_config,
|
|
num_token_non_padded=num_token_non_padded,
|
|
expert_location_dispatch_info=expert_location_dispatch_info,
|
|
)
|
|
return self._apply_waterfill(topk_output, hidden_states.shape[0])
|
|
|
|
def forward_cuda(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
router_logits: torch.Tensor,
|
|
*,
|
|
num_token_non_padded: Optional[torch.Tensor] = None,
|
|
expert_location_dispatch_info: Optional[ExpertLocationDispatchInfo] = None,
|
|
) -> TopKOutput:
|
|
if self.topk_config.output_format is not None:
|
|
output_format = self.topk_config.output_format
|
|
elif get_moe_runner_backend().is_triton_kernels():
|
|
output_format = TopKOutputFormat.TRITON_KERNEL
|
|
# ===== TO BE REFACTORED ====
|
|
elif get_moe_runner_backend().is_experimental_sgl_trtllm():
|
|
try:
|
|
|
|
use_standard_for_lora = bool(get_lora().enable_lora)
|
|
except ValueError:
|
|
use_standard_for_lora = False
|
|
output_format = (
|
|
TopKOutputFormat.STANDARD
|
|
if use_standard_for_lora
|
|
else TopKOutputFormat.BYPASSED
|
|
)
|
|
# ===== END TO BE REFACTORED ====
|
|
elif get_moe_runner_backend().is_flashinfer_trtllm() or (
|
|
get_moe_runner_backend().is_flashinfer_mxfp4() and not self.is_fp4_experts
|
|
):
|
|
output_format = TopKOutputFormat.BYPASSED
|
|
else:
|
|
output_format = TopKOutputFormat.STANDARD
|
|
|
|
if output_format == TopKOutputFormat.TRITON_KERNEL:
|
|
# renormalize=True is equivalent to sm_first=False
|
|
(
|
|
a_ragged_metadata,
|
|
gather_idx,
|
|
scatter_idx,
|
|
gate_scal,
|
|
n_expts_act,
|
|
) = routing(
|
|
router_logits,
|
|
self.topk_config.top_k,
|
|
sm_first=not self.topk_config.renormalize,
|
|
)
|
|
return TritonKernelTopKOutput(
|
|
a_ragged_metadata,
|
|
gather_idx,
|
|
scatter_idx,
|
|
gate_scal,
|
|
n_expts_act,
|
|
)
|
|
elif output_format == TopKOutputFormat.BYPASSED:
|
|
return BypassedTopKOutput(
|
|
hidden_states=hidden_states,
|
|
router_logits=router_logits,
|
|
topk_config=self.topk_config,
|
|
num_token_non_padded=num_token_non_padded,
|
|
expert_location_dispatch_info=expert_location_dispatch_info,
|
|
)
|
|
else:
|
|
self.topk_config.torch_native = False
|
|
with use_symmetric_memory(
|
|
get_tp_group(), disabled=not is_allocation_symmetric()
|
|
):
|
|
topk_output = select_experts(
|
|
hidden_states=hidden_states,
|
|
layer_id=self.layer_id,
|
|
router_logits=router_logits,
|
|
topk_config=self.topk_config,
|
|
num_token_non_padded=num_token_non_padded,
|
|
expert_location_dispatch_info=expert_location_dispatch_info,
|
|
)
|
|
return self._apply_waterfill(topk_output, hidden_states.shape[0])
|
|
|
|
def forward_cpu(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
router_logits: torch.Tensor,
|
|
*,
|
|
num_token_non_padded: Optional[torch.Tensor] = None,
|
|
expert_location_dispatch_info: Optional[ExpertLocationDispatchInfo] = None,
|
|
) -> TopKOutput:
|
|
topk_output = select_experts(
|
|
hidden_states=hidden_states,
|
|
layer_id=self.layer_id,
|
|
router_logits=router_logits,
|
|
topk_config=self.topk_config,
|
|
num_token_non_padded=num_token_non_padded,
|
|
expert_location_dispatch_info=expert_location_dispatch_info,
|
|
)
|
|
return self._apply_waterfill(topk_output, hidden_states.shape[0])
|
|
|
|
def forward_npu(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
router_logits: torch.Tensor,
|
|
*,
|
|
num_token_non_padded: Optional[torch.Tensor] = None,
|
|
expert_location_dispatch_info: Optional[ExpertLocationDispatchInfo] = None,
|
|
) -> TopKOutput:
|
|
|
|
from sglang.srt.hardware_backend.npu.moe.topk import fused_topk_npu
|
|
|
|
return fused_topk_npu(
|
|
hidden_states=hidden_states,
|
|
router_logits=router_logits,
|
|
topk_config=self.topk_config,
|
|
num_token_non_padded=num_token_non_padded,
|
|
expert_location_dispatch_info=expert_location_dispatch_info,
|
|
layer_id=self.layer_id,
|
|
)
|
|
|
|
def empty_topk_output(
|
|
self, device: torch.device, *, layer_id: Optional[int] = None
|
|
) -> TopKOutput:
|
|
"""Return an empty topk output for a rank with zero tokens this forward.
|
|
|
|
When ``layer_id`` is provided and the active dispatch algorithm is LP,
|
|
also calls ``LPLBSolver.solve(empty)`` so that this rank participates
|
|
in the EP all-reduce. Without this, an empty rank would skip the
|
|
collective and deadlock under DP-attention.
|
|
"""
|
|
if layer_id is not None:
|
|
# Skip the full ExpertLocationDispatchInfo allocation — we only
|
|
# need the per-layer solver to participate in the EP all-reduce.
|
|
from sglang.srt.eplb.lplb_solver import get_global_lplb_solver
|
|
|
|
lplb_solver = get_global_lplb_solver(layer_id)
|
|
if lplb_solver is not None:
|
|
lplb_solver.solve(
|
|
torch.empty(
|
|
(0, self.topk_config.top_k),
|
|
dtype=torch.int32,
|
|
device=device,
|
|
)
|
|
)
|
|
topk = self.topk_config.top_k - self.topk_config.num_fused_shared_experts
|
|
with use_symmetric_memory(
|
|
get_tp_group(), disabled=not is_allocation_symmetric()
|
|
):
|
|
topk_weights = torch.empty((0, topk), dtype=torch.float32, device=device)
|
|
topk_ids = torch.full((0, topk), -1, dtype=torch.int32, device=device)
|
|
# FIXME: router_logits should be of size (0, num_experts)
|
|
router_logits = torch.empty((0, topk), dtype=torch.float32, device=device)
|
|
topk_output = StandardTopKOutput(topk_weights, topk_ids, router_logits)
|
|
if has_per_rank_fused_shared_slots(self.topk_config.num_fused_shared_experts):
|
|
n = self.topk_config.num_fused_shared_experts
|
|
topk_output = topk_output._replace(
|
|
topk_ids=topk_output.topk_ids.new_empty(
|
|
(0, topk_output.topk_ids.shape[-1] + n)
|
|
),
|
|
topk_weights=topk_output.topk_weights.new_empty(
|
|
(0, topk_output.topk_weights.shape[-1] + n)
|
|
),
|
|
)
|
|
return self._apply_waterfill(topk_output, 0)
|
|
|
|
def forward_xpu(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
router_logits: torch.Tensor,
|
|
*,
|
|
num_token_non_padded: Optional[torch.Tensor] = None,
|
|
expert_location_dispatch_info: Optional[ExpertLocationDispatchInfo] = None,
|
|
) -> TopKOutput:
|
|
self.topk_config.torch_native = True
|
|
# [NOTE] XPU device support for topk kernels
|
|
# - support 'topk_softmax' and 'topk_sigmoid'
|
|
# - support up to 8 top-k and 256 experts
|
|
self.topk_config.torch_native = not (
|
|
self.topk_config.top_k <= 8 and router_logits.shape[1] <= 256
|
|
)
|
|
|
|
return select_experts(
|
|
hidden_states=hidden_states,
|
|
layer_id=self.layer_id,
|
|
router_logits=router_logits,
|
|
topk_config=self.topk_config,
|
|
num_token_non_padded=num_token_non_padded,
|
|
expert_location_dispatch_info=expert_location_dispatch_info,
|
|
)
|
|
|
|
|
|
# ------------------------------- TopK implementation -------------------------------------
|
|
|
|
|
|
def fused_topk_torch_native(
|
|
hidden_states: torch.Tensor,
|
|
gating_output: torch.Tensor,
|
|
topk: int,
|
|
renormalize: bool,
|
|
correction_bias: torch.Tensor = None,
|
|
scoring_func: str = "softmax",
|
|
):
|
|
def scoring_func_impl(gating_output: torch.Tensor) -> torch.Tensor:
|
|
if scoring_func == "softmax":
|
|
return gating_output.softmax(dim=-1)
|
|
elif scoring_func == "sigmoid":
|
|
return gating_output.sigmoid()
|
|
elif scoring_func == "sqrtsoftplus":
|
|
return F.softplus(gating_output).sqrt()
|
|
else:
|
|
raise ValueError(f"Invalid scoring function: {scoring_func}")
|
|
|
|
if correction_bias is not None:
|
|
n_routed_experts = gating_output.shape[-1]
|
|
scores = scoring_func_impl(gating_output)
|
|
scores_for_choice = scores.view(
|
|
-1, n_routed_experts
|
|
) + correction_bias.unsqueeze(0)
|
|
topk_ids = torch.topk(scores_for_choice, k=topk, dim=-1, sorted=False)[1]
|
|
topk_weights = scores.gather(1, topk_ids)
|
|
else:
|
|
assert (
|
|
hidden_states.shape[0] == gating_output.shape[0]
|
|
), f"Number of tokens mismatch, {hidden_states.shape=} vs {gating_output.shape=}"
|
|
M, _ = hidden_states.shape
|
|
topk_weights = torch.empty(
|
|
M, topk, dtype=torch.float32, device=hidden_states.device
|
|
)
|
|
topk_ids = torch.empty(M, topk, dtype=torch.int32, device=hidden_states.device)
|
|
topk_weights = scoring_func_impl(gating_output.float())
|
|
topk_weights, topk_ids = torch.topk(topk_weights, topk, dim=-1)
|
|
|
|
if renormalize:
|
|
# fp32 like the reference gate (the epsilon is not representable in
|
|
# fp16); the sum dtype and the division's type promotion upcast inside
|
|
# the existing kernels, so no extra cast launch is needed
|
|
topk_weights = topk_weights / (
|
|
topk_weights.sum(dim=-1, keepdim=True, dtype=torch.float32)
|
|
+ _RENORMALIZE_SUM_EPSILON
|
|
)
|
|
return topk_weights, topk_ids
|
|
|
|
|
|
def fused_topk_softmax_torch_raw_logits(
|
|
hidden_states: torch.Tensor,
|
|
gating_output: torch.Tensor,
|
|
topk: int,
|
|
renormalize: bool,
|
|
):
|
|
assert (
|
|
hidden_states.shape[0] == gating_output.shape[0]
|
|
), f"Number of tokens mismatch, {hidden_states.shape=} vs {gating_output.shape=}"
|
|
|
|
_, topk_ids = torch.topk(gating_output, k=topk, dim=-1, sorted=False)
|
|
logits = gating_output.float()
|
|
topk_weights = logits.gather(1, topk_ids)
|
|
if renormalize:
|
|
topk_weights = F.softmax(topk_weights, dim=-1, dtype=torch.float32)
|
|
|
|
return topk_weights.to(torch.float32), topk_ids.to(torch.int32)
|
|
|
|
|
|
def fused_topk_cpu(
|
|
hidden_states: torch.Tensor,
|
|
gating_output: torch.Tensor,
|
|
topk: int,
|
|
renormalize: bool,
|
|
correction_bias: torch.Tensor = None,
|
|
scoring_func: str = "softmax",
|
|
routed_scaling_factor: Optional[float] = None,
|
|
apply_routed_scaling_factor_on_output: Optional[bool] = False,
|
|
num_fused_shared_experts: int = 0,
|
|
packed_out: Optional[torch.Tensor] = None,
|
|
num_token_non_padded: Optional[torch.Tensor] = None,
|
|
):
|
|
if num_fused_shared_experts != 0:
|
|
raise ValueError(
|
|
f"num_fused_shared_experts must be 0 for CPU fused topk, got: {num_fused_shared_experts}"
|
|
)
|
|
if apply_routed_scaling_factor_on_output:
|
|
raise ValueError(
|
|
"apply_routed_scaling_factor_on_output is not supported for CPU fused topk"
|
|
)
|
|
if packed_out is not None:
|
|
raise ValueError("packed_out is not supported for CPU fused topk")
|
|
if num_token_non_padded is not None:
|
|
raise ValueError("num_token_non_padded is not supported for CPU fused topk")
|
|
if correction_bias is not None and correction_bias.dtype != torch.float32:
|
|
correction_bias = correction_bias.to(torch.float32)
|
|
|
|
if scoring_func == "softmax":
|
|
topk_weights, topk_ids = torch.ops.sgl_kernel.topk_softmax_cpu(
|
|
hidden_states=hidden_states,
|
|
gating_output=gating_output,
|
|
topk=topk,
|
|
renormalize=renormalize,
|
|
correction_bias=correction_bias,
|
|
)
|
|
elif scoring_func == "sigmoid":
|
|
topk_weights, topk_ids = torch.ops.sgl_kernel.topk_sigmoid_cpu(
|
|
hidden_states=hidden_states,
|
|
gating_output=gating_output,
|
|
topk=topk,
|
|
renormalize=renormalize,
|
|
correction_bias=correction_bias,
|
|
)
|
|
else:
|
|
# Fall back to the torch-native impl for the rest
|
|
return fused_topk_torch_native(
|
|
hidden_states,
|
|
gating_output,
|
|
topk,
|
|
renormalize,
|
|
correction_bias=correction_bias,
|
|
scoring_func=scoring_func,
|
|
)
|
|
|
|
return topk_weights, topk_ids
|
|
|
|
|
|
def apply_topk_weights_cpu(need_apply, topk_weights, inputs):
|
|
if not need_apply:
|
|
return inputs, topk_weights
|
|
|
|
# TODO: fuse below processing in fused_experts_cpu kernel
|
|
inputs = inputs * topk_weights.to(inputs.dtype)
|
|
topk_weights = torch.ones_like(
|
|
topk_weights, dtype=torch.float32
|
|
) # clear topk_weights as already applied
|
|
|
|
return inputs, topk_weights
|
|
|
|
|
|
def fused_topk(
|
|
hidden_states: torch.Tensor,
|
|
gating_output: torch.Tensor,
|
|
topk: int,
|
|
renormalize: bool,
|
|
correction_bias: Optional[torch.Tensor] = None,
|
|
scoring_func: str = "softmax",
|
|
routed_scaling_factor: Optional[float] = None,
|
|
apply_routed_scaling_factor_on_output: Optional[bool] = False,
|
|
num_fused_shared_experts: int = 0,
|
|
packed_out: Optional[torch.Tensor] = None,
|
|
num_token_non_padded: Optional[torch.Tensor] = None,
|
|
):
|
|
assert hidden_states.shape[0] == gating_output.shape[0], "Number of tokens mismatch"
|
|
|
|
M, _ = hidden_states.shape
|
|
|
|
topk_weights = torch.empty(
|
|
M, topk, dtype=torch.float32, device=hidden_states.device
|
|
)
|
|
topk_ids = torch.empty(M, topk, dtype=torch.int32, device=hidden_states.device)
|
|
|
|
if scoring_func == "softmax":
|
|
if _use_aiter:
|
|
|
|
# Use fused_topk instead of topk_softmax to auto dispatch to the correct kernel
|
|
topk_weights, topk_ids = aiter_fused_topk(
|
|
hidden_states,
|
|
gating_output,
|
|
topk,
|
|
renormalize,
|
|
topk_ids=topk_ids,
|
|
topk_weights=topk_weights,
|
|
)
|
|
# ===== TO BE REFACTORED ====
|
|
elif packed_out is not None:
|
|
# Fused gating + routed pack (SGLANG_OPT_LORA_FUSED_TOPK_PACK): one JIT kernel
|
|
# writes topk_weights/topk_ids AND the FlashInfer packed topk in one launch.
|
|
from sglang.kernels.ops.moe.trtllm_lora_temp.topk_softmax_pack import (
|
|
topk_softmax_pack,
|
|
)
|
|
|
|
topk_softmax_pack(
|
|
topk_weights,
|
|
topk_ids,
|
|
packed_out,
|
|
gating_output,
|
|
renormalize,
|
|
num_token_non_padded=num_token_non_padded,
|
|
)
|
|
# ===== END TO BE REFACTORED ====
|
|
elif _is_cuda:
|
|
# Unified Triton router (subsumes the AOT topk_softmax CUDA kernel).
|
|
from sglang.kernels.ops.moe.moe_fused_gate import (
|
|
moe_fused_gate as _jit_moe_fused_gate,
|
|
)
|
|
|
|
zero_bias = torch.zeros(
|
|
gating_output.shape[1],
|
|
dtype=torch.float32,
|
|
device=gating_output.device,
|
|
)
|
|
topk_weights, topk_ids = _jit_moe_fused_gate(
|
|
gating_output,
|
|
zero_bias,
|
|
topk,
|
|
scoring_func="softmax",
|
|
renormalize=renormalize,
|
|
)
|
|
else:
|
|
topk_softmax(
|
|
topk_weights,
|
|
topk_ids,
|
|
gating_output,
|
|
renormalize,
|
|
)
|
|
elif scoring_func == "sigmoid":
|
|
if _use_aiter and correction_bias is not None:
|
|
aiter_biased_grouped_topk(
|
|
gating_output,
|
|
correction_bias.to(dtype=gating_output.dtype),
|
|
topk_weights,
|
|
topk_ids,
|
|
num_expert_group=1,
|
|
topk_group=1,
|
|
need_renorm=renormalize,
|
|
)
|
|
if apply_routed_scaling_factor_on_output:
|
|
topk_weights *= (
|
|
routed_scaling_factor if routed_scaling_factor is not None else 1.0
|
|
)
|
|
elif _is_cuda:
|
|
# Unified Triton router (subsumes the AOT topk_sigmoid CUDA kernel).
|
|
from sglang.kernels.ops.moe.moe_fused_gate import (
|
|
moe_fused_gate as _jit_moe_fused_gate,
|
|
)
|
|
|
|
bias_fp32 = (
|
|
correction_bias.to(torch.float32)
|
|
if correction_bias is not None
|
|
else torch.zeros(
|
|
gating_output.shape[1],
|
|
dtype=torch.float32,
|
|
device=gating_output.device,
|
|
)
|
|
)
|
|
topk_weights, topk_ids = _jit_moe_fused_gate(
|
|
gating_output,
|
|
bias_fp32,
|
|
topk,
|
|
scoring_func="sigmoid",
|
|
renormalize=renormalize,
|
|
routed_scaling_factor=routed_scaling_factor,
|
|
apply_routed_scaling_factor_on_output=apply_routed_scaling_factor_on_output,
|
|
)
|
|
else:
|
|
if num_fused_shared_experts > 1:
|
|
raise ValueError(
|
|
"sigmoid topk supports at most one fused shared expert"
|
|
)
|
|
scale = (
|
|
routed_scaling_factor
|
|
if (
|
|
apply_routed_scaling_factor_on_output
|
|
and routed_scaling_factor is not None
|
|
)
|
|
else 1.0
|
|
)
|
|
topk_sigmoid(
|
|
topk_weights,
|
|
topk_ids,
|
|
gating_output,
|
|
renormalize,
|
|
correction_bias,
|
|
scale,
|
|
num_fused_shared_experts,
|
|
)
|
|
else:
|
|
raise ValueError(f"Invalid scoring function: {scoring_func}")
|
|
|
|
return topk_weights, topk_ids
|
|
|
|
|
|
# This is used by the Deepseek V2/V3/R1 series models
|
|
@torch.compile(dynamic=True, backend=get_compiler_backend(), disable=_is_npu)
|
|
def grouped_topk_gpu(
|
|
hidden_states: torch.Tensor,
|
|
gating_output: torch.Tensor,
|
|
topk: int,
|
|
renormalize: bool,
|
|
num_expert_group: Optional[int] = None,
|
|
topk_group: Optional[int] = None,
|
|
num_fused_shared_experts: int = 0,
|
|
routed_scaling_factor: Optional[float] = None,
|
|
apply_routed_scaling_factor_on_output: Optional[bool] = False,
|
|
scoring_func: str = "softmax",
|
|
):
|
|
assert hidden_states.shape[0] == gating_output.shape[0], "Number of tokens mismatch"
|
|
|
|
if scoring_func == "softmax":
|
|
scores = torch.softmax(gating_output, dim=-1)
|
|
elif scoring_func == "sigmoid":
|
|
scores = gating_output.sigmoid()
|
|
else:
|
|
raise ValueError(f"Unsupported scoring function: {scoring_func}")
|
|
|
|
num_token = scores.shape[0]
|
|
num_experts = scores.shape[1]
|
|
group_scores = (
|
|
scores.view(num_token, num_expert_group, -1).max(dim=-1).values
|
|
) # [n, n_group]
|
|
group_idx = torch.topk(group_scores, k=topk_group, dim=-1, sorted=False)[
|
|
1
|
|
] # [n, top_k_group]
|
|
group_mask = torch.zeros_like(group_scores) # [n, n_group]
|
|
group_mask.scatter_(1, group_idx, 1) # [n, n_group]
|
|
score_mask = (
|
|
group_mask.unsqueeze(-1)
|
|
.expand(num_token, num_expert_group, scores.shape[-1] // num_expert_group)
|
|
.reshape(num_token, -1)
|
|
) # [n, e]
|
|
tmp_scores = scores.masked_fill(~score_mask.bool(), 0.0) # [n, e]
|
|
topk_weights, topk_ids = torch.topk(
|
|
tmp_scores,
|
|
k=topk,
|
|
dim=-1,
|
|
sorted=(True if num_fused_shared_experts > 0 else False),
|
|
)
|
|
if num_fused_shared_experts:
|
|
topk_ids[:, -1] = torch.randint(
|
|
low=num_experts,
|
|
high=num_experts + num_fused_shared_experts,
|
|
size=(topk_ids.size(0),),
|
|
dtype=topk_ids.dtype,
|
|
device=topk_ids.device,
|
|
)
|
|
if routed_scaling_factor is not None:
|
|
topk_weights[:, -1] = (
|
|
topk_weights[:, :-1].sum(dim=-1) / routed_scaling_factor
|
|
)
|
|
|
|
if renormalize:
|
|
# fp32 like the reference gate (the epsilon is not representable in
|
|
# fp16); the sum dtype and the division's type promotion upcast inside
|
|
# the existing kernels, so no extra cast launch is needed
|
|
topk_weights_sum = (
|
|
topk_weights.sum(dim=-1, keepdim=True, dtype=torch.float32)
|
|
if num_fused_shared_experts == 0
|
|
else topk_weights[:, :-1].sum(dim=-1, keepdim=True, dtype=torch.float32)
|
|
)
|
|
topk_weights = topk_weights / (topk_weights_sum + _RENORMALIZE_SUM_EPSILON)
|
|
if apply_routed_scaling_factor_on_output:
|
|
topk_weights *= routed_scaling_factor
|
|
|
|
topk_weights, topk_ids = topk_weights.to(torch.float32), topk_ids.to(torch.int32)
|
|
|
|
return topk_weights, topk_ids
|
|
|
|
|
|
def grouped_topk_cpu(
|
|
hidden_states: torch.Tensor,
|
|
gating_output: torch.Tensor,
|
|
topk: int,
|
|
renormalize: bool,
|
|
num_expert_group: Optional[int] = None,
|
|
topk_group: Optional[int] = None,
|
|
num_fused_shared_experts: int = 0,
|
|
routed_scaling_factor: Optional[float] = None,
|
|
apply_routed_scaling_factor_on_output: Optional[bool] = False,
|
|
scoring_func: str = "softmax",
|
|
):
|
|
assert not apply_routed_scaling_factor_on_output
|
|
if scoring_func != "softmax":
|
|
raise ValueError(f"Unsupported scoring function: {scoring_func}")
|
|
|
|
return torch.ops.sgl_kernel.grouped_topk_cpu(
|
|
hidden_states,
|
|
gating_output,
|
|
topk,
|
|
renormalize,
|
|
num_expert_group,
|
|
topk_group,
|
|
num_fused_shared_experts,
|
|
routed_scaling_factor,
|
|
# num_token_non_padded must be None since it is not supported in kernel
|
|
num_token_non_padded=None,
|
|
)
|
|
|
|
|
|
def grouped_topk_xpu(
|
|
hidden_states: torch.Tensor,
|
|
gating_output: torch.Tensor,
|
|
topk: int,
|
|
renormalize: bool,
|
|
num_expert_group: Optional[int] = None,
|
|
topk_group: Optional[int] = None,
|
|
num_fused_shared_experts: int = 0,
|
|
routed_scaling_factor: Optional[float] = None,
|
|
apply_routed_scaling_factor_on_output: Optional[bool] = False,
|
|
scoring_func: str = "softmax",
|
|
):
|
|
num_experts = gating_output.shape[1]
|
|
experts_per_group = (
|
|
num_experts // num_expert_group if num_expert_group else num_experts
|
|
)
|
|
|
|
# moe_fused_gate kernel ensures that num_experts/num_expert_group does not exceed MAX_VPT=32 now.
|
|
if experts_per_group <= 32 and is_power_of_two(num_experts):
|
|
from sgl_kernel import moe_fused_gate
|
|
|
|
return moe_fused_gate(
|
|
gating_output.to(torch.float32),
|
|
None, # without bias
|
|
num_expert_group,
|
|
topk_group,
|
|
topk,
|
|
renormalize=renormalize,
|
|
scoring_func=scoring_func,
|
|
num_fused_shared_experts=num_fused_shared_experts,
|
|
routed_scaling_factor=(
|
|
routed_scaling_factor if routed_scaling_factor is not None else 1.0
|
|
),
|
|
apply_routed_scaling_factor_on_output=bool(
|
|
apply_routed_scaling_factor_on_output
|
|
),
|
|
)
|
|
|
|
# use default implementation
|
|
return grouped_topk_gpu(
|
|
hidden_states,
|
|
gating_output,
|
|
topk,
|
|
renormalize,
|
|
num_expert_group,
|
|
topk_group,
|
|
num_fused_shared_experts,
|
|
routed_scaling_factor,
|
|
apply_routed_scaling_factor_on_output,
|
|
scoring_func,
|
|
)
|
|
|
|
|
|
@torch.compile(dynamic=True, backend=get_compiler_backend(), disable=_is_npu)
|
|
def kimi_k2_biased_topk_impl(
|
|
hidden_states: torch.Tensor,
|
|
gating_output: torch.Tensor,
|
|
correction_bias: torch.Tensor,
|
|
topk: int,
|
|
renormalize: bool,
|
|
routed_scaling_factor: Optional[float] = None,
|
|
apply_routed_scaling_factor_on_output: Optional[bool] = False,
|
|
):
|
|
"""
|
|
Optimized version for num_expert_group=1 case (e.g., Kimi K2 with 384 experts).
|
|
Simplifies the grouped topk logic by removing unnecessary group masking operations.
|
|
Note: This function assumes num_fused_shared_experts=0.
|
|
"""
|
|
assert hidden_states.shape[0] == gating_output.shape[0], "Number of tokens mismatch"
|
|
|
|
scores = gating_output.sigmoid()
|
|
num_token = scores.shape[0]
|
|
|
|
# When num_expert_group=1, no need for group masking
|
|
# Directly compute scores with correction bias
|
|
tmp_scores = scores.view(num_token, -1) + correction_bias.unsqueeze(0)
|
|
|
|
# Directly select topk experts (no need to sort since num_fused_shared_experts=0)
|
|
_, topk_ids = torch.topk(tmp_scores, k=topk, dim=-1, sorted=False)
|
|
topk_weights = scores.gather(1, topk_ids)
|
|
|
|
if renormalize:
|
|
# fp32 like the reference gate (the epsilon is not representable in
|
|
# fp16); the sum dtype and the division's type promotion upcast inside
|
|
# the existing kernels, so no extra cast launch is needed
|
|
topk_weights_sum = topk_weights.sum(dim=-1, keepdim=True, dtype=torch.float32)
|
|
topk_weights = topk_weights / (topk_weights_sum + _RENORMALIZE_SUM_EPSILON)
|
|
if apply_routed_scaling_factor_on_output:
|
|
topk_weights *= routed_scaling_factor
|
|
|
|
topk_weights, topk_ids = topk_weights.to(torch.float32), topk_ids.to(torch.int32)
|
|
return topk_weights, topk_ids
|
|
|
|
|
|
@torch.compile(dynamic=True, backend=get_compiler_backend(), disable=_is_npu)
|
|
def biased_topk_impl(
|
|
hidden_states: torch.Tensor,
|
|
gating_output: torch.Tensor,
|
|
correction_bias: torch.Tensor,
|
|
topk: int,
|
|
renormalize: bool,
|
|
scoring_func: str = "sigmoid",
|
|
num_fused_shared_experts: int = 0,
|
|
routed_scaling_factor: Optional[float] = None,
|
|
num_token_non_padded: Optional[torch.Tensor] = None,
|
|
expert_location_dispatch_info: Optional[ExpertLocationDispatchInfo] = None,
|
|
apply_routed_scaling_factor_on_output: Optional[bool] = False,
|
|
):
|
|
assert hidden_states.shape[0] == gating_output.shape[0], "Number of tokens mismatch"
|
|
|
|
if scoring_func == "sigmoid":
|
|
scores = gating_output.sigmoid()
|
|
elif scoring_func == "sqrtsoftplus":
|
|
scores = torch.nn.functional.softplus(gating_output).sqrt()
|
|
|
|
num_token = scores.shape[0]
|
|
num_experts = scores.shape[1]
|
|
|
|
scores_for_choice = scores.view(num_token, -1) + correction_bias.unsqueeze(0)
|
|
_, topk_ids = torch.topk(
|
|
scores_for_choice,
|
|
k=topk,
|
|
dim=-1,
|
|
sorted=(True if num_fused_shared_experts > 0 else False),
|
|
)
|
|
topk_weights = scores.gather(1, topk_ids)
|
|
|
|
if num_fused_shared_experts:
|
|
topk_ids[:, -1] = torch.randint(
|
|
low=num_experts,
|
|
high=num_experts + num_fused_shared_experts,
|
|
size=(topk_ids.size(0),),
|
|
dtype=topk_ids.dtype,
|
|
device=topk_ids.device,
|
|
)
|
|
if routed_scaling_factor is not None:
|
|
topk_weights[:, -1] = (
|
|
topk_weights[:, :-1].sum(dim=-1) / routed_scaling_factor
|
|
)
|
|
|
|
if renormalize:
|
|
# fp32 like the reference gate (the epsilon is not representable in
|
|
# fp16); the sum dtype and the division's type promotion upcast inside
|
|
# the existing kernels, so no extra cast launch is needed
|
|
topk_weights_sum = (
|
|
topk_weights.sum(dim=-1, keepdim=True, dtype=torch.float32)
|
|
if num_fused_shared_experts == 0
|
|
else topk_weights[:, :-1].sum(dim=-1, keepdim=True, dtype=torch.float32)
|
|
)
|
|
topk_weights = topk_weights / (topk_weights_sum + _RENORMALIZE_SUM_EPSILON)
|
|
if apply_routed_scaling_factor_on_output:
|
|
topk_weights *= routed_scaling_factor
|
|
|
|
topk_weights, topk_ids = topk_weights.to(torch.float32), topk_ids.to(torch.int32)
|
|
return topk_weights, topk_ids
|
|
|
|
|
|
def biased_topk_jit_kernel_impl(
|
|
hidden_states: torch.Tensor,
|
|
gating_output: torch.Tensor,
|
|
correction_bias: torch.Tensor,
|
|
topk: int,
|
|
renormalize: bool,
|
|
scoring_func: str = "sigmoid",
|
|
num_fused_shared_experts: int = 0,
|
|
routed_scaling_factor: Optional[float] = None,
|
|
num_token_non_padded: Optional[torch.Tensor] = None,
|
|
expert_location_dispatch_info: Optional[ExpertLocationDispatchInfo] = None,
|
|
apply_routed_scaling_factor_on_output: Optional[bool] = False,
|
|
) -> Tuple[torch.Tensor, torch.Tensor]:
|
|
assert hidden_states.shape[0] == gating_output.shape[0], "Number of tokens mismatch"
|
|
|
|
if _use_aiter and scoring_func == "sqrtsoftplus" and num_fused_shared_experts == 0:
|
|
from aiter import topk_gating
|
|
|
|
num_tokens = gating_output.shape[0]
|
|
topk_weights = torch.empty(
|
|
(num_tokens, topk), dtype=torch.float32, device=gating_output.device
|
|
)
|
|
topk_ids = torch.empty(
|
|
(num_tokens, topk), dtype=torch.int32, device=gating_output.device
|
|
)
|
|
|
|
topk_gating(
|
|
topk_weights,
|
|
topk_ids,
|
|
gating_output,
|
|
correction_bias,
|
|
renormalize,
|
|
routed_scaling_factor,
|
|
score_func="sqrtsoftplus",
|
|
)
|
|
|
|
return topk_weights, topk_ids
|
|
|
|
else:
|
|
from sglang.kernels.ops.moe.moe_fused_gate import moe_fused_gate
|
|
|
|
# DeepSeek-V4 stores e_score_correction_bias in bf16 (for the aiter
|
|
# sqrtsoftplus topk path). moe_fused_gate upcasts the bias to fp32
|
|
# in-register, so pass it through directly rather than allocating a fresh
|
|
# fp32 copy of this static routing bias on every MoE invocation.
|
|
topk_weights, topk_ids = moe_fused_gate(
|
|
gating_output,
|
|
correction_bias,
|
|
topk=topk,
|
|
scoring_func=scoring_func,
|
|
num_fused_shared_experts=num_fused_shared_experts,
|
|
renormalize=renormalize,
|
|
routed_scaling_factor=routed_scaling_factor,
|
|
apply_routed_scaling_factor_on_output=apply_routed_scaling_factor_on_output,
|
|
)
|
|
topk_weights, topk_ids = topk_weights.to(torch.float32), topk_ids.to(
|
|
torch.int32
|
|
)
|
|
return topk_weights, topk_ids
|
|
|
|
|
|
def biased_topk_xpu(
|
|
hidden_states: torch.Tensor,
|
|
gating_output: torch.Tensor,
|
|
correction_bias: torch.Tensor,
|
|
topk: int,
|
|
renormalize: bool,
|
|
scoring_func: str = "sigmoid",
|
|
num_fused_shared_experts: int = 0,
|
|
routed_scaling_factor: Optional[float] = None,
|
|
num_token_non_padded: Optional[torch.Tensor] = None,
|
|
expert_location_dispatch_info: Optional[ExpertLocationDispatchInfo] = None,
|
|
apply_routed_scaling_factor_on_output: Optional[bool] = False,
|
|
) -> Tuple[torch.Tensor, torch.Tensor]:
|
|
assert hidden_states.shape[0] == gating_output.shape[0], "Number of tokens mismatch"
|
|
|
|
num_rows, _ = gating_output.shape
|
|
device = gating_output.device
|
|
|
|
output = torch.empty(num_rows, topk, dtype=torch.float32, device=device)
|
|
indices = torch.empty(num_rows, topk, dtype=torch.int32, device=device)
|
|
|
|
from sgl_kernel import biased_topk
|
|
|
|
biased_topk(
|
|
gating_output,
|
|
correction_bias,
|
|
output,
|
|
indices,
|
|
topk,
|
|
scoring_func,
|
|
num_fused_shared_experts,
|
|
renormalize,
|
|
routed_scaling_factor=(routed_scaling_factor if routed_scaling_factor else 1.0),
|
|
apply_routed_scaling_factor_on_output=bool(
|
|
apply_routed_scaling_factor_on_output
|
|
),
|
|
)
|
|
|
|
return output, indices
|
|
|
|
|
|
@torch.compile(dynamic=True, backend=get_compiler_backend(), disable=_is_npu)
|
|
def biased_grouped_topk_impl(
|
|
hidden_states: torch.Tensor,
|
|
gating_output: torch.Tensor,
|
|
correction_bias: torch.Tensor,
|
|
topk: int,
|
|
renormalize: bool,
|
|
num_expert_group: Optional[int] = None,
|
|
topk_group: Optional[int] = None,
|
|
num_fused_shared_experts: int = 0,
|
|
routed_scaling_factor: Optional[float] = None,
|
|
apply_routed_scaling_factor_on_output: Optional[bool] = False,
|
|
):
|
|
assert hidden_states.shape[0] == gating_output.shape[0], "Number of tokens mismatch"
|
|
|
|
scores = gating_output.sigmoid()
|
|
num_token = scores.shape[0]
|
|
num_experts = scores.shape[1]
|
|
scores_for_choice = scores.view(num_token, -1) + correction_bias.unsqueeze(0)
|
|
group_scores = (
|
|
scores_for_choice.view(num_token, num_expert_group, -1)
|
|
.topk(2, dim=-1)[0]
|
|
.sum(dim=-1)
|
|
) # [n, n_group]
|
|
group_idx = torch.topk(group_scores, k=topk_group, dim=-1, sorted=False)[
|
|
1
|
|
] # [n, top_k_group]
|
|
group_mask = torch.zeros_like(group_scores) # [n, n_group]
|
|
group_mask.scatter_(1, group_idx, 1) # [n, n_group]
|
|
score_mask = (
|
|
group_mask.unsqueeze(-1)
|
|
.expand(num_token, num_expert_group, scores.shape[-1] // num_expert_group)
|
|
.reshape(num_token, -1)
|
|
) # [n, e]
|
|
tmp_scores = scores_for_choice.masked_fill(
|
|
~score_mask.bool(), float("-inf")
|
|
) # [n, e]
|
|
_, topk_ids = torch.topk(
|
|
tmp_scores,
|
|
k=topk,
|
|
dim=-1,
|
|
sorted=(True if num_fused_shared_experts > 0 else False),
|
|
)
|
|
topk_weights = scores.gather(1, topk_ids)
|
|
|
|
if num_fused_shared_experts:
|
|
topk_ids[:, -1] = torch.randint(
|
|
low=num_experts,
|
|
high=num_experts + num_fused_shared_experts,
|
|
size=(topk_ids.size(0),),
|
|
dtype=topk_ids.dtype,
|
|
device=topk_ids.device,
|
|
)
|
|
if routed_scaling_factor is not None:
|
|
topk_weights[:, -1] = (
|
|
topk_weights[:, :-1].sum(dim=-1) / routed_scaling_factor
|
|
)
|
|
|
|
if renormalize:
|
|
# fp32 like the reference gate (the epsilon is not representable in
|
|
# fp16); the sum dtype and the division's type promotion upcast inside
|
|
# the existing kernels, so no extra cast launch is needed
|
|
topk_weights_sum = (
|
|
topk_weights.sum(dim=-1, keepdim=True, dtype=torch.float32)
|
|
if num_fused_shared_experts == 0
|
|
else topk_weights[:, :-1].sum(dim=-1, keepdim=True, dtype=torch.float32)
|
|
)
|
|
topk_weights = topk_weights / (topk_weights_sum + _RENORMALIZE_SUM_EPSILON)
|
|
if apply_routed_scaling_factor_on_output:
|
|
topk_weights *= routed_scaling_factor
|
|
|
|
topk_weights, topk_ids = topk_weights.to(torch.float32), topk_ids.to(torch.int32)
|
|
|
|
return topk_weights, topk_ids
|
|
|
|
|
|
from sglang.kernels.ops.moe.fill_padded_rows import (
|
|
_can_fuse_padded_region,
|
|
_fill_padded_rows,
|
|
)
|
|
|
|
|
|
def is_power_of_two(n):
|
|
return n > 0 and math.log2(n).is_integer()
|
|
|
|
|
|
def _eplb_remap_enabled() -> bool:
|
|
# A real logical->physical mapping only exists when EPLB is enabled, the
|
|
# initial expert placement is non-trivial, or there are redundant physical
|
|
# experts. Otherwise the map is identity and the remap must be skipped (it is
|
|
# both unnecessary and not well-defined over the padded region of topk_ids).
|
|
from sglang.srt.runtime_context import get_server_args
|
|
|
|
try:
|
|
get_server_args() # probes that a config is published
|
|
except ValueError:
|
|
# Global server args are not initialized outside the server runtime
|
|
# (e.g. in unit tests that call select_experts directly). In that case
|
|
# there is no EPLB mapping, so the remap must be skipped.
|
|
return False
|
|
return (
|
|
get_exec().moe.enable_eplb
|
|
or get_exec().moe.init_expert_location != "trivial"
|
|
or get_exec().moe.ep_num_redundant_experts > 0
|
|
)
|
|
|
|
|
|
def _mask_topk_ids_padded_region(
|
|
topk_ids: torch.Tensor,
|
|
num_token_non_padded: Optional[torch.Tensor] = None,
|
|
fill_value: int = -1,
|
|
) -> None:
|
|
if num_token_non_padded is None:
|
|
return
|
|
# TODO: let the kernel support other dtypes
|
|
if _is_cuda and topk_ids.dtype == torch.int32 and fill_value == -1:
|
|
mask_topk_ids(topk_ids, num_token_non_padded)
|
|
elif _is_npu:
|
|
return
|
|
elif _can_fuse_padded_region(topk_ids):
|
|
_fill_padded_rows(topk_ids, num_token_non_padded, fill_value)
|
|
else:
|
|
indices = torch.arange(0, topk_ids.shape[0], device=topk_ids.device)
|
|
topk_ids[indices >= num_token_non_padded, :] = fill_value
|
|
|
|
|
|
def _zero_topk_weights_padded_region(
|
|
topk_weights: torch.Tensor,
|
|
num_token_non_padded: Optional[torch.Tensor] = None,
|
|
):
|
|
if num_token_non_padded is None:
|
|
return
|
|
if _can_fuse_padded_region(topk_weights):
|
|
_fill_padded_rows(topk_weights, num_token_non_padded, 0.0)
|
|
return
|
|
indices = torch.arange(0, topk_weights.shape[0], device=topk_weights.device)
|
|
topk_weights[indices >= num_token_non_padded, :] = 0.0
|
|
|
|
|
|
@torch.compile(dynamic=True, backend=get_compiler_backend())
|
|
def _biased_grouped_topk_postprocess(
|
|
topk_ids, expert_location_dispatch_info, num_token_non_padded
|
|
):
|
|
topk_ids = topk_ids_logical_to_physical(topk_ids, expert_location_dispatch_info)
|
|
_mask_topk_ids_padded_region(topk_ids, num_token_non_padded)
|
|
return topk_ids
|
|
|
|
|
|
def biased_grouped_topk_gpu(
|
|
hidden_states: torch.Tensor,
|
|
gating_output: torch.Tensor,
|
|
correction_bias: torch.Tensor,
|
|
topk: int,
|
|
renormalize: bool,
|
|
num_expert_group: Optional[int] = None,
|
|
topk_group: Optional[int] = None,
|
|
num_fused_shared_experts: int = 0,
|
|
routed_scaling_factor: Optional[float] = None,
|
|
apply_routed_scaling_factor_on_output: Optional[bool] = False,
|
|
) -> Tuple[torch.Tensor, torch.Tensor]:
|
|
num_tokens = gating_output.shape[0]
|
|
num_experts = gating_output.shape[1]
|
|
experts_per_group = (
|
|
num_experts // num_expert_group if num_expert_group else num_experts
|
|
)
|
|
|
|
# topk for routed experts only (shared experts are appended separately below)
|
|
topk_routed = topk - num_fused_shared_experts
|
|
if (
|
|
_is_cuda
|
|
and num_expert_group
|
|
and num_expert_group > 1
|
|
and envs.SGLANG_OPT_USE_JIT_KERNEL_GROUPED_TOPK.get()
|
|
):
|
|
# Opt-in: unified Triton router for DeepSeek-V3 grouped routing. Bit-exact
|
|
# with the flashinfer/AOT paths on DeepSeek-V3.2 e2e (validated); handles any
|
|
# experts-per-group (no <=32 cap). Off by default — see the env-var comment.
|
|
from sglang.kernels.ops.moe.moe_fused_gate import (
|
|
moe_fused_gate as jit_grouped_gate,
|
|
)
|
|
|
|
return jit_grouped_gate(
|
|
gating_output.to(dtype=torch.float32),
|
|
correction_bias.to(dtype=torch.float32),
|
|
topk,
|
|
scoring_func="sigmoid",
|
|
num_fused_shared_experts=num_fused_shared_experts,
|
|
renormalize=renormalize,
|
|
routed_scaling_factor=(
|
|
routed_scaling_factor if routed_scaling_factor is not None else 1.0
|
|
),
|
|
apply_routed_scaling_factor_on_output=bool(
|
|
apply_routed_scaling_factor_on_output
|
|
),
|
|
num_expert_group=num_expert_group,
|
|
topk_group=topk_group,
|
|
)
|
|
if (
|
|
_is_cuda
|
|
and fused_topk_deepseek is not None
|
|
and is_power_of_two(num_experts)
|
|
# flashinfer constraints (applied to routed experts only)
|
|
and topk_routed <= 8
|
|
and topk_group <= num_expert_group
|
|
and topk_group * num_expert_group >= topk_routed
|
|
and (
|
|
(experts_per_group <= 32 and experts_per_group * topk_group <= 128)
|
|
if num_expert_group > 1
|
|
else num_experts <= 384
|
|
)
|
|
):
|
|
# Pre-allocate output tensors (flashinfer mutates them in-place)
|
|
topk_weights = torch.empty(
|
|
(num_tokens, topk_routed), dtype=torch.float32, device=gating_output.device
|
|
)
|
|
topk_ids = torch.empty(
|
|
(num_tokens, topk_routed), dtype=torch.int32, device=gating_output.device
|
|
)
|
|
|
|
# flashinfer always applies the scaling_factor internally
|
|
scaling_factor = 1.0
|
|
if routed_scaling_factor is not None and apply_routed_scaling_factor_on_output:
|
|
scaling_factor = routed_scaling_factor
|
|
|
|
# flashinfer's fused_topk_deepseek
|
|
fused_topk_deepseek(
|
|
gating_output.to(dtype=torch.float32),
|
|
correction_bias,
|
|
num_expert_group,
|
|
topk_group,
|
|
topk_routed,
|
|
scaling_factor,
|
|
topk_weights,
|
|
topk_ids,
|
|
True,
|
|
)
|
|
|
|
if num_fused_shared_experts > 0:
|
|
# Append shared expert columns: ID = num_experts (first shared slot),
|
|
# weight = sum(routed) / scaling_factor (matching biased_grouped_topk_impl).
|
|
# For DeepEP/MegaMOE per-rank shared-slot layout, post-process remaps
|
|
# this placeholder ID and overwrites the shared weight for the active scaling path.
|
|
topk_ids = F.pad(topk_ids, (0, num_fused_shared_experts), value=num_experts)
|
|
topk_weights = F.pad(topk_weights, (0, num_fused_shared_experts))
|
|
if routed_scaling_factor is not None:
|
|
topk_weights[:, topk_routed:] = (
|
|
topk_weights[:, :topk_routed].sum(dim=-1, keepdim=True)
|
|
/ routed_scaling_factor
|
|
)
|
|
|
|
return topk_weights, topk_ids
|
|
|
|
elif _is_cuda and num_expert_group > 1:
|
|
# CUDA grouped fallback (flashinfer unavailable / constraints unmet): the
|
|
# unified Triton router replaces the retired AOT moe_fused_gate kernel. It
|
|
# handles any experts-per-group (no MAX_VPT=32 cap) and any num_experts.
|
|
from sglang.kernels.ops.moe.moe_fused_gate import (
|
|
moe_fused_gate as jit_grouped_gate,
|
|
)
|
|
|
|
return jit_grouped_gate(
|
|
gating_output.to(dtype=torch.float32),
|
|
correction_bias.to(dtype=torch.float32),
|
|
topk,
|
|
scoring_func="sigmoid",
|
|
num_fused_shared_experts=num_fused_shared_experts,
|
|
renormalize=renormalize,
|
|
routed_scaling_factor=(
|
|
routed_scaling_factor if routed_scaling_factor is not None else 1.0
|
|
),
|
|
apply_routed_scaling_factor_on_output=bool(
|
|
apply_routed_scaling_factor_on_output
|
|
),
|
|
num_expert_group=num_expert_group,
|
|
topk_group=topk_group,
|
|
)
|
|
|
|
elif _use_aiter:
|
|
assert not apply_routed_scaling_factor_on_output, "Not implemented"
|
|
token = gating_output.shape[0]
|
|
device = gating_output.device
|
|
assert (
|
|
hidden_states.shape[0] == gating_output.shape[0]
|
|
), f"Number of tokens mismatch: hidden_states.shape[0] = {hidden_states.shape[0]}, gating_output.shape[0] = {gating_output.shape[0]}"
|
|
bias = correction_bias.to(dtype=gating_output.dtype)
|
|
scaling = routed_scaling_factor if routed_scaling_factor is not None else 1.0
|
|
|
|
if envs.SGLANG_K3_RADIX4_TOPK.get():
|
|
from sglang.kernels.ops.moe import moe_route_radix4
|
|
|
|
# Gated on the routing shape. Kimi-K3 (896 experts, top-16,
|
|
# ungrouped) is the only config covered for now; anything else
|
|
# falls back to aiter.
|
|
if moe_route_radix4.available() and moe_route_radix4.covered(
|
|
gating_output, bias, topk, num_expert_group, topk_group
|
|
):
|
|
return moe_route_radix4.route_radix4(
|
|
gating_output, bias, topk, renormalize, scaling
|
|
)
|
|
|
|
topk_weights = torch.empty((token, topk), dtype=torch.float32, device=device)
|
|
topk_ids = torch.empty((token, topk), dtype=torch.int32, device=device)
|
|
aiter_biased_grouped_topk(
|
|
gating_output,
|
|
bias,
|
|
topk_weights,
|
|
topk_ids,
|
|
num_expert_group,
|
|
topk_group,
|
|
renormalize,
|
|
scaling,
|
|
)
|
|
return topk_weights, topk_ids
|
|
elif _is_musa and (
|
|
gating_output.shape[1] // num_expert_group <= 32
|
|
or (num_expert_group == 1 and gating_output.shape[1] in {160, 256, 384})
|
|
):
|
|
topk_weights, topk_ids = moe_fused_gate(
|
|
gating_output.to(dtype=torch.float32),
|
|
correction_bias,
|
|
num_expert_group,
|
|
topk_group,
|
|
topk,
|
|
num_fused_shared_experts,
|
|
routed_scaling_factor if routed_scaling_factor is not None else 1.0,
|
|
True,
|
|
apply_routed_scaling_factor_on_output,
|
|
)
|
|
return topk_weights, topk_ids
|
|
else:
|
|
num_experts = gating_output.shape[1]
|
|
# The JIT triton router (single fused kernel: scoring + bias + topk +
|
|
# renorm) handles the ungrouped case with arbitrary num_experts/topk.
|
|
# Original user: Kimi K2 (384 experts). Also dispatch shapes the other
|
|
# fused kernels cannot cover, e.g. Kimi K3 (896 experts, top-16):
|
|
# fused_topk_deepseek needs pow2 experts + topk<=8, jit_grouped_topk
|
|
# needs experts<=512 + topk<=8.
|
|
_jit_gate_ok = (
|
|
_is_cuda
|
|
and num_expert_group == 1
|
|
and (topk_group is None or topk_group == 1)
|
|
and (
|
|
num_experts == 384
|
|
or (num_experts <= 1024 and (num_experts > 512 or topk > 8))
|
|
)
|
|
)
|
|
if _jit_gate_ok:
|
|
# ===== TO BE REFACTORED ====
|
|
_use_jit_bf16_gate = False
|
|
if _SGLANG_EXPERIMENTAL_LORA_OPTI:
|
|
from sglang.srt.lora.trtllm_lora_temp.environ import lora_envs
|
|
|
|
_use_jit_bf16_gate = (
|
|
lora_envs.SGLANG_OPT_USE_JIT_KERNEL_KIMI_GATE.get()
|
|
and lora_envs.SGLANG_OPT_KIMI_GATE_BF16_INPUT.get()
|
|
)
|
|
if _use_jit_bf16_gate:
|
|
from sglang.kernels.ops.moe.trtllm_lora_temp.kimi_k2_moe_fused_gate import (
|
|
kimi_k2_moe_fused_gate as _kimi_k2_moe_fused_gate,
|
|
)
|
|
|
|
# bf16 pass-through: skip the two host-side fp32 upcast kernels.
|
|
return _kimi_k2_moe_fused_gate(
|
|
gating_output,
|
|
correction_bias,
|
|
topk=topk,
|
|
renormalize=renormalize,
|
|
routed_scaling_factor=routed_scaling_factor,
|
|
apply_routed_scaling_factor_on_output=apply_routed_scaling_factor_on_output,
|
|
)
|
|
# ===== END TO BE REFACTORED ====
|
|
from sglang.kernels.ops.moe.moe_fused_gate import moe_fused_gate as jit_gate
|
|
|
|
# Pass BF16 logits through untouched for the automatic radix-select
|
|
# fast path. BF16 -> FP32 is exact, and unsupported shapes fall back
|
|
# inside moe_fused_gate.
|
|
_gating = (
|
|
gating_output
|
|
if gating_output.dtype == torch.bfloat16
|
|
else gating_output.to(dtype=torch.float32)
|
|
)
|
|
# K3 staged fusion: when the model layer staged the routed
|
|
# activations (route_quant_handoff), the radix route, the trtllm id
|
|
# pack and the mxfp8 quant run as one launch. Bit-identical
|
|
# (weights, ids); a miss falls through to the unfused router.
|
|
from sglang.srt.layers.moe import route_quant_handoff
|
|
|
|
fused = route_quant_handoff.try_route_quant_fused(
|
|
_gating,
|
|
correction_bias.to(dtype=torch.float32),
|
|
topk,
|
|
num_fused_shared_experts=num_fused_shared_experts,
|
|
renormalize=renormalize,
|
|
routed_scaling_factor=routed_scaling_factor,
|
|
apply_routed_scaling_factor_on_output=bool(
|
|
apply_routed_scaling_factor_on_output
|
|
),
|
|
)
|
|
if fused is not None:
|
|
return fused
|
|
return jit_gate(
|
|
_gating,
|
|
correction_bias.to(dtype=torch.float32),
|
|
topk=topk,
|
|
scoring_func="sigmoid",
|
|
num_fused_shared_experts=num_fused_shared_experts,
|
|
renormalize=renormalize,
|
|
routed_scaling_factor=(
|
|
routed_scaling_factor if routed_scaling_factor is not None else 1.0
|
|
),
|
|
apply_routed_scaling_factor_on_output=apply_routed_scaling_factor_on_output,
|
|
)
|
|
elif (
|
|
_is_cuda
|
|
and num_expert_group == 1
|
|
and topk_group == 1
|
|
and num_fused_shared_experts == 0
|
|
and num_experts <= 512
|
|
and topk <= 32
|
|
):
|
|
# Ungrouped sigmoid (num_expert_group == 1): use the unified Triton
|
|
# router, which subsumes the jit grouped_topk.cuh kernel here.
|
|
from sglang.kernels.ops.moe.moe_fused_gate import moe_fused_gate as jit_gate
|
|
|
|
return jit_gate(
|
|
gating_output,
|
|
correction_bias.to(torch.float32),
|
|
topk,
|
|
scoring_func="sigmoid",
|
|
renormalize=renormalize,
|
|
routed_scaling_factor=(
|
|
routed_scaling_factor if routed_scaling_factor is not None else 1.0
|
|
),
|
|
apply_routed_scaling_factor_on_output=apply_routed_scaling_factor_on_output,
|
|
)
|
|
elif (
|
|
_is_xpu
|
|
and num_expert_group == 1
|
|
and topk_group == 1
|
|
and num_fused_shared_experts == 0
|
|
and num_experts <= 256
|
|
and topk <= 8
|
|
):
|
|
|
|
scale = (
|
|
routed_scaling_factor
|
|
if (
|
|
apply_routed_scaling_factor_on_output
|
|
and routed_scaling_factor is not None
|
|
)
|
|
else 1.0
|
|
)
|
|
|
|
num_tokens = gating_output.shape[0]
|
|
|
|
topk_values = torch.empty(
|
|
(num_tokens, topk), dtype=torch.float32, device=gating_output.device
|
|
)
|
|
topk_indices = torch.empty(
|
|
(num_tokens, topk), dtype=torch.int32, device=gating_output.device
|
|
)
|
|
|
|
if num_tokens == 0:
|
|
return topk_values, topk_indices
|
|
|
|
topk_sigmoid(
|
|
topk_values,
|
|
topk_indices,
|
|
gating_output,
|
|
renormalize,
|
|
# The XPU topk_sigmoid AOT kernel requires an fp32 correction bias.
|
|
correction_bias.to(torch.float32),
|
|
scale,
|
|
)
|
|
|
|
return topk_values, topk_indices
|
|
elif (
|
|
_is_xpu
|
|
# moe_fused_gate kernel ensures that num_experts/num_expert_group does not exceed MAX_VPT=32 now.
|
|
and experts_per_group <= 32
|
|
and is_power_of_two(num_experts)
|
|
):
|
|
from sgl_kernel import moe_fused_gate
|
|
|
|
return moe_fused_gate(
|
|
gating_output.to(torch.float32),
|
|
correction_bias.to(torch.float32),
|
|
num_expert_group,
|
|
topk_group,
|
|
topk,
|
|
renormalize=renormalize,
|
|
scoring_func="sigmoid",
|
|
num_fused_shared_experts=num_fused_shared_experts,
|
|
routed_scaling_factor=(
|
|
routed_scaling_factor if routed_scaling_factor is not None else 1.0
|
|
),
|
|
apply_routed_scaling_factor_on_output=bool(
|
|
apply_routed_scaling_factor_on_output
|
|
),
|
|
)
|
|
|
|
else:
|
|
return biased_grouped_topk_impl(
|
|
hidden_states,
|
|
gating_output,
|
|
correction_bias,
|
|
topk,
|
|
renormalize,
|
|
num_expert_group,
|
|
topk_group,
|
|
num_fused_shared_experts=num_fused_shared_experts,
|
|
routed_scaling_factor=routed_scaling_factor,
|
|
apply_routed_scaling_factor_on_output=apply_routed_scaling_factor_on_output,
|
|
)
|
|
|
|
|
|
def biased_grouped_topk_cpu(
|
|
hidden_states: torch.Tensor,
|
|
gating_output: torch.Tensor,
|
|
correction_bias: torch.Tensor,
|
|
topk: int,
|
|
renormalize: bool,
|
|
num_expert_group: Optional[int] = None,
|
|
topk_group: Optional[int] = None,
|
|
compiled: bool = True,
|
|
num_fused_shared_experts: int = 0,
|
|
routed_scaling_factor: Optional[float] = None,
|
|
apply_routed_scaling_factor_on_output: Optional[bool] = False,
|
|
):
|
|
return torch.ops.sgl_kernel.biased_grouped_topk_cpu(
|
|
hidden_states,
|
|
gating_output,
|
|
correction_bias,
|
|
topk,
|
|
renormalize,
|
|
num_expert_group,
|
|
topk_group,
|
|
num_fused_shared_experts,
|
|
routed_scaling_factor if apply_routed_scaling_factor_on_output else None,
|
|
# num_token_non_padded must be None since it is not supported in kernel
|
|
num_token_non_padded=None,
|
|
)
|
|
|
|
|
|
if _is_cpu and _is_cpu_amx_available:
|
|
biased_grouped_topk = biased_grouped_topk_cpu
|
|
grouped_topk = grouped_topk_cpu
|
|
fused_topk_native = fused_topk_cpu
|
|
fused_topk = fused_topk_cpu
|
|
else:
|
|
biased_grouped_topk = biased_grouped_topk_gpu
|
|
grouped_topk = grouped_topk_xpu if _is_xpu else grouped_topk_gpu
|
|
fused_topk_native = fused_topk_torch_native
|
|
|
|
|
|
def remap_topk_for_per_rank_shared_slots(
|
|
topk_ids: torch.Tensor,
|
|
topk_weights: torch.Tensor,
|
|
num_fused_shared_experts: int,
|
|
num_physical_routed_experts: int,
|
|
topk_config: TopKConfig,
|
|
) -> tuple[torch.Tensor, torch.Tensor]:
|
|
"""Remap TopK IDs to a per-rank shared-slot layout.
|
|
|
|
DeepEP and MegaMoE dispatch need each rank's shared expert at a unique ID
|
|
so tokens route to the correct rank. The layout is ordered by rank:
|
|
[rank0 routed..., rank0 shared, rank1 routed..., rank1 shared, ...].
|
|
|
|
Routed IDs: e -> e + (e // num_local_routed) * num_fused_shared_experts
|
|
Shared IDs: ep_rank * num_local_experts + num_local_routed
|
|
Shared weight: 1.0 on the aiter path, else 1/routed_scaling_factor (see below).
|
|
"""
|
|
if topk_ids.shape[0] == 0:
|
|
return topk_ids, topk_weights
|
|
|
|
ep_size = get_parallel().moe_ep_size
|
|
ep_rank = get_parallel().moe_ep_rank
|
|
# Static EPLB may add redundant physical experts. At this point routed
|
|
# topk_ids have already been remapped from logical to physical ids, so the
|
|
# per-rank shared-slot layout must use the physical routed count.
|
|
num_local_routed = num_physical_routed_experts // ep_size
|
|
num_local_experts = num_local_routed + num_fused_shared_experts
|
|
|
|
# Remap routed IDs: insert gaps for shared expert slots (single fused op).
|
|
# Each rank r == e // num_local_routed is preceded by r shared-slot blocks of
|
|
# width num_fused_shared_experts, so shift by (e // num_local_routed) * S --
|
|
# a single-slot shift (S == 1) would let routed ids collide with an earlier
|
|
# rank's shared slots once S > 1.
|
|
routed = topk_ids[:, :-num_fused_shared_experts]
|
|
topk_ids[:, :-num_fused_shared_experts] = (
|
|
routed + (routed // num_local_routed) * num_fused_shared_experts
|
|
)
|
|
|
|
# Set shared expert IDs to route to home rank (vectorized)
|
|
topk_ids[:, -num_fused_shared_experts:] = (
|
|
ep_rank * num_local_experts
|
|
+ num_local_routed
|
|
+ torch.arange(num_fused_shared_experts, device=topk_ids.device)
|
|
)
|
|
|
|
# Override the fused shared expert's weight so its net contribution is 1.0x.
|
|
#
|
|
# The correct value depends on whether routed_scaling_factor is applied to
|
|
# the MoE output AFTER the experts run, or already folded into the routed
|
|
# topk weights BEFORE dispatch:
|
|
#
|
|
# * Post-MoE scaling path (default): DeepseekV2MoE.forward_deepep later
|
|
# multiplies the whole MoE output by routed_scaling_factor, so the shared
|
|
# weight must be 1/routed_scaling_factor for (1/rsf) * rsf = 1.0.
|
|
# * aiter (HIP) path: aiter_biased_grouped_topk folds routed_scaling_factor
|
|
# into each routed topk weight, and forward_deepep SKIPS the post-MoE
|
|
# multiply for _use_aiter (see its `not (... or _use_aiter)` guard). The
|
|
# shared weight must therefore be 1.0 -- applying 1/rsf here would
|
|
# under-weight the always-on shared expert by routed_scaling_factor and
|
|
# corrupt every MoE layer.
|
|
#
|
|
# NOTE: forward_deepep also skips the post-MoE multiply for the non-aiter
|
|
# families where routed_scaling_factor is pre-folded in topk
|
|
# (should_fuse_routed_scaling_factor_in_topk / apply_routed_scaling_factor_on_output:
|
|
# ModelOpt NVFP4, cutlass/trtllm-routed fp8), so those would likewise need a
|
|
# 1.0 shared weight. This fix is deliberately scoped to the aiter path (the
|
|
# one validated on AMD MI355X); those other backends are left at their
|
|
# existing behavior and can be addressed by their maintainers.
|
|
routed_scaling_factor = topk_config.routed_scaling_factor
|
|
if _use_aiter:
|
|
topk_weights[:, -num_fused_shared_experts:] = 1.0
|
|
elif routed_scaling_factor is not None and routed_scaling_factor != 0:
|
|
topk_weights[:, -num_fused_shared_experts:] = 1.0 / routed_scaling_factor
|
|
|
|
return topk_ids, topk_weights
|
|
|
|
|
|
def capture_routed_experts_if_allowed(
|
|
topk_config: TopKConfig,
|
|
layer_id: Optional[int],
|
|
topk_ids: torch.Tensor,
|
|
) -> None:
|
|
"""Single capture site for every backend, gated by the per-config opt-out.
|
|
|
|
Routing all backends through here keeps the draft-side opt-out from being
|
|
bypassed by an inlined capturer call.
|
|
"""
|
|
if not topk_config.allow_routed_experts_capture:
|
|
return
|
|
if (cap := get_global_experts_capturer()) is not None:
|
|
cap.capture(
|
|
layer_id=layer_id,
|
|
topk_indices=topk_ids,
|
|
)
|
|
|
|
|
|
def _post_process_topk_ids(
|
|
topk_ids: torch.Tensor,
|
|
topk_weights: torch.Tensor,
|
|
topk_config: TopKConfig,
|
|
router_logits: torch.Tensor,
|
|
layer_id: int,
|
|
num_token_non_padded: Optional[torch.Tensor] = None,
|
|
expert_location_dispatch_info: Optional[ExpertLocationDispatchInfo] = None,
|
|
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
|
num_fused_shared_experts = topk_config.num_fused_shared_experts
|
|
use_per_rank_shared_slots = has_per_rank_fused_shared_slots(
|
|
num_fused_shared_experts
|
|
)
|
|
fused_shared_experts_scaling_factor = (
|
|
topk_config.fused_shared_experts_scaling_factor
|
|
)
|
|
capture_routed_experts_if_allowed(topk_config, layer_id, topk_ids)
|
|
recorder_topk_ids = None
|
|
_fold_pad_into_append = False
|
|
if _is_cuda:
|
|
# LP path: solve LP outside torch.compile (the solver contains an
|
|
# EP all-reduce that can't run inside compiled regions).
|
|
log2phy_prob = None
|
|
if (
|
|
expert_location_dispatch_info is not None
|
|
and getattr(expert_location_dispatch_info, "ep_dispatch_algorithm", None)
|
|
== "lp"
|
|
):
|
|
from sglang.srt.eplb.lplb_solver import get_global_lplb_solver
|
|
|
|
lplb_solver = get_global_lplb_solver(layer_id)
|
|
if lplb_solver is not None:
|
|
log2phy_prob = lplb_solver.solve(topk_ids)
|
|
|
|
if log2phy_prob is not None:
|
|
topk_ids = topk_ids_logical_to_physical(
|
|
topk_ids, expert_location_dispatch_info, log2phy_prob
|
|
)
|
|
_mask_topk_ids_padded_region(topk_ids, num_token_non_padded)
|
|
elif use_per_rank_shared_slots:
|
|
# Shared experts appended as extra columns in topk_ids: their value
|
|
# would be out-of-bounds for the logical-to-physical dispatch table,
|
|
# so split, dispatch the routed cols, recombine.
|
|
shared_cols = topk_ids[:, -num_fused_shared_experts:]
|
|
routed_cols = topk_ids[:, :-num_fused_shared_experts]
|
|
routed_cols = _biased_grouped_topk_postprocess(
|
|
routed_cols, expert_location_dispatch_info, num_token_non_padded
|
|
)
|
|
topk_ids = torch.cat([routed_cols, shared_cols], dim=-1)
|
|
# ExpertDistributionRecorder tracks EPLB physical routed experts.
|
|
# Per-rank shared-slot remap later adds shared slots to the topk ID
|
|
# space, so keep the routed physical ids separately for statistics.
|
|
recorder_topk_ids = routed_cols
|
|
else:
|
|
topk_ids = _biased_grouped_topk_postprocess(
|
|
topk_ids, expert_location_dispatch_info, num_token_non_padded
|
|
)
|
|
elif _is_hip:
|
|
# On AMD HIP the aiter MoE kernels do not handle topk_ids=-1 safely
|
|
# (negative indices cause illegal memory access). Always fill the padded
|
|
# region with 0 so every kernel sees a valid in-range expert id.
|
|
# Routing weights for padded tokens are zeroed below so their
|
|
# contribution to the hidden state is still zero regardless of the id.
|
|
# Regression: skipping this mask when EPLB is disabled caused garbage
|
|
# MoE routing for models like DeepSeek-R1-MXFP4 (accuracy ~0.09 vs 0.94+).
|
|
#
|
|
# Fold: when the fused append+remap kernel runs below (aiter per-rank
|
|
# shared-slot path, EPLB off) it folds this padded fill itself
|
|
# (pad_fill_id=0 -> remap(0)=0, bit-identical), so skip the separate
|
|
# _fill_padded_rows launch here.
|
|
_fold_pad_into_append = (
|
|
num_fused_shared_experts > 0
|
|
and _use_aiter
|
|
and use_per_rank_shared_slots
|
|
and not _eplb_remap_enabled()
|
|
)
|
|
if not _fold_pad_into_append:
|
|
_mask_topk_ids_padded_region(topk_ids, num_token_non_padded, fill_value=0)
|
|
# The logical->physical remap is only meaningful when a real
|
|
# expert-location mapping exists. With a trivial placement and EPLB off
|
|
# the map is identity so the remap can be skipped safely.
|
|
if _eplb_remap_enabled():
|
|
topk_ids = topk_ids_logical_to_physical(
|
|
topk_ids, expert_location_dispatch_info
|
|
)
|
|
# NOTE (HIP): padded-token routing-weight zeroing is deferred to the
|
|
# single pass at the end of this function (gated by SGLANG_MORI_NO_PAD_MASK).
|
|
# That final pass re-zeros after any shared-expert append/remap, so a
|
|
# second zeroing here would be redundant (zeroing is idempotent).
|
|
|
|
if recorder_topk_ids is None:
|
|
recorder_topk_ids = topk_ids
|
|
|
|
_aiter_append = num_fused_shared_experts > 0 and _use_aiter
|
|
|
|
if _aiter_append and use_per_rank_shared_slots:
|
|
# Fused path: append shared experts AND apply the per-rank shared-slot
|
|
# remap in a single Triton kernel. This replaces the original
|
|
# fused_append_shared_experts() + eager per-rank shared-slot remap pair,
|
|
# collapsing ~6 launch-bound elementwise kernels/layer (div_floor / add /
|
|
# arange / fill / copy) into the one append kernel that already runs.
|
|
#
|
|
# Shared weight is 1.0 here because this branch is aiter-only:
|
|
# aiter_biased_grouped_topk folds routed_scaling_factor into the routed
|
|
# weights and forward_deepep skips the post-MoE multiply for _use_aiter,
|
|
# so the always-on shared expert must contribute 1.0x. (The eager
|
|
# per-rank shared-slot remap instead sets shared weight to
|
|
# 1/routed_scaling_factor to compensate a post-MoE scale that the aiter
|
|
# path does not apply; see PR #28237.)
|
|
num_physical_routed_experts = (
|
|
expert_location_dispatch_info.num_physical_experts
|
|
if expert_location_dispatch_info is not None
|
|
else router_logits.shape[1]
|
|
)
|
|
ep_size = get_parallel().moe_ep_size
|
|
ep_rank = get_parallel().moe_ep_rank
|
|
num_local_routed = num_physical_routed_experts // ep_size
|
|
num_local_experts = num_local_routed + num_fused_shared_experts
|
|
shared_id_base = ep_rank * num_local_experts + num_local_routed
|
|
|
|
# Lazy import to avoid circular-import issues
|
|
from sglang.kernels.ops.moe.fused_moe_triton_kernels import (
|
|
fused_append_remap_shared_experts_deepep,
|
|
)
|
|
|
|
topk_ids, topk_weights = fused_append_remap_shared_experts_deepep(
|
|
topk_ids,
|
|
topk_weights,
|
|
num_fused_shared_experts,
|
|
1.0, # shared-expert weight on the aiter path
|
|
shared_id_base,
|
|
num_local_routed,
|
|
num_token_non_padded=(
|
|
num_token_non_padded if _fold_pad_into_append else None
|
|
),
|
|
)
|
|
elif _aiter_append:
|
|
M, N = router_logits.shape
|
|
scale_factor = (
|
|
1.0
|
|
if fused_shared_experts_scaling_factor is None
|
|
else fused_shared_experts_scaling_factor
|
|
)
|
|
|
|
# Lazy import to avoid circular-import issues
|
|
from sglang.kernels.ops.moe.fused_moe_triton_kernels import (
|
|
fused_append_shared_experts,
|
|
)
|
|
|
|
topk_ids, topk_weights = fused_append_shared_experts(
|
|
topk_ids,
|
|
topk_weights,
|
|
num_fused_shared_experts,
|
|
scale_factor,
|
|
N, # base id for shared experts
|
|
)
|
|
|
|
elif use_per_rank_shared_slots:
|
|
# DeepEP/MegaMOE: remap to per-rank shared-slot layout where each
|
|
# rank's shared expert has a unique ID for dispatch routing.
|
|
num_physical_routed_experts = (
|
|
expert_location_dispatch_info.num_physical_experts
|
|
if expert_location_dispatch_info is not None
|
|
else router_logits.shape[1]
|
|
)
|
|
topk_ids, topk_weights = remap_topk_for_per_rank_shared_slots(
|
|
topk_ids,
|
|
topk_weights,
|
|
num_fused_shared_experts,
|
|
num_physical_routed_experts,
|
|
topk_config,
|
|
)
|
|
elif (
|
|
num_fused_shared_experts > 0 and fused_shared_experts_scaling_factor is not None
|
|
):
|
|
# Standard EP all-reduces the per-rank shared-expert outputs; without the
|
|
# supplied 1/ep_size factor the shared contribution is summed ep_size times.
|
|
topk_weights[
|
|
:, -num_fused_shared_experts:
|
|
] *= fused_shared_experts_scaling_factor
|
|
|
|
if _is_hip and not _skip_hip_pad_mask:
|
|
# Shared-expert append/remap can introduce non-zero weights after the
|
|
# initial HIP padding mask above. Ensure padded tokens leave this helper
|
|
# with all expert weights zeroed.
|
|
_zero_topk_weights_padded_region(topk_weights, num_token_non_padded)
|
|
|
|
return topk_ids, topk_weights, recorder_topk_ids
|
|
|
|
|
|
def select_experts(
|
|
hidden_states: torch.Tensor,
|
|
router_logits: torch.Tensor,
|
|
topk_config: TopKConfig,
|
|
*,
|
|
layer_id: Optional[int] = None,
|
|
num_token_non_padded: Optional[torch.Tensor] = None,
|
|
expert_location_dispatch_info: Optional[ExpertLocationDispatchInfo] = None,
|
|
) -> StandardTopKOutput:
|
|
top_k = topk_config.top_k
|
|
use_grouped_topk = topk_config.use_grouped_topk
|
|
topk_group = topk_config.topk_group
|
|
num_expert_group = topk_config.num_expert_group
|
|
renormalize = topk_config.renormalize
|
|
num_fused_shared_experts = topk_config.num_fused_shared_experts
|
|
custom_routing_function = topk_config.custom_routing_function
|
|
correction_bias = topk_config.correction_bias
|
|
torch_native = topk_config.torch_native
|
|
routed_scaling_factor = topk_config.routed_scaling_factor
|
|
apply_routed_scaling_factor_on_output = (
|
|
topk_config.apply_routed_scaling_factor_on_output
|
|
)
|
|
|
|
scoring_func = topk_config.scoring_func
|
|
|
|
# Set by the fused-gating+pack branch below; None everywhere else.
|
|
packed_topk = None
|
|
|
|
(
|
|
router_logits,
|
|
correction_bias,
|
|
) = expert_location_dispatch.transform_select_experts_inputs(
|
|
router_logits=router_logits,
|
|
correction_bias=correction_bias,
|
|
info=expert_location_dispatch_info,
|
|
)
|
|
|
|
# DeepSeek V2/V3/R1 series models use grouped_top_k
|
|
# remove num_fused_shared_experts from grouped_topk/biased_grouped_topk
|
|
num_routed_topk = top_k - num_fused_shared_experts
|
|
if use_grouped_topk:
|
|
assert topk_group is not None
|
|
assert num_expert_group is not None
|
|
if correction_bias is None:
|
|
topk_weights, topk_ids = grouped_topk(
|
|
hidden_states=hidden_states,
|
|
gating_output=router_logits,
|
|
topk=num_routed_topk if _use_aiter else top_k,
|
|
renormalize=renormalize,
|
|
num_expert_group=num_expert_group,
|
|
topk_group=topk_group,
|
|
num_fused_shared_experts=num_fused_shared_experts,
|
|
routed_scaling_factor=routed_scaling_factor,
|
|
apply_routed_scaling_factor_on_output=apply_routed_scaling_factor_on_output,
|
|
scoring_func=scoring_func,
|
|
)
|
|
else:
|
|
topk_weights, topk_ids = biased_grouped_topk(
|
|
hidden_states=hidden_states,
|
|
gating_output=router_logits,
|
|
correction_bias=correction_bias,
|
|
topk=num_routed_topk if _use_aiter else top_k,
|
|
renormalize=renormalize,
|
|
num_expert_group=num_expert_group,
|
|
topk_group=topk_group,
|
|
num_fused_shared_experts=num_fused_shared_experts,
|
|
routed_scaling_factor=routed_scaling_factor,
|
|
apply_routed_scaling_factor_on_output=apply_routed_scaling_factor_on_output,
|
|
)
|
|
elif torch_native and custom_routing_function is None:
|
|
assert (
|
|
num_token_non_padded is None
|
|
), "num_token_non_padded is not yet supported in fused_topk_native"
|
|
assert expert_location_dispatch_info is None
|
|
assert not apply_routed_scaling_factor_on_output, "Not implemented"
|
|
topk_weights, topk_ids = fused_topk_native(
|
|
hidden_states=hidden_states,
|
|
gating_output=router_logits,
|
|
topk=num_routed_topk if _use_aiter else top_k,
|
|
renormalize=renormalize,
|
|
correction_bias=correction_bias,
|
|
scoring_func=scoring_func,
|
|
)
|
|
elif custom_routing_function is None:
|
|
if scoring_func not in ("sqrtsoftplus", "sigmoid"):
|
|
assert not apply_routed_scaling_factor_on_output, "Not implemented"
|
|
|
|
# The JIT route depends on GPU-only topk_sigmoid/topk_softmax imports
|
|
_can_use_jit_kernel = not _is_cpu
|
|
|
|
if _can_use_jit_kernel and (
|
|
scoring_func == "sqrtsoftplus" or scoring_func == "sigmoid"
|
|
):
|
|
_biased_topk = biased_topk_xpu if _is_xpu else biased_topk_jit_kernel_impl
|
|
topk_weights, topk_ids = _biased_topk(
|
|
hidden_states=hidden_states,
|
|
gating_output=router_logits,
|
|
correction_bias=correction_bias,
|
|
topk=num_routed_topk if _use_aiter else top_k,
|
|
renormalize=renormalize,
|
|
scoring_func=scoring_func,
|
|
num_fused_shared_experts=num_fused_shared_experts,
|
|
routed_scaling_factor=routed_scaling_factor,
|
|
num_token_non_padded=num_token_non_padded,
|
|
expert_location_dispatch_info=expert_location_dispatch_info,
|
|
apply_routed_scaling_factor_on_output=apply_routed_scaling_factor_on_output,
|
|
)
|
|
elif (
|
|
get_moe_runner_backend().is_flashinfer_trtllm_routed()
|
|
and scoring_func == "softmax"
|
|
and correction_bias is None
|
|
):
|
|
# flashinfer_trtllm_routed uses raw-logits topk
|
|
topk_weights, topk_ids = fused_topk_softmax_torch_raw_logits(
|
|
hidden_states=hidden_states,
|
|
gating_output=router_logits,
|
|
topk=num_routed_topk if _use_aiter else top_k,
|
|
renormalize=renormalize,
|
|
)
|
|
else:
|
|
# Fused gating + routed pack (SGLANG_OPT_LORA_FUSED_TOPK_PACK): only on the plain
|
|
# CUDA softmax path with no EPLB remap / shared experts / bias / routing overrides.
|
|
_fused_topk_pack = False
|
|
if _SGLANG_EXPERIMENTAL_LORA_OPTI:
|
|
from sglang.srt.lora.trtllm_lora_temp.environ import lora_envs
|
|
|
|
_fused_topk_pack = lora_envs.SGLANG_OPT_LORA_FUSED_TOPK_PACK.get()
|
|
if (
|
|
_fused_topk_pack
|
|
and _is_cuda
|
|
and not _use_aiter
|
|
and scoring_func == "softmax"
|
|
and correction_bias is None
|
|
and expert_location_dispatch_info is None
|
|
and num_fused_shared_experts == 0
|
|
and not envs.SGLANG_SIMULATE_UNIFORM_EXPERTS.get()
|
|
and not envs.SGLANG_SIMULATE_ROUND_ROBIN_EXPERTS.get()
|
|
):
|
|
num_experts = router_logits.shape[-1]
|
|
if num_experts & (num_experts - 1) == 0 and num_experts <= 512:
|
|
packed_topk = torch.empty(
|
|
(hidden_states.shape[0], top_k),
|
|
dtype=torch.int32,
|
|
device=hidden_states.device,
|
|
)
|
|
|
|
# Qwen3MOE uses fused_topk
|
|
_fused_topk_kwargs = {}
|
|
# ===== TO BE REFACTORED ====
|
|
# Only the experimental fused topk+pack passes packed_out/num_token_non_padded;
|
|
# the default call keeps the upstream signature (fused_topk_cpu lacks these).
|
|
if packed_topk is not None:
|
|
_fused_topk_kwargs = dict(
|
|
packed_out=packed_topk,
|
|
num_token_non_padded=num_token_non_padded,
|
|
)
|
|
# ===== END TO BE REFACTORED ====
|
|
topk_weights, topk_ids = fused_topk(
|
|
hidden_states=hidden_states,
|
|
gating_output=router_logits,
|
|
topk=num_routed_topk if _use_aiter else top_k,
|
|
renormalize=renormalize,
|
|
correction_bias=correction_bias,
|
|
scoring_func=scoring_func,
|
|
num_fused_shared_experts=num_fused_shared_experts,
|
|
routed_scaling_factor=routed_scaling_factor,
|
|
apply_routed_scaling_factor_on_output=apply_routed_scaling_factor_on_output,
|
|
**_fused_topk_kwargs,
|
|
)
|
|
else:
|
|
# custom_routing_function itself is padding-unaware; its output on
|
|
# padded rows is garbage. That is fine because _post_process_topk_ids
|
|
# below masks rows >= num_token_non_padded (-1 on CUDA, 0 + zeroed
|
|
# weights on HIP) after the logical->physical remap.
|
|
assert not apply_routed_scaling_factor_on_output, "Not implemented"
|
|
topk_weights, topk_ids = custom_routing_function(
|
|
hidden_states=hidden_states,
|
|
gating_output=router_logits,
|
|
topk=num_routed_topk if _use_aiter else top_k,
|
|
renormalize=renormalize,
|
|
)
|
|
|
|
simulate_uniform_experts = envs.SGLANG_SIMULATE_UNIFORM_EXPERTS.get()
|
|
simulate_round_robin_experts = envs.SGLANG_SIMULATE_ROUND_ROBIN_EXPERTS.get()
|
|
if simulate_uniform_experts and simulate_round_robin_experts:
|
|
raise ValueError(
|
|
"SGLANG_SIMULATE_UNIFORM_EXPERTS and "
|
|
"SGLANG_SIMULATE_ROUND_ROBIN_EXPERTS are mutually exclusive"
|
|
)
|
|
|
|
if simulate_uniform_experts:
|
|
# Benchmark-only: override gating with random-offset uniform expert assignment
|
|
# to avoid expert imbalance from dummy/random weights. Do NOT use in production.
|
|
num_tokens, k = topk_ids.shape
|
|
num_experts = router_logits.shape[1]
|
|
if k > 0:
|
|
offsets = torch.randint(
|
|
0, num_experts, (num_tokens, 1), device=topk_ids.device
|
|
)
|
|
steps = torch.arange(k, device=topk_ids.device).unsqueeze(0)
|
|
step = max(num_experts // k, 1)
|
|
topk_ids = ((offsets + steps * step) % num_experts).to(topk_ids.dtype)
|
|
topk_weights = torch.ones_like(topk_weights) / k
|
|
elif simulate_round_robin_experts:
|
|
# Benchmark-only: override gating with deterministic expert assignment
|
|
# to avoid routing noise from dummy/random weights. Do NOT use in production.
|
|
num_tokens, k = topk_ids.shape
|
|
num_experts = router_logits.shape[1]
|
|
topk_ids = _make_round_robin_expert_ids(
|
|
num_tokens,
|
|
k,
|
|
num_experts,
|
|
device=topk_ids.device,
|
|
dtype=topk_ids.dtype,
|
|
layer_id=layer_id,
|
|
)
|
|
if k > 0:
|
|
topk_weights = torch.full_like(topk_weights, 1.0 / k)
|
|
|
|
topk_ids, topk_weights, recorder_topk_ids = _post_process_topk_ids(
|
|
topk_ids=topk_ids,
|
|
topk_weights=topk_weights,
|
|
topk_config=topk_config,
|
|
router_logits=router_logits,
|
|
num_token_non_padded=num_token_non_padded,
|
|
layer_id=layer_id,
|
|
expert_location_dispatch_info=expert_location_dispatch_info,
|
|
)
|
|
|
|
get_global_expert_distribution_recorder().on_select_experts(
|
|
topk_ids=recorder_topk_ids
|
|
)
|
|
|
|
# ===== TO BE REFACTORED ====
|
|
if packed_topk is not None:
|
|
return StandardTopKOutputPacked(
|
|
topk_weights, topk_ids, router_logits, packed_topk
|
|
)
|
|
# ===== END TO BE REFACTORED ====
|
|
return StandardTopKOutput(topk_weights, topk_ids, router_logits)
|
|
|
|
|
|
def precomputed_topk_postprocess_is_noop(
|
|
topk_config: TopKConfig,
|
|
num_token_non_padded: Optional[torch.Tensor] = None,
|
|
expert_location_dispatch_info: Optional[ExpertLocationDispatchInfo] = None,
|
|
) -> bool:
|
|
"""Whether :func:`build_precomputed_topk_output` can stand in for
|
|
:func:`select_experts`' post-processing.
|
|
|
|
A router that produces (weights, ids) itself -- e.g. K3's fused gate+top-k
|
|
kernel -- skips select_experts entirely, so it must not skip the work
|
|
select_experts does *after* the top-k: the EPLB logical->physical remap, the
|
|
padded-region mask, and the shared-expert append. This returns True only
|
|
when all of those reduce to no-ops, leaving just the capture hook and the
|
|
distribution recorder (which the builder below still runs).
|
|
"""
|
|
return (
|
|
_is_cuda
|
|
and topk_config.num_fused_shared_experts == 0
|
|
and num_token_non_padded is None
|
|
and expert_location_dispatch_info is None
|
|
and not envs.SGLANG_SIMULATE_UNIFORM_EXPERTS.get()
|
|
and not envs.SGLANG_SIMULATE_ROUND_ROBIN_EXPERTS.get()
|
|
)
|
|
|
|
|
|
def build_precomputed_topk_output(
|
|
topk_weights: torch.Tensor,
|
|
topk_ids: torch.Tensor,
|
|
topk_config: TopKConfig,
|
|
layer_id: int,
|
|
) -> StandardTopKOutput:
|
|
"""Wrap a router's own (weights, ids) as a STANDARD top-k output, running the
|
|
capture hook and the expert-distribution recorder that select_experts would.
|
|
|
|
Only valid when :func:`precomputed_topk_postprocess_is_noop` holds.
|
|
"""
|
|
capture_routed_experts_if_allowed(topk_config, layer_id, topk_ids)
|
|
get_global_expert_distribution_recorder().on_select_experts(topk_ids=topk_ids)
|
|
# router_logits is only read by the BYPASSED formats and by the
|
|
# shared-expert append (excluded above); STANDARD consumers take ids/weights.
|
|
return StandardTopKOutput(topk_weights, topk_ids, None)
|
|
|
|
|
|
# NOTE: the AOT sgl_kernel::moe_fused_gate and sgl_kernel::kimi_k2_moe_fused_gate
|
|
# ops (and their torch.compile fake impls) were retired here — both CUDA gate
|
|
# paths now route through the unified Triton router (kernels/ops/moe/moe_fused_gate.py),
|
|
# whose Python impl is traceable directly, so no register_fake shim is needed.
|