openbmb/MiniCPM-o-4_5
1.5k710k
1#!/usr/bin/env python2# -*- coding: utf-8 -*-3# Copyright 2026 The OpenBMB Team. All rights reserved.4#5# Licensed under the Apache License, Version 2.0 (the "License");6# you may not use this file except in compliance with the License.7# You may obtain a copy of the License at8#9# http://www.apache.org/licenses/LICENSE-2.010#11# Unless required by applicable law or agreed to in writing, software12# distributed under the License is distributed on an "AS IS" BASIS,13# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.14# See the License for the specific language governing permissions and15# limitations under the License.16 17import logging18from dataclasses import dataclass19from typing import Any20from typing import Dict21from typing import List22from typing import Literal23from typing import Optional24from typing import Tuple25from typing import Union26 27import torch28import torch.nn.functional as F29import torch.nn.utils.parametrize as P30from transformers.cache_utils import DynamicCache31 32logger = logging.getLogger(__name__)33 34 35# text36@dataclass37class GenerateChunkOutput:38 chunk_token_ids: torch.Tensor39 current_inputs_embeds: torch.Tensor40 input_last_hidden_states: Optional[torch.Tensor] # for tts use_speaker_embedding41 last_hidden_states: Optional[torch.Tensor] # for tts input feature (projector_semantic)42 past_key_values: Optional[torch.Tensor]43 finished: bool44 45 46class ChunkPrefillChunkGenerate:47 def __init__(self, model, tokenizer, terminators):48 self.tokenizer = tokenizer49 self.model = model50 self.terminators = terminators51 self.terminators_ids = [tokenizer.convert_tokens_to_ids(i) for i in self.terminators]52 self.embedding_layer = self.model.get_input_embeddings()53 54 self.forbidden_tokens = [55 ":",56 "：",57 "；",58 "#",59 "“",60 "”",61 "‘",62 "’",63 "@",64 "*",65 "【",66 "】",67 "「",68 "」",69 "(",70 ")",71 "（",72 "）",73 "[",74 "]",75 "&",76 "/",77 "$",78 ]79 80 self.forbidden_token_ids = [tokenizer.convert_tokens_to_ids(i) for i in self.forbidden_tokens]81 bad_token_ids = getattr(tokenizer, "bad_token_ids", [])82 if bad_token_ids:83 self.forbidden_token_ids.extend(bad_token_ids)84 85 @staticmethod86 def prepare_generation_config(do_sample, max_new_tokens=50, min_new_tokens=0, **kwargs):87 num_beams = kwargs.get("num_beams", 3)88 generation_config = {89 "num_beams": num_beams,90 "top_p": 0.8,91 "top_k": 100,92 "temperature": 0.7,93 "do_sample": True,94 "repetition_penalty": 1.05,95 }96 97 if do_sample:98 generation_config.update(99 {100 "top_p": 0.8,101 "top_k": 100,102 "temperature": 0.7,103 "do_sample": True,104 "repetition_penalty": 1.05,105 }106 )107 elif num_beams > 1:108 generation_config.update({"num_beams": num_beams, "repetition_penalty": 1.2, "do_sample": False})109 else:110 generation_config.update({"do_sample": False, "repetition_penalty": 1.05})111 112 generation_config.update((k, kwargs[k]) for k in generation_config.keys() & kwargs.keys())113 generation_config["min_new_tokens"] = min_new_tokens114 generation_config["max_new_tokens"] = max_new_tokens115 116 return generation_config117 118 def chunk_generate(119 self,120 inputs_embeds: torch.Tensor,121 past_key_values,122 is_first_generate_chunk: bool,123 chunk_size: int,124 return_hidden_states: bool,125 do_sample: bool,126 temperature: float,127 top_p: float,128 top_k: int,129 repetition_penalty: float = 1.05,130 length_penalty: float = 1.0,131 all_input_ids: Optional[torch.Tensor] = None,132 ) -> GenerateChunkOutput:133 """134 Args:135 inputs_embeds: [1, seq_len, hidden_dim], Input embeddings of current chunk.136 past_key_values: [num_layers, 2, batch_size, num_heads, seq_len, head_dim], Past key values for llm.137 is_first_generate_chunk: bool, Whether this is the first generate chunk.138 chunk_size: int, The size of the current chunk, default is 10, and it is fixed during training.139 return_hidden_states: bool Whether to return the hidden states, default is True.140 do_sample: bool Whether to sample from the model, default is True.141 temperature: float The temperature for the model, default is 0.7.142 top_p: float The top-p for the model, default is 0.8.143 top_k: int The top-k for the model, default is 100.144 repetition_penalty: float, The repetition penalty for the model, default is 1.05.145 length_penalty: float, The length penalty for the model, default is 1.0. Higher value means more detailed generation.146 all_input_ids: Optional[torch.Tensor], The input ids for the current chunk.147 """148 149 finished = False150 current_inputs_embeds = inputs_embeds.clone()151 input_last_hidden_states = []152 last_hidden_states = []153 generated_tokens = []154 155 for token_idx in range(chunk_size):156 if is_first_generate_chunk and token_idx == 0:157 # first generate chunk, prefill inputs_embeds158 model_inputs = {159 "inputs_embeds": current_inputs_embeds,160 "past_key_values": past_key_values,161 "use_cache": True,162 "output_hidden_states": return_hidden_states,163 }164 else: # for all other cases: prefill the latest generated token165 model_inputs = {166 "inputs_embeds": current_inputs_embeds[:, -1:, :],167 "past_key_values": past_key_values,168 "use_cache": True,169 "output_hidden_states": return_hidden_states,170 }171 172 with torch.no_grad():173 outputs = self.model(**model_inputs)174 175 # last token's logits176 logits = outputs.logits[:, -1, :].to(copy=True, dtype=torch.float32, device=inputs_embeds.device)177 178 # forbid specific tokens decoding = model.generate@suppress_tokens179 if self.forbidden_token_ids:180 logits[:, self.forbidden_token_ids] = float("-inf")181 182 past_key_values = outputs.past_key_values183 184 PENALTY_WINDOW_SIZE = 128185 186 # apply repetition penalty187 if repetition_penalty != 1.0:188 # get token ids for repetition penalty189 if all_input_ids is not None:190 # use global input ids (including original input and generated part)191 if len(generated_tokens) > 0:192 generated_token_ids = torch.cat(generated_tokens, dim=1)193 current_sequence = torch.cat(194 [195 all_input_ids[:, -PENALTY_WINDOW_SIZE:],196 generated_token_ids,197 ],198 dim=1,199 )200 else:201 current_sequence = all_input_ids[:, -PENALTY_WINDOW_SIZE:]202 unique_token_ids = torch.unique(current_sequence.squeeze(0))203 elif len(generated_tokens) > 0:204 # revert to original logic: only use generated tokens205 generated_token_ids = torch.cat(generated_tokens, dim=1).squeeze(0)206 unique_token_ids = torch.unique(generated_token_ids)207 else:208 unique_token_ids = torch.tensor([], dtype=torch.long, device=logits.device)209 210 # apply repetition penalty211 for token_id in unique_token_ids:212 if logits[0, token_id] > 0:213 logits[0, token_id] = logits[0, token_id] / repetition_penalty214 else:215 logits[0, token_id] = logits[0, token_id] * repetition_penalty216 217 # apply length penalty, higher value means more detailed generation218 if length_penalty != 1.0:219 for eos_token_id in self.terminators_ids:220 if logits[0, eos_token_id] > 0:221 logits[0, eos_token_id] = logits[0, eos_token_id] / length_penalty222 else:223 logits[0, eos_token_id] = logits[0, eos_token_id] * length_penalty224 225 # apply temperature226 if temperature != 1.0:227 logits = logits / temperature228 229 if do_sample:230 # Top-k filtering231 if top_k > 0:232 top_k_logits, top_k_indices = torch.topk(logits, min(top_k, logits.size(-1)))233 logits_filtered = torch.full_like(logits, float("-inf"))234 logits_filtered.scatter_(1, top_k_indices, top_k_logits)235 logits = logits_filtered236 237 # Top-p filtering238 if top_p < 1.0:239 sorted_logits, sorted_indices = torch.sort(logits, descending=True)240 cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)241 242 # remove tokens with cumulative probability greater than top_p243 sorted_indices_to_remove = cumulative_probs > top_p244 sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()245 sorted_indices_to_remove[..., 0] = 0246 247 indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)248 logits[indices_to_remove] = float("-inf")249 250 # sampling251 probs = F.softmax(logits, dim=-1)252 next_token = torch.multinomial(probs, num_samples=1)253 else:254 next_token = torch.argmax(logits, dim=-1, keepdim=True)255 256 if return_hidden_states:257 if is_first_generate_chunk and token_idx == 0:258 input_last_hidden_states.append(outputs.hidden_states[-1])259 else:260 last_hidden_states.append(outputs.hidden_states[-1])261 262 # if terminator token, stop generating263 if next_token.item() in self.terminators_ids:264 finished = True265 break266 267 generated_tokens.append(next_token)268 269 # convert new token to embeddings and concatenate270 next_token_embed = self.embedding_layer(next_token)271 272 # update inputs_embeds, add one273 current_inputs_embeds = torch.cat([current_inputs_embeds, next_token_embed], dim=1)274 275 if len(generated_tokens) > 0:276 chunk_token_ids = torch.cat(generated_tokens, dim=1)277 else:278 # special case: if last chunk and first predict is eos token, return last token of previous chunk. return a tensor with shape (1, 0)279 if finished:280 chunk_token_ids = torch.zeros((1, 0), dtype=torch.long, device=current_inputs_embeds.device)281 else:282 raise Exception("this should not happen")283 284 if len(last_hidden_states) > 0:285 last_hidden_states = torch.cat(last_hidden_states, dim=1)286 else:287 # special case: if last chunk, return last token of previous chunk.288 if finished:289 last_hidden_states = torch.cat(last_hidden_states, dim=1)290 else:291 raise Exception("this should not happen")292 293 if len(input_last_hidden_states) > 0:294 input_last_hidden_states = torch.cat(input_last_hidden_states, dim=1)295 else:296 input_last_hidden_states = None297 298 return GenerateChunkOutput(299 chunk_token_ids=chunk_token_ids,300 current_inputs_embeds=current_inputs_embeds,301 input_last_hidden_states=input_last_hidden_states,302 last_hidden_states=last_hidden_states,303 past_key_values=past_key_values,304 finished=finished,305 )306 307 308def streaming_token_decoder(token_iterator, tokenizer, skip_special_tokens=False):309 """310 Incrementally decode tokens from an iterator, handling partial multi-byte characters.311 312 When streaming tokens, multi-byte characters (like Chinese) may be split across multiple313 tokens. Decoding partial tokens results in replacement characters (U+FFFD). This function314 buffers tokens and only yields complete characters.315 316 Args:317 token_iterator: An iterator yielding (token_ids, is_finished) tuples.318 token_ids can be torch.Tensor or any iterable of integers.319 tokenizer: The tokenizer to use for decoding.320 skip_special_tokens: Whether to skip special tokens during decoding.321 322 Yields:323 (decoded_text, is_finished) tuples where decoded_text is the new text since last yield.324 """325 accumulated_token_ids = []326 yielded_text_len = 0327 328 for token_ids, is_finished in token_iterator:329 # Accumulate token IDs330 if torch.is_tensor(token_ids):331 accumulated_token_ids.extend(token_ids.reshape(-1).tolist())332 else:333 accumulated_token_ids.extend(list(token_ids) if hasattr(token_ids, "__iter__") else [token_ids])334 335 # Decode all accumulated tokens336 full_decoded = tokenizer.decode(accumulated_token_ids, skip_special_tokens=skip_special_tokens)337 338 if is_finished:339 # Final chunk - yield all remaining text340 new_text = full_decoded[yielded_text_len:]341 yield new_text, is_finished342 else:343 # Find safe prefix without incomplete multi-byte characters344 # The replacement character '�' (U+FFFD) indicates incomplete decoding345 new_text = full_decoded[yielded_text_len:]346 347 # Hold back text ending with replacement character (incomplete UTF-8 sequence)348 safe_end = len(new_text)349 while safe_end > 0 and new_text[safe_end - 1] == "\ufffd":350 safe_end -= 1351 352 safe_text = new_text[:safe_end] if safe_end > 0 else ""353 yielded_text_len += len(safe_text)354 yield safe_text, is_finished355 356 357def torch_clone_recursive(obj):358 """Recursively clone nested containers of torch.Tensors.359 360 Supported container types: dict, list, tuple. Non-container non-Tensor361 objects are returned as-is.362 """363 if torch.is_tensor(obj):364 return obj.clone()365 elif isinstance(obj, dict):366 return {k: torch_clone_recursive(v) for k, v in obj.items()}367 elif isinstance(obj, list):368 return [torch_clone_recursive(v) for v in obj]369 elif isinstance(obj, tuple):370 return tuple(torch_clone_recursive(v) for v in obj)371 else:372 raise ValueError(f"Unsupported type: {type(obj)}")373 374 375def rotate_half(x: torch.Tensor) -> torch.Tensor:376 """Rotate half the hidden dims of the input for RoPE."""377 dim = x.shape[-1]378 x1 = x[..., : dim // 2]379 x2 = x[..., dim // 2 :]380 return torch.cat((-x2, x1), dim=-1)381 382 383@dataclass384class SpeculativeSnapshot:385 """Speculative snapshot for VAD speculative rollback.386 387 Used in VAD speculative execution: creates a snapshot after streaming_prefill388 and before streaming_generate. If speculation fails (user continues speaking),389 the state can be restored to continue streaming_prefill.390 391 Implementation:392 - LLM KV Cache: only record length, restore by truncation (zero extra VRAM)393 - Audio KV Cache: requires cloning, as generate sets it to None394 - Mel processor: save full state snapshot (including buffer)395 """396 397 # KV Cache length (for truncation recovery)398 llm_cache_length: int399 audio_cache_length: int400 401 # session state402 new_user_msg: bool403 llm_generated: bool404 llm_generate_completed: bool405 406 # Round management407 next_round_id: int408 pending_round_id: Optional[int]409 omni_chunk_history_length: int410 411 # TTS state (requires cloning, but usually small)412 tts_last_turn_tokens: Optional[torch.Tensor]413 414 # Streaming processor state415 audio_chunk_idx: int416 417 # Mel processor state snapshot (including buffer)418 mel_processor_snapshot: Optional[dict] = None419 420 # Audio encoder KV cache (requires cloning to ensure determinism after recovery)421 audio_past_key_values: Optional[tuple] = None422 423 # timestamp (for debugging)424 timestamp: float = 0.0425 426 # debug field: for verifying correctness of recovery427 llm_cache_checksum: Optional[float] = None # LLM KV Cache first layer K sum428 audio_cache_checksum: Optional[float] = None # Audio KV Cache first layer K sum429 mel_buffer_checksum: Optional[float] = None # Mel buffer sum430 431 # RNG state (key: for ensuring determinism of dithering etc. after recovery)432 rng_state_cpu: Optional[torch.Tensor] = None # torch CPU RNG state433 rng_state_cuda: Optional[torch.Tensor] = None # torch CUDA RNG state (if on GPU)434 435 def summary(self) -> str:436 mel_buf_len = 0437 if self.mel_processor_snapshot:438 buf = self.mel_processor_snapshot.get("buffer")439 if buf is not None:440 mel_buf_len = len(buf)441 return (442 f"llm_cache={self.llm_cache_length}, "443 f"audio_cache={self.audio_cache_length}, "444 f"audio_chunk_idx={self.audio_chunk_idx}, "445 f"mel_buffer={mel_buf_len}, "446 f"history_len={self.omni_chunk_history_length}, "447 f"new_user_msg={self.new_user_msg}, "448 f"llm_generated={self.llm_generated}"449 )450 451 452# tts453@dataclass454class TTSSamplingParams:455 top_p: float = 0.85456 min_p: float = 0.01457 top_k: int = 25458 repetition_penalty: float = 1.05459 temperature: float = 0.8460 win_size: int = 16461 tau_r: float = 0.1462 463 464class TTSStreamingGenerator:465 """466 Streaming generator for TTS that processes chunks and yields audio tokens in real-time.467 468 Supported attention types:469 - full_attention: Full attention, all tokens can attend to each other470 - sliding_window: Sliding window attention, KV cache is truncated to fixed size (token_window_size)471 - sliding_recompute: Sliding recompute, only keep previous chunk and recompute with current chunk472 - reindex: Keep first chunk as sink, reindex sliding window positions via RoPE rotation473 """474 475 def __init__(476 self,477 model,478 temperature: float,479 eos_token: Union[int, torch.Tensor],480 chunk_size: int = 25, # s3tokenizer 1s = 25token481 tts_last_turn_tokens: torch.Tensor = None,482 logits_processors=None,483 logits_warpers=None,484 ):485 self.tts = model486 self.device = model.device487 self.temperature = torch.tensor([temperature], dtype=torch.float, device=self.device)488 self.eos_token = (489 torch.tensor(eos_token, device=self.device) if isinstance(eos_token, int) else eos_token.to(self.device)490 )491 492 self.num_vq = model.num_vq493 self.num_audio_tokens = model.num_audio_tokens494 self.recomputed_chunks = model.recomputed_chunks495 self.emb_code = model.emb_code496 self.head_code = model.head_code497 498 # Attention type and window sizes499 self.attention_type = model.attention_type # "full_attention", "sliding_window", "sliding_recompute", "reindex"500 self.chunk_window_size = model.chunk_window_size # chunk-level window for sliding_recompute (default 2)501 self.token_window_size = model.token_window_size # token-level window for sliding_window/reindex (default 300)502 503 # RoPE config (for reindex mode)504 self.rope_theta = model.model.config.rope_theta505 self.head_dim = model.model.config.hidden_size // model.model.config.num_attention_heads506 507 # Logits processors508 self.logits_processors = logits_processors if logits_processors is not None else []509 # Logits warpers (like TopP/TopK), separate from processors510 self.logits_warpers = logits_warpers if logits_warpers is not None else []511 512 # initialize state513 self.past_key_values = None514 self.text_start_pos = 0515 self.idx = -1 # start from -1, become 0 when first called516 self.all_conditions = []517 self.all_generated_tokens = []518 self.tts_last_turn_tokens = tts_last_turn_tokens519 self.spk_emb = None520 521 audio_bos = [self.tts.audio_bos_token_id]522 audio_bos = torch.Tensor(audio_bos).to(self.tts.emb_text.weight.device, dtype=torch.long)523 524 self.audio_bos_embeds = self.tts.emb_text(audio_bos).unsqueeze(0)525 self.text_eos_embed = self.tts.emb_text(526 torch.tensor(527 [self.tts.config.text_eos_token_id],528 device=self.tts.emb_text.weight.device,529 dtype=torch.long,530 )531 ).unsqueeze(0)532 533 # buffer related, used to fill up chunk_size and yield to outside534 self.chunk_size = chunk_size535 self._token_buffer: List[torch.Tensor] = []536 537 # Chunk info tracking for sliding_recompute and reindex538 self._chunk_info: List[dict] = []539 self._total_seq_len = 0540 541 # Reindex mode: track sink (first chunk) length542 self._sink_kv_len = 0543 544 def _build_recompute_inputs(self, current_condition: torch.Tensor) -> torch.Tensor:545 """Build recompute inputs for sliding_recompute mode."""546 if len(self._chunk_info) == 0:547 return current_condition548 549 prev_chunk = self._chunk_info[-1]550 prev_condition = prev_chunk["condition"]551 prev_audio_tokens = prev_chunk["audio_tokens"]552 553 recompute_list = [prev_condition]554 if len(prev_audio_tokens) > 0:555 prev_audio_embeds = torch.cat([self.emb_code[0](tok) for tok in prev_audio_tokens], dim=1)556 recompute_list.append(prev_audio_embeds)557 558 recompute_list.append(current_condition)559 return torch.cat(recompute_list, dim=1)560 561 def _truncate_kv_cache_sliding_window(self):562 """Truncate KV cache for sliding_window mode."""563 if self.past_key_values is None:564 return565 566 if hasattr(self.past_key_values, "get_seq_length"):567 current_kv_len = self.past_key_values.get_seq_length()568 else:569 current_kv_len = self.past_key_values[0][0].shape[2]570 571 if current_kv_len <= self.token_window_size:572 return573 574 new_cache = DynamicCache()575 num_layers = (576 len(self.past_key_values.key_cache)577 if hasattr(self.past_key_values, "key_cache")578 else len(self.past_key_values)579 )580 581 for layer_idx in range(num_layers):582 if hasattr(self.past_key_values, "key_cache"):583 key = self.past_key_values.key_cache[layer_idx][:, :, -self.token_window_size :, :]584 value = self.past_key_values.value_cache[layer_idx][:, :, -self.token_window_size :, :]585 else:586 key = self.past_key_values[layer_idx][0][:, :, -self.token_window_size :, :]587 value = self.past_key_values[layer_idx][1][:, :, -self.token_window_size :, :]588 new_cache.update(key, value, layer_idx)589 590 self.past_key_values = new_cache591 592 @staticmethod593 def _apply_rope_rotation(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:594 """Apply RoPE rotation to tensor."""595 return x * cos + rotate_half(x) * sin596 597 def _compute_rope_cos_sin(self, positions: torch.Tensor, device: torch.device, dtype: torch.dtype):598 """Compute RoPE cos and sin for given positions."""599 dim_half = self.head_dim // 2600 freq_seq = torch.arange(0, dim_half, dtype=torch.float32, device=device)601 inv_freq = 1.0 / (self.rope_theta ** (freq_seq / dim_half))602 603 # positions: [seq_len]604 angles = positions.float().unsqueeze(-1) * inv_freq.unsqueeze(0) # [seq_len, dim_half]605 angles = torch.cat([angles, angles], dim=-1) # [seq_len, head_dim]606 607 cos = angles.cos().to(dtype)608 sin = angles.sin().to(dtype)609 return cos, sin610 611 def _reindex_kv_cache(self):612 """613 Reindex KV cache for reindex mode:614 1. Keep first chunk as attention sink615 2. Keep last chunk616 3. Discard middle chunks617 4. Reindex the last chunk's key positions to be right after sink via RoPE rotation618 """619 if self.past_key_values is None or len(self._chunk_info) < 2:620 return621 622 # Get current KV cache length623 if hasattr(self.past_key_values, "get_seq_length"):624 current_kv_len = self.past_key_values.get_seq_length()625 else:626 current_kv_len = self.past_key_values[0][0].shape[2]627 628 # Calculate sink length (first chunk)629 sink_len = self._chunk_info[0]["condition_len"] + self._chunk_info[0]["audio_token_count"]630 631 # Last chunk length632 last_chunk = self._chunk_info[-1]633 last_chunk_len = last_chunk["condition_len"] + last_chunk["audio_token_count"]634 635 keep_len = sink_len + last_chunk_len636 637 # Get device and dtype638 device = self.past_key_values.key_cache[0].device639 dtype = self.past_key_values.key_cache[0].dtype640 641 if current_kv_len <= keep_len:642 last_chunk_kv_len = current_kv_len - sink_len643 if last_chunk_kv_len <= 0:644 return645 self.text_start_pos = current_kv_len646 return647 648 # Step 1: Truncate KV cache - keep sink and last chunk649 new_cache = DynamicCache()650 num_layers = len(self.past_key_values.key_cache)651 652 original_start_pos = current_kv_len - last_chunk_len653 new_start_pos = sink_len654 delta = new_start_pos - original_start_pos # This is a scalar constant655 delta_positions = torch.full((last_chunk_len,), delta, dtype=torch.float32, device=device)656 657 # Compute rotation cos/sin658 cos, sin = self._compute_rope_cos_sin(delta_positions, device, dtype)659 cos = cos.unsqueeze(0).unsqueeze(0) # [1, 1, seq_len, head_dim]660 sin = sin.unsqueeze(0).unsqueeze(0)661 662 for layer_idx in range(num_layers):663 key_full = self.past_key_values.key_cache[layer_idx]664 value_full = self.past_key_values.value_cache[layer_idx]665 666 # Extract sink and last chunk667 key_sink = key_full[:, :, :sink_len, :]668 value_sink = value_full[:, :, :sink_len, :]669 key_last = key_full[:, :, -last_chunk_len:, :]670 value_last = value_full[:, :, -last_chunk_len:, :]671 672 # Apply RoPE rotation to reindex key positions673 key_last_reindexed = self._apply_rope_rotation(key_last, cos, sin)674 675 # Concatenate sink and reindexed last chunk676 key = torch.cat([key_sink, key_last_reindexed], dim=2)677 value = torch.cat([value_sink, value_last], dim=2)678 679 new_cache.update(key, value, layer_idx)680 681 self.past_key_values = new_cache682 683 # Update text_start_pos to reflect new positions684 self.text_start_pos = sink_len + last_chunk_len685 686 @torch.inference_mode()687 def generate_with_buffer(688 self,689 condition: torch.Tensor,690 text_finished: bool = False,691 max_new_token: int = 500,692 ):693 """input a condition embedding chunk, generate audio token each time,694 and accumulate to buffer, only yield when buffer satisfies chunk_size.695 696 Yields:697 torch.Tensor of shape [chunk_size] (2D: [1, chunk_size])698 """699 self.idx += 1700 self.device = self.tts.device701 702 # if text finished, first concatenate Text EOS703 if text_finished:704 condition = torch.cat([condition, self.text_eos_embed], dim=1)705 706 # always concatenate Audio BOS707 condition = torch.cat([condition, self.audio_bos_embeds], dim=1).to(self.device)708 709 self.all_conditions.append(condition)710 711 # Initialize current chunk info712 current_chunk_info = {713 "condition_len": condition.shape[1],714 "audio_token_count": 0,715 "condition": condition.clone(),716 "audio_tokens": [],717 }718 719 # Handle different attention types720 if self.attention_type == "sliding_recompute" and self.idx >= 1:721 # sliding_recompute: discard KV cache, recompute with previous + current chunk722 self.past_key_values = None723 current_condition = self._build_recompute_inputs(condition)724 self.text_start_pos = 0725 elif self.attention_type == "reindex" and self.idx >= 1:726 # reindex: truncate KV cache keeping sink + last chunk, reindex positions via RoPE727 self._reindex_kv_cache()728 current_condition = condition729 # Always update text_start_pos based on actual KV cache length (like reference code)730 if self.past_key_values is not None:731 if hasattr(self.past_key_values, "get_seq_length"):732 kv_len = self.past_key_values.get_seq_length()733 else:734 kv_len = self.past_key_values[0][0].shape[2]735 self.text_start_pos = kv_len736 else:737 current_condition = condition738 739 condition_length = current_condition.shape[1]740 prefill_len = condition_length741 finished = torch.zeros(1, dtype=torch.bool, device=self.device)742 chunk_generated_tokens = []743 744 for t in range(max_new_token):745 if t == 0:746 inputs_embeds = current_condition747 pos_ids = torch.arange(748 self.text_start_pos,749 self.text_start_pos + condition_length,750 dtype=torch.long,751 device=self.device,752 ).unsqueeze(0)753 else:754 last = self.all_generated_tokens[-1]755 # last: [1,1], directly as code id756 inputs_embeds = self.emb_code[0](last)757 pos_ids = torch.tensor(758 [self.text_start_pos + prefill_len + t - 1],759 dtype=torch.long,760 device=self.device,761 ).unsqueeze(0)762 763 outputs = self.tts.model(764 position_ids=pos_ids,765 past_key_values=self.past_key_values,766 inputs_embeds=inputs_embeds,767 use_cache=True,768 )769 hidden_states = outputs.last_hidden_state770 771 # Handle KV cache based on attention type772 if self.attention_type == "sliding_window":773 self.past_key_values = outputs.past_key_values774 self._truncate_kv_cache_sliding_window()775 else:776 self.past_key_values = outputs.past_key_values777 778 with P.cached():779 logits = torch.empty(780 hidden_states.size(0),781 hidden_states.size(1),782 self.num_audio_tokens,783 self.num_vq,784 dtype=torch.float,785 device=self.device,786 )787 for num_vq_iter in range(self.num_vq):788 x: torch.Tensor = self.head_code[num_vq_iter](hidden_states)789 logits[..., num_vq_iter] = x790 del x791 792 del hidden_states793 794 logits = logits[:, -1].float()795 796 logits = logits.permute(0, 2, 1)797 logits = logits.reshape(-1, logits.size(2))798 799 logits /= self.temperature800 801 audio_bos = len(self.all_generated_tokens) == 0 and t == 0802 803 if not audio_bos:804 # use generated tokens (current chunk) as input for processor/warper (align with modeling_minicpmo)805 all_generated_tokens = torch.cat(self.all_generated_tokens, dim=1).to(self.device) # [1, T]806 for processor in self.logits_processors:807 logits = processor(all_generated_tokens, logits)808 809 for warper in self.logits_warpers:810 logits = warper(all_generated_tokens, logits)811 del all_generated_tokens812 813 # sample next token (only use first codebook, same as generate)814 scores = F.softmax(logits, dim=-1)815 idx_next = torch.multinomial(scores, num_samples=1) # [(B*num_vq), 1]816 next_id = idx_next.view(-1, self.num_vq)[:, 0:1] # only take first codebook → [B, 1]817 del scores818 819 if next_id.eq(820 self.eos_token821 ).any(): # generated audio eos token, means this chunk is finished, no longer generate new tokens822 finished[:] = True823 else: # eos token cannot be added to buffer, he does not speak.824 # convert next_id to correct shape [1, 1], no num_vq dimension825 if next_id.dim() == 0: # if scalar826 next_tok = next_id.unsqueeze(0).unsqueeze(0) # [1, 1]827 elif next_id.dim() == 1: # if 1D [1]828 next_tok = next_id.unsqueeze(0) # [1, 1]829 else:830 next_tok = next_id831 832 self.all_generated_tokens.append(next_tok)833 chunk_generated_tokens.append(next_tok)834 835 # Update chunk info for sliding_recompute836 current_chunk_info["audio_tokens"].append(next_tok.clone())837 current_chunk_info["audio_token_count"] += 1838 839 self._token_buffer.append(next_tok)840 841 if len(self._token_buffer) == 0:842 # case 1: if last text chunk, yield None843 if text_finished:844 yield torch.empty(1, 0, dtype=torch.long, device=self.device), True845 break846 # case 2: if not last text chunk, break directly847 else:848 break849 else: # buffer has something850 # case 1: if buffer is larger/equal to chunk_size, yield out851 if len(self._token_buffer) >= self.chunk_size:852 batch = torch.cat(self._token_buffer[: self.chunk_size], dim=1) # [1, chunk_size]853 yield batch, False # → [1, chunk_size]854 # discard yielded part855 self._token_buffer = self._token_buffer[self.chunk_size :]856 857 # case 2: if buffer is smaller than chunk_size858 else:859 # if generation finished, and is the last text chunk, yield all remaining tokens, then break860 if finished.all():861 if text_finished:862 batch = torch.cat(self._token_buffer, dim=1) # [1, chunk_size]863 yield batch, True # → [1, chunk_size]864 self._token_buffer = []865 break866 else:867 # not the last text chunk, need to wait for next text chunk to fill up buffer, then this call ends868 break869 else: # generation of this audio chunk is not finished, continue generating870 continue871 872 # Save current chunk info for sliding_recompute and reindex873 self._chunk_info.append(current_chunk_info)874 self._total_seq_len += condition.shape[1] + len(chunk_generated_tokens)875 876 # Update text_start_pos based on attention type877 if self.attention_type == "sliding_recompute":878 # sliding_recompute: will be reset at next chunk start, update normally here879 self.text_start_pos += prefill_len + len(chunk_generated_tokens)880 elif self.attention_type == "reindex":881 # reindex: position based on actual KV cache length (positions have been reindexed to be continuous)882 if self.past_key_values is not None:883 if hasattr(self.past_key_values, "get_seq_length"):884 self.text_start_pos = self.past_key_values.get_seq_length()885 else:886 self.text_start_pos = self.past_key_values[0][0].shape[2]887 else:888 self.text_start_pos += condition.shape[1] + len(chunk_generated_tokens)889 else:890 self.text_start_pos += condition.shape[1] + len(chunk_generated_tokens)891 # note: remaining tokens in buffer will be kept, and accumulated next time892 893 894# sliding window895@dataclass896class StreamingWindowConfig:897 text_window_high_tokens: int = 8000898 text_window_low_tokens: int = 6000899 900 901@dataclass902class DuplexWindowConfig:903 """duplex sliding window configuration904 905 sliding window mode:906 - "off": disable sliding window907 - "basic": basic sliding window (trigger by cache length)908 - "context": sliding window with context (trigger by unit number, preserve generated text to previous)909 """910 911 # sliding window mode912 sliding_window_mode: str = "off" # "off" / "basic" / "context"913 914 # basic sliding window parameters915 basic_window_high_tokens: int = 8000 # high watermark: trigger sliding window when exceeded916 basic_window_low_tokens: int = 6000 # low watermark: keep to this value after sliding window917 918 # context sliding window parameters919 context_previous_max_tokens: int = 500 # previous maximum token number920 context_max_units: int = 24 # maximum unit number (trigger sliding window when exceeded)921 922 # verification mode (for comparison test)923 verify_mode: bool = False # whether to enable verification log924 925 926def as_dynamic_cache(past_key_values):927 """Convert legacy tuple cache to DynamicCache if needed."""928 if isinstance(past_key_values, DynamicCache):929 return past_key_values930 931 if isinstance(past_key_values, tuple):932 return DynamicCache.from_legacy_cache(past_key_values)933 934 return past_key_values935 936 937def get_kv_cache_length(cache) -> int:938 """Get the sequence length of a KV cache.939 940 Args:941 cache: DynamicCache or tuple-based cache942 943 Returns:944 The number of tokens in the cache945 """946 if cache is None:947 return 0948 949 if isinstance(cache, DynamicCache):950 if not cache.key_cache or not cache.key_cache[0].numel():951 return 0952 return cache.key_cache[0].shape[-2]953 954 if isinstance(cache, tuple):955 return cache[0][0].shape[2]956 957 return 0958 959 960def get_rotary_cos_sin(961 head_dim: int,962 positions: torch.Tensor,963 device: torch.device,964 dtype: torch.dtype,965 rope_theta: float = 10000.0,966 inv_freq_cache: Optional[Dict[Tuple, torch.Tensor]] = None,967) -> Tuple[torch.Tensor, torch.Tensor]:968 """Compute RoPE cos and sin components for given positions.969 970 Args:971 head_dim: Dimension of each attention head972 positions: Position indices tensor973 device: Target device974 dtype: Target dtype975 rope_theta: RoPE base frequency (default 10000.0)976 inv_freq_cache: Optional cache dict for inverse frequencies977 978 Returns:979 Tuple of (cos, sin) tensors with shape [1, 1, seq_len, head_dim]980 """981 cache_key = (head_dim, device)982 983 inv_freq = inv_freq_cache.get(cache_key) if inv_freq_cache is not None else None984 if inv_freq is None or inv_freq.device != device or inv_freq.shape[0] != head_dim // 2:985 exponent = torch.arange(0, head_dim, 2, device=device, dtype=torch.float32) / head_dim986 inv_freq = 1.0 / (rope_theta**exponent)987 if inv_freq_cache is not None:988 inv_freq_cache[cache_key] = inv_freq989 990 positions = positions.to(device=device, dtype=torch.float32)991 angles = torch.einsum("i,j->ij", positions, inv_freq)992 cos = torch.cos(angles)993 sin = torch.sin(angles)994 995 # Use cat instead of repeat_interleave, consistent with model's original RotaryEmbedding996 # Original: emb = torch.cat((freqs, freqs), dim=-1) -> [f0, f1, ..., f_{d/2}, f0, f1, ..., f_{d/2}]997 cos_full = torch.cat([cos, cos], dim=-1).to(dtype=dtype)998 sin_full = torch.cat([sin, sin], dim=-1).to(dtype=dtype)999 cos_full = cos_full.unsqueeze(0).unsqueeze(0)1000 sin_full = sin_full.unsqueeze(0).unsqueeze(0)1001 return cos_full, sin_full1002 1003 1004def realign_rotary_suffix(1005 suffix_keys: torch.Tensor,1006 old_positions: torch.Tensor,1007 new_positions: torch.Tensor,1008 rope_theta: float = 10000.0,1009 inv_freq_cache: Optional[Dict[Tuple, torch.Tensor]] = None,1010) -> torch.Tensor:1011 """Realign RoPE position encoding after cache eviction.1012 1013 When tokens are dropped from the middle of a cache, the suffix tokens1014 need their RoPE embeddings recalculated with new position indices.1015 1016 Args:1017 suffix_keys: Key tensor to realign, shape [batch, heads, seq_len, head_dim]1018 old_positions: Original position indices1019 new_positions: New position indices after eviction1020 rope_theta: RoPE base frequency1021 inv_freq_cache: Optional cache dict for inverse frequencies1022 1023 Returns:1024 Realigned key tensor with same shape as input1025 """1026 if suffix_keys.numel() == 0:1027 return suffix_keys1028 1029 head_dim = suffix_keys.shape[-1]1030 device = suffix_keys.device1031 dtype = suffix_keys.dtype1032 1033 # Compute old position cos/sin1034 cos_old, sin_old = get_rotary_cos_sin(head_dim, old_positions, device, dtype, rope_theta, inv_freq_cache)1035 1036 # Inverse transform: recover original key1037 base = cos_old * suffix_keys - sin_old * rotate_half(suffix_keys)1038 1039 # Compute new position cos/sin1040 cos_new, sin_new = get_rotary_cos_sin(head_dim, new_positions, device, dtype, rope_theta, inv_freq_cache)1041 1042 # Forward transform: re-encode with new positions1043 return cos_new * base + sin_new * rotate_half(base)1044 1045 1046def drop_tokens_from_cache(1047 cache: Optional[DynamicCache | Tuple],1048 length: int,1049 preserve: int,1050 position_offset: int,1051 rope_theta: float = 10000.0,1052 inv_freq_cache: Optional[Dict[Tuple, torch.Tensor]] = None,1053) -> Tuple[Optional[DynamicCache], int, bool]:1054 """Drop tokens from a KV cache while preserving system prompt.1055 1056 Removes tokens in the range [preserve, preserve + length) from the cache,1057 realigning RoPE embeddings for the suffix.1058 1059 Args:1060 cache: DynamicCache or tuple-based cache (will be converted to DynamicCache)1061 length: Number of tokens to drop1062 preserve: Number of tokens to preserve at the start (system prompt)1063 position_offset: Current position offset for RoPE calculation1064 rope_theta: RoPE base frequency1065 inv_freq_cache: Optional cache dict for inverse frequencies1066 1067 Returns:1068 Tuple of (cache, new_position_offset, success)1069 Note: Tuple cache will be converted to DynamicCache. Modification is in-place.1070 """1071 if cache is None or length <= 0:1072 return cache, position_offset, False1073 1074 cache = as_dynamic_cache(cache)1075 1076 total_len = get_kv_cache_length(cache)1077 if total_len <= 0:1078 return cache, position_offset, False1079 1080 preserve = min(preserve, total_len)1081 available = total_len - preserve1082 1083 if available < length:1084 logger.warning(1085 "Cannot drop %d tokens: only %d available (total=%d, preserve=%d)",1086 length,1087 available,1088 total_len,1089 preserve,1090 )1091 return cache, position_offset, False1092 1093 suffix_len = total_len - preserve - length1094 # note: after RoPE reindex, the position of cache has been compressed (from preserve start)1095 # so here should not add position_offset, but use the actual layout of current cache1096 suffix_offset = preserve + length # suffix current position in cache1097 prefix_offset = preserve # suffix new position (follow preserve)1098 1099 # Prepare position tensors for RoPE realignment1100 old_positions = None1101 new_positions = None1102 if suffix_len > 0:1103 device = cache.key_cache[0].device1104 old_positions = torch.arange(1105 suffix_offset,1106 suffix_offset + suffix_len,1107 device=device,1108 dtype=torch.long,1109 )1110 new_positions = torch.arange(1111 prefix_offset,1112 prefix_offset + suffix_len,1113 device=device,1114 dtype=torch.long,1115 )1116 1117 keep_len = total_len - length1118 1119 # Process each layer (in-place modification)1120 for layer_idx in range(len(cache.key_cache)):1121 key_tensor = cache.key_cache[layer_idx]1122 value_tensor = cache.value_cache[layer_idx]1123 1124 if not key_tensor.numel():1125 continue1126 1127 # Preserve prefix (system prompt)1128 prefix_keys = key_tensor[:, :, :preserve, :]1129 prefix_values = value_tensor[:, :, :preserve, :]1130 1131 if suffix_len > 0:1132 # Keep and realign suffix1133 suffix_keys = key_tensor[:, :, preserve + length :, :]1134 suffix_values = value_tensor[:, :, preserve + length :, :]1135 1136 if old_positions is not None and new_positions is not None and suffix_keys.numel():1137 suffix_keys = realign_rotary_suffix(1138 suffix_keys,1139 old_positions,1140 new_positions,1141 rope_theta,1142 inv_freq_cache,1143 )1144 1145 cache.key_cache[layer_idx] = torch.cat([prefix_keys, suffix_keys], dim=-2).contiguous()1146 cache.value_cache[layer_idx] = torch.cat([prefix_values, suffix_values], dim=-2).contiguous()1147 else:1148 cache.key_cache[layer_idx] = prefix_keys.contiguous()1149 cache.value_cache[layer_idx] = prefix_values.contiguous()1150 1151 cache.crop(keep_len)1152 cache._seen_tokens = max(keep_len, 0)1153 1154 new_offset = position_offset + length1155 logger.debug("Dropped %d tokens from cache, new length=%d", length, keep_len)1156 1157 return cache, new_offset, True1158 1159 1160# stream decoder1161def top_k_top_p_filtering(logits, top_k=0, top_p=0.0, filter_value=-float("inf")):1162 logits = logits.clone()1163 1164 # Top-k filtering1165 if top_k > 0:1166 top_k = min(top_k, logits.size(-1))1167 indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1, None]1168 logits[indices_to_remove] = filter_value1169 1170 # Top-p (nucleus) filtering1171 if top_p > 0.0:1172 sorted_logits, sorted_indices = torch.sort(logits, descending=True)1173 probs = F.softmax(sorted_logits, dim=-1)1174 cumulative_probs = torch.cumsum(probs, dim=-1)1175 1176 sorted_indices_to_remove = cumulative_probs > top_p1177 # keep the first token that exceeds top_p1178 sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()1179 sorted_indices_to_remove[..., 0] = 01180 1181 indices_to_remove = sorted_indices[sorted_indices_to_remove]1182 logits[0, indices_to_remove] = filter_value1183 1184 return logits1185 1186 1187class StreamDecoder:1188 def __init__(self, llm, tokenizer, special_token_ids=None, forbidden_token_ids=None):1189 self.m = llm1190 self.tokenizer = tokenizer1191 self.listen_id = self.tokenizer.eos_token_id1192 1193 self.chunk_eos_id = self.tokenizer.convert_tokens_to_ids("<|chunk_eos|>")1194 self.chunk_tts_eos_id = self.tokenizer.convert_tokens_to_ids("<|chunk_tts_eos|>")1195 self.turn_eos_id = self.tokenizer.convert_tokens_to_ids("<|turn_eos|>")1196 self.speak_id = self.tokenizer.convert_tokens_to_ids("<|speak|>")1197 1198 self.special_token_ids = special_token_ids if special_token_ids is not None else []1199 1200 # cache special tokens (used for context sliding window filtering)