已合并
attention/ffn的aicpu算子提供torch_npu._afd的调用方式 #28865
pwc10490创建于 2025年12月30日
attention/ffn的aicpu算子提供torch_npu._afd的调用方式 #28865
已合并
共 4 个文件变更+226-0
| @@ -0,0 +1,109 @@ | |||
| 1 | +import unittest | ||
| 2 | +import torch | ||
| 3 | +import torch.nn as nn | ||
| 4 | +import torchair | ||
| 5 | +from torchair.configs.compiler_config import CompilerConfig | ||
| 6 | + | ||
| 7 | +import torch_npu | ||
| 8 | +from torch_npu.testing.common_utils import SupportedDevices | ||
| 9 | +from torch_npu.testing.testcase import TestCase, run_tests | ||
| 10 | + | ||
| 11 | +window_size = 209715200 | ||
| 12 | +attn_window_tensor = torch.zeros([window_size], dtype=torch.int8).npu() | ||
| 13 | + | ||
| 14 | +attn_workers = 144 | ||
| 15 | +micro_batch_number = 3 | ||
| 16 | +batch_size = 30 | ||
| 17 | +top_k = 8 | ||
| 18 | +hidden_size = 7168 | ||
| 19 | +expert_num = 288 | ||
| 20 | +attn_to_ffn_token_size = (7168 + 4 + 511) // 512 * 512 | ||
| 21 | +ffn_to_attn_token_size = 7168 * 2 | ||
| 22 | +attn_window = attn_window_tensor.data_ptr() | ||
| 23 | + | ||
| 24 | + | ||
| 25 | +def _set_all_flags(): | ||
| 26 | + num_int8 = batch_size * (top_k + 1) * 4 * micro_batch_number | ||
| 27 | + | ||
| 28 | + int32_view = attn_window_tensor[:num_int8].view(torch.int32) | ||
| 29 | + | ||
| 30 | + int32_view[:] = 1 | ||
| 31 | + | ||
| 32 | + | ||
| 33 | +class TestModelInplace(nn.Module): | ||
| 34 | + def __init__(self): | ||
| 35 | + super().__init__() | ||
| 36 | + | ||
| 37 | + def forward(self, schedule_context): | ||
| 38 | + torch_npu._afd.attention_worker_scheduler_(schedule_context) | ||
| 39 | + | ||
| 40 | + | ||
| 41 | +class TestModel(nn.Module): | ||
| 42 | + def __init__(self): | ||
| 43 | + super().__init__() | ||
| 44 | + | ||
| 45 | + def forward(self, schedule_context): | ||
| 46 | + return torch_npu._afd.attention_worker_scheduler(schedule_context) | ||
| 47 | + | ||
| 48 | + | ||
| 49 | +class TestAttentionWorkerScheduler(TestCase): | ||
| 50 | + def setUp(self): | ||
| 51 | + self.context_holder = torch_npu._afd.create_schedule_context_holder(schedule_mode=1, session_num=attn_workers, | ||
| 52 | + micro_batch_num=micro_batch_number, | ||
| 53 | + micro_batch_size=batch_size, | ||
| 54 | + selected_expert_num=top_k + 1, | ||
| 55 | + expert_num=expert_num, | ||
| 56 | + attn_to_ffn_token_size=attn_to_ffn_token_size, | ||
| 57 | + ffn_to_attn_token_size=ffn_to_attn_token_size, | ||
| 58 | + attention_window=attn_window, | ||
| 59 | + attention_window_size=window_size) | ||
| 60 | + | ||
| 61 | + self.schedule_context = self.context_holder.get_schedule_context_tensor() | ||
| 62 | + _set_all_flags() | ||
| 63 | + | ||
| 64 | + | ||
| 65 | + | ||
| 66 | + def test_attention_worker_scheduler_(self): | ||
| 67 | + schedule_context1 = self.schedule_context.clone() | ||
| 68 | + torch_npu._afd.attention_worker_scheduler_(self.schedule_context) | ||
| 69 | + self.assertNotEqual(schedule_context1, self.schedule_context) | ||
| 70 | + | ||
| 71 | + | ||
| 72 | + | ||
| 73 | + def test_attention_worker_scheduler(self): | ||
| 74 | + _set_all_flags() | ||
| 75 | + schedule_context1 = self.schedule_context.clone() | ||
| 76 | + schedule_context2 = torch_npu._afd.attention_worker_scheduler(self.schedule_context) | ||
| 77 | + self.assertEqual(schedule_context1, self.schedule_context) | ||
| 78 | + self.assertNotEqual(schedule_context2, self.schedule_context) | ||
| 79 | + | ||
| 80 | + | ||
| 81 | + | ||
| 82 | + def test_attention_worker_scheduler__graph(self): | ||
| 83 | + _set_all_flags() | ||
| 84 | + config = CompilerConfig() | ||
| 85 | + npu_backend = torchair.get_npu_backend(compiler_config=config) | ||
| 86 | + model = TestModelInplace().npu() | ||
| 87 | + model = torch.compile(model, backend=npu_backend) | ||
| 88 | + schedule_context1 = self.schedule_context.clone() | ||
| 89 | + model(self.schedule_context) | ||
| 90 | + self.assertNotEqual(schedule_context1, self.schedule_context) | ||
| 91 | + torch._dynamo.reset() | ||
| 92 | + | ||
| 93 | + | ||
| 94 | + | ||
| 95 | + def test_attention_worker_scheduler_graph(self): | ||
| 96 | + _set_all_flags() | ||
| 97 | + config = CompilerConfig() | ||
| 98 | + npu_backend = torchair.get_npu_backend(compiler_config=config) | ||
| 99 | + model = TestModel().npu() | ||
| 100 | + model = torch.compile(model, backend=npu_backend) | ||
| 101 | + schedule_context1 = self.schedule_context.clone() | ||
| 102 | + schedule_context2 = model(self.schedule_context) | ||
| 103 | + self.assertEqual(schedule_context1, self.schedule_context) | ||
| 104 | + self.assertNotEqual(schedule_context2, self.schedule_context) | ||
| 105 | + torch._dynamo.reset() | ||
| 106 | + | ||
| 107 | + | ||
| 108 | +if __name__ == '__main__': | ||
| 109 | + run_tests() | ||
| @@ -0,0 +1,108 @@ | |||
| 1 | +import unittest | ||
| 2 | +import torch | ||
| 3 | +import torch.nn as nn | ||
| 4 | +import torchair | ||
| 5 | +from torchair.configs.compiler_config import CompilerConfig | ||
| 6 | + | ||
| 7 | +import torch_npu | ||
| 8 | +from torch_npu.testing.common_utils import SupportedDevices | ||
| 9 | +from torch_npu.testing.testcase import TestCase, run_tests | ||
| 10 | + | ||
| 11 | +window_size = 209715200 | ||
| 12 | +ffn_window_tensor = torch.zeros([window_size], dtype=torch.int8).npu() | ||
| 13 | + | ||
| 14 | +attn_workers = 2 | ||
| 15 | +micro_batch_number = 3 | ||
| 16 | +batch_size = 6 | ||
| 17 | +top_k = 8 | ||
| 18 | +hidden_size = 7168 | ||
| 19 | +expert_num = 288 | ||
| 20 | +attn_to_ffn_token_size = (7168 + 4 + 511) // 512 * 512 | ||
| 21 | +ffn_to_attn_token_size = 7168 * 2 | ||
| 22 | +ffn_window = ffn_window_tensor.data_ptr() | ||
| 23 | + | ||
| 24 | + | ||
| 25 | +def _set_all_flags(): | ||
| 26 | + num_int8 = attn_workers * micro_batch_number * (8 + batch_size * top_k * 4) | ||
| 27 | + int32_view = ffn_window_tensor[:num_int8].view(torch.int32) | ||
| 28 | + int32_view[:] = 1 | ||
| 29 | + | ||
| 30 | + | ||
| 31 | +class TestModelInplace(nn.Module): | ||
| 32 | + def __init__(self): | ||
| 33 | + super().__init__() | ||
| 34 | + | ||
| 35 | + def forward(self, schedule_context): | ||
| 36 | + torch_npu._afd.ffn_worker_scheduler_(schedule_context, sync_group_size=1, execute_mode=0) | ||
| 37 | + | ||
| 38 | + | ||
| 39 | +class TestModel(nn.Module): | ||
| 40 | + def __init__(self): | ||
| 41 | + super().__init__() | ||
| 42 | + | ||
| 43 | + def forward(self, schedule_context): | ||
| 44 | + return torch_npu._afd.ffn_worker_scheduler(schedule_context, sync_group_size=1) | ||
| 45 | + | ||
| 46 | + | ||
| 47 | +class TestFfnWorkerScheduler(TestCase): | ||
| 48 | + def setUp(self): | ||
| 49 | + self.context_holder = torch_npu._afd.create_schedule_context_holder(schedule_mode=0, session_num=attn_workers, | ||
| 50 | + micro_batch_num=micro_batch_number, | ||
| 51 | + micro_batch_size=batch_size, | ||
| 52 | + selected_expert_num=top_k + 1, | ||
| 53 | + expert_num=expert_num, | ||
| 54 | + attn_to_ffn_token_size=attn_to_ffn_token_size, | ||
| 55 | + ffn_to_attn_token_size=ffn_to_attn_token_size, | ||
| 56 | + ffn_window=ffn_window, | ||
| 57 | + ffn_window_size=window_size) | ||
| 58 | + | ||
| 59 | + self.schedule_context = self.context_holder.get_schedule_context_tensor() | ||
| 60 | + _set_all_flags() | ||
| 61 | + | ||
| 62 | + | ||
| 63 | + | ||
| 64 | + def test_ffn_worker_scheduler_(self): | ||
| 65 | + _set_all_flags() | ||
| 66 | + schedule_context1 = self.schedule_context.clone() | ||
| 67 | + torch_npu._afd.ffn_worker_scheduler_(self.schedule_context, sync_group_size=2) | ||
| 68 | + self.assertNotEqual(schedule_context1, self.schedule_context) | ||
| 69 | + | ||
| 70 | + | ||
| 71 | + | ||
| 72 | + def test_ffn_worker_scheduler(self): | ||
| 73 | + _set_all_flags() | ||
| 74 | + schedule_context1 = self.schedule_context.clone() | ||
| 75 | + schedule_context2 = torch_npu._afd.ffn_worker_scheduler(self.schedule_context, sync_group_size=2) | ||
| 76 | + self.assertEqual(schedule_context1, self.schedule_context) | ||
| 77 | + self.assertNotEqual(schedule_context2, self.schedule_context) | ||
| 78 | + | ||
| 79 | + | ||
| 80 | + | ||
| 81 | + def test_ffn_worker_scheduler__graph(self): | ||
| 82 | + _set_all_flags() | ||
| 83 | + config = CompilerConfig() | ||
| 84 | + npu_backend = torchair.get_npu_backend(compiler_config=config) | ||
| 85 | + model = TestModelInplace().npu() | ||
| 86 | + model = torch.compile(model, backend=npu_backend) | ||
| 87 | + schedule_context1 = self.schedule_context.clone() | ||
| 88 | + model(self.schedule_context) | ||
| 89 | + self.assertNotEqual(schedule_context1, self.schedule_context) | ||
| 90 | + torch._dynamo.reset() | ||
| 91 | + | ||
| 92 | + | ||
| 93 | + | ||
| 94 | + def test_ffn_worker_scheduler_graph(self): | ||
| 95 | + _set_all_flags() | ||
| 96 | + config = CompilerConfig() | ||
| 97 | + npu_backend = torchair.get_npu_backend(compiler_config=config) | ||
| 98 | + model = TestModel().npu() | ||
| 99 | + model = torch.compile(model, backend=npu_backend) | ||
| 100 | + schedule_context1 = self.schedule_context.clone() | ||
| 101 | + schedule_context2 = model(self.schedule_context) | ||
| 102 | + self.assertEqual(schedule_context1, self.schedule_context) | ||
| 103 | + self.assertNotEqual(schedule_context2, self.schedule_context) | ||
| 104 | + torch._dynamo.reset() | ||
| 105 | + | ||
| 106 | + | ||
| 107 | +if __name__ == '__main__': | ||
| 108 | + run_tests() | ||
| @@ -76,6 +76,7 @@ from torch_npu.utils import _apply_module_patch, _add_tensor_methods, _add_colle | |||
| 76 | _apply_npu_show_warning, _apply_npugraph_tree_methods, _apply_dlpack_patch | 76 | _apply_npu_show_warning, _apply_npugraph_tree_methods, _apply_dlpack_patch |
| 77 | from torch_npu.utils._dynamo_device import _dynamo_register_interface_for_device | 77 | from torch_npu.utils._dynamo_device import _dynamo_register_interface_for_device |
| 78 | from torch_npu.npu._format import _apply_npu_format_patch | 78 | from torch_npu.npu._format import _apply_npu_format_patch |
| 79 | +import torch_npu.utils._afd_ops | ||
| 79 | import torch_npu.utils.custom_ops | 80 | import torch_npu.utils.custom_ops |
| 80 | import torch_npu.distributed.rpc | 81 | import torch_npu.distributed.rpc |
| 81 | import torch_npu.op_plugin | 82 | import torch_npu.op_plugin |
| @@ -0,0 +1,8 @@ | |||
| 1 | +import torch | ||
| 2 | +import torch_npu | ||
| 3 | + | ||
| 4 | + | ||
| 5 | +torch_npu._afd.attention_worker_scheduler_ = torch.ops.npu.attention_worker_scheduler_ | ||
| 6 | +torch_npu._afd.attention_worker_scheduler = torch.ops.npu.attention_worker_scheduler | ||
| 7 | +torch_npu._afd.ffn_worker_scheduler_ = torch.ops.npu.ffn_worker_scheduler_ | ||
| 8 | +torch_npu._afd.ffn_worker_scheduler = torch.ops.npu.ffn_worker_scheduler | ||