import math import gc from collections import deque from typing import Callable, List, Optional import torch from torch import nn from tqdm import tqdm from .attention import ( ForwardContext, get_forward_context, reset_forward_context, set_forward_context, ) from .kv_manager import KVCacheManager, Seq class Sampler(nn.Module): def __init__(self): super().__init__() # @torch.compile def forward(self, logits: torch.Tensor, temperatures: torch.Tensor): temperatures = temperatures.to(logits.device).clamp(min=1e-8) greedy_mask = temperatures < 1e-5 temp_for_scaling = torch.where(greedy_mask, 1.0, temperatures) scaled_logits = logits / temp_for_scaling.unsqueeze(-1) probs = torch.softmax(scaled_logits, dim=-1, dtype=torch.float32) q = torch.empty_like(probs) q.exponential_() sampled_tokens = probs.div_(q).argmax(dim=-1) greedy_tokens = logits.argmax(dim=-1) return torch.where(greedy_mask, greedy_tokens, sampled_tokens) class AccelInferenceEngine: def __init__( self, model, lm_head, num_layers: int, num_heads: int, head_dim: int, block_size: int = 256, num_blocks: int = 128, use_cuda_graph: bool = True, ): """ Args: model: The GPT transformer model (should have accel attention) lm_head: Language model head for generating logits num_layers: Number of transformer layers num_heads: Number of attention heads head_dim: Dimension per head block_size: KV cache block size num_blocks: Total number of KV cache blocks use_cuda_graph: Whether to use CUDA Graph for decode optimization """ self.model = model self.lm_head = lm_head self.block_size = block_size self.num_layers = num_layers self.num_heads = num_heads self.head_dim = head_dim self._default_num_blocks = max(1, int(num_blocks)) self.num_blocks = 1 self.use_cuda_graph = use_cuda_graph and torch.cuda.is_available() self.hidden_size = ( model.config.hidden_size if hasattr(model, "config") else head_dim * num_heads ) self.kv_manager = self._new_kv_manager(self.num_blocks, device=torch.device("cpu")) self.kv_manager.wire_kv_cache_to_model(model) self.sampler = Sampler() self.current_sequences = [] self.graphs = {} self.graph_vars = None self.graph_pool = None self.graph_captured = False self.graph_num_blocks = 0 self.graph_signature = None def _runtime_kv_device(self) -> torch.device: first_param = next(self.model.parameters(), None) if first_param is not None and first_param.device.type == "cuda": return first_param.device if torch.cuda.is_available(): return torch.device(f"cuda:{torch.cuda.current_device()}") if first_param is not None: return first_param.device return torch.device("cpu") def _new_kv_manager(self, num_blocks: int, device: torch.device) -> KVCacheManager: return KVCacheManager( num_layers=self.num_layers, num_heads=self.num_heads, head_dim=self.head_dim, block_size=self.block_size, num_blocks=int(max(1, num_blocks)), dtype=torch.float16, # Force fp16 for FlashAttention device=device, ) @staticmethod def _filter_logits(logits: torch.Tensor, top_k: int, top_p: float) -> torch.Tensor: filtered = logits vocab_size = int(filtered.size(-1)) top_k = int(top_k) if top_k is not None else 0 top_p = float(top_p) if top_p is not None else 1.0 if top_k > 0 and top_k < vocab_size: kth = torch.topk(filtered, k=top_k, dim=-1).values[..., -1, None] filtered = filtered.masked_fill(filtered < kth, float("-inf")) if top_p < 1.0: sorted_logits, sorted_indices = torch.sort(filtered, descending=True, dim=-1) sorted_probs = torch.softmax(sorted_logits.float(), dim=-1) cumulative_probs = torch.cumsum(sorted_probs, dim=-1) sorted_remove = cumulative_probs > top_p sorted_remove[..., 0] = False remove_mask = torch.zeros_like(sorted_remove, dtype=torch.bool) remove_mask.scatter_(dim=-1, index=sorted_indices, src=sorted_remove) filtered = filtered.masked_fill(remove_mask, float("-inf")) return filtered @staticmethod def _tensor_sig(tensor: Optional[torch.Tensor]): if tensor is None: return None return ( int(tensor.data_ptr()), str(tensor.dtype), tuple(int(x) for x in tensor.shape), int(tensor.device.index if tensor.device.index is not None else -1), ) def _module_first_param_sig(self, module: Optional[torch.nn.Module]): if module is None: return None try: first_param = next(module.parameters()) except StopIteration: return None return self._tensor_sig(first_param) def _make_capture_signature( self, tts_mel_embedding: Optional[torch.nn.Module] = None, tts_text_pos_embedding: Optional[torch.nn.Module] = None, ): mel_weight = getattr(tts_mel_embedding, "weight", None) if tts_mel_embedding is not None else None pos_emb = None if tts_text_pos_embedding is not None: pos_emb = getattr(tts_text_pos_embedding, "emb", tts_text_pos_embedding) pos_weight = getattr(pos_emb, "weight", None) if pos_emb is not None else None return ( self._module_first_param_sig(self.model), self._tensor_sig(self.kv_manager.kv_cache), self._tensor_sig(mel_weight), self._tensor_sig(pos_weight), int(self.num_blocks), int(self.block_size), ) def _compute_tts_embeds( self, input_ids: torch.Tensor, positions: torch.Tensor, tts_mel_embedding: Optional[torch.nn.Module] = None, tts_text_pos_embedding: Optional[torch.nn.Module] = None, ) -> torch.Tensor: if tts_mel_embedding is None or tts_text_pos_embedding is None: raise RuntimeError("TTS embedding modules are required for accel decode.") pos_emb_module = getattr(tts_text_pos_embedding, "emb", tts_text_pos_embedding) if not hasattr(pos_emb_module, "weight"): raise RuntimeError("TTS positional embedding module is missing a '.weight' tensor.") pos_clamped = torch.clamp(positions, min=0, max=pos_emb_module.weight.shape[0] - 1) mel_emb = tts_mel_embedding(input_ids) pos_emb = pos_emb_module(pos_clamped) return mel_emb + pos_emb def _required_blocks(self, total_tokens: int) -> int: # Keep one spare block to avoid edge overflows from token-length drift. return max(1, int(math.ceil(float(max(1, total_tokens)) / float(self.block_size))) + 1) def _reset_decode_graph(self): self.graphs = {} self.graph_vars = None self.graph_pool = None self.graph_captured = False self.graph_num_blocks = 0 self.graph_signature = None def _resize_kv_cache_if_needed(self, required_blocks: int): target_blocks = int(max(1, required_blocks)) target_device = self._runtime_kv_device() current_device = self.kv_manager.kv_cache.device same_device = ( current_device.type == target_device.type and (current_device.type != "cuda" or current_device.index == target_device.index) ) if target_blocks == int(self.num_blocks) and same_device: return self.num_blocks = target_blocks self.kv_manager = self._new_kv_manager(self.num_blocks, target_device) self.kv_manager.wire_kv_cache_to_model(self.model) self._reset_decode_graph() def release_runtime_cache(self): self.current_sequences = [] reset_forward_context() self._reset_decode_graph() old_kv_manager = self.kv_manager self.num_blocks = 1 self.kv_manager = self._new_kv_manager(self.num_blocks, torch.device("cpu")) self.kv_manager.wire_kv_cache_to_model(self.model) del old_kv_manager if torch.cuda.is_available(): torch.cuda.synchronize() torch.cuda.empty_cache() if hasattr(torch.cuda, "ipc_collect"): torch.cuda.ipc_collect() gc.collect() def prepare_decode_graph( self, max_total_tokens: int, tts_mel_embedding: Optional[torch.nn.Module] = None, tts_text_pos_embedding: Optional[torch.nn.Module] = None, ): required_blocks = self._required_blocks(int(max_total_tokens)) self._resize_kv_cache_if_needed(required_blocks) if not self.use_cuda_graph: return signature = self._make_capture_signature( tts_mel_embedding=tts_mel_embedding, tts_text_pos_embedding=tts_text_pos_embedding, ) if ( self.graph_captured and int(self.graph_num_blocks) == int(self.num_blocks) and self.graph_num_blocks >= required_blocks and self.graph_signature == signature ): return print( f"[CAPTURE] use_cuda_graph={self.use_cuda_graph}, graph_captured={self.graph_captured}, " f"graph_num_blocks={self.graph_num_blocks}, required_blocks={required_blocks}, cache_blocks={self.num_blocks}", flush=True, ) self._reset_decode_graph() self._capture_cuda_graphs( tts_mel_embedding=tts_mel_embedding, tts_text_pos_embedding=tts_text_pos_embedding, max_num_blocks=self.num_blocks, ) self.graph_captured = True self.graph_num_blocks = int(self.num_blocks) self.graph_signature = signature print(f"[CAPTURE] Completed! graphs={list(self.graphs.keys())}, num_blocks={self.graph_num_blocks}", flush=True) def _prepare_prefill(self, requests: List[Seq]): input_ids = [] positions = [] cu_seqlens_q = [0] cu_seqlens_k = [0] max_seqlen_q = 0 max_seqlen_k = 0 slot_mapping = [] for req in requests: seqlen = len(req) input_ids.extend(req[req.num_cached_tokens :]) positions.extend(list(range(req.num_cached_tokens, seqlen))) seqlen_q = seqlen - req.num_cached_tokens seqlen_k = seqlen cu_seqlens_q.append(cu_seqlens_q[-1] + seqlen_q) cu_seqlens_k.append(cu_seqlens_k[-1] + seqlen_k) max_seqlen_q = max(seqlen_q, max_seqlen_q) max_seqlen_k = max(seqlen_k, max_seqlen_k) if req.block_table: num_cached = req.num_cached_tokens num_total = len(req) for token_idx in range(num_cached, num_total): block_idx = token_idx // self.block_size block_offset = token_idx % self.block_size block_id = req.block_table[block_idx] slot_idx = block_id * self.block_size + block_offset slot_mapping.append(slot_idx) input_ids = torch.tensor(input_ids, dtype=torch.int64, pin_memory=True).cuda( non_blocking=True ) positions = torch.tensor(positions, dtype=torch.int64, pin_memory=True).cuda( non_blocking=True ) cu_seqlens_q = torch.tensor( cu_seqlens_q, dtype=torch.int32, pin_memory=True ).cuda(non_blocking=True) cu_seqlens_k = torch.tensor( cu_seqlens_k, dtype=torch.int32, pin_memory=True ).cuda(non_blocking=True) slot_mapping = torch.tensor( slot_mapping, dtype=torch.int32, pin_memory=True ).cuda(non_blocking=True) block_tables = None if cu_seqlens_k[-1] > cu_seqlens_q[-1]: max_len = max(len(req.block_table) for req in requests) block_tables_list = [] for req in requests: table = req.block_table + [-1] * (max_len - len(req.block_table)) block_tables_list.append(table) block_tables = torch.tensor( block_tables_list, dtype=torch.int32, pin_memory=True ).cuda(non_blocking=True) set_forward_context( True, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k, slot_mapping, None, block_tables, ) return input_ids, positions def _reset_kv_allocator_state(self): # Keep allocated KV tensors, but reset allocator metadata to avoid # stale block reuse across independent generation calls. self.kv_manager.block_hash_to_id.clear() self.kv_manager.free_block_ids = deque(range(self.num_blocks)) self.kv_manager.used_block_ids.clear() for block in self.kv_manager.blocks: block.ref_cnt = 0 block._block_hash = None block.token_ids = [] def _prepare_decode(self, requests: List[Seq]): if not requests: raise RuntimeError("FATAL: No requests provided to _prepare_decode!") input_ids = [] positions = [] slot_mapping = [] context_lens = [] for req in requests: input_ids.append(req.last_token) pos = len(req) - 1 if hasattr(self, "_tts_mode") and self._tts_mode: pos = pos - (self._tts_prompt_len - 1) positions.append(pos) context_lens.append(len(req)) slot_mapping.append( req.block_table[-1] * self.block_size + req.last_block_num_tokens - 1 ) input_ids = torch.tensor(input_ids, dtype=torch.int64, pin_memory=True).cuda( non_blocking=True ) positions = torch.tensor(positions, dtype=torch.int64, pin_memory=True).cuda( non_blocking=True ) slot_mapping = torch.tensor( slot_mapping, dtype=torch.int32, pin_memory=True ).cuda(non_blocking=True) context_lens = torch.tensor( context_lens, dtype=torch.int32, pin_memory=True ).cuda(non_blocking=True) max_len = max(len(req.block_table) for req in requests) block_tables_list = [] for req in requests: table = req.block_table + [-1] * (max_len - len(req.block_table)) block_tables_list.append(table) block_tables = torch.tensor( block_tables_list, dtype=torch.int32, pin_memory=True ).cuda(non_blocking=True) assert block_tables.dim() == 2, ( f"block_tables must be 2D, got shape {block_tables.shape}" ) assert block_tables.size(0) == len(requests), ( f"block_tables batch size mismatch: {block_tables.size(0)} vs {len(requests)}" ) set_forward_context( False, slot_mapping=slot_mapping, context_lens=context_lens, block_tables=block_tables, ) return input_ids, positions def _prepare_sample(self, requests: List[Seq], temperature: float): temperatures = [temperature] * len(requests) temperatures = torch.tensor( temperatures, dtype=torch.float32, pin_memory=True ).cuda(non_blocking=True) return temperatures def _capture_cuda_graphs( self, tts_mel_embedding=None, tts_text_pos_embedding=None, max_num_blocks: Optional[int] = None, ): print("Capturing CUDA graphs for decode optimization...") max_bs = 8 # Support up to batch size 8 if max_num_blocks is None: max_num_blocks = self.num_blocks max_num_blocks = max(1, min(int(max_num_blocks), int(self.num_blocks))) model_dtype = next(self.model.parameters()).dtype input_ids = torch.ones(max_bs, dtype=torch.int64, device="cuda") positions = torch.ones(max_bs, dtype=torch.int64, device="cuda") slot_mapping = torch.zeros(max_bs, dtype=torch.int32, device="cuda") context_lens = torch.zeros(max_bs, dtype=torch.int32, device="cuda") block_tables = torch.zeros(max_bs, max_num_blocks, dtype=torch.int32, device="cuda") outputs = torch.zeros(max_bs, self.hidden_size, dtype=model_dtype, device="cuda") inputs_embeds_buffer = torch.zeros(max_bs, self.hidden_size, dtype=model_dtype, device="cuda") self.graph_bs = [1, 2, 4, 8] use_tts = tts_mel_embedding is not None and tts_text_pos_embedding is not None for bs in reversed(self.graph_bs): graph = torch.cuda.CUDAGraph() slot_mapping[:bs].copy_(torch.arange(bs, dtype=torch.int32, device="cuda")) context_lens[:bs].fill_(bs + 1) block_tables[:bs, :].zero_() set_forward_context( False, slot_mapping=slot_mapping[:bs], context_lens=context_lens[:bs], block_tables=block_tables[:bs], ) # warmup if use_tts: inputs_embeds_buffer[:bs].copy_( self._compute_tts_embeds( input_ids[:bs], positions[:bs], tts_mel_embedding=tts_mel_embedding, tts_text_pos_embedding=tts_text_pos_embedding, ) ) out = self.model( inputs_embeds=inputs_embeds_buffer[:bs].unsqueeze(1), return_dict=True, ).last_hidden_state else: out = self.model( input_ids=input_ids[:bs].unsqueeze(1), return_dict=True ).last_hidden_state outputs[:bs].copy_(out.squeeze(1) if out.dim() == 3 else out) with torch.cuda.graph(graph, self.graph_pool, capture_error_mode="thread_local"): if use_tts: inputs_embeds_buffer[:bs].copy_( self._compute_tts_embeds( input_ids[:bs], positions[:bs], tts_mel_embedding=tts_mel_embedding, tts_text_pos_embedding=tts_text_pos_embedding, ) ) out = self.model( inputs_embeds=inputs_embeds_buffer[:bs].unsqueeze(1), return_dict=True, ).last_hidden_state else: out = self.model( input_ids=input_ids[:bs].unsqueeze(1), return_dict=True ).last_hidden_state outputs[:bs].copy_(out.squeeze(1) if out.dim() == 3 else out) if self.graph_pool is None: self.graph_pool = graph.pool() self.graphs[bs] = graph torch.cuda.synchronize() reset_forward_context() self.graph_vars = { "input_ids": input_ids, "positions": positions, "slot_mapping": slot_mapping, "context_lens": context_lens, "block_tables": block_tables, "outputs": outputs, "inputs_embeds": inputs_embeds_buffer, } print(f"CUDA graphs captured for batch sizes: {self.graph_bs}") def _run_decode_with_graph( self, input_ids: torch.Tensor, positions: torch.Tensor, context: ForwardContext, tts_mel_embedding: Optional[torch.nn.Module] = None, tts_text_pos_embedding: Optional[torch.nn.Module] = None, ) -> torch.Tensor: bs = input_ids.size(0) use_tts_embedding = hasattr(self, "_tts_mode") and self._tts_mode if not self.use_cuda_graph or not self.graphs: if use_tts_embedding: inputs_embeds = self._compute_tts_embeds( input_ids, positions, tts_mel_embedding=tts_mel_embedding, tts_text_pos_embedding=tts_text_pos_embedding, ) out = self.model( inputs_embeds=inputs_embeds.unsqueeze(1), return_dict=True ).last_hidden_state else: out = self.model( input_ids=input_ids.unsqueeze(1), return_dict=True ).last_hidden_state return out.squeeze(1) if out.dim() == 3 else out graph_bs = next((x for x in self.graph_bs if x >= bs), None) if graph_bs is None: if use_tts_embedding: inputs_embeds = self._compute_tts_embeds( input_ids, positions, tts_mel_embedding=tts_mel_embedding, tts_text_pos_embedding=tts_text_pos_embedding, ) out = self.model( inputs_embeds=inputs_embeds.unsqueeze(1), return_dict=True ).last_hidden_state else: out = self.model( input_ids=input_ids.unsqueeze(1), return_dict=True ).last_hidden_state return out.squeeze(1) if out.dim() == 3 else out graph = self.graphs[graph_bs] graph_vars = self.graph_vars if graph_vars is None: raise RuntimeError("Graph variables not initialized") graph_vars["input_ids"][:bs].copy_(input_ids) graph_vars["positions"][:bs].copy_(positions) graph_vars["slot_mapping"].fill_(-1) graph_vars["slot_mapping"][:bs].copy_(context.slot_mapping) graph_vars["context_lens"].zero_() graph_vars["context_lens"][:bs].copy_(context.context_lens) graph_vars["block_tables"][:bs, :].fill_(-1) graph_vars["block_tables"][:bs, : context.block_tables.size(1)].copy_(context.block_tables) graph.replay() return graph_vars["outputs"][:bs] def generate( self, input_ids: torch.Tensor, max_new_tokens: int = 100, temperature: float = 1.0, top_k: int = 50, top_p: float = 1.0, stop_tokens: Optional[List[int]] = None, attention_mask: Optional[torch.Tensor] = None, tts_embeddings: Optional[ torch.Tensor ] = None, # TTS: [pad][cond][text] embeddings (87 tokens, NO start_mel) tts_mel_embedding: Optional[torch.nn.Module] = None, # TTS: mel_embedding layer tts_text_pos_embedding: Optional[ torch.nn.Module ] = None, # TTS: text_pos_embedding layer cg_max_total_tokens: Optional[int] = None, stop_checker: Optional[Callable[[], bool]] = None, ) -> torch.Tensor: """ Generate tokens. Args: input_ids: Input token IDs [batch_size, seq_len] max_new_tokens: Maximum number of tokens to generate temperature: Sampling temperature top_k: Top-k sampling top_p: Nucleus sampling threshold stop_tokens: List of token IDs that stop generation Returns: Generated token IDs [batch_size, total_len] """ batch_size = input_ids.size(0) device = input_ids.device self._tts_mode = tts_embeddings is not None self._tts_prompt_len = input_ids.size(1) if self._tts_mode else 0 self._reset_kv_allocator_state() prompt_tokens = (tts_embeddings.size(1) + 1) if tts_embeddings is not None else input_ids.size(1) required_total_tokens = int(prompt_tokens + max(1, int(max_new_tokens))) if cg_max_total_tokens is not None: required_total_tokens = max(required_total_tokens, int(cg_max_total_tokens)) self.prepare_decode_graph( required_total_tokens, tts_mel_embedding=tts_mel_embedding, tts_text_pos_embedding=tts_text_pos_embedding, ) if tts_embeddings is not None: actual_seq_len = tts_embeddings.size(1) + 1 # embeddings + start_mel_token else: actual_seq_len = input_ids.size(1) is_varlen_batch = ( tts_embeddings is not None and attention_mask is not None and batch_size > 1 and (attention_mask.sum(dim=1) != attention_mask.size(1)).any() ) if is_varlen_batch: seq_lens = [attention_mask[i].sum().item() for i in range(batch_size)] else: seq_lens = [actual_seq_len] * batch_size sequences = [] for i in range(batch_size): seq_len = seq_lens[i] token_ids = [1] * seq_len if tts_embeddings is not None and seq_len > 0: token_ids[-1] = input_ids[i, -1].item() if input_ids.size(1) > 0 else 1 else: token_ids = input_ids[i].tolist() req = Seq(token_ids) self.kv_manager.allocate(req) sequences.append(req) self.current_sequences = sequences prefill_ids, prefill_pos = self._prepare_prefill(sequences) if ( tts_embeddings is not None and tts_mel_embedding is not None and tts_text_pos_embedding is not None ): start_token_id = input_ids[0, -1] if input_ids.size(1) > 0 else 8192 start_emb = tts_mel_embedding( torch.tensor([[start_token_id]], device="cuda") ) # [1, 1, hidden_dim] start_pos = torch.tensor( [[tts_embeddings.size(1)]], device="cuda", dtype=torch.long ) pos_emb_module = getattr(tts_text_pos_embedding, "emb", tts_text_pos_embedding) pos_emb = pos_emb_module(start_pos) start_emb = start_emb + pos_emb start_emb = start_emb.repeat(batch_size, 1, 1) if is_varlen_batch: valid_embeddings = [] for i in range(batch_size): emb_len = seq_lens[i] - 1 padding_len = tts_embeddings.size(1) - emb_len valid_emb = tts_embeddings[i, padding_len:].unsqueeze( 0 ) # [1, emb_len, hidden_dim] valid_embeddings.append( torch.cat([valid_emb, start_emb[i : i + 1]], dim=1) ) full_embeddings = torch.cat( valid_embeddings, dim=1 ) # [1, total_tokens, hidden_dim] else: full_embeddings = torch.cat( [tts_embeddings, start_emb], dim=1 ) # [batch_size, seq_len, hidden_dim] model_dtype = next(self.model.parameters()).dtype if full_embeddings.dtype != model_dtype: full_embeddings = full_embeddings.to(model_dtype) hidden_states = self.model( inputs_embeds=full_embeddings, return_dict=True ).last_hidden_state else: hidden_states = self.model( input_ids=input_ids, attention_mask=attention_mask, return_dict=True ).last_hidden_state if is_varlen_batch: context = get_forward_context() cu_seqlens = context.cu_seqlens_q.cpu().tolist() last_hidden = torch.stack( [hidden_states[0, cu_seqlens[i + 1] - 1] for i in range(batch_size)] ) else: last_hidden = hidden_states[:, -1, :] # [batch_size, hidden_size] reset_forward_context() if self.lm_head is not None: if last_hidden.dtype != next(self.lm_head.parameters()).dtype: last_hidden = last_hidden.to(next(self.lm_head.parameters()).dtype) logits = self.lm_head(last_hidden) # [batch_size, vocab_size] else: logits = self.model.compute_logits(last_hidden) # [batch_size, vocab_size] temperatures = self._prepare_sample(sequences, temperature) if temperature > 0: sampling_logits = self._filter_logits(logits, top_k=top_k, top_p=top_p) first_token = self.sampler(sampling_logits, temperatures) else: first_token = torch.argmax(logits, dim=-1) first_token_list = first_token.tolist() generated_tokens = [[] for _ in range(batch_size)] is_finished = [False] * batch_size token_progress = tqdm(total=int(max_new_tokens), desc="transformer_tokens", unit="tok", leave=True) def _should_stop_early(): if stop_checker is None: return False try: return bool(stop_checker()) except Exception: return False try: for i, token_id in enumerate(first_token_list): if stop_tokens and token_id in stop_tokens: is_finished[i] = True generated_tokens[i].append(token_id) else: generated_tokens[i].append(token_id) sequences[i].append_token(token_id) self.kv_manager.append_to_seq(sequences[i]) token_progress.update(1) stop_early = _should_stop_early() if all(is_finished) and not stop_early: for req in sequences: self.kv_manager.remove_seq(req) self.current_sequences = [] output_ids = [] for i in range(batch_size): full_sequence = input_ids[i].tolist() + generated_tokens[i] output_ids.append(full_sequence) output = torch.tensor(output_ids, dtype=torch.long, device=device) return output remaining_tokens = 0 if stop_early else (max_new_tokens - 1) for step in range(remaining_tokens): if _should_stop_early(): break decode_ids, decode_pos = self._prepare_decode(sequences) context = get_forward_context() hidden_states = self._run_decode_with_graph( decode_ids, decode_pos, context, tts_mel_embedding=tts_mel_embedding, tts_text_pos_embedding=tts_text_pos_embedding, ) # Get logits if self.lm_head is not None: logits = self.lm_head(hidden_states) # [batch_size, vocab_size] else: logits = self.model.compute_logits( hidden_states ) # [batch_size, vocab_size] reset_forward_context() temperatures = self._prepare_sample(sequences, temperature) if temperature > 0: sampling_logits = self._filter_logits(logits, top_k=top_k, top_p=top_p) next_token = self.sampler(sampling_logits, temperatures) else: next_token = torch.argmax(logits, dim=-1) next_token_list = next_token.tolist() for i, token_id in enumerate(next_token_list): if is_finished[i]: continue elif stop_tokens and token_id in stop_tokens: is_finished[i] = True generated_tokens[i].append(token_id) else: sequences[i].append_token(token_id) self.kv_manager.append_to_seq(sequences[i]) generated_tokens[i].append(token_id) token_progress.update(1) if all(is_finished): break for req in sequences: self.kv_manager.remove_seq(req) self.current_sequences = [] pad_token = stop_tokens[0] if stop_tokens else 0 if is_varlen_batch: max_prompt_len = attention_mask.size(1) output_ids = [] for i in range(batch_size): padding_len = max_prompt_len - seq_lens[i] initial_tokens = sequences[i].token_ids[ : sequences[i].num_prompt_tokens ] padded_prompt = [pad_token] * padding_len + initial_tokens full_sequence = padded_prompt + generated_tokens[i] output_ids.append(full_sequence) else: output_ids = [ sequences[i].token_ids[: sequences[i].num_prompt_tokens] + generated_tokens[i] for i in range(batch_size) ] max_length = max(len(seq) for seq in output_ids) padded_output_ids = [ seq + [pad_token] * (max_length - len(seq)) for seq in output_ids ] output = torch.tensor(padded_output_ids, dtype=torch.long, device=device) assert output.size(0) == batch_size, ( f"Output batch size mismatch: {output.size(0)} != {batch_size}" ) return output finally: token_progress.close() class Sampler(nn.Module): def __init__(self): super().__init__() # @torch.compile def forward(self, logits: torch.Tensor, temperatures: torch.Tensor): logits = logits.float().div_(temperatures.unsqueeze(dim=1)) probs = torch.softmax(logits, dim=-1) sample_tokens = probs.div_( torch.empty_like(probs).exponential_(1).clamp_min_(1e-10) ).argmax(dim=-1) return sample_tokens