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