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

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"""NPU grouped matrix multiplication function.""" 

16 

17from typing import Any, Optional, Union 

18 

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

20import torch_npu 

21 

22GroupList = Optional[Union[torch.Tensor, list[int], tuple[int, ...]]] 

23 

24 

25class GMMFunction(torch.autograd.Function): 

26 """Autograd bridge for NPU grouped matrix multiplication.""" 

27 

28 @staticmethod 

29 def forward( 

30 ctx: Any, 

31 original_weight: torch.Tensor, 

32 x: torch.Tensor, 

33 weight: torch.Tensor, 

34 bias: Optional[torch.Tensor], 

35 group_args: tuple, 

36 ) -> torch.Tensor: 

37 """Run grouped matrix multiplication and save its backward state.""" 

38 group_list, group_type, gemm_fusion, group_list_type, group_list_data_type = group_args 

39 

40 npu_bias = [bias] if bias is not None else [] 

41 

42 outputs = torch_npu.npu_grouped_matmul( 

43 x=[x], 

44 weight=[weight], 

45 bias=npu_bias, 

46 group_list=group_list, 

47 split_item=3, 

48 group_type=group_type, 

49 group_list_type=group_list_type 

50 ) 

51 

52 if group_list_data_type == 0: 

53 ctx.save_for_backward(x, weight) 

54 ctx.group_list = group_list 

55 else: 

56 ctx.save_for_backward(x, weight, group_list) 

57 

58 ctx.original_weight = original_weight 

59 ctx.gemm_fusion = gemm_fusion 

60 ctx.group_list_type = group_list_type 

61 ctx.group_list_data_type = group_list_data_type 

62 return outputs[0] 

63 

64 @staticmethod 

65 def backward(ctx: Any, grad_outputs: torch.Tensor) -> tuple: 

66 """Compute the input and grouped-weight gradients.""" 

67 grad_outputs = grad_outputs.contiguous() 

68 

69 if ctx.group_list_data_type == 0: 

70 x, weight = ctx.saved_tensors 

71 group_list = ctx.group_list 

72 else: 

73 x, weight, group_list = ctx.saved_tensors 

74 

75 original_weight = ctx.original_weight 

76 

77 if isinstance(group_list, (list, tuple)): 

78 group_list_tensor = torch.tensor(group_list, device=x.device, dtype=torch.int64) 

79 else: 

80 group_list_tensor = group_list 

81 

82 dx_list = torch_npu.npu_grouped_matmul( 

83 x=[grad_outputs], 

84 weight=[weight.transpose(-1, -2)], 

85 bias=[], 

86 group_list=group_list_tensor, 

87 split_item=3, 

88 group_type=0, 

89 group_list_type=ctx.group_list_type 

90 ) 

91 dx = dx_list[0] 

92 

93 grad_weight = None 

94 use_fusion = False 

95 if ctx.gemm_fusion: 

96 if hasattr(original_weight, 'main_grad'): 

97 if original_weight.main_grad is not None: 

98 use_fusion = True 

99 

100 if use_fusion: 

101 n_size = original_weight.main_grad.shape[-1] 

102 main_grad_view = original_weight.main_grad.view(-1, n_size) 

103 

104 torch_npu.npu_grouped_matmul_add_( 

105 main_grad_view, 

106 x, 

107 grad_outputs, 

108 group_list_tensor, 

109 transpose_x=True, 

110 transpose_weight=False, 

111 group_type=2 

112 ) 

113 

114 if getattr(weight, 'zero_out_wgrad', False): 

115 grad_weight = torch.zeros_like(weight) 

116 else: 

117 grad_weight = torch.empty_like(weight) 

118 

119 original_weight.grad_added_to_main_grad = True 

120 else: 

121 numel = x.nelement() 

122 if ( 

123 not torch.compiler.is_compiling() 

124 and isinstance(numel, int) 

125 and numel == 0 

126 ): 

127 grad_weight = torch.zeros_like(weight) 

128 else: 

129 grad_weight = torch_npu.npu_grouped_matmul( 

130 x=[x.transpose(-1, -2)], 

131 weight=[grad_outputs], 

132 bias=[], 

133 group_list=group_list_tensor, 

134 split_item=2, 

135 group_type=2, 

136 group_list_type=ctx.group_list_type 

137 )[0] 

138 

139 return None, dx, grad_weight, None, None 

140 

141 

142def grouped_matmul( 

143 x: torch.Tensor, 

144 weight: torch.Tensor, 

145 bias: Optional[torch.Tensor] = None, 

146 group_list: GroupList = None, 

147 group_type: int = 0, 

148 group_list_type: int = 0, 

149 gemm_fusion: bool = False, 

150 original_weight: Optional[torch.Tensor] = None, 

151) -> torch.Tensor: 

152 """Apply NPU grouped matrix multiplication. 

153 

154 Args: 

155 x: Expert-major input tensor. 

156 weight: Grouped weight tensor. 

157 bias: Optional grouped bias tensor. 

158 group_list: Cumulative group boundaries as a tensor or integer sequence. 

159 group_type: Grouping axis accepted by ``npu_grouped_matmul``. 

160 group_list_type: Group-list representation accepted by the NPU kernel. 

161 gemm_fusion: Whether to accumulate the weight gradient into ``main_grad``. 

162 original_weight: Parameter that owns ``main_grad``. Defaults to ``weight``. 

163 

164 Returns: 

165 The grouped matrix multiplication output. 

166 """ 

167 group_list_data_type = 1 if isinstance(group_list, torch.Tensor) else 0 

168 group_args = (group_list, group_type, gemm_fusion, group_list_type, group_list_data_type) 

169 

170 if original_weight is None: 

171 original_weight = weight 

172 

173 return GMMFunction.apply(original_weight, x, weight, bias, group_args)