Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / components / quantization / tensor / mxfp8_tensor.py: 81%

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 directional storage and logical Tensor wrapper.""" 

16 

17from typing import TYPE_CHECKING, Any, Callable, Iterable, Optional 

18 

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

20 

21from hyper_parallel.components.quantization.tensor.base import ( 

22 QuantizedTensorStorage, 

23) 

24from hyper_parallel.components.quantization.tensor.quantized_tensor import ( 

25 QuantizedTensor, 

26) 

27 

28if TYPE_CHECKING: 

29 from hyper_parallel.components.quantization.quantizers.mxfp8 import ( 

30 MXFP8Quantizer, 

31 ) 

32 

33 

34class MXFP8TensorStorage(QuantizedTensorStorage): 

35 """Store row-wise and column-wise MXFP8 data with paired scales.""" 

36 

37 def __init__( 

38 self, 

39 shape: Iterable[int], 

40 dtype: torch.dtype, 

41 *, 

42 quantizer: "MXFP8Quantizer", 

43 row_data: Optional[torch.Tensor] = None, 

44 row_scale: Optional[torch.Tensor] = None, 

45 col_data: Optional[torch.Tensor] = None, 

46 col_scale: Optional[torch.Tensor] = None, 

47 ) -> None: 

48 """Initialize one typed MXFP8 representation.""" 

49 

50 self._validate_pair("row", row_data, row_scale) 

51 self._validate_pair("column", col_data, col_scale) 

52 if row_data is None and col_data is None: 

53 raise ValueError( 

54 "MXFP8TensorStorage requires a row-wise or column-wise " 

55 "representation." 

56 ) 

57 if not isinstance(self, torch.Tensor): 

58 self.shape = tuple(shape) 

59 self.dtype = dtype 

60 self.quantizer = quantizer 

61 self.row_data = row_data 

62 self.row_scale = row_scale 

63 self.col_data = col_data 

64 self.col_scale = col_scale 

65 

66 @staticmethod 

67 def _validate_pair( 

68 direction: str, 

69 data: Optional[torch.Tensor], 

70 scale: Optional[torch.Tensor], 

71 ) -> None: 

72 if (data is None) != (scale is None): 

73 raise ValueError( 

74 f"MXFP8 {direction}-wise data and scale must either both be " 

75 "present or both be absent." 

76 ) 

77 

78 def update_usage(self, rowwise: bool = True, colwise: bool = True) -> None: 

79 """Release directional data that subsequent computation will not use.""" 

80 

81 if not rowwise: 

82 self.row_data = None 

83 self.row_scale = None 

84 if not colwise: 

85 self.col_data = None 

86 self.col_scale = None 

87 

88 def is_rowwise(self) -> bool: 

89 """Return whether row-wise data and scale are available.""" 

90 

91 return self.row_data is not None and self.row_scale is not None 

92 

93 def is_colwise(self) -> bool: 

94 """Return whether column-wise data and scale are available.""" 

95 

96 return self.col_data is not None and self.col_scale is not None 

97 

98 def get_metadata(self) -> dict[str, Any]: 

99 """Return directional storage and quantizer metadata.""" 

100 

101 return { 

102 "row_data": self.row_data, 

103 "row_scale": self.row_scale, 

104 "col_data": self.col_data, 

105 "col_scale": self.col_scale, 

106 "quantizer": self.quantizer, 

107 } 

108 

109 

110class MXFP8Tensor(MXFP8TensorStorage, QuantizedTensor): 

111 """Expose MXFP8 storage as a logical Tensor with the original dtype.""" 

112 

113 def __new__( 

114 cls, 

115 shape: Iterable[int], 

116 dtype: torch.dtype, 

117 *, 

118 quantizer: "MXFP8Quantizer", 

119 row_data: Optional[torch.Tensor] = None, 

120 row_scale: Optional[torch.Tensor] = None, 

121 col_data: Optional[torch.Tensor] = None, 

122 col_scale: Optional[torch.Tensor] = None, 

123 device: Optional[torch.device] = None, 

124 requires_grad: bool = False, 

125 ) -> "MXFP8Tensor": 

126 """Create an MXFP8 wrapper from at least one directional pair.""" 

127 

128 del quantizer, row_scale, col_scale 

129 physical_data = row_data if row_data is not None else col_data 

130 if physical_data is None: 

131 raise ValueError("MXFP8Tensor requires row_data or col_data.") 

132 wrapper_device = physical_data.device if device is None else device 

133 return QuantizedTensor.__new__( 

134 cls, 

135 shape, 

136 dtype, 

137 device=wrapper_device, 

138 requires_grad=requires_grad, 

139 ) 

140 

141 def __init__( 

142 self, 

143 shape: Iterable[int], 

144 dtype: torch.dtype, 

145 *, 

146 quantizer: "MXFP8Quantizer", 

147 row_data: Optional[torch.Tensor] = None, 

148 row_scale: Optional[torch.Tensor] = None, 

149 col_data: Optional[torch.Tensor] = None, 

150 col_scale: Optional[torch.Tensor] = None, 

151 device: Optional[torch.device] = None, 

152 requires_grad: bool = False, 

153 ) -> None: 

154 """Initialize directional storage after wrapper allocation.""" 

155 

156 del device, requires_grad 

157 MXFP8TensorStorage.__init__( 

158 self, 

159 shape, 

160 dtype, 

161 quantizer=quantizer, 

162 row_data=row_data, 

163 row_scale=row_scale, 

164 col_data=col_data, 

165 col_scale=col_scale, 

166 ) 

167 

168 def _new_with_transformed_data( 

169 self, 

170 transform: Callable[[torch.Tensor], torch.Tensor], 

171 *, 

172 requires_grad: bool, 

173 ) -> "MXFP8Tensor": 

174 def _transform_optional( 

175 value: Optional[torch.Tensor], 

176 ) -> Optional[torch.Tensor]: 

177 return None if value is None else transform(value) 

178 

179 return type(self)( 

180 shape=self.shape, 

181 dtype=self.dtype, 

182 quantizer=self.quantizer, 

183 row_data=_transform_optional(self.row_data), 

184 row_scale=_transform_optional(self.row_scale), 

185 col_data=_transform_optional(self.col_data), 

186 col_scale=_transform_optional(self.col_scale), 

187 requires_grad=requires_grad, 

188 ) 

189 

190 def __repr__(self) -> str: 

191 """Return logical shape and directional storage state.""" 

192 

193 return ( 

194 "MXFP8Tensor(" 

195 f"shape={tuple(self.shape)}, dtype={self.dtype}, " 

196 f"device={self.device}, rowwise={self.is_rowwise()}, " 

197 f"colwise={self.is_colwise()})" 

198 )