« 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"""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 GroupedLinear
23from hyper_parallel.components.quantization.ops.npu_hifloat8 import (
24 hifloat8_grouped_matmul,
25)
26from hyper_parallel.components.quantization.quantizers.hifloat8 import (
27 GRADIENT_FORMAT_MAX,
28 INPUT_WEIGHT_FORMAT_MAX,
29 HiFloat8Quantizer,
30)
31from hyper_parallel.components.quantization.tensor import (
32 HiFloat8Tensor,
33)
36class HiFloat8GroupedLinear(GroupedLinear):
37 """Bind HiFloat8 role quantizers and GMM lowering to the common flow."""
39 FORMAT_NAME = "HiFloat8"
41 def __init__(
42 self,
43 input_quantizer: Optional[HiFloat8Quantizer] = None,
44 weight_quantizer: Optional[HiFloat8Quantizer] = None,
45 grad_output_quantizer: Optional[HiFloat8Quantizer] = None,
46 ) -> None:
47 """Create or accept the three independent HiFloat8 role quantizers."""
49 self.input_quantizer = (
50 input_quantizer
51 if input_quantizer is not None
52 else HiFloat8Quantizer(fp8_max=INPUT_WEIGHT_FORMAT_MAX)
53 )
54 self.weight_quantizer = (
55 weight_quantizer
56 if weight_quantizer is not None
57 else HiFloat8Quantizer(fp8_max=INPUT_WEIGHT_FORMAT_MAX)
58 )
59 self.grad_output_quantizer = (
60 grad_output_quantizer
61 if grad_output_quantizer is not None
62 else HiFloat8Quantizer(fp8_max=GRADIENT_FORMAT_MAX)
63 )
65 def _validate_inputs(
66 self,
67 inputs: torch.Tensor,
68 weight: torch.Tensor,
69 group_list: torch.Tensor,
70 group_list_type: int,
71 ) -> None:
72 """Preserve the HiFloat8 dtype, device, and empty-group contract."""
74 super()._validate_inputs(
75 inputs,
76 weight,
77 group_list,
78 group_list_type,
79 )
80 if inputs.dtype not in (torch.float16, torch.bfloat16) or weight.dtype not in (
81 torch.float16,
82 torch.bfloat16,
83 ):
84 raise TypeError(
85 "HiFloat8 grouped linear inputs and weight must use float16 or bfloat16."
86 )
87 if group_list.dtype != torch.int64:
88 raise TypeError("HiFloat8 grouped linear group_list must use torch.int64.")
89 if inputs.device != weight.device or group_list.device != inputs.device:
90 raise ValueError(
91 "HiFloat8 grouped linear inputs, weight, and group_list must be on the same device."
92 )
93 if inputs.shape[0] == 0 and torch.any(group_list != 0).item():
94 raise ValueError(
95 "An empty HiFloat8 grouped input requires an all-zero group_list."
96 )
98 def normalize_group_list(
99 self,
100 group_list: torch.Tensor,
101 group_list_type: int,
102 ) -> tuple[torch.Tensor, int]:
103 """Preserve counts/type 1 or offsets/type 0 for HiFloat8.
105 Both representations are accepted by its grouped-matmul path, so the
106 tensor and its interpretation tag pass through together unchanged.
107 """
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 @property
117 def weight_forward_is_rowwise(self) -> bool:
118 """Use the column-wise weight view in the forward NN GMM."""
119 return False
121 def quantize_input(
122 self,
123 inputs: torch.Tensor,
124 *,
125 rowwise: bool,
126 colwise: bool,
127 group_list: torch.Tensor,
128 group_list_type: int,
129 ) -> HiFloat8Tensor:
130 """Quantize expert-major activations with the input-role recipe."""
132 return self.input_quantizer.quantize(
133 inputs,
134 group_list=group_list,
135 group_list_type=group_list_type,
136 rowwise=rowwise,
137 colwise=colwise,
138 )
140 def quantize_weight(
141 self,
142 weight: torch.Tensor,
143 *,
144 rowwise: bool,
145 colwise: bool,
146 ) -> HiFloat8Tensor:
147 """Quantize GMM-ready expert weights with the weight-role recipe."""
149 return self.weight_quantizer.quantize(
150 weight,
151 rowwise=rowwise,
152 colwise=colwise,
153 )
155 def quantize_grad_output(
156 self,
157 grad_output: torch.Tensor,
158 *,
159 rowwise: bool,
160 colwise: bool,
161 group_list: torch.Tensor,
162 group_list_type: int,
163 ) -> HiFloat8Tensor:
164 """Quantize output gradients with the gradient-role recipe."""
166 return self.grad_output_quantizer.quantize(
167 grad_output,
168 group_list=group_list,
169 group_list_type=group_list_type,
170 rowwise=rowwise,
171 colwise=colwise,
172 )
174 def grouped_matmul(
175 self,
176 left: HiFloat8Tensor,
177 right: HiFloat8Tensor,
178 *,
179 layout: str,
180 group_list: torch.Tensor,
181 group_type: int,
182 group_list_type: int,
183 output_dtype: torch.dtype,
184 ) -> torch.Tensor:
185 """Execute one HiFloat8 NN, NT, or TN grouped matrix multiplication."""
187 return hifloat8_grouped_matmul(
188 left,
189 right,
190 layout=layout,
191 group_list=group_list,
192 group_type=group_type,
193 group_list_type=group_list_type,
194 output_dtype=output_dtype,
195 )
197 @once_differentiable
198 def backward(
199 self,
200 ctx: torch.autograd.function.FunctionCtx,
201 grad_output: torch.Tensor,
202 ) -> tuple[Optional[torch.Tensor], Optional[torch.Tensor]]:
203 """Preserve HiFloat8's first-order-only backward contract."""
205 return super().backward(ctx, grad_output)
208__all__ = ["HiFloat8GroupedLinear"]