Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / components / quantization / ops / npu_fake_w4a8.py: 37%

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"""NPU operators used by the fake-QAT W4A8 grouped-linear strategy. 

16 

17The native W4A8 strategy is intentionally not used here. Fake W4A8 keeps the 

18master parameter in its source dtype, quantizes it to MXFP4 for every forward, 

19decodes that result to an MXFP8 operand, and executes the regular all-MXFP8 

20grouped-matmul contract. This mirrors the QAT path used by the original 

21grouped W4A8 implementation while keeping the native Packed-MXFP4 storage 

22and WeightNz adapter isolated in ``npu_w4a8.py``. 

23""" 

24 

25from functools import lru_cache 

26from types import ModuleType 

27from typing import Callable, Optional 

28 

29import torch # pylint: disable=forbidden-backend-import 

30 

31from hyper_parallel.components.quantization.config import ( 

32 GMM_SPLIT_SCALE, 

33 GROUP_LIST_OFFSETS, 

34 MX_BLOCK_SIZE_32, 

35) 

36from hyper_parallel.components.quantization.ops.npu_w4a8 import W4A8CapabilityError 

37 

38 

39_A5_DEVICE_MARKERS = ("Ascend950", "Ascend910_95") 

40 

41 

42class FakeW4A8CapabilityError(W4A8CapabilityError): 

43 """Report an unavailable fake-W4A8 NPU operator contract.""" 

44 

45 

46class FakeW4A8NpuOps: 

47 """Adapt MX quantization, MXFP4 conversion, and MXFP8 grouped matmul.""" 

48 

49 def __init__(self, torch_npu_module: ModuleType) -> None: 

50 """Validate and retain the active torch_npu operator module.""" 

51 required = ( 

52 "npu_dynamic_mx_quant", 

53 "npu_grouped_dynamic_mx_quant", 

54 "npu_grouped_matmul", 

55 "float4_e2m1fn_x2", 

56 "float8_e4m3fn", 

57 "float8_e8m0fnu", 

58 ) 

59 missing = [name for name in required if not hasattr(torch_npu_module, name)] 

60 if missing: 

61 version = getattr(torch_npu_module, "__version__", "unknown") 

62 raise FakeW4A8CapabilityError( 

63 "The active torch_npu stack does not provide fake W4A8 operators; " 

64 f"missing {missing}, torch_npu={version}." 

65 ) 

66 npu_handle = getattr(torch_npu_module, "npu", None) 

67 try: 

68 device_name = ( 

69 npu_handle.get_device_name() 

70 if npu_handle is not None and hasattr(npu_handle, "get_device_name") 

71 else "" 

72 ) 

73 except (RuntimeError, TypeError, OSError): 

74 device_name = "" 

75 if not any(marker in device_name for marker in _A5_DEVICE_MARKERS): 

76 raise FakeW4A8CapabilityError( 

77 "Fake W4A8 requires an Ascend 950/A5 runtime, " 

78 f"got {device_name or 'unknown'}." 

79 ) 

80 self._torch_npu = torch_npu_module 

81 

82 @property 

83 def activation_dtype(self) -> torch.dtype: 

84 """Return the MXFP8 activation dtype exposed by torch_npu.""" 

85 return self._torch_npu.float8_e4m3fn 

86 

87 @property 

88 def weight_dtype(self) -> torch.dtype: 

89 """Return the packed MXFP4 weight dtype exposed by torch_npu.""" 

90 return self._torch_npu.float4_e2m1fn_x2 

91 

92 @property 

93 def e8m0_dtype(self) -> torch.dtype: 

94 """Return the E8M0 scale dtype used by the grouped-matmul API.""" 

95 return self._torch_npu.float8_e8m0fnu 

96 

97 def dynamic_mx_quant( 

98 self, 

99 tensor: torch.Tensor, 

100 *, 

101 axis: int, 

102 block_size: int, 

103 quant_dtype: torch.dtype, 

104 round_mode: Optional[str] = None, 

105 scale_alg: Optional[int] = None, 

106 ) -> tuple[torch.Tensor, torch.Tensor]: 

107 """Quantize a tensor with the requested dynamic-MX orientation.""" 

108 kwargs = {"axis": axis, "dst_type": quant_dtype, "block_size": block_size} 

109 if round_mode is not None: 

110 kwargs["round_mode"] = round_mode 

111 if scale_alg is not None: 

112 kwargs["scale_alg"] = scale_alg 

113 return self._torch_npu.npu_dynamic_mx_quant(tensor, **kwargs) 

114 

