Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_multicore / modules / mega_moe / backward / tiling_tables.py: 92%

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

16Backward tiling lookup tables. 

17Each dict maps split_value -> CSV tiling string (for GMM ops) or 

18struct fields (for SwiGLU-grad). 

19 

20C++ source functions: 

21 GMM1 (pos 20) → get_tiling_data_first_gmm(split_value) 

22 GMM2 (pos 21) → get_tiling_data_second_gmm(split_value) 

23 GMM3 (pos 22) → get_tiling_data_third_gmm(split_value) 

24 GMM4 (pos 23) → get_tiling_data_fourth_gmm(split_value) 

25 SwiGLU-grad (pos 24) → get_tiling_data_swiglu_grad(split_value) 

26""" 

27# pylint: disable=line-too-long 

28import copy 

29 

30from hyper_multicore._core_spec import DEFAULT_CORE_SPEC 

31from hyper_multicore.scheduler.config import ( 

32 SwiGluTilingDataC, 

33) 

34 

35from .._gmm_tiling import _expand_gmm_string, _patch_gmm_dims 

36 

37 

38# ── helpers ─────────────────────────────────────────────────────────────────── 

39 

40def _expand_swiglu_struct(td: SwiGluTilingDataC) -> bytes: 

41 """Tile SwiGluTilingDataC 49 times (= 2 * 24 AI Cube cores + 1).""" 

42 return bytes(td) * 49 

43 

44 

45def _make_swiglu(row_len, col_len, base_row_len, base_col_len): 

46 """Build a SwiGluTilingDataC for SwiGLU-grad with the given tile dimensions.""" 

47 # These member names match the CANN tiling ABI and must retain their native spellings. 

48 # pylint: disable=invalid-name 

49 td = SwiGluTilingDataC() 

50 td.is32BAligned = 1 

51 td.isDoubleBuffer = 1 

52 td.rowLen = row_len 

53 td.colLen = col_len 

54 td.baseRowLen = base_row_len 

55 td.baseColLen = base_col_len 

56 td.activateLeft = 0 

57 td.biasIsEmpty = 0 

58 td.quantScaleIsEmpty = 0 

59 td.activateScaleIsEmpty = 0 

60 td.swiColLen = 0 

61 td.perRowLen = 0 

62 td.modRowLen = 0 

63 td.usedCoreNum = 0 

64 return td 

65 

66 

67# ── Backward GMM1: x=[per_rank_seq,7168], weight=[E,2048,7168] ─────────────── 

68# C++: get_tiling_data_first_gmm(split_value) 

69# x input [*, 7168] → y output [*, 2048] 

70 

71FIRST_GMM_TABLE = { 

72 4096: "8,24,0,0,0,0,0,1,1,1,0,0,0,0,4096,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,4294967295,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,7168,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,2048,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,24,4096,2048,7168,7168,4096,256,7168,128,256,64,8,8,1,1,0,0,0,0,393216,131072,0,1,1,1,1,4,4,0,0,2,2,1,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0", 

73} 

74 

75 

76# ── Backward GMM2: x=[per_rank_seq,4096], weight=[E,7168,4096] ─────────────── 

77# C++: get_tiling_data_second_gmm(split_value) 

78# x input [*, 4096] → y output [*, 7168] 

79 

80SECOND_GMM_TABLE = { 

81 4096: "8,24,0,0,0,0,0,1,1,1,0,0,0,0,4096,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,4294967295,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,4096,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,7168,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,24,4096,7168,4096,4096,4096,256,4096,128,256,64,8,8,1,1,0,0,0,0,393216,131072,0,1,1,1,1,4,4,0,0,2,2,1,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0", 

82} 

83 

84 

85# ── Backward GMM3: x1^T=[sre,7168], x2=[sre,4096], weight grad ─────────────── 

86# C++: get_tiling_data_third_gmm(split_value) 

87 

88THIRD_GMM_TABLE = { 

89 4096: "8,24,0,0,0,0,0,1,1,1,2,0,0,0,7168,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,7168,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,4294967295,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,4096,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,24,7168,4096,4096,4096,7168,256,4096,128,256,64,8,8,1,1,0,0,0,0,393216,131072,0,1,1,1,1,4,4,0,0,2,2,1,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0", 

90} 

91 

92 

93# ── Backward GMM4: x1^T=[sre,2048], x2=[sre,7168], weight grad ─────────────── 

94# C++: get_tiling_data_fourth_gmm(split_value) 

95 

96FOURTH_GMM_TABLE = { 

97 4096: "8,24,0,0,0,0,0,1,1,1,2,0,0,0,2048,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,2048,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,4294967295,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,7168,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,24,2048,7168,4096,4096,2048,256,4096,128,256,64,8,8,1,1,0,0,0,0,393216,131072,0,1,1,1,1,4,4,0,0,2,2,1,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0", 

98} 

99 

100 

101# ── Backward SwiGLU-grad tiling ─────────────────────────────────────────────── 

102# C++: get_tiling_data_swiglu_grad(split_value) → baseColLen=256 

103 

104SWIGLU_GRAD_TABLE = { 

105 32: _make_swiglu(row_len=32, col_len=2048, base_row_len=19, base_col_len=256), 

106 64: _make_swiglu(row_len=64, col_len=2048, base_row_len=19, base_col_len=256), 

107 128: _make_swiglu(row_len=128, col_len=2048, base_row_len=19, base_col_len=256), 

108} 

109 

110 

111# ── Public API ──────────────────────────────────────────────────────────────── 

112 

113def get_act_grad_tiling_bytes(split_value: int, *, 

114 hidden_size: int = 7168, 

115 intermediate_size: int = 2048, 

116 num_groups: int = 8, 

117 num_cube_cores: int = DEFAULT_CORE_SPEC.num_cube_cores) -> bytes: 

118 """Backward GMM1 [x,hidden]×[E,intermediate,hidden] → y=[x,intermediate].""" 

119 csv = FIRST_GMM_TABLE.get(split_value) 

120 if csv is None: 

121 raise KeyError(f"No first_gmm tiling for split_value={split_value}") 

122 csv = _patch_gmm_dims(csv, 

123 {7168: hidden_size, 2048: intermediate_size}, 

124 num_groups, split_value=split_value, 

125 num_cube_cores=num_cube_cores) 

126 return _expand_gmm_string(csv, repeat=num_cube_cores) 

127 

128 

129def get_gate_grad_tiling_bytes(split_value: int, *, 

130 hidden_size: int = 7168, 

131 intermediate_size: int = 2048, 

132 num_groups: int = 8, 

133 num_cube_cores: int = DEFAULT_CORE_SPEC.num_cube_cores) -> bytes: 

134 """Backward GMM2 [x,intermediate*2]×[E,hidden,intermediate*2] → y=[x,hidden].""" 

135 csv = SECOND_GMM_TABLE.get(split_value) 

136 if csv is None: 

137 raise KeyError(f"No second_gmm tiling for split_value={split_value}") 

138 csv = _patch_gmm_dims(csv, 

139 {4096: intermediate_size * 2, 7168: hidden_size}, 

140 num_groups, split_value=split_value, 

141 num_cube_cores=num_cube_cores) 

142 return _expand_gmm_string(csv, repeat=num_cube_cores) 

143 

144 

145def get_w1_grad_tiling_bytes(split_value: int, *, 

146 hidden_size: int = 7168, 

147 intermediate_size: int = 2048, 

148 num_groups: int = 8, 

149 num_cube_cores: int = DEFAULT_CORE_SPEC.num_cube_cores, 

150 expert_weight_layout: str = "EIO") -> bytes: 

151 """Build W13-gradient tiling with caller-layout output axes and dynamic reduction K.""" 

152 csv = THIRD_GMM_TABLE.get(split_value) 

153 if csv is None: 

154 raise KeyError(f"No third_gmm tiling for split_value={split_value}") 

155 rows, columns = ((intermediate_size * 2, hidden_size) if expert_weight_layout == "EOI" 

156 else (hidden_size, intermediate_size * 2)) 

157 csv = _patch_gmm_dims(csv, 

158 {7168: rows, 4096: columns}, 

159 num_groups, split_value=split_value, 

160 num_cube_cores=num_cube_cores) 

161 return _expand_gmm_string(csv, repeat=num_cube_cores) 

162 

163 

164def get_w2_grad_tiling_bytes(split_value: int, *, 

165 hidden_size: int = 7168, 

166 intermediate_size: int = 2048, 

167 num_groups: int = 8, 

168 num_cube_cores: int = DEFAULT_CORE_SPEC.num_cube_cores, 

169 expert_weight_layout: str = "EIO") -> bytes: 

170 """Build W2-gradient tiling with caller-layout output axes and dynamic reduction K.""" 

171 csv = FOURTH_GMM_TABLE.get(split_value) 

172 if csv is None: 

173 raise KeyError(f"No fourth_gmm tiling for split_value={split_value}") 

174 rows, columns = ((hidden_size, intermediate_size) if expert_weight_layout == "EOI" 

175 else (intermediate_size, hidden_size)) 

176 csv = _patch_gmm_dims(csv, 

177 {2048: rows, 7168: columns, 

178 4096: intermediate_size * 2}, 

179 num_groups, split_value=split_value, 

180 num_cube_cores=num_cube_cores) 

181 return _expand_gmm_string(csv, repeat=num_cube_cores) 

182 

183 

184def get_swiglu_grad_tiling_bytes(split_value: int, *, 

185 intermediate_size: int = 2048) -> bytes: 

186 """Backward SwiGLU-grad tiling (colLen=intermediate_size).""" 

187 td = SWIGLU_GRAD_TABLE.get(split_value) 

188 if td is None: 

189 raise KeyError(f"No swiglu_grad tiling for split_value={split_value}") 

190 td = copy.copy(td) 

191 td.colLen = intermediate_size 

192 return _expand_swiglu_struct(td)