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

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

16Forward tiling lookup tables. 

17Each function mirrors the C++ get_tiling_data_* functions. 

18String-based GMM tiling is expanded num_cube_cores times (default 24); 

19struct-based SwiGLU tiling is expanded 49 times (= 2 * 24 + 1 for Ascend 910B). 

20""" 

21# pylint: disable=line-too-long 

22import copy 

23 

24from hyper_multicore._core_spec import DEFAULT_CORE_SPEC 

25from hyper_multicore.scheduler.config import ( 

26 SwiGluTilingDataC, 

27) 

28 

29from .._gmm_tiling import _expand_gmm_string, _patch_gmm_dims 

30 

31 

32# ── helpers ─────────────────────────────────────────────────────────────────── 

33 

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

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

36 one = bytes(td) 

37 return one * 49 

38 

39 

40# ── Forward GMM1: x=[per_rank_seq,7168], weight=[E,7168,4096] ───────────────── 

41# C++: get_tiling_data_gmm_x_7168_g2(split_value) 

42 

43GMM1_TABLE = { 

44 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,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,4096,4096,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", 

45} 

46 

47 

48# ── Forward GMM2: x=[per_rank_seq,2048], weight=[E,2048,7168] ───────────────── 

49# C++: get_tiling_data_gmm_x_2048_g2(split_value) 

50 

51GMM2_TABLE = { 

52 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,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,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,2048,2048,4096,256,2048,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", 

53} 

54 

55 

56# ── Forward SwiGLU: colLen=2048 ─────────────────────────────────────────────── 

57# C++: get_tiling_data_swiglu(split_value) → SwiGluTilingData struct 

58 

59def _make_swiglu(is32b, idbl, row_len, col_len, base_row_len, base_col_len, 

60 act_left, bias_empty, qs_empty, as_empty, 

61 swi_col_len, per_row_len, mod_row_len, used_core_num): 

62 """Build a SwiGluTilingDataC struct with the given tile parameters.""" 

63 td = SwiGluTilingDataC() 

64 td.is32BAligned = is32b 

65 td.isDoubleBuffer = idbl 

66 td.rowLen = row_len 

67 td.colLen = col_len 

68 td.baseRowLen = base_row_len 

69 td.baseColLen = base_col_len 

70 td.activateLeft = act_left 

71 td.biasIsEmpty = bias_empty 

72 td.quantScaleIsEmpty = qs_empty 

73 td.activateScaleIsEmpty = as_empty 

74 td.swiColLen = swi_col_len 

75 td.perRowLen = per_row_len 

76 td.modRowLen = mod_row_len 

77 td.usedCoreNum = used_core_num 

78 return td 

79 

80 

81SWIGLU_TABLE = { 

82 16: _make_swiglu(1,1,16, 2048, 16, 512, 0,0,0,0, 0,0,0,0), 

83 32: _make_swiglu(1,1,32, 2048, 32, 256, 0,0,0,0, 0,0,0,0), 

84 64: _make_swiglu(1,1,64, 2048, 19, 512, 0,0,0,0, 0,0,0,0), 

85 128: _make_swiglu(1,1,128, 2048, 19, 512, 0,0,0,0, 0,0,0,0), 

86 256: _make_swiglu(1,1,256, 2048, 19, 512, 0,0,0,0, 0,0,0,0), 

87 512: _make_swiglu(1,1,512, 2048, 19, 512, 0,0,0,0, 0,0,0,0), 

88 1024: _make_swiglu(1,1,1024,2048, 19, 512, 0,0,0,0, 0,0,0,0), 

89} 

90 

91 

92# ── Public API ──────────────────────────────────────────────────────────────── 

93 

94def get_up_proj_tiling_bytes(split_value: int, *, 

95 hidden_size: int = 7168, 

96 intermediate_size: int = 2048, 

97 num_groups: int = 8, 

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

99 """Forward GMM1 tiling [x,hidden]×[E,hidden,intermediate*2], expanded num_cube_cores times.""" 

100 if split_value not in GMM1_TABLE: 

101 raise KeyError(f"GMM1 tiling: split_value={split_value} not found") 

102 csv = _patch_gmm_dims(GMM1_TABLE[split_value], 

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

104 num_groups, split_value=split_value, 

105 num_cube_cores=num_cube_cores) 

106 return _expand_gmm_string(csv, repeat=num_cube_cores) 

107 

108 

109def get_down_proj_tiling_bytes(split_value: int, *, 

110 hidden_size: int = 7168, 

111 intermediate_size: int = 2048, 

112 num_groups: int = 8, 

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

114 """Forward GMM2 tiling [x,intermediate]×[E,intermediate,hidden], expanded num_cube_cores times.""" 

115 if split_value not in GMM2_TABLE: 

116 raise KeyError(f"GMM2 tiling: split_value={split_value} not found") 

117 csv = _patch_gmm_dims(GMM2_TABLE[split_value], 

118 {2048: intermediate_size, 7168: hidden_size}, 

119 num_groups, split_value=split_value, 

120 num_cube_cores=num_cube_cores) 

121 return _expand_gmm_string(csv, repeat=num_cube_cores) 

122 

123 

124def get_swiglu_tiling_bytes(split_value: int, *, 

125 intermediate_size: int = 2048) -> bytes: 

126 """Forward SwiGLU tiling (colLen=intermediate_size), expanded 49 times.""" 

127 if split_value not in SWIGLU_TABLE: 

128 raise KeyError(f"SwiGLU tiling: split_value={split_value} not found") 

129 td = copy.copy(SWIGLU_TABLE[split_value]) 

130 td.colLen = intermediate_size 

131 return _expand_swiglu_struct(td)