Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / shard / ops / parallel_mhc_pre_sinkhorn.py: 73%

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"""Distributed implementation for npu_mhc_pre_sinkhorn operator.""" 

16from typing import Tuple, Dict, Any 

17 

18from hyper_parallel.core.dtensor.layout import Layout 

19from hyper_parallel.platform import get_platform 

20from hyper_parallel.platform.platform import PlatformType 

21from .parallel_ops import DistributedOp 

22 

23platform = get_platform() 

24 

25_HC_MULT_DEFAULT = 4 

26_NUM_ITERS_DEFAULT = 20 

27_HC_EPS_DEFAULT = 1e-6 

28_NORM_EPS_DEFAULT = 1e-6 

29_MHC_PRE_CLAMP_ARG_NAMES = ( 

30 "x", "phi", "alpha", "bias", "hc_mult", "num_iters", 

31 "hc_eps", "norm_eps", "out_flag", "clamp_min", "clamp_max", 

32) 

33_MHC_PRE_CLAMP_DEFAULTS = { 

34 "hc_mult": _HC_MULT_DEFAULT, 

35 "num_iters": _NUM_ITERS_DEFAULT, 

36 "hc_eps": _HC_EPS_DEFAULT, 

37 "norm_eps": _NORM_EPS_DEFAULT, 

38 "out_flag": True, 

39 "clamp_min": 0.0, 

40 "clamp_max": 0.0, 

41} 

42 

43 

44def _normalize_mhc_pre_sinkhorn_args( 

45 x, 

46 phi, 

47 alpha, 

48 bias, 

49 hc_mult=_HC_MULT_DEFAULT, 

50 num_iters=_NUM_ITERS_DEFAULT, 

51 hc_eps=_HC_EPS_DEFAULT, 

52 norm_eps=_NORM_EPS_DEFAULT, 

53 out_flag=True): 

54 """Normalize positional and keyword arguments into a canonical positional tuple. 

55 

56 Args: 

57 x: Input tensor [B,S,N,C] or [T,N,C]. 

58 phi: mHC parameter matrix [N*N+2*N, N*C]. 

59 alpha: mHC scaling parameters [3]. 

60 bias: mHC bias parameters [N*N+2*N]. 

61 hc_mult: HC dimension size (currently only 4 supported). 

62 num_iters: Sinkhorn iteration count. 

63 hc_eps: H_pre sigmoid eps parameter. 

64 norm_eps: RmsNorm eps parameter. 

65 out_flag: Whether to output intermediate gradients. 

66 

67 Returns: 

68 tuple: (positional_args_tuple, empty_kwargs_dict) 

69 """ 

70 return ( 

71 x, phi, alpha, bias, 

72 hc_mult, num_iters, hc_eps, norm_eps, out_flag, 

73 ), {} 

74 

75 

76def _normalize_mhc_pre_clamp_sinkhorn_args(*args, **kwargs): 

77 """Normalize npu_mhc_pre_clamp_sinkhorn arguments.""" 

78 values = dict(_MHC_PRE_CLAMP_DEFAULTS) 

79 if len(args) > len(_MHC_PRE_CLAMP_ARG_NAMES): 

80 raise TypeError( 

81 f"npu_mhc_pre_clamp_sinkhorn expected at most {len(_MHC_PRE_CLAMP_ARG_NAMES)} arguments" 

82 ) 

83 for name, value in zip(_MHC_PRE_CLAMP_ARG_NAMES, args): 

84 values[name] = value 

85 for name, value in kwargs.items(): 

86 if name not in _MHC_PRE_CLAMP_ARG_NAMES: 

87 raise TypeError(f"npu_mhc_pre_clamp_sinkhorn got an unexpected keyword argument '{name}'") 

88 if name in _MHC_PRE_CLAMP_ARG_NAMES[:len(args)]: 

89 raise TypeError(f"npu_mhc_pre_clamp_sinkhorn got multiple values for argument '{name}'") 

90 values[name] = value 

91 missing = [name for name in _MHC_PRE_CLAMP_ARG_NAMES[:4] if name not in values] 

92 if missing: 

93 raise TypeError(f"npu_mhc_pre_clamp_sinkhorn missing required arguments: {missing}") 

94 return tuple(values[name] for name in _MHC_PRE_CLAMP_ARG_NAMES), {} 

95 

96 

97# Validation rules table for npu_mhc_pre_sinkhorn 

98# Key: tensor_map length (format identifier) 

99# Value: validation rules for that format 

