大融合算子-MindIE-LLM-AtomGit

Mmindie_llm/runtime/layers/fused_moe/fused_moe.py+19-8

	@@ -14,7 +14,6 @@ from typing import List
14	import torch	14	import torch
15	import torch_npu	15	import torch_npu
16	import torch.distributed as dist	16	import torch.distributed as dist
17	-import mie_ops
18		17
19	from mindie_llm.runtime.layers.custom_layer import CustomLayer	18	from mindie_llm.runtime.layers.custom_layer import CustomLayer
20	from mindie_llm.runtime.layers.fused_moe.moe_comm_method import (	19	from mindie_llm.runtime.layers.fused_moe.moe_comm_method import (
	@@ -34,6 +33,10 @@ from mindie_llm.runtime.utils.distributed.parallel_info_manager import ParallelT
34	from mindie_llm.runtime.utils.distributed.utils import even_divide	33	from mindie_llm.runtime.utils.distributed.utils import even_divide
35	from mindie_llm.runtime.model_runner.forward_context import get_forward_context	34	from mindie_llm.runtime.model_runner.forward_context import get_forward_context
36		35
		36	+# Recommended value for fused operator output buffer size (in elements).
		37	+# Determined empirically based on current npu_dispatch_ffn_combine operator constraints.
		38	+MAX_OUTPUT_SIZE: int = 65536
		39	+
37		40
38	class FusedMoE(CustomLayer):	41	class FusedMoE(CustomLayer):
39	def __init__(	42	def __init__(
	@@ -49,7 +52,6 @@ class FusedMoE(CustomLayer):
49	prefix: str = "",	52	prefix: str = "",
50	suffix: List[str] = None,	53	suffix: List[str] = None,
51	activation: str = "silu",	54	activation: str = "silu",
52	- is_fused_op_enabled: bool = False,
53	):	55	):
54	super().__init__()	56	super().__init__()
55		57
	@@ -69,6 +71,11 @@ class FusedMoE(CustomLayer):
69	self.activation = activation	71	self.activation = activation
70	self.num_experts = num_experts	72	self.num_experts = num_experts
71	self.topk = topk_num	73	self.topk = topk_num
		74	+ # Max receivable tokens per device.
		75	+ # Formula is MAX_OUTPUT_SIZE / load_balance_factor / topk
		76	+ # Note: load_balance_factor=2 is recommended by the fused operator spec
		77	+ # to account for load balancing skew (reserves 2x the average per-card capacity).
		78	+ self.max_num_tokens_per_device = MAX_OUTPUT_SIZE // 2 // self.topk
72	self.expert_list = assign_experts(self.num_experts, self.moe_ep_size)[self.moe_ep_rank]	79	self.expert_list = assign_experts(self.num_experts, self.moe_ep_size)[self.moe_ep_rank]
73	self.num_local_experts = len(self.expert_list)	80	self.num_local_experts = len(self.expert_list)
74	self.expert_map = torch.full(size=(self.num_experts,), fill_value=-1, device='npu')	81	self.expert_map = torch.full(size=(self.num_experts,), fill_value=-1, device='npu')
	@@ -78,10 +85,9 @@ class FusedMoE(CustomLayer):
78	else:	85	else:
79	# Get moe quant method through gate proj weights of expert 0	86	# Get moe quant method through gate proj weights of expert 0
80	self.quant_method = self.quant_config.get_quant_method(self, prefix=f"{self.prefix}.0.{self.suffix[0]}")	87	self.quant_method = self.quant_config.get_quant_method(self, prefix=f"{self.prefix}.0.{self.suffix[0]}")
81	-
82	- self.is_fused_op_enabled = is_fused_op_enabled
83	# Only created when MC2 used fused op	88	# Only created when MC2 used fused op
84	self._moe_ep_group = None	89	self._moe_ep_group = None
		90	+
85	self._create_weights()	91	self._create_weights()
86	self._post_init()	92	self._post_init()
87		93
	@@ -117,9 +123,13 @@ class FusedMoE(CustomLayer):
117	topk_ids: torch.Tensor,	123	topk_ids: torch.Tensor,
118	) -> torch.Tensor:	124	) -> torch.Tensor:
119	moe_comm_type = select_moe_comm_method(	125	moe_comm_type = select_moe_comm_method(
120	- quant_type=None	126	+ quant_type=None,
		127	+ max_num_tokens_per_device=self.max_num_tokens_per_device
121	)	128	)
122	- if moe_comm_type == MoECommType.MC2 and self.is_fused_op_enabled:	129	+
		130	+ if moe_comm_type == MoECommType.FUSED_MC2:
		131	+ import mie_ops
		132	+ # FUSED_MC2 mode: Directly invoke the fused dispatch + FFN + combine operator
123	self._create_moe_ep_group()	133	self._create_moe_ep_group()
124	final_hidden_states = torch.empty_like(hidden_states)	134	final_hidden_states = torch.empty_like(hidden_states)
125	torch.ops.mie_ops.npu_dispatch_ffn_combine(	135	torch.ops.mie_ops.npu_dispatch_ffn_combine(
	@@ -131,7 +141,7 @@ class FusedMoE(CustomLayer):
131	scale2=[self.fused_down_weight_scale],	141	scale2=[self.fused_down_weight_scale],
132	probs=topk_weights.to(self.weight_dtype).to(torch.float32),	142	probs=topk_weights.to(self.weight_dtype).to(torch.float32),
133	group=self._moe_ep_group,	143	group=self._moe_ep_group,
134	- max_output_size=65536,	144	+ max_output_size=MAX_OUTPUT_SIZE,
135	out=final_hidden_states,	145	out=final_hidden_states,
136	)	146	)
137	return final_hidden_states	147	return final_hidden_states
	@@ -222,11 +232,12 @@ class FusedMoE(CustomLayer):
