Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / components / quantization / quantizers / fake_w4a8.py: 24%
92 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"""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 Optional
25import torch # pylint: disable=forbidden-backend-import
27from hyper_parallel.components.quantization.ops.npu_fake_w4a8 import (
28 convert_mx_scale_to_axis2,
29 expand_scale_to_per_row,
30 fake_mxfp4_to_mxfp8,
31 transform_grouped_scale,
32 _get_fake_w4a8_npu_ops,
33 FakeW4A8NpuOps,
34)
35from hyper_parallel.components.quantization.quantizers.base import Quantizer
36from hyper_parallel.components.quantization.tensor import MXFP8Tensor, QuantizedTensor
39_SUPPORTED_BLOCK_SIZES = (32, 128)
42class FakeW4A8Quantizer(Quantizer):
43 """Build fake MXFP4 weights and MXFP8 activation/backward operands."""
45 def __init__(self, block_size: int = 32) -> None:
46 if block_size not in _SUPPORTED_BLOCK_SIZES:
47 raise ValueError(
48 f"Fake W4A8 block_size must be 32 or 128, got {block_size}."
49 )
50 self.block_size = block_size
52 @property
53 def npu_ops(self) -> FakeW4A8NpuOps:
54 return _get_fake_w4a8_npu_ops()
56 @staticmethod
57 def _boundaries(
58 group_list: Optional[torch.Tensor],
59 group_list_type: int,
60 ) -> Optional[torch.Tensor]:
61 if group_list is None:
62 return None
63 return group_list if group_list_type == 0 else torch.cumsum(group_list, dim=0)
65 @staticmethod
66 def _scale_to_fp32(scale: torch.Tensor) -> torch.Tensor:
67 return torch.pow(2.0, scale.to(torch.float32) - 127.0)
69 def _quantize_activation(
70 self,
71 tensor: torch.Tensor,
72 *,
73 rowwise: bool,
74 colwise: bool,
75 group_list: Optional[torch.Tensor],
76 group_list_type: int,
77 ) -> MXFP8Tensor:
78 row_data = row_scale = col_data = col_scale = None
79 boundaries = self._boundaries(group_list, group_list_type)
80 if rowwise:
81 row_data, row_scale = self.npu_ops.dynamic_mx_quant(
82 tensor,
83 axis=-1,
84 block_size=self.block_size,
85 quant_dtype=self.npu_ops.activation_dtype,
86 )
87 if self.block_size == 128:
88 row_blocks = (tensor.shape[-1] + self.block_size - 1) // self.block_size
89 row_scale = self._scale_to_fp32(row_scale)
90 row_scale = row_scale.reshape(tensor.shape[0], -1)[..., :row_blocks]
91 if colwise:
92 if self.block_size == 32 and boundaries is not None:
93 # This is the G-B grouped backward quantizer. The operator
94 # requires int32 group indices and has no scale_alg argument.
95 col_data, col_scale = self.npu_ops.grouped_dynamic_mx_quant(
96 tensor,
97 boundaries,
98 block_size=32,
99 quant_dtype=self.npu_ops.activation_dtype,
100 )
101 else:
102 # Block-128 has no grouped_dynamic_mx_quant contract. Quantize
103 # along the token axis and insert the expert-boundary rows in
104 # the scale tensor consumed by group_type=2 GMM.
105 col_data, col_scale = self.npu_ops.dynamic_mx_quant(
106 tensor,
107 axis=-2,
108 block_size=self.block_size,
109 quant_dtype=self.npu_ops.activation_dtype,
110 )
111 if boundaries is not None:
112 col_scale = transform_grouped_scale(
113 tensor,
114 col_scale,
115 boundaries,
116 self.block_size,
117 )
118 return MXFP8Tensor(
119 shape=tensor.shape,
120 dtype=tensor.dtype,
121 quantizer=self,
122 row_data=row_data,
123 row_scale=row_scale,
124 col_data=col_data,
125 col_scale=col_scale,
126 device=tensor.device,
127 )
129 def _quantize_weight(
130 self,
131 tensor: torch.Tensor,
132 *,
133 rowwise: bool,
134 colwise: bool,
135 ) -> MXFP8Tensor:
136 if tensor.ndim != 3:
137 raise ValueError(
138 "Fake W4A8 GMM weights must be [experts,K,N], "
139 f"got shape {tuple(tensor.shape)}."
140 )
141 if any(dimension % 32 for dimension in tensor.shape[1:]):
142 raise ValueError(
143 "Fake MXFP4 weight dimensions must be divisible by 32, "
144 f"got shape {tuple(tensor.shape)}."
145 )
146 if any(dimension % self.block_size for dimension in tensor.shape[1:]):
147 raise ValueError(
148 "Fake MXFP8 GMM dimensions must be divisible by block_size, "
149 f"got shape {tuple(tensor.shape)}, block_size={self.block_size}."
150 )
151 row_data = row_scale = col_data = col_scale = None
153 def _quantize_source(source: torch.Tensor):
154 packed, packed_scale = self.npu_ops.dynamic_mx_quant(
155 source,
156 axis=-1,
157 block_size=32,
158 quant_dtype=self.npu_ops.weight_dtype,
159 round_mode="rint",
160 scale_alg=0,
161 )
162 return fake_mxfp4_to_mxfp8(
163 packed,
164 packed_scale,
165 mx8_block_h=self.block_size,
166 mx8_block_w=self.block_size,
167 )
169 # Block-128 uses the A5 G-B contract directly: the forward view is
170 # generated from the [E,K,N] source and dgrad consumes its transposed
171 # view. Its FP32 square scale grid is exactly what split_item=2
172 # grouped matmul expects.
173 if self.block_size == 128:
174 fp8_data, fp8_grid = _quantize_source(tensor)
175 if rowwise:
176 row_data = fp8_data
177 row_scale = fp8_grid.to(torch.float32)
178 if colwise:
179 col_data = fp8_data.transpose(-2, -1).contiguous()
180 col_scale = fp8_grid.transpose(1, 2).contiguous().to(torch.float32)
181 return MXFP8Tensor(
182 shape=tensor.shape,
183 dtype=tensor.dtype,
184 quantizer=self,
185 row_data=row_data,
186 row_scale=row_scale,
187 col_data=col_data,
188 col_scale=col_scale,
189 device=tensor.device,
190 )
192 # The ordinary MXFP8 GMM contract has two directional views:
193 # row-wise data/scales for the transposed-right dgrad and column-wise
194 # data/scales for the forward NN matmul. Quantize each view along its
195 # actual contracting axis instead of reusing a transposed scale grid;
196 # the latter makes the A5 kernel see mismatched weight/scale
197 # transposition flags for NT.
198 if rowwise:
199 row_data, row_grid = _quantize_source(tensor)
200 if self.block_size == 32:
201 row_scale = expand_scale_to_per_row(row_grid)
202 else:
203 row_scale = row_grid.to(torch.float32)
204 if colwise:
205 col_fp8, col_grid = _quantize_source(tensor.transpose(-2, -1).contiguous())
206 col_data = col_fp8.transpose(-2, -1).contiguous()
207 if self.block_size == 32:
208 col_grid = col_grid.permute(0, 2, 1).contiguous()
209 col_scale = convert_mx_scale_to_axis2(
210 col_grid,
211 tuple(tensor.shape),
212 )
213 else:
214 col_scale = col_grid.transpose(1, 2).contiguous().to(torch.float32)
215 return MXFP8Tensor(
216 shape=tensor.shape,
217 dtype=tensor.dtype,
218 quantizer=self,
219 row_data=row_data,
220 row_scale=row_scale,
221 col_data=col_data,
222 col_scale=col_scale,
223 device=tensor.device,
224 )
226 def quantize(
227 self,
228 tensor: torch.Tensor,
229 *,
230 rowwise: bool = False,
231 colwise: bool = False,
232 group_list: Optional[torch.Tensor] = None,
233 group_list_type: int = 0,
234 role: str = "activation",
235 ) -> MXFP8Tensor:
236 if not isinstance(tensor, torch.Tensor) or isinstance(tensor, QuantizedTensor):
237 raise TypeError("FakeW4A8Quantizer expects a high-precision Tensor.")
238 if not rowwise and not colwise:
239 raise ValueError("FakeW4A8Quantizer requires rowwise or colwise=True.")
240 if group_list_type not in (0, 1):
241 raise ValueError("Fake W4A8 group_list_type must be 0 or 1.")
242 if group_list is not None and (
243 not isinstance(group_list, torch.Tensor) or group_list.ndim != 1
244 ):
245 raise ValueError("Fake W4A8 group_list must be one-dimensional.")
246 if role == "weight":
247 if group_list is not None:
248 raise ValueError("Fake W4A8 weight quantization does not accept groups.")
249 return self._quantize_weight(tensor, rowwise=rowwise, colwise=colwise)
250 if role != "activation" or tensor.ndim != 2:
251 raise ValueError("Fake W4A8 activation quantization expects a 2D Tensor.")
252 return self._quantize_activation(
253 tensor,
254 rowwise=rowwise,
255 colwise=colwise,
256 group_list=group_list,
257 group_list_type=group_list_type,
258 )
260 def quantize_activation(self, tensor: torch.Tensor, **kwargs) -> MXFP8Tensor:
261 return self.quantize(tensor, role="activation", **kwargs)
263 def quantize_weight(self, tensor: torch.Tensor, **kwargs) -> MXFP8Tensor:
264 return self.quantize(tensor, role="weight", **kwargs)
267__all__ = ["FakeW4A8Quantizer"]