Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_multicore / modules / mega_gate / plan.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"""Build shared RuntimeConfigs for HyperMegaGate Route and RouteGrad.""" 

16 

17from __future__ import annotations 

18 

19from dataclasses import dataclass 

20from functools import partial 

21 

22import torch 

23import torch.distributed as dist 

24 

25from hyper_multicore.diagnostics import MegaKernelDiagnostics, prepare_mega_kernel_diagnostics 

26from hyper_multicore.scheduler.config import ( 

27 EVENT_INVALID_ID, 

28 INVALID_PROFILE_OWNER_ID, 

29 TaskAiCoreType, 

30 TaskDescC, 

31 TaskType, 

32 RuntimeConfigC, 

33) 

34from hyper_multicore.scheduler.runtime import allocate_runtime_config 

35 

36from .diagnostics import ( 

37 MEGA_GATE_ROUTE_GRAD_K1_STAGE_NAMES, 

38 MEGA_GATE_ROUTE_GRAD_STAGE_NAMES, 

39 MEGA_GATE_ROUTE_STAGE_NAMES, 

40 bind_pipeline_graph, 

41) 

42 

43 

44# AIV slot capacity shared with Route/RouteGrad tiling. 

45MEGA_GATE_AIV_WORKER_CAPACITY = 48 

46MEGA_GATE_ROUTE_TASK_TYPES = ( 

47 TaskType.TASK_GATE_SOFTPLUS, 

48 TaskType.TASK_GATE_SQRT, 

49 TaskType.TASK_GATE_ADD_BIAS, 

50 TaskType.TASK_GATE_TOPK, 

51 TaskType.TASK_GATE_GATHER, 

52 TaskType.TASK_GATE_REDUCE_SUM, 

53 TaskType.TASK_GATE_ADD_EPSILON, 

54 TaskType.TASK_GATE_DIV, 

55 TaskType.TASK_GATE_MUL_SCALE, 

56 TaskType.TASK_GATE_CAST_INDEX, 

57) 

58MEGA_GATE_ROUTE_GRAD_TASK_TYPES = ( 

59 TaskType.TASK_GATE_GRAD_MULS_SCALE, 

60 TaskType.TASK_GATE_GRAD_BROADCAST_DENOMINATOR, 

61 TaskType.TASK_GATE_GRAD_NEG, 

62 TaskType.TASK_GATE_GRAD_DIV_SELECTED, 

63 TaskType.TASK_GATE_GRAD_DIV_SELECTED_RATIO, 

64 TaskType.TASK_GATE_GRAD_MUL_CROSS, 

65 TaskType.TASK_GATE_GRAD_DIV_DIRECT, 

66 TaskType.TASK_GATE_GRAD_REDUCE_SUM, 

67 TaskType.TASK_GATE_GRAD_BROADCAST_ROW_SUM, 

68 TaskType.TASK_GATE_GRAD_ADD_SELECTED, 

69 TaskType.TASK_GATE_GRAD_ZEROS, 

70) 

71MEGA_GATE_ROUTE_GRAD_K1_TASK_TYPES = ( 

72 TaskType.TASK_GATE_GRAD_MULS_SCALE, 

73 TaskType.TASK_GATE_GRAD_ZEROS, 

74) 

75 

76 

77@dataclass(frozen=True) 

78class MegaGatePlan: 

79 """Capacity-bounded RuntimeConfig shared by all input shapes on one NPU.""" 

80 

81 route_runtime: MegaKernelDiagnostics 

82 route_grad_runtime: MegaKernelDiagnostics 

83 route_grad_k1_runtime: MegaKernelDiagnostics 

84 vision_mask_placeholder: torch.Tensor 

85 

86 

87def _build_pipeline_runtime_config( 

88 task_types: tuple[TaskType, ...], 

89 stage_names: tuple[str, ...], 

90 *, 

91 kernel_name: str, 

92) -> RuntimeConfigC: 