222	bias_dtype=self.weight_dtype,	232	bias_dtype=self.weight_dtype,
223	weight_loader=None,	233	weight_loader=None,
224	)	234	)
		235	+
225	def _create_moe_ep_group(self):	236	def _create_moe_ep_group(self):
226	if self._moe_ep_group is not None:	237	if self._moe_ep_group is not None:
227	return	238	return
228	parallel_info_manager = get_parallel_info_manager()	239	parallel_info_manager = get_parallel_info_manager()
229	- device_group = parallel_info_manager.moe_ep_mc2.process_group	240	+ device_group = parallel_info_manager.get(ParallelType.MOE_EP_MC2).process_group
230	local_rank = dist.get_rank(group=device_group)	241	local_rank = dist.get_rank(group=device_group)
231	backend = device_group._get_backend(torch.device("npu"))	242	backend = device_group._get_backend(torch.device("npu"))
232	self._moe_ep_group = backend.get_hccl_comm_name(local_rank)	243	self._moe_ep_group = backend.get_hccl_comm_name(local_rank)

Mmindie_llm/runtime/layers/fused_moe/moe_comm_method.py+40-88

	@@ -1,7 +1,3 @@
1	-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
2	-# SPDX-License-Identifier: Apache-2.0
3	-#
4	-# Implement part of this file based on vllm-project/vllm
5	# Copyright (c) Huawei Technologies Co., Ltd. 2025-2026. All rights reserved.	1	# Copyright (c) Huawei Technologies Co., Ltd. 2025-2026. All rights reserved.
6	# MindIE is licensed under Mulan PSL v2.	2	# MindIE is licensed under Mulan PSL v2.
7	# You can use this software according to the terms and conditions of the Mulan PSL v2.	3	# You can use this software according to the terms and conditions of the Mulan PSL v2.
	@@ -12,32 +8,21 @@
