[Hardware][Ascend] Adapt Qwen3 TTS for 310P (#4283)

Signed-off-by: zouyizhou <zouyizhou@huawei.com>
Co-authored-by: Canlin Guo <canlinguosdu@gmail.com>

vllm_omni/model_executor/models/qwen3_tts/prompt_embeds_builder.py

@@ -329,6 +329,7 @@ class Qwen3TTSPromptEmbedsBuilder:
         self._tts_pad_embed_buffer = tts_pad_embed
         self._encode_ref_audio_batch_fn = encode_ref_audio_batch
         self._speaker_cache = speaker_cache
+        self._embedding_dtype = torch.bfloat16
 
         self._text_tokenizer: Any | None = None
 
@@ -598,7 +599,7 @@ class Qwen3TTSPromptEmbedsBuilder:
         if isinstance(ref_code, torch.Tensor):
             entry["ref_code"] = ref_code.detach().to("cpu", dtype=torch.long).contiguous()
         if isinstance(ref_spk_embedding, torch.Tensor):
-            entry["ref_spk_embedding"] = ref_spk_embedding.detach().to("cpu", dtype=torch.bfloat16).reshape(-1)
+            entry["ref_spk_embedding"] = ref_spk_embedding.detach().to("cpu", dtype=self._embedding_dtype).reshape(-1)
         if not entry:
             return
         self._ref_audio_artifact_cache[cache_key] = entry
@@ -613,10 +614,11 @@ class Qwen3TTSPromptEmbedsBuilder:
         # CUDA. Ensure the speaker encoder is on the same device/dtype as the
         # main model before running it.
         dev = self._device()
+        dtype = self._embedding_dtype
         try:
             spk_param = next(self._speaker_encoder.parameters())
-            if spk_param.device != dev or spk_param.dtype != torch.bfloat16:
-                self._speaker_encoder.to(device=dev, dtype=torch.bfloat16)
+            if spk_param.device != dev or spk_param.dtype != dtype:
+                self._speaker_encoder.to(device=dev, dtype=dtype)
         except StopIteration:
             pass
 
@@ -641,8 +643,8 @@ class Qwen3TTSPromptEmbedsBuilder:
             fmin=0,
             fmax=12000,
         ).transpose(1, 2)
-        spk = self._speaker_encoder(mels.to(dtype=torch.bfloat16))[0]
-        return spk.to(dtype=torch.bfloat16)
+        spk = self._speaker_encoder(mels.to(dtype=dtype))[0]
+        return spk.to(dtype=dtype)
 
     def encode_ref_audio_batch(
         self,
@@ -1187,17 +1189,18 @@ class Qwen3TTSPromptEmbedsBuilder:
             # NOTE: Do NOT use _as_singleton here — the embedding may be a plain
             # float list (from API via msgspec IPC) that _as_singleton would
             # destructively unwrap to a single scalar.
+            speaker_dtype = self._embedding_dtype
             spk = None
             if voice_clone_prompt is not None:
                 spk = voice_clone_prompt.get("ref_spk_embedding")
             if isinstance(spk, torch.Tensor):
-                speaker_embed = spk.to(device=input_ids.device, dtype=torch.bfloat16).view(1, 1, -1)
+                speaker_embed = spk.to(device=input_ids.device, dtype=speaker_dtype).view(1, 1, -1)
             elif isinstance(spk, (list, np.ndarray)):
-                speaker_embed = torch.tensor(spk, dtype=torch.bfloat16, device=input_ids.device).view(1, 1, -1)
+                speaker_embed = torch.tensor(spk, dtype=speaker_dtype, device=input_ids.device).view(1, 1, -1)
             elif cached_artifacts is not None and isinstance(cached_artifacts.get("ref_spk_embedding"), torch.Tensor):
                 speaker_embed = (
                     cached_artifacts["ref_spk_embedding"]
-                    .to(device=input_ids.device, dtype=torch.bfloat16)
+                    .to(device=input_ids.device, dtype=speaker_dtype)
                     .view(1, 1, -1)
                 )
             elif artifact_only:

vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py

@@ -400,6 +400,7 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
             torch.zeros(1, int(self.talker_config.hidden_size), dtype=model_dtype),
             persistent=False,
         )
