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_npu_rotary_mul_args(x, cos, sin, rotary_mode=None): 

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

24 

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

26 

27 Mapping: 

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

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

30 

31 Args: 

32 x: Input tensor. 

33 cos: Cosine position encoding tensor. 

34 sin: Sine position encoding tensor. 

35 rotary_mode: Rotation mode string, optional. 

36 

37 Returns: 

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

39 

40 Raises: 

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

42 """ 

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

44 mode = 0 

45 elif rotary_mode == "interleave": 

46 mode = 1 

47 else: 

48 raise ValueError( 

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

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

51 ) 

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

53 

54 

55class RotaryPositionEmbeddingDistributedOp(DistributedOp): 

56 """Distributed operator for npu_rotary_mul. 

57 

58 Computes rotary position embedding element-wise: 

59 y = x * cos + x_rotate * sin 

60 

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

62 Output shape equals x shape exactly. 

63 

64 Serves ``torch_npu.npu_rotary_mul`` (rotary_mode keyword-only). 

65 

66 Sharding constraints: 

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

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

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

70 sharded. 

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

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

73 on that dimension. 

74 

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

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

77 

78 Output: 

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

80 """ 

81 

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

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

84 

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

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

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

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

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

90 

91 Args: 

92 x_layout: Layout of the x tensor. 

93 cos_layout: Layout of the cos tensor. 

94 sin_layout: Layout of the sin tensor. 

95 

96 Raises: 

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

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

99 """ 

100 x_tm = x_layout.tensor_map 

101 

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

103 raise ValueError( 

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

105 f"but got tensor_map={x_tm}" 

106 ) 

107 

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

109 tm = layout.tensor_map 

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

111 raise ValueError( 

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

113 f"but got tensor_map={tm}" 

114 ) 

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

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

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

118 raise ValueError( 

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

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

121 ) 

122 

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

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

125 

126 Normalizes with ``_normalize_npu_rotary_mul_args``, mapping the 

127 ``rotary_mode`` keyword onto ``mode``; ``rotary_mode`` (keyword-only) 

128 stays in ``local_kwargs``. 

129 

130 Args: 

131 args: Positional arguments, may include DTensors. 

132 kwargs: Keyword arguments. 

133 

134 Returns: 

135 tuple: (local_args, local_kwargs, cache_values) where 

136 cache_values = [x_layout, cos_layout, sin_layout]. 

137 """ 

138 # Step 1: normalize to canonical (x, cos, sin, mode_int), {}. 

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

140 

141 x, cos, sin, mode = norm_args 

142 

143 # Step 2: assemble local_args / local_kwargs; keyword-only params go 

144 # in kwargs. 

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

146 local_kwargs = {} 

147 if mode == 1: 

148 local_kwargs['rotary_mode'] = 'interleave' 

149 

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

151 return local_args, local_kwargs, cache_values 

152 

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

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

155 

156 Rules: 

157 1. Partial inputs are rejected. 

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

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

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

161 

162 Args: 

163 cache_values: [x_layout, cos_layout, sin_layout] 

164 

165 Returns: 

166 tuple: ((output_layout,), None) 

167 

168 Raises: 

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

170 or cos/sin sharding is inconsistent with x. 

171 """ 

172 x_layout = cache_values[0] 

173 cos_layout = cache_values[1] 

174 sin_layout = cache_values[2] 

175 

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

177 self._validate_input_layouts(x_layout, cos_layout, sin_layout) 

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