Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / components / quantization / ops / npu_hifloat8.py: 57%
90 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-11-15 07:18 +0800
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"""A5 HiFloat8 current-quantization, MM, and GMM adapters."""
17from functools import lru_cache
18from typing import Any, Optional
20import torch # pylint: disable=forbidden-backend-import
22from hyper_parallel.components.quantization.config import (
23 GMM_GROUP_K,
24 GMM_GROUP_M,
25 GMM_SPLIT_SCALE,
26 GROUP_LIST_COUNTS,
27 GROUP_LIST_OFFSETS,
28)
29from hyper_parallel.components.quantization.tensor.hifloat8_tensor import (
30 HiFloat8Tensor,
31)
33_A5_DEVICE_MARKERS = ("Ascend950", "Ascend910_95")
34_SUPPORTED_LAYOUTS = ("NN", "NT", "TN")
37class HiFloat8NpuOps:
38 """Adapt the verified A5 HiFloat8 torch_npu operator contract."""
40 def __init__(self, torch_npu_module: Any) -> None:
41 """Validate logical dtype, required operators, and target device."""
43 required_operators = ("npu_dynamic_quant", "npu_quant_matmul")
44 missing = [
45 name
46 for name in required_operators
47 if not callable(getattr(torch_npu_module, name, None))
48 ]
49 if getattr(torch_npu_module, "hifloat8", None) is None:
50 missing.append("hifloat8")
51 if missing:
52 version = getattr(torch_npu_module, "__version__", "unknown")
53 raise RuntimeError(
54 "The active torch_npu stack does not provide the A5 HiFloat8 "
55 f"operator contract; missing {missing}, torch_npu={version}."
56 )
58 npu_handle = getattr(torch_npu_module, "npu", None)
59 device_name = (
60 npu_handle.get_device_name()
61 if npu_handle is not None
62 and hasattr(npu_handle, "get_device_name")
63 else ""
64 )
65 if not any(marker in device_name for marker in _A5_DEVICE_MARKERS):
66 raise RuntimeError(
67 "HiFloat8 is supported only on Ascend 950PR/950DT (A5), "
68 f"but the active device is {device_name or 'unknown'}."
69 )
70 self._torch_npu = torch_npu_module
72 @property
73 def quant_dtype(self) -> Any:
74 """Return the runtime's logical HiFloat8 dtype identifier."""
76 return self._torch_npu.hifloat8
78 def dynamic_quant(
79 self,
80 tensor: torch.Tensor,
81 *,
82 dst_type_max: float,
83 quant_mode: str = "pertensor",
84 ) -> tuple[torch.Tensor, torch.Tensor]:
85 """Quantize one FP16/BF16 Tensor with the selected current scale."""
87 return self._torch_npu.npu_dynamic_quant(
88 tensor,
89 dst_type=self._torch_npu.hifloat8,
90 quant_mode=quant_mode,
91 dst_type_max=float(dst_type_max),
92 )
94 def quant_matmul(
95 self,
96 x1: torch.Tensor,
97 x2: torch.Tensor,
98 scale: torch.Tensor,
99 *,
100 pertoken_scale: torch.Tensor,
101 output_dtype: torch.dtype,
102 ) -> torch.Tensor:
103 """Execute an A5 HiFloat8 MM with explicit logical input dtypes."""
105 return self._torch_npu.npu_quant_matmul(
106 x1,
107 x2,
108 scale,
109 pertoken_scale=pertoken_scale,
110 output_dtype=output_dtype,
111 x1_dtype=self._torch_npu.hifloat8,
112 x2_dtype=self._torch_npu.hifloat8,
113 )
115 def validate_grouped_matmul(self) -> None:
116 """Fail when the runtime lacks the HiFloat8 GMM operator."""
118 if not callable(getattr(self._torch_npu, "npu_grouped_matmul", None)):
119 version = getattr(self._torch_npu, "__version__", "unknown")
120 raise RuntimeError(
121 "The active torch_npu stack does not provide "
122 f"npu_grouped_matmul; torch_npu={version}."
123 )
125 def quant_grouped_matmul(
126 self,
127 x1: torch.Tensor,
128 x2: torch.Tensor,
129 x2_scale: torch.Tensor,
130 *,
131 x1_scale: torch.Tensor,
132 group_list: torch.Tensor,
133 group_type: int,
134 group_list_type: int,
135 output_dtype: torch.dtype,
136 bias: Optional[list[torch.Tensor]] = None,
137 ) -> torch.Tensor:
138 """Execute the verified A5 HiFloat8 grouped matmul contract."""
140 self.validate_grouped_matmul()
141 return self._torch_npu.npu_grouped_matmul(
142 [x1],
143 [x2],
144 bias=[] if bias is None else bias,
145 scale=[x2_scale],
146 per_token_scale=[x1_scale],
147 group_list=group_list,
148 group_type=group_type,
149 output_dtype=output_dtype,
150 group_list_type=group_list_type,
151 x_dtype=self._torch_npu.hifloat8,
152 weight_dtype=self._torch_npu.hifloat8,
153 split_item=GMM_SPLIT_SCALE,
154 )[0]
157@lru_cache(maxsize=1)
158def _get_hifloat8_npu_ops() -> HiFloat8NpuOps:
159 """Return the process-local A5 HiFloat8 operator adapter."""
161 try:
162 import torch_npu # pylint: disable=C0415
163 except ImportError as exc:
164 raise RuntimeError(
165 "A5 HiFloat8 requires torch_npu, but it is not importable."
166 ) from exc
167 return HiFloat8NpuOps(torch_npu)
170def validate_hifloat8_runtime() -> None:
171 """Fail during model setup when the HiFloat8 runtime is unavailable."""
173 _get_hifloat8_npu_ops()
176def validate_hifloat8_gmm_runtime() -> None:
177 """Fail during expert setup when the HiFloat8 GMM runtime is unavailable."""
179 _get_hifloat8_npu_ops().validate_grouped_matmul()
182def _left_operand(
183 tensor: HiFloat8Tensor,
184 *,
185 transpose: bool,
186) -> tuple[torch.Tensor, torch.Tensor]:
187 """Select the Pangu-compatible left operand direction."""
189 if transpose:
190 if not tensor.is_colwise():
191 raise ValueError(
192 "Transposed HiFloat8 left operand requires column-wise data."
193 )
194 return tensor.col_data.transpose(-1, -2), tensor.col_scale
195 if not tensor.is_rowwise():
196 raise ValueError(
197 "Non-transposed HiFloat8 left operand requires row-wise data."
198 )
199 return tensor.row_data, tensor.row_scale
202def _right_operand(
203 tensor: HiFloat8Tensor,
204 *,
205 transpose: bool,
206) -> tuple[torch.Tensor, torch.Tensor]:
207 """Select the Pangu-compatible right operand direction."""
209 if transpose:
210 if not tensor.is_rowwise():
211 raise ValueError(
212 "Transposed HiFloat8 right operand requires row-wise data."
213 )
214 return tensor.row_data.transpose(-1, -2), tensor.row_scale
215 if not tensor.is_colwise():
216 raise ValueError(
217 "Non-transposed HiFloat8 right operand requires column-wise data."
218 )
219 return tensor.col_data, tensor.col_scale
222def hifloat8_matmul(
223 left: HiFloat8Tensor,
224 right: HiFloat8Tensor,
225 *,
226 layout: str,
227 output_dtype: Optional[torch.dtype] = None,
228) -> torch.Tensor:
229 """Multiply typed HiFloat8 operands and return a high-precision Tensor."""
231 if not isinstance(left, HiFloat8Tensor) or not isinstance(
232 right,
233 HiFloat8Tensor,
234 ):
235 raise TypeError("hifloat8_matmul requires two HiFloat8Tensor operands.")
236 if layout not in _SUPPORTED_LAYOUTS:
237 raise ValueError(
238 f"Unsupported HiFloat8 layout {layout!r}; "
239 f"expected one of {_SUPPORTED_LAYOUTS}."
240 )
242 result_dtype = output_dtype or left.dtype
243 left_data, left_scale = _left_operand(
244 left,
245 transpose=layout[0] == "T",
246 )
247 right_data, right_scale = _right_operand(
248 right,
249 transpose=layout[1] == "T",
250 )
251 return left.quantizer.npu_ops.quant_matmul(
252 left_data,
253 right_data,
254 right_scale,
255 pertoken_scale=left_scale,
256 output_dtype=result_dtype,
257 )
260def hifloat8_grouped_matmul(
261 left: HiFloat8Tensor,
262 right: HiFloat8Tensor,
263 *,
264 layout: str,
265 group_list: torch.Tensor,
266 group_type: int = GMM_GROUP_M,
267 group_list_type: int = GROUP_LIST_OFFSETS,
268 output_dtype: Optional[torch.dtype] = None,
269 bias: Optional[torch.Tensor] = None,
270) -> torch.Tensor:
271 """Multiply expert-grouped HiFloat8 operands on A5.
273 ``group_type=0`` splits the token/M axis for forward and dgrad, while
274 ``group_type=2`` splits the contracting axis for wgrad.
275 """
277 if not isinstance(left, HiFloat8Tensor) or not isinstance(
278 right,
279 HiFloat8Tensor,
280 ):
281 raise TypeError(
282 "hifloat8_grouped_matmul requires two HiFloat8Tensor operands."
283 )
284 if layout not in _SUPPORTED_LAYOUTS:
285 raise ValueError(
286 f"Unsupported HiFloat8 GMM layout {layout!r}; "
287 f"expected one of {_SUPPORTED_LAYOUTS}."
288 )
289 if not isinstance(group_list, torch.Tensor) or group_list.ndim != 1:
290 raise ValueError("HiFloat8 GMM group_list must be one-dimensional.")
291 if group_list.dtype != torch.int64:
292 raise TypeError("HiFloat8 GMM group_list must use torch.int64.")
293 if group_type not in (GMM_GROUP_M, GMM_GROUP_K):
294 raise ValueError(
295 f"HiFloat8 GMM group_type must be {GMM_GROUP_M} or {GMM_GROUP_K}, "
296 f"got {group_type}."
297 )
298 if group_list_type not in (GROUP_LIST_OFFSETS, GROUP_LIST_COUNTS):
299 raise ValueError(
300 "HiFloat8 GMM group_list_type must be "
301 f"{GROUP_LIST_OFFSETS} or {GROUP_LIST_COUNTS}, "
302 f"got {group_list_type}."
303 )
305 left_data, left_scale = _left_operand(
306 left,
307 transpose=layout[0] == "T",
308 )
309 right_data, right_scale = _right_operand(
310 right,
311 transpose=layout[1] == "T",
312 )
313 if left_scale.dtype != torch.float32 or right_scale.dtype != torch.float32:
314 raise TypeError("HiFloat8 GMM scales must use torch.float32.")
315 return left.quantizer.npu_ops.quant_grouped_matmul(
316 left_data,
317 right_data,
318 right_scale,
319 x1_scale=left_scale,
320 group_list=group_list,
321 group_type=group_type,
322 group_list_type=group_list_type,
323 output_dtype=output_dtype or left.dtype,
324 bias=None if bias is None else [bias],
325 )
328__all__ = [
329 "HiFloat8NpuOps",
330 "hifloat8_grouped_matmul",
331 "hifloat8_matmul",
332 "validate_hifloat8_gmm_runtime",
333 "validate_hifloat8_runtime",
334]