Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / components / quantization / ops / npu_mxfp8.py: 51%
114 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 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.config import (
23 GMM_GROUP_K,
24 GMM_GROUP_M,
25 GMM_SPLIT_MXFP8,
26 GROUP_LIST_COUNTS,
27 GROUP_LIST_OFFSETS,
28 MATRIX_NDIM,
29 MX_BLOCK_SIZE_32,
30 MX_SCALE_ALG_CUBALS,
31)
32from hyper_parallel.components.quantization.tensor import MXFP8Tensor
34_A5_DEVICE_MARKERS = ("Ascend950", "Ascend910_95")
35_SUPPORTED_LAYOUTS = ("NN", "NT", "TN")
38class LowPrecisionCapabilityError(RuntimeError):
39 """Report that the active torch_npu stack lacks the A5 MXFP8 contract."""
42class MXFP8NpuOps:
43 """Adapt the A5 MXFP8 torch_npu operator contract."""
45 def __init__(self, torch_npu_module: Any) -> None:
46 """Validate and retain an imported torch_npu module."""
48 required = (
49 "npu_dynamic_mx_quant",
50 "npu_dynamic_mx_quant_with_dual_axis",
51 "npu_quant_matmul",
52 "float8_e8m0fnu",
53 )
54 missing = [
55 name for name in required if not hasattr(torch_npu_module, name)
56 ]
57 if missing:
58 version = getattr(torch_npu_module, "__version__", "unknown")
59 raise LowPrecisionCapabilityError(
60 "The active torch_npu stack does not provide the A5 MXFP8 "
61 f"operator contract; missing {missing}, torch_npu={version}."
62 )
63 npu_handle = getattr(torch_npu_module, "npu", None)
64 device_name = (
65 npu_handle.get_device_name()
66 if npu_handle is not None
67 and hasattr(npu_handle, "get_device_name")
68 else ""
69 )
70 if not any(marker in device_name for marker in _A5_DEVICE_MARKERS):
71 raise LowPrecisionCapabilityError(
72 "MXFP8 is supported only on Ascend 950PR/950DT (A5), "
73 f"but the active device is {device_name or 'unknown'}."
74 )
75 self._torch_npu = torch_npu_module
77 def validate_grouped_matmul(self) -> None:
78 """Fail when the runtime lacks either MXFP8 GMM-specific operator."""
80 required = (
81 "npu_grouped_dynamic_mx_quant",
82 "npu_grouped_matmul",
83 )
84 missing = [
85 name for name in required if not hasattr(self._torch_npu, name)
86 ]
87 if missing:
88 version = getattr(self._torch_npu, "__version__", "unknown")
89 raise LowPrecisionCapabilityError(
90 "The active torch_npu stack does not provide the A5 MXFP8 "
91 f"GMM contract; missing {missing}, torch_npu={version}."
92 )
94 def dynamic_mx_quant(
95 self,
96 tensor: torch.Tensor,
97 *,
98 axis: int,
99 quant_dtype: Optional[torch.dtype] = None,
100 ) -> tuple[torch.Tensor, torch.Tensor]:
101 """Quantize one tensor along an MX block axis."""
103 return self._torch_npu.npu_dynamic_mx_quant(
104 tensor,
105 axis=axis,
106 dst_type=quant_dtype or torch.float8_e4m3fn,
107 scale_alg=MX_SCALE_ALG_CUBALS,
108 )
110 def dynamic_mx_quant_dual_axis(
111 self,
112 tensor: torch.Tensor,
113 *,
114 quant_dtype: Optional[torch.dtype] = None,
115 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
116 """Quantize one tensor into row-wise and column-wise MX storage."""
118 return self._torch_npu.npu_dynamic_mx_quant_with_dual_axis(
119 tensor,
120 dst_type=quant_dtype or torch.float8_e4m3fn,
121 scale_alg=MX_SCALE_ALG_CUBALS,
122 )
124 def validate_swiglu_fusion(self) -> None:
125 """Require both fused kernels; never silently fall back to plain SwiGLU."""
126 names = ("npu_swiglu_mx_quant_with_dual_axis", "_npu_swiglu_backward_mx_quant_with_dual_axis")
127 missing = [name for name in names if not hasattr(self._torch_npu, name)]
128 if missing:
129 raise LowPrecisionCapabilityError(f"MXFP8 SwiGLU fusion requires {missing}")
131 def swiglu_mx_quant_dual_axis(
132 self,
133 tensor: torch.Tensor,
134 *,
135 quant_dtype: torch.dtype = torch.float8_e4m3fn,
136 group_index: Optional[torch.Tensor] = None,
137 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
138 """Apply SiLU(gate) * up and quantize both MX directions.
140 Args:
141 tensor: BF16/FP16 matrix containing [gate | up] along its last axis.
142 quant_dtype: Payload dtype, E4M3 for the current MXFP8 recipe.
143 group_index: Cumulative expert token boundaries, or None for Dense.
145 Returns:
146 Row data, row scale, column data, and column scale.
147 """
148 return self._torch_npu.npu_swiglu_mx_quant_with_dual_axis(
149 tensor, group_index=group_index, activate_left=True,
150 dst_type=quant_dtype, scale_alg=1,
151 )
153 def swiglu_backward_mx_quant_dual_axis(
154 self,
155 tensor: torch.Tensor,
156 grad_output: torch.Tensor,
157 *,
158 quant_dtype: torch.dtype = torch.float8_e4m3fn,
159 group_index: Optional[torch.Tensor] = None,
160 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
161 """Compute the SwiGLU derivative and quantize both MX directions.
163 Args:
164 tensor: Saved BF16/FP16 packed Gate/Up matrix.
165 grad_output: BF16/FP16 gradient of the SwiGLU output.
166 quant_dtype: Payload dtype, E4M3 for the current MXFP8 recipe.
167 group_index: The same cumulative boundaries used by forward.
169 Returns:
170 Row data, row scale, column data, and column scale of the gradient.
171 """
172 return self._torch_npu._npu_swiglu_backward_mx_quant_with_dual_axis( # pylint: disable=protected-access
173 tensor, grad_output, group_index=group_index, activate_left=True,
174 dst_type=quant_dtype, scale_alg=1,
175 )
177 def quant_matmul(
178 self,
179 x1: torch.Tensor,
180 x2: torch.Tensor,
181 scale: torch.Tensor,
182 *,
183 pertoken_scale: torch.Tensor,
184 output_dtype: torch.dtype,
185 ) -> torch.Tensor:
186 """Execute an A5 MXFP8 matrix multiplication."""
188 return self._torch_npu.npu_quant_matmul(
189 x1,
190 x2,
191 scale,
192 pertoken_scale=pertoken_scale,
193 output_dtype=output_dtype,
194 scale_dtype=self._torch_npu.float8_e8m0fnu,
195 pertoken_scale_dtype=self._torch_npu.float8_e8m0fnu,
196 group_sizes=[1, 1, MX_BLOCK_SIZE_32],
197 )
199 def grouped_dynamic_mx_quant(
200 self,
201 tensor: torch.Tensor,
202 group_list: torch.Tensor,
203 *,
204 quant_dtype: Optional[torch.dtype] = None,
205 ) -> tuple[torch.Tensor, torch.Tensor]:
206 """Quantize the column direction independently for each expert."""
208 grouped_quant = getattr(
209 self._torch_npu,
210 "npu_grouped_dynamic_mx_quant",
211 None,
212 )
213 if grouped_quant is None:
214 version = getattr(self._torch_npu, "__version__", "unknown")
215 raise LowPrecisionCapabilityError(
216 "The active torch_npu stack does not provide "
217 "npu_grouped_dynamic_mx_quant, which MXFP8 GMM requires; "
218 f"torch_npu={version}."
219 )
220 return grouped_quant(
221 tensor,
222 # A5 GroupedDynamicMxQuant's groupIndex contract is int32.
223 group_list.to(torch.int32),
224 round_mode="rint",
225 dst_type=quant_dtype or torch.float8_e4m3fn,
226 blocksize=MX_BLOCK_SIZE_32,
227 )
229 def quant_grouped_matmul(
230 self,
231 x1: torch.Tensor,
232 x2: torch.Tensor,
233 x2_scale: torch.Tensor,
234 *,
235 x1_scale: torch.Tensor,
236 group_list: torch.Tensor,
237 group_type: int,
238 group_list_type: int,
239 output_dtype: torch.dtype,
240 bias: Optional[list[torch.Tensor]] = None,
241 ) -> torch.Tensor:
242 """Execute one A5 MXFP8 grouped matrix multiplication."""
244 grouped_matmul = getattr(self._torch_npu, "npu_grouped_matmul", None)
245 if grouped_matmul is None:
246 version = getattr(self._torch_npu, "__version__", "unknown")
247 raise LowPrecisionCapabilityError(
248 "The active torch_npu stack does not provide "
249 f"npu_grouped_matmul; torch_npu={version}."
250 )
251 # A5 GroupedMatmul's groupList contract is int64.
252 group_list = group_list.to(torch.int64)
253 return grouped_matmul(
254 [x1],
255 [x2],
256 bias=[] if bias is None else bias,
257 scale=[x2_scale],
258 per_token_scale=[x1_scale],
259 group_list=group_list,
260 group_type=group_type,
261 output_dtype=output_dtype,
262 group_list_type=group_list_type,
263 scale_dtype=self._torch_npu.float8_e8m0fnu,
264 per_token_scale_dtype=self._torch_npu.float8_e8m0fnu,
265 split_item=GMM_SPLIT_MXFP8,
266 )[0]
268 def is_e8m0_dtype(self, dtype: torch.dtype) -> bool:
269 """Return whether a scale dtype is this runtime's E8M0 representation."""
270 return dtype == self._torch_npu.float8_e8m0fnu
273@lru_cache(maxsize=1)
274def _get_npu_ops() -> MXFP8NpuOps:
275 """Return the process-local A5 MXFP8 operator adapter."""
277 try:
278 import torch_npu # pylint: disable=C0415
279 except ImportError as exc:
280 raise LowPrecisionCapabilityError(
281 "A5 MXFP8 requires torch_npu, but it is not importable."
282 ) from exc
283 return MXFP8NpuOps(torch_npu)
286def validate_npu_runtime() -> None:
287 """Fail during model setup when the MXFP8 runtime is unavailable."""
289 _get_npu_ops()
292def validate_npu_swiglu_runtime() -> None:
293 """Validate fusion without storing runtime modules on a model's quantizer."""
294 _get_npu_ops().validate_swiglu_fusion()
297def validate_npu_gmm_runtime() -> None:
298 """Fail during expert setup when the MXFP8 GMM runtime is unavailable."""
300 _get_npu_ops().validate_grouped_matmul()
303def _transpose_scale(scale: torch.Tensor) -> torch.Tensor:
304 """Transpose the two matrix dimensions of an MX scale tensor."""
306 if scale.ndim < MATRIX_NDIM:
307 raise ValueError(
308 "MXFP8 scale transpose requires at least two dimensions, got "
309 f"shape {tuple(scale.shape)}."
310 )
311 if scale.ndim == MATRIX_NDIM:
312 return scale.transpose(0, 1)
313 return scale.transpose(-3, -2)
316def _left_operand(
317 tensor: MXFP8Tensor,
318 *,
319 transpose: bool,
320) -> tuple[torch.Tensor, torch.Tensor]:
321 """Select MX data and scales for the requested left-operand layout."""
322 if transpose:
323 if not tensor.is_colwise():
324 raise ValueError(
325 "Transposed MXFP8 left operand requires column-wise data."
326 )
327 return (
328 tensor.col_data.transpose(-1, -2),
329 _transpose_scale(tensor.col_scale),
330 )
331 if not tensor.is_rowwise():
332 raise ValueError(
333 "Non-transposed MXFP8 left operand requires row-wise data."
334 )
335 return tensor.row_data, tensor.row_scale
338def _right_operand(
339 tensor: MXFP8Tensor,
340 *,
341 transpose: bool,
342) -> tuple[torch.Tensor, torch.Tensor]:
343 """Select MX data and scales for the requested right-operand layout."""
344 if transpose:
345 if not tensor.is_rowwise():
346 raise ValueError(
347 "Transposed MXFP8 right operand requires row-wise data."
348 )
349 return (
350 tensor.row_data.transpose(-1, -2),
351 _transpose_scale(tensor.row_scale),
352 )
353 if not tensor.is_colwise():
354 raise ValueError(
355 "Non-transposed MXFP8 right operand requires column-wise data."
356 )
357 return tensor.col_data, tensor.col_scale
360def mxfp8_matmul(
361 left: MXFP8Tensor,
362 right: MXFP8Tensor,
363 *,
364 layout: str,
365 output_dtype: Optional[torch.dtype] = None,
366) -> torch.Tensor:
367 """Multiply typed MXFP8 operands and return a high-precision Tensor."""
369 if not isinstance(left, MXFP8Tensor) or not isinstance(
370 right,
371 MXFP8Tensor,
372 ):
373 raise TypeError("mxfp8_matmul requires two MXFP8Tensor operands.")
374 if layout not in _SUPPORTED_LAYOUTS:
375 raise ValueError(
376 f"Unsupported MXFP8 layout {layout!r}; "
377 f"expected one of {_SUPPORTED_LAYOUTS}."
378 )
380 left_data, left_scale = _left_operand(
381 left,
382 transpose=layout[0] == "T",
383 )
384 right_data, right_scale = _right_operand(
385 right,
386 transpose=layout[1] == "T",
387 )
388 return left.quantizer.npu_ops.quant_matmul(
389 left_data,
390 right_data,
391 right_scale,
392 pertoken_scale=left_scale,
393 output_dtype=output_dtype or left.dtype,
394 )
397def mxfp8_grouped_matmul(
398 left: MXFP8Tensor,
399 right: MXFP8Tensor,
400 *,
401 layout: str,
402 group_list: torch.Tensor,
403 group_type: int = GMM_GROUP_M,
404 group_list_type: int = GROUP_LIST_OFFSETS,
405 output_dtype: Optional[torch.dtype] = None,
406 bias: Optional[torch.Tensor] = None,
407) -> torch.Tensor:
408 """Multiply expert-grouped MXFP8 operands and return a high-precision Tensor.
410 The layout follows ``Y = op(left) @ op(right)``. Forward and input-gradient
411 GMM split the M axis with ``group_type=0``; weight-gradient GMM splits the
412 contracting axis with ``group_type=2``.
414 Args:
415 left: Quantized left operand.
416 right: Quantized right operand.
417 layout: One of ``"NN"``, ``"NT"``, or ``"TN"``.
418 group_list: Expert boundaries or per-expert token counts.
419 group_type: Axis grouped by the NPU kernel.
420 group_list_type: ``0`` for cumulative boundaries or ``1`` for counts.
421 output_dtype: Optional high-precision output dtype.
422 bias: Optional grouped bias tensor.
424 Returns:
425 The grouped matrix-multiplication result.
426 """
428 if not isinstance(left, MXFP8Tensor) or not isinstance(right, MXFP8Tensor):
429 raise TypeError("mxfp8_grouped_matmul requires two MXFP8Tensor operands.")
430 if layout not in _SUPPORTED_LAYOUTS:
431 raise ValueError(
432 f"Unsupported MXFP8 GMM layout {layout!r}; "
433 f"expected one of {_SUPPORTED_LAYOUTS}."
434 )
435 if not isinstance(group_list, torch.Tensor) or group_list.ndim != 1:
436 raise ValueError("MXFP8 GMM group_list must be a one-dimensional Tensor.")
437 if group_type not in (GMM_GROUP_M, GMM_GROUP_K):
438 raise ValueError(
439 f"MXFP8 GMM group_type must be {GMM_GROUP_M} or {GMM_GROUP_K}, "
440 f"but got {group_type}."
441 )
442 if group_list_type not in (GROUP_LIST_OFFSETS, GROUP_LIST_COUNTS):
443 raise ValueError(
444 "MXFP8 GMM group_list_type must be "
445 f"{GROUP_LIST_OFFSETS} or {GROUP_LIST_COUNTS}, "
446 f"but got {group_list_type}."
447 )
449 left_data, left_scale = _left_operand(
450 left,
451 transpose=layout[0] == "T",
452 )
453 right_data, right_scale = _right_operand(
454 right,
455 transpose=layout[1] == "T",
456 )
457 bias_list = None if bias is None else [bias]
458 return left.quantizer.npu_ops.quant_grouped_matmul(
459 left_data,
460 right_data,
461 right_scale,
462 x1_scale=left_scale,
463 group_list=group_list,
464 group_type=group_type,
465 group_list_type=group_list_type,
466 output_dtype=output_dtype or left.dtype,
467 bias=bias_list,
468 )