[WIP]Fix merge pr_494 HCCL added-MindIE-LLM-AtomGit

Mmindie_llm/modeling/model_wrapper/aclgraph/aclgraph_model_wrapper_exp.py+1-1

	@@ -204,7 +204,7 @@ class AclGraphModelWrapperExp(ModelWrapper):
204	model_inputs.last_hidden_states = hidden_states.to(self.device)	204	model_inputs.last_hidden_states = hidden_states.to(self.device)
205	kwargs["hidden_states"] = model_inputs.last_hidden_states	205	kwargs["hidden_states"] = model_inputs.last_hidden_states
206		206
207	- forward_context = self.model_runner.build_forward_context(model_inputs)	207	+ forward_context = self.model_runner.build_forward_context(model_inputs, **kwargs)
208	# NOTE: `PluginManager` will modify `model_inputs.input_lengths`.	208	# NOTE: `PluginManager` will modify `model_inputs.input_lengths`.
209	# Assign `forward_context.seq_lens` to `model_inputs.input_lengths` so that they can share same address.	209	# Assign `forward_context.seq_lens` to `model_inputs.input_lengths` so that they can share same address.
210	model_inputs.input_lengths = forward_context.attn_metadata.seq_lens	210	model_inputs.input_lengths = forward_context.attn_metadata.seq_lens

Mmindie_llm/runtime/layers/attention/backend/sparse_attention.py+15-16

	@@ -199,6 +199,7 @@ class SfaMetadata(AttentionMetadata):
199	# NOTE: This class is used for building attention metadata in the future.	199	# NOTE: This class is used for building attention metadata in the future.
200	actual_seq_lengths_kv: torch.Tensor \| None = None	200	actual_seq_lengths_kv: torch.Tensor \| None = None
201	actual_seq_lengths_query: torch.Tensor \| None = None	201	actual_seq_lengths_query: torch.Tensor \| None = None
		202	+ seq_lens_list: np.ndarray = None
202	cp_input_dict: Dict \| None = None	203	cp_input_dict: Dict \| None = None
203	num_speculative_tokens: int = 0	204	num_speculative_tokens: int = 0
204	num_actual_tokens: int = 0	205	num_actual_tokens: int = 0
	@@ -207,14 +208,15 @@ class SfaMetadata(AttentionMetadata):
207	@staticmethod	208	@staticmethod
208	def from_model_input(model_inputs, cos_table, sin_table, mask, num_speculative_tokens=0):	209	def from_model_input(model_inputs, cos_table, sin_table, mask, num_speculative_tokens=0):
209	if model_inputs.is_prefill:	210	if model_inputs.is_prefill:
210	- actual_seq_lengths_kv = torch.tensor(model_inputs.context_length, dtype=torch.int32)	211	+ actual_seq_lengths_query = torch.tensor(model_inputs.context_length, dtype=torch.int32)
211	else:	212	else:
212	- actual_seq_lengths_kv = (	213	+ actual_seq_lengths_query = (
213	torch.tensor(model_inputs.q_lens, dtype=torch.int32)	214	torch.tensor(model_inputs.q_lens, dtype=torch.int32)
214	if num_speculative_tokens > 0	215	if num_speculative_tokens > 0
215	else torch.tensor([1] * model_inputs.block_tables.shape[0], dtype=torch.int32)	216	else torch.tensor([1] * model_inputs.block_tables.shape[0], dtype=torch.int32)
216	)	217	)
217	- actual_seq_lengths_query = torch.cumsum(actual_seq_lengths_kv, dim=0, dtype=torch.int32)	218	+ actual_seq_lengths_query = torch.cumsum(actual_seq_lengths_query, dim=0, dtype=torch.int32)
		219	+ actual_seq_lengths_kv = torch.tensor(model_inputs.context_length, dtype=torch.int32)
218		220
219	cp_size = get_parallel_info_manager().attn_cp.group_size	221	cp_size = get_parallel_info_manager().attn_cp.group_size
220	cp_input_dict = None	222	cp_input_dict = None
	@@ -230,6 +232,7 @@ class SfaMetadata(AttentionMetadata):
230	main_stream = model_inputs.main_stream	232	main_stream = model_inputs.main_stream
231	return SfaMetadata(	233	return SfaMetadata(
232	seq_lens=model_inputs.context_length,	234	seq_lens=model_inputs.context_length,
		235	+ seq_lens_list=model_inputs.context_length,
233	slot_mapping=model_inputs.slots,	236	slot_mapping=model_inputs.slots,
234	block_tables=model_inputs.block_tables,	237	block_tables=model_inputs.block_tables,
235	attn_mask=mask,	238	attn_mask=mask,
	@@ -260,7 +263,6 @@ class SfaMetadata(AttentionMetadata):
260	self.block_tables.record_stream(main_stream)	263	self.block_tables.record_stream(main_stream)
261	self.slot_mapping = torch.tensor(self.slot_mapping, dtype=torch.int32).to(device)	264	self.slot_mapping = torch.tensor(self.slot_mapping, dtype=torch.int32).to(device)
262	self.slot_mapping.record_stream(main_stream)	265	self.slot_mapping.record_stream(main_stream)
263	- self.seq_lens_list = self.seq_lens
264	self.seq_lens = torch.tensor(self.seq_lens, dtype=torch.int32).to(device)	266	self.seq_lens = torch.tensor(self.seq_lens, dtype=torch.int32).to(device)
265	self.seq_lens.record_stream(main_stream)	267	self.seq_lens.record_stream(main_stream)
266		268
	@@ -280,8 +282,7 @@ class SfaMetadata(AttentionMetadata):
280		282
281	def copy(self, num_actual_tokens, num_tokens):	283	def copy(self, num_actual_tokens, num_tokens):
282	# NOTE: only D2D operation is allowed, should be refactored later	284	# NOTE: only D2D operation is allowed, should be refactored later
283	-	285	+ max_len = self.seq_lens_list.max()
284	- max_len = self.seq_lens_list.max().item()
285	max_seq_pages = (max_len + 128 - 1) // 128	286	max_seq_pages = (max_len + 128 - 1) // 128
286	num_reqs = num_actual_tokens // (self.num_speculative_tokens + 1)	287	num_reqs = num_actual_tokens // (self.num_speculative_tokens + 1)
287		288
	@@ -293,7 +294,7 @@ class SfaMetadata(AttentionMetadata):
293	num_actual_tokens=self.num_speculative_tokens + 1)	294	num_actual_tokens=self.num_speculative_tokens + 1)
294	reqs_padding_length = actual_len - self.actual_seq_lengths_kv.shape[0]	295	reqs_padding_length = actual_len - self.actual_seq_lengths_kv.shape[0]
295		296
296	- seq_lens_pad = torch.tensor([0] * reqs_padding_length, dtype=torch.int32).npu()	297	+ seq_lens_pad = torch.tensor([0] * reqs_padding_length, dtype=torch.int32, device="npu")
297	self.seq_lens = torch.cat([self.seq_lens, seq_lens_pad])	298	self.seq_lens = torch.cat([self.seq_lens, seq_lens_pad])
298	input_buffer_seq_lens = input_buffer.get("seq_lens")	299	input_buffer_seq_lens = input_buffer.get("seq_lens")
299	input_buffer_seq_lens[:actual_len].copy_(self.seq_lens[:actual_len])	300	input_buffer_seq_lens[:actual_len].copy_(self.seq_lens[:actual_len])
	@@ -303,25 +304,24 @@ class SfaMetadata(AttentionMetadata):
303	self.seq_lens_list = self.seq_lens_list.tolist() + [0] * (actual_len - num_actual_tokens)	304	self.seq_lens_list = self.seq_lens_list.tolist() + [0] * (actual_len - num_actual_tokens)
304		305
305	actual_seq_lengths_kv_pad = torch.tensor(	306	actual_seq_lengths_kv_pad = torch.tensor(
306	- [self.num_speculative_tokens + 1] * reqs_padding_length, dtype=torch.int32).npu()	307	+ [self.num_speculative_tokens + 1] * reqs_padding_length, dtype=torch.int32, device="npu")
307		308
308	self.actual_seq_lengths_kv = torch.cat([self.actual_seq_lengths_kv, actual_seq_lengths_kv_pad])	309	self.actual_seq_lengths_kv = torch.cat([self.actual_seq_lengths_kv, actual_seq_lengths_kv_pad])
309	if last_req_tokens > 0:	310	if last_req_tokens > 0:
310	self.actual_seq_lengths_kv[-1] = last_req_tokens	311	self.actual_seq_lengths_kv[-1] = last_req_tokens
311	- self.actual_seq_lengths_query = torch.cumsum(self.actual_seq_lengths_kv, dim=0, dtype=torch.int32).npu()
312		312
313	input_buffer_actual_seq_lengths_kv = input_buffer.get("actual_seq_lengths_kv")	313	input_buffer_actual_seq_lengths_kv = input_buffer.get("actual_seq_lengths_kv")
314	- input_buffer_actual_seq_lengths_kv[:actual_len].copy_(self.actual_seq_lengths_kv[:actual_len])	314	+ input_buffer_actual_seq_lengths_kv[:self.actual_seq_lengths_kv.shape[-1]].copy_(self.actual_seq_lengths_kv[:self.actual_seq_lengths_kv.shape[-1]])
315	self.actual_seq_lengths_kv = input_buffer_actual_seq_lengths_kv[:actual_len]	315	self.actual_seq_lengths_kv = input_buffer_actual_seq_lengths_kv[:actual_len]
316		316
317	input_buffer_actual_seq_lengths_query = input_buffer.get("actual_seq_lengths_query")	317	input_buffer_actual_seq_lengths_query = input_buffer.get("actual_seq_lengths_query")
318	- input_buffer_actual_seq_lengths_query[:actual_len].copy_(self.actual_seq_lengths_query[:actual_len])	318	+
		319	+ input_buffer_actual_seq_lengths_query[:self.actual_seq_lengths_query.shape[-1]].copy_(self.actual_seq_lengths_query[:self.actual_seq_lengths_query.shape[-1]])
