Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_multicore / tasks / swiglu.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"""SwiGLU fill config: forward (TASK_SWI_GLU) and backward gradient (TASK_SWI_GLU_GRAD).""" 

16 

17import math 

18import struct 

19from dataclasses import dataclass 

20 

21from hyper_multicore.scheduler.config import ( 

22 TaskDescC, TensorDescC, RuntimeConfigC, 

23 TaskAiCoreType, TaskType, 

24 MAX_TENSOR_DIMS, 

25) 

26from hyper_multicore.scheduler.graph import OpType 

27from hyper_multicore.tasks.task_base import FillConfig 

28from hyper_multicore.tasks.utils import advance_tsv_vector 

29 

30 

31def _encode_clamp_limit(clamp_limit: float | None) -> int: 

32 """Encode a positive float32 clamp limit, reserving zero for no clamp.""" 

33 if clamp_limit is None: 

34 return 0 

35 valid_type = isinstance(clamp_limit, (int, float)) and not isinstance( 

36 clamp_limit, bool 

37 ) 

38 try: 

39 clamp_bits = struct.unpack( 

40 "<I", struct.pack("<f", float(clamp_limit)) 

41 )[0] 

42 encoded_limit = struct.unpack("<f", struct.pack("<I", clamp_bits))[0] 

43 except (OverflowError, TypeError, ValueError, struct.error): 

44 encoded_limit = 0.0 

45 clamp_bits = 0 

46 if not valid_type or not math.isfinite(encoded_limit) or encoded_limit <= 0: 

47 raise ValueError( 

48 "clamp_limit must be finite, positive, and float32-representable " 

49 "when provided" 

50 ) 

51 return clamp_bits 

52 

53 

54@dataclass 

55class SwiGLUFillConfig(FillConfig): 

56 """ 

57 SwiGLU fill config — forward (TASK_SWI_GLU) and backward (TASK_SWI_GLU_GRAD). 

58 

59 ``clamp_limit`` is serialized in ``TaskDesc.extra_value_0``. Zero keeps 

60 the legacy unclamped path; a positive value enables clipped SwiGLU. 

61 """ 

62 

63 clamp_limit: float | None = None 

64 

65 def fill( 

66 self, 

67 cfg: RuntimeConfigC, 

68 op, 

69 tsv, 

70 ) -> None: 

71 """Serialize SwiGLU tasks into the runtime configuration. 

72 

73 Args: 

74 cfg: Runtime configuration receiving task descriptors. 

75 op: SwiGLU operator node to serialize. 

76 tsv: Topology and task split values for the current graph. 

77 """ 

78 task_type = (TaskType.TASK_SWI_GLU if op.op_type == OpType.SWIGLU 

79 else TaskType.TASK_SWI_GLU_GRAD) 

80 clamp_bits = _encode_clamp_limit(self.clamp_limit) 

81 

82 num_triggers = tsv.per_expert_seq // op.split_value 

83 task_num = op.task_num 

84 param = op.param_positions 

85 

86 for i in range(task_num): 

87 task = TaskDescC() 

88 task.task_type = task_type 

89 task.task_aicore_type = TaskAiCoreType.TASK_AICORE_VECTOR 

90 task.num_inputs = len(op.inputs) 

91 task.num_outputs = len(op.outputs) 

92 

93 for j, spec in enumerate(op.inputs): 

94 td = TensorDescC() 

95 td.tensor_type = 1 # all SwiGLU inputs are tensor lists 

96 td.data_type = spec.dtype_size 

97 td.input_position = param[j] 

98 td.base_ptr_offset = i * spec.shape[1] * op.split_value 

99 td.dynamic_shape = int(spec.is_dynamic) 

100 for k in range(min(len(spec.shape), MAX_TENSOR_DIMS)): 

101 td.dim[k] = spec.shape[k] 

102 task.inputs[j] = td 

103 

104 for j, spec in enumerate(op.outputs): 

105 td = TensorDescC() 

106 td.tensor_type = 1 # all SwiGLU outputs are tensor lists 

107 td.data_type = spec.dtype_size 

108 td.input_position = param[task.num_inputs + j] 

109 td.base_ptr_offset = i * spec.shape[1] * op.split_value 

110 td.dynamic_shape = int(spec.is_dynamic) 

111 for k in range(min(len(spec.shape), MAX_TENSOR_DIMS)): 

112 td.dim[k] = spec.shape[k] 

113 task.outputs[j] = td 

114 

115 ev_idx = i // num_triggers 

116 task.dependent_event = tsv.pre_pre_event_num + ev_idx + 1 

117 task.trigger_event = tsv.pre_event_num + ev_idx + 1 

118 cfg.all_event_num_triggers[task.trigger_event] = num_triggers 

119 

120 task.task_index = i 

121 task.task_split_num = task_num 

122 task.task_split_value = op.split_value 

123 task.extra_value_0 = clamp_bits 

124 task.tiling_data_position = op.tiling_position 

125 

126 cfg.all_tasks[tsv.pre_task_num + i] = task 

127 cfg.vector_task_indices[tsv.pre_vector_task_num + i] = tsv.pre_task_num + i 

128 

129 cfg.task_index_num[1] += task_num 

130 advance_tsv_vector(tsv, task_num, event_group_size=tsv.single_rank_expert_num)