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

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"""Pangu-compatible directional storage for HiFloat8 tensors.""" 

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

30 HiFloat8Quantizer, 

31 ) 

32 

33 

34class HiFloat8TensorStorage(QuantizedTensorStorage): 

35 """Store Pangu-compatible row/column views of one HiFloat8 payload.""" 

36 

37 shape: Iterable[int] 

38 dtype: torch.dtype 

39 quantizer: "HiFloat8Quantizer" 

40 row_data: Optional[torch.Tensor] 

41 col_data: Optional[torch.Tensor] 

42 row_scale: Optional[torch.Tensor] 

43 col_scale: Optional[torch.Tensor] 

44 scale: Optional[torch.Tensor] 

45 

46 def __init__( 

47 self, 

48 shape: Iterable[int], 

49 dtype: torch.dtype, 

50 *, 

51 quantizer: "HiFloat8Quantizer", 

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

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

54 scale: Optional[torch.Tensor] = None, 

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

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

57 ) -> None: 

58 """Initialize directional fields using Pangu's HiF8 storage contract.""" 

59 

60 self.row_data = row_data 

61 self.col_data = col_data 

62 self.row_scale = row_scale if row_scale is not None else ( 

63 scale if self.row_data is not None else None 

64 ) 

65 self.col_scale = col_scale if col_scale is not None else ( 

66 scale if self.col_data is not None else None 

67 ) 

68 if (self.row_data is None) != (self.row_scale is None): 

69 raise ValueError("HiFloat8 row_data and row_scale must be provided together.") 

70 if (self.col_data is None) != (self.col_scale is None): 

71 raise ValueError("HiFloat8 col_data and col_scale must be provided together.") 

72 self.scale = ( 

73 scale 

74 if scale is not None 

75 else self.row_scale 

76 if self.row_scale is not None 

77 else self.col_scale 

78 ) 

79 if not hasattr(self, "shape"): 

80 self.shape = shape 

81 self.dtype = dtype 

82 self.quantizer = quantizer 

83 

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

85 """Release row/column references independently, matching Pangu. 

86 

87 Args: 

88 rowwise: Keep the row payload and scale when True, release otherwise. 

89 colwise: Keep the column payload and scale when True, release otherwise. 

90 """ 

91 

92 if not rowwise: 

93 self.row_data = None 

94 self.row_scale = None 

95 if not colwise: 

96 self.col_data = None 

97 self.col_scale = None 

98 self.scale = self.row_scale if self.row_scale is not None else self.col_scale 

99 

100 def is_rowwise(self) -> bool: 

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

102 

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

104 

105 def is_colwise(self) -> bool: 

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

107 

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

109 

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

111 """Return the same directional metadata fields as Pangu HiF8Tensor.""" 

112 

113 return { 

114 "row_data": self.row_data, 

115 "row_scale": self.row_scale, 

116 "col_data": self.col_data, 

117 "col_scale": self.col_scale, 

118 "quantizer": self.quantizer, 

119 } 

120 

121 

122class HiFloat8Tensor(HiFloat8TensorStorage, QuantizedTensor): 

123 """Expose Pangu-compatible HiFloat8 storage with a logical source dtype.""" 

124 

125 def __new__( 

126 cls, 

127 shape: Iterable[int], 

128 dtype: torch.dtype, 

129 *, 

130 quantizer: "HiFloat8Quantizer", 

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

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

133 scale: Optional[torch.Tensor] = None, 

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

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

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

137 requires_grad: bool = False, 

138 ) -> "HiFloat8Tensor": 

139 """Create a wrapper from at least one directional HiFloat8 view.""" 

140 

141 del quantizer, scale, row_scale, col_scale 

142 physical_data = row_data if row_data is not None else col_data 

143 if physical_data is None: 

144 raise ValueError("HiFloat8Tensor requires row_data or col_data.") 

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

146 return QuantizedTensor.__new__( 

147 cls, 

148 shape, 

149 dtype, 

150 device=wrapper_device, 

151 requires_grad=requires_grad, 

152 ) 

153 

154 def __init__( 

155 self, 

156 shape: Iterable[int], 

157 dtype: torch.dtype, 

158 *, 

159 quantizer: "HiFloat8Quantizer", 

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

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

162 scale: Optional[torch.Tensor] = None, 

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

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

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

166 requires_grad: bool = False, 

167 ) -> None: 

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

169 

170 del device, requires_grad 

171 QuantizedTensor.__init__(self) 

172 HiFloat8TensorStorage.__init__( 

173 self, 

174 shape, 

175 dtype, 

176 quantizer=quantizer, 

177 row_data=row_data, 

178 col_data=col_data, 

179 scale=scale, 

180 row_scale=row_scale, 

181 col_scale=col_scale, 

182 ) 

183 

184 def _new_with_transformed_data( 

185 self, 

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

187 *, 

188 requires_grad: bool, 

189 ) -> "HiFloat8Tensor": 

190 transformed: dict[int, torch.Tensor] = {} 

191 

192 def _transform_optional( 

193 value: Optional[torch.Tensor], 

194 ) -> Optional[torch.Tensor]: 

195 if value is None: 

196 return None 

197 key = id(value) 

198 if key not in transformed: 

199 transformed[key] = transform(value) 

200 return transformed[key] 

201 

202 return type(self)( 

203 shape=self.shape, 

204 dtype=self.dtype, 

205 quantizer=self.quantizer, 

206 row_data=_transform_optional(self.row_data), 

207 row_scale=_transform_optional(self.row_scale), 

208 col_data=_transform_optional(self.col_data), 

209 col_scale=_transform_optional(self.col_scale), 

210 requires_grad=requires_grad, 

211 ) 

212 

213 def __repr__(self) -> str: 

214 """Return logical shape and directional storage availability.""" 

215 

216 return ( 

217 "HiFloat8Tensor(" 

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

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

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

221 )