Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / multicore / __init__.py: 0%

1# Copyright 2026 Huawei Technologies Co., Ltd 

2# 

3# Licensed under the Apache License, Version 2.0 (the "License"); 

4# you may not use this file except in compliance with the License. 

5# You may obtain a copy of the License at 

6# 

7# http://www.apache.org/licenses/LICENSE-2.0 

8# 

9# Unless required by applicable law or agreed to in writing, software 

10# distributed under the License is distributed on an "AS IS" BASIS, 

11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 

12# See the License for the specific language governing permissions and 

13# limitations under the License. 

14# ============================================================================ 

15"""MoE-FFN multicore operator entry point with platform dispatch.""" 

16from hyper_parallel.platform import get_platform 

17 

18_platform = get_platform() 

19_multicore_handler = _platform.get_multicore_handler() 

20 

21 

22def moe_ffn_fwd( 

23 dispatch_target, dispatch_target_off, 

24 dispatch_src, dispatch_src_off, dispatch_size, 

25 up_proj_weight, up_proj_glist, 

26 up_proj_y, swiglu_out, 

27 down_proj_weight, down_proj_glist, down_proj_y, 

28 combine_target, combine_target_off, combine_src_off, combine_size, 

29 gmm_workspace, up_proj_tiling, swiglu_tiling, down_proj_tiling, 

30 runtime_config, all_event_counters, 

31 rank_id: int, ep: int, expert_num: int, 

32 hidden_size: int, seq_size: int, 

33): 

34 """MoE-FFN forward operator (platform-dispatched).""" 

35 return _multicore_handler.moe_ffn_fwd( 

36 dispatch_target, dispatch_target_off, 

37 dispatch_src, dispatch_src_off, dispatch_size, 

38 up_proj_weight, up_proj_glist, 

39 up_proj_y, swiglu_out, 

40 down_proj_weight, down_proj_glist, down_proj_y, 

41 combine_target, combine_target_off, combine_src_off, combine_size, 

42 gmm_workspace, up_proj_tiling, swiglu_tiling, down_proj_tiling, 

43 runtime_config, all_event_counters, 

44 rank_id, ep, expert_num, hidden_size, seq_size, 

45 ) 

46 

47 

48def moe_ffn_bwd( 

49 dispatch_target, dispatch_target_off, 

50 dy, dispatch_src_off, dispatch_size, 

51 hidden, hidden_dw, 

52 w2, act_grad_y, gate, grad_gate, w1, gate_dx, grad_x, 

53 combine_target_off, combine_src_off, combine_size, 

54 permute_out, gate_dw, group_list, 

55 act_grad_tiling, gate_grad_tiling, w2_grad_tiling, w1_grad_tiling, 

56 swiglu_grad_tiling, gmm_workspace, swiglu_grad_workspace, 

57 runtime_config, all_event_counters, 

58 rank_id: int, ep: int, expert_num: int, 

59 hidden_size: int, seq_size: int, 

60): 

61 """MoE-FFN backward operator (platform-dispatched).""" 

62 return _multicore_handler.moe_ffn_bwd( 

63 dispatch_target, dispatch_target_off, 

64 dy, dispatch_src_off, dispatch_size, 

65 hidden, hidden_dw, 

66 w2, act_grad_y, gate, grad_gate, w1, gate_dx, grad_x, 

67 combine_target_off, combine_src_off, combine_size, 

68 permute_out, gate_dw, group_list, 

69 act_grad_tiling, gate_grad_tiling, w2_grad_tiling, w1_grad_tiling, 

70 swiglu_grad_tiling, gmm_workspace, swiglu_grad_workspace, 

71 runtime_config, all_event_counters, 

72 rank_id, ep, expert_num, hidden_size, seq_size, 

73 ) 

74 

75 

76__all__ = ["moe_ffn_fwd", "moe_ffn_bwd"]