319	self.actual_seq_lengths_query = input_buffer_actual_seq_lengths_query[:actual_len]	320	self.actual_seq_lengths_query = input_buffer_actual_seq_lengths_query[:actual_len]
320		321
321	input_buffer_block_tables = input_buffer.get("block_tables")	322	input_buffer_block_tables = input_buffer.get("block_tables")
		323	+ input_buffer_block_tables.fill_(0)
322	input_buffer_block_tables[:num_reqs, :self.block_tables.shape[-1]].copy_(self.block_tables)	324	input_buffer_block_tables[:num_reqs, :self.block_tables.shape[-1]].copy_(self.block_tables)
323	- input_buffer_block_tables[:num_reqs, max_seq_pages:].fill_(0)
324	- input_buffer_block_tables[num_reqs:, :].fill_(0)
325	self.block_tables = input_buffer_block_tables[:actual_len, :]	325	self.block_tables = input_buffer_block_tables[:actual_len, :]
326		326
327	unit_size = len(self.mc2_mask)	327	unit_size = len(self.mc2_mask)
	@@ -679,8 +679,7 @@ class SfaBackendImpl(SelectAttentionImpl):
679	k.view(-1, k.shape[-1]))	679	k.view(-1, k.shape[-1]))
680		680
681	weights = self.indexer.weights_proj(hidden_state)	681	weights = self.indexer.weights_proj(hidden_state)
682	- actual_seq_lengths_key = attn_metadata.actual_seq_lengths_kv \	682	+ actual_seq_lengths_key = attn_metadata.actual_seq_lengths_kv
683	- if forward_context.is_prefill else attn_metadata.seq_lens
684		683
685	if forward_context.is_prefill and self.cp_size > 1:	684	if forward_context.is_prefill and self.cp_size > 1:
686	q = gather_tensor(q, cp_input_dict["cp_load_balance_idx"])	685	q = gather_tensor(q, cp_input_dict["cp_load_balance_idx"])
	@@ -1114,4 +1113,4 @@ class SfaBackendImpl(SelectAttentionImpl):
1114	attn_metadata=attn_metadata	1113	attn_metadata=attn_metadata
1115	)	1114	)
1116	output = self.mla_epilog(o_proj_input)	1115	output = self.mla_epilog(o_proj_input)
1117	- return output	1116	+ return output

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/layers/sampling/sampler.py+2-1

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/forward_metadata/mtp_metadata.py+1-0

	@@ -19,6 +19,7 @@ from mindie_llm.runtime.model_runner.forward_metadata.module_metadata import Mod
19	@dataclass	19	@dataclass
20	class MtpMetadata(ModuleMetadata):	20	class MtpMetadata(ModuleMetadata):
21	last_hidden_states: torch.Tensor = None	21	last_hidden_states: torch.Tensor = None
		22	+ draft_token_indices = None
22		23
23	@staticmethod	24	@staticmethod
24	def from_model_input(model_inputs):	25	def from_model_input(model_inputs):

Mmindie_llm/runtime/model_runner/model_runner_exp.py+7-7

	@@ -306,11 +306,6 @@ class ModelRunnerExp:
306	Logits tensor, or tuple of (logits, hidden_states) if speculative tokens enabled.	306	Logits tensor, or tuple of (logits, hidden_states) if speculative tokens enabled.
307	"""	307	"""
308	set_forward_context(forward_context)	308	set_forward_context(forward_context)
309	- if not self.is_draft_model:
310	- torch.npu.current_stream().synchronize()
311	- preprocess_done = torch.npu.Event()
312	- preprocess_done.record(torch.npu.current_stream())
313	- torch.npu.current_stream().wait_event(preprocess_done)
314		309
315	hidden_states = self.model(input_ids, position_ids)	310	hidden_states = self.model(input_ids, position_ids)
316	if forward_context.dp_metadata is not None:	311	if forward_context.dp_metadata is not None:
	@@ -327,7 +322,7 @@ class ModelRunnerExp:
327	return logits, hidden_states	322	return logits, hidden_states
328	return logits	323	return logits
329		324
330	- def build_forward_context(self, model_inputs: ModelInput) -> 'ForwardContext':	325	+ def build_forward_context(self, model_inputs: ModelInput, **kwargs) -> 'ForwardContext':
331	"""Build forward context from model inputs.	326	"""Build forward context from model inputs.
332		327
333	Do operations like H2D to prepare for the forward function.	328	Do operations like H2D to prepare for the forward function.
	@@ -345,7 +340,10 @@ class ModelRunnerExp:
345	padding_tokens = forward_context.num_actual_tokens	340	padding_tokens = forward_context.num_actual_tokens
346	if forward_context.dp_metadata is not None:	341	if forward_context.dp_metadata is not None:
347	padding_tokens = forward_context.dp_metadata.max_tokens_across_dp_cpu	342	padding_tokens = forward_context.dp_metadata.max_tokens_across_dp_cpu
348	- 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
349	forward_context.batch_descriptor = BatchDescriptor(num_tokens,	347	forward_context.batch_descriptor = BatchDescriptor(num_tokens,
350	forward_context.batch_descriptor.is_flash_comm_enabled)	348	forward_context.batch_descriptor.is_flash_comm_enabled)
351	forward_context.attn_metadata.num_tokens = num_tokens	349	forward_context.attn_metadata.num_tokens = num_tokens
	@@ -415,6 +413,8 @@ class ModelRunnerExp:
415		413
416	def _init_buffer(self) -> None:	414	def _init_buffer(self) -> None:
417	"""Initialize input buffers for graph mode."""	415	"""Initialize input buffers for graph mode."""
		416	+ if not self._is_aclgraph_enabled:
		417	+ return
418	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)
419	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))
420	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/model_runner/spec_worker.py+24-21

	@@ -337,6 +337,8 @@ class MtpWorkerExp(BaseWorkerProxy):
337	self.mapping = get_parallel_info_manager()	337	self.mapping = get_parallel_info_manager()
338	self.distributed_enable = kwargs.get("distributed_enable", False)	338	self.distributed_enable = kwargs.get("distributed_enable", False)
339	self.is_dp_and_server_centralized = self.mapping.has_dp() and not self.distributed_enable	339	self.is_dp_and_server_centralized = self.mapping.has_dp() and not self.distributed_enable
		340	+
		341	+ self._offsets = torch.arange(self.num_speculative_tokens + 1, device="npu")
340		342
341	def __getattr__(self, name):	343	def __getattr__(self, name):
342	return getattr(self.main_model_runner, name)	344	return getattr(self.main_model_runner, name)
	@@ -443,15 +445,16 @@ class MtpWorkerExp(BaseWorkerProxy):
443	return slot_list	445	return slot_list
444	slot_num_per_batch = self.num_speculative_tokens * 2	446	slot_num_per_batch = self.num_speculative_tokens * 2
445	used_slot_num_per_iter = self.num_speculative_tokens + 1	447	used_slot_num_per_iter = self.num_speculative_tokens + 1
446	- offsets = torch.arange(used_slot_num_per_iter).npu()
447	shift_num = (slot_input.size(0) - used_slot_num_per_iter) // slot_num_per_batch + 1	448	shift_num = (slot_input.size(0) - used_slot_num_per_iter) // slot_num_per_batch + 1
448	- shift_values = torch.arange(shift_num).npu()	449	+ shift_values = torch.arange(shift_num, device="npu")
449	for mtp_idx in range(self.num_speculative_tokens):	450	for mtp_idx in range(self.num_speculative_tokens):
450	if slot_input.size(0) > 1:	451	if slot_input.size(0) > 1:
451	starts = mtp_idx + shift_values * slot_num_per_batch	452	starts = mtp_idx + shift_values * slot_num_per_batch
452	- valid_starts = starts[starts + used_slot_num_per_iter <= slot_input.size(0)]	453	+ indices = starts.view(-1, 1) + self._offsets
453	- indices = valid_starts.view(-1, 1) + offsets	454	+ mask = (indices >= 0) & (indices < slot_input.size(0))
454	- slot_new = slot_input[indices].flatten()	455	+ safe_indices = indices.clamp(0, slot_input.size(0) - 1)
		456	+ slot_all = slot_input[safe_indices]
		457	+ slot_new = (slot_all * mask.to(slot_all.dtype)).flatten()
455	slot_list.append(slot_new)	458	slot_list.append(slot_new)
456	else:	459	else:
457	slot_list.append(slot_input)	460	slot_list.append(slot_input)
	@@ -479,9 +482,7 @@ class MtpWorkerExp(BaseWorkerProxy):
479	forward_context.attn_metadata.block_tables = sub_model_inputs.block_tables	482	forward_context.attn_metadata.block_tables = sub_model_inputs.block_tables
480	forward_context.mtp_metadata.last_hidden_states = hidden_states_mtp_input	483	forward_context.mtp_metadata.last_hidden_states = hidden_states_mtp_input
481	forward_context.lm_head_indices = lm_head_indices	484	forward_context.lm_head_indices = lm_head_indices
482	- forward_context.attn_metadata.actual_seq_lengths_kv = q_lens	485	+ padding_tokens = num_actual_tokens = input_ids_mtp.shape[0]
483	- forward_context.attn_metadata.actual_seq_lengths_query = \
484	- torch.cumsum(forward_context.attn_metadata.actual_seq_lengths_kv, dim=0, dtype=torch.int32)
485	# draft forward	486	# draft forward
486	for mtp_idx in range(self.num_speculative_tokens):	487	for mtp_idx in range(self.num_speculative_tokens):
487	slot_mapping = slot_list[mtp_idx]	488	slot_mapping = slot_list[mtp_idx]
	@@ -493,7 +494,6 @@ class MtpWorkerExp(BaseWorkerProxy):
