# # Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved. # This file is a part of the vllm-ascend project. # # 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. # # ruff: noqa: E501 from collections.abc import Callable from dataclasses import dataclass, field from functools import wraps import torch import torch.nn.functional as F import torch_npu from vllm.config import get_current_vllm_config from vllm.distributed import get_dp_group, get_ep_group, get_tp_group, tensor_model_parallel_all_reduce from vllm.forward_context import get_forward_context from vllm.logger import logger from vllm.model_executor.layers.fused_moe.config import FusedMoEConfig from vllm.model_executor.layers.fused_moe.layer import ( FusedMoE, # noqa: F401 MoERunner, ) from vllm_ascend.ascend_config import get_ascend_config from vllm_ascend.ascend_forward_context import _EXTRA_CTX, MoECommType from vllm_ascend.distributed.parallel_state import get_mc2_group from vllm_ascend.eplb.adaptor.vllm_adaptor import VllmEplbAdaptor from vllm_ascend.eplb.core.eplb_utils import init_eplb_config from vllm_ascend.flash_common3_context import get_flash_common3_context, set_flash_common3_context from vllm_ascend.ops.fused_moe.experts_selector import select_experts, zero_experts_compute from vllm_ascend.ops.fused_moe.moe_comm_method import AllGatherCommImpl, FusedExpertsResult, setup_moe_comm_method from vllm_ascend.ops.fused_moe.moe_runtime_args import build_fused_experts_input from vllm_ascend.quantization.methods.base import get_moe_num_logical_experts from vllm_ascend.quantization.quant_type import QuantType from vllm_ascend.utils import ( ACL_FORMAT_FRACTAL_NZ, maybe_trans_nz, npu_stream_switch, shared_expert_dp_enabled, shared_experts_calculation_stream, vllm_version_is, ) if vllm_version_is("0.23.0"): from vllm.model_executor.layers.fused_moe.layer import UnquantizedFusedMoEMethod else: from vllm.model_executor.layers.fused_moe.unquantized_fused_moe_method import UnquantizedFusedMoEMethod def get_compressed_expert_map(expert_map: torch.Tensor) -> str: global_indices = torch.where(expert_map != -1)[0] local_indices = expert_map[global_indices] return ", ".join( f"{local_index.item()}->{global_index.item()}" for local_index, global_index in zip(local_indices, global_indices) ) @dataclass class FusedMoEResult: routed_out: torch.Tensor before_dispatch_evt: torch.npu.Event | None = None before_gmm2_evt: torch.npu.Event | None = None before_combine_evt: torch.npu.Event | None = None swiglu_limit: float = 0.0 @dataclass class FusedMoEEvents: before_routed_experts: torch.npu.Event after_routed_experts: torch.npu.Event | None = field(default=None) before_dispatch: torch.npu.Event | None = field(default=None) before_gmm2: torch.npu.Event | None = field(default=None) before_combine: torch.npu.Event | None = field(default=None) swiglu_limit: float = 0.0 def mock_false(): return False def mock_true(): return True class AscendUnquantizedFusedMoEMethod(UnquantizedFusedMoEMethod): def __init__(self, moe: FusedMoEConfig = None, tid2eid=None): super().__init__(moe=moe) self.dynamic_eplb = get_ascend_config().eplb_config.dynamic_eplb self.tid2eid = tid2eid @property def is_monolithic(self) -> bool: return False def maybe_make_prepare_finalize(self, routing_tables=None): # Ascend uses its own MoE communication and forward_impl path. # Do not let upstream modular-kernel initialization replace it. return None def process_weights_after_loading(self, layer): super(UnquantizedFusedMoEMethod, self).process_weights_after_loading(layer) w13_data = self._maybe_pad_weight(layer.w13_weight.data).transpose(1, 2).contiguous() layer.w13_weight = torch.nn.Parameter(w13_data, requires_grad=False) w2_data = self._maybe_pad_weight(layer.w2_weight.data).transpose(1, 2).contiguous() layer.w2_weight = torch.nn.Parameter(w2_data, requires_grad=False) # TODO: Current dispatch_ffn_combine fusion operator ONLY supports NZ format. # Therefore, we must cast weights to NZ when fusion is enabled. # Once the underlying dispatch_ffn_combine operator is updated to support # ND format (or other formats), remove this specific 'if' check and the forced # npu_format_cast. At that point, the operator should be able to handle weights # in their native format without explicit casting here. enable_fused_mc2 = get_ascend_config().enable_fused_mc2 if enable_fused_mc2: layer.w13_weight.data = torch_npu.npu_format_cast(layer.w13_weight.data, ACL_FORMAT_FRACTAL_NZ) layer.w2_weight.data = torch_npu.npu_format_cast(layer.w2_weight.data, ACL_FORMAT_FRACTAL_NZ) if enable_fused_mc2 == 1 and self.dynamic_eplb: layer.w13_weight_list = [weight.clone() for weight in layer.w13_weight.data.unbind(dim=0)] layer.w2_weight_list = [weight.clone() for weight in layer.w2_weight.data.unbind(dim=0)] del layer.w13_weight del layer.w2_weight torch.npu.empty_cache() else: layer.w13_weight.data = maybe_trans_nz(layer.w13_weight.data) layer.w2_weight.data = maybe_trans_nz(layer.w2_weight.data) def apply( self, layer: torch.nn.Module, x: torch.Tensor, use_grouped_topk: bool, top_k: int, router_logits: torch.Tensor, renormalize: bool, topk_group: int | None = None, num_expert_group: int | None = None, custom_routing_function: Callable | None = None, scoring_func: str = "softmax", routed_scaling_factor: float = 1.0, e_score_correction_bias: torch.Tensor | None = None, num_experts: int = -1, expert_map: torch.Tensor | None = None, apply_router_weight_on_input: bool = False, activation: str = "silu", enable_force_load_balance: bool = False, log2phy: torch.Tensor = None, global_redundant_expert_num: int = 0, pertoken_scale: torch.Tensor | None = None, mc2_mask: torch.Tensor | None = None, ) -> torch.Tensor: zero_expert_num = getattr(layer, "zero_expert_num", 0) zero_expert_type = getattr(layer, "zero_expert_type", None) input_ids = getattr(get_forward_context(), "input_ids", None) num_shared_experts = getattr(layer, "n_shared_experts", 0) if num_shared_experts is None: num_shared_experts = 0 num_logical_experts = get_moe_num_logical_experts( layer, num_experts, global_redundant_expert_num=global_redundant_expert_num, num_shared_experts=num_shared_experts, ) topk_weights, topk_ids = select_experts( hidden_states=x, router_logits=router_logits, top_k=top_k, use_grouped_topk=use_grouped_topk, renormalize=renormalize, topk_group=topk_group, num_expert_group=num_expert_group, custom_routing_function=custom_routing_function, scoring_func=scoring_func, routed_scaling_factor=routed_scaling_factor, e_score_correction_bias=e_score_correction_bias, num_experts=num_logical_experts, tid2eid=self.tid2eid, input_ids=input_ids, ) if vllm_version_is("0.23.0"): model_config = layer.vllm_config.model_config else: try: _vllm_config = get_current_vllm_config() except AssertionError: _vllm_config = None model_config = None if _vllm_config is None else _vllm_config.model_config if model_config is not None and model_config.enable_return_routed_experts: capturer = getattr(layer, "_ascend_routed_experts_capturer", None) if capturer is not None: capturer.capture(layer_id=layer.layer_id, topk_ids=topk_ids) if zero_expert_num > 0 and zero_expert_type is not None: topk_ids, topk_weights, zero_expert_result = zero_experts_compute( expert_indices=topk_ids, expert_scales=topk_weights, num_experts=num_logical_experts, zero_expert_type=zero_expert_type, hidden_states=x, ) topk_weights = topk_weights.to(x.dtype) # this is a naive implementation for experts load balance so as # to avoid accumulating too much tokens on a single rank. # currently it is only activated when doing profile runs. if enable_force_load_balance: random_matrix = torch.rand(topk_ids.size(0), num_logical_experts, device=topk_ids.device) topk_ids = torch.argsort(random_matrix, dim=1)[:, : topk_ids.size(1)].to(topk_ids.dtype) moe_comm_method = _EXTRA_CTX.moe_comm_method # NOTE: In the MoECommType.FUSED_MC2 branch, we wrap weights (w1, w2) into lists # and provide dummy scales (w1_scale, w2_scale). This is required because: # The underlying Ascend fused operator (e.g., dispatch_ffn_combine) expects # inputs in a list format. # TODO: Passing an empty tensor as scale for float (BF16) cases is semantically # incorrect. The ideal solution is to pass None. However, if the underlying # dispatch_ffn_combine C++ operator does not support None for the scale argument # (due to signature constraints), we are forced to use a placeholder empty tensor. # This TODO tracks the requirement to update the C++ operator to accept Optional[Tensor] # or None for scales in non-quantized scenarios. w13_weight_list = getattr(layer, "w13_weight_list", None) w2_weight_list = getattr(layer, "w2_weight_list", None) has_split_weight_lists = isinstance(w13_weight_list, list) and isinstance(w2_weight_list, list) if _EXTRA_CTX.moe_comm_type == MoECommType.FUSED_MC2: if self.dynamic_eplb and not has_split_weight_lists: logger.warning_once( "FUSED_MC2 is enabled with dynamic EPLB, but unquantized MoE weights are not split into " "tensor lists. This may cause accuracy issues or communication hangs." ) w1 = w13_weight_list if isinstance(w13_weight_list, list) else [layer.w13_weight] w2 = w2_weight_list if isinstance(w2_weight_list, list) else [layer.w2_weight] w1_scale = [torch.tensor([], dtype=torch.int64)] w2_scale = [torch.tensor([], dtype=torch.int64)] w1_scale_bias = [torch.tensor([], dtype=torch.float32)] w2_scale_bias = [torch.tensor([], dtype=torch.float32)] else: w1 = w13_weight_list if isinstance(w13_weight_list, list) else layer.w13_weight w1_scale = None w2 = w2_weight_list if isinstance(w2_weight_list, list) else layer.w2_weight w2_scale = None w1_scale_bias = None w2_scale_bias = None final_hidden_states = moe_comm_method.fused_experts( fused_experts_input=build_fused_experts_input( hidden_states=x, topk_weights=topk_weights, topk_ids=topk_ids, w1=w1, w2=w2, w1_bias=layer.w13_bias if self.moe.has_bias else None, w2_bias=layer.w2_bias if self.moe.has_bias else None, quant_type=QuantType.NONE, dynamic_eplb=self.dynamic_eplb, expert_map=expert_map, global_redundant_expert_num=global_redundant_expert_num, mc2_mask=mc2_mask, apply_router_weight_on_input=apply_router_weight_on_input, log2phy=log2phy, pertoken_scale=pertoken_scale, activation=activation, w1_scale=w1_scale, w2_scale=w2_scale, w1_scale_bias=w1_scale_bias, w2_scale_bias=w2_scale_bias, swiglu_limit=layer.swiglu_limit, # Per-layer MoE LoRA state, set once by AscendFusedMoEWithLoRA # when an adapter wraps this layer; None for non-LoRA layers. lora_context=getattr(layer, "_ascend_moe_lora_context", None), ) ) if zero_expert_num > 0 and zero_expert_type is not None: final_hidden_states += zero_expert_result return final_hidden_states if vllm_version_is("0.23.0"): from vllm_ascend.ops.fused_moe.fused_moe_0_23_0 import AscendFusedMoE, AscendMoERunner AscendFusedMoE.__module__ = __name__ AscendMoERunner.__module__ = __name__ else: class AscendMoERunner(MoERunner): # type: ignore[no-redef] moe_counter = -1 gate_stream: torch.npu.Stream | None = None def __init__( self, layer_name, moe_config, router, routed_experts, enable_dbo=False, gate=None, shared_experts=None, shared_expert_gate=None, routed_input_transform=None, routed_output_transform=None, routed_scaling_factor=1, tid2eid=None, n_shared_experts: int = 0, ): super().__init__( layer_name, moe_config, router, routed_experts, enable_dbo, gate, shared_experts, shared_expert_gate, routed_input_transform, routed_output_transform, routed_scaling_factor, ) self.top_k = moe_config.experts_per_token self._gate = gate self.hidden_size = moe_config.hidden_dim # Routing params — all stored on routed_experts by the factory function. self.use_grouped_topk = routed_experts.use_grouped_topk self.renormalize = routed_experts.renormalize self.topk_group = routed_experts.topk_group self.num_expert_group = routed_experts.num_expert_group self.custom_routing_function = routed_experts.custom_routing_function self.scoring_func = routed_experts.scoring_func self._original_routed_scaling_factor = routed_experts.routed_scaling_factor self.e_score_correction_bias = routed_experts.e_score_correction_bias self.apply_router_weight_on_input = routed_experts.apply_router_weight_on_input # Ascend-specific: not stored on RoutedExperts, passed via runner_args. self.tid2eid = tid2eid # Replace quant_method on routed_experts with the Ascend version. # Must NOT set self.quant_method (instance attr) — vllm loader scans all # nn.Module children for quant_method and would call process_weights_after_loading # on the runner with wrong expectations. Set on routed_experts so that # self._quant_method (the MoERunner property) returns the Ascend version. if routed_experts.quant_config is None: routed_experts.quant_method = AscendUnquantizedFusedMoEMethod(self.moe_config, tid2eid=self.tid2eid) else: routed_experts.quant_method = routed_experts.quant_config.get_quant_method( routed_experts, self.layer_name, tid2eid=self.tid2eid ) self.quant_type = self._get_quant_type() # Can be removed after vllm fixes the issue. if self._needs_routed_expert_parameter_aliases(): self._register_routed_expert_parameter_aliases() self.moe_config.tp_group = get_tp_group() self.moe_config.dp_group = get_dp_group() if self.moe_config.ep_size > 1: self.moe_config.ep_group = get_ep_group() self.moe_config.mc2_group = get_mc2_group() ascend_config = get_ascend_config() self._shared_experts = shared_experts has_shared_experts = self._shared_experts is not None self.shared_multistream_overlap_gate = ascend_config.multistream_overlap_gate and has_shared_experts self.enable_npugraph_ex_static_kernel = ascend_config.ascend_compilation_config.enable_static_kernel # flashcommon3 gate stream self.multistream_overlap_gate = ascend_config.multistream_overlap_gate if self.multistream_overlap_gate and AscendMoERunner.gate_stream is None: AscendMoERunner.gate_stream = torch.npu.Stream() vllm_config = get_current_vllm_config() if ( self.custom_routing_function is None and self.e_score_correction_bias is not None and not vllm_config.model_config.is_deepseek_mla ): self.e_score_correction_bias.data = self.e_score_correction_bias.data.to( dtype=vllm_config.model_config.dtype ) self.enable_shared_expert_dp = ascend_config.enable_shared_expert_dp self.multistream_overlap_shared_expert = ( ascend_config.multistream_overlap_shared_expert and shared_experts is not None ) mix_placement = getattr(ascend_config, "mix_placement", False) # EPLB initialization (Ascend-specific; mirrors old AscendFusedMoE logic). AscendMoERunner.moe_counter += 1 self.moe_instance_id = AscendMoERunner.moe_counter eplb_config = ascend_config.eplb_config if mix_placement: moe_config.num_experts += n_shared_experts ( self.global_expert_map, self._expert_map, self.log2phy, self.global_redundant_expert_num, ) = init_eplb_config( eplb_config, AscendMoERunner.moe_counter, moe_config, mix_placement, n_shared_experts, tp_size=vllm_config.parallel_config.tensor_parallel_size, ) moe_config.global_redundant_expert_num = self.global_redundant_expert_num local_num_experts = (moe_config.num_experts + self.global_redundant_expert_num) // moe_config.ep_size moe_config.num_local_experts = local_num_experts routed_experts.expert_map_manager._local_num_experts = local_num_experts routed_experts.expert_map_manager._expert_map = self._expert_map self.dynamic_eplb = eplb_config.dynamic_eplb and (self.log2phy is not None) self.multi_stage = False self.moe_load = torch.zeros(local_num_experts, dtype=torch.int64).npu() if self.dynamic_eplb and eplb_config.expert_heat_collection_interval > 1: self.multi_stage = True self.load_counter = torch.tensor(0, dtype=torch.int32, device="npu") self.num_iter = eplb_config.expert_heat_collection_interval self.moe_load = torch.zeros((self.num_iter, local_num_experts), dtype=torch.int32, device="npu") setup_moe_comm_method(self.moe_config) if self.multistream_overlap_shared_expert: # Wrap the quant_method's process_weights_after_loading to validate that # splitting shared expert computation (gate_up projection + activation, # then down projection) yields identical results to integrated # computation after weight loading. original_process_weights = self._quant_method.process_weights_after_loading @wraps(original_process_weights) def wrapped_process_weights(*args, **kwargs): result = original_process_weights(*args, **kwargs) self._validate_shared_expert_consistency() return result self._quant_method.process_weights_after_loading = wrapped_process_weights # type: ignore # Register this MoE layer with EPLB for PP compatibility. # PPMissingLayer (nn.Identity) never calls AscendFusedMoE.__init__, # so only real MoE layers on this rank are registered. VllmEplbAdaptor.register_layer(self) def _validate_shared_expert_consistency(self): """Validate that split shared expert computation matches integrated computation.""" test_input = ( torch.rand(10, self.hidden_size, device="npu", dtype=self.moe_config.in_dtype) * 2 - 1 ) # Random input for testing, scoped to [-1, 1] assert self._shared_experts is not None integrated_out = self._shared_experts(test_input) part1_out = self._shared_experts_part1(test_input) split_out = self._shared_experts_part2(test_input, part1_out) if not torch.allclose(integrated_out, split_out): diff = (integrated_out - split_out).abs() logger.error( "[fused_moe/layer] Shared expert split computation validation failed." " The split-path computation does not match the integrated-path result." " max_abs_diff=%s, integrated_sum=%s, integrated_norm=%s," " split_sum=%s, split_norm=%s, hidden_size=%s, dtype=%s.", diff.max().item(), integrated_out.sum().item(), integrated_out.norm().item(), split_out.sum().item(), split_out.norm().item(), self.hidden_size, self.moe_config.in_dtype, ) raise ValueError("FusedMoE shared experts split computation does not match the integrated computation.") logger.info_once( "[fused_moe/layer] Shared expert split computation validation passed." " Integrated and split-path results are consistent." ) def _shared_experts_part1(self, hidden_states: torch.Tensor): shared_gate_up, _ = self._shared_experts.gate_up_proj(hidden_states) # type: ignore return shared_gate_up def _shared_experts_part2(self, hidden_states: torch.Tensor, shared_gate_up: torch.Tensor): shared_act = self._shared_experts.act_fn(shared_gate_up) # type: ignore shared_out, _ = self._shared_experts.down_proj(shared_act) # type: ignore # Qwen3-Next specific gating mechanism assert self._shared_experts is not None if hasattr(self._shared_experts, "expert_gate") and self._shared_experts.expert_gate is not None: gate_out, _ = self._shared_experts.expert_gate(hidden_states) # type: ignore shared_out = F.sigmoid(gate_out) * shared_out return shared_out def _get_quant_type(self) -> QuantType: quant_type = QuantType.NONE method = getattr(self._quant_method, "quant_method", None) if method is not None: quant_type = getattr(method, "quant_type", QuantType.NONE) return quant_type def _register_routed_expert_parameter_aliases(self) -> None: alias_names = [] for name, param in self.routed_experts.named_parameters(recurse=False): alias_param = torch.nn.Parameter(param.data, requires_grad=param.requires_grad) alias_param.__dict__.update(param.__dict__) self.register_parameter(name, alias_param) alias_names.append(name) original_process_weights = self._quant_method.process_weights_after_loading @wraps(original_process_weights) def wrapped_process_weights(layer, *args, **kwargs): for name in alias_names: self._parameters.pop(name, None) return original_process_weights(layer, *args, **kwargs) self._quant_method.process_weights_after_loading = wrapped_process_weights # type: ignore[method-assign] def _needs_routed_expert_parameter_aliases(self) -> bool: vllm_config = get_current_vllm_config() hf_config = getattr(vllm_config.model_config, "hf_config", None) return getattr(hf_config, "model_type", None) == "gpt_oss" @property def is_internal_router(self) -> bool: gate = self.gate return gate is not None and hasattr(gate, "weight_fp32") @property def use_dp_chunking(self) -> bool: """Ascend uses its own forward_impl path, not the FlashInfer Cutlass chunked path. Always return False to stay on forward_impl.""" return False @property def _fused_output_is_reduced(self) -> bool: # For MC2/ALLTOALL/FUSED_MC2 comm types, finalize() already includes # TP all-reduce for the routed output, and _forward_shared_experts # handles it for the shared output. Signal this to the upstream # MoERunner.forward() so _maybe_reduce_final_output does not apply a # second TP all-reduce (which would double-count the contributions). moe_comm_type = _EXTRA_CTX.moe_comm_type return moe_comm_type in { MoECommType.ALLTOALL, MoECommType.MC2, MoECommType.FUSED_MC2, } or (moe_comm_type == MoECommType.ALLGATHER and _EXTRA_CTX.flash_comm_v1_enabled) def _maybe_reduce_shared_expert_output( self, shared_output: torch.Tensor | None, ) -> torch.Tensor | None: # _forward_shared_experts already handles shared expert TP all-reduce # for MC2/ALLTOALL/FUSED_MC2. For AllGather the reduction is done # via _maybe_reduce_final_output on the combined (shared + routed) # output. Skip any additional reduction here. return shared_output def _maybe_reduce_final_output( self, states: torch.Tensor, trunc_size: int, ) -> torch.Tensor: states = torch.ops.vllm.maybe_all_reduce_tensor_model_parallel(states) return states[..., :trunc_size] def set_lora_context(self, lora_context): self.routed_experts._ascend_moe_lora_context = lora_context def no_shared_forward_impl( # type: ignore[override] self, hidden_states: torch.Tensor, router_logits: torch.Tensor, return_with_event: bool = False ) -> torch.Tensor | FusedMoEResult: forward_context = get_forward_context() # When static kernels are enabled, the forward pass runs twice (compilation + capture), # causing moe_layer_index to overflow. Wrap the index to prevent out-of-bounds errors. if self.enable_npugraph_ex_static_kernel and forward_context.all_moe_layers: moe_layer_index = forward_context.moe_layer_index % (len(forward_context.all_moe_layers)) forward_context.moe_layer_index = moe_layer_index # Load balancing for token distribution among experts in dummy_run # TODO: The community only considers load balancing when DP > 1. # This approach may overlook some extreme scenarios. enable_force_load_balance = _EXTRA_CTX.in_profile_run forward_context = get_forward_context() if self.multistream_overlap_gate: fc3_context = get_flash_common3_context() assert fc3_context is not None assert AscendMoERunner.gate_stream is not None AscendMoERunner.gate_stream.wait_stream(torch.npu.current_stream()) with npu_stream_switch(AscendMoERunner.gate_stream, enabled=self.multistream_overlap_gate): # share_expert assert fc3_context.shared_experts is not None shared_out = fc3_context.shared_experts(hidden_states) # NOTE: This is exactly the opposite of `maybe_all_reduce_tensor_model_parallel` moe_comm_type = _EXTRA_CTX.moe_comm_type if ( moe_comm_type in {MoECommType.ALLTOALL, MoECommType.MC2, MoECommType.FUSED_MC2} and not shared_expert_dp_enabled() ): shared_out = tensor_model_parallel_all_reduce(shared_out) set_flash_common3_context(shared_out=shared_out) input_ids = getattr(get_forward_context(), "input_ids", None) topk_weights, topk_ids = select_experts( hidden_states=hidden_states, router_logits=router_logits, top_k=self.top_k, use_grouped_topk=self.use_grouped_topk, renormalize=self.renormalize, topk_group=self.topk_group, num_expert_group=self.num_expert_group, custom_routing_function=self.custom_routing_function, scoring_func=self.scoring_func, routed_scaling_factor=self._original_routed_scaling_factor, e_score_correction_bias=self.e_score_correction_bias, num_experts=self.moe_config.num_experts, input_ids=input_ids, tid2eid=self.tid2eid, ) if isinstance(_EXTRA_CTX.moe_comm_method, AllGatherCommImpl): topk_weights = torch.ops.vllm.maybe_all_gather_and_maybe_unpad(topk_weights, True, True) topk_ids = torch.ops.vllm.maybe_all_gather_and_maybe_unpad(topk_ids, True, True) set_flash_common3_context(topk_weights=topk_weights, topk_ids=topk_ids) prepare_output = _EXTRA_CTX.moe_comm_method.prepare( hidden_states=hidden_states, router_logits=router_logits, replace_allreduce=_EXTRA_CTX.flash_comm_v1_enabled, enable_shared_expert_dp=self.enable_shared_expert_dp, quant_type=self.quant_type, ) hidden_states = prepare_output.hidden_states router_logits = prepare_output.router_logits mc2_mask = prepare_output.mc2_mask padded_hidden_states_shape = prepare_output.padded_hidden_states_shape pertoken_scale = prepare_output.pertoken_scale # Make sure the default stream waits for the gate stream to finish. if self.multistream_overlap_gate: assert AscendMoERunner.gate_stream is not None torch.npu.current_stream().wait_stream(AscendMoERunner.gate_stream) # Matrix multiply. # apply() expects a RoutedExperts-like layer for weight access # (w13_weight, w2_weight, swiglu_limit, etc.). Pass routed_experts, # not self; the routing params come through the other kwargs. fused_experts_results: FusedExpertsResult = self._quant_method.apply( layer=self.routed_experts, x=hidden_states, router_logits=router_logits, pertoken_scale=pertoken_scale, top_k=self.top_k, renormalize=self.renormalize, use_grouped_topk=self.use_grouped_topk, num_experts=self.moe_config.num_experts, expert_map=self._expert_map, topk_group=self.topk_group, num_expert_group=self.num_expert_group, custom_routing_function=self.custom_routing_function, scoring_func=self.scoring_func, routed_scaling_factor=self._original_routed_scaling_factor, e_score_correction_bias=self.e_score_correction_bias, activation=self.activation, apply_router_weight_on_input=self.apply_router_weight_on_input, enable_force_load_balance=enable_force_load_balance, log2phy=self.log2phy, global_redundant_expert_num=self.global_redundant_expert_num, mc2_mask=mc2_mask, ) if self.dynamic_eplb and _EXTRA_CTX.eplb_heat_collection_status: expert_tokens = fused_experts_results.expert_tokens group_list_type = fused_experts_results.group_list_type assert expert_tokens is not None and group_list_type is not None, ( "expert_tokens and group_list_type should not be None when dynamic_eplb is enabled." ) local_load = ( expert_tokens if group_list_type == 1 else torch.cat([expert_tokens[:1], expert_tokens[1:] - expert_tokens[:-1]]) ) if self.multi_stage: cur_iter = torch.remainder(self.load_counter, self.num_iter) self.moe_load.index_add_( dim=0, index=cur_iter, source=local_load.to(torch.int32, non_blocking=True).view(1, -1) ) self.load_counter.add_(1) else: self.moe_load.add_(local_load) routed_out = _EXTRA_CTX.moe_comm_method.finalize( hidden_states=fused_experts_results.routed_out, reduce_results=isinstance(_EXTRA_CTX.moe_comm_method, AllGatherCommImpl), padded_hidden_states_shape=padded_hidden_states_shape, ) if return_with_event: return FusedMoEResult( routed_out=routed_out, before_dispatch_evt=fused_experts_results.before_dispatch_evt, before_gmm2_evt=fused_experts_results.before_gmm2_evt, before_combine_evt=fused_experts_results.before_combine_evt, swiglu_limit=fused_experts_results.swiglu_limit, ) else: # The vLLM FusedMoE forward_impl does not return events. return routed_out def _forward_shared_experts(self, hidden_states: torch.Tensor, fused_moe_evts: FusedMoEEvents): if self._shared_experts is None: return None def maybe_wait_event(evt: torch.npu.Event | None): if evt is not None: torch.npu.current_stream().wait_event(evt) with npu_stream_switch(shared_experts_calculation_stream(), enabled=self.multistream_overlap_shared_expert): # Only used for int quantization has_quantized_shared = hasattr(self._shared_experts.gate_up_proj, "weight_scale") and hasattr( self._shared_experts.down_proj, "weight_scale" ) if has_quantized_shared and self.quant_type in (QuantType.W8A8, QuantType.W4A8): original_dtype = hidden_states.dtype # Execute dynamic quant concurrently with MoE gate. torch.npu.current_stream().wait_event(fused_moe_evts.before_routed_experts) quantized_x, pertoken_scale = torch_npu.npu_dynamic_quant(hidden_states) # Execute the gate projection and activation concurrently with the # dispatch communication. maybe_wait_event(fused_moe_evts.after_routed_experts) hidden_states = torch_npu.npu_quant_matmul( quantized_x, self._shared_experts.gate_up_proj.weight, self._shared_experts.gate_up_proj.weight_scale, pertoken_scale=None, bias=None, output_dtype=torch.int32, ) # Execute activation concurrently with gmm2. maybe_wait_event(fused_moe_evts.before_gmm2) quantized_x, swiglu_out_scale = torch.ops._C_ascend.npu_dequant_swiglu_quant( x=hidden_states, weight_scale=self._shared_experts.gate_up_proj.weight_scale_fp32, activation_scale=pertoken_scale, bias=None, quant_scale=None, quant_offset=None, group_index=None, activate_left=True, quant_mode=1, swiglu_mode=1, clamp_limit=fused_moe_evts.swiglu_limit, ) # Execute the down projection concurrently with the combine # communication. maybe_wait_event(fused_moe_evts.before_combine) shared_out = torch_npu.npu_quant_matmul( quantized_x, self._shared_experts.down_proj.weight, self._shared_experts.down_proj.weight_scale, pertoken_scale=swiglu_out_scale, bias=None, output_dtype=original_dtype, ) elif has_quantized_shared and self.quant_type == QuantType.W4A8MXFP: original_dtype = hidden_states.dtype # Execute dynamic quant concurrently with MoE gate. torch.npu.current_stream().wait_event(fused_moe_evts.before_routed_experts) quantized_x, pertoken_scale = torch_npu.npu_dynamic_mx_quant( hidden_states, dst_type=torch.float8_e4m3fn ) # Execute the gate projection and activation concurrently with the # dispatch communication. maybe_wait_event(fused_moe_evts.before_dispatch) hidden_states = self._shared_experts.gate_up_proj((quantized_x, pertoken_scale))[0] # Execute activation concurrently with gmm2. maybe_wait_event(fused_moe_evts.before_gmm2) quantized_x, swiglu_out_scale, _ = torch.ops._C_ascend.npu_swiglu_group_quant( hidden_states, topk_weight=None, group_index=None, dst_type=torch.float8_e4m3fn, quant_mode=2, clamp_value=fused_moe_evts.swiglu_limit, ) # Execute the down projection concurrently with the combine # communication. maybe_wait_event(fused_moe_evts.before_combine) shared_out = self._shared_experts.down_proj((quantized_x, swiglu_out_scale))[0] else: # Ensure the shared experts wait for hidden_states to be ready. torch.npu.current_stream().wait_event(fused_moe_evts.before_routed_experts) # Execute the gate projection and activation concurrently with the # dispatch communication. maybe_wait_event(fused_moe_evts.before_dispatch) part1_out = self._shared_experts_part1(hidden_states) # Execute the down projection concurrently with the combine # communication. maybe_wait_event(fused_moe_evts.before_combine) shared_out = self._shared_experts_part2(hidden_states, part1_out) # Make sure the default stream waits for the shared experts stream to # finish. if self.multistream_overlap_shared_expert: torch.npu.current_stream().wait_stream(shared_experts_calculation_stream()) # NOTE: This is exactly the opposite of # `maybe_all_reduce_tensor_model_parallel` moe_comm_type = _EXTRA_CTX.moe_comm_type if ( moe_comm_type in {MoECommType.ALLTOALL, MoECommType.MC2, MoECommType.FUSED_MC2} and not shared_expert_dp_enabled() ): shared_out = tensor_model_parallel_all_reduce(shared_out) return shared_out def shared_forward_impl( # type: ignore[override] self, hidden_states: torch.Tensor, router_logits: torch.Tensor ): if self.shared_multistream_overlap_gate: set_flash_common3_context(shared_experts=self._shared_experts) if self.is_internal_router: gate = self.gate assert gate is not None # NOTE(Angazenn): To make this cast explicitly, the hbm usage might # increase with extra hidden states. We also assume that all gate # linear is unquantized so that we the weight is pre-casted in # process_weights_after_loading of AscendUnquantizedLinearMethod. hidden_states_fp32 = hidden_states.float() before_routed_experts = torch.npu.current_stream().record_event() router_logits = F.linear(hidden_states_fp32, gate.weight_fp32) after_routed_experts = torch.npu.current_stream().record_event() else: before_routed_experts = torch.npu.current_stream().record_event() after_routed_experts = None fused_moe_results = self.no_shared_forward_impl( hidden_states, router_logits, return_with_event=True, ) routed_out = fused_moe_results.routed_out if self._shared_experts is None: return routed_out if self.shared_multistream_overlap_gate: fc3_context = get_flash_common3_context() assert fc3_context is not None shared_out = fc3_context.shared_out else: shared_out = self._forward_shared_experts( hidden_states, FusedMoEEvents( after_routed_experts=after_routed_experts, before_routed_experts=before_routed_experts, before_dispatch=fused_moe_results.before_dispatch_evt, before_gmm2=fused_moe_results.before_gmm2_evt, before_combine=fused_moe_results.before_combine_evt, swiglu_limit=fused_moe_results.swiglu_limit, ), ) return shared_out, routed_out def _forward_impl( self, hidden_states: torch.Tensor, router_logits: torch.Tensor, shared_experts_input: torch.Tensor | None, input_ids: torch.Tensor | None = None, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: with self._sequence_parallel_context(): if self.shared_experts is None: return self.no_shared_forward_impl(hidden_states, router_logits) else: return self.shared_forward_impl(hidden_states, router_logits)