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

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"""Functional entry point for the BF16 HyperMegaMhcGrad mega kernel.""" 

16 

17from __future__ import annotations 

18 

19from functools import lru_cache 

20 

21import torch 

22 

23from hyper_multicore._core_spec import get_device_core_spec 

24from hyper_multicore.torch import ops as multicore_ops 

25 

26from ..mega_mhc_common import ( 

27 validate_mega_mhc_core_spec, validate_mega_mhc_native_inputs, validate_mega_mhc_token_count, 

28) 

29from .common import MegaMhcGradCache, MegaMhcGradOutputs, validate_mega_mhc_grad_tensors 

30from .graph import DEFAULT_GRAD_TOKEN_TILE 

31from .plan import MegaMhcGradPlan, build_mega_mhc_grad_plan 

32 

33 

34@lru_cache(maxsize=16) 

35def _get_cached_plan( 

36 token_count: int, 

37 hidden_size: int, 

38 device_index: int, 

39 num_cube_cores: int, 

40 num_vector_cores: int, 

41 token_tile: int, 

42) -> MegaMhcGradPlan: 

43 """Reuse immutable RuntimeConfig tensors for recurring static shapes.""" 

44 return build_mega_mhc_grad_plan( 

45 token_count, 

46 hidden_size, 

47 torch.device("npu", device_index), 

48 num_cube_cores, 

49 num_vector_cores, 

50 token_tile, 

51 ) 

52 

53 

54def _get_plan_for_shape( 

55 token_count: int, 

56 hidden_size: int, 

57 device: torch.device, 

58 token_tile: int, 

59) -> MegaMhcGradPlan: 

60 """Read current Device core limits and retrieve the immutable plan for this shape.""" 

61 device_index = device.index 

62 if device_index is None: 

63 device_index = torch.npu.current_device() 

64 core_spec = get_device_core_spec(device_index) 

65 validate_mega_mhc_core_spec(core_spec) 

66 num_cube_cores, num_vector_cores = core_spec.num_cube_cores, core_spec.num_vector_cores 

67 validate_mega_mhc_token_count(token_count, num_vector_cores) 

68 return _get_cached_plan( 

69 token_count, 

70 hidden_size, 

71 device_index, 

72 num_cube_cores, 

73 num_vector_cores, 

74 token_tile, 

75 ) 

76 

77 

78def _restore_output_shapes( 

79 outputs: tuple[torch.Tensor, ...], 

80 leading_shape: tuple[int, ...], 

81 hidden_size: int, 

82) -> MegaMhcGradOutputs: 

83 """Restore public leading dimensions and reorder native outputs. 

84 

85 Args: 

86 outputs: Native operator outputs in ACLNN order. 

87 leading_shape: Original activation leading dimensions. 

88 hidden_size: Hidden dimension of every token. 

89 

90 Returns: 

91 Public gradient tuple in forward-input order. 

92 """ 

93 ( 

94 grad_residual, 

95 grad_phi, 

96 grad_alpha, 

97 grad_bias, 

98 grad_previous_output, 

99 grad_previous_pre, 

100 grad_previous_post, 

101 grad_previous_residual, 

102 grad_norm_weight, 

103 ) = outputs 

104 return ( 

105 grad_previous_output.reshape(*leading_shape, hidden_size), 

106 grad_residual.reshape(*leading_shape, 4, hidden_size), 

107 grad_previous_pre.reshape(*leading_shape, 4), 

108 grad_previous_post.reshape(*leading_shape, 4), 

109 grad_previous_residual.reshape(*leading_shape, 4, 4), 

110 grad_phi, 

111 grad_alpha, 

112 grad_bias, 

113 grad_norm_weight, 

114 ) 

115 

116 

117def hyper_mega_mhc_grad( 

118 grad_new_residual: torch.Tensor, 

119 grad_next_pre_mix: torch.Tensor, 

120 grad_next_post_mix: torch.Tensor, 

121 grad_next_residual_mix: torch.Tensor, 

122 grad_block_input: torch.Tensor, 

123 previous_output: torch.Tensor, 

124 residual: torch.Tensor, 

125 previous_pre_mix: torch.Tensor, 

126 previous_post_mix: torch.Tensor, 

127 previous_residual_mix: torch.Tensor, 

128 phi: torch.Tensor, 

129 alpha: torch.Tensor, 

130 bias: torch.Tensor, 

131 norm_weight: torch.Tensor, 

132 cache: MegaMhcGradCache, 

133 *, 

134 hc_eps: float = 1e-6, 

135 token_tile: int = DEFAULT_GRAD_TOKEN_TILE, 

136) -> MegaMhcGradOutputs: 

