Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_multicore / torch / ops.py: 36%
125 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"""
16hyper_multicore.torch.ops
17=============================================
18Out-of-tree PyTorch bindings for MegaMoE, MegaGate and MegaMHC operators.
20Registers into the ``hyper_multicore`` PyTorch namespace — does NOT modify
21op-plugin or any PyTorch source. The operators are accessible via:
23 torch.ops.hyper_multicore.mega_moe(...)
24 torch.ops.hyper_multicore.mega_moe_grad(...)
26Or via the Python wrappers in this module:
28 from hyper_multicore.torch.ops import mega_moe, mega_moe_grad
30Forward and backward ACLNN symbols are packaged in one component-owned
31``hyper_multicore_nn`` vendor. Source the packaged ``set_env.bash``
32before starting the application or framework Python process so CANN can discover that vendor.
33"""
34from __future__ import annotations
36__all__ = [
37 "cann_mhc_pre_sinkhorn",
38 "cann_mhc_pre_sinkhorn_with_cache",
39 "hyper_mega_mhc",
40 "hyper_mega_mhc_grad",
41 "mega_moe",
42 "mega_moe_grad",
43 "moe_token_permute_grad",
44 "moe_token_permute_out",
45]
47import torch
49from hyper_multicore._loader import (
50 NativeComponentUnavailableError,
51 load_torch_npu,
52 get_multicore_paths,
53 preload_vendor_library,
54 native_loading, native_context, native_ready,
55)
56from hyper_multicore.diagnostics import MegaKernelDiagnostics
59def _launch_diagnostic_out(
60 operator, inputs, outputs, runtime_config, event_counters, diagnostic_buffer,
61 diagnostics, direction, attributes,
62):
63 """Bind preallocated native out operands before the selected tool prepares."""
64 parameters = (*inputs, runtime_config, event_counters, diagnostic_buffer, *outputs)
65 with diagnostics.launch(parameters, direction=direction, fallback_event_counters=event_counters) as call:
66 return operator(*inputs, call.runtime_config, call.event_counters, call.diagnostic_buffer,
67 *outputs, *attributes)
70def _load_native() -> None:
71 """Register the ABI-specific native adapter once on first operation."""
72 if native_ready("torch"):
73 return
74 load_torch_npu()
75 vendor_root, adapter_path = get_multicore_paths()
76 with native_loading("torch") as needed:
77 if not needed:
78 return
79 preload_vendor_library(vendor_root)
80 try:
81 torch.ops.load_library(str(adapter_path))
82 if torch.ops.hyper_multicore.mega_moe_transport_version() != 1:
83 raise RuntimeError("adapter transport ABI mismatch; rebuild the Torch adapter")
84 # Torch exposes operator alias metadata through this private schema attribute.
85 schema = torch.ops.hyper_multicore.mega_moe_grad.default._schema # pylint: disable=protected-access
86 dispatch_alias = schema.arguments[0].alias_info.before_set
87 input_grad_alias = schema.arguments[12].alias_info.before_set
88 if not dispatch_alias or dispatch_alias != input_grad_alias:
89 raise RuntimeError("adapter does not declare backward dispatch storage reuse; rebuild the Torch adapter")
90 except (AttributeError, OSError, RuntimeError) as error:
91 raise NativeComponentUnavailableError(
92 "[HP-NATIVE-FRAMEWORK-ADAPTER-LOAD-FAILED] component=multicore framework=torch "
93 f"library={adapter_path} error={error}; {native_context()}. "
94 "Check the Python/Torch/torch_npu/CANN version combination and rebuild the Torch adapter."
95 ) from error
98# ---------------------------------------------------------------------------
101# Python wrappers — thin pass-through to the registered C++ ops
102# ---------------------------------------------------------------------------
105def cann_mhc_pre_sinkhorn(
106 x: torch.Tensor,
107 phi: torch.Tensor,
108 alpha: torch.Tensor,
109 bias: torch.Tensor,
110 *,
111 hc_eps: float = 1e-6,
112 norm_eps: float = 1e-6,
113 num_iters: int = 20,
114) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
115 """Call CANN MhcPreSinkhorn while retaining mappings needed by shifting.
117 Args:
118 x: Token-major residual input.
119 phi: Projection weights.
120 alpha: MhcPre scaling parameters.
121 bias: MhcPre bias parameters.
122 hc_eps: Sinkhorn numerical-stability epsilon.
123 norm_eps: RMS normalization epsilon.
124 num_iters: Sinkhorn iteration count.
125 """
126 _load_native()
127 if x.ndim != 3:
128 raise ValueError(f"x must be token-major [M, 4, H], got {tuple(x.shape)}.")
129 token_count, hc_mult, hidden_size = x.shape
130 hc_mix = hc_mult * hc_mult + 2 * hc_mult
131 x_bsnd = x.unsqueeze(0)
132 hin = torch.empty((1, token_count, hidden_size), dtype=x.dtype, device=x.device)
133 h_post = torch.empty((1, token_count, hc_mult), dtype=torch.float32, device=x.device)
134 h_res = torch.empty((1, token_count, hc_mult * hc_mult), dtype=torch.float32, device=x.device)
135 h_pre = torch.empty((1, token_count, hc_mult), dtype=torch.float32, device=x.device)
136 hc_before_norm = torch.empty((1, token_count, hc_mix), dtype=torch.float32, device=x.device)
137 inv_rms = torch.empty((1, token_count, 1), dtype=torch.float32, device=x.device)
138 sum_out = torch.empty((2 * num_iters, 1, token_count, hc_mult), dtype=torch.float32, device=x.device)
139 norm_out = torch.empty(
140 (2 * num_iters, 1, token_count, hc_mult, hc_mult),
141 dtype=torch.float32,
142 device=x.device,
143 )
144 torch.ops.hyper_multicore.cann_mhc_pre_sinkhorn(
145 x_bsnd,
146 phi,
147 alpha,
148 bias,
149 hin,
150 h_post,
151 h_res,
152 h_pre,
153 hc_before_norm,
154 inv_rms,
155 sum_out,
156 norm_out,
157 hc_eps,
158 norm_eps,
159 num_iters,
160 )
161 return (
162 hin.squeeze(0),
163 h_post.squeeze(0),
164 h_res.squeeze(0).reshape(token_count, hc_mult, hc_mult),
165 h_pre.squeeze(0),
166 )
169def cann_mhc_pre_sinkhorn_with_cache(
170 x: torch.Tensor,
171 phi: torch.Tensor,
172 alpha: torch.Tensor,
173 bias: torch.Tensor,
174 *,
175 hc_eps: float = 1e-6,
176 norm_eps: float = 1e-6,
177 num_iters: int = 20,
178) -> tuple[torch.Tensor, ...]:
179 """Call CANN MhcPreSinkhorn and expose all native backward caches.
181 Args:
182 x: Token-major residual input.
183 phi: Projection weights.
184 alpha: MhcPre scaling parameters.
185 bias: MhcPre bias parameters.
186 hc_eps: Sinkhorn numerical-stability epsilon.
187 norm_eps: RMS normalization epsilon.
188 num_iters: Sinkhorn iteration count.
189 """
190 _load_native()
191 if x.ndim != 3:
192 raise ValueError(f"x must be token-major [M, 4, H], got {tuple(x.shape)}.")
193 token_count, hc_mult, hidden_size = x.shape
194 hc_mix = hc_mult * hc_mult + 2 * hc_mult
195 x_bsnd = x.unsqueeze(0)
196 hin = torch.empty((1, token_count, hidden_size), dtype=x.dtype, device=x.device)
197 h_post = torch.empty((1, token_count, hc_mult), dtype=torch.float32, device=x.device)
198 h_res = torch.empty((1, token_count, hc_mult * hc_mult), dtype=torch.float32, device=x.device)
199 h_pre = torch.empty((1, token_count, hc_mult), dtype=torch.float32, device=x.device)
200 hc_before_norm = torch.empty((1, token_count, hc_mix), dtype=torch.float32, device=x.device)
201 inv_rms = torch.empty((1, token_count, 1), dtype=torch.float32, device=x.device)
202 sum_out = torch.empty((2 * num_iters, 1, token_count, hc_mult), dtype=torch.float32, device=x.device)
203 norm_out = torch.empty(
204 (2 * num_iters, 1, token_count, hc_mult, hc_mult), dtype=torch.float32, device=x.device
205 )
206 torch.ops.hyper_multicore.cann_mhc_pre_sinkhorn(
207 x_bsnd,
208 phi,
209 alpha,
210 bias,
211 hin,
212 h_post,
213 h_res,
214 h_pre,
215 hc_before_norm,
216 inv_rms,
217 sum_out,
218 norm_out,
219 hc_eps,
220 norm_eps,
221 num_iters,
222 )
223 return hin, h_post, h_res, h_pre, hc_before_norm, inv_rms, sum_out, norm_out
226def _allocate_hyper_mega_mhc_forward_cache(
227 residual: torch.Tensor,
228 num_iters: int,
229) -> tuple[torch.Tensor, ...]:
230 """Allocate the existing six native backward intermediates in ACLNN order."""
231 batch_size, sequence_length, hc_mult, hidden_size = residual.shape
232 hc_mix = hc_mult * hc_mult + 2 * hc_mult
233 sinkhorn_steps = 2 * num_iters
234 return (
235 torch.empty(
236 (batch_size, sequence_length, hc_mix),
237 dtype=torch.float32,
238 device=residual.device,
239 ),
240 torch.empty(
241 (batch_size, sequence_length, 1),
242 dtype=torch.float32,
243 device=residual.device,
244 ),
245 torch.empty(
246 (sinkhorn_steps, batch_size, sequence_length, hc_mult),
247 dtype=torch.float32,
248 device=residual.device,
249 ),
250 torch.empty(
251 (sinkhorn_steps, batch_size, sequence_length, hc_mult, hc_mult),
252 dtype=torch.float32,
253 device=residual.device,
254 ),
255 torch.empty(
256 (batch_size, sequence_length, hidden_size),
257 dtype=residual.dtype,
258 device=residual.device,
259 ),
260 torch.empty(
261 (batch_size, sequence_length, 1),
262 dtype=torch.float32,
263 device=residual.device,
264 ),
265 )
268def hyper_mega_mhc(
269 previous_output: torch.Tensor,
270 residual: torch.Tensor,
271 previous_pre_mix: torch.Tensor,
272 previous_post_mix: torch.Tensor,
273 previous_residual_mix: torch.Tensor,
274 phi: torch.Tensor,
275 alpha: torch.Tensor,
276 bias: torch.Tensor,
277 norm_weight: torch.Tensor,
278 runtime_config: torch.Tensor,
279 all_event_counters: torch.Tensor,
280 diagnostic_buffer: torch.Tensor,
281 *,
282 hc_eps: float = 1e-6,
283 norm_eps: float = 1e-6,
284 num_iters: int = 20,
285 need_backward: bool,
286 diagnostics: MegaKernelDiagnostics | None = None,
287) -> tuple[torch.Tensor, ...]:
288 """Allocate outputs and launch the single-kernel shifted mHC implementation.
290 Args:
291 previous_output: Previous block output used by MhcPost.
292 residual: Previous block residual state.
293 previous_pre_mix: Previous block pre-mix coefficients.
294 previous_post_mix: Previous block post-mix coefficients.
295 previous_residual_mix: Previous block residual-mix coefficients.
296 phi: Projection weights.
297 alpha: MhcPre scaling parameters.
298 bias: MhcPre bias parameters.
299 norm_weight: Shifted RMSNorm weight.
300 runtime_config: Serialized multicore schedule.
301 all_event_counters: Device event-counter storage.
302 diagnostic_buffer: Shared Device diagnostics storage or a disabled placeholder.
303 hc_eps: Sinkhorn numerical-stability epsilon.
304 norm_eps: RMS normalization epsilon.
305 num_iters: Sinkhorn iteration count.
306 need_backward: Whether MhcPre should populate backward caches.
307 diagnostics: Active generic launch boundary, omitted on the ordinary path.
309 Returns:
310 Five public forward outputs followed by the six tensors consumed by
311 ``HyperMegaMhcGrad``.
312 """
313 _load_native()
314 if residual.ndim != 4:
315 raise ValueError(f"residual must be BSND [B, S, 4, H], got {tuple(residual.shape)}.")
316 batch_size, sequence_length, hc_mult, hidden_size = residual.shape
317 new_residual = torch.empty_like(residual)
318 next_pre_mix = torch.empty(
319 (batch_size, sequence_length, hc_mult),
320 dtype=torch.float32,
321 device=residual.device,
322 )
323 next_post_mix = torch.empty_like(next_pre_mix)
324 next_residual_mix = torch.empty(
325 (batch_size, sequence_length, hc_mult * hc_mult),
326 dtype=torch.float32,
327 device=residual.device,
328 )
329 block_input = torch.empty(
330 (batch_size, sequence_length, hidden_size),
331 dtype=residual.dtype,
332 device=residual.device,
333 )
334 cache_outputs = _allocate_hyper_mega_mhc_forward_cache(residual, num_iters)
335 if diagnostics is not None:
336 diagnostic_outputs = _launch_diagnostic_out(
337 torch.ops.hyper_multicore.mega_mhc,
338 (previous_output, residual, previous_pre_mix, previous_post_mix, previous_residual_mix,
339 phi, alpha, bias, norm_weight),
340 (new_residual, next_pre_mix, next_post_mix, next_residual_mix, block_input, *cache_outputs),
341 runtime_config, all_event_counters, diagnostic_buffer, diagnostics, "forward",
342 (hc_eps, norm_eps, num_iters, need_backward),
343 )
344 return (*diagnostic_outputs[:3],
345 diagnostic_outputs[3].reshape(batch_size, sequence_length, hc_mult, hc_mult),
346 *diagnostic_outputs[4:])
347 torch.ops.hyper_multicore.mega_mhc(
348 previous_output, residual,
349 previous_pre_mix, previous_post_mix,
350 previous_residual_mix, phi,
351 alpha, bias,
352 norm_weight, runtime_config,
353 all_event_counters, diagnostic_buffer,
354 new_residual, next_pre_mix,
355 next_post_mix, next_residual_mix,
356 block_input, *cache_outputs,
357 hc_eps, norm_eps,
358 num_iters, need_backward,
359 )
360 return (
361 new_residual,
362 next_pre_mix,
363 next_post_mix,
364 next_residual_mix.reshape(batch_size, sequence_length, hc_mult, hc_mult),
365 block_input,
366 *cache_outputs,
367 )
370def _allocate_hyper_mega_mhc_grad_outputs(
371 previous_residual: torch.Tensor,
372 phi: torch.Tensor,
373 alpha: torch.Tensor,
374 bias: torch.Tensor,
375 previous_output: torch.Tensor,
376 previous_pre: torch.Tensor,
377 previous_post: torch.Tensor,
378 previous_residual_mix: torch.Tensor,
379 norm_weight: torch.Tensor,
380) -> tuple[torch.Tensor, ...]:
381 """Allocate output gradients required by the native backward ABI."""
382 return (
383 torch.empty_like(previous_residual),
384 torch.empty_like(phi),
385 torch.empty_like(alpha),
386 torch.empty_like(bias),
387 torch.empty_like(previous_output),
388 torch.empty_like(previous_pre),
389 torch.empty_like(previous_post),
390 torch.empty_like(previous_residual_mix),
391 torch.empty(norm_weight.shape, dtype=torch.float32, device=norm_weight.device),
392 )
395def hyper_mega_mhc_grad(
396 grad_hin_placeholder: torch.Tensor,
397 grad_h_post: torch.Tensor,
398 grad_h_res: torch.Tensor,
399 x: torch.Tensor,
400 phi: torch.Tensor,
401 alpha: torch.Tensor,
402 bias: torch.Tensor,
403 previous_pre: torch.Tensor,
404 hc_before_norm: torch.Tensor,
405 inv_rms: torch.Tensor,
406 sum_out: torch.Tensor,
407 norm_out: torch.Tensor,
408 grad_current_pre: torch.Tensor,
409 mixed_input: torch.Tensor,
410 rms_rstd: torch.Tensor,
411 norm_weight: torch.Tensor,
412 direct_grad_x: torch.Tensor,
413 previous_residual: torch.Tensor,
414 previous_output: torch.Tensor,
415 previous_post: torch.Tensor,
416 previous_residual_mix: torch.Tensor,
417 runtime_config: torch.Tensor,
418 all_event_counters: torch.Tensor,
419 diagnostic_buffer: torch.Tensor,
420 *,
421 hc_eps: float = 1e-6,
422 diagnostics: MegaKernelDiagnostics | None = None,
423) -> tuple[torch.Tensor, ...]:
424 """Allocate gradients and launch the single-kernel shifted mHC backward.
426 Args:
427 grad_hin_placeholder: Placeholder gradient matching native ABI order.
428 grad_h_post: Gradient of current post-mix coefficients.
429 grad_h_res: Gradient of current residual-mix coefficients.
430 x: Forward MhcPre input.
431 phi: Projection weights.
432 alpha: MhcPre scaling parameters.
433 bias: MhcPre bias parameters.
434 previous_pre: Previous pre-mix coefficients.
435 hc_before_norm: Cached pre-normalization coefficients.
436 inv_rms: Cached coefficient inverse RMS.
437 sum_out: Cached Sinkhorn reductions.
438 norm_out: Cached Sinkhorn normalization values.
439 grad_current_pre: Gradient of current pre-mix coefficients.
440 mixed_input: Cached shifted input mix.
441 rms_rstd: Cached shifted RMSNorm reciprocal standard deviation.
442 norm_weight: Shifted RMSNorm weight.
443 direct_grad_x: Direct gradient of the current residual state.
444 previous_residual: Previous residual state.
445 previous_output: Previous block output.
446 previous_post: Previous post-mix coefficients.
447 previous_residual_mix: Previous residual-mix coefficients.
448 runtime_config: Serialized multicore schedule.
449 all_event_counters: Device event-counter storage.
450 diagnostic_buffer: Shared Device diagnostics storage or a disabled placeholder.
451 hc_eps: Sinkhorn numerical-stability epsilon.
452 diagnostics: Active generic launch boundary, omitted on the ordinary path.
453 """
454 _load_native()
455 grad_outputs = _allocate_hyper_mega_mhc_grad_outputs(
456 previous_residual,
457 phi,
458 alpha,
459 bias,
460 previous_output,
461 previous_pre,
462 previous_post,
463 previous_residual_mix,
464 norm_weight,
465 )
466 if diagnostics is not None:
467 return _launch_diagnostic_out(
468 torch.ops.hyper_multicore.mega_mhc_grad,
469 (grad_hin_placeholder, grad_h_post, grad_h_res, x, phi, alpha, bias, previous_pre,
470 hc_before_norm, inv_rms, sum_out, norm_out, grad_current_pre, mixed_input, rms_rstd,
471 norm_weight, direct_grad_x, previous_residual, previous_output, previous_post, previous_residual_mix),
472 grad_outputs, runtime_config, all_event_counters, diagnostic_buffer, diagnostics, "backward", (hc_eps,),
473 )
474 return torch.ops.hyper_multicore.mega_mhc_grad(
475 grad_hin_placeholder, grad_h_post,
476 grad_h_res, x,
477 phi, alpha,
478 bias, previous_pre,
479 hc_before_norm, inv_rms,
480 sum_out, norm_out,
481 grad_current_pre, mixed_input,
482 rms_rstd, norm_weight,
483 direct_grad_x, previous_residual,
484 previous_output, previous_post,
485 previous_residual_mix, runtime_config,
486 all_event_counters, diagnostic_buffer,
487 *grad_outputs, hc_eps,
488 )
491def moe_token_permute_out(
492 tokens: torch.Tensor,
493 indices: torch.Tensor,
494 output: torch.Tensor,
495 mapping: torch.Tensor,
496) -> None:
497 """Permute all routed rows directly into caller-owned contiguous storage.
499 Args:
500 tokens: Contiguous token states shaped ``[tokens, hidden]``.
501 indices: Contiguous INT32/INT64 expert IDs shaped ``[tokens, top_k]``.
502 output: Output shaped ``[tokens * top_k, hidden]``, matching token dtype.
503 mapping: Independent INT32 inverse mapping shaped ``[tokens * top_k]``.
504 """
505 _load_native()
506 torch.ops.hyper_multicore.moe_token_permute_out(tokens, indices, output, mapping)
509def mega_moe_unpermute_grad_out(
510 permuted_tokens: torch.Tensor,
511 grad_output: torch.Tensor,
512 sorted_indices: torch.Tensor,
513 probs: torch.Tensor,
514 grad_permuted: torch.Tensor,
515 grad_probs: torch.Tensor,
516) -> None:
517 """Write routing gradients into owned outputs inside a workspace lease.
519 Args:
520 permuted_tokens: Independently saved expert-major forward outputs.
521 grad_output: Token-major output gradients.
522 sorted_indices: Flat INT32 inverse mapping from input permutation.
523 probs: Contiguous FP32 top-k probabilities saved by forward.
524 grad_permuted: Contiguous symmetric source receiving expert-major dY.
525 grad_probs: Independent FP32 output matching the probabilities.
526 """
527 _load_native()
528 torch.ops.hyper_multicore.mega_moe_unpermute_grad_out(
529 permuted_tokens, grad_output, sorted_indices, probs, grad_permuted, grad_probs,
530 )
533def mega_moe(
534 dispatch_target: torch.Tensor,
535 dispatch_target_off: torch.Tensor,
536 dispatch_src: torch.Tensor,
537 dispatch_src_off: torch.Tensor,
538 dispatch_size: torch.Tensor,
539 up_proj_weight: torch.Tensor,
540 up_proj_glist: torch.Tensor,
541 up_proj_y: torch.Tensor,
542 swiglu_out: torch.Tensor,
543 down_proj_weight: torch.Tensor,
544 down_proj_glist: torch.Tensor,
545 down_proj_y: torch.Tensor,
546 combine_target: torch.Tensor,
547 combine_target_off: torch.Tensor,
548 combine_src_off: torch.Tensor,
549 combine_size: torch.Tensor,
550 gmm_workspace: torch.Tensor,
551 up_proj_tiling: torch.Tensor,
552 swiglu_tiling: torch.Tensor,
553 down_proj_tiling: torch.Tensor,
554 runtime_config: torch.Tensor,
555 all_event_counters: torch.Tensor,
556 rank_id: int,
557 ep: int,
558 expert_num: int,
559 hidden_size: int,
560 seq_size: int,
561) -> None:
562 """
563 MoE-FFN forward operator.
565 Args:
566 dispatch_target: First operator tensor; the detailed ABI is documented below.
568 Writes in-place to: dispatch_target, up_proj_y, swiglu_out, down_proj_y,
569 combine_target.
570 All output tensors must be pre-allocated with correct shapes.
572 Args:
573 dispatch_target: First tensor in the fixed flat forward ABI; the complete parameter groups are documented
574 below and retain their registered schema order.
576 Parameters
577 ----------
578 dispatch_target, dispatch_target_off, dispatch_src, dispatch_src_off,
579 dispatch_size :
580 AllToAll dispatch buffers — dispatch_target written in-place.
581 up_proj_weight, up_proj_glist :
582 Expert weight and cumulative group sizes for GMM1 (up-projection).
583 up_proj_y, swiglu_out :
584 GMM1 output and SwiGLU output — written in-place.
585 down_proj_weight, down_proj_glist, down_proj_y :
586 Expert weight, cumulative group sizes, and output for GMM2 (down-projection).
587 combine_target, combine_target_off, combine_src_off, combine_size :
588 AllToAll combine buffers — combine_target written in-place.
589 gmm_workspace, up_proj_tiling, swiglu_tiling, down_proj_tiling :
590 Pre-computed tiling tensors (from gen_runtime_data.py).
591 runtime_config :
592 Per-rank runtime config tensor (from gen_runtime_data.py).
593 all_event_counters :
594 Event synchronization counter tensor.
595 rank_id, ep, expert_num, hidden_size, seq_size :
596 Topology / shape attributes.
597 """
598 _load_native()
599 torch.ops.hyper_multicore.mega_moe(
600 dispatch_target,
601 dispatch_target_off,
602 dispatch_src,
603 dispatch_src_off,
604 dispatch_size,
605 up_proj_weight,
606 up_proj_glist,
607 up_proj_y,
608 swiglu_out,
609 down_proj_weight,
610 down_proj_glist,
611 down_proj_y,
612 combine_target,
613 combine_target_off,
614 combine_src_off,
615 combine_size,
616 gmm_workspace,
617 up_proj_tiling,
618 swiglu_tiling,
619 down_proj_tiling,
620 runtime_config,
621 all_event_counters,
622 all_event_counters,
623 rank_id,
624 ep,
625 expert_num,
626 hidden_size,
627 seq_size,
628 )
631def mega_moe_grad(
632 dispatch_target: torch.Tensor,
633 dispatch_target_off: torch.Tensor,
634 dy: torch.Tensor,
635 dispatch_src_off: torch.Tensor,
636 dispatch_size: torch.Tensor,
637 hidden: torch.Tensor,
638 hidden_dw: torch.Tensor,
639 w2: torch.Tensor,
640 act_grad_y: torch.Tensor,
641 gate: torch.Tensor,
642 grad_gate: torch.Tensor,
643 w1: torch.Tensor,
644 gate_dx: torch.Tensor,
645 grad_x: torch.Tensor,
646 combine_target_off: torch.Tensor,
647 combine_src_off: torch.Tensor,
648 combine_size: torch.Tensor,
649 permute_out: torch.Tensor,
650 gate_dw: torch.Tensor,
651 group_list: torch.Tensor,
652 act_grad_tiling: torch.Tensor,
653 gate_grad_tiling: torch.Tensor,
654 w1_grad_tiling: torch.Tensor,
655 w2_grad_tiling: torch.Tensor,
656 swiglu_grad_tiling: torch.Tensor,
657 gmm_workspace: torch.Tensor,
658 swiglu_grad_workspace: torch.Tensor,
659 runtime_config: torch.Tensor,
660 all_event_counters: torch.Tensor,
661 rank_id: int,
662 ep: int,
663 expert_num: int,
664 hidden_size: int,
665 seq_size: int,
666) -> None:
667 """Launch the public MoE-FFN backward operator.
669 Tensor arguments describe dispatch/combine buffers, saved activations,
670 output gradients, weights, tiling data, workspaces, RuntimeConfig and
671 event counters. Integer arguments describe the rank-local topology and
672 shape. Output and workspace tensors must be pre-allocated; the operator
673 writes gradients and communication results in place.
675 Args:
676 dispatch_target: First tensor in the fixed flat backward ABI; remaining tensor and topology arguments retain
677 their registered schema order.
678 """
679 _load_native()
680 torch.ops.hyper_multicore.mega_moe_grad(
681 dispatch_target,
682 dispatch_target_off,
683 dy,
684 dispatch_src_off,
685 dispatch_size,
686 hidden,
687 hidden_dw,
688 w2,
689 act_grad_y,
690 gate,
691 grad_gate,
692 w1,
693 gate_dx,
694 grad_x,
695 combine_target_off,
696 combine_src_off,
697 combine_size,
698 permute_out,
699 gate_dw,
700 group_list,
701 act_grad_tiling,
702 gate_grad_tiling,
703 w1_grad_tiling,
704 w2_grad_tiling,
705 swiglu_grad_tiling,
706 gmm_workspace,
707 swiglu_grad_workspace,
708 runtime_config,
709 all_event_counters,
710 all_event_counters,
711 rank_id,
712 ep,
713 expert_num,
714 hidden_size,
715 seq_size,
716 )
719def mega_moe_with_diagnostic_buffer(
720 dispatch_target: torch.Tensor,
721 dispatch_target_off: torch.Tensor,
722 dispatch_src: torch.Tensor,
723 dispatch_src_off: torch.Tensor,
724 dispatch_size: torch.Tensor,
725 up_proj_weight: torch.Tensor,
726 up_proj_glist: torch.Tensor,
727 up_proj_y: torch.Tensor,
728 swiglu_out: torch.Tensor,
729 down_proj_weight: torch.Tensor,
730 down_proj_glist: torch.Tensor,
731 down_proj_y: torch.Tensor,
732 combine_target: torch.Tensor,
733 combine_target_off: torch.Tensor,
734 combine_src_off: torch.Tensor,
735 combine_size: torch.Tensor,
736 gmm_workspace: torch.Tensor,
737 up_proj_tiling: torch.Tensor,
738 swiglu_tiling: torch.Tensor,
739 down_proj_tiling: torch.Tensor,
740 runtime_config: torch.Tensor,
741 all_event_counters: torch.Tensor,
742 diagnostic_buffer: torch.Tensor,
743 rank_id: int,
744 ep: int,
745 expert_num: int,
746 hidden_size: int,
747 seq_size: int,
748) -> None:
749 """
750 Launch the internal forward ABI with a shared diagnostic-buffer argument.
752 Args:
753 dispatch_target: First operator tensor; the detailed ABI is documented below.
755 Writes in-place to: dispatch_target, up_proj_y, swiglu_out, down_proj_y,
756 combine_target.
757 All output tensors must be pre-allocated with correct shapes.
759 Args:
760 dispatch_target: First tensor in the fixed diagnostic-forward ABI; the complete parameter groups are documented
761 below and retain their registered schema order.
763 Parameters
764 ----------
765 dispatch_target, dispatch_target_off, dispatch_src, dispatch_src_off,
766 dispatch_size :
767 AllToAll dispatch buffers — dispatch_target written in-place.
768 up_proj_weight, up_proj_glist :
769 Expert weight and cumulative group sizes for GMM1 (up-projection).
770 up_proj_y, swiglu_out :
771 GMM1 output and SwiGLU output — written in-place.
772 down_proj_weight, down_proj_glist, down_proj_y :
773 Expert weight, cumulative group sizes, and output for GMM2 (down-projection).
774 combine_target, combine_target_off, combine_src_off, combine_size :
775 AllToAll combine buffers — combine_target written in-place.
776 gmm_workspace, up_proj_tiling, swiglu_tiling, down_proj_tiling :
777 Pre-computed tiling tensors (from gen_runtime_data.py).
778 runtime_config :
779 Per-rank runtime config tensor (from gen_runtime_data.py).
780 all_event_counters :
781 Event synchronization counter tensor.
782 diagnostic_buffer :
783 Ordinary NPU memory receiving either profiling records or L1 dump
784 snapshots, selected by runtime_config. Unused in ordinary/L0 mode.
785 rank_id, ep, expert_num, hidden_size, seq_size :
786 Topology / shape attributes.
787 """
788 _load_native()
789 torch.ops.hyper_multicore.mega_moe(
790 dispatch_target,
791 dispatch_target_off,
792 dispatch_src,
793 dispatch_src_off,
794 dispatch_size,
795 up_proj_weight,
796 up_proj_glist,
797 up_proj_y,
798 swiglu_out,
799 down_proj_weight,
800 down_proj_glist,
801 down_proj_y,
802 combine_target,
803 combine_target_off,
804 combine_src_off,
805 combine_size,
806 gmm_workspace,
807 up_proj_tiling,
808 swiglu_tiling,
809 down_proj_tiling,
810 runtime_config,
811 all_event_counters,
812 diagnostic_buffer,
813 rank_id,
814 ep,
815 expert_num,
816 hidden_size,
817 seq_size,
818 )
821def mega_moe_grad_with_diagnostic_buffer(
822 dispatch_target: torch.Tensor,
823 dispatch_target_off: torch.Tensor,
824 dy: torch.Tensor,
825 dispatch_src_off: torch.Tensor,
826 dispatch_size: torch.Tensor,
827 hidden: torch.Tensor,
828 hidden_dw: torch.Tensor,
829 w2: torch.Tensor,
830 act_grad_y: torch.Tensor,
831 gate: torch.Tensor,
832 grad_gate: torch.Tensor,
833 w1: torch.Tensor,
834 gate_dx: torch.Tensor,
835 grad_x: torch.Tensor,
836 combine_target_off: torch.Tensor,
837 combine_src_off: torch.Tensor,
838 combine_size: torch.Tensor,
839 permute_out: torch.Tensor,
840 gate_dw: torch.Tensor,
841 group_list: torch.Tensor,
842 act_grad_tiling: torch.Tensor,
843 gate_grad_tiling: torch.Tensor,
844 w1_grad_tiling: torch.Tensor,
845 w2_grad_tiling: torch.Tensor,
846 swiglu_grad_tiling: torch.Tensor,
847 gmm_workspace: torch.Tensor,
848 swiglu_grad_workspace: torch.Tensor,
849 runtime_config: torch.Tensor,
850 all_event_counters: torch.Tensor,
851 diagnostic_buffer: torch.Tensor,
852 rank_id: int,
853 ep: int,
854 expert_num: int,
855 hidden_size: int,
856 seq_size: int,
857) -> None:
858 """Launch the diagnostic MoE-FFN backward ABI.
860 The argument contract matches :func:`mega_moe_grad` with one additional
861 ordinary-NPU-memory ``diagnostic_buffer``. Device workers write profiling
862 records or L1 dump snapshots according to runtime_config; ordinary/L0 mode
863 does not access the buffer. All output and workspace tensors remain
864 caller-owned and are written in place.
866 Args:
867 dispatch_target: First tensor in the fixed diagnostic-backward ABI; remaining tensor and topology arguments
868 retain their registered schema order.
869 """
870 _load_native()
871 torch.ops.hyper_multicore.mega_moe_grad(
872 dispatch_target,
873 dispatch_target_off,
874 dy,
875 dispatch_src_off,
876 dispatch_size,
877 hidden,
878 hidden_dw,
879 w2,
880 act_grad_y,
881 gate,
882 grad_gate,
883 w1,
884 gate_dx,
885 grad_x,
886 combine_target_off,
887 combine_src_off,
888 combine_size,
889 permute_out,
890 gate_dw,
891 group_list,
892 act_grad_tiling,
893 gate_grad_tiling,
894 w1_grad_tiling,
895 w2_grad_tiling,
896 swiglu_grad_tiling,
897 gmm_workspace,
898 swiglu_grad_workspace,
899 runtime_config,
900 all_event_counters,
901 diagnostic_buffer,
902 rank_id,
903 ep,
904 expert_num,
905 hidden_size,
906 seq_size,
907 )
910def moe_token_permute_grad(
911 grad_permuted_tokens: torch.Tensor,
912 sorted_indices: torch.Tensor,
913 num_tokens: int,
914 top_k: int,
915) -> torch.Tensor:
916 """Reduce dropless permuted gradients without retaining original token values.
918 Args:
919 grad_permuted_tokens: Expert-major gradients with shape [T * K, H].
920 sorted_indices: Int32 inverse mapping returned by token permutation.
921 num_tokens: Original token count T.
922 top_k: Number of routes K per token.
924 Returns:
925 Owned [T, H] gradients with the input gradient's dtype and device.
926 """
927 _load_native()
928 return torch.ops.hyper_multicore.moe_token_permute_grad(
929 grad_permuted_tokens, sorted_indices, num_tokens, top_k
930 )
933def _mega_gate_route(
934 logits: torch.Tensor,
935 text_bias: torch.Tensor,
936 vision_bias: torch.Tensor,
937 image_mask: torch.Tensor,
938 runtime_config: torch.Tensor,
939 diagnostic_buffer: torch.Tensor,
940 *,
941 top_k: int,
942 routed_scaling_factor: float,
943 use_vision_bias: bool,
944 diagnostics: MegaKernelDiagnostics | None = None,
945) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
946 """Run HyperMegaGateRoute with text-only or per-token vision bias.
948 Args:
949 logits: Borrowed contiguous FP32 ``[tokens, expert_count]`` tensor.
950 text_bias: Borrowed contiguous FP32 ``[expert_count]`` text selection bias.
951 vision_bias: Borrowed contiguous FP32 ``[expert_count]`` vision selection bias.
952 image_mask: Borrowed contiguous BOOL token mask, or the shared ``[1]`` placeholder.
953 runtime_config: Internal UINT8 AIV pipeline descriptors.
954 diagnostic_buffer: Internal UINT8 cycle-record buffer or disabled placeholder.
955 top_k: Number of selected experts in ``[1, expert_count]``.
956 routed_scaling_factor: Finite factor applied to routed weights.
957 use_vision_bias: Whether ``image_mask`` selects between text and vision bias.
959 Returns:
960 FP32 routing weights, route scores, selected scores and normalization
961 denominator, plus INT64 expert indices. The final three FP32 tensors
962 are internal forward state consumed by MegaGate backward.
963 Completion follows current-stream order.
964 """
965 _load_native()
966 if diagnostics is not None:
967 tokens, experts = logits.shape
968 outputs = (
969 torch.empty((tokens, top_k), dtype=logits.dtype, device=logits.device),
970 torch.empty((tokens, top_k), dtype=torch.int64, device=logits.device),
971 torch.empty((tokens, experts), dtype=logits.dtype, device=logits.device),
972 torch.empty((tokens, top_k), dtype=logits.dtype, device=logits.device),
973 torch.empty((tokens, 1), dtype=logits.dtype, device=logits.device),
974 )
975 parameters = (logits, text_bias, vision_bias, image_mask, runtime_config, diagnostic_buffer, *outputs)
976 with diagnostics.launch(parameters, direction="forward") as call:
977 return torch.ops.hyper_multicore.mega_gate_route_out(
978 logits, text_bias, vision_bias, image_mask, call.runtime_config, call.diagnostic_buffer,
979 *outputs, top_k, routed_scaling_factor, use_vision_bias,
980 )
981 return torch.ops.hyper_multicore.mega_gate_route(
982 logits,
983 text_bias,
984 vision_bias,
985 image_mask,
986 runtime_config,
987 diagnostic_buffer,
988 top_k,
989 routed_scaling_factor,
990 use_vision_bias,
991 )
994def _mega_gate_route_grad(
995 logits: torch.Tensor,
996 route_scores: torch.Tensor,
997 selected_scores: torch.Tensor,
998 normalization_denominator: torch.Tensor,
999 expert_indices: torch.Tensor,
1000 grad_routing_weights: torch.Tensor,
1001 runtime_config: torch.Tensor,
1002 diagnostic_buffer: torch.Tensor,
1003 *,
1004 top_k: int,
1005 routed_scaling_factor: float,
1006) -> torch.Tensor:
1007 """Differentiate the sqrt-softplus Route with RouteGrad and CANN post kernels.
1009 All tensors are borrowed, contiguous, and complete in current-stream order.
1010 """
1011 _load_native()
1012 return torch.ops.hyper_multicore.mega_gate_route_grad(
1013 logits,
1014 route_scores,
1015 selected_scores,
1016 normalization_denominator,
1017 expert_indices,
1018 grad_routing_weights,
1019 runtime_config,
1020 diagnostic_buffer,
1021 top_k,
1022 routed_scaling_factor,
1023 )