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

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

16 

17from __future__ import annotations 

18 

19from dataclasses import dataclass 

20from typing import Any 

21 

22import torch.distributed as dist 

23 

24from hyper_multicore.diagnostics import MegaKernelDiagnostics, prepare_mega_kernel_diagnostics 

25from hyper_multicore.scheduler.builder import build_runtime_config 

26from hyper_multicore.scheduler.config import _validate_runtime_task_queues 

27 

28from .diagnostics import bind_backward_graph 

29from .graph import ( 

30 DEFAULT_GRAD_TOKEN_TILE, 

31 build_mega_mhc_grad_graph, 

32 order_vector_tasks_by_stage, 

33 resolve_grad_token_tile, 

34) 

35 

36 

37@dataclass(frozen=True) 

38class MegaMhcGradPlan: 

39 """Own one read-only backward schedule, without invocation state.""" 

40 

41 token_count: int 

42 hidden_size: int 

43 num_cube_cores: int 

44 num_vector_cores: int 

45 token_tile: int 

46 runtime: MegaKernelDiagnostics 

47 event_capacity: int 

48 

49 

50def build_mega_mhc_grad_plan( 

51 token_count: int, 

52 hidden_size: int, 

53 device: Any, 

54 num_cube_cores: int, 

55 num_vector_cores: int, 

56 token_tile: int = DEFAULT_GRAD_TOKEN_TILE, 

57) -> MegaMhcGradPlan: 

58 """Build and materialize the token-tiled fused backward schedule. 

59 

60 Args: 

61 token_count: Total number of flattened tokens. 

62 hidden_size: Hidden dimension of every token. 

63 device: NPU device used to materialize runtime tensors. 

64 num_cube_cores: Number of Cube cores. 

65 num_vector_cores: Number of Vector cores. 

66 token_tile: Requested AIV token tile. 

67 

68 Returns: 

69 Reusable backward plan bound to the requested shape and device. 

70 """ 

71 token_tile = resolve_grad_token_tile(token_count, token_tile) 

72 graph, topology = build_mega_mhc_grad_graph( 

73 token_count, 

74 hidden_size, 

75 token_tile, 

76 num_vector_cores=num_vector_cores, 

77 ) 

78 runtime_config = build_runtime_config( 

79 graph, topology, rank_id=0, num_cube_cores=num_cube_cores 

80 ) 

81 order_vector_tasks_by_stage(runtime_config) 

82 _validate_runtime_task_queues(runtime_config, runtime_config.event_capacity) 

83 runtime = prepare_mega_kernel_diagnostics( 

84 graph, runtime_config, kernel_name="HyperMegaMhcGrad", 

85 parameter_count=33, device=device, owner_label="NativeKernel", 

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

87 vector_worker_stride=1, graph_adapter=bind_backward_graph, 

88 ) 

89 return MegaMhcGradPlan( 

90 token_count, 

91 hidden_size, 

92 num_cube_cores, 

93 num_vector_cores, 

94 token_tile, 

95 runtime, 

96 runtime_config.event_capacity, 

97 )