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

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 if inputs.numel() == 0: 

64 grad_weight = torch.zeros_like(weight) 

65 else: 

66 grad_weight = torch_npu.npu_grouped_matmul( 

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

68 weight=[grad_output], 

69 bias=[], 

70 group_list=group_list, 

71 split_item=2, 

72 group_type=2, 

73 group_list_type=0, 

74 )[0] 

75 return grad_inputs, grad_weight, None 

76 

77 _NPU_GROUPED_MATMUL = _NPUGroupedMatmul 

78 return _NPU_GROUPED_MATMUL 

79 

80 

81def npu_grouped_swiglu( 

82 inputs: Any, 

83 gate_up_weight: Any, 

84 down_weight: Any, 

85 tokens_per_expert: Any, 

86) -> Any: 

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

88 

89 Args: 

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

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

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

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

94 

95 Returns: 

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

97 

98 Raises: 

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

100 """ 

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

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

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

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

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

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

107 

108 import torch # pylint: disable=C0415 

109 import torch_npu # pylint: disable=C0415 

110 

111 grouped_matmul = _get_npu_grouped_matmul() 

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

113 gate_up = grouped_matmul.apply( 

114 inputs, 

115 gate_up_weight.transpose(1, 2), 

116 group_list, 

117 ) 

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

119 return grouped_matmul.apply( 

120 intermediate, 

121 down_weight.transpose(1, 2), 

122 group_list, 

123 )