Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / shard / ops / parallel_rotary_position_embedding.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"""Distributed implementation for RotaryPositionEmbedding operator.""" 

16import copy 

17from typing import Optional, Tuple 

18 

19from .parallel_ops import DistributedOp 

20 

21 

22def _normalize_rpe_args(x, cos, sin, mode=0): 

23 """Normalize positional and keyword arguments into a canonical positional tuple. 

24 

25 Args: 

26 x: Input tensor. 

27 cos: Cosine position encoding tensor. 

28 sin: Sine position encoding tensor. 

29 mode: Rotation mode. 0=rotate_half, 1=rotate_interleaved, 2=quarter, 

30 3=interleave-half. Defaults to 0. 

31 

32 Returns: 

33 tuple: (positional_args_tuple, empty_kwargs_dict) 

34 """ 

35 return (x, cos, sin, mode), {} 

36 

37 

38def _normalize_npu_rotary_mul_args(x, cos, sin, rotary_mode=None): 

39 """Normalize npu_rotary_mul args to canonical positional form. 

40 

41 Maps ``rotary_mode`` (string, keyword-only) to ``mode`` (int, positional), 

42 matching the canonical form produced by :func:`_normalize_rpe_args`. 

43 

44 Mapping: 

45 - ``None`` / not specified / ``"half"`` → mode=0 (rotate_half) 

46 - ``"interleave"`` → mode=1 (rotate_interleaved) 

47 

48 Args: 

49 x: Input tensor. 

50 cos: Cosine position encoding tensor. 

51 sin: Sine position encoding tensor. 

52 rotary_mode: Rotation mode string, optional. 

53 

54 Returns: 

55 tuple: ((x, cos, sin, mode), {}) 

56 

57 Raises: 

58 ValueError: If rotary_mode is not None, 'half', or 'interleave'. 

59 """ 

60 if rotary_mode in (None, "half"): 

61 mode = 0 

62 elif rotary_mode == "interleave": 

63 mode = 1 

64 else: 

65 raise ValueError( 

66 f"npu_rotary_mul: unsupported rotary_mode '{rotary_mode}'. " 

67 f"Supported values: None, 'half', 'interleave'." 

68 ) 

69 return (x, cos, sin, mode), {} 

70 

71 

72class RotaryPositionEmbeddingDistributedOp(DistributedOp): 

73 """Distributed operator for RotaryPositionEmbedding and npu_rotary_mul. 

74 

75 Computes rotary position embedding element-wise: 

76 y = x * cos + x_rotate * sin 

77 

78 where x_rotate is obtained by rotating within the last (D) dimension. 

79 Output shape equals x shape exactly. 

80 

81 Serves both: 

82 - MindSpore Primitive ``RotaryPositionEmbedding`` (mode positional) 

83 - PyTorch ``torch_npu.npu_rotary_mul`` (rotary_mode keyword-only) 

84 

85 Sharding constraints: 

86 - D (last dim) must be replicated for x, cos, and sin: the kernel rotates 

87 within D and the operation is indivisible along that axis. 

88 - B, N, S dims are fully independent across positions and can be freely 

89 sharded. 

90 - cos/sin may have any subset of non-D dims replicated (broadcast case), 

91 but if cos/sin is sharded on a dimension, it must match x's sharding 

92 on that dimension. 

93 

94 MODE / rotary_mode does not affect layout inference: all modes produce 

95 output shape == x shape and leave B/N/S independence unchanged. 

96 

97 Output: 

98 Single tensor with the same shape and layout as x. 

99 """ 

100 

101 _MS_PRIMITIVE_OP_NAMES = frozenset({'RotaryPositionEmbedding'}) 

102 

103 def _validate_input_layouts(self, x_layout, cos_layout, sin_layout) -> None: 

104 """Validate sharding constraints for all input tensors. 

105 

106 Rules (applied to both 4-D BNSD/BSND/SBND and 3-D TND layouts): 

107 - x's last dim (D) must be replicated. 

108 - cos and sin's last dim (D) must be replicated. 

109 - For any non-D dimension d: if cos/sin is sharded there, the mesh 

110 axis must equal x's mesh axis on the same dimension. 

111 

112 Args: 

113 x_layout: Layout of the x tensor. 

114 cos_layout: Layout of the cos tensor. 

115 sin_layout: Layout of the sin tensor. 

116 

117 Raises: 

118 ValueError: If D is sharded for any input, or if cos/sin sharding 

119 is inconsistent with x on any non-D dimension. 

120 """ 

