Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / components / quantization / quantizers / mxfp8.py: 84%

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 quantizer for NPU dynamic block quantization.""" 

16 

17from typing import Optional 

18 

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

20 

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

22 MXFP8NpuOps, 

23 _get_npu_ops, 

24) 

25from hyper_parallel.components.quantization.quantizers.base import Quantizer 

26from hyper_parallel.components.quantization.tensor import ( 

27 MXFP8Tensor, 

28 QuantizedTensor, 

29) 

30 

31 

32class MXFP8Quantizer(Quantizer): 

33 """Build E4M3 data and E8M0 scales for MX block computation.""" 

34 

35 def __init__( 

36 self, 

37 quant_dtype: Optional[torch.dtype] = None, 

38 *, 

39 npu_ops: Optional[MXFP8NpuOps] = None, 

40 ) -> None: 

41 """Create a quantizer with lazy NPU operator resolution.""" 

42 

43 self.quant_dtype = quant_dtype or torch.float8_e4m3fn 

44 self._npu_ops = npu_ops 

45 

46 @property 

47 def npu_ops(self) -> MXFP8NpuOps: 

48 """Return the injected or lazily resolved NPU operator adapter.""" 

49 

50 if self._npu_ops is None: 

51 self._npu_ops = _get_npu_ops() 

52 return self._npu_ops 

53 

54 def quantize( 

55 self, 

56 tensor: torch.Tensor, 

57 *, 

58 group_list: Optional[torch.Tensor] = None, 

59 group_list_type: int = 0, 

60 rowwise: bool = False, 

61 colwise: bool = False, 

62 ) -> MXFP8Tensor: 

63 """Quantize a tensor into the requested MXFP8 directions.""" 

64 

65 if not isinstance(tensor, torch.Tensor) or isinstance( 

66 tensor, 

67 QuantizedTensor, 

68 ): 

69 raise TypeError( 

70 "MXFP8Quantizer expects a high-precision torch.Tensor." 

71 ) 

72 if tensor.ndim < 2: 

73 raise ValueError( 

74 "MXFP8Quantizer expects at least two dimensions, but got " 

75 f"shape {tuple(tensor.shape)}." 

76 ) 

77 if not rowwise and not colwise: 

78 raise ValueError( 

79 "MXFP8Quantizer requires rowwise=True or colwise=True." 

80 ) 

81 if group_list_type not in (0, 1): 

82 raise ValueError( 

83 "MXFP8Quantizer group_list_type must be 0 or 1, " 

84 f"but got {group_list_type}." 

85 ) 

86 if group_list is not None and ( 

87 not isinstance(group_list, torch.Tensor) or group_list.ndim != 1 

88 ): 

89 raise ValueError( 

90 "MXFP8Quantizer group_list must be a one-dimensional Tensor." 

91 ) 

92 

93 row_data = None 

94 row_scale = None 

95 col_data = None 

96 col_scale = None 

97 if group_list is not None: 

98 if not colwise: 

99 raise ValueError( 

100 "Grouped MXFP8 quantization requires colwise=True." 

101 ) 

102 quant_group_list = ( 

103 group_list 

104 if group_list_type == 0 

105 else torch.cumsum(group_list, dim=0) 

106 ) 

107 if rowwise: 

108 row_data, row_scale = self.npu_ops.dynamic_mx_quant( 

109 tensor, 

110 axis=-1, 

111 quant_dtype=self.quant_dtype, 

112 ) 

113 col_data, col_scale = self.npu_ops.grouped_dynamic_mx_quant( 

114 tensor, 

115 quant_group_list, 

116 quant_dtype=self.quant_dtype, 

117 ) 

118 elif rowwise and colwise: 

119 ( 

120 row_data, 

121 row_scale, 

122 col_data, 

123 col_scale, 

124 ) = self.npu_ops.dynamic_mx_quant_dual_axis( 

125 tensor, 

126 quant_dtype=self.quant_dtype, 

127 ) 

128 elif rowwise: 

129 row_data, row_scale = self.npu_ops.dynamic_mx_quant( 

130 tensor, 

131 axis=-1, 

132 quant_dtype=self.quant_dtype, 

133 ) 

134 elif colwise: 

135 col_data, col_scale = self.npu_ops.dynamic_mx_quant( 

136 tensor, 

137 axis=-2, 

138 quant_dtype=self.quant_dtype, 

139 ) 

140 return MXFP8Tensor( 

141 shape=tensor.shape, 

142 dtype=tensor.dtype, 

143 quantizer=self, 

144 row_data=row_data, 

145 row_scale=row_scale, 

146 col_data=col_data, 

147 col_scale=col_scale, 

148 device=tensor.device, 

149 )