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

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"""Directional storage for native W4A8 operands.""" 

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

30 W4A8Quantizer, 

31 ) 

32 

33 

34class W4A8TensorStorage(QuantizedTensorStorage): 

35 """Store native W4A8 data for the directions needed by backward.""" 

36 

37 def __init__( 

38 self, 

39 shape: Iterable[int], 

40 dtype: torch.dtype, 

41 *, 

42 quantizer: "W4A8Quantizer", 

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 directional physical data and scales.""" 

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("W4A8TensorStorage requires row or column data.") 

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

55 self.shape = tuple(shape) 

56 self.dtype = dtype 

57 self.quantizer = quantizer 

58 self.row_data = row_data 

59 self.row_scale = row_scale 

60 self.col_data = col_data 

61 self.col_scale = col_scale 

62 

63 @staticmethod 

64 def _validate_pair( 

65 direction: str, 

66 data: Optional[torch.Tensor], 

67 scale: Optional[torch.Tensor], 

68 ) -> None: 

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

70 raise ValueError( 

71 f"W4A8 {direction}-wise data and scale must be present together." 

72 ) 

73 

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

75 """Release directional data that is no longer needed.""" 

76 

77 if not rowwise: 

78 self.row_data = None 

79 self.row_scale = None 

80 if not colwise: 

81 self.col_data = None 

82 self.col_scale = None 

83 

84 def is_rowwise(self) -> bool: 

85 """Return whether row-wise data is available.""" 

86 

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

88 

89 def is_colwise(self) -> bool: 

90 """Return whether column-wise data is available.""" 

91 

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

93 

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

95 """Return physical data and quantizer metadata.""" 

96 

97 return { 

98 "row_data": self.row_data, 

99 "row_scale": self.row_scale, 

100 "col_data": self.col_data, 

101 "col_scale": self.col_scale, 

102 "quantizer": self.quantizer, 

103 } 

104 

105 

106class W4A8Tensor(W4A8TensorStorage, QuantizedTensor): 

107 """Expose native W4A8 storage with the source tensor's logical dtype.""" 

108 

109 def __new__( 

110 cls, 

111 shape: Iterable[int], 

112 dtype: torch.dtype, 

113 *, 

114 quantizer: "W4A8Quantizer", 

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

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

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

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

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

120 requires_grad: bool = False, 

121 ) -> "W4A8Tensor": 

122 """Create a wrapper from at least one directional representation.""" 

123 

124 del quantizer, row_scale, col_scale 

125 physical_data = row_data if row_data is not None else col_data 

126 if physical_data is None: 

127 raise ValueError("W4A8Tensor requires row_data or col_data.") 

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

129 return QuantizedTensor.__new__( 

130 cls, 

131 shape, 

132 dtype, 

133 device=wrapper_device, 

134 requires_grad=requires_grad, 

135 ) 

136 

137 def __init__( 

138 self, 

139 shape: Iterable[int], 

140 dtype: torch.dtype, 

141 *, 

142 quantizer: "W4A8Quantizer", 

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

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

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

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

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

148 requires_grad: bool = False, 

149 ) -> None: 

150 """Initialize the wrapper's directional storage.""" 

151 

152 del device, requires_grad 

153 W4A8TensorStorage.__init__( 

154 self, 

155 shape, 

156 dtype, 

157 quantizer=quantizer, 

158 row_data=row_data, 

159 row_scale=row_scale, 

160 col_data=col_data, 

161 col_scale=col_scale, 

162 ) 

163 

164 def _new_with_transformed_data( 

165 self, 

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

167 *, 

168 requires_grad: bool, 

169 ) -> "W4A8Tensor": 

170 """Rebuild the wrapper after transforming physical tensors.""" 

171 

172 def transform_optional(value: Optional[torch.Tensor]) -> Optional[torch.Tensor]: 

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

174 

175 return type(self)( 

176 shape=self.shape, 

177 dtype=self.dtype, 

178 quantizer=self.quantizer, 

179 row_data=transform_optional(self.row_data), 

180 row_scale=transform_optional(self.row_scale), 

181 col_data=transform_optional(self.col_data), 

182 col_scale=transform_optional(self.col_scale), 

183 requires_grad=requires_grad, 

184 ) 

185 

186 def __repr__(self) -> str: 

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

188 

189 return ( 

190 "W4A8Tensor(" 

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

192 f"rowwise={self.is_rowwise()}, colwise={self.is_colwise()})" 

193 ) 

194 

195 

196__all__ = ["W4A8Tensor", "W4A8TensorStorage"]