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

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"""Torch projection and autograd bridge for native HyperMegaGate Route.""" 

16 

17from __future__ import annotations 

18 

19from types import ModuleType 

20 

21import torch 

22from torch.autograd.function import once_differentiable 

23from torch.nn import functional 

24 

25 

26from .plan import MegaGatePlan 

27 

28 

29_NATIVE_OPS: ModuleType | None = None 

30 

31 

32def _native_ops() -> ModuleType: 

33 """Load the optional torch_npu adapter only when Device work is requested.""" 

34 global _NATIVE_OPS # pylint: disable=global-statement 

35 if _NATIVE_OPS is None: 

36 from hyper_multicore.torch import ops # pylint: disable=import-outside-toplevel 

37 

38 _NATIVE_OPS = ops 

39 return _NATIVE_OPS 

40 

41 

42def _launch_route( 

43 logits: torch.Tensor, 

44 text_bias: torch.Tensor, 

45 vision_bias: torch.Tensor, 

46 image_mask: torch.Tensor, 

47 use_vision_bias: bool, 

48 plan: MegaGatePlan, 

49 top_k: int, 

50 routed_scaling_factor: float, 

51) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: 

52 """Launch the eventless Route pipeline on FP32 logits.""" 

53 runtime = plan.route_runtime 

54 options = {} 

55 if runtime.is_active: 

56 options["diagnostics"] = runtime 

57 return _native_ops()._mega_gate_route( # pylint: disable=protected-access 

58 logits, text_bias, vision_bias, image_mask, runtime.runtime_config, runtime.runtime_config, 

59 top_k=top_k, routed_scaling_factor=routed_scaling_factor, use_vision_bias=use_vision_bias, **options, 

60 ) 

61 

62 

63def _launch_route_grad( 

64 plan: MegaGatePlan, 

65 logits: torch.Tensor, 

66 route_scores: torch.Tensor, 

67 selected_scores: torch.Tensor, 

68 normalization_denominator: torch.Tensor, 

69 expert_indices: torch.Tensor, 

70 grad_routing_weights: torch.Tensor, 

71 *, 

72 top_k: int, 

73 routed_scaling_factor: float, 

74) -> torch.Tensor: 

75 """Launch RouteGrad without combining the direct logits gradient.""" 

76 runtime = plan.route_grad_k1_runtime if top_k == 1 else plan.route_grad_runtime 

77 inputs = (logits, route_scores, selected_scores, normalization_denominator, expert_indices, grad_routing_weights) 

78 if not runtime.is_active: 

79 return _native_ops()._mega_gate_route_grad( # pylint: disable=protected-access 

80 *inputs, runtime.runtime_config, runtime.runtime_config, 

81 top_k=top_k, routed_scaling_factor=routed_scaling_factor, 

82 ) 

83 # Positions 0..6 mirror the internal RouteGrad MegaKernel ABI. The zero-task 

84 # ACLNN tail node only uses Host positions 9/10, never Device workspace pointers. 

85 parameters = [selected_scores, normalization_denominator, grad_routing_weights, route_scores, expert_indices, 

86 runtime.runtime_config, runtime.runtime_config, None, None, None, logits] 

87 with runtime.launch(parameters, direction="backward") as call: 

88 result = _native_ops()._mega_gate_route_grad( # pylint: disable=protected-access 

89 *inputs, call.runtime_config, call.diagnostic_buffer, 

90 top_k=top_k, routed_scaling_factor=routed_scaling_factor, 

91 ) 

92 parameters[9] = result 

93 return result 

94 

95 

96class _MegaGateRouteFunction(torch.autograd.Function): # pylint: disable=abstract-method,arguments-differ 

97 """Differentiate only the native Route pipeline with respect to logits.""" 

98 

99 @staticmethod 

100 def forward( 

101 ctx: torch.autograd.function.FunctionCtx, 

102 logits: torch.Tensor, 

103 text_bias: torch.Tensor, 

104 vision_bias: torch.Tensor, 

105 image_mask: torch.Tensor, 

106 use_vision_bias: bool, 

107 plan: MegaGatePlan, 

108 top_k: int, 

109 routed_scaling_factor: float, 

110 ) -> tuple[torch.Tensor, torch.Tensor]: 

