Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / components / quantization / functional / w4a8_gmm_func.py: 40%

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"""W4A8 format hooks for the shared native grouped-linear flow.""" 

16 

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

18 

19from hyper_parallel.components.quantization.functional.base_gmm_func import ( 

20 GroupedLinear, 

21) 

22from hyper_parallel.components.quantization.quantizers.w4a8 import ( 

23 W4A8Quantizer, 

24) 

25 

26 

27class W4A8GroupedLinear(GroupedLinear): 

28 """Bind W4A8's target-layout views and native GMM lowering.""" 

29 

30 FORMAT_NAME = "Native W4A8" 

31 

32 def __init__(self, block_size: int = 32) -> None: 

33 """Create one reusable W4A8 strategy and own its quantizer.""" 

34 

35 self.quantizer = W4A8Quantizer(block_size) 

36 self.block_size = self.quantizer.block_size 

37 

38 def normalize_group_list( 

39 self, 

40 group_list: torch.Tensor, 

41 group_list_type: int, 

42 ) -> tuple[torch.Tensor, int]: 

43 if group_list_type == 1: 

44 return torch.cumsum(group_list, dim=0), 0 

45 return group_list, 0 

46 

47 def output_features(self, weight: torch.Tensor) -> int: 

48 return weight.shape[-2] 

49 

50 def weight_quantization_directions( 

51 self, 

52 needs_grad_input: bool, 

53 ) -> tuple[bool, bool]: 

54 # W4A8 forward consumes the row-wise view (GMM-ready ``[E, K, N]``); 

55 # dgrad consumes the transposed column-wise view. 

56 return True, needs_grad_input 

57 

58 def retain_weight_backward(self, weight_quant, needs_grad_input: bool) -> None: 

59 weight_quant.update_usage(rowwise=False, colwise=needs_grad_input) 

60 

61 def quantize_input( 

62 self, 

63 inputs: torch.Tensor, 

64 *, 

65 rowwise: bool, 

66 colwise: bool, 

67 group_list: torch.Tensor, 

68 group_list_type: int, 

69 ): 

70 return self.quantizer.quantize_activation( 

71 inputs, 

72 rowwise=rowwise, 

73 colwise=colwise, 

74 group_list=group_list, 

75 group_list_type=group_list_type, 

76 ) 

77 

78 def quantize_weight( 

79 self, 

80 weight: torch.Tensor, 

81 *, 

82 rowwise: bool, 

83 colwise: bool, 

84 ): 

85 return self.quantizer.quantize_weight( 

86 weight, 

87 rowwise=rowwise, 

88 colwise=colwise, 

89 ) 

90 

91 def quantize_grad_output( 

92 self, 

93 grad_output: torch.Tensor, 

94 *, 

95 rowwise: bool, 

96 colwise: bool, 

97 group_list: torch.Tensor, 

98 group_list_type: int, 

99 ): 

100 return self.quantizer.quantize_activation( 

101 grad_output, 

102 rowwise=rowwise, 

103 colwise=colwise, 

104 group_list=group_list, 

105 group_list_type=group_list_type, 

106 ) 

107 

108 @staticmethod 

109 def _select_operands(left, right, layout: str): 

110 if layout == "NN": 

111 left_data, left_scale = left.row_data, left.row_scale 

112 right_data, right_scale = right.row_data, right.row_scale 

113 elif layout == "NT": 

114 left_data, left_scale = left.row_data, left.row_scale 

115 right_data, right_scale = right.col_data, right.col_scale 

116 elif layout == "TN": 

117 # Preserve transpose strides so the native MX kernel sees the 

118 # same transpose state on the operand and its scale tensor. 

119 left_data = left.col_data.transpose(0, 1) 

120 left_scale = left.col_scale.transpose(0, 1) 

121 right_data, right_scale = right.col_data, right.col_scale 

122 else: 

123 raise ValueError( 

124 f"Unsupported W4A8 GMM layout {layout!r}; expected NN, NT, or TN." 

125 ) 

126 if any(value is None for value in (left_data, left_scale, right_data, right_scale)): 

127 raise ValueError( 

128 f"W4A8 GMM layout {layout!r} requires both directional operands." 

129 ) 

130 return left_data, left_scale, right_data, right_scale 

131 

132 def grouped_matmul( 

133 self, 

134 left, 

135 right, 

136 *, 

137 layout: str, 

138 group_list: torch.Tensor, 

139 group_type: int, 

140 group_list_type: int, 

141 output_dtype: torch.dtype, 

142 ) -> torch.Tensor: 

143 del group_list_type 

144 left_data, left_scale, right_data, right_scale = self._select_operands( 

145 left, 

146 right, 

147 layout, 

148 ) 

149 # NN/NT are native MXFP8 x Packed-MXFP4 W4A8. TN computes the dense 

150 # master-weight gradient and therefore remains MXFP8 x MXFP8. 

151 native_w4 = layout in ("NN", "NT") 

152 use_e8m0_scale = self.quantizer.block_size == 32 or layout == "TN" 

153 return self.quantizer.npu_ops.grouped_matmul( 

154 left_data, 

155 right_data, 

156 right_scale, 

157 left_scale=left_scale, 

158 group_list=group_list, 

159 group_type=group_type, 

160 output_dtype=output_dtype, 

161 use_e8m0_scale=use_e8m0_scale, 

162 split_item=3 if use_e8m0_scale else 2, 

163 native_w4=native_w4, 

164 ) 

165 

166__all__ = ["W4A8GroupedLinear"]