Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / infer / utils.py: 0%

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"""Generation configuration and mask helpers.""" 

16from dataclasses import dataclass 

17from typing import Any, Callable, List, Optional 

18 

19import torch 

20import torch.distributed as dist 

21 

22 

23@dataclass 

24class GenerationConfig: 

25 """Runtime options for autoregressive generation.""" 

26 

27 max_new_tokens: int = 256 

28 temperature: float = 1.0 

29 top_k: int = 50 

30 top_p: float = 1.0 

31 do_sample: bool = False 

32 eos_token_id: Optional[int] = 2 

33 pad_token_id: int = 0 

34 repetition_penalty: float = 1.0 

35 use_cache: bool = True 

36 prefix_past_key_values: Optional[Any] = None 

37 prefix_attention_mask: Optional[torch.Tensor] = None 

38 prefix_sequence_shard_info: Optional[Any] = None 

39 prefix_cache_length: Optional[int] = None 

40 context_parallel_cache: bool = False 

41 context_parallel_rank: Optional[int] = None 

42 context_parallel_world_size: Optional[int] = None 

43 # context_logits_rank is local to context_process_group. 

44 context_logits_rank: Optional[Any] = None 

45 context_process_group: Optional[Any] = None 

46 gather_logits: bool = False 

47 logits_process_group: Optional[Any] = None 

48 logits_gather_dim: int = -1 

49 mask_dtype: Optional[torch.dtype] = None 

50 logits_processor: Optional[List[Callable]] = None 

51 stopping_criteria: Optional[List[Callable]] = None 

52 

53 def __post_init__(self): 

54 self._validate_sampling() 

55 self._validate_prefix() 

56 self._validate_context_parallel() 

57 if self.logits_gather_dim >= 0: 

58 raise ValueError("logits_gather_dim must be negative") 

59 if self.mask_dtype is not None and not isinstance(self.mask_dtype, torch.dtype): 

60 raise ValueError("mask_dtype must be a torch.dtype") 

61 self._validate_callables(self.logits_processor, "logits_processor") 

62 self._validate_callables(self.stopping_criteria, "stopping_criteria") 

63 

64 def _validate_sampling(self) -> None: 

65 """Validate scalar sampling and stopping options.""" 

66 if self.max_new_tokens < 0: 

67 raise ValueError("max_new_tokens must be >= 0") 

68 if self.temperature <= 0: 

69 raise ValueError("temperature must be > 0") 

70 if self.top_k < 0: 

71 raise ValueError("top_k must be >= 0") 

72 if not 0 < self.top_p <= 1.0: 

73 raise ValueError("top_p must be in (0, 1]") 

74 if self.repetition_penalty <= 0: 

75 raise ValueError("repetition_penalty must be > 0") 

76 

77 def _validate_prefix(self) -> None: 

78 """Validate optional prefix cache metadata.""" 

79 if self.prefix_past_key_values is None: 

80 if self.prefix_attention_mask is not None: 

81 raise ValueError("prefix_attention_mask requires prefix_past_key_values") 

82 if self.prefix_sequence_shard_info is not None: 

83 raise ValueError("prefix_sequence_shard_info requires prefix_past_key_values") 

84 if self.prefix_cache_length is not None: 

85 raise ValueError("prefix_cache_length requires prefix_past_key_values") 

86 else: 

87 if not self.use_cache: 

88 raise ValueError("prefix_past_key_values requires use_cache=True") 

89 if self.prefix_cache_length is not None and self.prefix_cache_length < 0: 

90 raise ValueError("prefix_cache_length must be >= 0") 

91 if self.prefix_sequence_shard_info is not None and not self.context_parallel_cache: 

92 raise ValueError("prefix_sequence_shard_info requires context_parallel_cache=True") 

93 

94 def _validate_context_parallel(self) -> None: 

95 """Validate context-parallel cache and logits metadata.""" 

96 if self.context_parallel_cache and not self.use_cache: 

97 raise ValueError("context_parallel_cache requires use_cache=True") 

98 if (self.context_parallel_rank is None) != (self.context_parallel_world_size is None): 

