« 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 strategy for the shared grouped-linear lifecycle."""
17import torch # pylint: disable=forbidden-backend-import
19from hyper_parallel.components.quantization.config import (
20 GROUP_LIST_COUNTS,
21 GROUP_LIST_OFFSETS,
22 MX_BLOCK_SIZE_32,
23)
24from hyper_parallel.components.quantization.functional.base_gmm_func import GroupedLinear
25from hyper_parallel.components.quantization.quantizers.fake_w4a8 import FakeW4A8Quantizer
26from hyper_parallel.components.quantization.tensor import MXFP8Tensor
29class FakeW4A8GroupedLinear(GroupedLinear):
30 """Use fake MXFP4 weights with the ordinary MXFP8 GMM contract.
32 The common flow owns the three autograd phases. This strategy declares its
33 forward weight view and provides the two backward quantizers:
35 * ``NN`` output and ``NT`` input-gradient (dA) use ``npu_dynamic_mx_quant``;
36 * ``TN`` weight-gradient (dX) uses grouped dynamic quantization for block32.
37 """
39 FORMAT_NAME = "Fake W4A8"
41 def __init__(self, block_size: int = MX_BLOCK_SIZE_32) -> None:
42 """Create a fake-W4A8 grouped strategy for one MX block size."""
43 self.quantizer = FakeW4A8Quantizer(block_size)
44 self.block_size = self.quantizer.block_size
46 def normalize_group_list(
47 self,
48 group_list: torch.Tensor,
49 group_list_type: int,
50 ) -> tuple[torch.Tensor, int]:
51 """Return cumulative offsets/type 0 for fake-W4A8 grouped matmul.
53 Router output is normally per-expert counts (type 1); convert those
54 counts once and pass existing cumulative offsets (type 0) through.
55 The fake-W4A8 NPU adapter also invokes grouped matmul with type 0.
56 """
57 if group_list_type == GROUP_LIST_COUNTS:
58 return torch.cumsum(group_list, dim=0), GROUP_LIST_OFFSETS
59 return group_list, GROUP_LIST_OFFSETS
61 def output_features(self, weight: torch.Tensor) -> int:
62 """Return the output width of a GMM-shaped ``[E,K,N]`` weight."""
63 return weight.shape[-2]
65 @property
66 def weight_forward_is_rowwise(self) -> bool:
67 """Use the column-wise weight view required by fake W4A8 GMM."""
68 return False
70 def quantize_input(
71 self,
72 inputs: torch.Tensor,
73 *,
74 rowwise: bool,
75 colwise: bool,
76 group_list: torch.Tensor,
77 group_list_type: int,
78 ) -> MXFP8Tensor:
79 """Quantize routed input activations in the requested directions."""
80 return self.quantizer.quantize_activation(
81 inputs,
82 rowwise=rowwise,
83 colwise=colwise,
84 group_list=group_list,
85 group_list_type=group_list_type,
86 )
88 def quantize_weight(
89 self,
90 weight: torch.Tensor,
91 *,
92 rowwise: bool,
93 colwise: bool,
94 ) -> MXFP8Tensor:
95 """Quantize grouped master weights through fake MXFP4."""
96 return self.quantizer.quantize_weight(
97 weight,
98 rowwise=rowwise,
99 colwise=colwise,
100 )
102 def quantize_grad_output(
103 self,
104 grad_output: torch.Tensor,
105 *,
106 rowwise: bool,
107 colwise: bool,
108 group_list: torch.Tensor,
109 group_list_type: int,
110 ) -> MXFP8Tensor:
111 """Quantize routed output gradients in the requested directions."""
112 return self.quantizer.quantize_activation(
113 grad_output,
114 rowwise=rowwise,
115 colwise=colwise,
116 group_list=group_list,
117 group_list_type=group_list_type,
118 )
120 @staticmethod
121 def _select_operands(left, right, layout: str):
122 """Select directional data and scales for one GMM layout."""
123 if layout == "NN":
124 left_data, left_scale = left.row_data, left.row_scale
125 right_data, right_scale = right.col_data, right.col_scale
126 elif layout == "NT":
127 left_data, left_scale = left.row_data, left.row_scale
128 right_data = right.row_data.transpose(-1, -2)
129 right_scale = right.row_scale.transpose(-3, -2)
130 elif layout == "TN":
131 left_data = left.col_data.transpose(0, 1)
132 left_scale = left.col_scale.transpose(0, 1)
133 right_data, right_scale = right.col_data, right.col_scale
134 else:
135 raise ValueError(
136 f"Unsupported fake W4A8 GMM layout {layout!r}; expected NN, NT, or TN."
137 )
138 if any(value is None for value in (left_data, left_scale, right_data, right_scale)):
139 raise ValueError(
140 f"Fake W4A8 GMM layout {layout!r} requires both directional operands."
141 )
142 return left_data, left_scale, right_data, right_scale
144 def grouped_matmul(
145 self,
146 left: MXFP8Tensor,
147 right: MXFP8Tensor,
148 *,
149 layout: str,
150 group_list: torch.Tensor,
151 group_type: int,
152 group_list_type: int,
153 output_dtype: torch.dtype,
154 ) -> torch.Tensor:
155 """Execute fake-W4A8 grouped matmul for NN, NT, or TN layout."""
156 del group_list_type
157 left_data, left_scale, right_data, right_scale = self._select_operands(
158 left, right, layout
159 )
160 return self.quantizer.npu_ops.grouped_matmul(
161 left_data,
162 right_data,
163 right_scale,
164 left_scale=left_scale,
165 group_list=group_list,
166 group_type=group_type,
167 block_size=self.block_size,
168 output_dtype=output_dtype,
169 )
172__all__ = ["FakeW4A8GroupedLinear"]