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.config import ( 

22 GROUP_LIST_COUNTS, 

23 GROUP_LIST_OFFSETS, 

24) 

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

26 MXFP8NpuOps, 

27 _get_npu_ops, 

28) 

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

30from hyper_parallel.components.quantization.tensor import ( 

31 MXFP8Tensor, 

32 QuantizedTensor, 

33) 

34 

35 

36class MXFP8Quantizer(Quantizer): 

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

38 

39 def __init__( 

40 self, 

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

42 *, 

43 npu_ops: Optional[MXFP8NpuOps] = None, 

44 ) -> None: 

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

46 

47 self.quant_dtype = quant_dtype or torch.float8_e4m3fn 

48 self._npu_ops = npu_ops 

49 

50 @property 

51 def npu_ops(self) -> MXFP8NpuOps: 

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

53 

54 if self._npu_ops is None: 

55 self._npu_ops = _get_npu_ops() 

56 return self._npu_ops 

57 

58 def quantize( 

59 self, 

60 tensor: torch.Tensor, 

61 *, 

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

63 group_list_type: int = GROUP_LIST_OFFSETS, 

64 rowwise: bool = False, 

65 colwise: bool = False, 

66 ) -> MXFP8Tensor: 

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

68 

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

70 tensor, 

71 QuantizedTensor, 

72 ): 

73 raise TypeError( 

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

75 ) 

76 if tensor.ndim < 2: 

77 raise ValueError( 

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

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

80 ) 

81 if not rowwise and not colwise: 

82 raise ValueError( 

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

84 ) 

85 if group_list_type not in (GROUP_LIST_OFFSETS, GROUP_LIST_COUNTS): 

86 raise ValueError( 

87 "MXFP8Quantizer group_list_type must be 0 or 1, " 

88 f"but got {group_list_type}." 

89 ) 

90 if group_list is not None and ( 

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

92 ): 

93 raise ValueError( 

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

95 ) 

96 

97 row_data = None 

98 row_scale = None 

99 col_data = None 

100 col_scale = None 

101 if group_list is not None: 

102 if not colwise: 

103 raise ValueError( 

104 "Grouped MXFP8 quantization requires colwise=True." 

105 ) 

106 quant_group_list = ( 

107 group_list 

108 if group_list_type == GROUP_LIST_OFFSETS 

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

110 ) 

111 if rowwise: 

112 row_data, row_scale = self.npu_ops.dynamic_mx_quant( 

113 tensor, 

114 axis=-1, 

115 quant_dtype=self.quant_dtype, 

116 ) 

117 col_data, col_scale = self.npu_ops.grouped_dynamic_mx_quant( 

118 tensor, 

119 quant_group_list, 

120 quant_dtype=self.quant_dtype, 

121 ) 

122 elif rowwise and colwise: 

123 ( 

124 row_data, 

125 row_scale, 

126 col_data, 

127 col_scale, 

128 ) = self.npu_ops.dynamic_mx_quant_dual_axis( 

129 tensor, 

130 quant_dtype=self.quant_dtype, 

131 ) 

132 elif rowwise: 

133 row_data, row_scale = self.npu_ops.dynamic_mx_quant( 

134 tensor, 

135 axis=-1, 

136 quant_dtype=self.quant_dtype, 

137 ) 

138 elif colwise: 

139 col_data, col_scale = self.npu_ops.dynamic_mx_quant( 

140 tensor, 

141 axis=-2, 

142 quant_dtype=self.quant_dtype, 

143 ) 

144 return MXFP8Tensor( 

145 shape=tensor.shape, 

146 dtype=tensor.dtype, 

147 quantizer=self, 

148 row_data=row_data, 

149 row_scale=row_scale, 

150 col_data=col_data, 

151 col_scale=col_scale, 

152 device=tensor.device, 

153 )