Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / components / quantization / ops / npu_w4a8.py: 29%
78 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"""Ascend native W4A8 operator adapter."""
17from functools import lru_cache
18from types import ModuleType
19from typing import Optional
21import torch # pylint: disable=forbidden-backend-import
23from hyper_parallel.components.quantization.config import (
24 GMM_SPLIT_ANTIQ,
25 GROUP_LIST_OFFSETS,
26 GROUPED_NDIM,
27 MX_BLOCK_SIZE_32,
28 MX_SCALE_ALG_OCP,
29 MX_SCALE_SLOTS,
30 MX_SCALE_TILE,
31 NPU_FORMAT_FRACTAL_NZ,
32 SCALE_NDIM,
33)
35_A5_DEVICE_MARKERS = ("Ascend950", "Ascend910_95")
38class W4A8CapabilityError(RuntimeError):
39 """Report that the active torch_npu stack lacks native W4A8 support."""
42class W4A8NpuOps:
43 """Adapt the native W4A8 torch_npu operator contract."""
45 def __init__(self, torch_npu_module: ModuleType) -> None:
46 """Validate and retain an imported torch_npu module."""
48 required = (
49 "npu_dynamic_mx_quant",
50 "npu_grouped_matmul",
51 "npu_format_cast",
52 "float4_e2m1fn_x2",
53 "float8_e4m3fn",
54 "float8_e8m0fnu",
55 )
56 missing = [name for name in required if not hasattr(torch_npu_module, name)]
57 if missing:
58 version = getattr(torch_npu_module, "__version__", "unknown")
59 raise W4A8CapabilityError(
60 "The active torch_npu stack does not provide the native W4A8 "
61 f"operator contract; missing {missing}, torch_npu={version}."
62 )
63 npu_handle = getattr(torch_npu_module, "npu", None)
64 try:
65 device_name = (
66 npu_handle.get_device_name()
67 if npu_handle is not None and hasattr(npu_handle, "get_device_name")
68 else ""
69 )
70 except (RuntimeError, TypeError, OSError):
71 device_name = ""
72 if not any(marker in device_name for marker in _A5_DEVICE_MARKERS):
73 raise W4A8CapabilityError(
74 "Native W4A8 is supported only on Ascend 950PR/950DT (A5), "
75 f"but the active device is {device_name or 'unknown'}."
76 )
77 self._torch_npu = torch_npu_module
79 @property
80 def e8m0_dtype(self) -> torch.dtype:
81 """Return the runtime E8M0 scale dtype."""
83 return self._torch_npu.float8_e8m0fnu
85 @property
86 def activation_dtype(self) -> torch.dtype:
87 """Return the runtime MXFP8 data dtype."""
89 return self._torch_npu.float8_e4m3fn
91 @property
92 def weight_dtype(self) -> torch.dtype:
93 """Return the runtime packed MXFP4 data dtype."""
95 return self._torch_npu.float4_e2m1fn_x2
97 def dynamic_mx_quant(
98 self,
99 tensor: torch.Tensor,
100 *,
101 axis: int,
102 block_size: int,
103 quant_dtype: torch.dtype,
104 round_mode: Optional[str] = None,
105 scale_alg: Optional[int] = None,
106 ) -> tuple[torch.Tensor, torch.Tensor]:
107 """Run dynamic MX quantization for one operand direction."""
109 kwargs = {"axis": axis, "dst_type": quant_dtype, "block_size": block_size}
110 if round_mode is not None:
111 kwargs["round_mode"] = round_mode
112 if scale_alg is not None:
113 kwargs["scale_alg"] = scale_alg
114 return self._torch_npu.npu_dynamic_mx_quant(tensor, **kwargs)
116 def grouped_dynamic_mx_quant(
117 self,
118 tensor: torch.Tensor,
119 group_list: torch.Tensor,
120 *,
121 block_size: int,
122 quant_dtype: torch.dtype,
123 ) -> tuple[torch.Tensor, torch.Tensor]:
124 """Run grouped dynamic MX quantization over expert token ranges."""
126 grouped_quant = getattr(
127 self._torch_npu,
128 "npu_grouped_dynamic_mx_quant",
129 None,
130 )
131 if grouped_quant is None:
132 raise W4A8CapabilityError(
133 "The active torch_npu stack does not provide "
134 "npu_grouped_dynamic_mx_quant, which native W4A8 block-32 "
135 "backward requires."
136 )
137 return grouped_quant(
138 tensor,
139 # A5 GroupedDynamicMxQuant's groupIndex contract is int32.
140 group_list.to(torch.int32),
141 round_mode="rint",
142 dst_type=quant_dtype,
143 blocksize=block_size,
144 )
146 def grouped_matmul(
147 self,
148 left: torch.Tensor,
149 right: torch.Tensor,
150 right_scale: torch.Tensor,
151 *,
152 left_scale: torch.Tensor,
153 group_list: torch.Tensor,
154 group_type: int,
155 output_dtype: torch.dtype,
156 use_e8m0_scale: bool,
157 split_item: int,
158 native_w4: bool = False,
159 ) -> torch.Tensor:
160 """Run one native grouped matrix multiplication.
162 ``native_w4`` selects the A5 ``MXFP8 x Packed MXFP4`` contract. The
163 W4A8 weight scale belongs to ``antiquant_scale``; the ordinary
164 ``scale`` input is reserved for the MXFP8 path used by wgrad.
165 """
167 # A5 GroupedMatmul's groupList contract is int64.
168 group_list = group_list.to(torch.int64)
170 if native_w4:
171 kwargs = {
172 "bias": [],
173 "scale": None,
174 "antiquant_scale": [right_scale],
175 "per_token_scale": [left_scale],
176 "group_list": group_list,
177 "group_type": group_type,
178 "output_dtype": output_dtype,
179 "group_list_type": GROUP_LIST_OFFSETS,
180 "split_item": GMM_SPLIT_ANTIQ,
181 "x_dtype": self.activation_dtype,
182 "weight_dtype": self.weight_dtype,
183 }
184 # Dynamic MX quantization returns raw uint8 E8M0 bytes for both
185 # block sizes. The explicit metadata tells CANN how to decode
186 # those bytes as float8_e8m0fnu in the grouped-matmul kernel.
187 if use_e8m0_scale:
188 kwargs["per_token_scale_dtype"] = self.e8m0_dtype
189 return self._torch_npu.npu_grouped_matmul([left], [right], **kwargs)[0]
191 kwargs = {
192 "bias": [],
193 "scale": [right_scale],
194 "per_token_scale": [left_scale],
195 "group_list": group_list,
196 "group_type": group_type,
197 "output_dtype": output_dtype,
198 "group_list_type": GROUP_LIST_OFFSETS,
199 "split_item": split_item,
200 }
201 if use_e8m0_scale:
202 kwargs["scale_dtype"] = self.e8m0_dtype
203 kwargs["per_token_scale_dtype"] = self.e8m0_dtype
204 return self._torch_npu.npu_grouped_matmul([left], [right], **kwargs)[0]
206 def pack_mxfp4_weight(
207 self,
208 tensor: torch.Tensor,
209 *,
210 block_size: int = MX_BLOCK_SIZE_32,
211 ) -> tuple[torch.Tensor, torch.Tensor]:
212 """Quantize and prepare one ``[E, N, K]`` weight for A5 W4 GMM.
214 The native API returns two FP4 values per byte plus raw E8M0 scales.
215 A5 grouped matmul requires the packed data in its NZ layout and the
216 scale metadata transposed in the matching
217 ``[E, ceil(K / 64), N, 2]`` contract (the final pair represents the
218 two 32-value MX blocks covered by one packed scale tile). This method
219 deliberately does not materialize an intermediate MXFP8 weight.
220 """
222 if tensor.ndim != GROUPED_NDIM:
223 raise ValueError(
224 "Native MXFP4 weight preparation expects [experts, rows, K], "
225 f"got shape {tuple(tensor.shape)}."
226 )
227 if block_size != MX_BLOCK_SIZE_32:
228 raise ValueError(
229 "Native Packed MXFP4 weights use block_size=32, "
230 f"got {block_size}."
231 )
232 if tensor.shape[-1] % block_size:
233 raise ValueError(
234 "Native MXFP4 weight K dimension must be divisible by 32, "
235 f"got shape {tuple(tensor.shape)}."
236 )
237 packed, scale = self.dynamic_mx_quant(
238 tensor,
239 axis=-1,
240 block_size=block_size,
241 quant_dtype=self.weight_dtype,
242 round_mode="rint",
243 scale_alg=MX_SCALE_ALG_OCP,
244 )
245 packed_nz = self._torch_npu.npu_format_cast(
246 packed,
247 NPU_FORMAT_FRACTAL_NZ,
248 customize_dtype=self.activation_dtype,
249 input_dtype=self.weight_dtype,
250 ).transpose(-2, -1)
251 expected_scale_tiles = (tensor.shape[-1] + MX_SCALE_TILE - 1) // MX_SCALE_TILE
252 if (
253 scale.ndim != SCALE_NDIM
254 or tuple(scale.shape[:2]) != tuple(tensor.shape[:2])
255 or scale.shape[2] != expected_scale_tiles
256 or scale.shape[-1] != MX_SCALE_SLOTS
257 ):
258 raise ValueError(
259 "Native MXFP4 scale must have [E, N, ceil(K/64), 2] layout, "
260 f"got shape {tuple(scale.shape)}."
261 )
262 return packed_nz, scale.permute(0, 2, 1, 3)
265@lru_cache(maxsize=1)
266def _get_w4a8_npu_ops() -> W4A8NpuOps:
267 """Return the process-local native W4A8 operator adapter."""
269 try:
270 import torch_npu # pylint: disable=C0415
271 except ImportError as exc:
272 raise W4A8CapabilityError(
273 "Native W4A8 requires torch_npu, but it is not importable."
274 ) from exc
275 return W4A8NpuOps(torch_npu)
278def validate_w4a8_gmm_runtime() -> None:
279 """Fail during model setup when native W4A8 is unavailable."""
281 _get_w4a8_npu_ops()
284__all__ = [
285 "W4A8CapabilityError",
286 "W4A8NpuOps",
287 "validate_w4a8_gmm_runtime",
288]