Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / components / quantization / ops / npu_w4a8.py: 25%
95 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"""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
23_A5_DEVICE_MARKERS = ("Ascend950", "Ascend910_95")
24_SUPPORTED_BLOCK_SIZES = (32, 128)
27class W4A8CapabilityError(RuntimeError):
28 """Report that the active torch_npu stack lacks native W4A8 support."""
31def transform_grouped_scale(
32 input_tensor: torch.Tensor,
33 input_scale_mxfp8: torch.Tensor,
34 group_list: torch.Tensor,
35 block_size: int,
36) -> torch.Tensor:
37 """Insert expert-boundary scale rows for block-128 grouped matmul."""
39 if input_scale_mxfp8.ndim != 3 or group_list.ndim != 1:
40 raise ValueError("Grouped MXFP8 scales require a 3D scale and 1D group list.")
41 if block_size not in _SUPPORTED_BLOCK_SIZES:
42 raise ValueError(f"Unsupported W4A8 block size {block_size}.")
43 indices = torch.arange(
44 group_list.shape[0],
45 dtype=group_list.dtype,
46 device=group_list.device,
47 )
48 zero_rows = group_list // (block_size // 2) + indices
49 scale_repeat = input_scale_mxfp8.reshape(-1, input_scale_mxfp8.shape[1]).repeat_interleave(4, dim=0)
50 scale_packed = scale_repeat.reshape(input_scale_mxfp8.shape[0] * 4, input_scale_mxfp8.shape[1], 2)
51 scale_shape = input_tensor.shape[0] // (block_size // 2) + group_list.shape[0]
52 all_indices = torch.arange(scale_shape, device=scale_packed.device)
53 non_zero_rows = all_indices[~torch.isin(all_indices, zero_rows)]
54 grouped_scale = torch.zeros(
55 scale_shape,
56 input_scale_mxfp8.shape[1],
57 2,
58 dtype=scale_packed.dtype,
59 device=scale_packed.device,
60 )
61 if scale_packed.shape[0] < non_zero_rows.shape[0]:
62 raise ValueError(
63 "Grouped MX scale rows do not cover all non-zero token blocks: "
64 f"required {non_zero_rows.shape[0]}, got {scale_packed.shape[0]}."
65 )
66 # A final partial 256-token source tile emits four physical rows even when
67 # the GMM contract only consumes complete 64-token rows; discard those
68 # excess rows while inserting expert-boundary zero rows.
69 grouped_scale[non_zero_rows] = scale_packed[: non_zero_rows.shape[0]]
70 return grouped_scale
73class W4A8NpuOps:
74 """Adapt the native W4A8 torch_npu operator contract."""
76 def __init__(self, torch_npu_module: ModuleType) -> None:
77 """Validate and retain an imported torch_npu module."""
79 required = (
80 "npu_dynamic_mx_quant",
81 "npu_grouped_matmul",
82 "npu_format_cast",
83 "float4_e2m1fn_x2",
84 "float8_e4m3fn",
85 "float8_e8m0fnu",
86 )
87 missing = [name for name in required if not hasattr(torch_npu_module, name)]
88 if missing:
89 version = getattr(torch_npu_module, "__version__", "unknown")
90 raise W4A8CapabilityError(
91 "The active torch_npu stack does not provide the native W4A8 "
92 f"operator contract; missing {missing}, torch_npu={version}."
93 )
94 npu_handle = getattr(torch_npu_module, "npu", None)
95 try:
96 device_name = (
97 npu_handle.get_device_name()
98 if npu_handle is not None and hasattr(npu_handle, "get_device_name")
99 else ""
100 )
101 except (RuntimeError, TypeError, OSError):
102 device_name = ""
103 if not any(marker in device_name for marker in _A5_DEVICE_MARKERS):
104 raise W4A8CapabilityError(
105 "Native W4A8 is supported only on Ascend 950PR/950DT (A5), "
106 f"but the active device is {device_name or 'unknown'}."
107 )
108 self._torch_npu = torch_npu_module
110 @property
111 def e8m0_dtype(self) -> torch.dtype:
112 """Return the runtime E8M0 scale dtype."""
114 return self._torch_npu.float8_e8m0fnu
116 @property
117 def activation_dtype(self) -> torch.dtype:
118 """Return the runtime MXFP8 data dtype."""
120 return self._torch_npu.float8_e4m3fn
122 @property
123 def weight_dtype(self) -> torch.dtype:
124 """Return the runtime packed MXFP4 data dtype."""
126 return self._torch_npu.float4_e2m1fn_x2
128 def dynamic_mx_quant(
129 self,
130 tensor: torch.Tensor,
131 *,
132 axis: int,
133 block_size: int,
134 quant_dtype: torch.dtype,
135 round_mode: Optional[str] = None,
136 scale_alg: Optional[int] = None,
137 ) -> tuple[torch.Tensor, torch.Tensor]:
138 """Run dynamic MX quantization for one operand direction."""
140 kwargs = {"axis": axis, "dst_type": quant_dtype, "block_size": block_size}
141 if round_mode is not None:
142 kwargs["round_mode"] = round_mode
143 if scale_alg is not None:
144 kwargs["scale_alg"] = scale_alg
145 return self._torch_npu.npu_dynamic_mx_quant(tensor, **kwargs)
147 def grouped_dynamic_mx_quant(
148 self,
149 tensor: torch.Tensor,
150 group_list: torch.Tensor,
151 *,
152 block_size: int,
153 quant_dtype: torch.dtype,
154 ) -> tuple[torch.Tensor, torch.Tensor]:
155 """Run grouped dynamic MX quantization over expert token ranges."""
157 grouped_quant = getattr(
158 self._torch_npu,
159 "npu_grouped_dynamic_mx_quant",
160 None,
161 )
162 if grouped_quant is None:
163 raise W4A8CapabilityError(
164 "The active torch_npu stack does not provide "
165 "npu_grouped_dynamic_mx_quant, which native W4A8 block-32 "
166 "backward requires."
167 )
168 return grouped_quant(
169 tensor,
170 # A5 GroupedDynamicMxQuant's groupIndex contract is int32.
171 group_list.to(torch.int32),
172 round_mode="rint",
173 dst_type=quant_dtype,
174 blocksize=block_size,
175 )
177 def grouped_matmul(
178 self,
179 left: torch.Tensor,
180 right: torch.Tensor,
181 right_scale: torch.Tensor,
182 *,
183 left_scale: torch.Tensor,
184 group_list: torch.Tensor,
185 group_type: int,
186 output_dtype: torch.dtype,
187 use_e8m0_scale: bool,
188 split_item: int,
189 native_w4: bool = False,
190 ) -> torch.Tensor:
191 """Run one native grouped matrix multiplication.
193 ``native_w4`` selects the A5 ``MXFP8 x Packed MXFP4`` contract. The
194 W4A8 weight scale belongs to ``antiquant_scale``; the ordinary
195 ``scale`` input is reserved for the MXFP8 path used by wgrad.
196 """
198 # A5 GroupedMatmul's groupList contract is int64.
199 group_list = group_list.to(torch.int64)
201 if native_w4:
202 kwargs = {
203 "bias": [],
204 "scale": None,
205 "antiquant_scale": [right_scale],
206 "per_token_scale": [left_scale],
207 "group_list": group_list,
208 "group_type": group_type,
209 "output_dtype": output_dtype,
210 "group_list_type": 0,
211 "split_item": 2,
212 "x_dtype": self.activation_dtype,
213 "weight_dtype": self.weight_dtype,
214 }
215 # Block-32 dynamic MX quantization returns raw uint8 E8M0 scales;
216 # block-128 activation scales are converted to FP32 by the
217 # quantizer and therefore must retain the default dtype contract.
218 if use_e8m0_scale:
219 kwargs["per_token_scale_dtype"] = self.e8m0_dtype
220 return self._torch_npu.npu_grouped_matmul([left], [right], **kwargs)[0]
222 kwargs = {
223 "bias": [],
224 "scale": [right_scale],
225 "per_token_scale": [left_scale],
226 "group_list": group_list,
227 "group_type": group_type,
228 "output_dtype": output_dtype,
229 "group_list_type": 0,
230 "split_item": split_item,
231 }
232 if use_e8m0_scale:
233 kwargs["scale_dtype"] = self.e8m0_dtype
234 kwargs["per_token_scale_dtype"] = self.e8m0_dtype
235 return self._torch_npu.npu_grouped_matmul([left], [right], **kwargs)[0]
237 def pack_mxfp4_weight(
238 self,
239 tensor: torch.Tensor,
240 *,
241 block_size: int = 32,
242 ) -> tuple[torch.Tensor, torch.Tensor]:
243 """Quantize and prepare one ``[E, N, K]`` weight for A5 W4 GMM.
245 The native API returns two FP4 values per byte plus raw E8M0 scales.
246 A5 grouped matmul requires the packed data in its NZ layout and the
247 scale metadata transposed in the matching
248 ``[E, ceil(K / 64), N, 2]`` contract (the final pair represents the
249 two 32-value MX blocks covered by one packed scale tile). This method
250 deliberately does not materialize an intermediate MXFP8 weight.
251 """
253 if tensor.ndim != 3:
254 raise ValueError(
255 "Native MXFP4 weight preparation expects [experts, rows, K], "
256 f"got shape {tuple(tensor.shape)}."
257 )
258 if block_size != 32:
259 raise ValueError(
260 "Native Packed MXFP4 weights use block_size=32, "
261 f"got {block_size}."
262 )
263 if tensor.shape[-1] % block_size:
264 raise ValueError(
265 "Native MXFP4 weight K dimension must be divisible by 32, "
266 f"got shape {tuple(tensor.shape)}."
267 )
268 packed, scale = self.dynamic_mx_quant(
269 tensor,
270 axis=-1,
271 block_size=block_size,
272 quant_dtype=self.weight_dtype,
273 round_mode="rint",
274 scale_alg=0,
275 )
276 packed_nz = self._torch_npu.npu_format_cast(
277 packed,
278 29,
279 customize_dtype=self.activation_dtype,
280 input_dtype=self.weight_dtype,
281 ).transpose(-2, -1)
282 expected_scale_tiles = (tensor.shape[-1] + 63) // 64
283 if (
284 scale.ndim != 4
285 or tuple(scale.shape[:2]) != tuple(tensor.shape[:2])
286 or scale.shape[2] != expected_scale_tiles
287 or scale.shape[-1] != 2
288 ):
289 raise ValueError(
290 "Native MXFP4 scale must have [E, N, ceil(K/64), 2] layout, "
291 f"got shape {tuple(scale.shape)}."
292 )
293 return packed_nz, scale.permute(0, 2, 1, 3)
296@lru_cache(maxsize=1)
297def _get_w4a8_npu_ops() -> W4A8NpuOps:
298 """Return the process-local native W4A8 operator adapter."""
300 try:
301 import torch_npu # pylint: disable=C0415
302 except ImportError as exc:
303 raise W4A8CapabilityError(
304 "Native W4A8 requires torch_npu, but it is not importable."
305 ) from exc
306 return W4A8NpuOps(torch_npu)
309def validate_w4a8_gmm_runtime() -> None:
310 """Fail during model setup when native W4A8 is unavailable."""
312 _get_w4a8_npu_ops()
315__all__ = [
316 "W4A8CapabilityError",
317 "W4A8NpuOps",
318 "transform_grouped_scale",
319 "validate_w4a8_gmm_runtime",
320]