Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / infer / generation.py: 0%
278 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-07-06 05:41 +0800
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"""Prefill + decode generation loop."""
16import inspect
17from typing import Optional
19import torch
20import torch.distributed as dist
22from hyper_parallel.infer.kv_cache import (
23 ContextParallelKVCache,
24 KVCache,
25 detach_and_validate_past_key_values,
26)
27from hyper_parallel.infer.sampler import sample_next_token
28from hyper_parallel.infer.utils import (
29 GenerationConfig,
30 append_attention_mask,
31 apply_logits_processors,
32 build_causal_mask,
33 build_position_ids,
34 prepare_logits_for_sampling,
35 should_stop_generation,
36)
39def _get_output(outputs, name: str):
40 """Read an output field from dict-like or object-like model outputs."""
41 if isinstance(outputs, dict):
42 return outputs.get(name)
43 return getattr(outputs, name, None)
46def _model_forward(
47 model,
48 input_ids: torch.Tensor,
49 position_ids: torch.Tensor,
50 attention_mask: Optional[torch.Tensor],
51 past_key_values,
52 use_cache: bool,
53 sequence_shard_info=None,
54 global_seq_len: Optional[int] = None,
55):
56 """Call model.forward with only the keyword arguments it accepts."""
57 kwargs = {
58 "input_ids": input_ids,
59 "position_ids": position_ids,
60 "attention_mask": attention_mask,
61 "past_key_values": past_key_values,
62 "use_cache": use_cache,
63 "sequence_shard_info": sequence_shard_info,
64 "global_seq_len": global_seq_len,
65 }
66 forward = getattr(model, "forward", model)
67 try:
68 signature = inspect.signature(forward)
69 except (TypeError, ValueError):
70 return forward(**kwargs)
71 parameters = signature.parameters
72 accepts_kwargs = any(
73 param.kind == inspect.Parameter.VAR_KEYWORD
74 for param in parameters.values()
75 )
76 if not accepts_kwargs:
77 for name in list(kwargs):
78 if name not in parameters:
79 kwargs.pop(name)
80 return forward(**kwargs)
83def _resolve_context_parallel_rank_world(config: GenerationConfig) -> tuple[int, int]:
84 if config.context_parallel_rank is not None:
85 return config.context_parallel_rank, config.context_parallel_world_size
86 if not dist.is_available() or not dist.is_initialized():
87 raise ValueError(
88 "context_parallel_cache requires initialized torch.distributed "
89 "or explicit context_parallel_rank/context_parallel_world_size",
90 )
91 return (
92 dist.get_rank(group=config.context_process_group),
93 dist.get_world_size(group=config.context_process_group),
94 )
97def _init_cache(config: GenerationConfig) -> KVCache:
98 if not config.context_parallel_cache:
99 return KVCache()
100 rank, world_size = _resolve_context_parallel_rank_world(config)
101 return ContextParallelKVCache(rank=rank, world_size=world_size)
104def _cache_shard_info(cache: KVCache):
105 return cache.shard_info if isinstance(cache, ContextParallelKVCache) else None
108def _cache_seq_len(past_key_values) -> Optional[int]:
109 """Resolve cached sequence length from tuple or opaque HF-style cache."""
110 if past_key_values is None:
111 return None
112 if hasattr(past_key_values, "get_seq_length") and not isinstance(
113 past_key_values, (list, tuple),
114 ):
115 return int(past_key_values.get_seq_length())
116 values = detach_and_validate_past_key_values(past_key_values)
117 if not values:
118 return 0
119 return int(values[0][0].shape[-2])
122def _cache_batch_size(past_key_values) -> Optional[int]:
123 """Resolve cache batch size when cache tensors are inspectable."""
124 if past_key_values is None:
125 return None
126 if hasattr(past_key_values, "get_seq_length") and not isinstance(
127 past_key_values, (list, tuple),
128 ):
129 return None
130 values = detach_and_validate_past_key_values(past_key_values)
131 if not values:
132 return None
133 return int(values[0][0].shape[0])
136def _resolve_prefix_length(config: GenerationConfig) -> int:
137 """Validate and resolve reusable prefix cache length."""
138 if config.prefix_past_key_values is None:
139 return 0
140 candidates = []
141 if config.prefix_cache_length is not None:
142 candidates.append(int(config.prefix_cache_length))
143 if config.prefix_attention_mask is not None:
144 candidates.append(int(config.prefix_attention_mask.shape[-1]))
145 if config.prefix_sequence_shard_info is not None:
146 candidates.append(int(config.prefix_sequence_shard_info.global_seq_len))
147 seq_len = _cache_seq_len(config.prefix_past_key_values)
148 if seq_len is not None and config.prefix_sequence_shard_info is None:
149 candidates.append(seq_len)
150 if not candidates:
151 raise ValueError(
152 "prefix_past_key_values requires prefix_cache_length for opaque caches",
153 )
154 prefix_len = candidates[0]
155 if any(length != prefix_len for length in candidates):
156 raise ValueError("prefix cache length metadata is inconsistent")
157 return prefix_len
160def _prepare_prefix_attention_mask(
161 config: GenerationConfig,
162 input_ids: torch.Tensor,
163 device,
164) -> tuple[Optional[torch.Tensor], int]:
165 """Prepare a 2-D attention mask for reusable prefix cache."""
166 prefix_len = _resolve_prefix_length(config)
167 if prefix_len == 0:
168 return None, 0
169 cache_batch_size = _cache_batch_size(config.prefix_past_key_values)
170 if cache_batch_size is not None and cache_batch_size != input_ids.size(0):
171 raise ValueError("prefix cache batch size must match input_ids batch size")
172 prefix_attention_mask = config.prefix_attention_mask
173 if prefix_attention_mask is None:
174 return torch.ones(
175 input_ids.size(0),
176 prefix_len,
177 device=device,
178 dtype=torch.long,
179 ), prefix_len
180 if prefix_attention_mask.ndim != 2:
181 raise ValueError("prefix_attention_mask must have shape (batch, prefix_seq)")
182 if prefix_attention_mask.shape != (input_ids.size(0), prefix_len):
183 raise ValueError("prefix_attention_mask batch/sequence length mismatch")
184 return prefix_attention_mask.to(device=device), prefix_len
187def _init_cache_with_prefix(config: GenerationConfig) -> tuple[KVCache, int]:
188 """Create the generation cache and preload prefix cache when present."""
189 cache = _init_cache(config)
190 prefix_len = _resolve_prefix_length(config)
191 if prefix_len == 0:
192 return cache, prefix_len
193 if isinstance(cache, ContextParallelKVCache):
194 if config.prefix_sequence_shard_info is None:
195 cache.update_full(config.prefix_past_key_values)
196 else:
197 cache.update_local(
198 config.prefix_past_key_values,
199 config.prefix_sequence_shard_info,
200 )
201 return cache, prefix_len
202 cache.update(config.prefix_past_key_values)
203 return cache, prefix_len
206def _update_cache(cache: KVCache, outputs) -> None:
207 """Update normal or context-parallel KV cache from model outputs."""
208 past_key_values = _get_output(outputs, "past_key_values")
209 if not isinstance(cache, ContextParallelKVCache):
210 cache.update(past_key_values)
211 return
212 if past_key_values is None:
213 return
214 sequence_shard_info = _get_output(outputs, "sequence_shard_info")
215 if sequence_shard_info is not None:
216 cache.update_local(past_key_values, sequence_shard_info)
217 return
218 if cache.is_empty:
219 cache.update_full(past_key_values)
220 return
221 raise ValueError(
222 "context-parallel cached decode requires model output sequence_shard_info",
223 )
226def _resolve_mask_dtype(model, config: GenerationConfig) -> torch.dtype:
227 """Choose additive-mask dtype from config or model floating state."""
228 if config.mask_dtype is not None:
229 return config.mask_dtype
230 for iterator_name in ("parameters", "buffers"):
231 iterator = getattr(model, iterator_name, None)
232 if iterator is None:
233 continue
234 for tensor in iterator():
235 if tensor.is_floating_point():
236 return tensor.dtype
237 return torch.float32
240def _build_decode_key_mask(
241 attention_mask: Optional[torch.Tensor],
242 dtype: torch.dtype,
243) -> Optional[torch.Tensor]:
244 """Build additive key padding mask for one-token cached decode."""
245 if attention_mask is None:
246 return None
247 if attention_mask.ndim != 2:
248 raise ValueError("attention_mask must have shape (batch, seq)")
249 batch_size, seq_len = attention_mask.shape
250 mask = torch.zeros(
251 batch_size,
252 1,
253 1,
254 seq_len,
255 device=attention_mask.device,
256 dtype=dtype,
257 )
258 padding = attention_mask == 0
259 return mask.masked_fill(padding.view(batch_size, 1, 1, seq_len), float("-inf"))
262def _combined_attention_mask(
263 prefix_attention_mask: Optional[torch.Tensor],
264 attention_mask: Optional[torch.Tensor],
265) -> Optional[torch.Tensor]:
266 if prefix_attention_mask is None:
267 return attention_mask
268 if attention_mask is None:
269 return prefix_attention_mask
270 return torch.cat([prefix_attention_mask, attention_mask], dim=-1)
273def _build_prefill_mask(
274 input_ids: torch.Tensor,
275 attention_mask: Optional[torch.Tensor],
276 prefix_attention_mask: Optional[torch.Tensor],
277 dtype: torch.dtype,
278) -> torch.Tensor:
279 """Build the prefill causal mask, including optional prefix keys."""
280 if prefix_attention_mask is None:
281 return build_causal_mask(input_ids, attention_mask, dtype=dtype)
282 batch_size, query_len = input_ids.shape
283 prefix_len = prefix_attention_mask.shape[-1]
284 device = input_ids.device
285 if attention_mask is None:
286 attention_mask = torch.ones(
287 batch_size,
288 query_len,
289 device=device,
290 dtype=prefix_attention_mask.dtype,
291 )
292 if attention_mask.shape != input_ids.shape:
293 raise ValueError("attention_mask must match input_ids shape")
294 mask = torch.zeros(
295 batch_size,
296 1,
297 query_len,
298 prefix_len + query_len,
299 device=device,
300 dtype=dtype,
301 )
302 causal = torch.triu(
303 torch.full((query_len, query_len), float("-inf"), device=device, dtype=dtype),
304 diagonal=1,
305 )
306 mask[:, :, :, prefix_len:] = causal.view(1, 1, query_len, query_len)
307 current_padding = attention_mask == 0
308 current_key_padding = torch.cat([prefix_attention_mask == 0, current_padding], dim=-1)
309 mask = mask.masked_fill(
310 current_key_padding.view(batch_size, 1, 1, prefix_len + query_len),
311 float("-inf"),
312 )
313 mask = mask.masked_fill(current_padding.view(batch_size, 1, query_len, 1), 0.0)
314 return mask
317def _build_prefill_position_ids(
318 input_ids: torch.Tensor,
319 attention_mask: Optional[torch.Tensor],
320 prefix_attention_mask: Optional[torch.Tensor],
321) -> torch.Tensor:
322 """Build position ids for prefill with optional prefix offset."""
323 position_ids = build_position_ids(input_ids, attention_mask)
324 if prefix_attention_mask is None:
325 return position_ids
326 prefix_lengths = prefix_attention_mask.long().sum(dim=-1).view(-1, 1)
327 return position_ids + prefix_lengths
330def _prompt_lengths(input_ids: torch.Tensor, attention_mask: Optional[torch.Tensor]):
331 """Count valid prompt tokens per batch row."""
332 if attention_mask is None:
333 return torch.full(
334 (input_ids.size(0),),
335 input_ids.size(1),
336 device=input_ids.device,
337 dtype=torch.long,
338 )
339 return attention_mask.long().sum(dim=-1)
342def _finalize_sequences(
343 sequences: torch.Tensor,
344 initial_attention_mask: Optional[torch.Tensor],
345 prompt_lengths: torch.Tensor,
346 generated_counts: torch.Tensor,
347 pad_token_id: int,
348) -> torch.Tensor:
349 """Strip left padding and right-pad finalized generated sequences."""
350 rows = []
351 max_len = 0
352 for batch_idx in range(sequences.size(0)):
353 if initial_attention_mask is None:
354 start = 0
355 else:
356 starts = torch.nonzero(
357 initial_attention_mask[batch_idx].bool(), as_tuple=False,
358 )
359 if starts.numel() == 0:
360 raise ValueError("attention_mask row must contain at least one valid token")
361 start = int(starts[0].item())
362 total_len = int(prompt_lengths[batch_idx].item() + generated_counts[batch_idx].item())
363 row = sequences[batch_idx, start:start + total_len]
364 rows.append(row)
365 max_len = max(max_len, row.numel())
366 output = sequences.new_full((len(rows), max_len), pad_token_id)
367 for idx, row in enumerate(rows):
368 output[idx, :row.numel()] = row
369 return output
372def _validate_generate_inputs(
373 input_ids: torch.Tensor,
374 attention_mask: Optional[torch.Tensor],
375) -> None:
376 if input_ids.ndim != 2:
377 raise ValueError("input_ids must have shape (batch, seq)")
378 if attention_mask is not None and attention_mask.shape != input_ids.shape:
379 raise ValueError("attention_mask must match input_ids shape")
380 if attention_mask is not None and torch.any(attention_mask.long().sum(dim=-1) == 0):
381 raise ValueError("attention_mask rows must contain at least one valid token")
384def _finalize_zero_new_tokens(
385 input_ids: torch.Tensor,
386 attention_mask: Optional[torch.Tensor],
387 config: GenerationConfig,
388) -> torch.Tensor:
389 """Finalize left-padded prompts when no new tokens are requested."""
390 prompt_lengths = _prompt_lengths(input_ids, attention_mask)
391 generated_counts = torch.zeros(
392 input_ids.size(0),
393 device=input_ids.device,
394 dtype=torch.long,
395 )
396 return _finalize_sequences(
397 input_ids.clone(),
398 initial_attention_mask=attention_mask,
399 prompt_lengths=prompt_lengths,
400 generated_counts=generated_counts,
401 pad_token_id=config.pad_token_id,
402 )
405def _prepare_generation_context(
406 model,
407 input_ids: torch.Tensor,
408 attention_mask: Optional[torch.Tensor],
409 config: GenerationConfig,
410):
411 """Create the mutable generation context used by the decode loop."""
412 mask_dtype = _resolve_mask_dtype(model, config)
413 sequences = input_ids.clone()
414 prefix_attention_mask, prefix_len = _prepare_prefix_attention_mask(
415 config,
416 input_ids,
417 input_ids.device,
418 )
419 current_attention_mask = attention_mask.clone() if attention_mask is not None else None
420 if prefix_attention_mask is not None and current_attention_mask is None:
421 current_attention_mask = torch.ones_like(sequences, dtype=torch.long)
422 prompt_lengths = _prompt_lengths(input_ids, current_attention_mask)
423 prefix_valid_lengths = (
424 prefix_attention_mask.long().sum(dim=-1)
425 if prefix_attention_mask is not None
426 else torch.zeros(input_ids.size(0), device=input_ids.device, dtype=torch.long)
427 )
428 cache, prefix_len = _init_cache_with_prefix(config)
429 return {
430 "mask_dtype": mask_dtype,
431 "sequences": sequences,
432 "prefix_attention_mask": prefix_attention_mask,
433 "current_attention_mask": current_attention_mask,
434 "initial_attention_mask": (
435 current_attention_mask.clone()
436 if current_attention_mask is not None
437 else None
438 ),
439 "prompt_lengths": prompt_lengths,
440 "generated_counts": torch.zeros(
441 input_ids.size(0), device=input_ids.device, dtype=torch.long,
442 ),
443 "unfinished": torch.ones(
444 input_ids.size(0), device=input_ids.device, dtype=torch.bool,
445 ),
446 "prefix_valid_lengths": prefix_valid_lengths,
447 "cache": cache,
448 "prefix_len": prefix_len,
449 }
452def _prefill(
453 model,
454 config: GenerationConfig,
455 context: dict,
456):
457 """Run the initial full-prompt forward pass."""
458 position_ids = _build_prefill_position_ids(
459 context["sequences"],
460 context["current_attention_mask"],
461 context["prefix_attention_mask"],
462 )
463 attention_mask = _build_prefill_mask(
464 context["sequences"],
465 context["current_attention_mask"],
466 context["prefix_attention_mask"],
467 dtype=context["mask_dtype"],
468 )
469 cache = context["cache"]
470 return _model_forward(
471 model,
472 input_ids=context["sequences"],
473 position_ids=position_ids,
474 attention_mask=attention_mask,
475 past_key_values=None if cache.is_empty else cache.past_key_values,
476 use_cache=config.use_cache,
477 sequence_shard_info=_cache_shard_info(cache),
478 global_seq_len=context["prefix_len"] + context["sequences"].shape[-1],
479 )
482def _required_logits(outputs) -> torch.Tensor:
483 logits = _get_output(outputs, "logits")
484 if logits is None:
485 raise ValueError("model output must contain logits")
486 return logits
489def _finalize_prefill_outputs(
490 config: GenerationConfig,
491 context: dict,
492 outputs,
493) -> tuple[torch.Tensor, bool]:
494 """Validate prefill output and decide whether cached decode can be used."""
495 logits = _required_logits(outputs)
496 if (
497 config.prefix_past_key_values is not None
498 and _get_output(outputs, "past_key_values") is None
499 ):
500 raise ValueError("prefix_past_key_values requires model to return past_key_values")
501 _update_cache(context["cache"], outputs)
502 return logits, config.use_cache and not context["cache"].is_empty
505def _append_next_token(context: dict, next_tokens: torch.Tensor, config: GenerationConfig):
506 """Append sampled tokens and advance per-row generation metadata."""
507 if config.eos_token_id is not None:
508 next_tokens = torch.where(
509 context["unfinished"].view(-1, 1),
510 next_tokens,
511 torch.full_like(next_tokens, config.pad_token_id),
512 )
513 context["sequences"] = torch.cat([context["sequences"], next_tokens], dim=-1)
514 context["generated_counts"] = (
515 context["generated_counts"] + context["unfinished"].long()
516 )
517 context["current_attention_mask"] = append_attention_mask(
518 context["current_attention_mask"],
519 next_tokens,
520 )
521 if config.eos_token_id is not None:
522 context["unfinished"] = (
523 context["unfinished"] & (next_tokens.squeeze(-1) != config.eos_token_id)
524 )
525 return next_tokens
528def _should_finish_generation(
529 context: dict,
530 logits: torch.Tensor,
531 config: GenerationConfig,
532 step: int,
533) -> bool:
534 """Check EOS, custom stopping criteria, and max token limit."""
535 if config.eos_token_id is not None and not context["unfinished"].any():
536 return True
537 if should_stop_generation(context["sequences"], logits, config):
538 return True
539 return step == config.max_new_tokens - 1
542def _decode(
543 model,
544 context: dict,
545 next_tokens: torch.Tensor,
546 use_cached_decode: bool,
547):
548 """Run one cached or no-cache decode step."""
549 model_attention_mask = _combined_attention_mask(
550 context["prefix_attention_mask"],
551 context["current_attention_mask"],
552 )
553 if use_cached_decode:
554 decode_pos = (
555 context["prefix_valid_lengths"]
556 + context["prompt_lengths"]
557 + context["generated_counts"]
558 - 1
559 )
560 return _model_forward(
561 model,
562 input_ids=next_tokens,
563 position_ids=decode_pos.view(-1, 1),
564 attention_mask=_build_decode_key_mask(
565 model_attention_mask,
566 context["mask_dtype"],
567 ),
568 past_key_values=context["cache"].past_key_values,
569 use_cache=True,
570 sequence_shard_info=_cache_shard_info(context["cache"]),
571 global_seq_len=context["prefix_len"] + context["sequences"].shape[-1],
572 )
573 decode_pos = _build_prefill_position_ids(
574 context["sequences"],
575 context["current_attention_mask"],
576 context["prefix_attention_mask"],
577 )
578 decode_mask = _build_prefill_mask(
579 context["sequences"],
580 context["current_attention_mask"],
581 context["prefix_attention_mask"],
582 dtype=context["mask_dtype"],
583 )
584 return _model_forward(
585 model,
586 input_ids=context["sequences"],
587 position_ids=decode_pos,
588 attention_mask=decode_mask,
589 past_key_values=None,
590 use_cache=False,
591 sequence_shard_info=_cache_shard_info(context["cache"]),
592 global_seq_len=context["prefix_len"] + context["sequences"].shape[-1],
593 )
596@torch.no_grad()
597def generate(
598 model,
599 input_ids: torch.Tensor,
600 generation_config: Optional[GenerationConfig] = None,
601 attention_mask: Optional[torch.Tensor] = None,
602 **kwargs,
603) -> torch.Tensor:
604 """Generate token ids from a causal language model."""
605 if kwargs:
606 raise TypeError(f"Unexpected generate kwargs: {sorted(kwargs)}")
607 config = generation_config or GenerationConfig()
608 _validate_generate_inputs(input_ids, attention_mask)
609 if config.max_new_tokens == 0:
610 return _finalize_zero_new_tokens(input_ids, attention_mask, config)
612 was_training = getattr(model, "training", False)
613 model.eval()
614 try:
615 context = _prepare_generation_context(model, input_ids, attention_mask, config)
616 outputs = _prefill(model, config, context)
617 logits, use_cached_decode = _finalize_prefill_outputs(
618 config,
619 context,
620 outputs,
621 )
623 for step in range(config.max_new_tokens):
624 next_logits = prepare_logits_for_sampling(logits[:, -1, :], config)
625 next_logits = apply_logits_processors(
626 context["sequences"],
627 next_logits,
628 config,
629 )
630 next_tokens = sample_next_token(next_logits, context["sequences"], config)
631 next_tokens = _append_next_token(context, next_tokens, config)
632 if _should_finish_generation(context, next_logits, config, step):
633 break
635 outputs = _decode(model, context, next_tokens, use_cached_decode)
636 if use_cached_decode:
637 _update_cache(context["cache"], outputs)
638 logits = _required_logits(outputs)
640 return _finalize_sequences(
641 context["sequences"],
642 initial_attention_mask=context["initial_attention_mask"],
643 prompt_lengths=context["prompt_lengths"],
644 generated_counts=context["generated_counts"],
645 pad_token_id=config.pad_token_id,
646 )
647 finally:
648 if was_training:
649 model.train()