Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_multicore / modules / mega_moe / backward / storage.py: 94%

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"""Check the physical backward schedule before sharing transient storage.""" 

16 

17from collections import defaultdict 

18 

19from hyper_multicore.scheduler.config import RuntimeConfigC, TaskDescC, TaskType 

20 

21 

22_DISPATCH = 0 

23_WEIGHT2_GRAD = 6 

24_ACTIVATION_GRAD = 8 

25_SWIGLU_GRAD = 10 

26_INPUT_GRAD = 12 

27 

28 

29def _worker_read_order( 

30 tasks: list[TaskDescC], num_experts: int, num_cores: int, 

31) -> bool: 

32 """Require every Cube to finish both receive readers before writing dX.""" 

33 for worker in range(num_cores): 

34 locations = {} 

35 for index, task in enumerate(tasks[worker::num_cores]): 

36 output = task.outputs[0].input_position 

37 if output not in (_WEIGHT2_GRAD, _ACTIVATION_GRAD, _INPUT_GRAD): 

38 continue 

39 key = (output, task.task_index // num_cores) 

40 if key in locations or task.task_type != TaskType.TASK_GROUPED_MATMUL: 

41 return False 

42 locations[key] = index 

43 if len(locations) != 3 * num_experts: 

44 return False 

45 for expert in range(num_experts): 

46 order = [locations.get((output, expert), -1) 

47 for output in (_WEIGHT2_GRAD, _ACTIVATION_GRAD, _INPUT_GRAD)] 

48 if not 0 <= order[0] < order[1] < order[2]: 

49 return False 

50 return True 

51 

52 

53def _expert_events_join_readers( 

54 cfg: RuntimeConfigC, cube_tasks: list[TaskDescC], vector_tasks: list[TaskDescC], 

55 num_experts: int, num_cores: int, 

56) -> bool: 

57 """Require the existing event chain to join all readers of each expert.""" 

58 acts = defaultdict(list) 

59 gates = defaultdict(list) 

60 swiglu = defaultdict(list) 

61 for task in cube_tasks: 

62 expert = task.task_index // num_cores 

63 if task.outputs[0].input_position == _ACTIVATION_GRAD: 

64 acts[expert].append(task) 

65 elif task.outputs[0].input_position == _INPUT_GRAD: 

66 gates[expert].append(task) 

67 for task in vector_tasks: 

68 if task.task_type != TaskType.TASK_SWI_GLU_GRAD: 

69 continue 

70 if task.task_split_num < num_experts or task.task_split_num % num_experts: 

71 return False 

72 expert = task.task_index // (task.task_split_num // num_experts) 

73 swiglu[expert].append(task) 

74 producers = defaultdict(list) 

75 for task in cube_tasks + vector_tasks: 

76 producers[task.trigger_event].append(task) 

77 return all(_expert_dependency_chain(cfg, acts[expert], gates[expert], swiglu[expert], num_cores, producers) 

78 for expert in range(num_experts)) 

79 

80 

81def _event_joins_tasks(cfg: RuntimeConfigC, tasks: list[TaskDescC], producers: dict) -> bool: 

82 """Require a completion event to include exactly these nonempty tasks.""" 

83 event = tasks[0].trigger_event 

84 return (len(producers[event]) == len(tasks) 

85 and cfg.all_event_num_triggers[event] == len(tasks) 

86 and all(task.trigger_event == event for task in tasks)) 

87 

88 

89def _expert_dependency_chain( 

90 cfg: RuntimeConfigC, acts: list[TaskDescC], gates: list[TaskDescC], 

91 swiglu: list[TaskDescC], num_cores: int, producers: dict, 

92) -> bool: 

93 """Join every expert reader before any task overwrites its receive rows.""" 

94 if len(acts) != num_cores or len(gates) != num_cores or not swiglu: 

95 return False 

96 if not _event_joins_tasks(cfg, acts, producers) or not _event_joins_tasks(cfg, swiglu, producers): 

97 return False 

98 act_event = acts[0].trigger_event 

99 swiglu_event = swiglu[0].trigger_event 

100 return (all(task.dependent_event == act_event and task.outputs[0].input_position == _SWIGLU_GRAD 

101 for task in swiglu) 

102 and all(task.dependent_event == swiglu_event for task in gates)) 

103 

104 

105def can_reuse_backward_dispatch(cfg: RuntimeConfigC, num_experts: int, num_cores: int) -> bool: 

106 """Check a validated runtime before reusing received dY storage for dX. 

107 

108 Args: 

109 cfg: Validated, fully reordered backward runtime descriptors. 

110 num_experts: Number of local experts. 

111 num_cores: Number of active Cube workers. 

112 

113 Returns: 

114 Whether every receive reader precedes dX through worker order and 

115 the existing full-expert completion events. Unknown readers or changed 

116 schedules disable reuse without adding dependencies to the hot path. 

117 """ 

118 if cfg.task_index_num[2]: 

119 return False 

120 cube_tasks = [cfg.all_tasks[index] for index in cfg.cube_task_indices[:cfg.task_index_num[0]]] 

121 vector_tasks = [cfg.all_tasks[index] for index in cfg.vector_task_indices[:cfg.task_index_num[1]]] 

122 for task in cube_tasks + vector_tasks: 

123 reads_dispatch = any(task.inputs[index].input_position == _DISPATCH for index in range(task.num_inputs)) 

124 if reads_dispatch and (task.task_type != TaskType.TASK_GROUPED_MATMUL 

125 or task.outputs[0].input_position not in (_WEIGHT2_GRAD, _ACTIVATION_GRAD)): 

126 return False 

127 return (_worker_read_order(cube_tasks, num_experts, num_cores) 

128 and _expert_events_join_readers(cfg, cube_tasks, vector_tasks, num_experts, num_cores))