137 """Execute the fused Prepare/PhiRms/PrevXPost backward pipeline. 

138 

139 Args: 

140 grad_new_residual: Gradient of the updated residual branches. 

141 grad_next_pre_mix: Gradient of the next pre-mix coefficients. 

142 grad_next_post_mix: Gradient of the next post-mix coefficients. 

143 grad_next_residual_mix: Gradient of the next residual-mix coefficients. 

144 grad_block_input: Gradient of the normalized block input. 

145 previous_output: Previous block output activation. 

146 residual: Current residual branches. 

147 previous_pre_mix: Previous pre-mix coefficients. 

148 previous_post_mix: Previous post-mix coefficients. 

149 previous_residual_mix: Previous residual-mix coefficients. 

150 phi: Projection weight. 

151 alpha: Projection scale. 

152 bias: Projection bias. 

153 norm_weight: RMSNorm weight. 

154 cache: Forward intermediates required by the backward kernel. 

155 hc_eps: RMSNorm epsilon. 

156 token_tile: Number of tokens assigned to each AIV task. 

157 

158 Returns: 

159 Gradients for the five block inputs and four trainable tensors. 

160 """ 

161 validate_mega_mhc_native_inputs( 

162 previous_output, residual, previous_pre_mix, previous_post_mix, previous_residual_mix, 

163 phi, alpha, bias, norm_weight, need_backward=True, token_tile=token_tile, 

164 ) 

165 validate_mega_mhc_grad_tensors( 

166 (grad_new_residual, grad_next_pre_mix, grad_next_post_mix, grad_next_residual_mix, grad_block_input), 

167 residual, 

168 cache, 

169 ) 

170 leading_shape, hidden_size = residual.shape[:-2], residual.shape[-1] 

171 token_count = residual.numel() // (4 * hidden_size) 

172 plan = _get_plan_for_shape(token_count, hidden_size, residual.device, token_tile) 

173 events = torch.zeros(plan.event_capacity * 4, dtype=torch.uint8, device=residual.device) 

174 kernel_inputs = ( 

175 grad_block_input.reshape(1, token_count, hidden_size).contiguous(), 

176 grad_next_post_mix.reshape(1, token_count, 4).contiguous(), 

177 grad_next_residual_mix.reshape(1, token_count, 4, 4).contiguous(), 

178 cache.new_residual.reshape(1, token_count, 4, hidden_size).contiguous(), 

179 phi.contiguous(), 

180 alpha.contiguous(), 

181 bias.contiguous(), 

182 previous_pre_mix.reshape(1, token_count, 4).contiguous(), 

183 cache.hc_before_norm.contiguous(), 

184 cache.inv_rms.contiguous(), 

185 cache.sum_out.contiguous(), 

186 cache.norm_out.contiguous(), 

187 grad_next_pre_mix.reshape(1, token_count, 4).contiguous(), 

188 cache.mixed_input.reshape(1, token_count, hidden_size).contiguous(), 

189 cache.rms_rstd.contiguous(), 

190 norm_weight.contiguous(), 

191 grad_new_residual.reshape(1, token_count, 4, hidden_size).contiguous(), 

192 residual.reshape(1, token_count, 4, hidden_size).contiguous(), 

193 previous_output.reshape(1, token_count, hidden_size).contiguous(), 

194 previous_post_mix.reshape(1, token_count, 4).contiguous(), 

195 previous_residual_mix.reshape(1, token_count, 4, 4).contiguous(), 

196 ) 

197 options = {"hc_eps": hc_eps} 

198 if plan.runtime.is_active: 

199 options["diagnostics"] = plan.runtime 

200 outputs = multicore_ops.hyper_mega_mhc_grad( 

201 *kernel_inputs, plan.runtime.runtime_config, events, events, **options, 

202 ) 

203 return _restore_output_shapes(outputs, leading_shape, hidden_size)