Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_multicore / modules / mega_mhc / plan.py: 96%

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"""Shape-bound RuntimeConfig plan for HyperMegaMhc.""" 

16 

17from __future__ import annotations 

18 

19from dataclasses import dataclass 

20from typing import Any 

21 

22import torch 

23import torch.distributed as dist 

24 

25from hyper_multicore.diagnostics import MegaKernelDiagnostics, prepare_mega_kernel_diagnostics 

26from hyper_multicore.scheduler.builder import build_runtime_config 

27from hyper_multicore.scheduler.config import _validate_runtime_task_queues 

28 

29from .diagnostics import bind_forward_graph 

30from .graph import DEFAULT_TOKEN_TILE, build_mega_mhc_graph, resolve_token_tile 

31 

32 

33@dataclass(frozen=True) 

34class MegaMhcPlan: 

35 """Own one read-only token-shape schedule, without invocation state.""" 

36 

37 token_count: int 

38 hidden_size: int 

39 num_cube_cores: int 

40 token_tile: int 

41 runtime: MegaKernelDiagnostics 

42 event_capacity: int 

43 

44 @property 

45 def runtime_config(self) -> torch.Tensor: 

46 """Return the profiling-disabled RuntimeConfig tensor.""" 

47 return self.runtime.runtime_config 

48 

49 

50def build_mega_mhc_plan( 

51 token_count: int, 

52 hidden_size: int, 

53 device: Any, 

54 num_cube_cores: int, 

55 token_tile: int = DEFAULT_TOKEN_TILE, 

56) -> MegaMhcPlan: 

57 """Build and materialize a token-owned HyperMegaMhc schedule. 

58 

59 Args: 

60 token_count: Total flattened token count. 

61 hidden_size: Hidden dimension. 

62 device: Target NPU device. 

63 num_cube_cores: Number of Cube cores. 

64 token_tile: Preferred token count per task. 

65 """ 

66 token_tile = resolve_token_tile(token_count, token_tile, num_cube_cores) 

67 graph, topology = build_mega_mhc_graph( 

68 token_count, 

69 hidden_size, 

70 num_cube_cores, 

71 token_tile, 

72 ) 

73 runtime_config = build_runtime_config( 

74 graph, 

75 topology, 

76 rank_id=0, 

77 num_cube_cores=num_cube_cores, 

78 ) 

79 _validate_runtime_task_queues(runtime_config, runtime_config.event_capacity) 

80 runtime = prepare_mega_kernel_diagnostics( 

81 graph, runtime_config, kernel_name="HyperMegaMhc", 

82 parameter_count=23, device=device, owner_label="TokenPartition", 

83 rank=dist.get_rank() if dist.is_initialized() else 0, 

84 vector_worker_stride=1, graph_adapter=bind_forward_graph, 

85 ) 

86 return MegaMhcPlan( 

87 token_count=token_count, 

88 hidden_size=hidden_size, 

89 num_cube_cores=num_cube_cores, 

90 token_tile=token_tile, 

91 runtime=runtime, 

92 event_capacity=runtime_config.event_capacity, 

93 )