493	npu_cache, input_ids_mtp, positions_mtp, forward_context, **draft_kwargs)	494	npu_cache, input_ids_mtp, positions_mtp, forward_context, **draft_kwargs)
494	else:	495	else:
495	# If an input is not changed, we can skip padding.	496	# If an input is not changed, we can skip padding.
496	- padding_tokens = num_actual_tokens = input_ids_mtp.shape[0]
497	if forward_context.dp_metadata is not None:	497	if forward_context.dp_metadata is not None:
498	padding_tokens = forward_context.dp_metadata.num_tokens_across_dp_cpu.max().item()	498	padding_tokens = forward_context.dp_metadata.num_tokens_across_dp_cpu.max().item()
499	num_tokens = self.draft_model_runner.model.get_padded_graph_size(padding_tokens)	499	num_tokens = self.draft_model_runner.model.get_padded_graph_size(padding_tokens)
	@@ -555,7 +555,7 @@ class MtpWorkerExp(BaseWorkerProxy):
555	def update_model_input_ids_by_draft(self,	555	def update_model_input_ids_by_draft(self,
556	input_ids,	556	input_ids,
557	input_ids_draft_out,	557	input_ids_draft_out,
558	- kwargs) -> tuple[torch.Tensor, torch.Tensor]:	558	+ forward_context) -> tuple[torch.Tensor, torch.Tensor]:
559	"""	559	"""
560	Update the input token id of main model by draft model output.	560	Update the input token id of main model by draft model output.
561		561
	@@ -571,20 +571,17 @@ class MtpWorkerExp(BaseWorkerProxy):
571	# 'draft_token_indices' indicates that the index of token for current dp group.	571	# 'draft_token_indices' indicates that the index of token for current dp group.
572	# This flag will be removed after support for deploying DP-in-and-DP-out is implemented.	572	# This flag will be removed after support for deploying DP-in-and-DP-out is implemented.
573	if self.is_dp_and_server_centralized:	573	if self.is_dp_and_server_centralized:
574	- dp_rank_ids = kwargs.get('dp_rank_ids')	574	+ if forward_context.mtp_metadata.draft_token_indices.shape[0] == 0:
575	- draft_token_indices = torch.where(dp_rank_ids == self.mapping.attn_dp.rank)[0]
576	- if draft_token_indices.numel() == 0:
577	- # It indicates the current device is running Dummy Data.
578	return input_ids, input_ids_draft_out	575	return input_ids, input_ids_draft_out
579	- else:	576	+
580	- draft_token_indices = None
581	input_ids_reshaped = input_ids.view(-1, self.num_speculative_tokens + 1)	577	input_ids_reshaped = input_ids.view(-1, self.num_speculative_tokens + 1)
582	input_ids_draft_out_reshaped = input_ids_draft_out.view(-1, self.num_speculative_tokens)	578	input_ids_draft_out_reshaped = input_ids_draft_out.view(-1, self.num_speculative_tokens)
583	- if draft_token_indices is not None:	579	+ if forward_context.mtp_metadata.draft_token_indices is not None:
584	- input_ids_draft_out_reshaped = input_ids_draft_out_reshaped[draft_token_indices]	580	+ input_ids_draft_out_reshaped = \
		581	+ input_ids_draft_out_reshaped[forward_context.mtp_metadata.draft_token_indices]