100_MHC_PRE_SINKHORN_VALIDATION_RULES: Dict[int, Dict[str, Any]] = { 

101 4: { 

102 "op_name": "npu_mhc_pre_sinkhorn", 

103 "forbidden_dims": {2: "N", 3: "C"}, 

104 "phi_forbidden_dims": {0: "dim0", 1: "dim1"}, 

105 "alpha_forbidden_dims": {0: "dim0"}, 

106 "bias_forbidden_dims": {0: "dim0"}, 

107 }, 

108 3: { 

109 "op_name": "npu_mhc_pre_sinkhorn", 

110 "forbidden_dims": {1: "N", 2: "C"}, 

111 "phi_forbidden_dims": {0: "dim0", 1: "dim1"}, 

112 "alpha_forbidden_dims": {0: "dim0"}, 

113 "bias_forbidden_dims": {0: "dim0"}, 

114 }, 

115} 

116 

117 

118def _create_output_layout(mesh: Any, tensor_map: tuple) -> Layout: 

119 """Create an output layout with placements derived from ``tensor_map``.""" 

120 output_layout = Layout.from_device_mesh(mesh) 

121 output_layout.set_tensor_map(tensor_map) 

122 output_layout.tensor_map_to_placement() 

123 return output_layout 

124 

125 

126def _validate_tensor_map_dims( 

127 tensor_map: tuple, 

128 op_name: str, 

129 forbidden_dims: Dict[int, str], 

130) -> None: 

131 """Check that specified dimensions are not sharded (replicated). 

132 

133 Args: 

134 tensor_map: The tensor_map to check. 

135 op_name: Operator name for error message. 

136 forbidden_dims: Dict mapping dim index to dim name. 

137 

138 Raises: 

139 ValueError: If any forbidden dimension is sharded. 

140 """ 

141 for dim_idx, dim_name in forbidden_dims.items(): 

142 dim_value = tensor_map[dim_idx] 

143 if dim_value != -1: 

144 raise ValueError( 

145 f"For {op_name}, {dim_name} dimension (dim {dim_idx}) of x " 

146 f"should be replicated, but got {dim_value}" 

147 ) 

148 

149 

150def _validate_input_layouts_mhc_pre_sinkhorn( 

151 x_layout: Layout, 

152 phi_layout: Layout, 

153 alpha_layout: Layout, 

154 bias_layout: Layout, 

155) -> None: 

156 """Validate input layouts for npu_mhc_pre_sinkhorn operator.""" 

157 x_tm = x_layout.tensor_map 

158 x_tm_len = len(x_tm) 

159 

160 rules = _MHC_PRE_SINKHORN_VALIDATION_RULES.get(x_tm_len) 

161 if rules is None: 

162 raise ValueError( 

163 f"For npu_mhc_pre_sinkhorn, tensor_map length should be 4 or 3, but got {x_tm_len}" 

164 ) 

165 

166 _validate_tensor_map_dims(x_tm, rules["op_name"], rules["forbidden_dims"]) 

167 _validate_tensor_map_dims(phi_layout.tensor_map, rules["op_name"], rules["phi_forbidden_dims"]) 

168 _validate_tensor_map_dims(alpha_layout.tensor_map, rules["op_name"], rules["alpha_forbidden_dims"]) 

169 _validate_tensor_map_dims(bias_layout.tensor_map, rules["op_name"], rules["bias_forbidden_dims"]) 

170 

171 

172class NpuMhcPreSinkhornDistributedOp(DistributedOp): 

173 """DistributedOp for npu_mhc_pre_sinkhorn operator. 

174 

175 Implements layout inference for the MHC pre-processing with Sinkhorn operation. 

176 Outputs 8 tensors: hin, h_post, h_res, h_pre, hc_before_norm, inv_rms, sum_out, norm_out. 

177 """ 

178 

179 def preprocess(self, args: tuple, kwargs: dict) -> tuple: 

180 """Unwrap DTensor inputs and cache their layouts for inference.""" 

181 norm_args, _ = _normalize_mhc_pre_sinkhorn_args(*args, **kwargs) 

182 dtensor_x = norm_args[0] 

183 dtensor_phi = norm_args[1] 

184 dtensor_alpha = norm_args[2] 

185 dtensor_bias = norm_args[3] 

186 

187 if platform.platform_type == PlatformType.MINDSPORE: 

