Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / components / modules / mxfp8_grouped_swiglu.py: 85%

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 local computation with optional fusion and the GroupedExperts contract.""" 

16 

17from __future__ import annotations 

18 

19from collections.abc import Mapping 

20from typing import Any 

21 

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

23from torch import nn # pylint: disable=forbidden-backend-import 

24from torch.nn import functional as F # pylint: disable=forbidden-backend-import 

25from transformers.activations import ACT2FN 

26 

27from hyper_parallel.components.modules.grouped_experts import GroupedExperts 

28from hyper_parallel.components.quantization.functional.mxfp8_fused_grouped_swiglu import ( 

29 fused_grouped_swiglu, 

30) 

31from hyper_parallel.components.quantization.functional.mxfp8_projection import mxfp8_projection 

32from hyper_parallel.components.quantization.ops.npu_mxfp8 import ( 

33 validate_npu_gmm_runtime, validate_npu_swiglu_runtime, 

34) 

35from hyper_parallel.components.quantization.quantizers.mxfp8 import MXFP8Quantizer 

36from hyper_parallel.models.replacement import module_replacement 

37 

38 

39@module_replacement 

40class MXFP8GroupedSwiGLU(GroupedExperts): 

41 """Retain routing and checkpoint conversion; replace only local expert math. 

42 

43 Persistent weights remain [E,H,2I] and [E,I,H]. Counts name tokens per 

44 local expert, including zero-token experts. Output retains expert-major 

45 order. Scores are accepted by the inherited protocol; outer routing applies 

46 them exactly once after GMM2, with its original dtype policy. 

47 

48 Note: 

49 EP recipes must explicitly set ``use_grouped_gemm: true`` to call this 

50 module's MXFP8 computation. EP binding rejects a missing or false option 

51 because the eager path does not support this module's weight layout. 

52 """ 

53 

54 requires_ep_grouped_gemm = True 

55 

56 def __init__(self, *, module: nn.Module, module_fqn: str = "", 

57 context: Mapping[str, Any] | None = None) -> None: 

58 """Validate the supported structure and initialize a non-fused variant.""" 

59 context = context or {} 

60 self.module_fqn = module_fqn 

61 prefix = f"{module_fqn or '<root>'}: MXFP8GroupedSwiGLU" 

62 if any(context.get(axis) for axis in ("tp", "cp", "pp")): 

63 raise NotImplementedError(f"{prefix} requires TP=CP=PP=1") 

64 try: 

65 super().__init__(module=module, module_fqn=module_fqn, context=context) 

66 except (ValueError, TypeError, KeyError) as exc: 

67 raise ValueError(f"{prefix}: {exc}") from exc 

68 self._validate_source_contract(module, prefix) 

69 try: 

70 validate_npu_gmm_runtime() 

71 except RuntimeError as exc: 

72 raise RuntimeError(f"{prefix}: {exc}") from exc 

73 self.quantizer = MXFP8Quantizer() 

74 

75 def _validate_source_contract(self, module: nn.Module, prefix: str) -> None: 

76 """Check the initial layout and activation before any device computation.""" 

77 hidden_act = (getattr(self.config, "hidden_act", None) 

78 or getattr(self.config, "hidden_activation", None) 

79 or getattr(self.config, "mlp_hidden_act", "silu")) 

80 source_activation = getattr(module, "act_fn", None) 

81 if (not self.has_gate or self.add_bias or hidden_act not in ("silu", "swiglu") 

82 or (source_activation is not None 

83 and not isinstance(source_activation, (nn.SiLU, type(ACT2FN["silu"]))))): 

84 raise ValueError(f"{prefix} requires standard bias-free SwiGLU") 

85 if self.use_2d_experts or not self.router_gating_in_fp32: 

86 raise ValueError(f"{prefix} requires 3D packed weights and routing weights after GMM2") 

87 if self.hidden_size % 32 or self.intermediate_size % 32: 

88 raise ValueError(f"{prefix} requires H and I divisible by 32") 

89 

90 def _grouped_gemm_expert_forward(self, gate_up_proj, down_proj, permuted, 

91 tokens_per_expert, permuted_probs): 

92 """Validate expert groups and dispatch local math without applying scores.""" 

93 del permuted_probs 

94 counts = tokens_per_expert.to(device=permuted.device) 

95 if counts.ndim != 1 or counts.shape[0] != gate_up_proj.shape[0]: 

96 raise ValueError(f"{self.module_fqn}: counts must contain one count per local expert") 

97 if counts.dtype not in (torch.int32, torch.int64): 

98 raise ValueError(f"{self.module_fqn}: counts must be integers") 

99 # Device counts come from routing; avoid copying them to the host per GMM. 

100 # CPU callers can validate values before handing counts to device kernels. 

101 if tokens_per_expert.device.type == "cpu": 

102 if bool((tokens_per_expert < 0).any()) or int(tokens_per_expert.sum()) != permuted.shape[0]: 

103 raise ValueError(f"{self.module_fqn}: counts must be non-negative and sum to token count") 

104 return self._local_compute(permuted, gate_up_proj, down_proj, counts) 

105 

106 def _local_compute(self, permuted, gate_up_proj, down_proj, counts): 

107 """Run the non-fused local chain after common protocol validation.""" 

108 projected = mxfp8_projection(permuted, gate_up_proj, counts, self.quantizer) 

109 gate, up = projected.chunk(2, dim=-1) 

110 activated = F.silu(gate) * up 

111 return mxfp8_projection(activated, down_proj, counts, self.quantizer) 

112 

113 

114@module_replacement 

115class FusedMXFP8GroupedSwiGLU(MXFP8GroupedSwiGLU): 

116 """Use fused activation/quantization while retaining the same family contract.""" 

117 

118 def __init__(self, *, module: nn.Module, module_fqn: str = "", 

119 context: Mapping[str, Any] | None = None) -> None: 

120 """Validate both fused kernels during construction, with module diagnostics.""" 

121 super().__init__(module=module, module_fqn=module_fqn, context=context) 

122 try: 

123 validate_npu_swiglu_runtime() 

124 except RuntimeError as exc: 

125 raise RuntimeError(f"{module_fqn or '<root>'}: {exc}") from exc 

126 

127 def _local_compute(self, permuted, gate_up_proj, down_proj, counts): 

128 """Use one autograd boundary without passing quantized carriers between nodes.""" 

129 return fused_grouped_swiglu(permuted, gate_up_proj, down_proj, counts, self.quantizer)