« 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"""MXFP8 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_mxfp8 import mxfp8_matmul
25from hyper_parallel.components.quantization.quantizers.mxfp8 import MXFP8Quantizer
26from hyper_parallel.components.quantization.tensor import QuantizedTensorStorage
29class MXFP8LinearStrategy(LinearStrategy):
30 """Bind one MXFP8 quantizer and Dense MM to the shared Linear flow."""
32 def __init__(self, quantizer: Optional[MXFP8Quantizer] = None) -> None:
33 """Create a strategy with an optional injected quantizer."""
35 self.quantizer = quantizer if quantizer is not None else MXFP8Quantizer()
37 def quantize_input(
38 self,
39 inputs: torch.Tensor,
40 *,
41 rowwise: bool,
42 colwise: bool,
43 ) -> QuantizedTensorStorage:
44 """Quantize the input with the shared MXFP8 quantizer."""
46 return self.quantizer.quantize(
47 inputs,
48 rowwise=rowwise,
49 colwise=colwise,
50 )
52 def quantize_weight(
53 self,
54 weight: torch.Tensor,
55 *,
56 rowwise: bool,
57 colwise: bool,
58 ) -> QuantizedTensorStorage:
59 """Quantize the weight with the shared MXFP8 quantizer."""
61 return self.quantizer.quantize(
62 weight,
63 rowwise=rowwise,
64 colwise=colwise,
65 )
67 def quantize_grad_output(
68 self,
69 grad_output: torch.Tensor,
70 *,
71 rowwise: bool,
72 colwise: bool,
73 ) -> QuantizedTensorStorage:
74 """Quantize the gradient output with the shared MXFP8 quantizer."""
76 return self.quantizer.quantize(
77 grad_output,
78 rowwise=rowwise,
79 colwise=colwise,
80 )
82 def matmul(
83 self,
84 left: QuantizedTensorStorage,
85 right: QuantizedTensorStorage,
86 *,
87 layout: str,
88 output_dtype: torch.dtype,
89 ) -> torch.Tensor:
90 """Execute the selected MXFP8 Dense MM layout."""
92 return mxfp8_matmul(
93 left,
94 right,
95 layout=layout,
96 output_dtype=output_dtype,
97 )
100__all__ = ["MXFP8LinearStrategy"]