Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / components / functional / mhc_post.py: 69%

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"""Manifold hyper-connection post-processing functions.""" 

16 

17from typing import Any, Tuple 

18 

19import torch # pylint: disable=forbidden-backend-import 

20 

21try: 

22 import omni_training_custom_ops # noqa: F401 # pylint: disable=unused-import 

23except ImportError: 

24 omni_training_custom_ops = None 

25 

26 

27class _MhcPost(torch.autograd.Function): 

28 """Autograd bridge for the NPU MHC post custom operator.""" 

29 

30 @staticmethod 

31 def forward( 

32 ctx: Any, 

33 x: torch.Tensor, 

34 h_res: torch.Tensor, 

35 h_out: torch.Tensor, 

36 h_post: torch.Tensor, 

37 ) -> torch.Tensor: 

38 """Run the NPU MHC post forward operator.""" 

39 if omni_training_custom_ops is None: 

40 raise ImportError("MHC AscendC requires omni_training_custom_ops") 

41 ctx.save_for_backward(x, h_res, h_out, h_post) 

42 return torch.ops.custom.npu_ai_infra_manifold_constrained_hyper_connection_post( 

43 x, 

44 h_res, 

45 h_out, 

46 h_post, 

47 ) 

48 

49 @staticmethod 

50 def backward( 

51 ctx: Any, 

52 grad_output: torch.Tensor, 

53 ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: 

54 """Run the NPU MHC post backward operator.""" 

55 x, h_res, h_out, h_post = ctx.saved_tensors 

56 grad_x, grad_h_res, grad_h_out, grad_h_post = torch.ops.custom.npu_ai_infra_mhc_post_grad( 

57 grad_output.contiguous(), 

58 x, 

59 h_res, 

60 h_out, 

61 h_post, 

62 ) 

63 return grad_x, grad_h_res, grad_h_out, grad_h_post 

64 

65 

66def mhc_post( 

67 x: torch.Tensor, 

68 residual: torch.Tensor, 

69 h_post: torch.Tensor, 

70 h_res: torch.Tensor, 

71 num_stream: int, 

72) -> torch.Tensor: 

73 """Combine transformed and residual streams after an MHC-wrapped block. 

74 

75 Args: 

76 x: Output of the wrapped block. 

77 residual: Flattened residual streams. 

78 h_post: Per-stream output mixing coefficients. 

79 h_res: Residual stream mixing matrix. 

80 num_stream: Number of residual streams. 

81 

82 Returns: 

83 Flattened mixed residual streams. 

84 """ 

85 x_shape = x.size() 

86 residual = residual.reshape(x_shape[0], x_shape[1], num_stream, -1) 

87 if x.device.type == "npu" and omni_training_custom_ops is not None: 

88 return _MhcPost.apply(residual, h_res, x, h_post).flatten(2) 

89 output = h_post.unsqueeze(-1) * x.unsqueeze(-2) 

90 output = output + torch.sum( 

91 h_res.unsqueeze(-1) * residual.unsqueeze(-2), dim=2 

92 ) 

93 return output.flatten(2).to(x.dtype)