Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / components / quantization / quantizers / w4a8.py: 34%
74 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"""Native W4A8 quantizer for grouped expert training."""
17from typing import Optional
19import torch # pylint: disable=forbidden-backend-import
21from hyper_parallel.components.quantization.ops.npu_w4a8 import (
22 transform_grouped_scale,
23 W4A8NpuOps,
24 _get_w4a8_npu_ops,
25)
26from hyper_parallel.components.quantization.quantizers.base import Quantizer
27from hyper_parallel.components.quantization.tensor import (
28 QuantizedTensor,
29 W4A8Tensor,
30)
32_SUPPORTED_BLOCK_SIZES = (32, 128)
33_ROLE_NAMES = ("activation", "weight")
36class W4A8Quantizer(Quantizer):
37 """Quantize MXFP4 weights and MXFP8 activations for native GMM."""
39 def __init__(
40 self,
41 block_size: int = 32,
42 ) -> None:
43 """Create a block-size-specific native W4A8 quantizer."""
45 if block_size not in _SUPPORTED_BLOCK_SIZES:
46 raise ValueError(
47 f"Native W4A8 block_size must be 32 or 128, got {block_size}."
48 )
49 self.block_size = block_size
51 @property
52 def npu_ops(self) -> W4A8NpuOps:
53 """Return the process-local native W4A8 ops adapter."""
55 return _get_w4a8_npu_ops()
57 @staticmethod
58 def _validate_group_list(
59 group_list: Optional[torch.Tensor],
60 group_list_type: int,
61 ) -> None:
62 if group_list_type not in (0, 1):
63 raise ValueError(
64 "W4A8 group_list_type must be 0 or 1, "
65 f"but got {group_list_type}."
66 )
67 if group_list is not None and (
68 not isinstance(group_list, torch.Tensor) or group_list.ndim != 1
69 ):
70 raise ValueError("W4A8 group_list must be a one-dimensional Tensor.")
72 @staticmethod
73 def _boundaries(
74 group_list: Optional[torch.Tensor],
75 group_list_type: int,
76 ) -> Optional[torch.Tensor]:
77 if group_list is None:
78 return None
79 return group_list if group_list_type == 0 else torch.cumsum(group_list, dim=0)
81 def _quantize_activation(
82 self,
83 tensor: torch.Tensor,
84 *,
85 rowwise: bool,
86 colwise: bool,
87 group_list: Optional[torch.Tensor],
88 group_list_type: int,
89 ) -> W4A8Tensor:
90 row_data = row_scale = col_data = col_scale = None
91 boundaries = self._boundaries(group_list, group_list_type)
92 if rowwise:
93 row_data, row_scale = self.npu_ops.dynamic_mx_quant(
94 tensor,
95 axis=-1,
96 block_size=self.block_size,
97 quant_dtype=self.npu_ops.activation_dtype,
98 )
99 if self.block_size == 128:
100 row_scale = self._scale_to_fp32(row_scale).reshape(
101 tensor.shape[0],
102 -1,
103 )
104 if colwise:
105 if self.block_size == 32 and boundaries is not None:
106 col_data, col_scale = self.npu_ops.grouped_dynamic_mx_quant(
107 tensor,
108 boundaries,
109 block_size=32,
110 quant_dtype=self.npu_ops.activation_dtype,
111 )
112 else:
113 col_data, col_scale = self.npu_ops.dynamic_mx_quant(
114 tensor,
115 axis=-2,
116 block_size=self.block_size,
117 quant_dtype=self.npu_ops.activation_dtype,
118 )
119 if boundaries is not None:
120 col_scale = transform_grouped_scale(
121 tensor,
122 col_scale,
123 boundaries,
124 self.block_size,
125 )
126 return W4A8Tensor(
127 shape=tensor.shape,
128 dtype=tensor.dtype,
129 quantizer=self,
130 row_data=row_data,
131 row_scale=row_scale,
132 col_data=col_data,
133 col_scale=col_scale,
134 device=tensor.device,
135 )
137 @staticmethod
138 def _scale_to_fp32(scale: torch.Tensor) -> torch.Tensor:
139 """Convert an E8M0 exponent tensor to the block-128 FP32 contract."""
141 return torch.pow(2.0, scale.to(torch.float32) - 127)
143 def _quantize_weight(
144 self,
145 tensor: torch.Tensor,
146 *,
147 rowwise: bool,
148 colwise: bool,
149 ) -> W4A8Tensor:
150 if any(dimension % 32 for dimension in tensor.shape[1:]):
151 raise ValueError(
152 "Native MXFP4 weight dimensions must be divisible by 32, "
153 f"got shape {tuple(tensor.shape)}."
154 )
155 row_data = row_scale = col_data = col_scale = None
156 if rowwise:
157 # Forward GMM consumes [E, K/2, N] packed NZ data. Quantizing the
158 # live [E, N, K] source along K gives that native contract.
159 row_data, row_scale = self.npu_ops.pack_mxfp4_weight(
160 tensor.transpose(-2, -1).contiguous(),
161 block_size=32,
162 )
163 if colwise:
164 # Dgrad uses the other matrix direction, so its FP4 scales must be
165 # generated independently from the [E, K, N] source.
166 col_data, col_scale = self.npu_ops.pack_mxfp4_weight(
167 tensor,
168 block_size=32,
169 )
170 return W4A8Tensor(
171 shape=tensor.shape,
172 dtype=tensor.dtype,
173 quantizer=self,
174 row_data=row_data,
175 row_scale=row_scale,
176 col_data=col_data,
177 col_scale=col_scale,
178 device=tensor.device,
179 )
181 def quantize(
182 self,
183 tensor: torch.Tensor,
184 *,
185 rowwise: bool = False,
186 colwise: bool = False,
187 group_list: Optional[torch.Tensor] = None,
188 group_list_type: int = 0,
189 role: str = "activation",
190 ) -> W4A8Tensor:
191 """Quantize an activation or target-layout expert weight.
193 Args:
194 tensor: High-precision source tensor.
195 rowwise: Keep the forward representation.
196 colwise: Keep the representation required by backward.
197 group_list: Optional cumulative expert boundaries or token counts.
198 group_list_type: ``0`` for boundaries and ``1`` for token counts.
199 role: ``"weight"`` quantizes a 3D MXFP4 expert weight with no
200 grouping; ``"activation"`` quantizes a 2D activation with
201 optional grouped boundaries.
202 """
204 if not isinstance(tensor, torch.Tensor) or isinstance(tensor, QuantizedTensor):
205 raise TypeError("W4A8Quantizer expects a high-precision torch.Tensor.")
206 if not rowwise and not colwise:
207 raise ValueError("W4A8Quantizer requires rowwise=True or colwise=True.")
208 if role not in _ROLE_NAMES:
209 raise ValueError(f"Unsupported W4A8 quantization role {role!r}.")
210 self._validate_group_list(group_list, group_list_type)
211 if role == "weight":
212 if tensor.ndim != 3:
213 raise ValueError(
214 "Native W4A8 weights must be three-dimensional, "
215 f"got shape {tuple(tensor.shape)}."
216 )
217 if group_list is not None:
218 raise ValueError("W4A8 weight quantization does not accept group_list.")
219 return self._quantize_weight(
220 tensor,
221 rowwise=rowwise,
222 colwise=colwise,
223 )
224 if tensor.ndim != 2:
225 raise ValueError(
226 "Native W4A8 activations must be two-dimensional, "
227 f"got shape {tuple(tensor.shape)}."
228 )
229 return self._quantize_activation(
230 tensor,
231 rowwise=rowwise,
232 colwise=colwise,
233 group_list=group_list,
234 group_list_type=group_list_type,
235 )
237 def quantize_activation(
238 self,
239 tensor: torch.Tensor,
240 *,
241 rowwise: bool = False,
242 colwise: bool = False,
243 group_list: Optional[torch.Tensor] = None,
244 group_list_type: int = 0,
245 ) -> W4A8Tensor:
246 """Quantize an MXFP8 activation representation."""
248 return self.quantize(
249 tensor,
250 rowwise=rowwise,
251 colwise=colwise,
252 group_list=group_list,
253 group_list_type=group_list_type,
254 role="activation",
255 )
257 def quantize_weight(
258 self,
259 tensor: torch.Tensor,
260 *,
261 rowwise: bool = False,
262 colwise: bool = False,
263 ) -> W4A8Tensor:
264 """Quantize an expert weight into native Packed MXFP4 views."""
266 return self.quantize(
267 tensor,
268 rowwise=rowwise,
269 colwise=colwise,
270 role="weight",
271 )
274__all__ = [
275 "W4A8Quantizer",
276]