+        self._embedding_dtype = torch.bfloat16
 
         tokenizer_config = Qwen3TTSTokenizerV2Config.from_pretrained(
             self.model_path,
@@ -642,7 +643,8 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
         # :meth:`_init_runtime_buffers`. Materialize once on the right
         # device/dtype and reuse for both the prefill placeholder padding
         # and the decode text-step fallback below.
-        tts_pad_embed = self._tts_pad_embed.to(device=input_ids.device, dtype=torch.bfloat16).reshape(1, -1)
+        dtype = self._embedding_dtype
+        tts_pad_embed = self._tts_pad_embed.to(device=input_ids.device, dtype=dtype).reshape(1, -1)
 
         if is_prefill:
             # Prefill (prompt embeddings)
@@ -689,7 +691,7 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
                 pad_n = int(span_len - int(take.shape[0]))
                 pad_rows = tts_pad_embed.reshape(1, -1).to("cpu").expand(pad_n, -1)
                 take = torch.cat([take, pad_rows], dim=0)
-            prompt_embeds = take.to(device=input_ids.device, dtype=torch.bfloat16)
+            prompt_embeds = take.to(device=input_ids.device, dtype=dtype)
             info_update["meta"]["talker_prefill_offset"] = int(offset + span_len)
 
             # When inputs_embeds is set, token ids are ignored by the model but must stay in-vocab for vLLM bookkeeping.
@@ -719,7 +721,7 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
                     tail[text_offset : text_offset + 1]
                     .to(
                         device=input_ids.device,
-                        dtype=torch.bfloat16,
+                        dtype=dtype,
                     )
                     .reshape(1, -1)
                 )
@@ -744,14 +746,14 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
 
         last_hidden = hs.get("last")
         if isinstance(last_hidden, torch.Tensor):
-            past_hidden = last_hidden.to(device=input_ids.device, dtype=torch.bfloat16).reshape(1, -1)
+            past_hidden = last_hidden.to(device=input_ids.device, dtype=dtype).reshape(1, -1)
         else:
             # Defensive: EOS step row is zeroed by the invalid-layer-0 mask and filtered downstream.
             past_hidden = torch.zeros_like(text_step)
 
         # Use OmniGPUModelRunner talker_mtp fast-path for residual codebooks and per-step inputs_embeds update.
         last_id_hidden = self.embed_input_ids(input_ids.reshape(1, 1).to(torch.long)).to(
-            device=input_ids.device, dtype=torch.bfloat16
+            device=input_ids.device, dtype=dtype
         )
         inputs_embeds_out = last_id_hidden.reshape(1, -1)
 
@@ -784,7 +786,7 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
             )
 
         device = input_ids_flat.device
-        dtype = torch.bfloat16
+        dtype = self._embedding_dtype
         # Request-independent constant (see :meth:`_init_runtime_buffers`) —
         # compute once for the batch instead of fetching it per request from
         # ``info_dict["embed"]["tts_pad"]``.
@@ -1043,11 +1045,12 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
         bsz = int(input_ids.shape[0])
         q = int(self.talker_config.num_code_groups)
         dev = input_embeds.device
+        dtype = self._embedding_dtype
 
         input_ids = input_ids.reshape(bsz, 1).to(dtype=torch.long, device=dev)
-        last_id_hidden = input_embeds.reshape(bsz, 1, -1).to(dtype=torch.bfloat16, device=dev)
-        past_hidden = last_talker_hidden.reshape(bsz, 1, -1).to(dtype=torch.bfloat16, device=dev)
-        text_step = text_step.reshape(bsz, 1, -1).to(dtype=torch.bfloat16, device=dev)
+        last_id_hidden = input_embeds.reshape(bsz, 1, -1).to(dtype=dtype, device=dev)
+        past_hidden = last_talker_hidden.reshape(bsz, 1, -1).to(dtype=dtype, device=dev)
+        text_step = text_step.reshape(bsz, 1, -1).to(dtype=dtype, device=dev)
 
         # Residual predictor runs fixed-length (Q-1) steps via the vLLM-native code_predictor.
         max_steps = q - 1