585	input_ids_reshaped[:, 1:] = input_ids_draft_out_reshaped	582	input_ids_reshaped[:, 1:] = input_ids_draft_out_reshaped
586	input_ids = input_ids_reshaped.flatten()	583	input_ids = input_ids_reshaped.flatten()
587	- return input_ids, input_ids_draft_out	584	+ return input_ids, input_ids_draft_out
588		585
589	def forward_mtp_decode(	586	def forward_mtp_decode(
590	self,	587	self,
	@@ -613,7 +610,7 @@ class MtpWorkerExp(BaseWorkerProxy):
613		610
614	# update args and main model forward	611	# update args and main model forward
615	input_ids, input_ids_draft_out = \	612	input_ids, input_ids_draft_out = \
616	- self.update_model_input_ids_by_draft(input_ids, input_ids_draft_out, kwargs)	613	+ self.update_model_input_ids_by_draft(input_ids, input_ids_draft_out, forward_context)
617		614
618	logits, hidden_states = self.main_model_runner.forward(	615	logits, hidden_states = self.main_model_runner.forward(
619	npu_cache, input_ids, position_ids, forward_context, **kwargs)	616	npu_cache, input_ids, position_ids, forward_context, **kwargs)
	@@ -635,7 +632,7 @@ class MtpWorkerExp(BaseWorkerProxy):
635	else:	632	else:
636	return self.forward_mtp_decode(npu_cache, input_ids, position_ids, forward_context, **kwargs)	633	return self.forward_mtp_decode(npu_cache, input_ids, position_ids, forward_context, **kwargs)
637		634
638	- def build_forward_context(self, model_inputs):	635	+ def build_forward_context(self, model_inputs, **kwargs):
639	'''	636	'''
640	Do operations like H2D to prepare for the forward function.	637	Do operations like H2D to prepare for the forward function.
641	It will be used by eager mode and aclgraph mode	638	It will be used by eager mode and aclgraph mode
	@@ -643,6 +640,12 @@ class MtpWorkerExp(BaseWorkerProxy):
643	forward_context = self.main_model_runner.build_forward_context(model_inputs)	640	forward_context = self.main_model_runner.build_forward_context(model_inputs)
644	sub_forward_context = self.draft_model_runner.build_forward_context(model_inputs)	641	sub_forward_context = self.draft_model_runner.build_forward_context(model_inputs)
645	forward_context.sub_forward_context = sub_forward_context	642	forward_context.sub_forward_context = sub_forward_context
		643	+ if self.is_dp_and_server_centralized:
		644	+ dp_rank_ids = kwargs.get('dp_rank_ids')
		645	+ forward_context.mtp_metadata.draft_token_indices = \
		646	+ torch.where(dp_rank_ids == self.mapping.attn_dp.rank)[0]
		647	+ else:
		648	+ forward_context.mtp_metadata.draft_token_indices = None
646	return forward_context	649	return forward_context
647		650
648		651

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+10-5

	@@ -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
	@@ -175,7 +176,7 @@ class ParallelInfoManager:
175	return None	176	return None
176		177
177	@staticmethod	178	@staticmethod
178	- def _create_npu_process_group(parallel_info: ParallelInfo) -> ProcessGroup:	179	+ def _create_npu_process_group(parallel_info: ParallelInfo, is_all_reduce: bool = False) -> ProcessGroup:
179	"""Creates an HCCL process group with custom buffer size.	180	"""Creates an HCCL process group with custom buffer size.
180	Need to traverse all groups to establish communication domains,	181	Need to traverse all groups to establish communication domains,
181	even if your rank is not present within them, otherwise it will	182	even if your rank is not present within them, otherwise it will
	@@ -184,7 +185,11 @@ class ParallelInfoManager:
184	cur_process_group = None	185	cur_process_group = None
185	for group_idx, group_ranks in enumerate(parallel_info.rank_per_group):	186	for group_idx, group_ranks in enumerate(parallel_info.rank_per_group):
186	options = dist_c.ProcessGroupHCCL.Options()	187	options = dist_c.ProcessGroupHCCL.Options()
187	- options.hccl_config = {"hccl_buffer_size": parallel_info.buffer_size}	188	+ if is_all_reduce:
		189	+ options.hccl_config = {"hccl_buffer_size": parallel_info.buffer_size,
		190	+ "hccl_op_expansion_mode": 1}
		191	+ else:
		192	+ options.hccl_config = {"hccl_buffer_size": parallel_info.buffer_size}
188	process_group = dist.new_group(ranks=group_ranks, pg_options=options)	193	process_group = dist.new_group(ranks=group_ranks, pg_options=options)
189	if group_idx == parallel_info.current_group_id:	194	if group_idx == parallel_info.current_group_id:
190	cur_process_group = process_group	195	cur_process_group = process_group
	@@ -280,5 +285,5 @@ class ParallelInfoManager:
280	parallel_info.rank = parallel_info.rank_per_group[parallel_info.current_group_id].index(self.rank)	285	parallel_info.rank = parallel_info.rank_per_group[parallel_info.current_group_id].index(self.rank)
281		286
282	parallel_info.process_group = self._create_npu_process_group(parallel_info)	287	parallel_info.process_group = self._create_npu_process_group(parallel_info)
283	- parallel_info.preprocess_group = self._create_npu_process_group(parallel_info)	288	+ parallel_info.preprocess_group = self._create_npu_process_group(parallel_info, is_all_reduce=True)
284	return parallel_info	289	return parallel_info

Mmindie_llm/text_generator/plugins/mtp/mtp_plugin.py+207-89

	@@ -14,6 +14,7 @@ from .decoding_policy import DecodingPolicy
14	from ..plugin import Plugin	14	from ..plugin import Plugin
15	from ....utils.log.logging import logger	15	from ....utils.log.logging import logger
16	from ....utils.tensor import backend	16	from ....utils.tensor import backend
		17	+from ....utils.env import ENV
17		18
18		19
19	class MtpPlugin(Plugin):	20	class MtpPlugin(Plugin):
	@@ -38,6 +39,9 @@ class MtpPlugin(Plugin):
38	self.rank = generator_backend.rank	39	self.rank = generator_backend.rank
39	self.max_seq_len = kwargs.get('max_seq_len')	40	self.max_seq_len = kwargs.get('max_seq_len')
40		41
		42	+ self.token_range = torch.arange(self.num_speculative_tokens + 1, device="npu")
		43	+ self.valid_gather_indices = torch.zeros((200, 1), device="npu")
		44	+
41	def model_inputs_update(self, model_inputs, input_metadata, sampling_metadata, cache_ids, input_len_mask, **kwargs):	45	def model_inputs_update(self, model_inputs, input_metadata, sampling_metadata, cache_ids, input_len_mask, **kwargs):
42	hit_mask = kwargs.get('hit_mask')	46	hit_mask = kwargs.get('hit_mask')
43	model_inputs, q_len, attention_mask = (	47	model_inputs, q_len, attention_mask = (
	@@ -54,6 +58,56 @@ class MtpPlugin(Plugin):
54	current_dp_sequence_ids,	58	current_dp_sequence_ids,
55	current_all_sequence_ids,	59	current_all_sequence_ids,
56	last_all_sequence_ids	60	last_all_sequence_ids
		61	+ ):
		62	+ if ENV.model_runner_exp:
		63	+ return self.prepare_masks_for_filling_exp(
		64	+ model_inputs,
		65	+ current_dp_sequence_ids,
		66	+ current_all_sequence_ids,
		67	+ last_all_sequence_ids)
		68	+ masks = {}
		69	+ if last_all_sequence_ids is not None:
		70	+ speculative_length = self.num_speculative_tokens + 1
		71	+ hit_mask = np.isin(current_dp_sequence_ids, last_all_sequence_ids)
		72	+ if hit_mask.any():
		73	+ hit_sequence_ids = current_dp_sequence_ids[hit_mask]
		74	+ hit_indices = np.where(hit_sequence_ids[:, None] == last_all_sequence_ids[None, :])[1]
		75	+ hit_mask_per_token = np.repeat(hit_mask, speculative_length)
		76	+ hit_size = len(hit_indices)
		77	+ hit_increments = np.repeat(
		78	+ np.arange(speculative_length).reshape(-1, 1), hit_size, axis=1).transpose().reshape(-1)
		79	+ hit_indices_per_token = np.repeat(hit_indices * speculative_length, speculative_length) + hit_increments
		80	+ masks['hit_mask'] = hit_mask
		81	+ masks['hit_mask_tensor'] = self.generator_backend.to_tensor(hit_mask)
		82	+ all_hit_mask = np.isin(current_all_sequence_ids, last_all_sequence_ids)
		83	+ all_hit_sequence_ids = current_all_sequence_ids[all_hit_mask]
		84	+ masks['all_hit_mask_tensor'] = self.generator_backend.to_tensor(all_hit_mask)
		85	+ all_hit_indices = np.where(all_hit_sequence_ids[:, None] == last_all_sequence_ids[None, :])[1]
		86	+ masks['all_hit_indices'] = all_hit_indices
		87	+ masks['hit_indices'] = hit_indices
		88	+ masks['hit_indices_per_token_tensor'] = self.generator_backend.to_tensor(hit_indices_per_token)
		89	+ masks['hit_mask_per_token_tensor'] = self.generator_backend.to_tensor(hit_mask_per_token)
		90	+ masks['hit_speculative_length'] = np.full(len(all_hit_indices), speculative_length)
		91	+ masks['hit_dp_speculative_length'] = np.full(len(hit_indices), speculative_length)
		92	+ hit_arange = np.arange(hit_size)
		93	+ masks['hit_arange_tensor'] = self.generator_backend.to_tensor(hit_arange)
		94	+ hit_mask_mod = self.generator_backend.to_tensor(
		95	+ np.arange(len(hit_mask_per_token)) % speculative_length != 0)
		96	+ masks['hit_mask_mod'] = hit_mask_mod
		97	+ hit_block_tables = model_inputs.block_tables_array[hit_mask]
		98	+ candidate_slots = self.infer_context.block_table_to_slots(hit_block_tables).reshape(hit_size, -1)
		99	+ hit_block_indices = np.repeat(hit_arange, speculative_length)
		100	+ masks['candidate_slots'] = self.generator_backend.to_tensor(candidate_slots)
		101	+ masks['hit_block_indices'] = self.generator_backend.to_tensor(hit_block_indices)
		102	+ masks['hit_increments'] = self.generator_backend.to_tensor(hit_increments)
		103	+ return masks
		104	+
		105	+ def prepare_masks_for_filling_exp(
		106	+ self,
		107	+ model_inputs,
		108	+ current_dp_sequence_ids,
		109	+ current_all_sequence_ids,
		110	+ last_all_sequence_ids
57	):	111	):
58	masks = {}	112	masks = {}
59	if last_all_sequence_ids is not None:	113	if last_all_sequence_ids is not None:
	@@ -70,15 +124,21 @@ class MtpPlugin(Plugin):
70	hit_indices_per_token = np.repeat(hit_indices * speculative_length, speculative_length) + hit_increments	124	hit_indices_per_token = np.repeat(hit_indices * speculative_length, speculative_length) + hit_increments
71	masks['hit_mask'] = hit_mask	125	masks['hit_mask'] = hit_mask
72	masks['hit_mask_tensor'] = self.generator_backend.to_tensor(hit_mask)	126	masks['hit_mask_tensor'] = self.generator_backend.to_tensor(hit_mask)
		127	+ masks['hit_mask_local_tensor'] = torch.masked_select(torch.arange(hit_mask.shape[0], device="npu"), masks['hit_mask_tensor'])
73	masks['hit_indices'] = hit_indices	128	masks['hit_indices'] = hit_indices
		129	+ masks['hit_indices_tensor'] = self.generator_backend.to_tensor(hit_indices)
74	masks['hit_indices_per_token_tensor'] = self.generator_backend.to_tensor(hit_indices_per_token)	130	masks['hit_indices_per_token_tensor'] = self.generator_backend.to_tensor(hit_indices_per_token)
75	masks['hit_mask_per_token_tensor'] = self.generator_backend.to_tensor(hit_mask_per_token)	131	masks['hit_mask_per_token_tensor'] = self.generator_backend.to_tensor(hit_mask_per_token)
		132	+ masks['hit_local_indices_per_token_tensor'] = torch.masked_select(torch.arange(hit_mask_per_token.shape[0], device="npu"), masks['hit_mask_per_token_tensor'])
76	masks['hit_dp_speculative_length'] = torch.tensor(np.full(len(hit_indices), speculative_length), device="npu")	133	masks['hit_dp_speculative_length'] = torch.tensor(np.full(len(hit_indices), speculative_length), device="npu")
77	hit_arange = np.arange(hit_size)	134	hit_arange = np.arange(hit_size)
78	masks['hit_arange_tensor'] = self.generator_backend.to_tensor(hit_arange)	135	masks['hit_arange_tensor'] = self.generator_backend.to_tensor(hit_arange)
79	hit_mask_mod = self.generator_backend.to_tensor(	136	hit_mask_mod = self.generator_backend.to_tensor(
80	np.arange(len(hit_mask_per_token)) % speculative_length != 0)	137	np.arange(len(hit_mask_per_token)) % speculative_length != 0)
81	masks['hit_mask_mod'] = hit_mask_mod	138	masks['hit_mask_mod'] = hit_mask_mod
		139	+ hit_mask_per_token[hit_mask_mod.cpu().numpy()] = False
		140	+ hit_mask_per_token_tensor = self.generator_backend.to_tensor(hit_mask_per_token)
		141	+ masks['hit_mask_per_token_mod_tensor'] = torch.masked_select(torch.arange(hit_mask_per_token_tensor.shape[0], device="npu"), hit_mask_per_token_tensor)
82	hit_block_tables = model_inputs.block_tables_array[hit_mask]	142	hit_block_tables = model_inputs.block_tables_array[hit_mask]
83	candidate_slots = self.infer_context.block_table_to_slots(hit_block_tables).reshape(hit_size, -1)	143	candidate_slots = self.infer_context.block_table_to_slots(hit_block_tables).reshape(hit_size, -1)
84	hit_block_indices = np.repeat(hit_arange, speculative_length)	144	hit_block_indices = np.repeat(hit_arange, speculative_length)
	@@ -86,11 +146,11 @@ class MtpPlugin(Plugin):
86	masks['hit_block_indices'] = self.generator_backend.to_tensor(hit_block_indices)	146	masks['hit_block_indices'] = self.generator_backend.to_tensor(hit_block_indices)
87	masks['hit_increments'] = self.generator_backend.to_tensor(hit_increments)	147	masks['hit_increments'] = self.generator_backend.to_tensor(hit_increments)
88	if all_hit_mask.any():	148	if all_hit_mask.any():
89	- masks['all_hit_mask'] = all_hit_mask	149	+ masks["all_hit_mask"] = all_hit_mask
90	masks['all_hit_mask_tensor'] = self.generator_backend.to_tensor(all_hit_mask)	150	masks['all_hit_mask_tensor'] = self.generator_backend.to_tensor(all_hit_mask)
91	all_hit_sequence_ids = current_all_sequence_ids[all_hit_mask]	151	all_hit_sequence_ids = current_all_sequence_ids[all_hit_mask]
92	all_hit_indices = np.where(all_hit_sequence_ids[:, None] == last_all_sequence_ids[None, :])[1]	152	all_hit_indices = np.where(all_hit_sequence_ids[:, None] == last_all_sequence_ids[None, :])[1]
93	- masks['all_hit_indices'] = torch.tensor(all_hit_indices)	153	+ masks['all_hit_indices'] = torch.tensor(all_hit_indices, device="npu")
94	masks['hit_speculative_length'] = torch.tensor(np.full(len(all_hit_indices), speculative_length), device="npu")	154	masks['hit_speculative_length'] = torch.tensor(np.full(len(all_hit_indices), speculative_length), device="npu")
95	return masks	155	return masks
96		156
	@@ -158,7 +218,12 @@ class MtpPlugin(Plugin):
158	block_offsets = backend.repeat_interleave(offset_start_indices, speculative_length) + hit_increments	218	block_offsets = backend.repeat_interleave(offset_start_indices, speculative_length) + hit_increments
159	candidate_slots = filling_masks.get('candidate_slots')	219	candidate_slots = filling_masks.get('candidate_slots')
160	hit_block_indices = filling_masks.get('hit_block_indices')	220	hit_block_indices = filling_masks.get('hit_block_indices')
161	- model_inputs.slots[hit_mask_per_token] = candidate_slots[hit_block_indices, block_offsets]	221	+ hit_local_indices_per_token_tensor = filling_masks.get('hit_local_indices_per_token_tensor')
		222	+ model_inputs.forward_context.attn_metadata.slot_mapping.scatter_(
		223	+ 0,
		224	+ hit_local_indices_per_token_tensor,
		225	+ candidate_slots[hit_block_indices, block_offsets]
		226	+ )
162	else:	227	else:
163	model_inputs.cached_context_length[hit_mask] += hit_num_tokens_tensor	228	model_inputs.cached_context_length[hit_mask] += hit_num_tokens_tensor
164	indices = np.where(hit_mask)[0]	229	indices = np.where(hit_mask)[0]
	@@ -180,6 +245,11 @@ class MtpPlugin(Plugin):
180	* self.infer_context.spcp_parallel_info.sp_size	245	* self.infer_context.spcp_parallel_info.sp_size
181	][self.infer_context.spcp_parallel_info.sp_rank]	246	][self.infer_context.spcp_parallel_info.sp_rank]
182		247
		248	+ actual_seq_lengths_query = torch.ones_like(model_inputs.input_lengths, dtype=torch.int32)
		249	+ actual_seq_lengths_query = torch.cumsum(actual_seq_lengths_query, dim=0, dtype=torch.int32)
		250	+ model_inputs.forward_context.attn_metadata.actual_seq_lengths_query = actual_seq_lengths_query * speculative_length
		251	+ model_inputs.forward_context.attn_metadata.actual_seq_lengths_kv = model_inputs.input_lengths
		252	+ model_inputs.forward_context.attn_metadata.seq_lens = model_inputs.input_lengths
