Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_multicore / modules / mega_gate / diagnostics.py: 96%

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"""MegaGate broadcast task identities and native GM operand bindings.""" 

16 

17from typing import Any 

18 

19from hyper_multicore._core_spec import get_device_core_spec 

20from hyper_multicore.diagnostics.common.graph import bind_graph_tensor_ports 

21from hyper_multicore.diagnostics.dump.l1 import L1TensorView, L1TensorViewContext 

22from hyper_multicore.scheduler.graph import ComputeGraph, OperatorNode, OpType, SplitSpec, TensorSpec 

23 

24 

25MEGA_GATE_OWNER_LABEL = "TokenRowShard" 

26MEGA_GATE_ROUTE_STAGE_NAMES = ( 

27 "Softplus", 

28 "Sqrt", 

29 "AddBias", 

30 "TopK", 

31 "Gather", 

32 "ReduceSum", 

33 "AddEpsilon", 

34 "Div", 

35 "MulScale", 

36 "CastIndex", 

37) 

38 

39MEGA_GATE_ROUTE_GRAD_STAGE_NAMES = ( 

40 "MulsScaleGrad", 

41 "BroadcastDenominator", 

42 "Neg", 

43 "DivSelected", 

44 "DivSelectedRatio", 

45 "MulCross", 

46 "DivDirect", 

47 "ReduceSum", 

48 "BroadcastRowSum", 

49 "AddSelectedGrad", 

50 "ZerosLike", 

51) 

52 

53MEGA_GATE_ROUTE_GRAD_K1_STAGE_NAMES = ( 

54 "MulsScaleGrad", 

55 "ZerosLike", 

56) 

57 

58 

59 

60def _port(name: str, position: int, dtype_size: int = 4) -> TensorSpec: 

61 """Bind one native GM operand; shapes are invocation-owned.""" 

62 return TensorSpec(name, [], position, dtype_size=dtype_size, is_dynamic=True) 

63 

64 

65def _row_view(context: L1TensorViewContext) -> L1TensorView: 

66 """Mirror Route/RouteGrad tiling without running a Device-side resolver.""" 

67 tensor = context.parameters[context.tensor_spec.param_position] 

68 shape = tuple(tensor.shape) 

69 tokens = int(context.parameters[0].shape[0]) 

70 if context.tensor_spec.name in ("text_bias", "vision_bias") or ( 

71 context.tensor_spec.name == "image_mask" and tensor.numel() == 1 

72 ): 

73 return context.make_view(shape, source_offset_bytes=0) 

74 cache_key = -1 

75 if cache_key not in context.parameter_cache: 

76 cores = int(context.runtime_config.num_workers) 

77 if tensor.device.type == "npu": 

78 cores = min(cores, get_device_core_spec(tensor.device.index).num_vector_cores) 

79 workers = min(tokens, cores) 

80 context.parameter_cache[cache_key] = (tokens + workers - 1) // workers 

81 rows_per_worker = context.parameter_cache[cache_key] 

82 offset = context.core_id * rows_per_worker 

83 rows = max(0, min(rows_per_worker, tokens - offset)) 

84 if rows == 0: 

85 return context.make_view((0, *shape[1:]), source_offset_bytes=0) 

86 row_elements = tensor.numel() // tokens 

87 return context.make_view((rows, *shape[1:]), 

88 source_offset_bytes=offset * row_elements * tensor.element_size()) 

89 

90 

91def bind_pipeline_graph(graph: Any, config: Any, *, backward: bool = False) -> ComputeGraph: 

92 """Build diagnostic ports only; native stages still use their original ABI. 

93 

94 RouteGrad's scatter/sqrt/softplus tail is launched by ACLNN after the MegaKernel. 

95 A zero-task boundary node describes its L0 result, not an L1 Device capture. 

96 Internal CANN workspace arrays are deliberately not advertised as ABI Tensors. 

97 

98 Args: 

99 graph: Graph placeholder replaced by the lazily constructed broadcast graph. 

100 config: Configuration consumed by this callback. 

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

102 """ 

103 p = _port 

104 if backward: 

105 names = MEGA_GATE_ROUTE_GRAD_K1_STAGE_NAMES if config.task_num == 2 else MEGA_GATE_ROUTE_GRAD_STAGE_NAMES 

106 bindings = {name: ([], []) for name in names} 

107 bindings["MulsScaleGrad"] = ([p("grad_routing_weights", 2)], []) 

108 if config.task_num != 2: 

109 bindings["BroadcastDenominator"] = ([p("normalization_denominator", 1)], []) 

110 bindings["DivSelected"] = ([p("selected_scores", 0)], []) 

111 bindings["RouteGradPost"] = ( 

112 [p("logits", 10), p("route_scores", 3), p("expert_indices", 4, 8), 

113 p("selected_scores", 0), p("normalization_denominator", 1)], [p("grad_logits", 9)], 

114 ) 

115 names = (*names, "RouteGradPost") 

116 else: 

117 names = MEGA_GATE_ROUTE_STAGE_NAMES 

118 score, selected, denominator = p("route_scores", 8), p("selected_scores", 9), p("normalization_denominator", 10) 

119 bindings = { 

120 "Softplus": ([p("logits", 0)], []), 

121 "Sqrt": ([], [score]), 

122 "AddBias": ([score, p("text_bias", 1), p("vision_bias", 2), p("image_mask", 3, 1)], []), 

123 "TopK": ([], []), "Gather": ([score], [selected]), "ReduceSum": ([selected], []), 

124 "AddEpsilon": ([], [denominator]), "Div": ([selected, denominator], []), 

125 "MulScale": ([], [p("routing_weights", 6)]), "CastIndex": ([], [p("expert_indices", 7, 8)]), 

126 } 

127 graph = ComputeGraph() 

128 previous = None 

129 for index, name in enumerate(names): 

130 node = OperatorNode(name, OpType.AIV_PIPELINE, [], [], [], 0, SplitSpec(None, lambda _: 1), 0, None, 

131 diagnostic_name=name) 

132 node.task_num = 1 if index < config.task_num else 0 

133 graph.add_op(node) 

134 if previous is not None: 

135 graph.add_edge(previous, node) 

136 previous = node 

137 return bind_graph_tensor_ports(graph, bindings, _row_view)