Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / components / quantization / config.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"""Typed low-precision policies and their FQN-independent dtype-scheme catalog.""" 

16 

17from collections.abc import Mapping 

18from dataclasses import dataclass, field 

19from typing import Literal, Optional, TypeAlias, cast 

20 

21 

22MXBlockSizeValue: TypeAlias = Literal[32] 

23FP4_BITS = 4 

24FP8_BITS = 8 

25MX_BLOCK_SIZE_32 = 32 

26SUPPORTED_MX_BLOCK_SIZES = (MX_BLOCK_SIZE_32,) 

27 

28# MX tensor and scale layout. 

29MX_SCALE_SLOTS = 2 

30MX_SCALE_TILE = MX_BLOCK_SIZE_32 * MX_SCALE_SLOTS 

31MXFP4_VALUES_PER_BYTE = 2 

32MXFP4_NIBBLE_BITS = 4 

33MXFP4_NIBBLE_MASK = (1 << MXFP4_NIBBLE_BITS) - 1 

34 

35# E8M0/FP8 conversion limits used by the fake-QAT path. 

36E8M0_BIAS = 127.0 

37E8M0_MAX_EXPONENT = 255.0 

38MXFP4_MAX_VALUE = 6.0 

39FP8_E4M3_MAX = 448.0 

40 

41# GroupedMatmul and DynamicMxQuant protocol values. 

42GROUP_LIST_OFFSETS = 0 

43GROUP_LIST_COUNTS = 1 

44GMM_GROUP_M = 0 

45GMM_GROUP_K = 2 

46GMM_SPLIT_ANTIQ = 2 

47GMM_SPLIT_MXFP8 = 2 

48GMM_SPLIT_SCALE = 3 

49MX_SCALE_ALG_OCP = 0 

50MX_SCALE_ALG_CUBALS = 1 

51NPU_FORMAT_FRACTAL_NZ = 29 

52 

53# Shape ranks used by grouped quantization contracts. 

54MATRIX_NDIM = 2 

55GROUPED_NDIM = 3 

56SCALE_NDIM = 4 

57_SUPPORTED_BLOCK_SIZE_TEXT = " or ".join( 

58 str(block_size) for block_size in SUPPORTED_MX_BLOCK_SIZES 

59) 

60 

61 

62def validate_mx_block_size(block_size: int) -> MXBlockSizeValue: 

63 """Validate a configured integer and preserve its numeric representation.""" 

64 

65 if isinstance(block_size, bool) or not isinstance(block_size, int): 

66 raise ValueError( 

67 f"MX block_size must be {_SUPPORTED_BLOCK_SIZE_TEXT}, " 

68 f"got {block_size!r}." 

69 ) 

70 if block_size not in SUPPORTED_MX_BLOCK_SIZES: 

71 raise ValueError( 

72 f"MX block_size must be {_SUPPORTED_BLOCK_SIZE_TEXT}, " 

73 f"got {block_size!r}." 

74 ) 

75 return cast(MXBlockSizeValue, block_size) 

76 

77 

78# One registry row per accepted dtype: (format family, bit width). 

79# weight_format and act_format must share one family (operators never mix 

80# families). The weight/activation bit widths are read directly from these 

81# rows to select the native strategy (e.g. mxfp4 weight + mxfp8 act => W4A8). 

82# Future dtypes (int/fp/bf, e.g. bf16) are one new row here; hif4 is 

83# pre-listed as part of the hifloat family (reserved for its flow). 

84_FORMAT_SPECS = { 

85 "mxfp4": ("mxfp", FP4_BITS), 

86 "mxfp8": ("mxfp", FP8_BITS), 

87 "hif4": ("hifloat", FP4_BITS), 

88 "hif8": ("hifloat", FP8_BITS), 

89} 

90_FORMATS = tuple(_FORMAT_SPECS) 

91 

92 

93def _format_spec(format_name: str) -> tuple[str, int]: 

94 """Look up one dtype format in ``_FORMAT_SPECS``. 

95 

96 Drives ``LowPrecisionDtypeScheme.__post_init__`` validation (family and 

97 bit width). Raises ``ValueError`` on an unregistered 

98 format so a bad ``weight_format``/``act_format`` fails at config time. 

99 """ 

100 

101 try: 

102 return _FORMAT_SPECS[format_name] 

103 except (KeyError, TypeError) as exc: 

104 raise ValueError( 

105 f"Unsupported low-precision format {format_name!r}; " 

106 f"expected {_FORMATS}." 

107 ) from exc 

108 

109 

110def _format_family(format_name: str) -> str: 