183	hit_mask_per_token[hit_mask_mod] = False	253	hit_mask_per_token[hit_mask_mod] = False
184	model_inputs.input_ids[hit_mask_per_token] = hit_token_ids[hit_arange_tensor, hit_num_tokens_tensor - 1]	254	model_inputs.input_ids[hit_mask_per_token] = hit_token_ids[hit_arange_tensor, hit_num_tokens_tensor - 1]
185		255
	@@ -200,18 +270,33 @@ class MtpPlugin(Plugin):
200	all_hit_mask = filling_masks.get('all_hit_mask')	270	all_hit_mask = filling_masks.get('all_hit_mask')
201	sampling_output = model_output_wrapper.sampling_output	271	sampling_output = model_output_wrapper.sampling_output
202	if all_hit_mask is not None:	272	if all_hit_mask is not None:
		273	+ # Get the device tensor of subtrahend of lmhead indices.
		274	+ # The default prefill_head_indices is calculated assuming the number of tokens is speculative_length.
		275	+ # The true head indices equals to the default prefill_head_indices minus head_indices_subtrahend.
203	all_hit_indices = filling_masks.get('all_hit_indices')	276	all_hit_indices = filling_masks.get('all_hit_indices')
204	hit_speculative_length = filling_masks.get('hit_speculative_length')	277	hit_speculative_length = filling_masks.get('hit_speculative_length')
205	- all_hit_num_tokens = sampling_output.num_new_tokens[all_hit_indices]	278	+ all_hit_num_tokens = torch.index_select(
206	- all_hit_mask_tensor = filling_masks.get('all_hit_mask_tensor')	279	+ sampling_output.num_new_tokens,
		280	+ dim=0,
		281	+ index=all_hit_indices
		282	+ )
207	head_indices_subtrahend = hit_speculative_length - all_hit_num_tokens	283	head_indices_subtrahend = hit_speculative_length - all_hit_num_tokens
208	- mtp_model_inputs.prefill_head_indices[all_hit_mask_tensor] -= head_indices_subtrahend	284	+ mtp_model_inputs.prefill_head_indices.scatter_add_(
		285	+ dim=0,
		286	+ index=all_hit_indices,
		287	+ src=-head_indices_subtrahend.to(torch.int32)
		288	+ )
209	if hit_mask is not None:	289	if hit_mask is not None:
210	model_output_hidden_states = model_output_wrapper.model_output.hidden_states	290	model_output_hidden_states = model_output_wrapper.model_output.hidden_states
211		291
212	# Move the hit_token_ids of mtp model to device.	292	# Move the hit_token_ids of mtp model to device.
213	hit_indices = filling_masks.get('hit_indices')	293	hit_indices = filling_masks.get('hit_indices')
214	- hit_token_ids = sampling_output.token_ids[hit_indices]	294	+ hit_indices_tensor = filling_masks.get('hit_indices_tensor')
		295	+ hit_token_ids = torch.index_select(
		296	+ sampling_output.token_ids,
		297	+ dim=0,
		298	+ index=hit_indices_tensor
		299	+ )
215	hit_token_ids_cols = hit_token_ids.shape[1]	300	hit_token_ids_cols = hit_token_ids.shape[1]
216	if hit_token_ids_cols < speculative_length:	301	if hit_token_ids_cols < speculative_length:
217	hit_token_ids = backend.nn.functional.pad(	302	hit_token_ids = backend.nn.functional.pad(
	@@ -222,38 +307,67 @@ class MtpPlugin(Plugin):
222	hit_token_ids = hit_token_ids[:, :speculative_length]	307	hit_token_ids = hit_token_ids[:, :speculative_length]
223		308
224	# Move the hit_num_tokens and hit_num_tokens_per_token to device.	309	# Move the hit_num_tokens and hit_num_tokens_per_token to device.
225	- hit_num_tokens_tensor = sampling_output.num_new_tokens[hit_indices]	310	+ hit_num_tokens_tensor = torch.index_select(
		311	+ sampling_output.num_new_tokens,
		312	+ dim=0,
		313	+ index=hit_indices_tensor
		314	+ )
226	hit_num_tokens_per_token = backend.repeat_interleave(hit_num_tokens_tensor, speculative_length)	315	hit_num_tokens_per_token = backend.repeat_interleave(hit_num_tokens_tensor, speculative_length)
		316	+ hit_num_tokens_numpy = sampling_output.num_new_tokens_numpy[hit_indices]
227		317
228	# Get masks.	318	# Get masks.
229	- hit_mask_tensor = filling_masks.get('hit_mask_tensor')
230	- hit_mask_per_token = filling_masks.get('hit_mask_per_token_tensor')
231	hit_indices_per_token = filling_masks.get('hit_indices_per_token_tensor')	319	hit_indices_per_token = filling_masks.get('hit_indices_per_token_tensor')
		320	+ hit_mask_per_token = filling_masks.get('hit_mask_per_token_tensor')
232	hit_arange_tensor = filling_masks.get('hit_arange_tensor')	321	hit_arange_tensor = filling_masks.get('hit_arange_tensor')
233	- hit_mask_mod = filling_masks.get('hit_mask_mod')	322	+ hit_mask_local_tensor = filling_masks.get('hit_mask_local_tensor')
234	-	323	+ hit_mask_tensor = filling_masks.get('hit_mask_tensor')
		324	+ hit_local_indices_per_token_tensor = filling_masks.get('hit_local_indices_per_token_tensor')
235	# assignation	325	# assignation
236	- hit_token_ids_tensor = hit_token_ids.reshape(-1)	326	+ hit_token_ids_flatten = hit_token_ids.flatten()
237	- mtp_model_inputs.input_ids[hit_mask_per_token] = hit_token_ids_tensor	327	+ mtp_model_inputs.input_ids.scatter_(
238	- hit_hidden_states = model_output_hidden_states[hit_indices_per_token]	328	+ dim=0,
239	- hidden_states[hit_mask_per_token] = hit_hidden_states	329	+ index=hit_local_indices_per_token_tensor,
		330	+ src=hit_token_ids_flatten
		331	+ )
		332	+
		333	+ if model_output_wrapper.input_metadata.is_prefill:
		334	+ model_output_hidden_states = model_output_hidden_states.repeat(self.num_speculative_tokens + 1, 1)
		335	+ hidden_states = hidden_states.repeat(self.num_speculative_tokens + 1, 1)
		336	+ hidden_dim = hidden_states.shape[-1]
		337	+ hit_hidden_states = backend.index_select(
		338	+ model_output_hidden_states, 0, hit_indices_per_token
		339	+ )
		340	+ scatter_index = hit_local_indices_per_token_tensor.unsqueeze(-1).expand(-1, hidden_dim)
		341	+ hidden_states.scatter_(0, scatter_index, hit_hidden_states)
		342	+ model_kwargs['hidden_states'] = hidden_states
		343	+
		344	+ # if lm_head_local_dp is not None and not input_metadata.is_dummy_batch:
240	if lm_head_local_dp is not None and not (len(lm_head_local_dp) == 1 and lm_head_local_dp[0] == 0):	345	if lm_head_local_dp is not None and not (len(lm_head_local_dp) == 1 and lm_head_local_dp[0] == 0):
241	hit_dp_speculative_length = filling_masks.get('hit_dp_speculative_length')	346	hit_dp_speculative_length = filling_masks.get('hit_dp_speculative_length')
242	- lm_head_local_dp[hit_mask] -= \	347	+ lm_head_local_dp.scatter_add_(
243	- hit_dp_speculative_length - hit_num_tokens_tensor	348	+ dim=0,
		349	+ index=hit_mask_local_tensor,
		350	+ src=hit_num_tokens_tensor - hit_dp_speculative_length
		351	+ )
244		352
245	- model_inputs.position_ids[hit_mask_per_token] += hit_num_tokens_per_token	353	+ model_inputs.position_ids.scatter_add_(0, hit_local_indices_per_token_tensor, hit_num_tokens_per_token)
246	- model_inputs.context_length[hit_mask] += hit_num_tokens_tensor.cpu().numpy()	354	+ model_inputs.context_length[hit_mask] += hit_num_tokens_numpy
247	- model_inputs.input_lengths[hit_mask_tensor] += hit_num_tokens_tensor
248	model_inputs.max_seq_len = max(model_inputs.context_length)	355	model_inputs.max_seq_len = max(model_inputs.context_length)
		356	+ model_inputs.input_lengths.scatter_add_(0, hit_mask_local_tensor, hit_num_tokens_tensor.to(torch.int32))
249	model_inputs.forward_context.attn_metadata.max_seq_len = model_inputs.max_seq_len	357	model_inputs.forward_context.attn_metadata.max_seq_len = model_inputs.max_seq_len
250	if self.infer_context.spcp_parallel_info.scp_size == 1:	358	if self.infer_context.spcp_parallel_info.scp_size == 1:
251	offset_start_indices = model_inputs.input_lengths[hit_mask_tensor] - speculative_length	359	offset_start_indices = model_inputs.input_lengths[hit_mask_tensor] - speculative_length
		360	+ offset_start_indices = torch.index_select(model_inputs.input_lengths, 0, hit_mask_local_tensor) - speculative_length
252	hit_increments = filling_masks.get('hit_increments')	361	hit_increments = filling_masks.get('hit_increments')
253	block_offsets = backend.repeat_interleave(offset_start_indices, speculative_length) + hit_increments	362	block_offsets = backend.repeat_interleave(offset_start_indices, speculative_length) + hit_increments
254	candidate_slots = filling_masks.get('candidate_slots')	363	candidate_slots = filling_masks.get('candidate_slots')
255	hit_block_indices = filling_masks.get('hit_block_indices')	364	hit_block_indices = filling_masks.get('hit_block_indices')
256	- model_inputs.slots[hit_mask_per_token] = candidate_slots[hit_block_indices, block_offsets]	365	+ hit_local_indices_per_token_tensor = filling_masks.get('hit_local_indices_per_token_tensor')
		366	+ model_inputs.forward_context.attn_metadata.slot_mapping.scatter_(
		367	+ 0,
		368	+ hit_local_indices_per_token_tensor,
		369	+ candidate_slots[hit_block_indices, block_offsets]
		370	+ )
257	else:	371	else:
258	model_inputs.cached_context_length[hit_mask] += hit_num_tokens_tensor	372	model_inputs.cached_context_length[hit_mask] += hit_num_tokens_tensor
259	indices = np.where(hit_mask)[0]	373	indices = np.where(hit_mask)[0]
	@@ -274,11 +388,21 @@ class MtpPlugin(Plugin):
274	)	388	)
275	* self.infer_context.spcp_parallel_info.sp_size	389	* self.infer_context.spcp_parallel_info.sp_size
276	][self.infer_context.spcp_parallel_info.sp_rank]	390	][self.infer_context.spcp_parallel_info.sp_rank]
		391	+ actual_seq_lengths_query = torch.ones_like(model_inputs.input_lengths, dtype=torch.int32)
		392	+ actual_seq_lengths_query = torch.cumsum(actual_seq_lengths_query, dim=0, dtype=torch.int32)
		393	+ model_inputs.forward_context.attn_metadata.actual_seq_lengths_query = actual_seq_lengths_query * speculative_length
		394	+ model_inputs.forward_context.attn_metadata.actual_seq_lengths_kv = model_inputs.input_lengths
		395	+ if not model_output_wrapper.input_metadata.is_prefill:
		396	+ model_inputs.forward_context.sub_forward_context.attn_metadata.actual_seq_lengths_kv = model_inputs.input_lengths #- self.num_speculative_tokens + 1
		397	+ model_inputs.forward_context.sub_forward_context.attn_metadata.actual_seq_lengths_query = model_inputs.forward_context.attn_metadata.actual_seq_lengths_query
		398	+ model_inputs.forward_context.attn_metadata.seq_lens = model_inputs.input_lengths
