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

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 _format_family, 

23) 

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

25 GroupedLinear, 

26) 

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

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

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

30 FakeW4A8GroupedLinear, 

31) 

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

33 HiFloat8GroupedLinear, 

34) 

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

36 validate_hifloat8_gmm_runtime, 

37) 

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

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

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

41 validate_fake_w4a8_gmm_runtime, 

42) 

43 

44 

45def build_low_precision_strategy( 

46 config: Optional[LowPrecisionDtypeScheme | LowPrecisionConfig], 

47 *, 

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

49) -> GroupedLinear: 

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

51 

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

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

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

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

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

57 enforce their real-weight tile-alignment contract. 

58 """ 

59 

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

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

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

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

64 # as ``LowPrecisionDtypeScheme`` instances. 

65 if isinstance(config, LowPrecisionConfig): 

66 config = config.resolve_dtype_scheme() 

67 

68 if config is None: 

69 weight_format = act_format = "mxfp8" 

70 else: 

71 weight_format = config.weight_format 

72 act_format = config.act_format 

73 

74 family = _format_family(weight_format) 

75 

76 if family == "mxfp": 

77 if config is not None and config.is_fake_quantize: 

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

79 tile_size = config.block_size 

80 validate_fake_w4a8_gmm_runtime() 

81 strategy = FakeW4A8GroupedLinear(tile_size) 

82 _validate_tile_alignment("fake_w4a8", tile_size, tile_shapes) 

83 return strategy 

84 raise NotImplementedError( 

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

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

87 ) 

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

89 validate_npu_gmm_runtime() 

90 strategy = MXFP8GroupedLinear() 

91 _validate_tile_alignment("w8a8", 32, tile_shapes) 

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

93 tile_size = config.block_size if config is not None else 32 

94 validate_w4a8_gmm_runtime() 

95 strategy = W4A8GroupedLinear(tile_size) 

96 _validate_tile_alignment("w4a8", tile_size, tile_shapes) 

97 else: 

98 raise NotImplementedError( 

99 "mxfp combination weight_format=" 

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

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

102 ) 

103 elif family == "hifloat": 

104 if config is not None and config.is_fake_quantize: 

105 raise NotImplementedError( 

106 "fake HiFloat QAT is not implemented; " 

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

108 ) 

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

110 validate_hifloat8_gmm_runtime() 

111 return HiFloat8GroupedLinear() 

112 raise NotImplementedError( 

113 "hifloat combination weight_format=" 

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

115 "supported: hif8/hif8." 

116 ) 

117 else: 

118 raise NotImplementedError( 

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

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

121 ) 

122 return strategy 

123 

124 

125def _validate_tile_alignment( 

126 strategy_name: str, 

127 tile_size: int, 

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

129) -> None: 

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

131 

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

133 the block size (32 for MXFP8, ``config.block_size`` 32/128 for W4A8), so a 

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

135 block for the NPU operator. 

136 """ 

137 

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

139 raise ValueError( 

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

141 f"{tile_shapes}." 

142 ) 

143 

144 

145__all__ = ["build_low_precision_strategy"]