188 local_args = ( 

189 dtensor_x.to_local(), 

190 dtensor_phi.to_local(), 

191 dtensor_alpha.to_local(), 

192 dtensor_bias.to_local(), 

193 norm_args[4], 

194 norm_args[5], 

195 norm_args[6], 

196 norm_args[7], 

197 norm_args[8], 

198 ) 

199 local_kwargs = {} 

200 else: 

201 local_args = ( 

202 dtensor_x.to_local(), 

203 dtensor_phi.to_local(), 

204 dtensor_alpha.to_local(), 

205 dtensor_bias.to_local(), 

206 ) 

207 local_kwargs = { 

208 'hc_mult': norm_args[4], 

209 'num_iters': norm_args[5], 

210 'hc_eps': norm_args[6], 

211 'norm_eps': norm_args[7], 

212 'out_flag': norm_args[8], 

213 } 

214 

215 cache_values = [ 

216 dtensor_x.layout, 

217 dtensor_phi.layout, 

218 dtensor_alpha.layout, 

219 dtensor_bias.layout, 

220 ] 

221 return local_args, local_kwargs, cache_values 

222 

223 def infer_layout(self, cache_values: list) -> Tuple[tuple, None]: 

224 """Validate input layouts and infer the eight output layouts.""" 

225 x_layout, phi_layout, alpha_layout, bias_layout = cache_values 

226 

227 self._check_partial_inputs([x_layout, phi_layout, alpha_layout, bias_layout]) 

228 

229 _validate_input_layouts_mhc_pre_sinkhorn( 

230 x_layout, phi_layout, alpha_layout, bias_layout 

231 ) 

232 

233 out_layouts = self.infer_output_layouts(x_layout) 

234 return out_layouts, None 

235 

236 @staticmethod 

237 def infer_output_layouts( 

238 x_layout: Layout, 

239 ) -> Tuple[Layout, Layout, Layout, Layout, Layout, Layout, Layout, Layout]: 

240 """Infer per-output layouts from the input x layout. 

241 

242 The input x is either 4-D (B, S, N, C) or 3-D (T, N, C), but the 

243 8 kernel outputs each have a different rank (2-D through 5-D). 

244 Blindly copying the input tensor_map to every output causes 

245 ``get_global_shape`` to raise because ``len(slice_shape) != 

246 len(tensor_map)``. 

247 

248 Each output preserves only the leading dims that remain from the 

249 input, with new kernel-internal axes set to replicated (``-1``). 

250 """ 

251 x_tm = x_layout.tensor_map 

252 x_tm_len = len(x_tm) 

253 mesh = x_layout.mesh 

254 

255 if x_tm_len == 4: 

256 # (B, S, N, C) → outputs preserve B (dim 0) and S (dim 1) 

257 b_map, s_map, _, c_map = x_tm 

258 

259 # 3-D: (B, S, …) 

260 tm_h_in = (b_map, s_map, c_map) # h_in: (B, S, C) 

261 tm_3d = (b_map, s_map, -1) # h_post/h_res/h_pre/hc_before_norm: (B, S, …) 

262 # 4-D: (2*iters, B, S, N) # sum_out 

263 tm_sum = (-1, b_map, s_map, -1) 

264 # 5-D: (2*iters, B, S, N, N) # norm_out 

265 tm_norm = (-1, b_map, s_map, -1, -1) 

266 

267 return ( 

268 _create_output_layout(mesh, tm_h_in), # h_in 

269 _create_output_layout(mesh, tm_3d), # h_post 

270 _create_output_layout(mesh, tm_3d), # h_res 

271 _create_output_layout(mesh, tm_3d), # h_pre 

272 _create_output_layout(mesh, tm_3d), # hc_before_norm 

273 _create_output_layout(mesh, tm_3d), # inv_rms 

274 _create_output_layout(mesh, tm_sum), # sum_out 

275 _create_output_layout(mesh, tm_norm), # norm_out 

276 ) 

277 

278 if x_tm_len == 3: 

279 # (T, N, C) → only T (dim 0) is batch-like 

280 t_map, _, c_map = x_tm 

281 

282 # 2-D: (T, …) 

283 tm_h_in = (t_map, c_map) # h_in: (T, C) 

284 tm_2d = (t_map, -1) # h_post/h_res/h_pre/hc_before_norm: (T, …) 

285 # 3-D: (2*iters, T, N) # sum_out 

286 tm_sum = (-1, t_map, -1) 

287 # 4-D: (2*iters, T, N, N) # norm_out 

288 tm_norm = (-1, t_map, -1, -1) 

289 