277		399
278	- hit_mask_per_token[hit_mask_mod] = False	400	+ hit_mask_per_token_mod_tensor = filling_masks.get('hit_mask_per_token_mod_tensor')
279	- model_inputs.input_ids[hit_mask_per_token] = hit_token_ids[hit_arange_tensor, hit_num_tokens_tensor - 1]	401	+ model_inputs.input_ids.scatter_(0, hit_mask_per_token_mod_tensor, hit_token_ids[hit_arange_tensor, hit_num_tokens_tensor - 1])
280		402
281	def sample_preprocess(self, logits, result, sampling_metadata, input_metadata):	403	def sample_preprocess(self, logits, result, sampling_metadata, input_metadata):
		404	+ if ENV.model_runner_exp:
		405	+ return self.sample_preprocess_exp(logits, result, sampling_metadata, input_metadata)
282	self.sampling_param = sampling_metadata	406	self.sampling_param = sampling_metadata
283	self.decoding_policy.sampling_param = self.sampling_param	407	self.decoding_policy.sampling_param = self.sampling_param
284	self.input_metadata = input_metadata	408	self.input_metadata = input_metadata
	@@ -303,6 +427,27 @@ class MtpPlugin(Plugin):
303		427
304	return logits	428	return logits
305		429
		430	+ def sample_preprocess_exp(self, logits, result, sampling_metadata, input_metadata):
		431	+ self.input_metadata = input_metadata
		432	+ self.sampling_param = sampling_metadata
		433	+ self.decoding_policy.sampling_param = self.sampling_param
		434	+ if isinstance(result, tuple):
		435	+ logits = result[0]
		436	+ else:
		437	+ logits = result
		438	+ if sampling_metadata is None or sampling_metadata.is_prefill:
		439	+ return logits
		440	+ return logits
		441	+
		442	+ def compose_model_inputs_exp(self, sampling_metadata):
		443	+ all_sequence_ids = sampling_metadata.all_sequence_ids
		444	+ batch_size = len(all_sequence_ids)
		445	+ logits_num_per_batch = [self.num_speculative_tokens + 1] * batch_size
		446	+ req_id_new = np.concatenate([[req_id] * n for req_id, n in zip(all_sequence_ids, logits_num_per_batch)])
		447	+ sampling_metadata.all_sequence_ids = req_id_new
		448	+ sampling_metadata.parent_sequence_ids = req_id_new
		449	+ return sampling_metadata
		450	+
306	def plugin_verify(self, sampling_output, cache_ids, result):	451	def plugin_verify(self, sampling_output, cache_ids, result):
307	sampling_output.repeating_indices = np.arange(len(cache_ids))	452	sampling_output.repeating_indices = np.arange(len(cache_ids))
308	if self.input_metadata.is_prefill:	453	if self.input_metadata.is_prefill:
	@@ -332,80 +477,53 @@ class MtpPlugin(Plugin):
332	self.decoding_policy.stop_criteria(sampling_output, output_space_left, next_tokens_indices)	477	self.decoding_policy.stop_criteria(sampling_output, output_space_left, next_tokens_indices)
333	self.reshape_speculative_outputs(sampling_output, next_tokens_indices)	478	self.reshape_speculative_outputs(sampling_output, next_tokens_indices)
334		479
335	- def plugin_verify_exp(self, sampling_output, cache_ids, result):	480	+ def plugin_verify_exp(self, sampling_output, cache_ids, result):
336	sampling_output.repeating_indices = np.arange(len(cache_ids))	481	sampling_output.repeating_indices = np.arange(len(cache_ids))
337	if self.input_metadata.is_prefill:	482	if self.input_metadata.is_prefill:
338	return	483	return
339	- draft_token = result[2].cpu()	484	+ draft_token = result[2]
340	next_tokens_uncheck = sampling_output.token_ids	485	next_tokens_uncheck = sampling_output.token_ids
341	input_metadata = self.input_metadata	486	input_metadata = self.input_metadata
342	- next_tokens_indices = []
343	out_seq_len = 1 if input_metadata.is_prefill else (self.num_speculative_tokens + 1)	487	out_seq_len = 1 if input_metadata.is_prefill else (self.num_speculative_tokens + 1)
344	- start_pos = 0
345	- draft_token_num_per_batch = self.num_speculative_tokens
346	- for batch in range(input_metadata.batch_size):
347	- end = start_pos + out_seq_len
348	- next_guess_by_batch = next_tokens_uncheck[start_pos:end].cpu()
349		488
350	- verify_guess_tokens = \	489	+ batch_size = next_tokens_uncheck.shape[0] // out_seq_len
351	- draft_token[batch * draft_token_num_per_batch: (batch + 1) * draft_token_num_per_batch].view(-1)	490	+ next_guess_all = next_tokens_uncheck.view(batch_size, out_seq_len)
		491	+ verify_guess_all = draft_token.view(batch_size, self.num_speculative_tokens)
		492	+ matches = (verify_guess_all == next_guess_all[:, :self.num_speculative_tokens])
		493	+ indices_counts = torch.cumprod(matches.to(torch.int32), dim=1).sum(dim=1) + 1
