« 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"""HiFloat8 format hooks for the shared grouped-linear lifecycle."""
17from typing import Optional
19import torch # pylint: disable=forbidden-backend-import
20from torch.autograd.function import once_differentiable # pylint: disable=forbidden-backend-import
22from hyper_parallel.components.quantization.functional.base_gmm_func import (
23 GroupedLinear,
24)
25from hyper_parallel.components.quantization.ops.npu_hifloat8 import (
26 hifloat8_grouped_matmul,
27)
28from hyper_parallel.components.quantization.quantizers.hifloat8 import (
29 GRADIENT_FORMAT_MAX,
30 INPUT_WEIGHT_FORMAT_MAX,
31 HiFloat8Quantizer,
32)
33from hyper_parallel.components.quantization.tensor import (
34 HiFloat8Tensor,
35 QuantizedTensorStorage,
36)
39class HiFloat8GroupedLinear(GroupedLinear):
40 """Bind HiFloat8 role quantizers and GMM lowering to the common flow."""
42 FORMAT_NAME = "HiFloat8"
43 CAST_ROUTED_OUTPUT_TO_INPUT_DTYPE = True
45 def __init__(
46 self,
47 input_quantizer: Optional[HiFloat8Quantizer] = None,
48 weight_quantizer: Optional[HiFloat8Quantizer] = None,
49 grad_output_quantizer: Optional[HiFloat8Quantizer] = None,
50 ) -> None:
51 """Create or accept the three independent HiFloat8 role quantizers."""
53 self.input_quantizer = (
54 input_quantizer
55 if input_quantizer is not None
56 else HiFloat8Quantizer(fp8_max=INPUT_WEIGHT_FORMAT_MAX)
57 )
58 self.weight_quantizer = (
59 weight_quantizer
60 if weight_quantizer is not None
61 else HiFloat8Quantizer(fp8_max=INPUT_WEIGHT_FORMAT_MAX)
62 )
63 self.grad_output_quantizer = (
64 grad_output_quantizer
65 if grad_output_quantizer is not None
66 else HiFloat8Quantizer(fp8_max=GRADIENT_FORMAT_MAX)
67 )
69 def _validate_inputs(
70 self,
71 inputs: torch.Tensor,
72 weight: torch.Tensor,
73 group_list: torch.Tensor,
74 group_list_type: int,
75 ) -> None:
76 """Preserve the HiFloat8 dtype, device, and empty-group contract."""
78 super()._validate_inputs(
79 inputs,
80 weight,
81 group_list,
82 group_list_type,
83 )
84 if inputs.dtype not in (torch.float16, torch.bfloat16) or weight.dtype not in (
85 torch.float16,
86 torch.bfloat16,
87 ):
88 raise TypeError(
89 "HiFloat8 grouped linear inputs and weight must use float16 or bfloat16."
90 )
91 if group_list.dtype != torch.int64:
92 raise TypeError("HiFloat8 grouped linear group_list must use torch.int64.")
93 if inputs.device != weight.device or group_list.device != inputs.device:
94 raise ValueError(
95 "HiFloat8 grouped linear inputs, weight, and group_list must be on the same device."
96 )
97 if inputs.shape[0] == 0 and torch.any(group_list != 0).item():
98 raise ValueError(
99 "An empty HiFloat8 grouped input requires an all-zero group_list."
100 )
102 def normalize_group_list(
103 self,
104 group_list: torch.Tensor,
105 group_list_type: int,
106 ) -> tuple[torch.Tensor, int]:
107 """Keep the caller's counts or offsets because HiFloat8 accepts both."""
109 return group_list, group_list_type
111 def output_features(self, weight: torch.Tensor) -> int:
112 """Return the output width from the live ``[E, O, K]`` weight."""
114 return weight.shape[-2]
116 def weight_quantization_directions(
117 self,
118 needs_grad_input: bool,
119 ) -> tuple[bool, bool]:
120 """Build the column view for forward and row view only for dgrad."""
122 return needs_grad_input, True
124 def retain_weight_backward(
125 self,
126 weight_quant: QuantizedTensorStorage,
127 needs_grad_input: bool,
128 ) -> None:
129 """Release the forward column view and retain the dgrad row view."""
131 weight_quant.update_usage(rowwise=needs_grad_input, colwise=False)
133 def quantize_input(
134 self,
135 inputs: torch.Tensor,
136 *,
137 rowwise: bool,
138 colwise: bool,
139 group_list: torch.Tensor,
140 group_list_type: int,
141 ) -> HiFloat8Tensor:
142 """Quantize expert-major activations with the input-role recipe."""
144 return self.input_quantizer.quantize(
145 inputs,
146 group_list=group_list,
147 group_list_type=group_list_type,
148 rowwise=rowwise,
149 colwise=colwise,
150 )
152 def quantize_weight(
153 self,
154 weight: torch.Tensor,
155 *,
156 rowwise: bool,
157 colwise: bool,
158 ) -> HiFloat8Tensor:
159 """Quantize GMM-ready expert weights with the weight-role recipe."""
161 return self.weight_quantizer.quantize(
162 weight,
163 rowwise=rowwise,
164 colwise=colwise,
165 )
167 def quantize_grad_output(
168 self,
169 grad_output: torch.Tensor,
170 *,
171 rowwise: bool,
172 colwise: bool,
173 group_list: torch.Tensor,
174 group_list_type: int,
175 ) -> HiFloat8Tensor:
176 """Quantize output gradients with the gradient-role recipe."""
178 return self.grad_output_quantizer.quantize(
179 grad_output,
180 group_list=group_list,
181 group_list_type=group_list_type,
182 rowwise=rowwise,
183 colwise=colwise,
184 )
186 def grouped_matmul(
187 self,
188 left: HiFloat8Tensor,
189 right: HiFloat8Tensor,
190 *,
191 layout: str,
192 group_list: torch.Tensor,
193 group_type: int,
194 group_list_type: int,
195 output_dtype: torch.dtype,
196 ) -> torch.Tensor:
197 """Execute one HiFloat8 NN, NT, or TN grouped matrix multiplication."""
199 return hifloat8_grouped_matmul(
200 left,
201 right,
202 layout=layout,
203 group_list=group_list,
204 group_type=group_type,
205 group_list_type=group_list_type,
206 output_dtype=output_dtype,
207 )
209 @once_differentiable
210 def backward(
211 self,
212 ctx: torch.autograd.function.FunctionCtx,
213 grad_output: torch.Tensor,
214 ) -> tuple[Optional[torch.Tensor], Optional[torch.Tensor]]:
215 """Preserve HiFloat8's first-order-only backward contract."""
217 return super().backward(ctx, grad_output)
220__all__ = ["HiFloat8GroupedLinear"]