290 return ( 

291 _create_output_layout(mesh, tm_h_in), # h_in 

292 _create_output_layout(mesh, tm_2d), # h_post 

293 _create_output_layout(mesh, tm_2d), # h_res 

294 _create_output_layout(mesh, tm_2d), # h_pre 

295 _create_output_layout(mesh, tm_2d), # hc_before_norm 

296 _create_output_layout(mesh, tm_2d), # inv_rms 

297 _create_output_layout(mesh, tm_sum), # sum_out 

298 _create_output_layout(mesh, tm_norm), # norm_out 

299 ) 

300 

301 raise ValueError( 

302 f"For npu_mhc_pre_sinkhorn, tensor_map length should be 4 or 3, " 

303 f"but got {x_tm_len}." 

304 ) 

305 

306 

307class NpuMhcPreClampSinkhornDistributedOp(DistributedOp): 

308 """DistributedOp for npu_mhc_pre_clamp_sinkhorn operator. 

309 

310 The clamp variant follows the same input layout rules as npu_mhc_pre_sinkhorn 

311 and emits one additional h_res_logits output. 

312 """ 

313 

314 def preprocess(self, args: tuple, kwargs: dict) -> tuple: 

315 """Unwrap clamp operator inputs and cache their layouts for inference.""" 

316 norm_args, _ = _normalize_mhc_pre_clamp_sinkhorn_args(*args, **kwargs) 

317 dtensor_x = norm_args[0] 

318 dtensor_phi = norm_args[1] 

319 dtensor_alpha = norm_args[2] 

320 dtensor_bias = norm_args[3] 

321 

322 if platform.platform_type == PlatformType.MINDSPORE: 

323 local_args = ( 

324 dtensor_x.to_local(), 

325 dtensor_phi.to_local(), 

326 dtensor_alpha.to_local(), 

327 dtensor_bias.to_local(), 

328 norm_args[4], 

329 norm_args[5], 

330 norm_args[6], 

331 norm_args[7], 

332 norm_args[8], 

333 norm_args[9], 

334 norm_args[10], 

335 ) 

336 local_kwargs = {} 

337 else: 

338 local_args = ( 

339 dtensor_x.to_local(), 

340 dtensor_phi.to_local(), 

341 dtensor_alpha.to_local(), 

342 dtensor_bias.to_local(), 

343 ) 

344 local_kwargs = { 

345 'hc_mult': norm_args[4], 

346 'num_iters': norm_args[5], 

347 'hc_eps': norm_args[6], 

348 'norm_eps': norm_args[7], 

349 'out_flag': norm_args[8], 

350 'clamp_min': norm_args[9], 

351 'clamp_max': norm_args[10], 

352 } 

353 

354 cache_values = [ 

355 dtensor_x.layout, 

356 dtensor_phi.layout, 

357 dtensor_alpha.layout, 

358 dtensor_bias.layout, 

359 ] 

360 return local_args, local_kwargs, cache_values 

361 

362 def infer_layout(self, cache_values: list) -> Tuple[tuple, None]: 

363 """Validate input layouts and infer the nine clamp output layouts.""" 

364 x_layout, phi_layout, alpha_layout, bias_layout = cache_values 

365 

366 self._check_partial_inputs([x_layout, phi_layout, alpha_layout, bias_layout]) 

367 _validate_input_layouts_mhc_pre_sinkhorn( 

368 x_layout, phi_layout, alpha_layout, bias_layout 

369 ) 

370 

371 # First 8 outputs use the same layout logic as the non-clamp variant. 

372 out_layouts = NpuMhcPreSinkhornDistributedOp.infer_output_layouts(x_layout) 

373 

374 # 9th output h_res_logits has the same shape as norm_out (5-D for 

375 # BSND, 4-D for TND). 

376 x_tm = x_layout.tensor_map 

377 x_tm_len = len(x_tm) 

378 mesh = x_layout.mesh 

379 

380 if x_tm_len == 4: 

381 b_map, s_map = x_tm[0], x_tm[1] 

382 tm_logits = (-1, b_map, s_map, -1, -1) 

383 elif x_tm_len == 3: 

384 t_map = x_tm[0] 

385 tm_logits = (-1, t_map, -1, -1) 

386 else: 

387 raise ValueError( 

388 f"For npu_mhc_pre_clamp_sinkhorn, tensor_map length should be " 

389 f"4 or 3, but got {x_tm_len}." 

390 ) 

391 

392 logits_layout = _create_output_layout(mesh, tm_logits) 

393 

394 return out_layouts + (logits_layout,), None