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

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"""Real GM operand bindings for the shifted mHC forward kernel.""" 

16 

17from math import prod 

18from typing import Any 

19 

20from hyper_multicore.diagnostics.common.graph import bind_graph_tensor_ports 

21from hyper_multicore.diagnostics.dump.l1 import L1TensorView, L1TensorViewContext 

22from hyper_multicore.scheduler.graph import TensorSpec 

23 

24 

25def tensor_port(name: str, position: int, dtype_size: int = 4) -> TensorSpec: 

26 """Describe a native operand whose actual shape is supplied at launch. 

27 

28 Args: 

29 name: Logical diagnostic name of the native operand. 

30 position: Zero-based position in the native ABI parameter list. 

31 dtype_size: Expected element size in bytes. 

32 """ 

33 return TensorSpec(name, [], position, dtype_size=dtype_size, is_dynamic=True) 

34 

35 

36def resolve_token_view(context: L1TensorViewContext) -> L1TensorView: 

37 """Resolve a contiguous token tile, not a native CANN workspace alias. 

38 

39 Args: 

40 context: Selected invocation or operand context supplied by the diagnostic framework. 

41 """ 

42 tensor = context.parameters[context.tensor_spec.param_position] 

43 shape = tuple(tensor.shape) 

44 # Weights and scalar parameters are shared by every task. 

45 if context.tensor_spec.name in ("phi", "alpha", "bias", "norm_weight"): 

46 return context.make_view(shape, source_offset_bytes=0) 

47 token_count = int(context.parameters[1].shape[1]) 

48 offset = int(context.task_desc.task_index) * int(context.task_desc.task_split_value) 

49 rows = max(0, min(int(context.task_desc.task_split_value), token_count - offset)) 

50 if len(shape) >= 2 and shape[0] == 1 and shape[1] == token_count: 

51 row_shape = shape[2:] 

52 elif shape and shape[0] == token_count: 

53 row_shape = shape[1:] 

54 else: 

55 raise ValueError(f"mHC operand {context.tensor_spec.name} has no contiguous token axis: {shape}") 

56 return context.make_view((rows, *row_shape), 

57 source_offset_bytes=offset * prod(row_shape) * tensor.element_size()) 

58 

59 

60def bind_forward_graph(graph: Any, config: Any) -> Any: 

61 """Keep scheduling edges but replace placeholder operands by real ABI ports. 

62 

63 The ring X-cast workspace is CANN-owned, not hc_before_norm. It is deliberately 

64 not exposed as a Tensor. Iteration-major Sinkhorn caches are also omitted: 

65 a token tile there is strided, whereas L1 currently copies contiguous ranges. 

66 

67 Args: 

68 graph: Business compute graph whose scheduling storage remains unchanged. 

69 config: Private scheduling configuration; port binding does not rewrite tasks. 

70 """ 

71 del config 

72 p = tensor_port 

73 previous = p("previous_output", 0, 2) 

74 residual = p("residual", 1, 2) 

75 updated = p("new_residual", 12, 2) 

76 projection = p("hc_before_norm", 17) 

77 mixed = p("mixed_input", 21, 2) 

78 ports = { 

79 "mhc_post": ([previous, residual, p("previous_post_mix", 3), p("previous_residual_mix", 4)], [updated]), 

80 "mhc_norm_cast": ([updated], [p("inv_rms", 18)]), 

81 "mhc_projection": ([p("phi", 5)], [projection]), 

82 "mhc_input_mix": ([updated, p("previous_pre_mix", 2)], [mixed]), 

83 "mhc_rms_norm": ([mixed, p("norm_weight", 8, 2)], [p("block_input", 16, 2), p("rms_rstd", 22)]), 

84 "mhc_mapping": ([projection, p("inv_rms", 18), p("alpha", 6), p("bias", 7)], 

85 [p("next_pre_mix", 13), p("next_post_mix", 14), p("next_residual_mix", 15)]), 

86 } 

87 return bind_graph_tensor_ports(graph, ports, resolve_token_view)