Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / components / quantization / ops / npu_fake_w4a8.py: 22%
129 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"""NPU operators used by the fake-QAT W4A8 grouped-linear strategy.
17The native W4A8 strategy is intentionally not used here. Fake W4A8 keeps the
18master parameter in its source dtype, quantizes it to MXFP4 for every forward,
19decodes that result to an MXFP8 operand, and executes the regular all-MXFP8
20grouped-matmul contract. This mirrors the QAT path used by the original
21grouped W4A8 implementation while keeping the native Packed-MXFP4 storage
22and WeightNz adapter isolated in ``npu_w4a8.py``.
23"""
25from functools import lru_cache
26from types import ModuleType
27from typing import Optional
29import torch # pylint: disable=forbidden-backend-import
31from hyper_parallel.components.quantization.ops.npu_w4a8 import (
32 W4A8CapabilityError,
33)
36_A5_DEVICE_MARKERS = ("Ascend950", "Ascend910_95")
37_SUPPORTED_BLOCK_SIZES = (32, 128)
40class FakeW4A8CapabilityError(W4A8CapabilityError):
41 """Report an unavailable fake-W4A8 NPU operator contract."""
44def _decode_packed_mxfp4(packed: torch.Tensor) -> torch.Tensor:
45 """Decode two E2M1 values in each packed byte without CPU round-trips."""
47 if packed.dtype != torch.uint8:
48 raise TypeError(
49 "Packed MXFP4 data must be uint8, "
50 f"got {packed.dtype}."
51 )
52 # E2M1 codes are ordered as positive values in the low/high nibbles and
53 # their sign-bit variants in the upper half. Keeping the table in the
54 # fake-only module prevents the native path from regressing to a LUT
55 # expansion while preserving the exact MXFP4 E2M1 decode.
56 values = torch.tensor(
57 (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0,
58 0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0),
59 dtype=torch.float16,
60 device=packed.device,
61 )
62 low = values[(packed & 0x0F).long()]
63 high = values[((packed >> 4) & 0x0F).long()]
64 return torch.stack((low, high), dim=-1).flatten(-2, -1)
67def fake_mxfp4_to_mxfp8(
68 packed: torch.Tensor,
69 scale: torch.Tensor,
70 *,
71 mx4_block_size: int = 32,
72 mx8_block_h: int = 32,
73 mx8_block_w: int = 32,
74) -> tuple[torch.Tensor, torch.Tensor]:
75 """Decode packed MXFP4 and requantize it into a full-MXFP8 operand.
77 ``npu_dynamic_mx_quant`` emits one E8M0 exponent for each 32-value MXFP4
78 block. Fake W4A8 groups those blocks into the block-size square used by
79 the ordinary MXFP8 GMM. The returned scale is raw E8M0 for block-32 and
80 is converted to FP32 for block-128, matching the active GMM contracts.
81 """
83 if packed.ndim != 3 or scale.ndim != 4 or scale.shape[-1] != 2:
84 raise ValueError(
85 "Fake W4A8 expects packed [E,H,K/2] and scale [E,H,K/64,2], "
86 f"got packed={tuple(packed.shape)}, scale={tuple(scale.shape)}."
87 )
88 if mx4_block_size != 32 or mx8_block_h not in _SUPPORTED_BLOCK_SIZES:
89 raise ValueError("Fake W4A8 only supports MXFP4=32 and MXFP8 blocks 32/128.")
90 if packed.shape[-1] * 2 % mx4_block_size:
91 raise ValueError("Packed MXFP4 contracting dimension must be divisible by 32.")
92 if packed.shape[1] % mx8_block_h or packed.shape[-1] * 2 % mx8_block_w:
93 raise ValueError(
94 "Fake MXFP4-to-MXFP8 conversion requires dimensions divisible by "
95 f"the output block ({mx8_block_h}, {mx8_block_w})."
96 )
98 decoded = _decode_packed_mxfp4(packed)
99 # Keep the original MXFP4 E8M0 exponents until the output MXFP8 block is
100 # formed. The MXFP4-to-MXFP8 conversion chooses ``s8 = max(s4) - 6``:
101 # the largest E2M1 value is 6, so the scaled E4M3 block stays within its
102 # finite range while the GMM's E8M0 scale restores the same magnitude.
103 # The native MXFP4 scale layout is ``[E, H, ceil(W/64), 2]``. The
104 # trailing pair stores the two 32-value exponents covered by each 64-value
105 # tile; for an odd final tile (notably ``W=32``) the second entry is a
106 # hardware padding value. Flatten the pair and discard only those padded
107 # entries before mapping scales to the actual 32-value MXFP4 blocks.
108 e4_scale = scale.reshape(scale.shape[0], scale.shape[1], -1)
109 e4_scale = e4_scale[..., : packed.shape[-1] * 2 // mx4_block_size]
110 e4_scale = e4_scale.to(torch.float32)
111 h_blocks = packed.shape[1] // mx8_block_h
112 w = packed.shape[-1] * 2
113 w_blocks = w // mx8_block_w
114 scale_blocks = e4_scale.reshape(
115 packed.shape[0], h_blocks, mx8_block_h, w_blocks,
116 mx8_block_w // mx4_block_size,
117 )
118 scale_blocks = scale_blocks.unsqueeze(-1).expand(
119 -1, -1, -1, -1, -1, mx4_block_size
120 ).reshape(
121 packed.shape[0], h_blocks, mx8_block_h, w_blocks, mx8_block_w
122 ).permute(0, 1, 3, 2, 4)
123 fp32_blocks = decoded.reshape(
124 packed.shape[0], h_blocks, mx8_block_h, w_blocks, mx8_block_w
125 ).permute(0, 1, 3, 2, 4)
126 exponent = (scale_blocks.amax(dim=(-2, -1)) - 6.0).clamp(0.0, 255.0)
127 exponent_diff = scale_blocks - exponent[..., None, None]
128 scaled = fp32_blocks * torch.pow(2.0, exponent_diff)
129 fp8_blocks = scaled.clamp(-448.0, 448.0).to(torch.float8_e4m3fn)
130 fp8 = fp8_blocks.permute(0, 1, 3, 2, 4).reshape(
131 packed.shape[0], packed.shape[1], w
132 )
133 if mx8_block_h == 32:
134 return fp8, exponent.to(torch.uint8)
135 return fp8, torch.pow(2.0, exponent.to(torch.float32) - 127.0)
138def expand_scale_to_per_row(scale_2d: torch.Tensor, block_size: int = 32) -> torch.Tensor:
139 """Expand ``[E,H/32,W/32]`` scales for the transposed dgrad operand."""
141 if scale_2d.ndim != 3 or block_size != 32:
142 raise ValueError("Fake block-32 scale expansion expects a 3D scale tensor.")
143 _, h_blocks, w_blocks = scale_2d.shape
144 expanded = scale_2d[:, :, None, :].expand(-1, -1, block_size, -1)
145 per_row = expanded.reshape(scale_2d.shape[0], h_blocks * block_size, w_blocks)
146 return per_row.reshape(scale_2d.shape[0], h_blocks * block_size, w_blocks // 2, 2)
149def convert_mx_scale_to_axis2(
150 scale_mxfp8: torch.Tensor,
151 weight_shape: tuple[int, ...],
152) -> torch.Tensor:
153 """Convert square block scales to the GMM ``[E,K/64,N,2]`` layout."""
155 if scale_mxfp8.ndim != 3 or len(weight_shape) != 3:
156 raise ValueError("Fake block-32 weight scales require 3D inputs.")
157 experts, rows, cols = weight_shape
158 if scale_mxfp8.shape[0] != experts or rows % 64 or cols % 32:
159 raise ValueError(
160 "Fake block-32 weight scale dimensions do not match the GMM weight."
161 )
162 if scale_mxfp8.shape[1] * 32 != rows or scale_mxfp8.shape[2] * 32 != cols:
163 raise ValueError("Fake block-32 scale grid does not match the weight shape.")
164 reshaped = scale_mxfp8.reshape(experts, rows // 64, 2, cols // 32)
165 top = reshaped[:, :, 0, :].unsqueeze(-1).expand(-1, -1, -1, 32)
166 bottom = reshaped[:, :, 1, :].unsqueeze(-1).expand(-1, -1, -1, 32)
167 return torch.stack((top.reshape(experts, rows // 64, cols),
168 bottom.reshape(experts, rows // 64, cols)), dim=-1)
171def transform_grouped_scale(
172 input_tensor: torch.Tensor,
173 input_scale_mxfp8: torch.Tensor,
174 group_list: torch.Tensor,
175 block_size: int,
176) -> torch.Tensor:
177 """Insert zero scale rows at expert boundaries for group-type 2 GMM."""
179 if input_scale_mxfp8.ndim != 3 or group_list.ndim != 1:
180 raise ValueError("Grouped fake-MX scales require a 3D scale and 1D groups.")
181 if block_size not in _SUPPORTED_BLOCK_SIZES:
182 raise ValueError(f"Unsupported fake W4A8 block size {block_size}.")
183 indices = torch.arange(group_list.shape[0], device=group_list.device, dtype=group_list.dtype)
184 zero_rows = group_list // (block_size // 2) + indices
185 repeated = input_scale_mxfp8.reshape(-1, input_scale_mxfp8.shape[1]).repeat_interleave(4, dim=0)
186 packed = repeated.reshape(input_scale_mxfp8.shape[0] * 4, input_scale_mxfp8.shape[1], 2)
187 scale_shape = input_tensor.shape[0] // (block_size // 2) + group_list.shape[0]
188 all_indices = torch.arange(scale_shape, device=packed.device)
189 non_zero = all_indices[~torch.isin(all_indices, zero_rows)]
190 grouped = torch.zeros(
191 scale_shape,
192 input_scale_mxfp8.shape[1],
193 2,
194 dtype=packed.dtype,
195 device=packed.device,
196 )
197 if packed.shape[0] < non_zero.shape[0]:
198 raise ValueError("Grouped fake-MX scales do not cover all token blocks.")
199 grouped[non_zero] = packed[: non_zero.shape[0]]
200 return grouped
203class FakeW4A8NpuOps:
204 """Adapt dynamic MX quantization and ordinary MXFP8 grouped matmul."""
206 def __init__(self, torch_npu_module: ModuleType) -> None:
207 required = (
208 "npu_dynamic_mx_quant",
209 "npu_grouped_dynamic_mx_quant",
210 "npu_grouped_matmul",
211 "float4_e2m1fn_x2",
212 "float8_e4m3fn",
213 "float8_e8m0fnu",
214 )
215 missing = [name for name in required if not hasattr(torch_npu_module, name)]
216 if missing:
217 version = getattr(torch_npu_module, "__version__", "unknown")
218 raise FakeW4A8CapabilityError(
219 "The active torch_npu stack does not provide fake W4A8 operators; "
220 f"missing {missing}, torch_npu={version}."
221 )
222 npu_handle = getattr(torch_npu_module, "npu", None)
223 try:
224 device_name = (
225 npu_handle.get_device_name()
226 if npu_handle is not None and hasattr(npu_handle, "get_device_name")
227 else ""
228 )
229 except (RuntimeError, TypeError, OSError):
230 device_name = ""
231 if not any(marker in device_name for marker in _A5_DEVICE_MARKERS):
232 raise FakeW4A8CapabilityError(
233 "Fake W4A8 requires an Ascend 950/A5 runtime, "
234 f"got {device_name or 'unknown'}."
235 )
236 self._torch_npu = torch_npu_module
238 @property
239 def activation_dtype(self) -> torch.dtype:
240 return self._torch_npu.float8_e4m3fn
242 @property
243 def weight_dtype(self) -> torch.dtype:
244 return self._torch_npu.float4_e2m1fn_x2
246 @property
247 def e8m0_dtype(self) -> torch.dtype:
248 return self._torch_npu.float8_e8m0fnu
250 def dynamic_mx_quant(
251 self,
252 tensor: torch.Tensor,
253 *,
254 axis: int,
255 block_size: int,
256 quant_dtype: torch.dtype,
257 round_mode: Optional[str] = None,
258 scale_alg: Optional[int] = None,
259 ) -> tuple[torch.Tensor, torch.Tensor]:
260 kwargs = {"axis": axis, "dst_type": quant_dtype, "block_size": block_size}
261 if round_mode is not None:
262 kwargs["round_mode"] = round_mode
263 if scale_alg is not None:
264 kwargs["scale_alg"] = scale_alg
265 return self._torch_npu.npu_dynamic_mx_quant(tensor, **kwargs)
267 def grouped_dynamic_mx_quant(
268 self,
269 tensor: torch.Tensor,
270 group_list: torch.Tensor,
271 *,
272 block_size: int,
273 quant_dtype: torch.dtype,
274 ) -> tuple[torch.Tensor, torch.Tensor]:
275 if block_size != 32:
276 raise ValueError("npu_grouped_dynamic_mx_quant is only valid for block_size=32.")
277 return self._torch_npu.npu_grouped_dynamic_mx_quant(
278 tensor,
279 group_list.to(torch.int32),
280 round_mode="rint",
281 dst_type=quant_dtype,
282 blocksize=32,
283 )
285 def grouped_matmul(
286 self,
287 left: torch.Tensor,
288 right: torch.Tensor,
289 right_scale: torch.Tensor,
290 *,
291 left_scale: torch.Tensor,
292 group_list: torch.Tensor,
293 group_type: int,
294 block_size: int,
295 output_dtype: torch.dtype,
296 ) -> torch.Tensor:
297 kwargs = {
298 "bias": [],
299 "scale": [right_scale],
300 "per_token_scale": [left_scale],
301 "group_list": group_list.to(torch.int64),
302 "group_type": group_type,
303 "output_dtype": output_dtype,
304 "group_list_type": 0,
305 "split_item": 3 if block_size == 32 or group_type == 2 else 2,
306 }
307 if block_size == 32 or group_type == 2:
308 kwargs["scale_dtype"] = self.e8m0_dtype
309 kwargs["per_token_scale_dtype"] = self.e8m0_dtype
310 return self._torch_npu.npu_grouped_matmul([left], [right], **kwargs)[0]
313@lru_cache(maxsize=1)
314def _get_fake_w4a8_npu_ops() -> FakeW4A8NpuOps:
315 """Resolve the process-local fake-W4A8 NPU adapter."""
317 try:
318 import torch_npu # pylint: disable=C0415
319 except ImportError as exc:
320 raise FakeW4A8CapabilityError(
321 "Fake W4A8 requires torch_npu, but it is not importable."
322 ) from exc
323 return FakeW4A8NpuOps(torch_npu)
326def validate_fake_w4a8_gmm_runtime() -> None:
327 """Fail during replacement if the fake-W4A8 operator contract is absent."""
329 _get_fake_w4a8_npu_ops()
332__all__ = [
333 "FakeW4A8CapabilityError",
334 "FakeW4A8NpuOps",
335 "convert_mx_scale_to_axis2",
336 "expand_scale_to_per_row",
337 "fake_mxfp4_to_mxfp8",
338 "transform_grouped_scale",
339 "validate_fake_w4a8_gmm_runtime",
340]