Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / components / quantization / functional / w4a8_gmm_func.py: 46%
48 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"""W4A8 format hooks for the shared native grouped-linear flow."""
17import torch # pylint: disable=forbidden-backend-import
19from hyper_parallel.components.quantization.config import (
20 GMM_SPLIT_ANTIQ,
21 GMM_SPLIT_SCALE,
22 GROUP_LIST_COUNTS,
23 GROUP_LIST_OFFSETS,
24 MX_BLOCK_SIZE_32,
25)
26from hyper_parallel.components.quantization.functional.base_gmm_func import GroupedLinear
27from hyper_parallel.components.quantization.quantizers.w4a8 import (
28 W4A8Quantizer,
29)
30from hyper_parallel.components.quantization.tensor import QuantizedTensor
33class W4A8GroupedLinear(GroupedLinear):
34 """Bind W4A8's target-layout views and native GMM lowering."""
36 FORMAT_NAME = "Native W4A8"
38 def __init__(self, block_size: int = MX_BLOCK_SIZE_32) -> None:
39 """Create one reusable W4A8 strategy and own its quantizer."""
41 self.quantizer = W4A8Quantizer(block_size)
42 self.block_size = self.quantizer.block_size
44 def normalize_group_list(
45 self,
46 group_list: torch.Tensor,
47 group_list_type: int,
48 ) -> tuple[torch.Tensor, int]:
49 """Return cumulative offsets/type 0 for W4A8 quantization and GMM.
51 Router output is normally per-expert counts (type 1), so convert it
52 once here. Existing cumulative offsets (type 0) pass through unchanged.
53 The native W4A8 grouped-matmul adapter uses the offsets contract.
54 """
55 if group_list_type == GROUP_LIST_COUNTS:
56 return torch.cumsum(group_list, dim=0), GROUP_LIST_OFFSETS
57 return group_list, GROUP_LIST_OFFSETS
59 def output_features(self, weight: torch.Tensor) -> int:
60 """Return the output width from the live ``[E, O, K]`` weight."""
61 return weight.shape[-2]
63 @property
64 def weight_forward_is_rowwise(self) -> bool:
65 """Use the row-wise weight view in the forward NN GMM."""
66 return True
68 def quantize_input(
69 self,
70 inputs: torch.Tensor,
71 *,
72 rowwise: bool,
73 colwise: bool,
74 group_list: torch.Tensor,
75 group_list_type: int,
76 ) -> QuantizedTensor:
77 """Quantize grouped activations for the native W4A8 input role."""
78 return self.quantizer.quantize_activation(
79 inputs,
80 rowwise=rowwise,
81 colwise=colwise,
82 group_list=group_list,
83 group_list_type=group_list_type,
84 )
86 def quantize_weight(
87 self,
88 weight: torch.Tensor,
89 *,
90 rowwise: bool,
91 colwise: bool,
92 ) -> QuantizedTensor:
93 """Quantize expert weights for the native W4A8 weight role."""
94 return self.quantizer.quantize_weight(
95 weight,
96 rowwise=rowwise,
97 colwise=colwise,
98 )
100 def quantize_grad_output(
101 self,
102 grad_output: torch.Tensor,
103 *,
104 rowwise: bool,
105 colwise: bool,
106 group_list: torch.Tensor,
107 group_list_type: int,
108 ) -> QuantizedTensor:
109 """Quantize output gradients for the native W4A8 gradient role."""
110 return self.quantizer.quantize_activation(
111 grad_output,
112 rowwise=rowwise,
113 colwise=colwise,
114 group_list=group_list,
115 group_list_type=group_list_type,
116 )
118 @staticmethod
119 def _select_operands(left, right, layout: str):
120 """Select directional data and scales for one native GMM layout."""
121 if layout == "NN":
122 left_data, left_scale = left.row_data, left.row_scale
123 right_data, right_scale = right.row_data, right.row_scale
124 elif layout == "NT":
125 left_data, left_scale = left.row_data, left.row_scale
126 right_data, right_scale = right.col_data, right.col_scale
127 elif layout == "TN":
128 # Preserve transpose strides so the native MX kernel sees the
129 # same transpose state on the operand and its scale tensor.
130 left_data = left.col_data.transpose(0, 1)
131 left_scale = left.col_scale.transpose(0, 1)
132 right_data, right_scale = right.col_data, right.col_scale
133 else:
134 raise ValueError(
135 f"Unsupported W4A8 GMM layout {layout!r}; expected NN, NT, or TN."
136 )
137 if any(value is None for value in (left_data, left_scale, right_data, right_scale)):
138 raise ValueError(
139 f"W4A8 GMM layout {layout!r} requires both directional operands."
140 )
141 return left_data, left_scale, right_data, right_scale
143 def grouped_matmul(
144 self,
145 left: QuantizedTensor,
146 right: QuantizedTensor,
147 *,
148 layout: str,
149 group_list: torch.Tensor,
150 group_type: int,
151 group_list_type: int,
152 output_dtype: torch.dtype,
153 ) -> torch.Tensor:
154 """Execute one native W4A8 grouped matrix multiplication."""
155 del group_list_type
156 left_data, left_scale, right_data, right_scale = self._select_operands(
157 left,
158 right,
159 layout,
160 )
161 # NN/NT are native MXFP8 x Packed-MXFP4 W4A8. TN computes the dense
162 # master-weight gradient and therefore remains MXFP8 x MXFP8.
163 native_w4 = layout in ("NN", "NT")
164 # A8W4 grouped matmul consumes E8M0 per-token scales for both MX
165 # block sizes. The TN weight-gradient path uses the same contract.
166 use_e8m0_scale = True
167 return self.quantizer.npu_ops.grouped_matmul(
168 left_data,
169 right_data,
170 right_scale,
171 left_scale=left_scale,
172 group_list=group_list,
173 group_type=group_type,
174 output_dtype=output_dtype,
175 use_e8m0_scale=use_e8m0_scale,
176 split_item=GMM_SPLIT_SCALE if use_e8m0_scale else GMM_SPLIT_ANTIQ,
177 native_w4=native_w4,
178 )
180__all__ = ["W4A8GroupedLinear"]