111 """Run Route while retaining only the state required by RouteGrad.""" 

112 ctx.set_materialize_grads(False) 

113 routing_weights, expert_indices, route_scores, selected_scores, normalization_denominator = _launch_route( 

114 logits, 

115 text_bias.detach(), 

116 vision_bias.detach(), 

117 image_mask, 

118 use_vision_bias, 

119 plan, 

120 top_k, 

121 routed_scaling_factor, 

122 ) 

123 ctx.save_for_backward( 

124 logits, 

125 expert_indices, 

126 route_scores, 

127 selected_scores, 

128 normalization_denominator, 

129 ) 

130 ctx.top_k = top_k 

131 ctx.routed_scaling_factor = routed_scaling_factor 

132 ctx.plan = plan 

133 ctx.mark_non_differentiable(expert_indices) 

134 return routing_weights, expert_indices 

135 

136 @staticmethod 

137 @once_differentiable 

138 def backward( 

139 ctx: torch.autograd.function.FunctionCtx, 

140 grad_routing_weights: torch.Tensor | None, 

141 grad_expert_indices: torch.Tensor | None, 

142 ) -> tuple[torch.Tensor | None, None, None, None, None, None, None, None]: 

143 """Return the Route contribution to the FP32 logits gradient.""" 

144 del grad_expert_indices 

145 if grad_routing_weights is None or not ctx.needs_input_grad[0]: 

146 return None, None, None, None, None, None, None, None 

147 

148 logits, expert_indices, route_scores, selected_scores, normalization_denominator = ctx.saved_tensors 

149 grad_logits = _launch_route_grad( 

150 ctx.plan, 

151 logits, 

152 route_scores, 

153 selected_scores, 

154 normalization_denominator, 

155 expert_indices, 

156 grad_routing_weights.contiguous(), 

157 top_k=ctx.top_k, 

158 routed_scaling_factor=ctx.routed_scaling_factor, 

159 ) 

160 return grad_logits, None, None, None, None, None, None, None 

161 

162 

163def mega_gate( 

164 hidden_states: torch.Tensor, 

165 weight: torch.Tensor, 

166 correction_bias: torch.Tensor, 

167 plan: MegaGatePlan, 

168 *, 

169 vision_bias: torch.Tensor | None = None, 

170 image_mask: torch.Tensor | None = None, 

171 top_k: int, 

172 routed_scaling_factor: float, 

173) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: 

174 """Apply Torch projection followed by the native HyperMegaGate Route.""" 

175 flattened = hidden_states.reshape(-1, hidden_states.shape[-1]) 

176 logits = functional.linear( # pylint: disable=not-callable 

177 flattened.float(), weight.float() 

178 ) 

179 route_logits = logits.float() 

180 use_vision_bias = vision_bias is not None and image_mask is not None 

181 # FSDP may present BF16 bias parameters; golden promotes them when adding FP32 scores. 

182 correction_bias = correction_bias.detach().float() 

183 resolved_vision_bias = vision_bias.detach().float() if use_vision_bias else correction_bias 

184 resolved_image_mask = plan.vision_mask_placeholder if not use_vision_bias else image_mask 

185 if not torch.is_grad_enabled() or not logits.requires_grad: 

186 routing_weights, expert_indices = _launch_route( 

187 route_logits, 

188 correction_bias, 

189 resolved_vision_bias, 

190 resolved_image_mask, 

191 use_vision_bias, 

192 plan, 

193 top_k, 

194 routed_scaling_factor, 

195 )[:2] 

196 return logits, routing_weights, expert_indices 

197 routing_weights, expert_indices = _MegaGateRouteFunction.apply( 

198 route_logits, 

199 correction_bias, 

200 resolved_vision_bias, 

201 resolved_image_mask, 

202 use_vision_bias, 

203 plan, 

204 top_k, 

205 routed_scaling_factor, 

206 ) 

207 return logits, routing_weights, expert_indices