Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_multicore / modules / mega_moe / diagnostics.py: 81%

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"""MegaMoe-owned expert semantics shared by its diagnostic tools.""" 

16 

17__all__ = [] 

18 

19from dataclasses import dataclass 

20from typing import Any 

21 

22from hyper_multicore.scheduler.config import ( 

23 TaskDescC, 

24 TaskSplitValue, 

25) 

26from hyper_multicore.scheduler.graph import OperatorNode, OpType 

27 

28 

29MEGA_MOE_OWNER_LABEL = "Expert" 

30 

31 

32@dataclass(frozen=True) 

33class _MegaMoeDiagnosticContext: 

34 """Host values needed to resolve expert ownership after task construction.""" 

35 

36 topology: TaskSplitValue 

37 num_cube_cores: int 

38 

39 

40def _require_diagnostic_context(context: Any) -> _MegaMoeDiagnosticContext: 

41 if not isinstance(context, _MegaMoeDiagnosticContext): 

42 raise TypeError( 

43 f"MegaMoe diagnostic context must be _MegaMoeDiagnosticContext, got {type(context).__name__}" 

44 ) 

45 return context 

46 

47 

48def _resolve_expert_owner( 

49 operator: OperatorNode, 

50 task_desc: TaskDescC, 

51 raw_context: Any, 

52) -> int: 

53 """Resolve the global expert ID from a graph operator and its concrete task.""" 

54 context = _require_diagnostic_context(raw_context) 

55 topology = context.topology 

56 rank_owner_base = topology.rank_id * topology.single_rank_expert_num 

57 

58 if operator.op_type == OpType.GMM: 

59 return rank_owner_base + task_desc.task_index // context.num_cube_cores 

60 

61 if operator.op_type in (OpType.SWIGLU, OpType.SWIGLU_GRAD): 

62 if task_desc.task_split_value == 0: 

63 raise ValueError("SwiGLU diagnostic owner requires a non-zero task_split_value") 

64 num_triggers = topology.per_expert_seq // task_desc.task_split_value 

65 if num_triggers == 0: 

66 raise ValueError("SwiGLU diagnostic owner resolved a zero task count per expert") 

67 return rank_owner_base + task_desc.task_index // num_triggers 

68 

69 if operator.op_type == OpType.ALLTOALL: 

70 task_count_per_expert = task_desc.task_split_num // topology.all_expert_num 

71 if task_count_per_expert == 0: 

72 raise ValueError("AllToAll diagnostic owner resolved a zero task count per expert") 

73 expert_index = task_desc.task_index // task_count_per_expert 

74 if operator.name == "dispatch": 

75 if topology.dispatch_mode == "pull": 

76 return rank_owner_base + expert_index % topology.single_rank_expert_num 

77 return expert_index 

78 if operator.name == "combine": 

79 return rank_owner_base + expert_index % topology.single_rank_expert_num 

80 raise ValueError(f"MegaMoe AllToAll operator has no expert owner rule: {operator.name!r}") 

81 

82 raise ValueError(f"MegaMoe operator type has no expert owner rule: {operator.op_type}") 

83 

84 

85def _mega_moe_owner_options( 

86 topology: TaskSplitValue, 

87 *, 

88 num_cube_cores: int, 

89) -> dict[str, Any]: 

90 """Supply only MegaMoe's owner semantics to the shared diagnostics adapter.""" 

91 if not isinstance(topology, TaskSplitValue): 

92 raise TypeError(f"topology must be TaskSplitValue, got {type(topology).__name__}") 

93 if not isinstance(num_cube_cores, int) or isinstance(num_cube_cores, bool) or num_cube_cores <= 0: 

94 raise ValueError(f"num_cube_cores must be a positive integer, got {num_cube_cores!r}") 

95 context = _MegaMoeDiagnosticContext(topology=topology, num_cube_cores=num_cube_cores) 

96 return { 

97 "owner_label": MEGA_MOE_OWNER_LABEL, 

98 "owner_resolver": _resolve_expert_owner, 

99 "owner_context": context, 

100 }