From c56d0da1bbd4382c838d807811b95a204d093710 Mon Sep 17 00:00:00 2001 From: Canlin Guo Date: Fri, 24 Jul 2026 19:08:17 +0800 Subject: [PATCH] [Perf] Refactor MOSS-TTS-Local v1.5 talker and local structure for CUDA Graph (#5197) Signed-off-by: gcanlin Signed-off-by: Canlin Guo Signed-off-by: boatman <109857087+sphinxkkkbc@users.noreply.github.com> Co-authored-by: boatman <109857087+sphinxkkkbc@users.noreply.github.com> --- vllm_omni/deploy/moss_tts_local.yaml | 7 +- .../moss_tts/modeling_moss_tts_codec.py | 8 +- .../moss_tts/modeling_moss_tts_talker.py | 227 ++++++++---------- .../stage_input_processors/moss_tts.py | 5 - 4 files changed, 105 insertions(+), 142 deletions(-) diff --git a/vllm_omni/deploy/moss_tts_local.yaml b/vllm_omni/deploy/moss_tts_local.yaml index ab66df300..deffe4a17 100644 --- a/vllm_omni/deploy/moss_tts_local.yaml +++ b/vllm_omni/deploy/moss_tts_local.yaml @@ -24,7 +24,7 @@ stages: gpu_memory_utilization: 0.60 enable_prefix_caching: false async_scheduling: true - max_num_batched_tokens: 4096 + max_num_batched_tokens: 65536 max_model_len: 4096 devices: "0" output_connectors: @@ -40,9 +40,12 @@ stages: - stage_id: 1 max_num_seqs: 8 gpu_memory_utilization: 0.15 - enforce_eager: true + enforce_eager: false enable_prefix_caching: false async_scheduling: true + compilation_config: + cudagraph_capture_sizes: [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15] + cudagraph_num_of_warmups: 1 max_num_batched_tokens: 65536 max_model_len: 65536 devices: "0" diff --git a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_codec.py b/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_codec.py index 8f1a72e9b..d0f9ea1bc 100644 --- a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_codec.py +++ b/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_codec.py @@ -307,11 +307,15 @@ class MossTTSCodecDecoder(nn.Module): "MossTTS codec requires seq_token_counts; otherwise concatenated " "codec tokens cannot be split per request." ) - if sum(token_counts) != int(ids_flat.shape[0]): + real_token_count = sum(token_counts) + input_token_count = int(ids_flat.shape[0]) + if real_token_count > input_token_count: raise RuntimeError( "MossTTS codec seq_token_counts mismatch: " - f"counts={token_counts}, sum={sum(token_counts)}, input_tokens={int(ids_flat.shape[0])}." + f"counts={token_counts}, sum={real_token_count}, input_tokens={input_token_count}." ) + if real_token_count < input_token_count: + ids_flat = ids_flat[:real_token_count] num_req = len(token_counts) if len(info_list) < num_req: diff --git a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_talker.py b/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_talker.py index 9758c922e..59313ecdf 100644 --- a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_talker.py +++ b/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_talker.py @@ -1321,7 +1321,14 @@ class MossTTSLocalTalkerForGeneration(nn.Module): # additive-fusion gathers (avoids an n_vq-iteration Python loop). self._stacked_audio_emb_w: torch.Tensor | None = None self.mtp_hidden_size = hidden_size - self.talker_mtp_accepts_req_infos = True + self.talker_mtp_graph_safe = True + self.talker_mtp_output_key = ("audio_codes", "current") + self.use_async_omni_output = False + self.eager_omni_postprocess_before_async_output = True + self.omni_pooler_payload_include_hidden = False + self.postprocess_uses_multimodal_outputs = False + self.postprocess_uses_req_infos = False + self._batch_should_continue: torch.Tensor | None = None self.gpu_resident_buffer_keys: set[tuple[str, str]] = { ("audio_codes", "current"), @@ -1373,6 +1380,38 @@ class MossTTSLocalTalkerForGeneration(nn.Module): logits = hidden_states.new_full((num_rows, self.text_vocab_size), float("-inf")) states = self._batch_state or [] + batch_continue_mask = self._batch_should_continue + if isinstance(batch_continue_mask, torch.Tensor) and batch_continue_mask.numel() > 0: + should_by_req = batch_continue_mask.reshape(-1) + if should_by_req.device != logits.device or should_by_req.dtype != torch.bool: + should_by_req = should_by_req.to(device=logits.device, dtype=torch.bool) + if int(should_by_req.numel()) == num_rows: + should_continue = should_by_req + else: + row_req_indices: list[int] = [] + for req_idx, (_state, row_start, row_end) in enumerate( + _iter_state_row_spans( + states, + getattr(self, "_batch_state_spans", None), + num_rows, + ) + ): + if req_idx >= int(should_by_req.shape[0]): + continue + row_req_indices.extend([req_idx] * max(0, row_end - row_start)) + if row_req_indices: + req_t = torch.tensor(row_req_indices, device=logits.device, dtype=torch.long) + should_continue = should_by_req.index_select(0, req_t) + else: + should_continue = torch.ones((num_rows,), device=logits.device, dtype=torch.bool) + if int(should_continue.numel()) != num_rows: + should_continue = torch.ones((num_rows,), device=logits.device, dtype=torch.bool) + zeros = logits.new_zeros((num_rows,)) + neg_inf = logits.new_full((num_rows,), float("-inf")) + logits[:, self.audio_assistant_slot_token_id] = torch.where(should_continue, zeros, neg_inf) + logits[:, self.im_end_token_id] = torch.where(should_continue, neg_inf, zeros) + return logits + for state, row_start, row_end in _iter_state_row_spans( states, getattr(self, "_batch_state_spans", None), @@ -1441,7 +1480,6 @@ class MossTTSLocalTalkerForGeneration(nn.Module): ref_codes = (info_dict.get("codes", {}) or {}).get("ref") ref_offset = int(info_dict.get("ref_offset", 0)) - last_ref_row = None if isinstance(ref_codes, torch.Tensor) and ref_codes.numel() > 0: if ref_codes.dim() == 1 and ref_codes.numel() % self.n_vq == 0: ref_codes = ref_codes.view(-1, self.n_vq) @@ -1451,23 +1489,9 @@ class MossTTSLocalTalkerForGeneration(nn.Module): if chunk.numel() > 0 and chunk.shape[0] == span_len: codes = chunk.to(device=device, dtype=torch.long) embeds = embeds + self._audio_embed(codes) - last_ref_row = codes[-1].detach() - - current_codes = ( - last_ref_row - if last_ref_row is not None - else torch.full((self.n_vq,), self.audio_pad_token_id, dtype=torch.long, device=device) - ) - - max_new_frames = info_dict.get("max_new_frames", [-1])[0] info_update: dict[str, Any] = { - "audio_state": { - "is_stopping": False, - "step": 0, - "max_new_frames": max_new_frames, - }, - "audio_codes": {"current": current_codes, "emit": False}, + "audio_state": {"is_stopping": False}, "ref_offset": ref_offset + span_len, } return input_ids, embeds, info_update @@ -1483,7 +1507,11 @@ class MossTTSLocalTalkerForGeneration(nn.Module): mtp_hidden = last_hidden.to(device=device, dtype=text_embed.dtype).reshape(1, -1) else: mtp_hidden = torch.zeros((1, self.hidden_size), device=device, dtype=text_embed.dtype) - return input_ids, text_embed, {"mtp_inputs": (mtp_hidden, torch.zeros_like(mtp_hidden))} + state = info_dict.get("audio_state", {}) or {} + active = isinstance(state, dict) and not bool(state.get("is_stopping")) + mtp_control = torch.zeros_like(mtp_hidden) + mtp_control[:, :1] = 1.0 if active else 0.0 + return input_ids, text_embed, {"mtp_inputs": (mtp_hidden, mtp_control)} def postprocess(self, hidden_states: torch.Tensor, **_: Any) -> dict[str, Any]: if hidden_states.numel() == 0: @@ -1502,21 +1530,15 @@ class MossTTSLocalTalkerForGeneration(nn.Module): top_k: int | None = None, top_p: float | None = None, generator: torch.Generator | None = None, - req_infos: list[dict[str, Any]] | None = None, **_: Any, - ) -> tuple[torch.Tensor, None]: + ) -> tuple[torch.Tensor, torch.Tensor]: bsz = int(input_embeds.shape[0]) - dev = input_embeds.device - input_embeds_out = input_embeds.reshape(bsz, -1).clone() - req_infos = req_infos or [{} for _ in range(bsz)] - - def write_update(info: dict[str, Any], update: dict[str, Any]) -> None: - for key, value in update.items(): - if isinstance(value, dict): - target = info.setdefault(key, {}) - target.update(value) - else: - info[key] = value + input_embeds_out = input_embeds.reshape(bsz, -1) + last_talker_hidden = last_talker_hidden.reshape(bsz, -1).to( + device=input_embeds.device, + dtype=input_embeds.dtype, + ) + active_mask = text_step.reshape(bsz, -1)[:, :1] > 0 # Match MOSS-TTS Local v1.5's model-card/SGLang defaults. The generic # stage SamplingParams are for the backbone token loop and should not @@ -1526,89 +1548,32 @@ class MossTTSLocalTalkerForGeneration(nn.Module): audio_top_p = 0.8 do_sample_val = True if do_sample is None else bool(do_sample) - for i in range(bsz): - info = req_infos[i] if i < len(req_infos) and isinstance(req_infos[i], dict) else {} - state = dict(info.get("audio_state", {}) or {}) - step = int(state.get("step", 0)) - max_new_frames = int(state.get("max_new_frames", -1)) - acc = (info.get("audio_codes", {}) or {}).get("accumulated") + should_continue_t, new_codes = self.local_transformer.generate_frame( + last_talker_hidden, + self.audio_lm_heads, + self.audio_embeddings, + self.local_text_lm_head, + n_vq=self.n_vq, + do_sample=do_sample_val, + temperature=audio_temperature, + top_k=audio_top_k, + top_p=audio_top_p, + repetition_penalty=1.0, + history_per_codebook=None, + generator=generator, + ) + new_codes = new_codes.to(device=input_embeds.device, dtype=torch.long) + emit_mask = active_mask & should_continue_t.reshape(bsz, 1) + audio_embed = self._audio_embed(new_codes).to(dtype=input_embeds_out.dtype) + input_embeds_out = input_embeds_out + audio_embed * emit_mask.to(dtype=input_embeds_out.dtype) - if state.get("is_stopping"): - write_update( - info, - { - "audio_state": state, - "audio_codes": { - "current": input_embeds.new_empty((0, self.n_vq), dtype=torch.long), - "emit": False, - }, - }, - ) - continue - - rep_window = 50 - if isinstance(acc, torch.Tensor) and acc.numel() > 0: - tail = acc[-rep_window:].long().cpu().tolist() - hist_per_cb = [[row[cb] for row in tail] for cb in range(self.n_vq)] - else: - hist_per_cb = [[] for _ in range(self.n_vq)] - - should_continue_t, new_codes_b = self.local_transformer.generate_frame( - last_talker_hidden[i : i + 1], - self.audio_lm_heads, - self.audio_embeddings, - self.local_text_lm_head, - n_vq=self.n_vq, - do_sample=do_sample_val, - temperature=audio_temperature, - top_k=audio_top_k, - top_p=audio_top_p, - repetition_penalty=1.0, - history_per_codebook=hist_per_cb, - generator=generator, - ) - should_continue = bool(should_continue_t[0].item()) - new_codes = new_codes_b.squeeze(0).to(device=dev, dtype=torch.long) - - if not should_continue or (0 <= max_new_frames <= step): - state["is_stopping"] = True - state["step"] = step + 1 - write_update( - info, - { - "audio_state": state, - "audio_codes": { - "current": input_embeds.new_empty((0, self.n_vq), dtype=torch.long), - "accumulated": acc, - "emit": False, - }, - }, - ) - continue - - if isinstance(acc, torch.Tensor) and acc.numel() > 0: - updated_acc = torch.cat([acc.to(dev), new_codes.unsqueeze(0)], dim=0) - else: - updated_acc = new_codes.unsqueeze(0) - - state["step"] = step + 1 - current = new_codes.unsqueeze(0) - input_embeds_out[i : i + 1] = input_embeds_out[i : i + 1] + self._audio_embed(current).to( - dtype=input_embeds_out.dtype - ) - write_update( - info, - { - "audio_state": state, - "audio_codes": { - "current": current, - "accumulated": updated_acc, - "emit": True, - }, - }, - ) - - return input_embeds_out, None + keep = emit_mask.reshape(-1) + output_codes = torch.where( + keep.unsqueeze(1), + new_codes, + torch.full_like(new_codes, self.audio_pad_token_id), + ) + return input_embeds_out, output_codes # ------------------------------------------------------------------ # Package runner-generated audio frames @@ -1622,6 +1587,7 @@ class MossTTSLocalTalkerForGeneration(nn.Module): if isinstance(model_outputs, OmniOutput): self._batch_state = None self._batch_state_spans = None + self._batch_should_continue = None return model_outputs hidden = model_outputs # (S, H) @@ -1631,46 +1597,41 @@ class MossTTSLocalTalkerForGeneration(nn.Module): for info in info_dicts: if isinstance(info, dict) and not isinstance(info.get("audio_state"), dict): - info["audio_state"] = {"is_stopping": False, "step": 0, "max_new_frames": -1} + info["audio_state"] = {"is_stopping": False} self._batch_state = [(info["audio_state"] if isinstance(info, dict) else {}) for info in info_dicts] self._batch_state_spans = kwargs.get("request_token_spans") - + should_values: list[torch.Tensor] = [] + have_mtp_output = False per_req_codes: list[torch.Tensor] = [] have_codes = False for info in info_dicts: - current = (info.get("audio_codes", {}) or {}).get("current") if isinstance(info, dict) else None - emit = bool((info.get("audio_codes", {}) or {}).get("emit")) if isinstance(info, dict) else False - if emit and isinstance(current, torch.Tensor) and current.numel() > 0: + audio_codes = (info.get("audio_codes", {}) or {}) if isinstance(info, dict) else {} + current = audio_codes.pop("current", None) + if isinstance(current, torch.Tensor) and current.numel() > 0: codes = current.to(device=hidden.device, dtype=torch.long) if codes.dim() == 1: codes = codes.unsqueeze(0) - per_req_codes.append(codes) - have_codes = True + should_continue = codes.ne(self.audio_pad_token_id).any(dim=-1) + should_values.append(should_continue.reshape(-1)[:1]) + have_mtp_output = True + emitted_codes = codes[should_continue] + per_req_codes.append(emitted_codes) + have_codes = have_codes or emitted_codes.numel() > 0 else: + should_values.append(torch.ones((1,), device=hidden.device, dtype=torch.bool)) per_req_codes.append(hidden.new_empty((0, self.n_vq), dtype=torch.long)) + self._batch_should_continue = torch.cat(should_values, dim=0) if have_mtp_output and should_values else None if not have_codes: return OmniOutput(text_hidden_states=hidden, multimodal_outputs={}) # Local-v1.5 emits the current raw code frame (no delay pattern); the # pipeline routes those rows through ``talker2codec_raw_async_chunk``. - per_req_ref_codes: list[torch.Tensor] = [] - for info in info_dicts: - ref_codes = (info.get("codes", {}) or {}).get("ref") if isinstance(info, dict) else None - if isinstance(ref_codes, torch.Tensor) and ref_codes.numel() > 0: - ref_codes = ref_codes.detach() - if ref_codes.dim() == 1 and ref_codes.numel() % self.n_vq == 0: - ref_codes = ref_codes.view(-1, self.n_vq) - if ref_codes.dim() == 2 and ref_codes.shape[1] == self.n_vq: - per_req_ref_codes.append(ref_codes.to(device=hidden.device, dtype=torch.long)) - continue - per_req_ref_codes.append(hidden.new_empty((0, self.n_vq), dtype=torch.long)) - return OmniOutput( text_hidden_states=hidden, multimodal_outputs={ - "codes": {"audio": per_req_codes, "ref": per_req_ref_codes}, + "codes": {"audio": per_req_codes}, }, ) diff --git a/vllm_omni/model_executor/stage_input_processors/moss_tts.py b/vllm_omni/model_executor/stage_input_processors/moss_tts.py index 3dfe2a603..2a7da4a4d 100644 --- a/vllm_omni/model_executor/stage_input_processors/moss_tts.py +++ b/vllm_omni/model_executor/stage_input_processors/moss_tts.py @@ -291,17 +291,12 @@ def talker2codec_raw_async_chunk( transfer_manager.code_prompt_token_ids.pop(req_id, None) transfer_manager.request_payload.pop(req_id, None) return OmniPayloadStruct( - # A non-empty sentinel is required so Stage-1 is scheduled. - # ``code_flat_numel=0`` tells the codec this is a control-only - # finish packet, not an audio code. - codes=CodesStruct(audio=torch.tensor([0], dtype=torch.long)), meta=MetaStruct( req_id=[req_id], left_context_size=0, codec_streaming=True, codec_chunk_frames=0, codec_left_context_frames=0, - code_flat_numel=0, stream_finished=torch.tensor(True, dtype=torch.bool), finished=torch.tensor(True, dtype=torch.bool), ),