Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_multicore / modules / mega_mhc / function.py: 73%
59 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-11-15 05:03 +0800
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 HyperMegaMhc scheduler kernel."""
17from __future__ import annotations
19from functools import lru_cache
20from typing import Any
22import torch
24from hyper_multicore._core_spec import get_device_core_spec
25from hyper_multicore.modules.mega_mhc_grad.common import MegaMhcGradCache
26from hyper_multicore.modules.mega_mhc_grad.function import hyper_mega_mhc_grad
27from hyper_multicore.torch import ops as multicore_ops
29from ..mega_mhc_common import (
30 MHC_SINKHORN_ITERS,
31 MegaMhcOutputs,
32 validate_mega_mhc_core_spec,
33 validate_mega_mhc_native_inputs,
34 validate_mega_mhc_token_count,
35)
36from .graph import DEFAULT_TOKEN_TILE
37from .plan import MegaMhcPlan, build_mega_mhc_plan
40@lru_cache(maxsize=16)
41def _get_cached_plan(
42 token_count: int,
43 hidden_size: int,
44 device_index: int,
45 num_cube_cores: int,
46 token_tile: int,
47) -> MegaMhcPlan:
48 """Reuse immutable RuntimeConfig tensors for recurring static shapes."""
49 return build_mega_mhc_plan(
50 token_count,
51 hidden_size,
52 torch.device("npu", device_index),
53 num_cube_cores,
54 token_tile,
55 )
58def _restore_forward_outputs(
59 outputs: tuple[torch.Tensor, ...],
60 leading_shape: tuple[int, ...],
61 hidden_size: int,
62) -> tuple[MegaMhcOutputs, MegaMhcGradCache]:
63 """Restore public dimensions and retain the native-layout backward cache."""
64 (
65 new_residual,
66 next_pre,
67 next_post,
68 next_res,
69 block_input,
70 hc_before_norm,
71 inv_rms,
72 sum_out,
73 norm_out,
74 mixed_input,
75 rms_rstd,
76 ) = outputs
77 public_outputs = (
78 new_residual.reshape(*leading_shape, 4, hidden_size),
79 next_pre.reshape(*leading_shape, 4),
80 next_post.reshape(*leading_shape, 4),
81 next_res.reshape(*leading_shape, 4, 4),
82 block_input.reshape(*leading_shape, hidden_size),
83 )
84 cache = MegaMhcGradCache(
85 new_residual=new_residual,
86 hc_before_norm=hc_before_norm,
87 inv_rms=inv_rms,
88 sum_out=sum_out,
89 norm_out=norm_out,
90 mixed_input=mixed_input,
91 rms_rstd=rms_rstd,
92 )
93 return public_outputs, cache
96def _hyper_mega_mhc_forward(
97 previous_output: torch.Tensor,
98 residual: torch.Tensor,
99 previous_pre_mix: torch.Tensor,
100 previous_post_mix: torch.Tensor,
101 previous_residual_mix: torch.Tensor,
102 phi: torch.Tensor,
103 alpha: torch.Tensor,
104 bias: torch.Tensor,
105 norm_weight: torch.Tensor,
106 *,
107 hc_eps: float = 1e-6,
108 norm_eps: float = 1e-6,
109 num_iters: int = MHC_SINKHORN_ITERS,
110 need_backward: bool = True,
111 token_tile: int = DEFAULT_TOKEN_TILE,
112) -> tuple[MegaMhcOutputs, MegaMhcGradCache]:
113 """Run one shifted boundary and retain the native backward cache.
115 Args:
116 previous_output: Previous block output.
117 residual: Previous residual streams.
118 previous_pre_mix: Previous pre-mix coefficients.
119 previous_post_mix: Previous post-mix coefficients.
120 previous_residual_mix: Previous residual-mix coefficients.
121 phi: Projection weights.
122 alpha: Mapping scale parameters.
123 bias: Mapping bias parameters.
124 norm_weight: Shifted RMSNorm weight.
125 hc_eps: Sinkhorn numerical-stability epsilon.
126 norm_eps: RMS normalization epsilon.
127 num_iters: Sinkhorn iteration count.
128 need_backward: Whether to retain native backward caches.
129 token_tile: Number of tokens assigned to each logical task.
130 """
131 validate_mega_mhc_native_inputs(
132 previous_output,
133 residual,
134 previous_pre_mix,
135 previous_post_mix,
136 previous_residual_mix,
137 phi,
138 alpha,
139 bias,
140 norm_weight,
141 num_iters=num_iters,
142 need_backward=need_backward,
143 token_tile=token_tile,
144 )
145 leading_shape = residual.shape[:-2]
146 hidden_size = residual.shape[-1]
147 token_count = residual.numel() // (4 * hidden_size)
148 device_index = residual.device.index
149 if device_index is None:
150 device_index = torch.npu.current_device()
151 core_spec = get_device_core_spec(device_index)
152 validate_mega_mhc_core_spec(core_spec)
153 num_cube_cores, num_vector_cores = core_spec.num_cube_cores, core_spec.num_vector_cores
154 if need_backward:
155 validate_mega_mhc_token_count(token_count, num_vector_cores)
156 plan = _get_cached_plan(
157 token_count,
158 hidden_size,
159 device_index,
160 num_cube_cores,
161 token_tile,
162 )
163 events = torch.zeros(plan.event_capacity * 4, dtype=torch.uint8, device=residual.device)
164 options = {}
165 if plan.runtime.is_active:
166 options["diagnostics"] = plan.runtime
167 outputs = multicore_ops.hyper_mega_mhc(
168 previous_output.reshape(1, token_count, hidden_size).contiguous(),
169 residual.reshape(1, token_count, 4, hidden_size).contiguous(),
170 previous_pre_mix.reshape(1, token_count, 4).contiguous(),
171 previous_post_mix.reshape(1, token_count, 4).contiguous(),
172 previous_residual_mix.reshape(1, token_count, 4, 4).contiguous(),
173 phi.contiguous(),
174 alpha.contiguous(),
175 bias.contiguous(),
176 norm_weight.contiguous(),
177 plan.runtime.runtime_config,
178 events,
179 events,
180 hc_eps=hc_eps,
181 norm_eps=norm_eps,
182 num_iters=num_iters,
183 need_backward=need_backward,
184 **options,
185 )
186 return _restore_forward_outputs(outputs, leading_shape, hidden_size)
189class _HyperMegaMhcFunction(torch.autograd.Function):
190 """Connect the fused forward and backward NPU operators to autograd."""
192 @staticmethod
193 def forward(
194 ctx: Any,
195 previous_output: torch.Tensor,
196 residual: torch.Tensor,
197 previous_pre_mix: torch.Tensor,
198 previous_post_mix: torch.Tensor,
199 previous_residual_mix: torch.Tensor,
200 phi: torch.Tensor,
201 alpha: torch.Tensor,
202 bias: torch.Tensor,
203 norm_weight: torch.Tensor,
204 hc_eps: float,
205 norm_eps: float,
206 num_iters: int,
207 token_tile: int,
208 ) -> MegaMhcOutputs:
209 """Launch HyperMegaMhc and save its explicit native cache.
211 Args:
212 ctx: Autograd context retaining only the state required for backward.
213 previous_output: Previous block output tensor.
214 residual: Residual streams supplied to the current block.
215 previous_pre_mix: Previous block's input mixing coefficients.
216 previous_post_mix: Previous block's output mixing coefficients.
217 previous_residual_mix: Previous block's residual mixing matrix.
218 phi: Projection weight used to generate mixing coefficients.
219 alpha: Scaling parameters for the projected mixing coefficients.
220 bias: Bias parameters for the projected mixing coefficients.
221 norm_weight: RMS normalization weight.
222 hc_eps: Numerical epsilon used by mixing coefficient normalization.
223 norm_eps: Numerical epsilon used by RMS normalization.
224 num_iters: Number of Sinkhorn normalization iterations.
225 token_tile: Number of token rows assigned to a task tile.
226 """
227 outputs, cache = _hyper_mega_mhc_forward(
228 previous_output,
229 residual,
230 previous_pre_mix,
231 previous_post_mix,
232 previous_residual_mix,
233 phi,
234 alpha,
235 bias,
236 norm_weight,
237 hc_eps=hc_eps,
238 norm_eps=norm_eps,
239 num_iters=num_iters,
240 need_backward=True,
241 token_tile=token_tile,
242 )
243 ctx.save_for_backward(
244 previous_output,
245 residual,
246 previous_pre_mix,
247 previous_post_mix,
248 previous_residual_mix,
249 phi,
250 alpha,
251 bias,
252 norm_weight,
253 cache.new_residual,
254 cache.hc_before_norm,
255 cache.inv_rms,
256 cache.sum_out,
257 cache.norm_out,
258 cache.mixed_input,
259 cache.rms_rstd,
260 )
261 ctx.hc_eps = hc_eps
262 ctx.token_tile = token_tile
263 return outputs
265 @staticmethod
266 def backward(ctx: Any, *grad_outputs: torch.Tensor) -> tuple[torch.Tensor | None, ...]:
267 """Differentiate the boundary with HyperMegaMhcGrad.
269 Args:
270 ctx: Autograd context retaining only the state required for backward.
271 grad_outputs: Upstream gradients for the differentiable forward results.
272 """
273 (
274 previous_output,
275 residual,
276 previous_pre_mix,
277 previous_post_mix,
278 previous_residual_mix,
279 phi,
280 alpha,
281 bias,
282 norm_weight,
283 new_residual,
284 hc_before_norm,
285 inv_rms,
286 sum_out,
287 norm_out,
288 mixed_input,
289 rms_rstd,
290 ) = ctx.saved_tensors
291 cache = MegaMhcGradCache(
292 new_residual=new_residual,
293 hc_before_norm=hc_before_norm,
294 inv_rms=inv_rms,
295 sum_out=sum_out,
296 norm_out=norm_out,
297 mixed_input=mixed_input,
298 rms_rstd=rms_rstd,
299 )
300 gradients = hyper_mega_mhc_grad(
301 *grad_outputs,
302 previous_output,
303 residual,
304 previous_pre_mix,
305 previous_post_mix,
306 previous_residual_mix,
307 phi,
308 alpha,
309 bias,
310 norm_weight,
311 cache,
312 hc_eps=ctx.hc_eps,
313 token_tile=ctx.token_tile,
314 )
315 return (*gradients, None, None, None, None)
318def hyper_mega_mhc(
319 previous_output: torch.Tensor,
320 residual: torch.Tensor,
321 previous_pre_mix: torch.Tensor,
322 previous_post_mix: torch.Tensor,
323 previous_residual_mix: torch.Tensor,
324 phi: torch.Tensor,
325 alpha: torch.Tensor,
326 bias: torch.Tensor,
327 norm_weight: torch.Tensor,
328 *,
329 hc_eps: float = 1e-6,
330 norm_eps: float = 1e-6,
331 num_iters: int = MHC_SINKHORN_ITERS,
332 need_backward: bool = True,
333 token_tile: int = DEFAULT_TOKEN_TILE,
334) -> MegaMhcOutputs:
335 """Run the fused shifted boundary with its matching autograd backward."""
336 differentiable_inputs = (
337 previous_output,
338 residual,
339 previous_pre_mix,
340 previous_post_mix,
341 previous_residual_mix,
342 phi,
343 alpha,
344 bias,
345 norm_weight,
346 )
347 if need_backward and torch.is_grad_enabled() and any(
348 isinstance(value, torch.Tensor) and value.requires_grad for value in differentiable_inputs
349 ):
350 return _HyperMegaMhcFunction.apply(
351 *differentiable_inputs,
352 hc_eps,
353 norm_eps,
354 num_iters,
355 token_tile,
356 )
357 outputs, _ = _hyper_mega_mhc_forward(
358 *differentiable_inputs,
359 hc_eps=hc_eps,
360 norm_eps=norm_eps,
361 num_iters=num_iters,
362 need_backward=False,
363 token_tile=token_tile,
364 )
365 return outputs