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