Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / multicore / platform / mindspore / __init__.py: 0%

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""" 

16hyper_parallel.core.multicore.platform.mindspore 

17================================================== 

18MindSpore platform adapter for MoE-FFN multicore operators. 

19 

20Exposes ``moe_ffn_fwd`` and ``moe_ffn_bwd`` backed by the 

21``hyper_parallel_multicore_moe_ffn_ms`` pybind11 extension. 

22 

23.. note:: 

24 **Only MindSpore PyNative mode is supported.** 

25 Graph mode (``ms.GRAPH_MODE``) is not yet implemented — the operator 

26 YAML definitions set ``function: disable: True`` and 

27 ``dispatch: enable: False``, so the MindSpore compiler cannot trace 

28 or lower these ops into a static graph. Calling either function 

29 inside a ``@ms.jit``-decorated function or when 

30 ``ms.set_context(mode=ms.GRAPH_MODE)`` is active will raise a 

31 ``RuntimeError`` at call time. 

32 

33Vendor library resolution order (highest priority first): 

34 

35 1. ``CANN_VENDOR_FWD_LIBDIR`` / ``CANN_VENDOR_BWD_LIBDIR`` 

36 2. ``HP_MULTICORE_DIR`` (vendors root; fwd/bwd derived automatically) 

37 3. ``CANN_VENDOR_LIBDIR`` (legacy single-lib fallback) 

38 4. Auto-detect from ``prebuild/multicore_moe_ffn/vendors/`` 

39""" 

40import ctypes 

41import os 

42import re 

43import sys 

44 

45_ASCEND_HOME = os.environ.get("ASCEND_HOME_PATH", 

46 "/usr/local/Ascend/ascend-toolkit/latest") 

47 

48_CANN_VENDOR_FWD_LIBDIR = os.environ.get("CANN_VENDOR_FWD_LIBDIR", "") 

49_CANN_VENDOR_BWD_LIBDIR = os.environ.get("CANN_VENDOR_BWD_LIBDIR", "") 

50 

51if not (_CANN_VENDOR_FWD_LIBDIR or _CANN_VENDOR_BWD_LIBDIR): 

52 _HP_MULTICORE_DIR = os.environ.get("HP_MULTICORE_DIR", "") 

53 if _HP_MULTICORE_DIR: 

54 _CANN_VENDOR_FWD_LIBDIR = os.path.join( 

55 _HP_MULTICORE_DIR, "multicore_moe_ffn_nn", "op_api", "lib") 

56 _CANN_VENDOR_BWD_LIBDIR = os.path.join( 

57 _HP_MULTICORE_DIR, "multicore_moe_ffn_grad_nn", "op_api", "lib") 

58 elif os.environ.get("CANN_VENDOR_LIBDIR", ""): 

59 _CANN_VENDOR_FWD_LIBDIR = os.environ["CANN_VENDOR_LIBDIR"] 

60 _CANN_VENDOR_BWD_LIBDIR = os.environ["CANN_VENDOR_LIBDIR"] 

61 else: 

62 _PREBUILD_DIR = os.path.normpath( 

63 os.path.join(os.path.dirname(__file__), "../../prebuild/multicore_moe_ffn")) 

64 if os.path.isdir(_PREBUILD_DIR): 

65 _vendors = os.path.join(_PREBUILD_DIR, "vendors") 

66 _fwd = os.path.join(_vendors, "multicore_moe_ffn_nn", "op_api", "lib") 

67 _bwd = os.path.join(_vendors, "multicore_moe_ffn_grad_nn", "op_api", "lib") 

68 if os.path.isdir(_fwd): 

69 _CANN_VENDOR_FWD_LIBDIR = _fwd 

70 if os.path.isdir(_bwd): 

71 _CANN_VENDOR_BWD_LIBDIR = _bwd 

72 

73# --------------------------------------------------------------------------- 

74# Set ASCEND_CUSTOM_OPP_PATH at module load time — BEFORE any MindSpore import. 