111 """Resolve one dtype format to its family name for the factory family gate. 

112 

113 Exported for the strategy factory, which must reject cross-family schemes 

114 without coupling to the raw ``_FORMAT_SPECS`` tuple layout. 

115 """ 

116 

117 return _format_spec(format_name)[0] 

118 

119 

120@dataclass(frozen=True) 

121class LowPrecisionDtypeScheme: 

122 """One complete low-precision dtype scheme referenced by a plan entry. 

123 Defines the quantization dtypes for weights and activations together with 

124 the quantization granularity and whether the scheme runs the native 

125 operator flow or a future fake-quantize (fake QAT) seam. The native 

126 strategy (W4A8/W8A8) is derived downstream from ``weight_format`` / 

127 ``act_format`` bit widths, so it is not a separate field here. 

128 """ 

129 

130 # Optional: True = future fake-QAT seam; False (default) = native flow. 

131 is_fake_quantize: bool = False 

132 weight_format: Literal["mxfp4", "mxfp8", "hif4", "hif8"] = "mxfp8" 

133 act_format: Literal["mxfp4", "mxfp8", "hif4", "hif8"] = "mxfp8" 

134 block_size: int = MX_BLOCK_SIZE_32 

135 

136 def __post_init__(self) -> None: 

137 """Validate that the scheme names a consistent native/QAT policy. 

138 ``weight_format``/``act_format`` must share one format family (the 

139 underlying operator never mixes families); their bit widths drive 

140 strategy selection downstream. ``block_size`` must be an integer; 

141 only MX formats restrict it to the currently supported block32 value. 

142 """ 

143 

144 if not isinstance(self.is_fake_quantize, bool): 

145 raise ValueError( 

146 "LowPrecisionDtypeScheme.is_fake_quantize must be a bool, " 

147 f"but got {self.is_fake_quantize!r}." 

148 ) 

149 if isinstance(self.block_size, bool) or not isinstance(self.block_size, int): 

150 raise ValueError("LowPrecisionDtypeScheme.block_size must be an integer.") 

151 

152 weight_family, _ = _format_spec(self.weight_format) 

153 act_family, _ = _format_spec(self.act_format) 

154 if weight_family != act_family: 

155 raise ValueError( 

156 "weight_format and act_format must belong to the same format " 

157 "family (the underlying operator never mixes families); got " 

158 f"weight={self.weight_format!r}/act={self.act_format!r}." 

159 ) 

160 if weight_family == "mxfp": 

161 try: 

162 block_size = validate_mx_block_size(self.block_size) 

163 except ValueError as exc: 

164 raise ValueError(f"LowPrecisionDtypeScheme.{exc}") from exc 

165 object.__setattr__(self, "block_size", block_size) 

166 

167 

168@dataclass 

169class LowPrecisionConfig: 

170 """Global switch plus a catalog of named dtype schemes. 

171 

172 A dtype scheme is selected by the surrounding ``PlanOverride`` entry. 

173 Matching module FQNs and selecting a dtype scheme are intentionally 

174 separate concerns: this object is only the low-precision catalog. 

175 

176 ``format``/``scaling`` retain the pre-catalog configuration surface. They 

177 are converted to one implicit scheme when no named catalog is supplied, 

178 so existing callers remain valid while mixed W8A8/W4A8 plans use 

179 ``dtype_schemes`` explicitly. 

180 """ 

181 

182 enabled: bool = False 

183 format: Literal["mxfp8_e4m3", "mxfp4_e2m1", "hif8"] = "mxfp8_e4m3" 

184 scaling: Literal["mx_block", "current"] = "mx_block" 

185 dtype_schemes: dict[str, LowPrecisionDtypeScheme] = field(default_factory=dict) 

186 default_dtype_scheme: Optional[str] = None 

187 

188 def __post_init__(self) -> None: 

189 """Validate catalog keys and the optional default dtype scheme.""" 

190 

191 if not isinstance(self.enabled, bool): 

192 raise ValueError("LowPrecisionConfig.enabled must be a bool.") 

193 legacy_pairs = { 

194 ("mxfp8_e4m3", "mx_block"), 

195 ("mxfp4_e2m1", "mx_block"), 

196 ("hif8", "current"), 

197 } 

198 if (self.format, self.scaling) not in legacy_pairs: 

199 raise ValueError( 

200 "Unsupported low-precision format/scaling combination " 

201 f"{self.format!r}/{self.scaling!r}." 

202 ) 

203 

204 if not isinstance(self.dtype_schemes, Mapping): 

205 raise ValueError( 

206 "LowPrecisionConfig.dtype_schemes must be a mapping of " 

207 "non-empty names to dtype schemes." 

208 ) 

