Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_multicore / scheduler / scheduler.py: 97%

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

16Task-queue reordering passes applied after the main fill loop. 

17 

18revise_task_queue — RATR: rank-aware tile reorder for dispatch/combine. 

19revise_gmm_task_queue_bwd — Backward GMM1/GMM4 expert interleave in cube_task_indices. 

20""" 

21from __future__ import annotations 

22 

23from hyper_multicore._core_spec import DEFAULT_CORE_SPEC 

24from hyper_multicore.scheduler.config import RuntimeConfigC, TaskSplitValue 

25 

26 

27def revise_task_queue(cfg: RuntimeConfigC, tsv: TaskSplitValue, 

28 dispatch_task_num: int, swiglu_task_num: int, combine_task_num: int | None = None) -> None: 

29 """ 

30 Reorder vector_task_indices for dispatch and combine based on rank_id (RATR). 

31 

32 Args: 

33 cfg: Runtime whose vector queue is reordered in place. 

34 tsv: Expert topology and local rank. 

35 dispatch_task_num: Number of dispatch tasks. 

36 swiglu_task_num: Number of intervening SwiGLU or SwiGLU-grad tasks. 

37 combine_task_num: Number of combine tasks, possibly using a different split. 

38 """ 

39 temp = list(cfg.vector_task_indices) 

40 ep_rank = [(peer + tsv.rank_id) % tsv.ep for peer in range(tsv.ep)] 

41 if combine_task_num is None: 

42 combine_task_num = dispatch_task_num 

43 segments = ((0, dispatch_task_num), (dispatch_task_num + swiglu_task_num, combine_task_num)) 

44 for start, task_count in segments: 

45 tasks_per_expert = task_count // tsv.all_expert_num 

46 tasks_per_peer = task_count // tsv.ep 

47 index = 0 

48 for expert in range(tsv.single_rank_expert_num): 

49 for tile in range(tasks_per_expert): 

50 for peer in ep_rank: 

51 source = peer * tasks_per_peer + expert * tasks_per_expert + tile 

52 cfg.vector_task_indices[start + index] = temp[start + source] 

53 index += 1 

54 

55 

56def revise_gmm_task_queue_bwd(cfg: RuntimeConfigC, tsv: TaskSplitValue, 

57 act_grad_task_num: int, 

58 num_cube_cores: int = DEFAULT_CORE_SPEC.num_cube_cores) -> None: 

59 """ 

60 Backward-only: interleave w2_grad and act_grad experts in cube_task_indices. 

61 

62 Result pattern: [w2_grad exp0, act_grad exp0, w2_grad exp1, act_grad exp1, ...] 

63 act_grad start offset = 0; w2_grad start offset = act_grad_task_num. 

64 

65 Args: 

66 cfg: Runtime image with the original Cube queue. 

67 tsv: Expert topology. 

68 act_grad_task_num: Offset at which weight-gradient tasks begin. 

69 num_cube_cores: Number of participating Cube workers. 

70 """ 

71 temp = list(cfg.cube_task_indices) 

72 expert_single = tsv.single_rank_expert_num 

73 changes_num = 2 # two streams: w2_grad (index=1) and act_grad (index=0) 

74 

75 for i in range(expert_single * changes_num): 

76 index = 1 - (i % changes_num) # alternates: 1, 0, 1, 0, ... 

77 m = i // changes_num # expert block: 0, 0, 1, 1, ... 

78 for j in range(num_cube_cores): 

79 dst = i * num_cube_cores + j 

80 src = index * act_grad_task_num + m * num_cube_cores + j 

81 cfg.cube_task_indices[dst] = temp[src]