75# 

76# GetOpApiFuncAddr (called by LAUNCH_ACLNN_FUNC at op invocation) scans 

77# ASCEND_CUSTOM_OPP_PATH to populate g_custom_lib_path, which may be cached 

78# as a static variable on first call. MindSpore's own initialization can 

79# trigger GetOpApiFuncAddr for standard CANN ops during `import mindspore`, 

80# which would cache an empty g_custom_lib_path if we set the env var too late 

81# (e.g., inside _get_ms_ops()). Setting it here ensures the correct paths are 

82# present before any such call. 

83# --------------------------------------------------------------------------- 

84for _libdir in (_CANN_VENDOR_FWD_LIBDIR, _CANN_VENDOR_BWD_LIBDIR): 

85 if _libdir: 

86 _vendor_root = re.sub(r'/op_api/lib/?$', '', _libdir) 

87 _cur_opp = os.environ.get('ASCEND_CUSTOM_OPP_PATH', '') 

88 if _vendor_root and _vendor_root not in _cur_opp: 

89 os.environ['ASCEND_CUSTOM_OPP_PATH'] = ( 

90 f"{_vendor_root}:{_cur_opp}" if _cur_opp else _vendor_root) 

91 

92# Pre-load CANN base libs with RTLD_GLOBAL. 

93for _ascendcl_path in [ 

94 os.path.join(_ASCEND_HOME, "lib64", "libascendcl.so"), 

95 os.path.join(_ASCEND_HOME, "fwkacllib", "lib64", "libascendcl.so"), 

96]: 

97 if os.path.exists(_ascendcl_path): 

98 ctypes.CDLL(_ascendcl_path, mode=ctypes.RTLD_GLOBAL) 

99 break 

100 

101_opapi_path = os.path.join(_ASCEND_HOME, "lib64", "libopapi.so") 

102if os.path.exists(_opapi_path): 

103 ctypes.CDLL(_opapi_path, mode=ctypes.RTLD_GLOBAL) 

104 

105# Pre-load both libcust_opapi.so with RTLD_GLOBAL so LAUNCH_ACLNN_FUNC can 

106# resolve aclnnMulticoreMoeFfn* / aclnnMulticoreMoeFfnGrad* from either package. 

107for _vendor_libdir in (_CANN_VENDOR_FWD_LIBDIR, _CANN_VENDOR_BWD_LIBDIR): 

108 if _vendor_libdir: 

109 _cust_opapi = os.path.join(_vendor_libdir, "libcust_opapi.so") 

110 if os.path.exists(_cust_opapi): 

111 ctypes.CDLL(_cust_opapi, mode=ctypes.RTLD_GLOBAL) 

112 

113# Lazily import the compiled MindSpore extension on first use to avoid a 

114# circular import triggered by MindSpore's custom-op scan at load time. 

115_ms_ops = None 

116 

117 

118def _check_pynative_mode(fn_name: str) -> None: 

119 """Raise RuntimeError if the current MindSpore execution mode is not PyNative. 

120 

121 Graph mode (GRAPH_MODE) is not yet implemented for MoE-FFN multicore ops. 

122 The operator YAML definitions disable function-level dispatch and static 

123 graph compilation, so these ops cannot be traced or lowered by the 

124 MindSpore compiler. 

125 """ 

126 import mindspore as ms # pylint: disable=import-outside-toplevel 

127 if ms.get_context("mode") != ms.PYNATIVE_MODE: 

128 raise RuntimeError( 

129 f"'{fn_name}' only supports MindSpore PyNative mode. " 

130 "Graph mode (ms.GRAPH_MODE) is not yet implemented for MoE-FFN " 

131 "multicore operators. " 

132 "Please call ms.set_context(mode=ms.PYNATIVE_MODE) before using " 

133 "this operator, and do not wrap it inside @ms.jit functions." 

134 ) 

135 

136 

137def _get_ms_ops(): 

138 """Lazily load and return the compiled MindSpore extension module.""" 

