Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / components / quantization / functional / mxfp8_gmm_func.py: 96%
26 statements
« 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 format hooks for the shared native grouped-linear flow."""
17import torch # pylint: disable=forbidden-backend-import
19from hyper_parallel.components.quantization.functional.base_gmm_func import (
20 GroupedLinear,
21)
22from hyper_parallel.components.quantization.ops.npu_mxfp8 import (
23 mxfp8_grouped_matmul,
24)
25from hyper_parallel.components.quantization.quantizers import (
26 MXFP8Quantizer,
27)
28from hyper_parallel.components.quantization.tensor import (
29 QuantizedTensorStorage,
30)
33class MXFP8GroupedLinear(GroupedLinear):
34 """Bind MXFP8 quantization and GMM lowering to the common native flow."""
36 FORMAT_NAME = "MXFP8"
38 def __init__(self, quantizer: MXFP8Quantizer | None = None) -> None:
39 """Create one reusable MXFP8 strategy and own its quantizer."""
41 self.quantizer = quantizer if quantizer is not None else MXFP8Quantizer()
43 def output_features(self, weight: torch.Tensor) -> int:
44 return weight.shape[-2]
46 def normalize_group_list(
47 self,
48 group_list: torch.Tensor,
49 group_list_type: int,
50 ) -> tuple[torch.Tensor, int]:
51 # MXFP8's operator accepts per-expert token counts as-is.
52 return group_list, group_list_type
54 def weight_quantization_directions(
55 self,
56 needs_grad_input: bool,
57 ) -> tuple[bool, bool]:
58 # MXFP8 weights use the column-wise representation for forward GMM
59 # and the row-wise representation for dgrad.
60 return needs_grad_input, True
62 def retain_weight_backward(
63 self,
64 weight_quant: QuantizedTensorStorage,
65 needs_grad_input: bool,
66 ) -> None:
67 weight_quant.update_usage(rowwise=needs_grad_input, colwise=False)
69 def quantize_input(
70 self,
71 inputs: torch.Tensor,
72 *,
73 rowwise: bool,
74 colwise: bool,
75 group_list: torch.Tensor,
76 group_list_type: int,
77 ):
78 return self.quantizer.quantize(
79 inputs,
80 group_list=group_list,
81 group_list_type=group_list_type,
82 rowwise=rowwise,
83 colwise=colwise,
84 )
86 def quantize_weight(
87 self,
88 weight: torch.Tensor,
89 *,
90 rowwise: bool,
91 colwise: bool,
92 ):
93 return self.quantizer.quantize(
94 weight,
95 rowwise=rowwise,
96 colwise=colwise,
97 )
99 def quantize_grad_output(
100 self,
101 grad_output: torch.Tensor,
102 *,
103 rowwise: bool,
104 colwise: bool,
105 group_list: torch.Tensor,
106 group_list_type: int,
107 ):
108 return self.quantizer.quantize(
109 grad_output,
110 group_list=group_list,
111 group_list_type=group_list_type,
112 rowwise=rowwise,
113 colwise=colwise,
114 )
116 def grouped_matmul(
117 self,
118 left,
119 right,
120 *,
121 layout: str,
122 group_list: torch.Tensor,
123 group_type: int,
124 group_list_type: int,
125 output_dtype: torch.dtype,
126 ) -> torch.Tensor:
127 return mxfp8_grouped_matmul(
128 left,
129 right,
130 layout=layout,
131 group_list=group_list,
132 group_type=group_type,
133 group_list_type=group_list_type,
134 output_dtype=output_dtype,
135 )
138__all__ = ["MXFP8GroupedLinear"]