12	# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.	8	# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
13	# See the Mulan PSL v2 for more details.	9	# See the Mulan PSL v2 for more details.
14		10
15	-from enum import Enum	11	+from typing import Optional, Dict, Type
16	-from typing import Optional
17		12
18	-from mindie_llm.runtime.utils.npu.device_utils import DeviceType, get_npu_node_info
19	from mindie_llm.runtime.layers.fused_moe.token_dispatcher import (	13	from mindie_llm.runtime.layers.fused_moe.token_dispatcher import (
20	MoETokenDispatcher,	14	MoETokenDispatcher,
21	TokenDispatcherWithAllGather,	15	TokenDispatcherWithAllGather,
22	TokenDispatcherWithMC2,	16	TokenDispatcherWithMC2,
23	TokenDispatcherWithAll2AllV	17	TokenDispatcherWithAll2AllV
24	)	18	)
25	-from mindie_llm.runtime.utils.distributed import get_parallel_info_manager	19	+from mindie_llm.runtime.layers.fused_moe.token_dispatcher import MoETokenDispatcher
26	-from mindie_llm.runtime.utils.distributed.parallel_info_manager import ParallelType	20	+from mindie_llm.runtime.layers.fused_moe.moe_comm_strategy import MOE_COMM_STRATEGIES, MoECommType
27	-from mindie_llm.runtime.layers.quantization.ms_model_slim.quant_type import QuantType
28	-from mindie_llm.runtime.model_runner.forward_context import get_forward_context
29	-from mindie_llm.runtime.model_runner.forward_context_exp import get_mc2_token_capacity
30	from mindie_llm.utils.log.logging import logger	21	from mindie_llm.utils.log.logging import logger
31		22
32		23
33	-class MoECommType(Enum):	24	+# Communication type to dispatcher class mapping
34	- ALLGATHER = "allgather"	25	+_COMM_TYPE_TO_DISPATCHER: Dict[MoECommType, Type[MoETokenDispatcher]] = {
35	- MC2 = "mc2"
36	- ALLTOALL = "all2all"
37	- FUSED_ALLTOALL = "fused_all2all"
38	-
39	-
40	-_COMM_TYPE_TO_DISPATCHER = {
41	MoECommType.ALLGATHER: TokenDispatcherWithAllGather,	26	MoECommType.ALLGATHER: TokenDispatcherWithAllGather,
42	MoECommType.MC2: TokenDispatcherWithMC2,	27	MoECommType.MC2: TokenDispatcherWithMC2,
43	MoECommType.ALLTOALL: TokenDispatcherWithAll2AllV,	28	MoECommType.ALLTOALL: TokenDispatcherWithAll2AllV,
	@@ -46,80 +31,47 @@ _COMM_TYPE_TO_DISPATCHER = {
46		31
47	def select_moe_comm_method(	32	def select_moe_comm_method(
48	quant_type: Optional[str] = None,	33	quant_type: Optional[str] = None,
49	- is_fused_op_enabled: bool = False) -> Optional[MoECommType]:	34	+ max_num_tokens_per_device: Optional[int] = None
50	- """1. If expert parallel is not enabled, we use all-gather since MC2 and all-to-all	35	+) -> Optional[MoECommType]:
51	- are designed for expert parallelism.	36	+ """Select optimal MoE communication method.
52	- 2. If expert parallel is enabled, we need to consider the soc version and the	37	+
53	- number of tokens. This is based on the observation that all-gather is more
54	- efficient than all-to-all when running on _910B.
55	-
56	- a. For _910B, we choose from MC2 and all-gather.
57	-
58	- b. For _910_93, we choose from MC2 and all-to-all.
59	-
60	- In both cases, we use MC2 when in prefill phase.
61	-
62	Args:	38	Args:
63	- quant_type (Optional[str], optional): The quantization type. Defaults to None.	39	+ quant_type: Quantization type (e.g., W4A8_DYNAMIC).
64	-	40	+ max_num_tokens_per_device: Max tokens per device limit.
65	- Raises:	41	+
66	- ValueError: If the soc version is unsupported.
67	-
68	Returns:	42	Returns:
69	- MoECommType: The selected MoE communication method.	43	+ Selected MoECommType.
		44	+
		45	+ Raises:
		46	+ RuntimeError: No applicable strategy found.
70	"""	47	"""
71	- ascend_device_type = get_npu_node_info().get_device_type()	48	+ # Strategy traversal: return first applicable strategy
72	- parallel_info_manager = get_parallel_info_manager()	49	+ for strategy_cls in MOE_COMM_STRATEGIES:
73	- is_prefill = get_forward_context().is_prefill	50	+ if strategy_cls.is_applicable(
74	- enable_expert_parallel = parallel_info_manager.get(ParallelType.MOE_EP).is_enabled()	51	+ quant_type=quant_type,
75	-	52	+ max_num_tokens_per_device=max_num_tokens_per_device
76	- # NOTE: Due to unsolved error for mc2 in prefill phase, we disable mc2 in prefill phase.	53	+ ):
77	- # NOTE: Temporarily not use parameter num_tokens to select MoECommType.	54	+ selected = strategy_cls.get_comm_type()
78	- if not enable_expert_parallel:	55	+ return selected
79	- moe_comm_type = MoECommType.ALLGATHER	56	+ error_msg = (
80	- elif ascend_device_type in {DeviceType.ASCEND_910B}:	57	+ "MoE strategy selection failed, please check the config."
81	- if not is_prefill and parallel_info_manager.world_size >= 16:	58	+ )
82	- moe_comm_type = MoECommType.MC2	59	+ logger.error(error_msg)
83	- else:	60	+ raise RuntimeError(error_msg)
84	- if quant_type == QuantType.W4A8_DYNAMIC:
85	- moe_comm_type = MoECommType.ALLTOALL
86	- else:
87	- moe_comm_type = MoECommType.ALLGATHER
88	- elif ascend_device_type in {DeviceType.ASCEND_910_93}:
89	- # NOTE: Fused op is not enabled in prefill stage
90	- if get_forward_context().is_prefill:
91	- moe_comm_type = MoECommType.ALLTOALL
92	- elif get_forward_context().batch_descriptor.num_tokens <= get_mc2_token_capacity():
93	- moe_comm_type = MoECommType.MC2
94	- else:
95	- moe_comm_type = MoECommType.ALLTOALL
96	- else:
97	- raise ValueError(f"Unsupported soc_version: {ascend_device_type}")
98	-
99	- # NOTE: MoECommType.ALLGATHER: Currently does not support enabling DP.
100	- if moe_comm_type == MoECommType.ALLGATHER and parallel_info_manager.get(ParallelType.ATTN_DP).is_enabled():
101	- moe_comm_type = MoECommType.ALLTOALL
102	- logger.debug('Currently do not support MoECommType.ALLGATHER with dp.')
103	-
104	- # NOTE: MoECommType.MC2 and MoECommType.ALLTOALL: Do not support moe_tp > 1,
105	- # and require flash_comm to be enabled.
106	- if moe_comm_type in (MoECommType.MC2, MoECommType.ALLTOALL):
107	- if parallel_info_manager.get(ParallelType.MOE_TP).is_enabled():
108	- err_msg = 'MoECommType.MC2 and MoECommType.ALLTOALL: Do not support moe_tp > 1,'
109	- # try to change to MoECommType.ALLGATHER
110	- if parallel_info_manager.get(ParallelType.ATTN_DP).is_enabled():
111	- logger.error(err_msg)
112	- raise RuntimeError(err_msg)
113	- else:
114	- logger.warning(err_msg)
115	- moe_comm_type = MoECommType.ALLGATHER
116	-
117	- return moe_comm_type
118		61
119		62
120	def get_cached_dispatcher(	63	def get_cached_dispatcher(
121	- moe_comm_type: Optional[MoECommType]) -> Optional[MoETokenDispatcher]:	64	+ moe_comm_type: Optional[MoECommType]
		65	+) -> Optional[MoETokenDispatcher]:
		66	+ """Get dispatcher instance for given communication type.
		67	+
		68	+ Args:
		69	+ moe_comm_type: Communication type enum.
		70	+
		71	+ Returns:
		72	+ Dispatcher instance or None if type unsupported.
		73	+ """
122	if moe_comm_type is None or moe_comm_type not in _COMM_TYPE_TO_DISPATCHER:	74	if moe_comm_type is None or moe_comm_type not in _COMM_TYPE_TO_DISPATCHER:
123	return None	75	return None
124	dispatcher_cls = _COMM_TYPE_TO_DISPATCHER[moe_comm_type]	76	dispatcher_cls = _COMM_TYPE_TO_DISPATCHER[moe_comm_type]
125	- return dispatcher_cls()	77	+ return dispatcher_cls()

Amindie_llm/runtime/layers/fused_moe/moe_comm_strategy.py+183-0

@@ -0,0 +1,183 @@
	1	+# Copyright (c) Huawei Technologies Co., Ltd. 2025-2026. All rights reserved.
	2	+# MindIE is licensed under Mulan PSL v2.
	3	+# You can use this software according to the terms and conditions of the Mulan PSL v2.
	4	+# You may obtain a copy of Mulan PSL v2 at:
	5	+# http://license.coscl.org.cn/MulanPSL2
	6	+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
	7	+# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
	8	+# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
	9	+# See the Mulan PSL v2 for more details.
	10	+
	11	+from typing import Optional
	12	+
	13	+from enum import Enum
	14	+from mindie_llm.runtime.utils.npu.device_utils import DeviceType
	15	+from mindie_llm.runtime.layers.quantization.ms_model_slim.quant_type import QuantType
	16	+from mindie_llm.runtime.model_runner.forward_context_exp import get_mc2_token_capacity
	17	+from mindie_llm.utils.log.logging import logger
	18	+
	19	+from mindie_llm.runtime.model_runner.forward_context import get_forward_context
	20	+from mindie_llm.runtime.utils.distributed import get_parallel_info_manager
	21	+from mindie_llm.runtime.utils.distributed.parallel_info_manager import ParallelType
	22	+from mindie_llm.runtime.utils.npu.device_utils import get_npu_node_info
	23	+
	24	+
	25	+class MoECommType(Enum):
	26	+ """Supported MoE communication operator types."""
	27	+ ALLGATHER = "allgather"
	28	+ MC2 = "mc2"
	29	+ ALLTOALL = "all2all"
	30	+ FUSED_ALLTOALL = "fused_all2all"
	31	+ FUSED_MC2 = "fused_mc2"
	32	+
	33	+
	34	+class MoECommStrategyBase:
	35	+ """Base class for MoE communication strategy selection."""
	36	+
	37	+ @staticmethod
	38	+ def is_applicable(quant_type: Optional[str], max_num_tokens_per_device: Optional[int]) -> bool:
	39	+ """Check if strategy applies to given context."""
	40	+ raise NotImplementedError
	41	+
	42	+ @staticmethod
	43	+ def get_comm_type() -> MoECommType:
	44	+ """Return communication type for this strategy."""
	45	+ raise NotImplementedError
	46	+
	47	+
	48	+class FusedMC2Strategy(MoECommStrategyBase):
	49	+ """Fused MC2 strategy: highest priority for Ascend 910C.
	50	+
	51	+ Requirements: EP group size <= 32, no MoE TP, tokens within capacity.
	52	+ """
	53	+
	54	+ @staticmethod
	55	+ def is_applicable(quant_type: Optional[str], max_num_tokens_per_device: Optional[int]) -> bool:
	56	+ device_type = get_npu_node_info().get_device_type()
	57	+ parallel_mgr = get_parallel_info_manager()
	58	+ forward_ctx = get_forward_context()
	59	+
	60	+ # EP must be enabled for any MoE communication strategy
	61	+ if not parallel_mgr.get(ParallelType.MOE_EP).is_enabled():
	62	+ return False
	63	+ # Hardware constraint: only Ascend 910C (910_93) supports fused MC2
	64	+ if device_type != DeviceType.ASCEND_910_93:
	65	+ return False
	66	+
	67	+ # EP group size limit for fused operator efficiency
	68	+ is_ep_valid = parallel_mgr.get(ParallelType.MOE_EP_MC2).group_size <= 32
	69	+ # Token count must fit within per-device capacity
	70	+ tp_size = parallel_mgr.get(ParallelType.ATTN_TP).group_size
	71	+ num_tokens_per_device = (forward_ctx.batch_descriptor.num_tokens + tp_size - 1) // tp_size
	72	+ is_num_tokens_valid = num_tokens_per_device <= max_num_tokens_per_device
	73	+ if not (is_ep_valid and is_num_tokens_valid):
	74	+ return False
	75	+ # Incompatible with MoE TP parallelism
	76	+ if parallel_mgr.get(ParallelType.MOE_TP).is_enabled():
	77	+ return False
	78	+ return True
	79	+
	80	+ @staticmethod
	81	+ def get_comm_type() -> MoECommType:
	82	+ return MoECommType.FUSED_MC2
	83	+
	84	+
	85	+class MC2Strategy(MoECommStrategyBase):
	86	+ """MC2 strategy: for Ascend 910B/910C decode phase.
	87	+
	88	+ 910B: requires world size >= 16. 910C: strict token capacity check.
	89	+ """
	90	+
	91	+ @staticmethod
	92	+ def is_applicable(quant_type: Optional[str], max_num_tokens_per_device: Optional[int]) -> bool:
	93	+ device_type = get_npu_node_info().get_device_type()
	94	+ parallel_mgr = get_parallel_info_manager()
	95	+ forward_ctx = get_forward_context()
	96	+ if not parallel_mgr.get(ParallelType.MOE_EP).is_enabled():
	97	+ return False
	98	+
	99	+ if device_type == DeviceType.ASCEND_910B:
	100	+ # 910B: decode phase only, large cluster required
	101	+ if forward_ctx.is_prefill or parallel_mgr.world_size < 16:
	102	+ return False
	103	+ elif device_type == DeviceType.ASCEND_910_93:
	104	+ # 910C: decode phase only
	105	+ if forward_ctx.is_prefill:
	106	+ return False
	107	+ # Strict capacity check: fail fast if exceeded
	108	+ if forward_ctx.batch_descriptor.num_tokens > get_mc2_token_capacity():
	109	+ error_msg = (
	110	+ f"MC2 operator limitation: tokens per card ({forward_ctx.batch_descriptor.num_tokens})"
	111	+ f" exceed the limit of mc2_token_capacity = {get_mc2_token_capacity()}. "
	112	+ f"Please reduce the batch_size during the decode stage."
	113	+ )
	114	+ logger.error(error_msg)
	115	+ raise ValueError(error_msg)
	116	+ else:
	117	+ logger.error("Unsupported device type: {device_type}")
	118	+ raise RuntimeError("Unsupported device type: {device_type}")
	119	+
	120	+ # Incompatible with MoE TP parallelism
	121	+ if parallel_mgr.get(ParallelType.MOE_TP).is_enabled():
	122	+ return False
	123	+ return True
	124	+
	125	+ @staticmethod
	126	+ def get_comm_type() -> MoECommType:
	127	+ return MoECommType.MC2
	128	+
	129	+
	130	+class All2AllStrategy(MoECommStrategyBase):
	131	+ """All2All strategy: fallback for specific scenarios.
	132	+
	133	+ 910B: W4A8_DYNAMIC quant or Attn DP enabled.
	134	+ 910C: prefill phase when fused MC2 unavailable.
	135	+ """
	136	+
	137	+ @staticmethod
	138	+ def is_applicable(quant_type: Optional[str], max_num_tokens_per_device: Optional[int]) -> bool:
	139	+ device_type = get_npu_node_info().get_device_type()
	140	+ parallel_mgr = get_parallel_info_manager()
	141	+ if not parallel_mgr.get(ParallelType.MOE_EP).is_enabled():
	142	+ return False
	143	+
	144	+ if device_type == DeviceType.ASCEND_910B:
	145	+ # 910B: prefer W4A8_DYNAMIC; allow fallback if Attn DP enabled
	146	+ if quant_type != QuantType.W4A8_DYNAMIC:
	147	+ return parallel_mgr.get(ParallelType.ATTN_DP).is_enabled()
	148	+
	149	+ # Compatibility rule: MoE TP + All2All requires special handling
	150	+ if parallel_mgr.get(ParallelType.MOE_TP).is_enabled():
	151	+ # Attn DP + MoE TP > 1 is not supported for MC2/All2All
	152	+ if parallel_mgr.get(ParallelType.ATTN_DP).is_enabled():
	153	+ err_msg = 'MoECommType.MC2 and MoECommType.ALLTOALL: Do not support moe_tp > 1'
	154	+ logger.error(err_msg)
	155	+ raise RuntimeError(err_msg)
	156	+ return False
	157	+ return True
	158	+
	159	+ @staticmethod
	160	+ def get_comm_type() -> MoECommType:
	161	+ return MoECommType.ALLTOALL
	162	+
	163	+
	164	+class AllGatherStrategy(MoECommStrategyBase):
	165	+ """AllGather strategy: universal fallback with lowest priority."""
	166	+
	167	+ @staticmethod
	168	+ def is_applicable(quant_type: Optional[str], max_num_tokens_per_device: Optional[int]) -> bool:
	169	+ # Fallback: applicable when no other strategy matches
	170	+ return True
	171	+
	172	+ @staticmethod
	173	+ def get_comm_type() -> MoECommType:
	174	+ return MoECommType.ALLGATHER
	175	+
	176	+
	177	+# Strategy priority order: first applicable strategy is selected
	178	+MOE_COMM_STRATEGIES = [
	179	+ FusedMC2Strategy, # P0: Optimal (910C Fused)
	180	+ MC2Strategy, # P1: High perf (large cluster/decode)
	181	+ All2AllStrategy, # P2: Specific quant/prefill
	182	+ AllGatherStrategy, # P3: Universal fallback
	183	+]

Mmindie_llm/runtime/layers/fused_moe/token_dispatcher.py+18-11

	@@ -1,3 +1,4 @@
		1	+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
1	# SPDX-License-Identifier: Apache-2.0	2	# SPDX-License-Identifier: Apache-2.0
2	# Copyright (c) 2024; NVIDIA CORPORATION. All rights reserved.	3	# Copyright (c) 2024; NVIDIA CORPORATION. All rights reserved.
3	# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved.	4	# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved.
	@@ -596,10 +597,11 @@ class TokenDispatcherWithAll2AllV(Singleton, MoETokenDispatcher):
596	global_input_tokens_local_experts_indices):	597	global_input_tokens_local_experts_indices):
597	# Early return if no local experts or no tokens	598	# Early return if no local experts or no tokens
598	if num_local_experts <= 1:	599	if num_local_experts <= 1:
599	- return global_input_tokens, dynamic_scale, None	600	+ return global_input_tokens, None
600		601
601	# Handle quantized case	602	# Handle quantized case
602	if self.with_quant:	603	if self.with_quant:
		604	+ expert_idx_2d = global_input_tokens_local_experts_indices.unsqueeze(-1)
603	active_num = global_input_tokens_local_experts_indices.numel()	605	active_num = global_input_tokens_local_experts_indices.numel()
604		606
605	# Handle case with no active tokens	607	# Handle case with no active tokens
	@@ -607,16 +609,21 @@ class TokenDispatcherWithAll2AllV(Singleton, MoETokenDispatcher):
607	reversed_global_input_permutation_mapping = global_input_tokens_local_experts_indices	609	reversed_global_input_permutation_mapping = global_input_tokens_local_experts_indices
608	return global_input_tokens, dynamic_scale, reversed_global_input_permutation_mapping	610	return global_input_tokens, dynamic_scale, reversed_global_input_permutation_mapping
609		611
610	- dynamic_scale, _ = torch_npu.npu_moe_token_permute(	612	+ # Process with active tokens
611	- dynamic_scale.unsqueeze(-1),	613	+ (global_input_tokens, reversed_global_input_permutation_mapping,
612	- global_input_tokens_local_experts_indices)	614	+ _, expanded_scale) = torch_npu.npu_moe_init_routing_v2(
613	- dynamic_scale = dynamic_scale.squeeze(-1)	615	+ global_input_tokens,
614	-	616	+ expert_idx_2d,
615	- # Handle non-quantized case	617	+ scale=dynamic_scale,
616	- global_input_tokens, reversed_global_input_permutation_mapping = torch_npu.npu_moe_token_permute(	618	+ active_num=active_num,
617	- global_input_tokens,	619	+ expert_capacity=0,
618	- global_input_tokens_local_experts_indices)	620	+ expert_num=num_local_experts,
619	- return global_input_tokens, dynamic_scale, reversed_global_input_permutation_mapping	621	+ expert_tokens_num_type=1,
		622	+ expert_tokens_num_flag=True,
		623	+ active_expert_range=[0, num_local_experts],
		624	+ quant_mode=-1,
		625	+ row_idx_type=0)
		626	+ return global_input_tokens, expanded_scale, reversed_global_input_permutation_mapping
620		627
621	# Handle non-quantized case	628	# Handle non-quantized case
622	global_input_tokens, reversed_global_input_permutation_mapping = torch_npu.npu_moe_token_permute(	629	global_input_tokens, reversed_global_input_permutation_mapping = torch_npu.npu_moe_token_permute(

Mmindie_llm/runtime/layers/quantization/ms_model_slim/w8a8.py+33-60

	@@ -25,7 +25,6 @@ from typing import List
25	import torch	25	import torch
26	import torch_npu	26	import torch_npu
27	from torch import nn	27	from torch import nn
28	-from torch.nn.functional import pad
29		28
30	from mindie_llm.runtime.layers.fused_moe.fused_moe_method_base import FusedMoEMethodBase	29	from mindie_llm.runtime.layers.fused_moe.fused_moe_method_base import FusedMoEMethodBase
31	from mindie_llm.runtime.layers.linear.linear_method_base import LinearMethodBase	30	from mindie_llm.runtime.layers.linear.linear_method_base import LinearMethodBase
	@@ -44,7 +43,7 @@ from mindie_llm.runtime.utils.distributed.utils import even_divide
44	from mindie_llm.utils.log.logging import logger	43	from mindie_llm.utils.log.logging import logger
45		44
46		45
47	-SUPPORT_NZ_NPU_LIST = ("Ascend910B3", "Ascend910B4_1", "Ascend910_9382", "Ascend910_9362")	46	+SUPPORT_NZ_NPU_LIST = ("Ascend910B3", "Ascend910B4-1", "Ascend910_9382", "Ascend910_9362")
48	MXFP8_GROUP_SIZE = 32	47	MXFP8_GROUP_SIZE = 32
49		48
50		49
	@@ -467,13 +466,26 @@ class W8A8PerTokenFusedMoEMethod(FusedMoEMethodBase):
467	else:	466	else:
468	pertoken_scale = dynamic_scale	467	pertoken_scale = dynamic_scale
469		468
470	- hidden_states, swiglu_out_scale, _ = torch_npu.npu_grouped_matmul_swiglu_quant(	469	+ hidden_states = torch_npu.npu_grouped_matmul(
471	- x=x,	470	+ x=[x],
472	- weight=layer.gate_up_weight,	471	+ weight=[layer.gate_up_weight],
473	- bias=None,	472	+ split_item=2, # output a single tensor
474	- group_list=cumsum_group_list(group_list, group_list_type, 0),	473	+ group_list_type=group_list_type,
		474	+ group_type=0, # the axis to group
		475	+ group_list=group_list,
		476	+ output_dtype=torch.int32)[0]
		477	+
		478	+ hidden_states, swiglu_out_scale = torch_npu.npu_dequant_swiglu_quant(
		479	+ x=hidden_states,
475	weight_scale=layer.gate_up_weight_scale,	480	weight_scale=layer.gate_up_weight_scale,
476	- x_scale=pertoken_scale)	481	+ activation_scale=pertoken_scale,
		482	+ bias=None,
		483	+ quant_scale=None,
		484	+ quant_offset=None,
		485	+ group_index=group_list,
		486	+ activate_left=True, # whether to left-activate
		487	+ quant_mode=1, # 0: static quant, 1: dynamic quant
		488	+ )
477		489
478	hidden_states = torch_npu.npu_grouped_matmul(	490	hidden_states = torch_npu.npu_grouped_matmul(
479	x=[hidden_states],	491	x=[hidden_states],
	@@ -492,22 +504,18 @@ class W8A8PerTokenFusedMoEMethod(FusedMoEMethodBase):
492	def process_weights_after_loading(self, layer):	504	def process_weights_after_loading(self, layer):
493	layer.gate_up_weight.data = layer.gate_up_weight.data.transpose(-2, -1).contiguous()	505	layer.gate_up_weight.data = layer.gate_up_weight.data.transpose(-2, -1).contiguous()
494	layer.down_weight.data = layer.down_weight.data.transpose(-2, -1).contiguous()	506	layer.down_weight.data = layer.down_weight.data.transpose(-2, -1).contiguous()
495	- layer.gate_up_weight_scale.data = layer.gate_up_weight_scale.data.squeeze(-1).to(torch.float32)	507	+ layer.gate_up_weight_scale.data = layer.gate_up_weight_scale.data.view(
496	- layer.down_weight_scale.data = layer.down_weight_scale.data.squeeze(-1)	508	+ layer.gate_up_weight_scale.data.shape[0], -1).to(torch.float32)
497	-	509	+ layer.down_weight_scale.data = layer.down_weight_scale.data.view(
498	- if layer.is_fused_op_enabled:	510	+ layer.down_weight_scale.data.shape[0], -1)
499	- layer.gate_up_weight_scale.data = layer.gate_up_weight_scale.data.view(	511	+ layer.register_parameter(
500	- layer.gate_up_weight_scale.data.shape[0], -1)	512	+ "fused_gate_up_weight_scale",
501	- layer.down_weight_scale.data = layer.down_weight_scale.data.view(	513	+ ScalerParameter(scale_from_float_to_int64(layer.gate_up_weight_scale.data))
502	- layer.down_weight_scale.data.shape[0], -1)	514	+ )
503	- layer.register_parameter(	515	+ layer.register_parameter(
504	- "fused_gate_up_weight_scale",	516	+ "fused_down_weight_scale",
505	- ScalerParameter(scale_from_float_to_int64(layer.gate_up_weight_scale.data))	517	+ ScalerParameter(scale_from_float_to_int64(layer.down_weight_scale.data))
506	- )	518	+ )
507	- layer.register_parameter(
508	- "fused_down_weight_scale",
509	- ScalerParameter(scale_from_float_to_int64(layer.down_weight_scale.data))
510	- )
511		519
512	soc_name = get_npu_node_info().soc_name	520	soc_name = get_npu_node_info().soc_name
513	if soc_name in SUPPORT_NZ_NPU_LIST:	521	if soc_name in SUPPORT_NZ_NPU_LIST:
	@@ -520,44 +528,9 @@ class W8A8PerTokenFusedMoEMethod(FusedMoEMethodBase):
520		528
521		529
522	def scale_from_float_to_int64(scale):	530	def scale_from_float_to_int64(scale):
		531	+ """Converts float32 tensor to int64"""
523	import numpy as np	532	import numpy as np
524	scale = torch.from_numpy(	533	scale = torch.from_numpy(
525	np.frombuffer(scale.cpu().to(torch.float32).numpy().tobytes(),	534	np.frombuffer(scale.cpu().to(torch.float32).numpy().tobytes(),
526	dtype=np.int32).astype(np.int64)).to(scale.device)	535	dtype=np.int32).astype(np.int64)).to(scale.device)
527	return scale	536	return scale
528	-
529	-
530	-def cumsum_group_list(group_list: torch.Tensor,
531	- src_list_type: int,
532	- dst_list_type: int,
533	- active_num: int = 0,
534	- expert_num: int = 0) -> torch.Tensor:
535	- if src_list_type not in [0, 1, 2]:
536	- raise ValueError(
537	- f"group_list_type should be in [0, 1, 2], but received {src_list_type}"
538	- )
539	-
540	- if src_list_type == dst_list_type:
541	- return group_list
542	- if src_list_type == 1 and dst_list_type == 0:
543	- return group_list.cumsum(dim=0)
544	- if src_list_type == 0 and dst_list_type == 1:
545	- group_diff = torch.diff(group_list)
546	- new_group = torch.cat([group_list[0].unsqueeze(0), group_diff], dim=0)
547	- return new_group
548	- if src_list_type == 2 and dst_list_type == 0:
549	- experts = pad(group_list[:, 0], (1, 0))
550	- tokens = pad(group_list[:, 1].cumsum(dim=0), (1, 0))
551	- target_list = torch.full(size=(expert_num, ),
552	- fill_value=active_num,
553	- dtype=group_list.dtype,
554	- device=group_list.device)
555	-
556	- for i, (start, end) in enumerate(zip(experts[:-1], experts[1:])):
557	- if end > start:
558	- target_list[start:end] = tokens[i]
559	-
560	- return target_list
561	- raise NotImplementedError(
562	- f"Conversion from src_list_type={src_list_type} to dst_list_type={dst_list_type} is not implemented yet. "
563	- "This feature is under development.")

Mmindie_llm/runtime/model_runner/forward_context_exp.py+3-3

	@@ -25,10 +25,10 @@ _MC2_TOKEN_CAPACITY = None
25		25
26		26
27	def set_mc2_token_capacity(max_num_reqs, uniform_decode_query_len):	27	def set_mc2_token_capacity(max_num_reqs, uniform_decode_query_len):
28	- max_num_tokens = min(max_num_reqs * uniform_decode_query_len, 512)	28	+ max_num_tokens = max_num_reqs * uniform_decode_query_len
29	- tp_size = get_parallel_info_manager().get(ParallelType.MOE_TP).group_size	29	+ tp_size = get_parallel_info_manager().get(ParallelType.ATTN_TP).group_size
30	# Use integer arithmetic for ceiling division.	30	# Use integer arithmetic for ceiling division.
31	- num_tokens_per_tp_rank = (max_num_tokens + tp_size - 1) // tp_size	31	+ num_tokens_per_tp_rank = min((max_num_tokens + tp_size - 1) // tp_size, 512)
32	global _MC2_TOKEN_CAPACITY	32	global _MC2_TOKEN_CAPACITY
33	_MC2_TOKEN_CAPACITY = num_tokens_per_tp_rank * tp_size	33	_MC2_TOKEN_CAPACITY = num_tokens_per_tp_rank * tp_size
34		34

Mmindie_llm/runtime/model_runner/model_runner_exp.py+6-1

	@@ -340,7 +340,10 @@ class ModelRunnerExp:
340	padding_tokens = forward_context.num_actual_tokens	340	padding_tokens = forward_context.num_actual_tokens
341	if forward_context.dp_metadata is not None:	341	if forward_context.dp_metadata is not None:
342	padding_tokens = forward_context.dp_metadata.max_tokens_across_dp_cpu	342	padding_tokens = forward_context.dp_metadata.max_tokens_across_dp_cpu
343	- num_tokens = self.model.get_padded_graph_size(padding_tokens)	343	+ if self._is_aclgraph_enabled:
		344	+ num_tokens = self.model.get_padded_graph_size(padding_tokens)
		345	+ else:
		346	+ num_tokens = forward_context.num_actual_tokens
344	forward_context.batch_descriptor = BatchDescriptor(num_tokens,	347	forward_context.batch_descriptor = BatchDescriptor(num_tokens,
345	forward_context.batch_descriptor.is_flash_comm_enabled)	348	forward_context.batch_descriptor.is_flash_comm_enabled)
346	forward_context.attn_metadata.num_tokens = num_tokens	349	forward_context.attn_metadata.num_tokens = num_tokens
	@@ -410,6 +413,8 @@ class ModelRunnerExp:
410		413
411	def _init_buffer(self) -> None:	414	def _init_buffer(self) -> None:
412	"""Initialize input buffers for graph mode."""	415	"""Initialize input buffers for graph mode."""
		416	+ if not self._is_aclgraph_enabled:
		417	+ return
413	ForwardContext.register(self._max_num_token, self.device, self._mindie_llm_config.hf_config)	418	ForwardContext.register(self._max_num_token, self.device, self._mindie_llm_config.hf_config)
414	input_buffer.register("input_ids", torch.zeros(self._max_num_token, dtype=torch.int32, device=self.device))	419	input_buffer.register("input_ids", torch.zeros(self._max_num_token, dtype=torch.int32, device=self.device))
415	input_buffer.register("position_ids", torch.zeros(self._max_num_token, dtype=torch.int64, device=self.device))	420	input_buffer.register("position_ids", torch.zeros(self._max_num_token, dtype=torch.int64, device=self.device))

Mmindie_llm/runtime/models/deepseek_v3/deepseek_v3.py+1-5

	@@ -23,7 +23,6 @@ from mindie_llm.runtime.layers.fused_moe.experts_selector import select_experts
23	from mindie_llm.runtime.layers.fused_moe.fused_moe import FusedMoE, assign_experts	23	from mindie_llm.runtime.layers.fused_moe.fused_moe import FusedMoE, assign_experts
24	from mindie_llm.runtime.model_runner.forward_context import get_forward_context	24	from mindie_llm.runtime.model_runner.forward_context import get_forward_context
25	from mindie_llm.runtime.utils.distributed import get_parallel_info_manager	25	from mindie_llm.runtime.utils.distributed import get_parallel_info_manager
26	-from mindie_llm.runtime.utils.npu.device_utils import DeviceType, get_npu_node_info
27	from mindie_llm.runtime.layers.attention import get_global_attn_dict	26	from mindie_llm.runtime.layers.attention import get_global_attn_dict
28	from mindie_llm.runtime.layers.parameter import BaseParameter	27	from mindie_llm.runtime.layers.parameter import BaseParameter
29		28
	@@ -131,8 +130,6 @@ class DeepseekV3Moe(nn.Module):
131	self.expert_num = config.n_routed_experts	130	self.expert_num = config.n_routed_experts
132	self.expert_list = assign_experts(config.n_routed_experts,	131	self.expert_list = assign_experts(config.n_routed_experts,
133	parallel_info.moe_ep.group_size)[parallel_info.moe_ep.rank]	132	parallel_info.moe_ep.group_size)[parallel_info.moe_ep.rank]
134	- ascend_device_type = get_npu_node_info().get_device_type()
135	- self.is_fused_op_enabled = True if ascend_device_type in {DeviceType.ASCEND_910_93} else False
136	self.experts = FusedMoE(	133	self.experts = FusedMoE(
137	num_experts=config.n_routed_experts,	134	num_experts=config.n_routed_experts,
138	topk_num=self.topk_num,	135	topk_num=self.topk_num,
	@@ -140,8 +137,7 @@ class DeepseekV3Moe(nn.Module):
140	intermediate_size=config.moe_intermediate_size,	137	intermediate_size=config.moe_intermediate_size,
141	quant_config=quant_config,	138	quant_config=quant_config,
142	prefix=f"{prefix}.experts",	139	prefix=f"{prefix}.experts",
143	- suffix=["gate_proj", "down_proj", "up_proj"],	140	+ suffix=["gate_proj", "down_proj", "up_proj"],)
144	- is_fused_op_enabled=self.is_fused_op_enabled)
145		141
146	self.gate = ReplicatedLinear(	142	self.gate = ReplicatedLinear(
147	config.hidden_size,	143	config.hidden_size,

Mmindie_llm/runtime/utils/distributed/parallel_info_manager.py+3-2

	@@ -38,7 +38,7 @@ class ParallelType(str, Enum):
38	LM_HEAD_TP = "lm_head_tp"	38	LM_HEAD_TP = "lm_head_tp"
39	MOE_TP = "moe_tp"	39	MOE_TP = "moe_tp"
40	MOE_EP = "moe_ep"	40	MOE_EP = "moe_ep"
41	- # DP_ALL_REDUCE = "dp_all_reduce"	41	+ MOE_EP_MC2 = "moe_ep_mc2"
42		42
43		43
44	@dataclass	44	@dataclass
	@@ -147,7 +147,8 @@ class ParallelInfoManager:
147	ParallelType.ATTN_INNER_SP: self.attn_inner_sp,	147	ParallelType.ATTN_INNER_SP: self.attn_inner_sp,
148	ParallelType.MOE_TP: self.moe_tp,	148	ParallelType.MOE_TP: self.moe_tp,
149	ParallelType.MOE_EP: self.moe_ep,	149	ParallelType.MOE_EP: self.moe_ep,
150	- ParallelType.WORLD: self.world	150	+ ParallelType.WORLD: self.world,
		151	+ ParallelType.MOE_EP_MC2: self.moe_ep_mc2,
151	}	152	}
152		153
153	@staticmethod	154	@staticmethod

Mmindie_llm/text_generator/plugins/plugin_manager.py+2-1文件内容审核中，请稍后刷新重试