Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / components / quantization / functional / strategy_factory.py: 93%

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"""Build concrete grouped-linear strategies from resolved policy data.""" 

16 

17from typing import Optional 

18 

19from hyper_parallel.components.quantization.config import ( 

20 LowPrecisionConfig, 

21 LowPrecisionDtypeScheme, 

22 MX_BLOCK_SIZE_32, 

23 MX_SCALE_TILE, 

24 _format_family, 

25) 

26from hyper_parallel.components.quantization.functional.base_gmm_func import ( 

27 GroupedLinear, 

28) 

29from hyper_parallel.components.quantization.functional.mxfp8_gmm_func import MXFP8GroupedLinear 

30from hyper_parallel.components.quantization.functional.w4a8_gmm_func import W4A8GroupedLinear 

31from hyper_parallel.components.quantization.functional.fake_w4a8_gmm_func import ( 

32 FakeW4A8GroupedLinear, 

33) 

34from hyper_parallel.components.quantization.functional.hifloat8_gmm_func import ( 

35 HiFloat8GroupedLinear, 

36) 

37from hyper_parallel.components.quantization.ops.npu_hifloat8 import ( 

38 validate_hifloat8_gmm_runtime, 

39) 

40from hyper_parallel.components.quantization.ops.npu_mxfp8 import validate_npu_gmm_runtime 

41from hyper_parallel.components.quantization.ops.npu_w4a8 import validate_w4a8_gmm_runtime 

42from hyper_parallel.components.quantization.ops.npu_fake_w4a8 import ( 

43 validate_fake_w4a8_gmm_runtime, 

44) 

45 

46 

47def build_low_precision_strategy( 

48 config: Optional[LowPrecisionDtypeScheme | LowPrecisionConfig], 

49 *, 

50 tile_shapes: tuple[tuple[int, ...], tuple[int, ...]], 

51) -> GroupedLinear: 

52 """Pick the concrete grouped-linear strategy for one replaced module. 

53 

54 ``config`` is the already-validated dtype scheme (``None`` = default native 

55 ``mxfp8``/``mxfp8``, i.e. W8A8). Discriminators run in priority order: the 

56 explicit fake-QAT flag selects the independent fake W4A8 strategy, then 

57 format family and the weight/activation profile select MXFP8, HiFloat8, 

58 or native W4A8 with the corresponding runtime probe. MX strategies also 

59 enforce their real-weight tile-alignment contract. 

60 """ 

61 

62 # Replacement factories historically receive the global ``LowPrecisionConfig`` 

63 # directly. Resolve its legacy/default catalog entry at this boundary so 

64 # all downstream strategy selection operates on one typed scheme. Named 

65 # schemes are resolved by the plan resolver before replacement and arrive 

66 # as ``LowPrecisionDtypeScheme`` instances. 

67 if isinstance(config, LowPrecisionConfig): 

68 config = config.resolve_dtype_scheme() 

69 

70 if config is None: 

71 weight_format = act_format = "mxfp8" 

72 else: 

73 weight_format = config.weight_format 

74 act_format = config.act_format 

75 

76 family = _format_family(weight_format) 

77 

78 if family == "mxfp": 

79 if config is not None and config.is_fake_quantize: 

80 if weight_format == "mxfp4" and act_format == "mxfp8": 

81 tile_size = config.block_size 

82 _validate_tile_alignment("fake_w4a8", tile_size, tile_shapes) 

83 validate_fake_w4a8_gmm_runtime(tile_size) 

84 return FakeW4A8GroupedLinear(tile_size) 

85 raise NotImplementedError( 

86 "fake QAT currently supports only mxfp4/mxfp8 (fake W4A8); " 

87 f"got weight_format={weight_format!r}/act_format={act_format!r}." 

88 ) 

89 if weight_format == "mxfp8" and act_format == "mxfp8": 

90 validate_npu_gmm_runtime() 

91 strategy = MXFP8GroupedLinear() 

92 _validate_tile_alignment("w8a8", MX_BLOCK_SIZE_32, tile_shapes) 

93 elif weight_format == "mxfp4" and act_format == "mxfp8": 

94 tile_size = config.block_size if config is not None else MX_BLOCK_SIZE_32 

95 validate_w4a8_gmm_runtime() 

96 strategy = W4A8GroupedLinear(tile_size) 

97 _validate_tile_alignment("w4a8", tile_size, tile_shapes) 

98 else: 

99 raise NotImplementedError( 

100 "mxfp combination weight_format=" 

101 f"{weight_format!r}/act_format={act_format!r} is not " 

102 "implemented; supported: mxfp8/mxfp8 (w8a8), mxfp4/mxfp8 (w4a8)." 

103 ) 

104 elif family == "hifloat": 

105 if config is not None and config.is_fake_quantize: 

106 raise NotImplementedError( 

107 "fake HiFloat QAT is not implemented; " 

108 f"got weight_format={weight_format!r}/act_format={act_format!r}." 

109 ) 

110 if weight_format == "hif8" and act_format == "hif8": 

111 validate_hifloat8_gmm_runtime() 

112 return HiFloat8GroupedLinear() 

113 raise NotImplementedError( 

114 "hifloat combination weight_format=" 

115 f"{weight_format!r}/act_format={act_format!r} is not implemented; " 

116 "supported: hif8/hif8." 

117 ) 

118 else: 

119 raise NotImplementedError( 

120 "grouped-linear format family is not implemented: " 

121 f"weight_format={weight_format!r}/act_format={act_format!r}." 

122 ) 

123 return strategy 

124 

125 

126def _validate_tile_alignment( 

127 strategy_name: str, 

128 tile_size: int, 

129 tile_shapes: tuple[tuple[int, ...], tuple[int, ...]], 

130) -> None: 

131 """Require both non-expert weight dims to be whole blocks of ``tile_size``. 

132 

133 Live ``[E, O, K]`` weights must have their ``-2``/``-1`` dims divisible by 

134 the block size (``MX_BLOCK_SIZE_32`` for both MXFP8 and W4A8), so a 

135 misaligned model fails at replacement time instead of leaving a partial 

136 block for the NPU operator. Block-32 fake W4A8 also needs one matrix axis 

137 divisible by 64 so the CANN converter can use it as its row axis. 

138 """ 

139 

140 if any(dimension % tile_size for shape in tile_shapes for dimension in shape[-2:]): 

141 raise ValueError( 

142 f"{strategy_name!r} is not aligned for block_size={tile_size}: " 

143 f"{tile_shapes}." 

144 ) 

145 if strategy_name == "fake_w4a8": 

146 unsupported = tuple( 

147 shape 

148 for shape in tile_shapes 

149 if all(dimension % MX_SCALE_TILE for dimension in shape[-2:]) 

150 ) 

151 if unsupported: 

152 raise ValueError( 

153 "'fake_w4a8' block_size=32 requires at least one matrix " 

154 "dimension divisible by 64 for CANN's packed-MXFP4 " 

155 f"to block-MXFP8 conversion; unsupported shapes: {unsupported}." 

156 ) 

157 

158 

159__all__ = ["build_low_precision_strategy"]