Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_multicore / tasks / alltoall.py: 100%

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"""AllToAll fill config: dispatch and combine (forward + backward).""" 

16from dataclasses import dataclass 

17from enum import Enum 

18 

19from hyper_multicore.scheduler.config import ( 

20 TaskDescC, TensorDescC, RuntimeConfigC, 

21 TaskAiCoreType, TaskType, TaskSplitValue, 

22) 

23from hyper_multicore.scheduler.graph import OperatorNode 

24from hyper_multicore.tasks.task_base import FillConfig 

25from hyper_multicore.tasks.utils import ( 

26 advance_tsv_vector, advance_tsv_vector_only, 

27) 

28 

29 

30class AllToAllType(Enum): 

31 """ 

32 MoE AllToAll semantic type — determines event wiring. 

33 

34 DISPATCH — transfer source tokens into expert storage via configured PUT or GET. 

35 dependent_event = pre_pre_event_num + 0 

36 trigger_event = pre_event_num + global_expert + 1 

37 trigger_count = task_num * ep // all_expert_num 

38 

39 COMBINE — gather expert results back to the originating rank. 

40 dependent_event = pre_pre_event_num + (i // per_g_e_num) % sre + 1 

41 trigger_event = all_event_num 

42 trigger_count = task_num 

43 

44 OTHER — reserved for AllToAll patterns outside the MoE dispatch/combine semantic. 

45 """ 

46 DISPATCH = 1 

47 COMBINE = 2 

48 OTHER = 3 

49 

50 

51@dataclass 

52class AllToAllFillConfig(FillConfig): 

53 """ 

54 AllToAll fill config covering dispatch and combine for fwd and bwd. 

55 

56 moe_type : AllToAllType 

57 Determines event wiring; see AllToAllType for details. 

58 

59 advance : "vector" | "vector_only" 

60 "vector" — advance_tsv_vector(tsv, task_num, event_group_size=event_group) 

61 Used by: dispatch (fwd/bwd), combine forward. 

62 "vector_only" — advance_tsv_vector_only(tsv, task_num) 

63 Only advances pre_task_num/pre_vector_task_num. 

64 Used by: combine backward. 

65 

66 event_group : int 

67 Only effective when advance="vector". 

68 dispatch (fwd/bwd): all_expert_num 

69 combine forward: 1 

70 """ 

71 moe_type: AllToAllType = AllToAllType.DISPATCH 

72 advance: str = "vector" # "vector" | "vector_only" 

73 event_group: int = 1 # only used when advance="vector" 

74 

75 def fill(self, cfg: RuntimeConfigC, op: OperatorNode, tsv: TaskSplitValue) -> None: 

76 """Emit configured dispatch and PUT combine tasks with completion events. 

77 

78 Args: 

79 cfg: Runtime receiving task descriptors and event thresholds. 

80 op: Communication node with fixed tensor argument positions. 

81 tsv: Topology and running graph offsets. 

82 """ 

83 task_num = op.task_num 

84 per_g_e_num = task_num // tsv.all_expert_num 

85 param = op.param_positions 

86 

87 for i in range(task_num): 

88 task = TaskDescC() 

89 task.task_type = ( 

90 TaskType.TASK_SHMEM_GET_MEM 

91 if self.moe_type == AllToAllType.DISPATCH and tsv.dispatch_mode == "pull" 

92 else TaskType.TASK_SHMEM_PUT_MEM_SIGNAL 

93 ) 

94 task.task_aicore_type = TaskAiCoreType.TASK_AICORE_CUBE 

95 task.num_inputs = len(op.inputs) 

96 task.num_outputs = len(op.outputs) 

97 

98 for j, spec in enumerate(op.inputs): 

99 td = TensorDescC() 

100 td.data_type = spec.dtype_size 

101 td.input_position = param[j] 

102 td.base_ptr_offset = 0 

103 

104 if j == 1: # src data: tensor_type and is_dynamic from TensorSpec 

105 td.tensor_type = spec.tensor_type 

106 td.dynamic_shape = int(spec.is_dynamic) 

107 else: # metadata (target_offset / src_offset / size): fixed type=0 

108 td.tensor_type = 0 

109 td.base_ptr_offset = i // per_g_e_num # expert-group index 

110 

111 task.inputs[j] = td 

112 

113 out_spec = op.outputs[0] 

114 out = TensorDescC() 

115 out.tensor_type = out_spec.tensor_type 

116 out.data_type = out_spec.dtype_size 

117 out.input_position = param[task.num_inputs] 

118 out.base_ptr_offset = 0 

119 out.dynamic_shape = int(out_spec.is_dynamic) 

120 task.outputs[0] = out 

121 

122 if self.moe_type == AllToAllType.DISPATCH: 

123 res = (tsv.rank_id * tsv.single_rank_expert_num + (i // per_g_e_num) % tsv.single_rank_expert_num 

124 if tsv.dispatch_mode == "pull" else i // per_g_e_num) 

125 task.dependent_event = tsv.pre_pre_event_num + 0 

126 task.trigger_event = tsv.pre_event_num + res + 1 

127 cfg.all_event_num_triggers[task.trigger_event] = ( 

128 task_num * tsv.ep // tsv.all_expert_num 

129 ) 

130 else: # COMBINE (or OTHER) 

131 current_dep = (i // per_g_e_num) % tsv.single_rank_expert_num 

132 task.dependent_event = tsv.pre_pre_event_num + current_dep + 1 

133 task.trigger_event = tsv.all_event_num 

134 cfg.all_event_num_triggers[task.trigger_event] = task_num 

135 

136 task.task_index = i 

137 task.task_split_num = task_num 

138 task.task_split_value = op.split_value 

139 task.tiling_data_position = 0xFFFFFFFF 

140 

141 cfg.all_tasks[tsv.pre_task_num + i] = task 

142 cfg.vector_task_indices[tsv.pre_vector_task_num + i] = tsv.pre_task_num + i 

143 

144 cfg.task_index_num[1] += task_num 

145 if self.advance == "vector": 

146 advance_tsv_vector(tsv, task_num, event_group_size=self.event_group) 

147 else: # "vector_only" 

148 advance_tsv_vector_only(tsv, task_num)