99 raise ValueError( 

100 "context_parallel_rank and context_parallel_world_size must be set together", 

101 ) 

102 if self.context_parallel_world_size is not None: 

103 if self.context_parallel_world_size <= 0: 

104 raise ValueError("context_parallel_world_size must be > 0") 

105 if ( 

106 self.context_parallel_rank < 0 

107 or self.context_parallel_rank >= self.context_parallel_world_size 

108 ): 

109 raise ValueError("context_parallel_rank must be in [0, context_parallel_world_size)") 

110 

111 @staticmethod 

112 def _validate_callables(values: Optional[List[Callable]], field_name: str) -> None: 

113 """Validate optional generation extension hooks.""" 

114 if values is None: 

115 return 

116 if not isinstance(values, list): 

117 raise ValueError(f"{field_name} must be a list of callables") 

118 if not all(callable(item) for item in values): 

119 raise ValueError(f"{field_name} must contain only callables") 

120 

121 

122def build_position_ids( 

123 input_ids: torch.Tensor, 

124 attention_mask: Optional[torch.Tensor] = None, 

125) -> torch.Tensor: 

126 """Build left-padding aware position ids.""" 

127 if input_ids.ndim != 2: 

128 raise ValueError("input_ids must have shape (batch, seq)") 

129 if attention_mask is None: 

130 seq_len = input_ids.size(1) 

131 return torch.arange( 

132 seq_len, device=input_ids.device, dtype=torch.long, 

133 ).view(1, -1).expand(input_ids.size(0), -1) 

134 if attention_mask.shape != input_ids.shape: 

135 raise ValueError("attention_mask must match input_ids shape") 

136 position_ids = attention_mask.long().cumsum(dim=-1) - 1 

137 return position_ids.clamp_min_(0) 

138 

139 

140def gather_context_parallel_logits(logits: torch.Tensor, config: GenerationConfig) -> torch.Tensor: 

141 """Select final-token logits from the owning CP rank before sampling.""" 

142 if config.context_logits_rank is None: 

143 return logits 

144 if not dist.is_available() or not dist.is_initialized(): 

145 return logits 

146 world_size = dist.get_world_size(group=config.context_process_group) 

147 if world_size == 1: 

148 return logits 

149 gathered = [torch.empty_like(logits) for _ in range(world_size)] 

150 dist.all_gather(gathered, logits, group=config.context_process_group) 

151 stacked = torch.stack(gathered, dim=0) 

152 owner = torch.as_tensor( 

153 config.context_logits_rank, 

154 device=logits.device, 

155 dtype=torch.long, 

156 ) 

157 if owner.ndim == 0: 

158 owner_rank = int(owner.item()) 

159 if owner_rank < 0 or owner_rank >= world_size: 

160 raise ValueError("context_logits_rank contains an invalid rank") 

161 return stacked[owner_rank] 

162 if owner.shape != (logits.shape[0],): 

163 raise ValueError("context_logits_rank must be a scalar or a batch-sized tensor") 

164 if torch.any((owner < 0) | (owner >= world_size)): 

165 raise ValueError("context_logits_rank contains an invalid rank") 

166 batch_indices = torch.arange(logits.shape[0], device=logits.device) 

167 return stacked[owner, batch_indices] 

168 

169 

170def gather_tensor_parallel_logits(logits: torch.Tensor, config: GenerationConfig) -> torch.Tensor: 

171 """Gather vocab-sharded logits before sampling when TP inference is active.""" 

172 if not config.gather_logits: 

173 return logits 

174 if not dist.is_available() or not dist.is_initialized(): 

175 return logits 

176 world_size = dist.get_world_size(group=config.logits_process_group) 

177 if world_size == 1: 

178 return logits 

179 gather_dim = logits.ndim + config.logits_gather_dim 

180 if gather_dim < 0 or gather_dim >= logits.ndim: 

181 raise ValueError("logits_gather_dim is out of range for logits") 

182 local_shard = torch.tensor( 

183 [logits.shape[gather_dim]], 

184 device=logits.device, 

185 dtype=torch.long, 

186 ) 