209 typed_schemes = {} 

210 for name, scheme in self.dtype_schemes.items(): 

211 if isinstance(scheme, Mapping): 

212 try: 

213 scheme = LowPrecisionDtypeScheme(**dict(scheme)) 

214 except (TypeError, ValueError) as exc: 

215 raise ValueError( 

216 f"Invalid low-precision dtype scheme {name!r}: {exc}" 

217 ) from exc 

218 typed_schemes[name] = scheme 

219 self.dtype_schemes = typed_schemes 

220 

221 if any( 

222 not isinstance(name, str) or not name 

223 or not isinstance(scheme, LowPrecisionDtypeScheme) 

224 for name, scheme in self.dtype_schemes.items() 

225 ): 

226 raise ValueError( 

227 "LowPrecisionConfig.dtype_schemes must map non-empty names to " 

228 "LowPrecisionDtypeScheme objects." 

229 ) 

230 if self.default_dtype_scheme is not None: 

231 if not isinstance(self.default_dtype_scheme, str) or not self.default_dtype_scheme: 

232 raise ValueError( 

233 "LowPrecisionConfig.default_dtype_scheme must be a dtype scheme name." 

234 ) 

235 if self.default_dtype_scheme not in self.dtype_schemes: 

236 raise ValueError( 

237 f"LowPrecisionConfig.default_dtype_scheme " 

238 f"{self.default_dtype_scheme!r} is not present in dtype_schemes." 

239 ) 

240 

241 def resolve_dtype_scheme( 

242 self, dtype_scheme_name: Optional[str] = None, 

243 ) -> LowPrecisionDtypeScheme: 

244 """Map a dtype-scheme *name* to the scheme object used downstream. 

245 

246 The input is either ``plan_overrides[].low_precision_dtype_scheme`` or 

247 ``default_dtype_scheme``; the returned ``LowPrecisionDtypeScheme`` is 

248 what the strategy factory and adapters actually consume. Only resolves 

249 catalog names -- which modules are matched is decided by the plan. 

250 """ 

251 

252 selected = dtype_scheme_name or self.default_dtype_scheme 

253 if selected is not None: 

254 if not isinstance(selected, str) or not selected: 

255 raise ValueError( 

256 "low-precision dtype scheme name must be a non-empty string." 

257 ) 

258 try: 

259 return self.dtype_schemes[selected] 

260 except KeyError as exc: 

261 raise ValueError( 

262 f"Unknown low-precision dtype scheme {selected!r}; " 

263 f"available dtype schemes: {sorted(self.dtype_schemes)}." 

264 ) from exc 

265 if self.dtype_schemes: 

266 # A catalog exists but nothing was selected: fail loudly instead of 

267 # silently running the model unquantized. 

268 raise ValueError( 

269 "LowPrecisionConfig.dtype_schemes is configured, but the replacement " 

270 "plan did not select low_precision_dtype_scheme and no " 

271 "default_dtype_scheme was declared." 

272 ) 

273 legacy_scheme = { 

274 ("mxfp8_e4m3", "mx_block"): LowPrecisionDtypeScheme(), 

275 ("mxfp4_e2m1", "mx_block"): LowPrecisionDtypeScheme( 

276 weight_format="mxfp4", 

277 act_format="mxfp8", 

278 ), 

279 ("hif8", "current"): LowPrecisionDtypeScheme( 

280 weight_format="hif8", 

281 act_format="hif8", 

282 ), 

283 } 

284 return legacy_scheme[(self.format, self.scaling)] 

285 

286 def to_dict(self) -> dict[str, object]: 

287 """Serialize legacy fields and the optional named scheme catalog.""" 

288 

289 result: dict[str, object] = { 

290 "enabled": self.enabled, 

291 "format": self.format, 

292 "scaling": self.scaling, 

293 } 

294 if self.dtype_schemes: 

295 result["dtype_schemes"] = { 

296 name: { 

297 "is_fake_quantize": scheme.is_fake_quantize, 

298 "weight_format": scheme.weight_format, 

299 "act_format": scheme.act_format, 

300 "block_size": scheme.block_size, 

301 } 

302 for name, scheme in self.dtype_schemes.items() 

303 } 

304 if self.default_dtype_scheme is not None: 

305 result["default_dtype_scheme"] = self.default_dtype_scheme 

306 return result 

307 

308 

309__all__ = [ 

310 "MXBlockSizeValue", 

311 "LowPrecisionConfig", 

312 "LowPrecisionDtypeScheme", 

313 "MX_BLOCK_SIZE_32", 

314 "SUPPORTED_MX_BLOCK_SIZES", 

315 "validate_mx_block_size", 

316]