352		494
353	- indices = self.decoding_policy.verify_greedy_one_batch(verify_guess_tokens, next_guess_by_batch)	495	+ max_num_tokens = out_seq_len
354	- next_tokens_indices.append(list(range(start_pos, start_pos + indices + 1)))	496	+ token_range = self.token_range[:max_num_tokens].unsqueeze(0).expand(batch_size, -1)
355	- start_pos += out_seq_len	497	+ mask = token_range < indices_counts.unsqueeze(1)
356	- output_token_len = self.infer_context.get_output_len_count(cache_ids)	498	+ batch_offsets_cpu = torch.arange(batch_size) * out_seq_len
357	- output_space_left1 = self.input_metadata.batch_max_output_lens - output_token_len	499	+ raw_token_ids = sampling_output.token_ids.view(batch_size, out_seq_len)
358	- output_space_left2 = self.max_seq_len - \	500	+ raw_logprobs = sampling_output.logprobs.view(batch_size, out_seq_len)
359	- self.infer_context.get_seq_lens(cache_ids)
360	- output_space_left = np.minimum(output_space_left1, output_space_left2)
361	- self.decoding_policy.stop_criteria(sampling_output, output_space_left, next_tokens_indices)
362	- self.reshape_speculative_outputs_exp(sampling_output, next_tokens_indices)
363		501
364	- def reshape_speculative_outputs_exp(	502	+ sampling_output.token_ids = torch.where(mask, raw_token_ids, self.pad_token_id)
365	- self,	503	+ sampling_output.logprobs = torch.where(mask, raw_logprobs, -9999.0)
366	- sampling_output,	504	+ if sampling_output.top_token_ids.shape[-1] > 0:
367	- next_tokens_indices,	505	+ k_dim = sampling_output.top_token_ids.shape[-1]
368	- num_new_tokens=None	506	+ raw_top_tokens = sampling_output.top_token_ids.view(batch_size, out_seq_len, k_dim)
369	- ) -> None:	507	+ raw_top_probs = sampling_output.top_logprobs.view(batch_size, out_seq_len, k_dim)
370	- num_new_tokens = [len(indices) for indices in next_tokens_indices]	508	+
371	- max_num_tokens = max(num_new_tokens)	509	+ full_mask_k = mask.unsqueeze(-1)
372	- sequence_ids = []	510	+ sampling_output.top_token_ids = torch.where(full_mask_k, raw_top_tokens, self.pad_token_id)
373	- parent_sequence_ids = []	511	+ sampling_output.top_logprobs = torch.where(full_mask_k, raw_top_probs, -9999.0)
374	- group_indices = []	512	+ else:
375	- token_ids = []	513	+ sampling_output.top_token_ids = torch.zeros((batch_size, max_num_tokens, 0),
376	- logprobs = []	514	+ dtype=torch.int64, device="npu")
377	- top_token_ids = []	515	+ sampling_output.top_logprobs = torch.zeros((batch_size, max_num_tokens, 0),
378	- top_logprobs = []	516	+ dtype=torch.float32, device="npu")
379	- for i, indices in enumerate(next_tokens_indices):	517	+
380	- if sampling_output.sequence_ids is not None:	518	+ sampling_output.num_new_tokens = indices_counts
381	- sequence_ids.append(sampling_output.sequence_ids[indices[0]])	519	+ if sampling_output.sequence_ids is not None:
382	- parent_sequence_ids.append(sampling_output.parent_sequence_ids[indices[0]])	520	+ seq_ids_npu = torch.as_tensor(sampling_output.sequence_ids)
383	- group_indices.append((i, i + 1))	521	+ parent_ids_npu = torch.as_tensor(sampling_output.parent_sequence_ids)
384	- padded_array = backend.nn.functional.pad(sampling_output.token_ids[indices],	522	+
385	- pad=(0, max_num_tokens - len(indices)), mode='constant', value=self.pad_token_id)	523	+ sampling_output.sequence_ids = seq_ids_npu[batch_offsets_cpu].reshape(batch_size).numpy()
386	- token_ids.append(padded_array)	524	+ sampling_output.parent_sequence_ids = parent_ids_npu[batch_offsets_cpu].reshape(batch_size).numpy()
387	- padded_array = backend.nn.functional.pad(sampling_output.logprobs[indices],	525	+
388	- pad=(0, max_num_tokens - len(indices)), mode='constant', value=-9999.0)	526	+ sampling_output.group_indices = [(i, i + 1) for i in range(batch_size)]
389	- logprobs.append(padded_array)
390	- if 0 not in sampling_output.top_token_ids.shape:
391	- padded_array = np.pad(sampling_output.top_token_ids[indices],
392	- pad_width=((0, max_num_tokens - len(indices)), (0, 0)),
393	- constant_values=self.pad_token_id)
394	- top_token_ids.append(padded_array)
395	- padded_array = np.pad(sampling_output.top_logprobs[indices],
396	- pad_width=((0, max_num_tokens - len(indices)), (0, 0)), constant_values=-9999.0)
397	- top_logprobs.append(padded_array)
398	- else:
399	- top_token_ids.append(np.zeros((max_num_tokens, 0), dtype=np.int64))
400	- top_logprobs.append(np.zeros((max_num_tokens, 0), dtype=np.float32))
401	- sampling_output.sequence_ids = np.array(sequence_ids)
402	- sampling_output.parent_sequence_ids = np.array(parent_sequence_ids)
403	- sampling_output.group_indices = group_indices
404	- sampling_output.token_ids = backend.stack(token_ids).to(backend.int64)
405	- sampling_output.logprobs = backend.stack(logprobs)
406	- sampling_output.top_token_ids = backend.from_numpy(np.array(top_token_ids))
407	- sampling_output.top_logprobs = backend.from_numpy(np.array(top_logprobs))
408	- sampling_output.num_new_tokens = backend.tensor(num_new_tokens, device="npu")
409		527
410	def plugin_cache_update(self, cache_ids, sampling_output, la_cache_input, is_prefill=False):	528	def plugin_cache_update(self, cache_ids, sampling_output, la_cache_input, is_prefill=False):
411	result, _ = la_cache_input	529	result, _ = la_cache_input

Mmindie_llm/text_generator/plugins/plugin_manager.py+22-5

	@@ -15,6 +15,7 @@ import copy
15	from typing import Iterable, Optional, Any	15	from typing import Iterable, Optional, Any
16	from dataclasses import fields	16	from dataclasses import fields
17	import torch	17	import torch
		18	+import torch.nn.functional as F
18	import numpy as np	19	import numpy as np
19	import numpy.typing as npt	20	import numpy.typing as npt
20	from mindie_llm.runtime.layers.attention.backend.sparse_attention import main_stream	21	from mindie_llm.runtime.layers.attention.backend.sparse_attention import main_stream
	@@ -86,11 +87,14 @@ class PluginManager:
86	def unsqueeze_sampling_output(sampling_output: SamplingOutput):	87	def unsqueeze_sampling_output(sampling_output: SamplingOutput):
87	if ENV.model_runner_exp and ENV.async_inference:	88	if ENV.model_runner_exp and ENV.async_inference:
88	sampling_output.token_ids = torch.unsqueeze(sampling_output.token_ids, 1)	89	sampling_output.token_ids = torch.unsqueeze(sampling_output.token_ids, 1)
		90	+ sampling_output.logprobs = torch.unsqueeze(sampling_output.logprobs, 1)
		91	+ sampling_output.top_token_ids = torch.unsqueeze(sampling_output.top_token_ids, 1)
		92	+ sampling_output.top_logprobs = torch.unsqueeze(sampling_output.top_logprobs, 1)
89	else:	93	else:
90	sampling_output.token_ids = np.expand_dims(sampling_output.token_ids, 1)	94	sampling_output.token_ids = np.expand_dims(sampling_output.token_ids, 1)
91	- sampling_output.logprobs = np.expand_dims(sampling_output.logprobs, 1)	95	+ sampling_output.logprobs = np.expand_dims(sampling_output.logprobs, 1)
92	- sampling_output.top_token_ids = np.expand_dims(sampling_output.top_token_ids, 1)	96	+ sampling_output.top_token_ids = np.expand_dims(sampling_output.top_token_ids, 1)
93	- sampling_output.top_logprobs = np.expand_dims(sampling_output.top_logprobs, 1)	97	+ sampling_output.top_logprobs = np.expand_dims(sampling_output.top_logprobs, 1)
94		98
95	@staticmethod	99	@staticmethod
96	def _to_host(data_instance: Any):	100	def _to_host(data_instance: Any):
	@@ -342,6 +346,12 @@ class PluginManager:
342	else:	346	else:
343	model_inputs, sampling_metadata, trace_ids = \	347	model_inputs, sampling_metadata, trace_ids = \
344	self.infer_context.compose_model_inputs(input_metadata, cache_ids, warmup=False, hit_mask=hit_mask)	348	self.infer_context.compose_model_inputs(input_metadata, cache_ids, warmup=False, hit_mask=hit_mask)
		349	+ if ENV.model_runner_exp and (sampling_metadata is not None) and (not sampling_metadata.is_prefill):
		350	+ for plugin in self.plugin_list:
		351	+ plugin_instance = getattr(self, plugin, None)
		352	+ method = getattr(plugin_instance, 'compose_model_inputs_exp', None)
		353	+ if method is not None:
		354	+ sampling_metadata = method(sampling_metadata)
345	res = (cache_ids, model_inputs, sampling_metadata, trace_ids)	355	res = (cache_ids, model_inputs, sampling_metadata, trace_ids)
346	return res	356	return res
347		357
	@@ -430,7 +440,8 @@ class PluginManager:
430	model_inputs, input_metadata, sampling_metadata, cache_ids, input_len_mask, **kwargs)	440	model_inputs, input_metadata, sampling_metadata, cache_ids, input_len_mask, **kwargs)
431	(q_len, spec_mask) = input_len_mask	441	(q_len, spec_mask) = input_len_mask
432	# 需要确保虚推 context_length 至少为 1，否则会导致模型内部维度计算错误 (coreDim = 0)	442	# 需要确保虚推 context_length 至少为 1，否则会导致模型内部维度计算错误 (coreDim = 0)
433	- if input_metadata.all_sequence_ids is not None and not input_metadata.is_prefill:	443	+ if (input_metadata.all_sequence_ids is not None) and \
		444	+ (input_metadata is not None) and (not input_metadata.is_prefill):
434	has_simulate = any(sid == SIMULATE_SEQUENCE_ID for sid in input_metadata.all_sequence_ids)	445	has_simulate = any(sid == SIMULATE_SEQUENCE_ID for sid in input_metadata.all_sequence_ids)
435	if has_simulate and model_inputs.context_length[0] == 0:	446	if has_simulate and model_inputs.context_length[0] == 0:
436	model_inputs.context_length[0] = 1	447	model_inputs.context_length[0] = 1
	@@ -670,6 +681,12 @@ class PluginManager:
670	model_inputs.context_length[hit_sequence_ids_mask] += 1	681	model_inputs.context_length[hit_sequence_ids_mask] += 1
671	model_inputs.max_seq_len = max(model_inputs.context_length)	682	model_inputs.max_seq_len = max(model_inputs.context_length)
672	model_inputs.forward_context.attn_metadata.max_seq_len = model_inputs.max_seq_len	683	model_inputs.forward_context.attn_metadata.max_seq_len = model_inputs.max_seq_len
		684	+ actual_seq_lengths_query = torch.ones_like(model_inputs.input_lengths, dtype=torch.int32)
		685	+ actual_seq_lengths_query = F.pad(torch.cumsum(actual_seq_lengths_query, dim=0, dtype=torch.int32), (0, 0), value=0)
		686	+ seq_lens_q = actual_seq_lengths_query[1:] - actual_seq_lengths_query[:-1]
		687	+ model_inputs.forward_context.attn_metadata.actual_seq_lengths_query = actual_seq_lengths_query
		688	+ model_inputs.forward_context.attn_metadata.actual_seq_lengths_kv = model_inputs.input_lengths
		689	+ model_inputs.forward_context.attn_metadata.seq_lens = model_inputs.input_lengths
