Star181
850
代码
代码
Issues154
Pull Requests331
流水线
项目讨论
Wiki
分析
Star181
850
已合并
适配支持两类算子：ffn_worker_scheduler_/ffn_worker_scheduler、 attention_worker_scheduler_/attention_worker_scheduler #3397
pwc10490创建于 2025年10月27日
适配支持两类算子：ffn_worker_scheduler_/ffn_worker_scheduler、 attention_worker_scheduler_/attention_worker_scheduler #3397
已合并
pwc10490创建于 2025年10月27日
从已删除 :master合入到Ascend/op-pluginmaster
共 9 个文件变更+484-1
Mcodegen/templates/_op_plugin_docs.py+256-0
@@ -11316,6 +11316,262 @@ out = torch_npu.npu_top_k_top_p(logits, p, k)
11316)11316)
11317 11317 
11318 11318 
11319+_add_torch_npu_docstr(
11320+ "ffn_worker_scheduler_",
11321+ """
11322+接口原型:
11323+ffn_worker_scheduler_(Tensor(a!) self, *, int sync_group_size=1, int execute_mode=0) -> Tensor(a!)
11324+ 
11325+功能描述:
11326+Attention和Ffn分离部署场景下，Ffn侧数据扫描功能，扫描并原地完成数据整理。
11327+ 
11328+参数说明:
11329+scheduler_context(torch.Tensor): 输入张量, scheduler_context的定义与生成可参照torch_npu._afd包；
11330+sync_group_size(int): 可选，默认值为1；
11331+execute_mode(int): 可选，默认值为0。
11332+ 
11333+约束说明:
11334+无
11335+ 
11336+支持版本:
11337+PyTorch 2.1及更高版本
11338+ 
11339+支持的型号:
11340+Atlas A3训练系列产品
11341+Atlas A3推理系列产品
11342+ 
11343+调用示例:
11344+import torch
11345+import torch_npu
11346+import os
11347+ 
11348+window_size = 209715200
11349+ffn_window_tensor = torch.zeros([window_size], dtype=torch.int8).npu()
11350+ 
11351+attn_workers = 2
11352+micro_batch_number = 3
11353+batch_size = 6
11354+top_k = 8
11355+hidden_size = 7168
11356+expert_num = 288
11357+attn_to_ffn_token_size = (7168 + 4 + 511) // 512 * 512
11358+ffn_to_attn_token_size = 7168 * 2
11359+ffn_window = ffn_window_tensor.data_ptr()
11360+ 
11361+context_holder = torch_npu._afd.create_schedule_context_holder(schedule_mode = 0, session_num = attn_workers,
11362+ micro_batch_num = micro_batch_number, micro_batch_size = batch_size, selected_expert_num = top_k + 1,
11363+ expert_num = expert_num, attn_to_ffn_token_size = attn_to_ffn_token_size, ffn_to_attn_token_size = ffn_to_attn_token_size,
11364+ ffn_window = ffn_window, ffn_window_size = window_size)
11365+
11366+schedule_context = context_holder.get_schedule_context_tensor()
11367+ 
11368+ 
11369+def _set_all_flags():
11370+ num_int8 = attn_workers * micro_batch_number * (8 + batch_size * top_k * 4)
11371+ per_session_num = micro_batch_number * (8 + batch_size * top_k * 4)
11372+ int32_view = ffn_window_tensor[:num_int8].view(torch.int32)
11373+ int32_view[:] = 1
11374+ 
11375+ 
11376+_set_all_flags()
11377+torch_npu.ffn_worker_scheduler_(schedule_context, sync_group_size = 2)
11378+"""
11379+)
11380+ 
11381+ 
11382+_add_torch_npu_docstr(
11383+ "ffn_worker_scheduler",
11384+ """
11385+接口原型:
11386+ffn_worker_scheduler(Tensor self, *, int sync_group_size=1, int execute_mode=0) -> Tensor
11387+ 
11388+功能描述:
11389+Attention和Ffn分离部署场景下，Ffn侧数据扫描功能，扫描并完成数据整理输出。
11390+ 
11391+参数说明:
11392+scheduler_context(torch.Tensor): 输入张量, scheduler_context的定义与生成可参照torch_npu._afd包；
11393+sync_group_size(int): 可选，默认值为1；
11394+execute_mode(int): 可选，默认值为0。
11395+ 
11396+输出说明:
11397+scheduler_context(torch.Tensor): 输出结果张量
11398+ 
11399+约束说明:
11400+无
11401+ 
11402+支持版本:
11403+PyTorch 2.1及更高版本
11404+ 
11405+支持的型号:
11406+Atlas A3训练系列产品
11407+Atlas A3推理系列产品
11408+ 
11409+调用示例:
11410+import torch
11411+import torch_npu
11412+import os
11413+ 
11414+window_size = 209715200
11415+ffn_window_tensor = torch.zeros([window_size], dtype=torch.int8).npu()
11416+ 
11417+attn_workers = 2
11418+micro_batch_number = 3
11419+batch_size = 6
11420+top_k = 8
11421+hidden_size = 7168
11422+expert_num = 288
11423+attn_to_ffn_token_size = (7168 + 4 + 511) // 512 * 512
11424+ffn_to_attn_token_size = 7168 * 2
11425+ffn_window = ffn_window_tensor.data_ptr()
11426+ 
11427+context_holder = torch_npu._afd.create_schedule_context_holder(schedule_mode = 0, session_num = attn_workers,
11428+ micro_batch_num = micro_batch_number, micro_batch_size = batch_size, selected_expert_num = top_k + 1,
11429+ expert_num = expert_num, attn_to_ffn_token_size = attn_to_ffn_token_size, ffn_to_attn_token_size = ffn_to_attn_token_size,
11430+ ffn_window = ffn_window, ffn_window_size = window_size)
11431+
11432+schedule_context = context_holder.get_schedule_context_tensor()
11433+ 
11434+ 
11435+def _set_all_flags():
11436+ num_int8 = attn_workers * micro_batch_number * (8 + batch_size * top_k * 4)
11437+ per_session_num = micro_batch_number * (8 + batch_size * top_k * 4)
11438+ int32_view = ffn_window_tensor[:num_int8].view(torch.int32)
11439+ int32_view[:] = 1
11440+ 
11441+ 
11442+_set_all_flags()
11443+schedule_context_out = torch_npu.ffn_worker_scheduler(schedule_context, sync_group_size = 2)
11444+"""
11445+)
11446+ 
11447+ 
11448+_add_torch_npu_docstr(
11449+ "attention_worker_scheduler_",
11450+ """
11451+接口原型:
11452+attention_worker_scheduler_(Tensor(a!) self) -> Tensor(a!)
11453+ 
11454+功能描述:
11455+Attention和Ffn分离部署场景下，Attention侧数据扫描功能，扫描并原地确保数据就绪。
11456+ 
11457+参数说明:
11458+scheduler_context(torch.Tensor): 输入张量, scheduler_context的定义与生成可参照torch_npu._afd包。
11459+ 
11460+约束说明:
11461+无
11462+ 
11463+支持版本:
11464+PyTorch 2.1及更高版本
11465+ 
11466+支持的型号:
11467+Atlas A3训练系列产品
11468+Atlas A3推理系列产品
11469+ 
11470+调用示例:
11471+import torch
11472+import torch_npu
11473+import os
11474+ 
11475+window_size = 209715200
11476+attn_window_tensor = torch.zeros([window_size], dtype=torch.int8).npu()
11477+ 
11478+attn_workers = 144
11479+micro_batch_number = 3
11480+batch_size = 30
11481+top_k = 8
11482+hidden_size = 7168
11483+expert_num = 288
11484+attn_to_ffn_token_size = (7168 + 4 + 511) // 512 * 512
11485+ffn_to_attn_token_size = 7168 * 2
11486+attn_window = attn_window_tensor.data_ptr()
11487+ 
11488+context_holder = torch_npu._afd.create_schedule_context_holder(schedule_mode = 1, session_num = attn_workers,
11489+ micro_batch_num = micro_batch_number, micro_batch_size = batch_size, selected_expert_num = top_k + 1,
11490+ expert_num = expert_num, attn_to_ffn_token_size = attn_to_ffn_token_size, ffn_to_attn_token_size = ffn_to_attn_token_size,
11491+ attention_window = attn_window, attention_window_size = window_size)
11492+
11493+schedule_context = context_holder.get_schedule_context_tensor()
11494+ 
11495+def _set_all_flags():
11496+ num_int8 = batch_size * (top_k + 1) * 4 * micro_batch_number
11497+ 
11498+ int32_view = attn_window_tensor[:num_int8].view(torch.int32)
11499+ 
11500+ int32_view[:] = 1
11501+ 
11502+ 
11503+_set_all_flags()
11504+torch_npu.attention_worker_scheduler_(schedule_context)
11505+ 
11506+"""
11507+)
11508+ 
11509+ 
11510+_add_torch_npu_docstr(
11511+ "attention_worker_scheduler",
11512+ """
11513+接口原型:
11514+attention_worker_scheduler(Tensor self) -> Tensor
11515+ 
11516+功能描述:
11517+Attention和Ffn分离部署场景下，Attention侧数据扫描功能，扫描并确保数据就绪。
11518+ 
11519+参数说明:
11520+scheduler_context(torch.Tensor): 输入张量, scheduler_context的定义与生成可参照torch_npu._afd包。
11521+ 
11522+输出说明:
11523+scheduler_context(torch.Tensor): 输出结果张量
11524+ 
11525+约束说明:
11526+无
11527+ 
11528+支持版本:
11529+PyTorch 2.1及更高版本
11530+ 
11531+支持的型号:
11532+Atlas A3训练系列产品
11533+Atlas A3推理系列产品
11534+ 
11535+调用示例:
11536+import torch
11537+import torch_npu
11538+import os
11539+ 
11540+window_size = 209715200
11541+attn_window_tensor = torch.zeros([window_size], dtype=torch.int8).npu()
11542+ 
11543+attn_workers = 144
11544+micro_batch_number = 3
11545+batch_size = 30
11546+top_k = 8
11547+hidden_size = 7168
11548+expert_num = 288
11549+attn_to_ffn_token_size = (7168 + 4 + 511) // 512 * 512
11550+ffn_to_attn_token_size = 7168 * 2
11551+attn_window = attn_window_tensor.data_ptr()
11552+ 
11553+context_holder = torch_npu._afd.create_schedule_context_holder(schedule_mode = 1, session_num = attn_workers,
11554+ micro_batch_num = micro_batch_number, micro_batch_size = batch_size, selected_expert_num = top_k + 1,
11555+ expert_num = expert_num, attn_to_ffn_token_size = attn_to_ffn_token_size, ffn_to_attn_token_size = ffn_to_attn_token_size,
11556+ attention_window = attn_window, attention_window_size = window_size)
11557+
11558+schedule_context = context_holder.get_schedule_context_tensor()
11559+ 
11560+def _set_all_flags():
11561+ num_int8 = batch_size * (top_k + 1) * 4 * micro_batch_number
11562+ 
11563+ int32_view = attn_window_tensor[:num_int8].view(torch.int32)
11564+ 
11565+ int32_view[:] = 1
11566+ 
11567+ 
11568+_set_all_flags()
11569+schedule_context_out = torch_npu.attention_worker_scheduler(schedule_context)
11570+ 
11571+"""
11572+)
11573+ 
11574+ 
11319_add_torch_npu_docstr(11575_add_torch_npu_docstr(
11320 "npu_top_k_top_p_sample",11576 "npu_top_k_top_p_sample",
11321 """11577 """
Mop_plugin/config/op_plugin_functions.yaml+16-0
@@ -5428,6 +5428,22 @@ autograd:
5428 op_api: [v2.5, newest]5428 op_api: [v2.5, newest]
5429 5429 
5430custom:5430custom:
5431+ - func: ffn_worker_scheduler_(Tensor(a!) self, *, int sync_group_size=1, int execute_mode=0) -> Tensor(a!)
5432+ op_api: [v2.1, newest]
5433+ gen_opapi:
5434+ exec: aclnnInplaceFfnWorkerScheduler
5435+ 
5436+ - func: ffn_worker_scheduler(Tensor self, *, int sync_group_size=1, int execute_mode=0) -> Tensor
5437+ op_api: [v2.1, newest]
5438+ 
5439+ - func: attention_worker_scheduler_(Tensor(a!) self) -> Tensor(a!)
5440+ op_api: [v2.1, newest]
5441+ gen_opapi:
5442+ exec: aclnnInplaceAttentionWorkerScheduler
5443+ 
5444+ - func: attention_worker_scheduler(Tensor self) -> Tensor
5445+ op_api: [v2.1, newest]
5446+ 
5431 - func: npu_top_k_top_p_sample(Tensor logits, Tensor top_k, Tensor top_p, Tensor? q=None, float eps=1e-8, bool is_need_logits=False, int top_k_guess=32) -> (Tensor, Tensor)5447 - func: npu_top_k_top_p_sample(Tensor logits, Tensor top_k, Tensor top_p, Tensor? q=None, float eps=1e-8, bool is_need_logits=False, int top_k_guess=32) -> (Tensor, Tensor)
5432 op_api: all_version5448 op_api: all_version
5433 gen_opapi:5449 gen_opapi:
Aop_plugin/ops/opapi/AttentionWorkerSchedulerKernelNpuOpApi.cpp+20-0
@@ -0,0 +1,20 @@
1+// Copyright (c) 2025 Huawei Technologies Co., Ltd
2+// All rights reserved.
3+ 
4+#include "op_plugin/OpApiInterface.h"
5+#include "op_plugin/utils/op_api_common.h"
6+ 
7+namespace op_api {
8+ using npu_preparation = at_npu::native::OpPreparation;
9+ 
10+ at::Tensor attention_worker_scheduler(const at::Tensor & self)
11+ {
12+ auto output_size_0 = self.sizes();
13+ auto output_dtype_0 = self.scalar_type();
14+ at::Tensor out = npu_preparation::apply_tensor_without_format(output_size_0, self.options().dtype(output_dtype_0));
15+ 
16+ out.copy_(self);
17+ EXEC_NPU_CMD(aclnnInplaceAttentionWorkerScheduler, out);
18+ return out;
19+ }
20+}
Aop_plugin/ops/opapi/FfnWorkerSchedulerKernelNpuOpApi.cpp+20-0
@@ -0,0 +1,20 @@
1+// Copyright (c) 2025 Huawei Technologies Co., Ltd
2+// All rights reserved.
3+ 
4+#include "op_plugin/OpApiInterface.h"
5+#include "op_plugin/utils/op_api_common.h"
6+ 
7+namespace op_api {
8+ using npu_preparation = at_npu::native::OpPreparation;
9+ 
10+ at::Tensor ffn_worker_scheduler(const at::Tensor & self, int64_t sync_group_size, int64_t execute_mode)
11+ {
12+ auto output_size_0 = self.sizes();
13+ auto output_dtype_0 = self.scalar_type();
14+ at::Tensor out = npu_preparation::apply_tensor_without_format(output_size_0, self.options().dtype(output_dtype_0));
15+ 
16+ out.copy_(self);
17+ EXEC_NPU_CMD(aclnnInplaceFfnWorkerScheduler, out, sync_group_size, execute_mode);
18+ return out;
19+ }
20+}
Mop_plugin/python/meta/_meta_registrations.py+20-0
@@ -687,6 +687,26 @@ def npu_moe_init_routing_v2_meta(x, expert_idx, *, scale=None, offset=None, acti
687 x.new_empty(tuple(expanded_scale_dim_list), dtype=torch.float32))687 x.new_empty(tuple(expanded_scale_dim_list), dtype=torch.float32))
688 688 
689 689 
690+@impl(m, "ffn_worker_scheduler_")
691+def ffn_worker_scheduler__meta(self, *, sync_group_size=1, execute_mode=0):
692+ return self
693+ 
694+ 
695+@impl(m, "attention_worker_scheduler_")
696+def attention_worker_scheduler__meta(self):
697+ return self
698+ 
699+ 
700+@impl(m, "ffn_worker_scheduler")
701+def ffn_worker_scheduler_meta(self, *, sync_group_size=1, execute_mode=0):
702+ return torch.empty_like(self)
703+ 
704+ 
705+@impl(m, "attention_worker_scheduler")
706+def attention_worker_scheduler_meta(self):
707+ return torch.empty_like(self)
708+ 
709+ 
690@impl(m, "npu_ffn_worker_batching")710@impl(m, "npu_ffn_worker_batching")
691def npu_ffn_worker_batching(schedule_context, expert_num, max_out_shape, *, token_dtype=0, need_schedule=0, layer_num=0):711def npu_ffn_worker_batching(schedule_context, expert_num, max_out_shape, *, token_dtype=0, need_schedule=0, layer_num=0):
692 Y_size = max_out_shape[0] * max_out_shape[1] * max_out_shape[2]712 Y_size = max_out_shape[0] * max_out_shape[1] * max_out_shape[2]
Mtest/allowlist_for_publicAPI.json+5-1
@@ -77,7 +77,11 @@
77 "npu_gather_sparse_index",77 "npu_gather_sparse_index",
78 "npu_moe_distribute_combine_add_rms_norm",78 "npu_moe_distribute_combine_add_rms_norm",
79 "npu_moe_update_expert",79 "npu_moe_update_expert",
80- "npu_dynamic_block_quant"80+ "npu_dynamic_block_quant",
81+ "attention_worker_scheduler_",
82+ "attention_worker_scheduler",
83+ "ffn_worker_scheduler_",
84+ "ffn_worker_scheduler"
81 ]85 ]
82 }86 }
83 87
Mtest/core_tests/torch_npu_OpApi_schema_all.json+36-0
@@ -128,6 +128,30 @@
128 "func: npu_add_rms_norm_dynamic_quant(Tensor x1, Tensor x2, Tensor gamma, *, Tensor? smooth_scale1=None, Tensor? smooth_scale2=None, Tensor? beta=None, float epsilon=1e-6, bool[2] output_mask=[]) -> (Tensor, Tensor, Tensor, Tensor, Tensor)": {128 "func: npu_add_rms_norm_dynamic_quant(Tensor x1, Tensor x2, Tensor gamma, *, Tensor? smooth_scale1=None, Tensor? smooth_scale2=None, Tensor? beta=None, float epsilon=1e-6, bool[2] output_mask=[]) -> (Tensor, Tensor, Tensor, Tensor, Tensor)": {
129 "version": ["all_version"]129 "version": ["all_version"]
130 },130 },
131+ "op_api: torch_npu.attention_worker_scheduler_(*args, **kwargs)": {
132+ "version": ["v2.1", "v2.5", "v2.6", "v2.7"]
133+ },
134+ "op_api: torch_npu.attention_worker_scheduler(*args, **kwargs)": {
135+ "version": ["v2.1", "v2.5", "v2.6", "v2.7"]
136+ },
137+ "op_api: torch_npu.ffn_worker_scheduler_(*args, **kwargs)": {
138+ "version": ["v2.1", "v2.5", "v2.6", "v2.7"]
139+ },
140+ "op_api: torch_npu.ffn_worker_scheduler(*args, **kwargs)": {
141+ "version": ["v2.1", "v2.5", "v2.6", "v2.7"]
142+ },
143+ "op_api: torch_npu.attention_worker_scheduler_(*args: _P.args, **kwargs: _P.kwargs) -> ~_T": {
144+ "version": ["v2.8", "newest"]
145+ },
146+ "op_api: torch_npu.attention_worker_scheduler(*args: _P.args, **kwargs: _P.kwargs) -> ~_T": {
147+ "version": ["v2.8", "newest"]
148+ },
149+ "op_api: torch_npu.ffn_worker_scheduler_(*args: _P.args, **kwargs: _P.kwargs) -> ~_T": {
150+ "version": ["v2.8", "newest"]
151+ },
152+ "op_api: torch_npu.ffn_worker_scheduler(*args: _P.args, **kwargs: _P.kwargs) -> ~_T": {
153+ "version": ["v2.8", "newest"]
154+ },
131 "op_api: torch_npu.npu_top_k_top_p_sample(*args, **kwargs)": {155 "op_api: torch_npu.npu_top_k_top_p_sample(*args, **kwargs)": {
132 "version": ["v2.1", "v2.5", "v2.6", "v2.7"]156 "version": ["v2.1", "v2.5", "v2.6", "v2.7"]
133 },157 },
@@ -451,6 +475,18 @@
451 },475 },
452 "func: npu_top_k_top_p(Tensor logits, Tensor? p=None, Tensor? k=None) -> Tensor": {476 "func: npu_top_k_top_p(Tensor logits, Tensor? p=None, Tensor? k=None) -> Tensor": {
453 "version": ["all_version"]477 "version": ["all_version"]
478+ },
479+ "func: ffn_worker_scheduler_(Tensor(a!) self, *, int sync_group_size=1, int execute_mode=0) -> Tensor(a!)": {
480+ "version": ["all_version"]
481+ },
482+ "func: ffn_worker_scheduler(Tensor self, *, int sync_group_size=1, int execute_mode=0) -> Tensor": {
483+ "version": ["all_version"]
484+ },
485+ "func: attention_worker_scheduler_(Tensor(a!) self) -> Tensor(a!)": {
486+ "version": ["all_version"]
487+ },
488+ "func: attention_worker_scheduler(Tensor self) -> Tensor": {
489+ "version": ["all_version"]
454 },490 },
455 "func: npu_top_k_top_p_sample(Tensor logits, Tensor top_k, Tensor top_p, Tensor? q=None, float eps=1e-8, bool is_need_logits=False, int top_k_guess=32) -> (Tensor, Tensor)": {491 "func: npu_top_k_top_p_sample(Tensor logits, Tensor top_k, Tensor top_p, Tensor? q=None, float eps=1e-8, bool is_need_logits=False, int top_k_guess=32) -> (Tensor, Tensor)": {
456 "version": ["all_version"]492 "version": ["all_version"]
Atest/test_custom_ops/test_attention_worker_schedule.py+56-0
@@ -0,0 +1,56 @@
1+import unittest
2+import torch
3+import torch_npu
4+import os
5+from torch_npu.testing.testcase import TestCase, run_tests
6+from torch_npu.testing.common_utils import SupportedDevices
7+ 
8+window_size = 209715200
9+attn_window_tensor = torch.zeros([window_size], dtype=torch.int8).npu()
10+ 
11+attn_workers = 144
12+micro_batch_number = 3
13+batch_size = 30
14+top_k = 8
15+hidden_size = 7168
16+expert_num = 288
17+attn_to_ffn_token_size = (7168 + 4 + 511) // 512 * 512
18+ffn_to_attn_token_size = 7168 * 2
19+attn_window = attn_window_tensor.data_ptr()
20+ 
21+context_holder = torch_npu._afd.create_schedule_context_holder(schedule_mode = 1, session_num = attn_workers,
22+ micro_batch_num = micro_batch_number, micro_batch_size = batch_size, selected_expert_num = top_k + 1,
23+ expert_num = expert_num, attn_to_ffn_token_size = attn_to_ffn_token_size, ffn_to_attn_token_size = ffn_to_attn_token_size,
24+ attention_window = attn_window, attention_window_size = window_size)
25+
26+schedule_context = context_holder.get_schedule_context_tensor()
27+ 
28+def _set_all_flags():
29+ num_int8 = batch_size * (top_k + 1) * 4 * micro_batch_number
30+ 
31+ int32_view = attn_window_tensor[:num_int8].view(torch.int32)
32+ 
33+ int32_view[:] = 1
34+ 
35+ 
36+class TestAttentionWorkerScheduler(TestCase):
37+ @unittest.skip("skip case until cann supported")
38+ @SupportedDevices(['Ascend910B'])
39+ def test_attention_worker_scheduler_(self):
40+ _set_all_flags()
41+ schedule_context1 = schedule_context.clone()
42+ torch_npu.attention_worker_scheduler_(schedule_context)
43+ self.assertNotEqual(schedule_context1, schedule_context)
44+ 
45+ @unittest.skip("skip case until cann supported")
46+ @SupportedDevices(['Ascend910B'])
47+ def test_attention_worker_scheduler(self):
48+ _set_all_flags()
49+ schedule_context1 = schedule_context.clone()
50+ schedule_context2 = torch_npu.attention_worker_scheduler(schedule_context)
51+ self.assertEqual(schedule_context1, schedule_context)
52+ self.assertNotEqual(schedule_context2, schedule_context)
53+ 
54+ 
55+if __name__ == '__main__':
56+ run_tests()
Atest/test_custom_ops/test_ffn_worker_schedule.py+55-0
@@ -0,0 +1,55 @@
1+import unittest
2+import torch
3+import torch_npu
4+import os
5+from torch_npu.testing.testcase import TestCase, run_tests
6+from torch_npu.testing.common_utils import SupportedDevices
7+ 
8+window_size = 209715200
9+ffn_window_tensor = torch.zeros([window_size], dtype=torch.int8).npu()
10+ 
11+attn_workers = 2
12+micro_batch_number = 3
13+batch_size = 6
14+top_k = 8
15+hidden_size = 7168
16+expert_num = 288
17+attn_to_ffn_token_size = (7168 + 4 + 511) // 512 * 512
18+ffn_to_attn_token_size = 7168 * 2
19+ffn_window = ffn_window_tensor.data_ptr()
20+ 
21+context_holder = torch_npu._afd.create_schedule_context_holder(schedule_mode = 0, session_num = attn_workers,
22+ micro_batch_num = micro_batch_number, micro_batch_size = batch_size, selected_expert_num = top_k + 1,
23+ expert_num = expert_num, attn_to_ffn_token_size = attn_to_ffn_token_size, ffn_to_attn_token_size = ffn_to_attn_token_size,
24+ ffn_window = ffn_window, ffn_window_size = window_size)
25+
26+schedule_context = context_holder.get_schedule_context_tensor()
27+ 
28+def _set_all_flags():
29+ num_int8 = attn_workers * micro_batch_number * (8 + batch_size * top_k * 4)
30+ per_session_num = micro_batch_number * (8 + batch_size * top_k * 4)
31+ int32_view = ffn_window_tensor[:num_int8].view(torch.int32)
32+ int32_view[:] = 1
33+ 
34+ 
35+class TestFfnWorkerScheduler(TestCase):
36+ @unittest.skip("skip case until cann supported")
37+ @SupportedDevices(['Ascend910B'])
38+ def test_ffn_worker_scheduler_(self):
39+ _set_all_flags()
40+ schedule_context1 = schedule_context.clone()
41+ torch_npu.ffn_worker_scheduler_(schedule_context, sync_group_size = 2)
42+ self.assertNotEqual(schedule_context1, schedule_context)
43+ 
44+ @unittest.skip("skip case until cann supported")
45+ @SupportedDevices(['Ascend910B'])
46+ def test_ffn_worker_scheduler(self):
47+ _set_all_flags()
48+ schedule_context1 = schedule_context.clone()
49+ schedule_context2 = torch_npu.ffn_worker_scheduler(schedule_context, sync_group_size = 2)
50+ self.assertEqual(schedule_context1, schedule_context)
51+ self.assertNotEqual(schedule_context2, schedule_context)
52+ 
53+ 
54+if __name__ == '__main__':
55+ run_tests()