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

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

16 

17from typing import Optional 

18 

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

20 

21from hyper_parallel.components.quantization.functional.base_gmm_func import GroupedLinear 

22from hyper_parallel.components.quantization.ops.npu_mxfp8 import ( 

23 mxfp8_grouped_matmul, 

24) 

25from hyper_parallel.components.quantization.quantizers import ( 

26 MXFP8Quantizer, 

27) 

28from hyper_parallel.components.quantization.tensor import QuantizedTensor 

29 

30 

31class MXFP8GroupedLinear(GroupedLinear): 

32 """Bind MXFP8 quantization and GMM lowering to the common native flow.""" 

33 

34 FORMAT_NAME = "MXFP8" 

35 

36 def __init__(self, quantizer: Optional[MXFP8Quantizer] = None) -> None: 

37 """Create one reusable MXFP8 strategy and own its quantizer.""" 

38 

39 self.quantizer = quantizer if quantizer is not None else MXFP8Quantizer() 

40 

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

42 """Return the output width from the live ``[E, O, K]`` weight.""" 

43 return weight.shape[-2] 

44 

45 def normalize_group_list( 

46 self, 

47 group_list: torch.Tensor, 

48 group_list_type: int, 

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

50 """Preserve counts/type 1 or offsets/type 0 for the MXFP8 path. 

51 

52 The quantizer derives cumulative offsets when grouped activation 

53 quantization needs them; the GMM operator receives the original list 

54 together with its matching type tag. 

55 """ 

56 return group_list, group_list_type 

57 

58 @property 

59 def weight_forward_is_rowwise(self) -> bool: 

60 """Use the column-wise weight view in the forward NN GMM.""" 

61 return False 

62 

63 def quantize_input( 

64 self, 

65 inputs: torch.Tensor, 

66 *, 

67 rowwise: bool, 

68 colwise: bool, 

69 group_list: torch.Tensor, 

70 group_list_type: int, 

71 ) -> QuantizedTensor: 

72 """Quantize grouped activations for the MXFP8 input role.""" 

73 return self.quantizer.quantize( 

74 inputs, 

75 group_list=group_list, 

76 group_list_type=group_list_type, 

77 rowwise=rowwise, 

78 colwise=colwise, 

79 ) 

80 

81 def quantize_weight( 

82 self, 

83 weight: torch.Tensor, 

84 *, 

85 rowwise: bool, 

86 colwise: bool, 

87 ) -> QuantizedTensor: 

88 """Quantize expert weights for the MXFP8 weight role.""" 

89 return self.quantizer.quantize( 

90 weight, 

91 rowwise=rowwise, 

92 colwise=colwise, 

93 ) 

94 

95 def quantize_grad_output( 

96 self, 

97 grad_output: torch.Tensor, 

98 *, 

99 rowwise: bool, 

100 colwise: bool, 

101 group_list: torch.Tensor, 

102 group_list_type: int, 

103 ) -> QuantizedTensor: 

104 """Quantize output gradients for the MXFP8 gradient role.""" 

105 return self.quantizer.quantize( 

106 grad_output, 

107 group_list=group_list, 

108 group_list_type=group_list_type, 

109 rowwise=rowwise, 

110 colwise=colwise, 

111 ) 

112 

113 def grouped_matmul( 

114 self, 

115 left: QuantizedTensor, 

116 right: QuantizedTensor, 

117 *, 

118 layout: str, 

119 group_list: torch.Tensor, 

120 group_type: int, 

121 group_list_type: int, 

122 output_dtype: torch.dtype, 

123 ) -> torch.Tensor: 

124 """Execute one MXFP8 grouped matrix multiplication.""" 

125 return mxfp8_grouped_matmul( 

126 left, 

127 right, 

128 layout=layout, 

129 group_list=group_list, 

130 group_type=group_type, 

131 group_list_type=group_list_type, 

132 output_dtype=output_dtype, 

133 ) 

134 

135 

136__all__ = ["MXFP8GroupedLinear"]