« 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 strategy for the shared grouped-linear lifecycle."""
17import torch # pylint: disable=forbidden-backend-import
19from hyper_parallel.components.quantization.functional.base_gmm_func import GroupedLinear
20from hyper_parallel.components.quantization.quantizers.fake_w4a8 import FakeW4A8Quantizer
23class FakeW4A8GroupedLinear(GroupedLinear):
24 """Use fake MXFP4 weights with the ordinary MXFP8 GMM contract.
26 The common flow still owns the three autograd phases. The format-specific
27 hooks below only select operand directions and the two backward quantizers:
29 * ``NN`` output and ``NT`` input-gradient (dA) use ``npu_dynamic_mx_quant``;
30 * ``TN`` weight-gradient (dX) uses grouped dynamic quantization for block32;
31 * block128 uses ordinary dynamic quantization for both backward directions.
32 """
34 FORMAT_NAME = "Fake W4A8"
36 def __init__(self, block_size: int = 32) -> None:
37 self.quantizer = FakeW4A8Quantizer(block_size)
38 self.block_size = self.quantizer.block_size
40 def normalize_group_list(
41 self,
42 group_list: torch.Tensor,
43 group_list_type: int,
44 ) -> tuple[torch.Tensor, int]:
45 if group_list_type == 1:
46 return torch.cumsum(group_list, dim=0), 0
47 return group_list, 0
49 def output_features(self, weight: torch.Tensor) -> int:
50 return weight.shape[-2]
52 def weight_quantization_directions(self, needs_grad_input: bool) -> tuple[bool, bool]:
53 if self.block_size == 128:
54 # Block-128 follows the native G-B layout: forward uses the
55 # source-direction view and dgrad uses its transposed companion.
56 return True, needs_grad_input
57 # Follow the ordinary MXFP8 GMM operand contract: forward consumes the
58 # column-wise weight view, while dgrad consumes the row-wise view after
59 # the right operand is transposed.
60 return needs_grad_input, True
62 def retain_weight_backward(self, weight_quant, needs_grad_input: bool) -> None:
63 if self.block_size == 128:
64 weight_quant.update_usage(rowwise=False, colwise=needs_grad_input)
65 return
66 weight_quant.update_usage(rowwise=needs_grad_input, colwise=False)
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 ):
77 return self.quantizer.quantize_activation(
78 inputs,
79 rowwise=rowwise,
80 colwise=colwise,
81 group_list=group_list,
82 group_list_type=group_list_type,
83 )
85 def quantize_weight(self, weight: torch.Tensor, *, rowwise: bool, colwise: bool):
86 return self.quantizer.quantize_weight(
87 weight,
88 rowwise=rowwise,
89 colwise=colwise,
90 )
92 def quantize_grad_output(
93 self,
94 grad_output: torch.Tensor,
95 *,
96 rowwise: bool,
97 colwise: bool,
98 group_list: torch.Tensor,
99 group_list_type: int,
100 ):
101 return self.quantizer.quantize_activation(
102 grad_output,
103 rowwise=rowwise,
104 colwise=colwise,
105 group_list=group_list,
106 group_list_type=group_list_type,
107 )
109 @staticmethod
110 def _select_operands(left, right, layout: str):
111 block128_layout = (
112 getattr(getattr(right, "quantizer", None), "block_size", None) == 128
113 )
114 if layout == "NN" and block128_layout:
115 left_data, left_scale = left.row_data, left.row_scale
116 right_data, right_scale = right.row_data, right.row_scale
117 elif layout == "NT" and block128_layout:
118 left_data, left_scale = left.row_data, left.row_scale
119 right_data, right_scale = right.col_data, right.col_scale
120 elif layout == "NN":
121 left_data, left_scale = left.row_data, left.row_scale
122 right_data, right_scale = right.col_data, right.col_scale
123 elif layout == "NT":
124 left_data, left_scale = left.row_data, left.row_scale
125 right_data = right.row_data.transpose(-1, -2)
126 right_scale = right.row_scale.transpose(-3, -2)
127 elif layout == "TN":
128 left_data = left.col_data.transpose(0, 1)
129 left_scale = left.col_scale.transpose(0, 1)
130 right_data, right_scale = right.col_data, right.col_scale
131 else:
132 raise ValueError(
133 f"Unsupported fake W4A8 GMM layout {layout!r}; expected NN, NT, or TN."
134 )
135 if any(value is None for value in (left_data, left_scale, right_data, right_scale)):
136 raise ValueError(
137 f"Fake W4A8 GMM layout {layout!r} requires both directional operands."
138 )
139 return left_data, left_scale, right_data, right_scale
141 def grouped_matmul(
142 self,
143 left,
144 right,
145 *,
146 layout: str,
147 group_list: torch.Tensor,
148 group_type: int,
149 group_list_type: int,
150 output_dtype: torch.dtype,
151 ) -> torch.Tensor:
152 del group_list_type
153 left_data, left_scale, right_data, right_scale = self._select_operands(
154 left, right, layout
155 )
156 return self.quantizer.npu_ops.grouped_matmul(
157 left_data,
158 right_data,
159 right_scale,
160 left_scale=left_scale,
161 group_list=group_list,
162 group_type=group_type,
163 block_size=self.block_size,
164 output_dtype=output_dtype,
165 )
168__all__ = ["FakeW4A8GroupedLinear"]