Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / components / quantization / ops / npu_mxfp8.py: 52%
102 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 MXFP8 quantization and matrix-multiplication operator adapters."""
17from functools import lru_cache
18from typing import Any, Optional
20import torch # pylint: disable=forbidden-backend-import
22from hyper_parallel.components.quantization.tensor import MXFP8Tensor
24_A5_DEVICE_MARKERS = ("Ascend950", "Ascend910_95")
25_SUPPORTED_LAYOUTS = ("NN", "NT", "TN")
28class LowPrecisionCapabilityError(RuntimeError):
29 """Report that the active torch_npu stack lacks the A5 MXFP8 contract."""
32class MXFP8NpuOps:
33 """Adapt the A5 MXFP8 torch_npu operator contract."""
35 def __init__(self, torch_npu_module: Any) -> None:
36 """Validate and retain an imported torch_npu module."""
38 required = (
39 "npu_dynamic_mx_quant",
40 "npu_dynamic_mx_quant_with_dual_axis",
41 "npu_quant_matmul",
42 "float8_e8m0fnu",
43 )
44 missing = [
45 name for name in required if not hasattr(torch_npu_module, name)
46 ]
47 if missing:
48 version = getattr(torch_npu_module, "__version__", "unknown")
49 raise LowPrecisionCapabilityError(
50 "The active torch_npu stack does not provide the A5 MXFP8 "
51 f"operator contract; missing {missing}, torch_npu={version}."
52 )
53 npu_handle = getattr(torch_npu_module, "npu", None)
54 device_name = (
55 npu_handle.get_device_name()
56 if npu_handle is not None
57 and hasattr(npu_handle, "get_device_name")
58 else ""
59 )
60 if not any(marker in device_name for marker in _A5_DEVICE_MARKERS):
61 raise LowPrecisionCapabilityError(
62 "MXFP8 is supported only on Ascend 950PR/950DT (A5), "
63 f"but the active device is {device_name or 'unknown'}."
64 )
65 self._torch_npu = torch_npu_module
67 def validate_grouped_matmul(self) -> None:
68 """Fail when the runtime lacks either MXFP8 GMM-specific operator."""
70 required = (
71 "npu_grouped_dynamic_mx_quant",
72 "npu_grouped_matmul",
73 )
74 missing = [
75 name for name in required if not hasattr(self._torch_npu, name)
76 ]
77 if missing:
78 version = getattr(self._torch_npu, "__version__", "unknown")
79 raise LowPrecisionCapabilityError(
80 "The active torch_npu stack does not provide the A5 MXFP8 "
81 f"GMM contract; missing {missing}, torch_npu={version}."
82 )
84 def dynamic_mx_quant(
85 self,
86 tensor: torch.Tensor,
87 *,
88 axis: int,
89 quant_dtype: Optional[torch.dtype] = None,
90 ) -> tuple[torch.Tensor, torch.Tensor]:
91 """Quantize one tensor along an MX block axis."""
93 return self._torch_npu.npu_dynamic_mx_quant(
94 tensor,
95 axis=axis,
96 dst_type=quant_dtype or torch.float8_e4m3fn,
97 scale_alg=1,
98 )
100 def dynamic_mx_quant_dual_axis(
101 self,
102 tensor: torch.Tensor,
103 *,
104 quant_dtype: Optional[torch.dtype] = None,
105 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
106 """Quantize one tensor into row-wise and column-wise MX storage."""
108 return self._torch_npu.npu_dynamic_mx_quant_with_dual_axis(
109 tensor,
110 dst_type=quant_dtype or torch.float8_e4m3fn,
111 scale_alg=1,
112 )
114 def quant_matmul(
115 self,
116 x1: torch.Tensor,
117 x2: torch.Tensor,
118 scale: torch.Tensor,
119 *,
120 pertoken_scale: torch.Tensor,
121 output_dtype: torch.dtype,
122 ) -> torch.Tensor:
123 """Execute an A5 MXFP8 matrix multiplication."""
125 return self._torch_npu.npu_quant_matmul(
126 x1,
127 x2,
128 scale,
129 pertoken_scale=pertoken_scale,
130 output_dtype=output_dtype,
131 scale_dtype=self._torch_npu.float8_e8m0fnu,
132 pertoken_scale_dtype=self._torch_npu.float8_e8m0fnu,
133 group_sizes=[1, 1, 32],
134 )
136 def grouped_dynamic_mx_quant(
137 self,
138 tensor: torch.Tensor,
139 group_list: torch.Tensor,
140 *,
141 quant_dtype: Optional[torch.dtype] = None,
142 ) -> tuple[torch.Tensor, torch.Tensor]:
143 """Quantize the column direction independently for each expert."""
145 grouped_quant = getattr(
146 self._torch_npu,
147 "npu_grouped_dynamic_mx_quant",
148 None,
149 )
150 if grouped_quant is None:
151 version = getattr(self._torch_npu, "__version__", "unknown")
152 raise LowPrecisionCapabilityError(
153 "The active torch_npu stack does not provide "
154 "npu_grouped_dynamic_mx_quant, which MXFP8 GMM requires; "
155 f"torch_npu={version}."
156 )
157 return grouped_quant(
158 tensor,
159 # A5 GroupedDynamicMxQuant's groupIndex contract is int32.
160 group_list.to(torch.int32),
161 round_mode="rint",
162 dst_type=quant_dtype or torch.float8_e4m3fn,
163 blocksize=32,
164 )
166 def quant_grouped_matmul(
167 self,
168 x1: torch.Tensor,
169 x2: torch.Tensor,
170 x2_scale: torch.Tensor,
171 *,
172 x1_scale: torch.Tensor,
173 group_list: torch.Tensor,
174 group_type: int,
175 group_list_type: int,
176 output_dtype: torch.dtype,
177 bias: Optional[list[torch.Tensor]] = None,
178 ) -> torch.Tensor:
179 """Execute one A5 MXFP8 grouped matrix multiplication."""
181 grouped_matmul = getattr(self._torch_npu, "npu_grouped_matmul", None)
182 if grouped_matmul is None:
183 version = getattr(self._torch_npu, "__version__", "unknown")
184 raise LowPrecisionCapabilityError(
185 "The active torch_npu stack does not provide "
186 f"npu_grouped_matmul; torch_npu={version}."
187 )
188 # A5 GroupedMatmul's groupList contract is int64.
189 group_list = group_list.to(torch.int64)
190 return grouped_matmul(
191 [x1],
192 [x2],
193 bias=[] if bias is None else bias,
194 scale=[x2_scale],
195 per_token_scale=[x1_scale],
196 group_list=group_list,
197 group_type=group_type,
198 output_dtype=output_dtype,
199 group_list_type=group_list_type,
200 scale_dtype=self._torch_npu.float8_e8m0fnu,
201 per_token_scale_dtype=self._torch_npu.float8_e8m0fnu,
202 split_item=2,
203 )[0]
205 def is_e8m0_dtype(self, dtype: torch.dtype) -> bool:
206 """Return whether a scale dtype is this runtime's E8M0 representation."""
207 return dtype == self._torch_npu.float8_e8m0fnu
210@lru_cache(maxsize=1)
211def _get_npu_ops() -> MXFP8NpuOps:
212 """Return the process-local A5 MXFP8 operator adapter."""
214 try:
215 import torch_npu # pylint: disable=C0415
216 except ImportError as exc:
217 raise LowPrecisionCapabilityError(
218 "A5 MXFP8 requires torch_npu, but it is not importable."
219 ) from exc
220 return MXFP8NpuOps(torch_npu)
223def validate_npu_runtime() -> None:
224 """Fail during model setup when the MXFP8 runtime is unavailable."""
226 _get_npu_ops()
229def validate_npu_gmm_runtime() -> None:
230 """Fail during expert setup when the MXFP8 GMM runtime is unavailable."""
232 _get_npu_ops().validate_grouped_matmul()
235def _transpose_scale(scale: torch.Tensor) -> torch.Tensor:
236 """Transpose the two matrix dimensions of an MX scale tensor."""
238 if scale.ndim < 2:
239 raise ValueError(
240 "MXFP8 scale transpose requires at least two dimensions, got "
241 f"shape {tuple(scale.shape)}."
242 )
243 if scale.ndim == 2:
244 return scale.transpose(0, 1)
245 return scale.transpose(-3, -2)
248def _left_operand(
249 tensor: MXFP8Tensor,
250 *,
251 transpose: bool,
252) -> tuple[torch.Tensor, torch.Tensor]:
253 if transpose:
254 if not tensor.is_colwise():
255 raise ValueError(
256 "Transposed MXFP8 left operand requires column-wise data."
257 )
258 return (
259 tensor.col_data.transpose(-1, -2),
260 _transpose_scale(tensor.col_scale),
261 )
262 if not tensor.is_rowwise():
263 raise ValueError(
264 "Non-transposed MXFP8 left operand requires row-wise data."
265 )
266 return tensor.row_data, tensor.row_scale
269def _right_operand(
270 tensor: MXFP8Tensor,
271 *,
272 transpose: bool,
273) -> tuple[torch.Tensor, torch.Tensor]:
274 if transpose:
275 if not tensor.is_rowwise():
276 raise ValueError(
277 "Transposed MXFP8 right operand requires row-wise data."
278 )
279 return (
280 tensor.row_data.transpose(-1, -2),
281 _transpose_scale(tensor.row_scale),
282 )
283 if not tensor.is_colwise():
284 raise ValueError(
285 "Non-transposed MXFP8 right operand requires column-wise data."
286 )
287 return tensor.col_data, tensor.col_scale
290def mxfp8_matmul(
291 left: MXFP8Tensor,
292 right: MXFP8Tensor,
293 *,
294 layout: str,
295 output_dtype: Optional[torch.dtype] = None,
296) -> torch.Tensor:
297 """Multiply typed MXFP8 operands and return a high-precision Tensor."""
299 if not isinstance(left, MXFP8Tensor) or not isinstance(
300 right,
301 MXFP8Tensor,
302 ):
303 raise TypeError("mxfp8_matmul requires two MXFP8Tensor operands.")
304 if layout not in _SUPPORTED_LAYOUTS:
305 raise ValueError(
306 f"Unsupported MXFP8 layout {layout!r}; "
307 f"expected one of {_SUPPORTED_LAYOUTS}."
308 )
310 left_data, left_scale = _left_operand(
311 left,
312 transpose=layout[0] == "T",
313 )
314 right_data, right_scale = _right_operand(
315 right,
316 transpose=layout[1] == "T",
317 )
318 return left.quantizer.npu_ops.quant_matmul(
319 left_data,
320 right_data,
321 right_scale,
322 pertoken_scale=left_scale,
323 output_dtype=output_dtype or left.dtype,
324 )
327def mxfp8_grouped_matmul(
328 left: MXFP8Tensor,
329 right: MXFP8Tensor,
330 *,
331 layout: str,
332 group_list: torch.Tensor,
333 group_type: int = 0,
334 group_list_type: int = 0,
335 output_dtype: Optional[torch.dtype] = None,
336 bias: Optional[torch.Tensor] = None,
337) -> torch.Tensor:
338 """Multiply expert-grouped MXFP8 operands and return a high-precision Tensor.
340 The layout follows ``Y = op(left) @ op(right)``. Forward and input-gradient
341 GMM split the M axis with ``group_type=0``; weight-gradient GMM splits the
342 contracting axis with ``group_type=2``.
344 Args:
345 left: Quantized left operand.
346 right: Quantized right operand.
347 layout: One of ``"NN"``, ``"NT"``, or ``"TN"``.
348 group_list: Expert boundaries or per-expert token counts.
349 group_type: Axis grouped by the NPU kernel.
350 group_list_type: ``0`` for cumulative boundaries or ``1`` for counts.
351 output_dtype: Optional high-precision output dtype.
352 bias: Optional grouped bias tensor.
354 Returns:
355 The grouped matrix-multiplication result.
356 """
358 if not isinstance(left, MXFP8Tensor) or not isinstance(right, MXFP8Tensor):
359 raise TypeError("mxfp8_grouped_matmul requires two MXFP8Tensor operands.")
360 if layout not in _SUPPORTED_LAYOUTS:
361 raise ValueError(
362 f"Unsupported MXFP8 GMM layout {layout!r}; "
363 f"expected one of {_SUPPORTED_LAYOUTS}."
364 )
365 if not isinstance(group_list, torch.Tensor) or group_list.ndim != 1:
366 raise ValueError("MXFP8 GMM group_list must be a one-dimensional Tensor.")
367 if group_type not in (0, 2):
368 raise ValueError(
369 f"MXFP8 GMM group_type must be 0 or 2, but got {group_type}."
370 )
371 if group_list_type not in (0, 1):
372 raise ValueError(
373 "MXFP8 GMM group_list_type must be 0 or 1, "
374 f"but got {group_list_type}."
375 )
377 left_data, left_scale = _left_operand(
378 left,
379 transpose=layout[0] == "T",
380 )
381 right_data, right_scale = _right_operand(
382 right,
383 transpose=layout[1] == "T",
384 )
385 bias_list = None if bias is None else [bias]
386 return left.quantizer.npu_ops.quant_grouped_matmul(
387 left_data,
388 right_data,
389 right_scale,
390 x1_scale=left_scale,
391 group_list=group_list,
392 group_type=group_type,
393 group_list_type=group_list_type,
394 output_dtype=output_dtype or left.dtype,
395 bias=bias_list,
396 )