187 shard_sizes = [torch.empty_like(local_shard) for _ in range(world_size)] 

188 dist.all_gather(shard_sizes, local_shard, group=config.logits_process_group) 

189 shard_sizes = torch.cat(shard_sizes) 

190 if torch.any(shard_sizes != shard_sizes[0]): 

191 raise ValueError( 

192 "tensor-parallel logits gather requires equal local vocab shard sizes; " 

193 "pad vocab shards before generation", 

194 ) 

195 gathered = [torch.empty_like(logits) for _ in range(world_size)] 

196 dist.all_gather(gathered, logits, group=config.logits_process_group) 

197 return torch.cat(gathered, dim=config.logits_gather_dim) 

198 

199 

200def prepare_logits_for_sampling(logits: torch.Tensor, config: GenerationConfig) -> torch.Tensor: 

201 """Apply distributed logits handoffs before sampling.""" 

202 logits = gather_context_parallel_logits(logits, config) 

203 return gather_tensor_parallel_logits(logits, config) 

204 

205 

206def apply_logits_processors( 

207 input_ids: torch.Tensor, 

208 logits: torch.Tensor, 

209 config: GenerationConfig, 

210) -> torch.Tensor: 

211 """Apply user-supplied logits processors in order.""" 

212 if config.logits_processor is None: 

213 return logits 

214 processed = logits 

215 for processor in config.logits_processor: 

216 processed = processor(input_ids, processed) 

217 if not isinstance(processed, torch.Tensor): 

218 raise ValueError("logits_processor must return a tensor") 

219 return processed 

220 

221 

222def should_stop_generation( 

223 input_ids: torch.Tensor, 

224 logits: torch.Tensor, 

225 config: GenerationConfig, 

226) -> bool: 

227 """Return whether any configured stopping criterion requests termination.""" 

228 if config.stopping_criteria is None: 

229 return False 

230 for criterion in config.stopping_criteria: 

231 result = criterion(input_ids, logits) 

232 if isinstance(result, torch.Tensor): 

233 if result.numel() != 1: 

234 raise ValueError("stopping_criteria tensor output must be scalar") 

235 result = bool(result.item()) 

236 if bool(result): 

237 return True 

238 return False 

239 

240 

241def build_causal_mask( 

242 input_ids: torch.Tensor, 

243 attention_mask: Optional[torch.Tensor] = None, 

244 dtype: torch.dtype = torch.float32, 

245) -> torch.Tensor: 

246 """Build additive causal + padding mask for prefill.""" 

247 if input_ids.ndim != 2: 

248 raise ValueError("input_ids must have shape (batch, seq)") 

249 batch_size, seq_len = input_ids.shape 

250 device = input_ids.device 

251 mask = torch.zeros( 

252 batch_size, 1, seq_len, seq_len, device=device, dtype=dtype, 

253 ) 

254 causal = torch.triu( 

255 torch.full((seq_len, seq_len), float("-inf"), device=device, dtype=dtype), 

256 diagonal=1, 

257 ) 

258 mask = mask + causal.view(1, 1, seq_len, seq_len) 

259 if attention_mask is not None: 

260 if attention_mask.shape != input_ids.shape: 

261 raise ValueError("attention_mask must match input_ids shape") 

262 padding = attention_mask.to(device=device) == 0 

263 mask = mask.masked_fill(padding.view(batch_size, 1, 1, seq_len), float("-inf")) 

264 mask = mask.masked_fill(padding.view(batch_size, 1, seq_len, 1), 0.0) 

265 return mask 

266 

267 

268def append_attention_mask( 

269 attention_mask: Optional[torch.Tensor], 

270 next_tokens: torch.Tensor, 

271) -> Optional[torch.Tensor]: 

272 """Append valid-token mask entries for generated tokens.""" 

273 if attention_mask is None: 

274 return None 

275 ones = torch.ones( 

276 next_tokens.shape, 

277 device=attention_mask.device, 

278 dtype=attention_mask.dtype, 

279 ) 

280 return torch.cat([attention_mask, ones], dim=-1)