93 """Create descriptors broadcast to every active token-row worker.""" 

94 if len(task_types) != len(stage_names): 

95 raise RuntimeError("MegaGate task types and profile stages must have equal length") 

96 if not isinstance(kernel_name, str) or not kernel_name: 

97 raise ValueError("MegaGate kernel_name must be a non-empty string") 

98 runtime_config = allocate_runtime_config(len(task_types)) 

99 runtime_config.task_num = len(task_types) 

100 runtime_config.num_workers = MEGA_GATE_AIV_WORKER_CAPACITY 

101 runtime_config.task_index_num[1] = len(task_types) 

102 for task_id, task_type in enumerate(task_types): 

103 task = TaskDescC() 

104 task.task_type = task_type 

105 task.task_aicore_type = TaskAiCoreType.TASK_AICORE_VECTOR 

106 task.trigger_event = EVENT_INVALID_ID 

107 task.dependent_event = EVENT_INVALID_ID 

108 task.task_index = 0 

109 task.task_split_num = MEGA_GATE_AIV_WORKER_CAPACITY 

110 task.tiling_data_offset = task_id 

111 task.profile_owner_id = INVALID_PROFILE_OWNER_ID 

112 runtime_config.all_tasks[task_id] = task 

113 runtime_config.vector_task_indices[task_id] = task_id 

114 return runtime_config 

115 

116 

117def build_mega_gate_plan(device: torch.device) -> MegaGatePlan: 

118 """Build an eventless Route runtime for one NPU device. 

119 

120 Args: 

121 device: NPU device receiving the serialized RuntimeConfig. 

122 

123 Returns: 

124 Immutable plan with normal and profiled broadcast stage descriptors. 

125 """ 

126 if device.type != "npu": 

127 raise ValueError(f"MegaGate plan requires an NPU device, got {device}") 

128 rank = dist.get_rank() if dist.is_initialized() else 0 

129 

130 def prepare(runtime_config: RuntimeConfigC, *, backward: bool = False) -> MegaKernelDiagnostics: 

131 """Prepare only ordinary storage; broadcast graph ports remain lazy. 

132 

133 Args: 

134 runtime_config: Complete Host configuration describing native task queues. 

135 backward: Whether to prepare the RouteGrad pipeline instead of Route. 

136 """ 

137 return prepare_mega_kernel_diagnostics( 

138 None, runtime_config, kernel_name="HyperMegaGateRouteGrad" if backward else "HyperMegaGateRoute", 

139 parameter_count=11, device=device, rank=rank, owner_label="TokenRowShard", 

140 execution_mode="aiv_pipeline", vector_worker_stride=1, 

141 graph_adapter=partial(bind_pipeline_graph, backward=backward), 

142 ) 

143 

144 return MegaGatePlan( 

145 route_runtime=prepare(_build_pipeline_runtime_config( 

146 MEGA_GATE_ROUTE_TASK_TYPES, 

147 MEGA_GATE_ROUTE_STAGE_NAMES, 

148 kernel_name="HyperMegaGateRoute", 

149 )), 

150 route_grad_runtime=prepare(_build_pipeline_runtime_config( 

151 MEGA_GATE_ROUTE_GRAD_TASK_TYPES, 

152 MEGA_GATE_ROUTE_GRAD_STAGE_NAMES, 

153 kernel_name="HyperMegaGateRouteGrad", 

154 ), backward=True), 

155 route_grad_k1_runtime=prepare(_build_pipeline_runtime_config( 

156 MEGA_GATE_ROUTE_GRAD_K1_TASK_TYPES, 

157 MEGA_GATE_ROUTE_GRAD_K1_STAGE_NAMES, 

158 kernel_name="HyperMegaGateRouteGrad", 

159 ), backward=True), 

160 vision_mask_placeholder=torch.zeros((1,), device=device, dtype=torch.bool), 

161 )