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

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"""Native backward operand bindings; reduction outputs are L0-only.""" 

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 resolve_backward_view(context: L1TensorViewContext) -> L1TensorView: 

26 """Capture local token ranges, never a concurrently accumulated parameter. 

27 

28 Args: 

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

30 """ 

31 position = context.tensor_spec.param_position 

32 tensor = context.parameters[position] 

33 if context.io == "output" and position in (25, 26, 27, 32): 

34 raise ValueError("mHC parameter gradients are multi-core reductions; use L0 or select token-local L1 outputs") 

35 shape = tuple(tensor.shape) 

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

37 if position in (4, 5, 6, 10, 11, 15): 

38 return context.make_view(shape, source_offset_bytes=0, note="shared read-only parameter or forward cache") 

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

40 row_shape = shape[2:] 

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

42 row_shape = shape[1:] 

43 else: 

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

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

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

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

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

49 

50 

51def bind_backward_graph(graph: Any, config: Any) -> Any: 

52 """Describe true ABI Tensors, leaving native workspace-only edges unexposed. 

53 

54 Args: 

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

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

57 """ 

58 del config 

59 names = ( 

60 "grad_block_input", "grad_next_post_mix", "grad_next_residual_mix", "new_residual", 

61 "phi", "alpha", "bias", "previous_pre_mix", "hc_before_norm", "inv_rms", "sum_out", "norm_out", 

62 "grad_next_pre_mix", "mixed_input", "rms_rstd", "norm_weight", "grad_new_residual", "residual", 

63 "previous_output", "previous_post_mix", "previous_residual_mix", "runtime", "events", "diagnostic", 

64 "grad_residual", "grad_phi", "grad_alpha", "grad_bias", "grad_previous_output", "grad_previous_pre_mix", 

65 "grad_previous_post_mix", "grad_previous_residual_mix", "grad_norm_weight", 

66 ) 

67 bf16 = {0, 3, 13, 15, 16, 17, 18, 24, 28} 

68 

69 def ports(indices: tuple[int, ...]) -> list[Any]: 

70 """Map native positions to actual operand dtypes without workspace aliases. 

71 

72 Args: 

73 indices: ABI positions selected from the backward operand table. 

74 """ 

75 return [ 

76 TensorSpec(names[index], [], index, dtype_size=2 if index in bf16 else 4, is_dynamic=True) 

77 for index in indices 

78 ] 

79 

80 bindings = { 

81 "grad_output_init": ([], []), 

82 "rms_norm_grad": (ports((0, 13, 14, 15)), ports((32,))), 

83 "mhc_grad_prev_a_and_mapping": (ports((3, 7, 12, 1, 2, 5, 6, 8, 9, 10, 11)), ports((29,))), 

84 "mhc_grad_phi_rms": (ports((3, 4, 8, 9)), ports((25, 26, 27))), 

85 "mhc_grad_prev_x_and_post": (ports((3, 4, 16, 17, 18, 19, 20)), ports((24, 28, 30, 31))), 

86 } 

87 return bind_graph_tensor_ports(graph, bindings, resolve_backward_view)