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

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"""Ascend grouped-GEMM implementation for packed SwiGLU experts.""" 

16 

17from typing import Any 

18 

19 

20_NPU_GROUPED_MATMUL = None 

21 

22 

23def _get_npu_grouped_matmul() -> Any: 

24 """Build and cache the autograd wrapper after the Torch backend is selected.""" 

25 global _NPU_GROUPED_MATMUL # pylint: disable=global-statement 

26 if _NPU_GROUPED_MATMUL is not None: 

27 return _NPU_GROUPED_MATMUL 

28 

29 import torch # pylint: disable=C0415 

30 import torch_npu # pylint: disable=C0415 

31 

32 class _NPUGroupedMatmul(torch.autograd.Function): # pylint: disable=abstract-method 

33 """Autograd wrapper around ``torch_npu.npu_grouped_matmul``.""" 

34 

35 @staticmethod 

36 def forward(ctx: Any, inputs: Any, weight: Any, group_list: Any) -> Any: 

37 """Run one grouped matrix multiplication and retain backward inputs.""" 

38 ctx.save_for_backward(inputs, weight, group_list) 

39 return torch_npu.npu_grouped_matmul( 

40 x=[inputs], 

41 weight=[weight], 

42 bias=[], 

43 group_list=group_list, 

44 split_item=3, 

45 group_type=0, 

46 group_list_type=0, 

47 )[0] 

48 

49 @staticmethod 

50 def backward(ctx: Any, grad_output: Any) -> tuple[Any, Any, None]: 

51 """Compute grouped input and weight gradients.""" 

52 inputs, weight, group_list = ctx.saved_tensors 

53 grad_output = grad_output.contiguous() 

54 grad_inputs = torch_npu.npu_grouped_matmul( 

55 x=[grad_output], 

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

57 bias=[], 

58 group_list=group_list, 

59 split_item=3, 

60 group_type=0, 

61 group_list_type=0, 

62 )[0] 

63 numel = inputs.numel() 

64 if ( 

65 not torch.compiler.is_compiling() 

66 and isinstance(numel, int) 

67 and numel == 0 

68 ): 

69 grad_weight = torch.zeros_like(weight) 

70 else: 

71 grad_weight = torch_npu.npu_grouped_matmul( 

72 x=[inputs.transpose(0, 1)], 

73 weight=[grad_output], 

74 bias=[], 

75 group_list=group_list, 

76 split_item=2, 

77 group_type=2, 

78 group_list_type=0, 

79 )[0] 

80 return grad_inputs, grad_weight, None 

81 

82 _NPU_GROUPED_MATMUL = _NPUGroupedMatmul 

83 return _NPU_GROUPED_MATMUL 

84 

85 

86def npu_grouped_swiglu( 

87 inputs: Any, 

88 gate_up_weight: Any, 

89 down_weight: Any, 

90 tokens_per_expert: Any, 

91) -> Any: 

92 """Run packed local SwiGLU experts with two grouped GEMMs. 

93 

94 Args: 

95 inputs: Expert-major routed tokens with shape ``[T, H]``. 

96 gate_up_weight: Packed projection weights with shape ``[E, 2I, H]``. 

97 down_weight: Down projection weights with shape ``[E, H, I]``. 

98 tokens_per_expert: Token count for each local expert, shape ``[E]``. 

99 

100 Returns: 

101 Expert-major output with shape ``[T, H]``. 

102 

103 Raises: 

104 ValueError: If called for a non-NPU tensor or incompatible shapes. 

105 """ 

106 if inputs.device.type != "npu": 

107 raise ValueError("npu_grouped_swiglu requires inputs on an Ascend NPU") 

108 if gate_up_weight.ndim != 3 or down_weight.ndim != 3: 

109 raise ValueError("grouped SwiGLU weights must be three-dimensional") 

110 if gate_up_weight.shape[0] != tokens_per_expert.numel(): 

111 raise ValueError("tokens_per_expert must contain one count per local expert") 

112 

113 import torch # pylint: disable=C0415 

114 import torch_npu # pylint: disable=C0415 

115 

116 grouped_matmul = _get_npu_grouped_matmul() 

117 group_list = torch.cumsum(tokens_per_expert.to(torch.int64), dim=0) 

118 gate_up = grouped_matmul.apply( 

119 inputs, 

120 gate_up_weight.transpose(1, 2), 

121 group_list, 

122 ) 

123 intermediate = torch_npu.npu_swiglu(gate_up, dim=-1) 

124 return grouped_matmul.apply( 

125 intermediate, 

126 down_weight.transpose(1, 2), 

127 group_list, 

128 )