121 x_tm = x_layout.tensor_map 

122 

123 if x_tm[-1] != -1: 

124 raise ValueError( 

125 f"For {self.op_name}, D (last dim) of x must be replicated, " 

126 f"but got tensor_map={x_tm}" 

127 ) 

128 

129 for name, layout in [('cos', cos_layout), ('sin', sin_layout)]: 

130 tm = layout.tensor_map 

131 if tm[-1] != -1: 

132 raise ValueError( 

133 f"For {self.op_name}, D (last dim) of {name} must be replicated, " 

134 f"but got tensor_map={tm}" 

135 ) 

136 for d in range(len(tm) - 1): 

137 x_d = x_tm[d] if d < len(x_tm) - 1 else -1 

138 if tm[d] != -1 and tm[d] != x_d: 

139 raise ValueError( 

140 f"For {self.op_name}, {name} sharding on dim {d} must match x " 

141 f"or be replicated, but got x={x_d}, {name}={tm[d]}" 

142 ) 

143 

144 def preprocess(self, args: tuple, kwargs: dict) -> Optional[tuple]: 

145 """Extract local tensors and build the layout cache. 

146 

147 Cross-platform routing via ``_MS_PRIMITIVE_OP_NAMES``: 

148 - Primitive (RotaryPositionEmbedding): normalize with 

149 _normalize_rpe_args, mode → local_args (all positional), 

150 local_kwargs = {}. 

151 - PyTorch (npu_rotary_mul): normalize with 

152 _normalize_npu_rotary_mul_args, rotary_mode → local_kwargs 

153 (keyword-only). 

154 

155 Args: 

156 args: Positional arguments, may include DTensors. 

157 kwargs: Keyword arguments. 

158 

159 Returns: 

160 tuple: (local_args, local_kwargs, cache_values) where 

161 cache_values = [x_layout, cos_layout, sin_layout]. 

162 """ 

163 # Step 1: normalize — both functions return canonical 

164 # (x, cos, sin, mode_int), {}. 

165 if self.op_name in self._MS_PRIMITIVE_OP_NAMES: 

166 norm_args, _ = _normalize_rpe_args(*args, **kwargs) 

167 else: 

168 norm_args, _ = _normalize_npu_rotary_mul_args(*args, **kwargs) 

169 

170 x, cos, sin, mode = norm_args 

171 

172 # Step 2: assemble local_args / local_kwargs by platform convention. 

173 if self.op_name in self._MS_PRIMITIVE_OP_NAMES: 

174 # MindSpore Primitive: all positional, no kwargs. 

175 local_args = (x.to_local(), cos.to_local(), sin.to_local(), mode) 

176 local_kwargs = {} 

177 else: 

178 # PyTorch: keyword-only params go in kwargs. 

179 local_args = (x.to_local(), cos.to_local(), sin.to_local()) 

180 local_kwargs = {} 

181 if mode == 1: 

182 local_kwargs['rotary_mode'] = 'interleave' 

183 

184 cache_values = [x.layout, cos.layout, sin.layout] 

185 return local_args, local_kwargs, cache_values 

186 

187 def infer_layout(self, cache_values: list) -> Tuple[tuple, None]: 

188 """Infer output layout for the single output tensor. 

189 

190 Rules: 

191 1. Partial inputs are rejected. 

192 2. D (last dim) must be replicated for x, cos, and sin. 

193 3. cos/sin sharding on non-D dims must match x or be replicated. 

194 4. Output layout = deep copy of x_layout (output shape == x shape). 

195 

196 Args: 

197 cache_values: [x_layout, cos_layout, sin_layout] 

198 

199 Returns: 

200 tuple: ((output_layout,), None) 

201 

202 Raises: 

203 ValueError: If any input has Partial status, D is sharded, 

204 or cos/sin sharding is inconsistent with x. 

205 """ 

206 x_layout = cache_values[0] 

207 cos_layout = cache_values[1] 

208 sin_layout = cache_values[2] 

209 

210 self._check_partial_inputs([x_layout, cos_layout, sin_layout]) 

211 self._validate_input_layouts(x_layout, cos_layout, sin_layout) 

212 return (copy.deepcopy(x_layout),), None