673		690
674	def _get_token_num_per_seq(self, input_metadata):	691	def _get_token_num_per_seq(self, input_metadata):
675	# computed_blocks为None时prefixcache无命中，batch_seq_len即为q_len	692	# computed_blocks为None时prefixcache无命中，batch_seq_len即为q_len
	@@ -684,4 +701,4 @@ class PluginManager:
684	).astype(np.int64)	701	).astype(np.int64)
685	cache_len_eaqual_mask = input_metadata.batch_is_prefill & (token_num_per_seq == 0)	702	cache_len_eaqual_mask = input_metadata.batch_is_prefill & (token_num_per_seq == 0)
686	token_num_per_seq[cache_len_eaqual_mask] = self.generator_backend.block_size	703	token_num_per_seq[cache_len_eaqual_mask] = self.generator_backend.block_size
687	- return token_num_per_seq	704	+ return token_num_per_seq

Mmindie_llm/text_generator/utils/output_filter.py+27-6

	@@ -86,15 +86,34 @@ class OutputFilter:
86	end_reason[eos_idx] = ResponseConfig.EOS	86	end_reason[eos_idx] = ResponseConfig.EOS
87	return filter_ids_arr	87	return filter_ids_arr
88		88
89	- def filter_by_length(self, cache_ids, batch_max_output_lens, num_new_tokens, filter_ids_arr, end_reason):	89	+ def filter_by_length(self, cache_ids, batch_max_output_lens, sampling_output, filter_ids_arr, end_reason):
		90	+ def adjust_num_new_tokens(cached_lens, num_new_tokens, max_seq_len):
		91	+ """
		92	+ 调整 num_new_tokens 的值，确保 cached_seq_lens + num_new_tokens <= max_seq_len。
		93	+
		94	+ 参数:
		95	+ cached_seq_lens (np.ndarray): 已缓存的序列长度数组。
		96	+ num_new_tokens (np.ndarray): 新的 token 数量数组。
		97	+ max_seq_len (int): 最大序列长度。
		98	+
		99	+ 返回:
		100	+ np.ndarray: 调整后的 num_new_tokens。
		101	+ """
		102	+ available_tokens = max_seq_len - cached_lens
		103	+ available_tokens = np.maximum(available_tokens, 0)
		104	+ adjusted_num_new_tokens = np.minimum(num_new_tokens, available_tokens)
		105	+ return adjusted_num_new_tokens
90	output_len_count = self.tg_infer_context.get_output_len_count(cache_ids)	106	output_len_count = self.tg_infer_context.get_output_len_count(cache_ids)
91	cached_seq_lens = self.tg_infer_context.get_seq_lens(cache_ids)	107	cached_seq_lens = self.tg_infer_context.get_seq_lens(cache_ids)
		108	+ sampling_output.num_new_tokens = adjust_num_new_tokens(cached_seq_lens, sampling_output.num_new_tokens, self.cache_config.max_seq_len)
		109	+ sampling_output.num_new_tokens = adjust_num_new_tokens(output_len_count, sampling_output.num_new_tokens, batch_max_output_lens)
		110	+ sampling_output.num_new_tokens = adjust_num_new_tokens(output_len_count, sampling_output.num_new_tokens, self.cache_config.max_gen_len)
92	exceed_seq_limit_idx = np.nonzero(	111	exceed_seq_limit_idx = np.nonzero(
93	- cached_seq_lens + num_new_tokens >= self.cache_config.max_seq_len)[0]	112	+ cached_seq_lens + sampling_output.num_new_tokens >= self.cache_config.max_seq_len)[0]
94	exceed_user_output_limit_idx = np.nonzero(	113	exceed_user_output_limit_idx = np.nonzero(
95	- output_len_count + num_new_tokens >= batch_max_output_lens)[0]	114	+ output_len_count + sampling_output.num_new_tokens >= batch_max_output_lens)[0]
96	exceed_global_output_limit_idx = np.nonzero(	115	exceed_global_output_limit_idx = np.nonzero(
97	- output_len_count + num_new_tokens >= self.cache_config.max_gen_len)[0]	116	+ output_len_count + sampling_output.num_new_tokens >= self.cache_config.max_gen_len)[0]
98		117
99	if exceed_seq_limit_idx.size != 0:	118	if exceed_seq_limit_idx.size != 0:
100	filter_ids_arr = np.union1d(filter_ids_arr, exceed_seq_limit_idx)	119	filter_ids_arr = np.union1d(filter_ids_arr, exceed_seq_limit_idx)
	@@ -168,6 +187,8 @@ class OutputFilter:
168	metadata: InputMetadata,	187	metadata: InputMetadata,
169	sampling_output: SamplingOutput	188	sampling_output: SamplingOutput
170	):	189	):
		190	+ for eos in self.eos_token_id:
		191	+ sampling_output.truncate_after_eos(eos)
171	filter_ids_arr = np.array([], dtype=np.int32)	192	filter_ids_arr = np.array([], dtype=np.int32)
172	if sampling_output.repeating_indices is not None:	193	if sampling_output.repeating_indices is not None:
173	cache_ids = cache_ids[sampling_output.repeating_indices]	194	cache_ids = cache_ids[sampling_output.repeating_indices]
	@@ -184,7 +205,7 @@ class OutputFilter:
184	filter_ids_arr, truncation_indices = self.filter_by_stop(	205	filter_ids_arr, truncation_indices = self.filter_by_stop(
185	cache_ids, sampling_output.token_ids, sampling_output.num_new_tokens, filter_ids_arr, end_reason)	206	cache_ids, sampling_output.token_ids, sampling_output.num_new_tokens, filter_ids_arr, end_reason)
186	filter_ids_arr = self.filter_by_length(	207	filter_ids_arr = self.filter_by_length(
187	- cache_ids, batch_max_output_lens, sampling_output.num_new_tokens, filter_ids_arr, end_reason)	208	+ cache_ids, batch_max_output_lens, sampling_output, filter_ids_arr, end_reason)
188		209
189	# 对于没有完成prefill的req，不可以清除后处理参数	210	# 对于没有完成prefill的req，不可以清除后处理参数
190	if metadata.is_mix is not None:	211	if metadata.is_mix is not None:
	@@ -194,4 +215,4 @@ class OutputFilter:
194		215
195	if filter_ids_arr.shape[0] == 0:	216	if filter_ids_arr.shape[0] == 0:
196	filter_ids_arr = None	217	filter_ids_arr = None
197	- return end_reason, filter_ids_arr, truncation_indices	218	+ return end_reason, filter_ids_arr, truncation_indices

Mmindie_llm/text_generator/utils/sampling_output.py+21-1

	@@ -25,8 +25,28 @@ class SamplingOutput:
25	top_logprobs: np.ndarray	25	top_logprobs: np.ndarray
26	cumulative_logprobs: np.ndarray	26	cumulative_logprobs: np.ndarray
27	num_new_tokens: np.ndarray	27	num_new_tokens: np.ndarray
		28	+ num_new_tokens_numpy: np.ndarray = None
28	num_top_tokens: np.ndarray = None	29	num_top_tokens: np.ndarray = None
29	seeds: np.ndarray = None	30	seeds: np.ndarray = None
30		31
31	def to_deprecated(self):	32	def to_deprecated(self):
32	- return self.token_ids, self.logprobs	33	+ return self.token_ids, self.logprobs
		34	+
		35	+ def truncate_after_eos(self, eos_token_id: int):
		36	+ """
		37	+ Truncates the token_ids after the first occurrence of the EOS token in each row,
		38	+ and updates num_new_tokens accordingly.
		39	+
		40	+ :param eos_token_id: The ID of the EOS (End Of Sequence) token.
		41	+ """
		42	+ for i in range(self.token_ids.shape[0]):
		43	+ # Find the index of the first occurrence of EOS token in the row
		44	+ eos_index = np.where(self.token_ids[i] == eos_token_id)[0]
		45	+ if eos_index.size > 0:
		46	+ eos_index = eos_index[0]
		47	+ # Truncate the token_ids after the first EOS
		48	+ self.token_ids[i, eos_index + 1:] = 0 # Optionally set to 0 or any padding value
		49	+ self.num_new_tokens[i] = eos_index + 1 # Update the num_new_tokens
		50	+ else:
		51	+ # If no EOS token is found, num_new_tokens remains unchanged
		52	+ pass

	@@ -93,7 +93,7 @@ class Sampler:
93	def greedy_select_on_device(logits: torch.Tensor, metadata: SamplingMetadata):	93	def greedy_select_on_device(logits: torch.Tensor, metadata: SamplingMetadata):
94	batch_size = len(logits)	94	batch_size = len(logits)
95	token_ids = logits.argmax(dim=-1).to(torch.int64)	95	token_ids = logits.argmax(dim=-1).to(torch.int64)
96	- logprobs = torch.full((batch_size,), -9999.0, dtype=torch.float32)	96	+ logprobs = torch.full((batch_size,), -9999.0, dtype=torch.float32, device="npu")
97	top_token_ids = torch.zeros((batch_size, 0), dtype=torch.int64)	97	top_token_ids = torch.zeros((batch_size, 0), dtype=torch.int64)
98	top_logprobs = torch.zeros((batch_size, 0), dtype=torch.float32)	98	top_logprobs = torch.zeros((batch_size, 0), dtype=torch.float32)
99	sampling_output = SamplingOutput(	99	sampling_output = SamplingOutput(
	@@ -107,6 +107,7 @@ class Sampler:
107	top_logprobs=top_logprobs,	107	top_logprobs=top_logprobs,
108	cumulative_logprobs=torch.zeros(batch_size, dtype=torch.float32),	108	cumulative_logprobs=torch.zeros(batch_size, dtype=torch.float32),
109	num_new_tokens=torch.ones(batch_size, dtype=torch.int64, device="npu"),	109	num_new_tokens=torch.ones(batch_size, dtype=torch.int64, device="npu"),
		110	+ num_new_tokens_numpy=np.ones(batch_size, dtype=np.int64),
110	num_top_tokens=metadata.num_top_tokens if metadata is not None else None	111	num_top_tokens=metadata.num_top_tokens if metadata is not None else None
111	)	112	)
112	return sampling_output	113	return sampling_output