Voxtral TTS: drop hardcoded CUDA in audio tokenizer; add XPU stage config (#2428)

Signed-off-by: Joshna Medisetty <joshna.medisetty@intel.com>
Signed-off-by: Joshna-Medisetty <joshna.medisetty@intel.com>

vllm_omni/model_executor/models/voxtral_tts/voxtral_tts_audio_tokenizer.py

@@ -17,6 +17,7 @@ from vllm_omni.model_executor.models.voxtral_tts.voxtral_tts_audio_generation im
     MultimodalAudioModelArgs,
     from_nested_dict,
 )
+from vllm_omni.platforms import current_omni_platform
 
 try:
     from flash_attn import flash_attn_func
@@ -954,7 +955,10 @@ class VoxtralTTSAudioTokenizer(nn.Module):
         if x.shape[-1] % self.patch_size != 0:
             pad_length = self.patch_size - (x.shape[-1] % self.patch_size)
             x = F.pad(x, (0, pad_length), mode="constant", value=0)
-        with torch.autocast(dtype=torch.bfloat16, device_type="cuda"):
+        with torch.autocast(
+            device_type=current_omni_platform.device_type,
+            dtype=torch.bfloat16,
+        ):
             # bf16 to use alibi bias in flash attn
             emb = self._forward_encoder(x)  # (b, d, t)
         codes = self.quantizer.encode(emb)  # (b, k, t)
@@ -1095,7 +1099,7 @@ class VoxtralTTSAudioTokenizer(nn.Module):
         for i, chunk in enumerate(all_chunks):
             padded[i, : len(chunk)] = chunk
 
-        audio_codes = padded.to(device=torch.device("cuda"))  # [B, T, K]
+        audio_codes = padded.to(device=current_omni_platform.device_type)  # [B, T, K]
         audio_values = self.decode(audio_codes.transpose(1, 2), dtype=torch.bfloat16)  # [B, 1, T_out]
         audio_values = audio_values.detach().cpu().float().squeeze(1)  # [B, T_out]
         if torch.min(audio_values) < -1.0:

vllm_omni/platforms/xpu/stage_configs/voxtral_tts.yaml

@@ -0,0 +1,111 @@
+# Voxtral TTS — Intel XPU (AR → audio tokenizer). Matches CUDA stage config knobs where noted.
+
+async_chunk: true
+stage_args:
+  - stage_id: 0
+    stage_type: llm
+    runtime:
+      process: true
+      devices: "0"
+      max_batch_size: 1
+    engine_args:
+      max_num_seqs: 32
+      model_stage: audio_generation
+      model_arch: VoxtralTTSForConditionalGeneration
+      worker_type: ar
+      worker_cls: vllm_omni.platforms.xpu.worker.xpu_ar_worker.XPUARWorker
+      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
+      gpu_memory_utilization: 0.7
+      enforce_eager: true
+      trust_remote_code: true
+      async_scheduling: true
+      engine_output_type: latent
+      enable_prefix_caching: false
+      tokenizer_mode: mistral
+      config_format: mistral
+      load_format: mistral
+      skip_mm_profiling: true
+      enable_chunked_prefill: false
+      max_model_len: 4096
+      distributed_executor_backend: "mp"
+      custom_process_next_stage_input_func: vllm_omni.model_executor.stage_input_processors.voxtral_tts.generator2tokenizer_async_chunk
+    output_connectors:
+      to_stage_1: connector_of_shared_memory
+    is_comprehension: true
+    final_output: false
+    final_output_type: text
+    default_sampling_params:
+      temperature: 0.0
+      top_p: 1.0
+      top_k: -1
+      max_tokens: 2048
+      seed: 42
+      detokenize: True
+      repetition_penalty: 1.1
+
+  - stage_id: 1
+    stage_type: llm
+    runtime:
+      process: true
+      devices: "0"
+      max_batch_size: 1
+    engine_args:
+      model_stage: audio_tokenizer
+      model_arch: VoxtralTTSForConditionalGeneration
+      worker_type: generation
+      worker_cls: vllm_omni.platforms.xpu.worker.xpu_generation_worker.XPUGenerationWorker
+      scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
+      async_scheduling: false
+      max_num_seqs: 1
+      gpu_memory_utilization: 0.28
+      enforce_eager: true
+      trust_remote_code: true
+      enable_prefix_caching: false
+      skip_mm_profiling: true
+      engine_output_type: audio
+      tokenizer_mode: mistral
+      config_format: mistral
+      load_format: mistral
+      max_num_batched_tokens: 16384
+      max_model_len: 16384
+      distributed_executor_backend: "mp"
+    engine_input_source: [0]
+    is_comprehension: false
+    final_output: true
+    final_output_type: audio
+    input_connectors:
+      from_stage_0: connector_of_shared_memory
+    tts_args:
+      max_instructions_length: 500
+    default_sampling_params:
+      temperature: 0.9
+      top_p: 0.8
+      top_k: 40
+      max_tokens: 2048
+      seed: 42
+      detokenize: True
+      repetition_penalty: 1.05
+
+runtime:
+  enabled: true
+  defaults:
+    window_size: -1
+    max_inflight: 1
+
+  connectors:
+    connector_of_shared_memory:
+      name: SharedMemoryConnector
+      extra:
+        shm_threshold_bytes: 65536
+        codec_streaming: true
+        connector_get_sleep_s: 0.01
+        connector_get_max_wait_first_chunk: 3000
+        connector_get_max_wait: 300
+        codec_chunk_frames: 25
+        codec_chunk_frames_at_begin: 5
+        codec_left_context_frames: 25
+
+  edges:
+    - from: 0
+      to: 1
+      window_size: -1