139 global _ms_ops 

140 if _ms_ops is None: 

141 build_lib = os.path.join(os.path.dirname(os.path.abspath(__file__)), "build", "lib") 

142 if build_lib not in sys.path: 

143 sys.path.insert(0, build_lib) 

144 import hyper_parallel_multicore_moe_ffn_ms as _m # noqa: E402 # pylint: disable=import-outside-toplevel 

145 _ms_ops = _m 

146 return _ms_ops 

147 

148 

149def moe_ffn_fwd( 

150 dispatch_target, dispatch_target_off, 

151 dispatch_src, dispatch_src_off, dispatch_size, 

152 up_proj_weight, up_proj_glist, 

153 up_proj_y, swiglu_out, 

154 down_proj_weight, down_proj_glist, down_proj_y, 

155 combine_target, combine_target_off, combine_src_off, combine_size, 

156 gmm_workspace, up_proj_tiling, swiglu_tiling, down_proj_tiling, 

157 runtime_config, all_event_counters, 

158 rank_id: int, ep: int, expert_num: int, 

159 hidden_size: int, seq_size: int, 

160): 

161 """MoE-FFN forward operator (MindSpore). 

162 

163 .. note:: Only PyNative mode is supported. Raises ``RuntimeError`` if 

164 called in Graph mode or inside an ``@ms.jit`` function. 

165 """ 

166 _check_pynative_mode("moe_ffn_fwd") 

167 _get_ms_ops().moe_ffn_fwd( 

168 dispatch_target, dispatch_target_off, 

169 dispatch_src, dispatch_src_off, dispatch_size, 

170 up_proj_weight, up_proj_glist, 

171 up_proj_y, swiglu_out, 

172 down_proj_weight, down_proj_glist, down_proj_y, 

173 combine_target, combine_target_off, combine_src_off, combine_size, 

174 gmm_workspace, up_proj_tiling, swiglu_tiling, down_proj_tiling, 

175 runtime_config, all_event_counters, 

176 rank_id, ep, expert_num, hidden_size, seq_size, 

177 ) 

178 

179 

180def moe_ffn_bwd( 

181 dispatch_target, dispatch_target_off, 

182 dy, dispatch_src_off, dispatch_size, 

183 hidden, hidden_dw, 

184 w2, act_grad_y, gate, grad_gate, w1, gate_dx, grad_x, 

185 combine_target_off, combine_src_off, combine_size, 

186 permute_out, gate_dw, group_list, 

187 act_grad_tiling, gate_grad_tiling, w2_grad_tiling, w1_grad_tiling, 

188 swiglu_grad_tiling, gmm_workspace, swiglu_grad_workspace, 

189 runtime_config, all_event_counters, 

190 rank_id: int, ep: int, expert_num: int, 

191 hidden_size: int, seq_size: int, 

192): 

193 """MoE-FFN backward operator (MindSpore). 

194 

195 .. note:: Only PyNative mode is supported. Raises ``RuntimeError`` if 

196 called in Graph mode or inside an ``@ms.jit`` function. 

197 """ 

198 _check_pynative_mode("moe_ffn_bwd") 

199 _get_ms_ops().moe_ffn_bwd( 

200 dispatch_target, dispatch_target_off, 

201 dy, dispatch_src_off, dispatch_size, 

202 hidden, hidden_dw, 

203 w2, act_grad_y, gate, grad_gate, w1, gate_dx, grad_x, 

204 combine_target_off, combine_src_off, combine_size, 

205 permute_out, gate_dw, group_list, 

206 act_grad_tiling, gate_grad_tiling, w2_grad_tiling, w1_grad_tiling, 

207 swiglu_grad_tiling, gmm_workspace, swiglu_grad_workspace, 

208 runtime_config, all_event_counters, 

209 rank_id, ep, expert_num, hidden_size, seq_size, 

210 ) 

211 

212 

213__all__ = ["moe_ffn_fwd", "moe_ffn_bwd"]