Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / components / quantization / quantizers / fake_w4a8.py: 29%
83 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"""Fake-QAT W4A8 quantizer.
17This quantizer is deliberately separate from ``W4A8Quantizer``. Native W4A8
18stores packed MXFP4/WeightNz data; fake W4A8 stores only the MXFP8 operands
19that the all-MXFP8 GMM consumes after an on-the-fly MXFP4 quantize/dequantize
20step.
21"""
23from typing import Any, Optional
25import torch # pylint: disable=forbidden-backend-import
27from hyper_parallel.components.quantization.config import (
28 GROUPED_NDIM,
29 GROUP_LIST_COUNTS,
30 GROUP_LIST_OFFSETS,
31 MATRIX_NDIM,
32 MX_BLOCK_SIZE_32,
33 MX_SCALE_ALG_OCP,
34 MX_SCALE_TILE,
35 validate_mx_block_size,
36)
37from hyper_parallel.components.quantization.ops.npu_fake_w4a8 import (
38 _get_fake_w4a8_npu_ops,
39 FakeW4A8NpuOps,
40)
41from hyper_parallel.components.quantization.quantizers.base import Quantizer
42from hyper_parallel.components.quantization.tensor import MXFP8Tensor, QuantizedTensor
45class FakeW4A8Quantizer(Quantizer):
46 """Build fake MXFP4 weights and MXFP8 activation/backward operands."""
48 def __init__(
49 self,
50 block_size: int = MX_BLOCK_SIZE_32,
51 ) -> None:
52 """Create a fake-W4A8 quantizer for one supported MX block size."""
53 self.block_size = validate_mx_block_size(block_size)
55 @property
56 def npu_ops(self) -> FakeW4A8NpuOps:
57 """Return the process-local fake-W4A8 NPU adapter."""
58 return _get_fake_w4a8_npu_ops()
60 @staticmethod
61 def _boundaries(
62 group_list: Optional[torch.Tensor],
63 group_list_type: int,
64 ) -> Optional[torch.Tensor]:
65 if group_list is None:
66 return None
67 return (
68 group_list
69 if group_list_type == GROUP_LIST_OFFSETS
70 else torch.cumsum(group_list, dim=0)
71 )
73 def _quantize_activation(
74 self,
75 tensor: torch.Tensor,
76 *,
77 rowwise: bool,
78 colwise: bool,
79 group_list: Optional[torch.Tensor],
80 group_list_type: int,
81 ) -> MXFP8Tensor:
82 """Quantize an activation in the requested row/column directions."""
83 boundaries = self._boundaries(group_list, group_list_type)
84 row_data, row_scale, col_data, col_scale = self._quantize_activation_block32(
85 tensor,
86 rowwise=rowwise,
87 colwise=colwise,
88 boundaries=boundaries,
89 )
90 return MXFP8Tensor(
91 shape=tensor.shape,
92 dtype=tensor.dtype,
93 quantizer=self,
94 row_data=row_data,
95 row_scale=row_scale,
96 col_data=col_data,
97 col_scale=col_scale,
98 device=tensor.device,
99 )
101 def _quantize_activation_axis(
102 self,
103 tensor: torch.Tensor,
104 *,
105 axis: int,
106 ) -> tuple[torch.Tensor, torch.Tensor]:
107 """Apply the shared dynamic-MX activation quantization operation."""
109 return self.npu_ops.dynamic_mx_quant(
110 tensor,
111 axis=axis,
112 block_size=self.block_size,
113 quant_dtype=self.npu_ops.activation_dtype,
114 )
116 def _quantize_activation_block32(
117 self,
118 tensor: torch.Tensor,
119 *,
120 rowwise: bool,
121 colwise: bool,
122 boundaries: Optional[torch.Tensor],
123 ) -> tuple[
124 Optional[torch.Tensor], Optional[torch.Tensor],
125 Optional[torch.Tensor], Optional[torch.Tensor],
126 ]:
127 """Quantize block-32 activations, including the grouped TN wgrad view."""
129 row_data = row_scale = col_data = col_scale = None
130 if rowwise:
131 row_data, row_scale = self._quantize_activation_axis(tensor, axis=-1)
132 if colwise and boundaries is not None:
133 # The TN weight-gradient path needs grouped quantization here.
134 # Its block-32 operator accepts cumulative boundaries as int32.
135 col_data, col_scale = self.npu_ops.grouped_dynamic_mx_quant(
136 tensor,
137 boundaries,
138 block_size=MX_BLOCK_SIZE_32,
139 quant_dtype=self.npu_ops.activation_dtype,
140 )
141 elif colwise:
142 col_data, col_scale = self._quantize_activation_axis(tensor, axis=-2)
143 return row_data, row_scale, col_data, col_scale
145 def _quantize_weight(
146 self,
147 tensor: torch.Tensor,
148 *,
149 rowwise: bool,
150 colwise: bool,
151 ) -> MXFP8Tensor:
152 """Quantize a grouped master weight through the fake W4A8 path."""
153 if tensor.ndim != GROUPED_NDIM:
154 raise ValueError(
155 "Fake W4A8 GMM weights must be [experts,K,N], "
156 f"got shape {tuple(tensor.shape)}."
157 )
158 if any(dimension % MX_BLOCK_SIZE_32 for dimension in tensor.shape[1:]):
159 raise ValueError(
160 "Fake MXFP4 weight dimensions must be divisible by 32, "
161 f"got shape {tuple(tensor.shape)}."
162 )
163 if any(dimension % self.block_size for dimension in tensor.shape[1:]):
164 raise ValueError(
165 "Fake MXFP8 GMM dimensions must be divisible by block_size, "
166 f"got shape {tuple(tensor.shape)}, block_size={self.block_size}."
167 )
168 source, transposed_source = self._prepare_block32_weight_source(tensor)
170 # Both layouts start by generating the same packed MXFP4 representation.
171 # Only the source orientation and the following MXFP8 conversion differ.
172 packed_mxfp4, mxfp4_scale = self._quantize_mxfp4(source)
173 row_data, row_scale, col_data, col_scale = self._convert_weight_block32(
174 packed_mxfp4,
175 mxfp4_scale,
176 transposed_source=transposed_source,
177 rowwise=rowwise,
178 colwise=colwise,
179 )
180 return MXFP8Tensor(
181 shape=tensor.shape,
182 dtype=tensor.dtype,
183 quantizer=self,
184 row_data=row_data,
185 row_scale=row_scale,
186 col_data=col_data,
187 col_scale=col_scale,
188 device=tensor.device,
189 )
191 def _quantize_mxfp4(self, tensor: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
192 """Quantize one source orientation to packed MXFP4 and E8M0 scales."""
194 return self.npu_ops.dynamic_mx_quant(
195 tensor,
196 axis=-1,
197 block_size=MX_BLOCK_SIZE_32,
198 quant_dtype=self.npu_ops.weight_dtype,
199 round_mode="rint",
200 scale_alg=MX_SCALE_ALG_OCP,
201 )
203 def _prepare_block32_weight_source(
204 self,
205 tensor: torch.Tensor,
206 ) -> tuple[torch.Tensor, bool]:
207 """Select the MXFP4 source orientation required by the CANN converter."""
209 rows, cols = tensor.shape[-2:]
210 # CANN requires its penultimate axis to be 64-aligned. Prefer [E,K,N]
211 # so MXFP4 quantization stays on N; transpose only when K is misaligned.
212 transposed_source = rows % MX_SCALE_TILE != 0
213 if transposed_source and cols % MX_SCALE_TILE != 0:
214 raise ValueError(
215 "Fake W4A8 block_size=32 requires at least one of the weight "
216 "matrix dimensions to be divisible by 64 for "
217 "the CANN MXFP4-to-block-MXFP8 conversion; "
218 f"got shape {tuple(tensor.shape)}."
219 )
220 if transposed_source:
221 return tensor.transpose(-2, -1).contiguous(), True
222 return tensor, False
224 def _convert_weight_block32(
225 self,
226 packed_mxfp4: torch.Tensor,
227 mxfp4_scale: torch.Tensor,
228 *,
229 transposed_source: bool,
230 rowwise: bool,
231 colwise: bool,
232 ) -> tuple[
233 Optional[torch.Tensor], Optional[torch.Tensor],
234 Optional[torch.Tensor], Optional[torch.Tensor],
235 ]:
236 """Convert packed MXFP4 through CANN's block-MXFP8 operator."""
238 (
239 fp8_data,
240 scale1,
241 scale2,
242 ) = self.npu_ops.convert_packed_mxfp4_to_block_mxfp8(
243 packed_mxfp4,
244 mxfp4_scale,
245 )
247 if transposed_source:
248 # For source [E,N,K], scale1 follows K (the original contracting
249 # axis) and scale2 follows N. Map them back to [E,K,N] GMM views.
250 fp8_data = fp8_data.transpose(-2, -1).contiguous()
251 row_scale = scale2.transpose(1, 2).contiguous()
252 col_scale = scale1.transpose(1, 2).contiguous()
253 else:
254 # For source [E,K,N], scale1 is [E,K,ceil(N/64),2] and scale2 is
255 # [E,ceil(K/64),N,2], the GMM's rowwise and colwise layouts.
256 row_scale = scale1
257 col_scale = scale2
258 return (
259 fp8_data if rowwise else None,
260 row_scale if rowwise else None,
261 fp8_data if colwise else None,
262 col_scale if colwise else None,
263 )
265 def quantize(
266 self,
267 tensor: torch.Tensor,
268 *,
269 rowwise: bool = False,
270 colwise: bool = False,
271 group_list: Optional[torch.Tensor] = None,
272 group_list_type: int = GROUP_LIST_OFFSETS,
273 role: str = "activation",
274 ) -> MXFP8Tensor:
275 """Build directional fake-MX storage for an activation or weight."""
276 if not isinstance(tensor, torch.Tensor) or isinstance(tensor, QuantizedTensor):
277 raise TypeError("FakeW4A8Quantizer expects a high-precision Tensor.")
278 if not rowwise and not colwise:
279 raise ValueError("FakeW4A8Quantizer requires rowwise or colwise=True.")
280 if group_list_type not in (GROUP_LIST_OFFSETS, GROUP_LIST_COUNTS):
281 raise ValueError("Fake W4A8 group_list_type must be 0 or 1.")
282 if group_list is not None and (
283 not isinstance(group_list, torch.Tensor) or group_list.ndim != 1
284 ):
285 raise ValueError("Fake W4A8 group_list must be one-dimensional.")
286 if role == "weight":
287 if group_list is not None:
288 raise ValueError("Fake W4A8 weight quantization does not accept groups.")
289 return self._quantize_weight(tensor, rowwise=rowwise, colwise=colwise)
290 if role != "activation" or tensor.ndim != MATRIX_NDIM:
291 raise ValueError("Fake W4A8 activation quantization expects a 2D Tensor.")
292 return self._quantize_activation(
293 tensor,
294 rowwise=rowwise,
295 colwise=colwise,
296 group_list=group_list,
297 group_list_type=group_list_type,
298 )
300 def quantize_activation(
301 self, tensor: torch.Tensor, **kwargs: Any
302 ) -> MXFP8Tensor:
303 """Quantize a two-dimensional activation for fake-W4A8 GMM."""
304 return self.quantize(tensor, role="activation", **kwargs)
306 def quantize_weight(
307 self, tensor: torch.Tensor, **kwargs: Any
308 ) -> MXFP8Tensor:
309 """Quantize grouped weights through the fake-MXFP4 path."""
310 return self.quantize(tensor, role="weight", **kwargs)
313__all__ = ["FakeW4A8Quantizer"]