[Hardware][Ascend] Use incremental KV-cache decode for MOSS-TTS depth transformer (#6967)

Signed-off-by: jingchengtian <tjc1995@126.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Co-authored-by: amy-why-3459 <wuhaiyan17@huawei.com>

vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_local_depth.py

@@ -222,6 +222,42 @@ class MossTTSLocalDepthTransformer(nn.Module):
         forward_prefix = self._compiled_forward_prefix or self._forward_prefix
         return forward_prefix(seq_embeds, kv_cache, position)
 
+    @staticmethod
+    def _sample_channel(
+        audio_lm_heads: nn.ModuleList,
+        local_hidden: torch.Tensor,
+        codes: torch.Tensor,
+        *,
+        channel_index: int,
+        repetition_penalty: float,
+        history_per_codebook: list[list[int]] | None,
+        temperature: float,
+        top_k: int,
+        top_p: float,
+        do_sample: bool,
+        generator: torch.Generator | None,
+    ) -> torch.Tensor:
+        """Compute channel logits, apply repetition penalty, sample, store."""
+        channel_logits = audio_lm_heads[channel_index](local_hidden).float()
+        if repetition_penalty != 1.0 and history_per_codebook is not None and channel_index < len(history_per_codebook):
+            hist = history_per_codebook[channel_index]
+            if hist:
+                hist_t = torch.tensor(hist, dtype=torch.long, device=channel_logits.device)
+                sel = channel_logits.index_select(-1, hist_t)
+                pos = sel > 0
+                sel = torch.where(pos, sel / repetition_penalty, sel * repetition_penalty)
+                channel_logits.index_copy_(-1, hist_t, sel)
+        channel_token = _sample_token(
+            channel_logits,
+            temperature,
+            top_k,
+            top_p,
+            do_sample,
+            generator=generator,
+        )
+        codes[:, channel_index] = channel_token
+        return channel_token
+
     @torch.no_grad()
     def generate_frame(
         self,
@@ -296,28 +332,19 @@ class MossTTSLocalDepthTransformer(nn.Module):
 
         codes = backbone_last_hidden.new_zeros((batch_size, n_vq), dtype=torch.long)
         for channel_index in range(n_vq):
-            channel_logits = audio_lm_heads[channel_index](local_hidden).float()
-            if (
-                repetition_penalty != 1.0
-                and history_per_codebook is not None
-                and channel_index < len(history_per_codebook)
-            ):
-                hist = history_per_codebook[channel_index]
-                if hist:
-                    hist_t = torch.tensor(hist, dtype=torch.long, device=channel_logits.device)
-                    sel = channel_logits.index_select(-1, hist_t)
-                    pos = sel > 0
-                    sel = torch.where(pos, sel / repetition_penalty, sel * repetition_penalty)
-                    channel_logits.index_copy_(-1, hist_t, sel)
-            channel_token = _sample_token(
-                channel_logits,
-                temperature,
-                top_k,
-                top_p,
-                do_sample,
+            channel_token = self._sample_channel(
+                audio_lm_heads,
+                local_hidden,
+                codes,
+                channel_index=channel_index,
+                repetition_penalty=repetition_penalty,
+                history_per_codebook=history_per_codebook,
+                temperature=temperature,
+                top_k=top_k,
+                top_p=top_p,
+                do_sample=do_sample,
                 generator=generator,
             )
-            codes[:, channel_index] = channel_token
 
             if channel_index + 1 < n_vq:
                 embeds = audio_embeddings[channel_index](channel_token).to(dtype)[:, None, :]

vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_talker.py

@@ -35,6 +35,7 @@ from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_local_depth impo
     MossTTSLocalDepthTransformer,
 )
 from vllm_omni.model_executor.models.output_templates import OmniOutput
+from vllm_omni.platforms import current_omni_platform
 
 logger = init_logger(__name__)
 
@@ -1353,7 +1354,7 @@ 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_graph_safe = True
+        self.talker_mtp_graph_safe = not current_omni_platform.is_npu()
         self.talker_mtp_output_key = ("audio_codes", "current")
         # ``make_omni_output`` keeps code rows fixed-shape and performs all
         # state updates eagerly, so the runner can safely pack/snapshot them