115 def convert_packed_mxfp4_to_block_mxfp8( 

116 self, 

117 packed_mxfp4: torch.Tensor, 

118 mxfp4_scale: torch.Tensor, 

119 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: 

120 """Run CANN's converter and return MXFP8 data with two axis scales.""" 

121 

122 return _get_cann_mx_to_block_mx_quant_op()( 

123 packed_mxfp4, 

124 mxfp4_scale, 

125 dst_type=self.activation_dtype, 

126 x_type=self.weight_dtype, 

127 ) 

128 

129 def grouped_dynamic_mx_quant( 

130 self, 

131 tensor: torch.Tensor, 

132 group_list: torch.Tensor, 

133 *, 

134 block_size: int, 

135 quant_dtype: torch.dtype, 

136 ) -> tuple[torch.Tensor, torch.Tensor]: 

137 """Quantize grouped activation ranges with block-32 MX metadata.""" 

138 if block_size != MX_BLOCK_SIZE_32: 

139 raise ValueError("npu_grouped_dynamic_mx_quant is only valid for block_size=32.") 

140 return self._torch_npu.npu_grouped_dynamic_mx_quant( 

141 tensor, 

142 group_list.to(torch.int32), 

143 round_mode="rint", 

144 dst_type=quant_dtype, 

145 blocksize=MX_BLOCK_SIZE_32, 

146 ) 

147 

148 def grouped_matmul( 

149 self, 

150 left: torch.Tensor, 

151 right: torch.Tensor, 

152 right_scale: torch.Tensor, 

153 *, 

154 left_scale: torch.Tensor, 

155 group_list: torch.Tensor, 

156 group_type: int, 

157 block_size: int, 

158 output_dtype: torch.dtype, 

159 ) -> torch.Tensor: 

160 """Execute one grouped matrix multiplication with fake-MX operands.""" 

161 if block_size != MX_BLOCK_SIZE_32: 

162 raise ValueError("Fake W4A8 supports only block_size=32.") 

163 kwargs = { 

164 "bias": [], 

165 "scale": [right_scale], 

166 "per_token_scale": [left_scale], 

167 "group_list": group_list.to(torch.int64), 

168 "group_type": group_type, 

169 "output_dtype": output_dtype, 

170 "group_list_type": GROUP_LIST_OFFSETS, 

171 "split_item": GMM_SPLIT_SCALE, 

172 "scale_dtype": self.e8m0_dtype, 

173 "per_token_scale_dtype": self.e8m0_dtype, 

174 } 

175 return self._torch_npu.npu_grouped_matmul([left], [right], **kwargs)[0] 

176 

177 

178@lru_cache(maxsize=1) 

179def _get_fake_w4a8_npu_ops() -> FakeW4A8NpuOps: 

180 """Resolve the process-local fake-W4A8 NPU adapter.""" 

181 

182 try: 

183 import torch_npu # pylint: disable=C0415 

184 except ImportError as exc: 

185 raise FakeW4A8CapabilityError( 

186 "Fake W4A8 requires torch_npu, but it is not importable." 

187 ) from exc 

188 return FakeW4A8NpuOps(torch_npu) 

189 

190 

191@lru_cache(maxsize=1) 

192def _get_cann_mx_to_block_mx_quant_op() -> Callable[ 

193 ..., tuple[torch.Tensor, torch.Tensor, torch.Tensor] 

194]: 

195 """Load the CANN operator behind the project-level conversion method.""" 

196 

197 try: 

198 from cann_ops_nn import ( # pylint: disable=C0415 

199 mx_to_block_mx_quant as cann_mx_to_block_mx_quant, 

200 ) 

201 except (ImportError, OSError) as exc: 

202 raise FakeW4A8CapabilityError( 

203 "Fake W4A8 block_size=32 requires CANN 9.2 with " 

204 "cann_ops_nn.mx_to_block_mx_quant importable." 

205 ) from exc 

206 return cann_mx_to_block_mx_quant 

207 

208 

209def validate_fake_w4a8_gmm_runtime(block_size: int) -> None: 

210 """Validate fake-W4A8 operators and the block-32 CANN converter.""" 

211 

212 _get_fake_w4a8_npu_ops() 

213 if block_size != MX_BLOCK_SIZE_32: 

214 raise ValueError("Fake W4A8 supports only block_size=32.") 

215 _get_cann_mx_to_block_mx_quant_op() 

216 

217 

218__all__ = [ 

219 "FakeW4A8CapabilityError", 

220 "FakeW4A8NpuOps", 

221 "validate_fake_w4a8_gmm_runtime", 

222]