« 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 hooks for the shared Dense low-precision autograd flow."""
17from typing import Optional
19import torch # pylint: disable=forbidden-backend-import
21from hyper_parallel.components.quantization.functional.base_linear_func import (
22 LinearStrategy,
23)
24from hyper_parallel.components.quantization.ops.npu_hifloat8 import hifloat8_matmul
25from hyper_parallel.components.quantization.quantizers.hifloat8 import (
26 GRADIENT_FORMAT_MAX,
27 INPUT_WEIGHT_FORMAT_MAX,
28 HiFloat8Quantizer,
29)
30from hyper_parallel.components.quantization.tensor import QuantizedTensorStorage
33class HiFloat8LinearStrategy(LinearStrategy):
34 """Bind HiFloat8 role quantizers and Dense MM to the shared Linear flow."""
36 def __init__(
37 self,
38 input_quantizer: Optional[HiFloat8Quantizer] = None,
39 weight_quantizer: Optional[HiFloat8Quantizer] = None,
40 grad_output_quantizer: Optional[HiFloat8Quantizer] = None,
41 ) -> None:
42 """Create or accept the three independent HiFloat8 role quantizers."""
44 self.input_quantizer = (
45 input_quantizer
46 if input_quantizer is not None
47 else HiFloat8Quantizer(fp8_max=INPUT_WEIGHT_FORMAT_MAX)
48 )
49 self.weight_quantizer = (
50 weight_quantizer
51 if weight_quantizer is not None
52 else HiFloat8Quantizer(fp8_max=INPUT_WEIGHT_FORMAT_MAX)
53 )
54 self.grad_output_quantizer = (
55 grad_output_quantizer
56 if grad_output_quantizer is not None
57 else HiFloat8Quantizer(fp8_max=GRADIENT_FORMAT_MAX)
58 )
60 def quantize_input(
61 self,
62 inputs: torch.Tensor,
63 *,
64 rowwise: bool,
65 colwise: bool,
66 ) -> QuantizedTensorStorage:
67 """Quantize the input with the input-role recipe."""
69 return self.input_quantizer.quantize(
70 inputs,
71 rowwise=rowwise,
72 colwise=colwise,
73 )
75 def quantize_weight(
76 self,
77 weight: torch.Tensor,
78 *,
79 rowwise: bool,
80 colwise: bool,
81 ) -> QuantizedTensorStorage:
82 """Quantize the weight with the weight-role recipe."""
84 return self.weight_quantizer.quantize(
85 weight,
86 rowwise=rowwise,
87 colwise=colwise,
88 )
90 def quantize_grad_output(
91 self,
92 grad_output: torch.Tensor,
93 *,
94 rowwise: bool,
95 colwise: bool,
96 ) -> QuantizedTensorStorage:
97 """Quantize the gradient output with the gradient-role recipe."""
99 return self.grad_output_quantizer.quantize(
100 grad_output,
101 rowwise=rowwise,
102 colwise=colwise,
103 )
105 def matmul(
106 self,
107 left: QuantizedTensorStorage,
108 right: QuantizedTensorStorage,
109 *,
110 layout: str,
111 output_dtype: torch.dtype,
112 ) -> torch.Tensor:
113 """Execute the selected HiFloat8 Dense MM layout."""
115 return hifloat8_matmul(
116 left,
117 right,
118 layout=layout,
119 output_dtype=output_dtype,
120 )
123__all__ = ["HiFloat8LinearStrategy"]