from __future__ import annotations from typing import TYPE_CHECKING, List, Optional import torch import torch.nn.functional as F from torch.nn.parameter import Parameter from sglang.srt.layers.amx_utils import ( CPUQuantMethod, _amx_process_weight_after_loading, ) from sglang.srt.layers.moe import ( MoeRunner, MoeRunnerBackend, MoeRunnerConfig, get_moe_runner_backend, ) from sglang.srt.layers.moe.moe_runner.triton import TritonMoeQuantInfo from sglang.srt.layers.quantization.base_config import ( FusedMoEMethodBase, LinearMethodBase, QuantizeMethodBase, ) from sglang.srt.layers.utils import MultiPlatformOp from sglang.srt.utils import ( cpu_has_amx_support, get_bool_env_var, is_cpu, is_hip, is_npu, next_power_of_2, set_weight_attrs, use_intel_amx_backend, use_intel_xpu_backend, ) if TYPE_CHECKING: from sglang.srt.layers.moe.token_dispatcher import ( CombineInput, StandardDispatchOutput, ) _is_cpu_amx_available = cpu_has_amx_support() _is_hip = is_hip() _is_cpu = is_cpu() _is_npu = is_npu() _use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip if _use_aiter: from aiter import ActivationType from aiter.fused_moe import fused_moe from aiter.ops.shuffle import shuffle_weight if _is_npu: from sglang.srt.hardware_backend.npu.utils import npu_format_cast try: from flashinfer.fused_moe import cutlass_fused_moe as flashinfer_cutlass_fused_moe except ImportError: flashinfer_cutlass_fused_moe = None class UnquantizedEmbeddingMethod(QuantizeMethodBase): """Unquantized method for embeddings.""" def create_weights( self, layer: torch.nn.Module, input_size_per_partition: int, output_partition_sizes: List[int], input_size: int, output_size: int, params_dtype: torch.dtype, **extra_weight_attrs, ): """Create weights for embedding layer.""" weight = Parameter( torch.empty( sum(output_partition_sizes), input_size_per_partition, dtype=params_dtype, ), requires_grad=False, ) set_weight_attrs(weight, {"input_dim": 1, "output_dim": 0}) layer.register_parameter("weight", weight) set_weight_attrs(weight, extra_weight_attrs) def apply( self, layer: torch.nn.Module, x: torch.Tensor, bias: Optional[torch.Tensor] = None, ) -> torch.Tensor: return F.linear(x, layer.weight, bias) def embedding(self, layer: torch.nn.Module, input_: torch.Tensor) -> torch.Tensor: return F.embedding(input_, layer.weight) class UnquantizedLinearMethod(LinearMethodBase): """Linear method without quantization.""" def create_weights( self, layer: torch.nn.Module, input_size_per_partition: int, output_partition_sizes: List[int], input_size: int, output_size: int, params_dtype: torch.dtype, **extra_weight_attrs, ): weight = Parameter( torch.empty( sum(output_partition_sizes), input_size_per_partition, dtype=params_dtype, ), requires_grad=False, ) set_weight_attrs(weight, {"input_dim": 1, "output_dim": 0}) layer.register_parameter("weight", weight) set_weight_attrs(weight, extra_weight_attrs) def process_weights_after_loading(self, layer: torch.nn.Module) -> None: if _is_cpu and _is_cpu_amx_available: _amx_process_weight_after_loading(layer, ["weight"]) def apply( self, layer: torch.nn.Module, x: torch.Tensor, bias: Optional[torch.Tensor] = None, ) -> torch.Tensor: if use_intel_amx_backend(layer): x_shapes = x.shape if len(x_shapes) == 3: x = x.view(-1, x.shape[-1]) output = torch.ops.sgl_kernel.weight_packed_linear( x, layer.weight, bias, True, # is_vnni ) if len(x_shapes) == 3: output = output.view(x_shapes[0], x_shapes[1], -1) return output return F.linear(x, layer.weight, bias) class UnquantizedFusedMoEMethod(FusedMoEMethodBase, MultiPlatformOp): """MoE method without quantization.""" def __init__( self, use_triton_kernels: bool = False, use_flashinfer_trtllm_moe: bool = False ): super().__init__() self.use_flashinfer_cutlass = get_moe_runner_backend().is_flashinfer_cutlass() self.use_triton_kernels = use_triton_kernels self.with_bias = False self.use_flashinfer_trtllm_moe = use_flashinfer_trtllm_moe self._cache_permute_indices = dict({}) def create_weights( self, layer: torch.nn.Module, num_experts: int, hidden_size: int, intermediate_size_per_partition: int, params_dtype: torch.dtype, with_bias: bool = False, **extra_weight_attrs, ): self.with_bias = with_bias # Fused gate_up_proj (column parallel) w13_up_dim = ( 2 * intermediate_size_per_partition if layer.moe_runner_config.is_gated else intermediate_size_per_partition ) w13_weight_n, w13_weight_k = (w13_up_dim, hidden_size) if self.use_triton_kernels: w13_weight_n, w13_weight_k = w13_weight_k, w13_weight_n w13_weight = torch.nn.Parameter( torch.empty(num_experts, w13_weight_n, w13_weight_k, dtype=params_dtype), requires_grad=False, ) layer.register_parameter("w13_weight", w13_weight) set_weight_attrs(w13_weight, extra_weight_attrs) if self.with_bias: w13_weight_bias = torch.nn.Parameter( torch.empty(num_experts, w13_up_dim, dtype=torch.float32), requires_grad=False, ) layer.register_parameter("w13_weight_bias", w13_weight_bias) set_weight_attrs(w13_weight_bias, extra_weight_attrs) # down_proj (row parallel) w2_weight_n, w2_weight_k = ( hidden_size, intermediate_size_per_partition, ) if self.use_triton_kernels: w2_weight_n, w2_weight_k = w2_weight_k, w2_weight_n w2_weight = torch.nn.Parameter( torch.empty(num_experts, w2_weight_n, w2_weight_k, dtype=params_dtype), requires_grad=False, ) layer.register_parameter("w2_weight", w2_weight) set_weight_attrs(w2_weight, extra_weight_attrs) if self.with_bias: w2_weight_bias = torch.nn.Parameter( torch.empty(num_experts, hidden_size, dtype=torch.float32), requires_grad=False, ) layer.register_parameter("w2_weight_bias", w2_weight_bias) set_weight_attrs(w2_weight_bias, extra_weight_attrs) def process_weights_after_loading(self, layer: torch.nn.Module) -> None: # Skip aiter weight shuffle when using non-auto MoE backend (e.g., triton, triton_kernels) # because aiter CK kernels don't support all GEMM dimensions _should_use_aiter_moe = _use_aiter and get_moe_runner_backend().is_auto() if _should_use_aiter_moe: layer.w13_weight = torch.nn.Parameter( shuffle_weight(layer.w13_weight.data, (16, 16)), requires_grad=False, ) torch.cuda.empty_cache() layer.w2_weight = torch.nn.Parameter( shuffle_weight(layer.w2_weight.data, (16, 16)), requires_grad=False, ) torch.cuda.empty_cache() # Pack weight for get better performance on CPU if _is_cpu and _is_cpu_amx_available: _amx_process_weight_after_loading(layer, ["w13_weight", "w2_weight"]) # Reorder rows of W1 for fused gated activation if self.use_flashinfer_trtllm_moe: from flashinfer.fused_moe.core import ( _maybe_get_cached_w3_w1_permute_indices, convert_to_block_layout, get_w2_permute_indices_with_cache, ) # w1 and w3 have been swapped, so we don't need do that here epilogue_tile_m = 128 block_k = 128 old_shape_w13 = layer.w13_weight.data[0].shape old_shape_w2 = layer.w2_weight.data[0].shape new_shape_w13 = None new_shape_w2 = None for i in range(layer.num_local_experts): permute_indices = _maybe_get_cached_w3_w1_permute_indices( self._cache_permute_indices, layer.w13_weight.data[i].view(torch.uint8), epilogue_tile_m, ) tmp_weights1 = ( layer.w13_weight.data[i] .clone() .view(torch.uint8)[permute_indices.to(layer.w13_weight.data.device)] .contiguous() ) permute_indices = get_w2_permute_indices_with_cache( self._cache_permute_indices, layer.w2_weight.data[i].view(torch.uint8), epilogue_tile_m, ) tmp_weights2 = ( layer.w2_weight.data[i] .clone() .view(torch.uint8)[permute_indices.to(layer.w2_weight.data.device)] .contiguous() ) tmp_weights1 = convert_to_block_layout( tmp_weights1.view(torch.uint8), block_k ) tmp_weights2 = convert_to_block_layout( tmp_weights2.view(torch.uint8), block_k ) new_shape_w13 = tmp_weights1.view(torch.bfloat16).shape new_shape_w2 = tmp_weights2.view(torch.bfloat16).shape layer.w13_weight.data[i] = ( tmp_weights1.view(torch.bfloat16) .contiguous() .reshape(old_shape_w13) ) layer.w2_weight.data[i] = ( tmp_weights2.view(torch.bfloat16).contiguous().reshape(old_shape_w2) ) layer.w13_weight.data = layer.w13_weight.data.reshape( layer.num_local_experts, *new_shape_w13 ) layer.w2_weight.data = layer.w2_weight.data.reshape( layer.num_local_experts, *new_shape_w2 ) if _is_npu: for weight_name in ["w13_weight", "w2_weight"]: weight = getattr(layer, weight_name) weight.data = weight.data.transpose(1, 2) weight.data = npu_format_cast( weight.data, ) return def create_moe_runner( self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig ): self.moe_runner_config = moe_runner_config backend = ( MoeRunnerBackend.TRITON_KERNELS if self.use_triton_kernels else MoeRunnerBackend.TRITON ) self.runner = MoeRunner(backend, moe_runner_config) @property def load_up_proj_weight_first(self) -> bool: # FlashInfer CUTLASS kernel assumes [Up, Gate] Proj as W13 return self.use_flashinfer_cutlass def apply( self, layer: torch.nn.Module, dispatch_output: StandardDispatchOutput, ) -> CombineInput: return self.forward( layer=layer, dispatch_output=dispatch_output, ) def forward_cuda( self, layer: torch.nn.Module, dispatch_output: StandardDispatchOutput, ) -> CombineInput: from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput x = dispatch_output.hidden_states topk_output = dispatch_output.topk_output moe_runner_config = self.moe_runner_config backend = self.runner.runner_backend if backend.is_triton_kernels(): from sglang.srt.layers.moe.moe_runner.triton_kernels import ( TritonKernelsQuantInfo, ) quant_info = TritonKernelsQuantInfo( w13_weight=layer.w13_weight, w2_weight=layer.w2_weight, w13_bias=getattr(layer, "w13_weight_bias", None), w2_bias=getattr(layer, "w2_weight_bias", None), ) return self.runner.run(dispatch_output, quant_info) elif self.use_flashinfer_cutlass: output = flashinfer_cutlass_fused_moe( input=x, token_selected_experts=topk_output.topk_ids, token_final_scales=topk_output.topk_weights, fc1_expert_weights=layer.w13_weight, fc2_expert_weights=layer.w2_weight, output_dtype=x.dtype, quant_scales=None, ep_size=layer.moe_ep_size, ep_rank=layer.moe_ep_rank, tp_size=layer.moe_tp_size, tp_rank=layer.moe_tp_rank, tune_max_num_tokens=next_power_of_2(x.shape[0]), )[0] return StandardCombineInput(hidden_states=output) else: # Skip aiter fused_moe when using non-auto MoE backend (e.g., triton, triton_kernels) # because aiter CK kernels don't support all GEMM dimensions _should_use_aiter_moe = _use_aiter and get_moe_runner_backend().is_auto() if _should_use_aiter_moe: assert not moe_runner_config.no_combine, "unsupported" topk_weights, topk_ids, _ = topk_output if moe_runner_config.apply_router_weight_on_input: assert ( topk_weights.dim() == 2 ), "`topk_weights` should be in shape (num_tokens, topk)" _, topk = topk_weights.shape assert ( topk == 1 ), "Only support topk=1 when `apply_router_weight_on_input` is True" x = x * topk_weights.to(x.dtype) topk_weights = torch.ones_like( topk_weights, dtype=torch.float32 ) # topk_weights must be FP32 (float32) output = fused_moe( x, layer.w13_weight, layer.w2_weight, topk_weights, topk_ids, activation=( ActivationType.Silu if moe_runner_config.activation == "silu" else ActivationType.Gelu ), expert_mask=layer.expert_mask_gpu, ) return StandardCombineInput(hidden_states=output) else: quant_info = TritonMoeQuantInfo( w13_weight=layer.w13_weight, w2_weight=layer.w2_weight, b13=getattr(layer, "w13_weight_bias", None), b2=getattr(layer, "w2_weight_bias", None), ) return self.runner.run(dispatch_output, quant_info) def forward_cpu( self, layer: torch.nn.Module, dispatch_output: StandardDispatchOutput, ) -> CombineInput: from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput x = dispatch_output.hidden_states topk_output = dispatch_output.topk_output moe_runner_config = self.moe_runner_config assert ( moe_runner_config.activation == "silu" ), f"activation = {moe_runner_config.activation} is not supported." if use_intel_amx_backend(layer): from sglang.srt.layers.moe.topk import apply_topk_weights_cpu topk_weights, topk_ids, _ = topk_output x, topk_weights = apply_topk_weights_cpu( moe_runner_config.apply_router_weight_on_input, topk_weights, x ) output = torch.ops.sgl_kernel.fused_experts_cpu( x, layer.w13_weight, layer.w2_weight, topk_weights, topk_ids, False, # inplace # See [Note] inplace should be False in fused_experts. CPUQuantMethod.UNQUANT, None, # w1_scale None, # w2_scale None, # w1_zp None, # w2_zp None, # block_size True, # is_vnni ) return StandardCombineInput(hidden_states=output) else: from sglang.srt.layers.moe.fused_moe_native import moe_forward_native output = moe_forward_native( layer, x, topk_output, moe_runner_config, ) return StandardCombineInput(hidden_states=output) def forward_xpu( self, layer: torch.nn.Module, dispatch_output: StandardDispatchOutput, ) -> CombineInput: from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput x = dispatch_output.hidden_states topk_output = dispatch_output.topk_output moe_runner_config = self.moe_runner_config assert moe_runner_config.activation in [ "silu", "gelu", ], f"activation = {moe_runner_config.activation} is not supported." backend = self.runner.runner_backend if use_intel_xpu_backend(): # sgl-kernel-xpu path from sgl_kernel import fused_experts topk_weights, topk_ids, _ = topk_output output = fused_experts( x, layer.w13_weight, layer.w2_weight, topk_weights, topk_ids, b1=getattr(layer, "w13_weight_bias", None), b2=getattr(layer, "w2_weight_bias", None), activation=moe_runner_config.activation, ) return StandardCombineInput(hidden_states=output) else: assert backend.is_triton() assert ( moe_runner_config.activation == "silu" ), f"activation = {moe_runner_config.activation} is not supported \ for Triton PATH, please set ENV SGLANG_USE_SGL_XPU=1." quant_info = TritonMoeQuantInfo( w13_weight=layer.w13_weight, w2_weight=layer.w2_weight, b13=getattr(layer, "w13_weight_bias", None), b2=getattr(layer, "w2_weight_bias", None), ) return self.runner.run(dispatch_output, quant_info) def forward_npu( self, layer: torch.nn.Module, dispatch_output: StandardDispatchOutput, ) -> CombineInput: import torch_npu from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput x = dispatch_output.hidden_states topk_weights, topk_ids, _ = dispatch_output.topk_output original_dtype = x.dtype num_tokens = x.shape[0] topk_weights = topk_weights.to(x.dtype) topk_ids = topk_ids.to(torch.int32) num_experts = layer.num_experts top_k = layer.top_k row_idx_len = num_tokens * top_k row_idx = ( torch.arange(0, row_idx_len, dtype=torch.int32, device=topk_weights.device) .view(top_k, -1) .permute(1, 0) .contiguous() ) hidden_states, expanded_row_idx, expanded_expert_idx = ( torch_npu.npu_moe_init_routing( x, row_idx=row_idx, expert_idx=topk_ids, active_num=num_tokens ) ) expert_tokens = torch_npu.npu_moe_compute_expert_tokens( expanded_expert_idx, num_experts ) expert_tokens = expert_tokens.to(torch.int64) w13_bias = [layer.w13_weight_bias] if self.with_bias else None w2_bias = [layer.w2_weight_bias] if self.with_bias else None # gmm1: gate_up_proj hidden_states = torch_npu.npu_grouped_matmul( x=[hidden_states], weight=[layer.w13_weight], bias=w13_bias, split_item=2, group_list_type=0, group_type=0, group_list=expert_tokens, output_dtype=original_dtype, )[0] # act_fn: if self.moe_runner_config.activation == "npu_swiglu_oai": from sgl_kernel_npu.activation.swiglu_oai import swiglu_oai hidden_states = swiglu_oai(layer, hidden_states) elif self.moe_runner_config.activation == "silu": hidden_states = torch_npu.npu_swiglu(hidden_states) else: from sglang.srt.layers.activation import GeluAndMul hidden_states = GeluAndMul()(hidden_states) # gmm2: down_proj hidden_states = torch_npu.npu_grouped_matmul( x=[hidden_states], weight=[layer.w2_weight], bias=w2_bias, split_item=2, group_list_type=0, group_type=0, group_list=expert_tokens, output_dtype=original_dtype, )[0] final_hidden_states = torch_npu.npu_moe_finalize_routing( hidden_states, skip1=None, skip2=None, bias=None, scales=topk_weights, expanded_src_to_dst_row=expanded_row_idx, export_for_source_row=topk_ids, ) return StandardCombineInput(hidden_states=final_hidden_states) def forward_tpu(self, *args, **kwargs) -> CombineInput: raise NotImplementedError("The TPU backend currently does not support MoE.") forward_native = forward_cpu