[Feature] Add FP8 quantization for Voxtral TTS (#3036)

Signed-off-by: akshatvishu <akshatnayak197@gmail.com>
Signed-off-by: Chen-Yo Sun <chenyo.sun@mistral.ai>
Signed-off-by: akshatvishu <33392262+akshatvishu@users.noreply.github.com>
Co-authored-by: Chen-Yo Sun <chenyo.sun@mistral.ai>
Co-authored-by: SYLAR <125541396+lishunyang12@users.noreply.github.com>

examples/offline_inference/text_to_speech/voxtral_tts/end2end.py

@@ -42,6 +42,7 @@ async def run_streaming(inputs, sampling_params_list, model_name, args, output_d
         model=model_name,
         deploy_config=args.deploy_config,
         log_stats=args.log_stats,
+        quantization=args.quantization,
     )
 
     # Normalize to a list so batch and single-input share the same code path
@@ -194,6 +195,7 @@ def run_non_streaming(inputs, sampling_params_list, model_name, args, output_dir
         model=model_name,
         log_stats=args.log_stats,
         deploy_config=args.deploy_config,
+        quantization=args.quantization,
     )
 
     if args.profiling_mode:
@@ -305,6 +307,12 @@ def parse_args() -> Namespace:
         default=None,
         help="CFG alpha for flow-matching guidance (default: use value from stage config, typically 1.2).",
     )
+    parser.add_argument(
+        "--quantization",
+        type=str,
+        default=None,
+        help="Quantization method (e.g. 'fp8'). Applied to the language model only.",
+    )
     nullify_stage_engine_defaults(parser)
     return parser.parse_args()
 

tests/e2e/offline_inference/test_voxtral_tts.py

@@ -31,7 +31,7 @@ from vllm import SamplingParams
 
 from tests.helpers.mark import hardware_test
 from tests.helpers.runtime import OmniRunner
-from tests.helpers.stage_config import get_deploy_config_path
+from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config
 from vllm_omni.entrypoints.async_omni import AsyncOmni
 
 MODEL = "mistralai/Voxtral-4B-TTS-2603"
@@ -42,14 +42,6 @@ MIN_AUDIO_SAMPLES = 1000
 VOICE = "casual_female"
 TEST_TEXT = "Hello, how are you?"
 
-# (model, stage_config_path, extra_omni_kwargs) for indirect parametrize on
-# ``omni_runner_function`` (function-scoped: must exit before ``AsyncOmni`` test).
-_OMNI_RUNNER_PARAM = (
-    MODEL,
-    STAGE_CONFIG,
-    {"enforce_eager": True},
-)
-
 
 def _compose_request(model_name: str, text: str, voice: str) -> dict:
     """Build the TTS input dict using mistral tokenizer."""
@@ -66,60 +58,83 @@ def _compose_request(model_name: str, text: str, voice: str) -> dict:
     }
 
 
-@pytest.mark.advanced_model
-@pytest.mark.tts
-@pytest.mark.parametrize("omni_runner_function", [_OMNI_RUNNER_PARAM], indirect=True)
-@hardware_test(res={"cuda": "H100"}, num_cards=1)
-def test_voxtral_tts_offline_basic(omni_runner_function: OmniRunner) -> None:
-    """Offline sync path; function-scoped runner so AsyncOmni test does not overlap a live Omni."""
-    omni = omni_runner_function.omni
-    inputs = _compose_request(MODEL, TEST_TEXT, VOICE)
-
-    sampling_params = SamplingParams(max_tokens=2500)
-    sampling_params_list = [sampling_params, sampling_params]
-
-    outputs = list(omni.generate([inputs], sampling_params_list))
-
-    assert len(outputs) > 0, "No outputs generated"
-
-    # Find audio output from the final stage
-    audio_data = None
-    for o in outputs:
-        mm = getattr(o, "multimodal_output", None)
-        if mm and "audio" in mm:
-            audio_data = mm["audio"]
-            break
-
-    assert audio_data is not None, "No audio output found in any stage output"
-
-    # Concatenate audio chunks if returned as a list of tensors
-    if isinstance(audio_data, list):
-        audio_tensor = torch.cat(audio_data)
-    elif isinstance(audio_data, torch.Tensor):
-        audio_tensor = audio_data
-    else:
-        audio_tensor = torch.tensor(audio_data)
-
-    audio_array = audio_tensor.float().cpu().numpy()
-
-    assert len(audio_array) > MIN_AUDIO_SAMPLES, (
-        f"Audio too short: {len(audio_array)} samples, expected > {MIN_AUDIO_SAMPLES}"
-    )
-
-    # Verify audio isn't all zeros / silence
-    assert np.max(np.abs(audio_array)) > 0.01, "Audio appears to be silence"
+def _resolve_stage_config(run_level: str) -> str:
+    """Resolve stage config: strip load_format for advanced_model (real weights)."""
+    if run_level == "advanced_model":
+        return modify_stage_config(
+            STAGE_CONFIG,
+            deletes={
+                "stages": {
+                    0: ["load_format"],
+                    1: ["load_format"],
+                }
+            },
+        )
+    return STAGE_CONFIG
 
 
 @pytest.mark.advanced_model
-@pytest.mark.tts
+@pytest.mark.omni
 @hardware_test(res={"cuda": "H100"}, num_cards=1)
-def test_voxtral_tts_offline_streaming():
+def test_voxtral_tts_offline_basic(run_level):
+    """Test basic Voxtral TTS offline inference with a voice preset."""
+    stage_config = _resolve_stage_config(run_level)
+
+    with OmniRunner(
+        MODEL,
+        stage_configs_path=stage_config,
+        enforce_eager=True,
+    ) as runner:
+        omni = runner.omni
+        inputs = _compose_request(MODEL, TEST_TEXT, VOICE)
+
+        sampling_params = SamplingParams(max_tokens=2500)
+        sampling_params_list = [sampling_params, sampling_params]
+
+        outputs = list(omni.generate([inputs], sampling_params_list))
+
+        assert len(outputs) > 0, "No outputs generated"
+
+        # Find audio output from the final stage
+        audio_data = None
+        for o in outputs:
+            mm = getattr(o, "multimodal_output", None)
+            if mm and "audio" in mm:
+                audio_data = mm["audio"]
+                break
+
+        assert audio_data is not None, "No audio output found in any stage output"
+
+        # Concatenate audio chunks if returned as a list of tensors
+        if isinstance(audio_data, list):
+            audio_tensor = torch.cat(audio_data)
+        elif isinstance(audio_data, torch.Tensor):
+            audio_tensor = audio_data
+        else:
+            audio_tensor = torch.tensor(audio_data)
+
+        audio_array = audio_tensor.float().cpu().numpy()
+
+        assert len(audio_array) > MIN_AUDIO_SAMPLES, (
+            f"Audio too short: {len(audio_array)} samples, expected > {MIN_AUDIO_SAMPLES}"
+        )
+
+        # Verify audio isn't all zeros / silence
+        assert np.max(np.abs(audio_array)) > 0.01, "Audio appears to be silence"
+
+
+@pytest.mark.advanced_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
+def test_voxtral_tts_offline_streaming(run_level):
     """Test AsyncOmni streaming inference for Voxtral TTS."""
 
     async def _run():
+        stage_config = _resolve_stage_config(run_level)
+
         async_omni = AsyncOmni(
             model=MODEL,
-            stage_configs_path=STAGE_CONFIG,
+            stage_configs_path=stage_config,
             stage_init_timeout=300,
             enforce_eager=True,
         )

vllm_omni/model_executor/models/voxtral_tts/voxtral_tts.py

@@ -1,4 +1,5 @@
 from collections.abc import Iterable
+from dataclasses import replace
 from functools import cached_property
 from pathlib import Path
 
@@ -144,6 +145,9 @@ class VoxtralTTSForConditionalGeneration(
         elif self.model_stage == "audio_tokenizer":
             self.requires_raw_input_tokens = True
             self.audio_generation = None
+            if vllm_config.quant_config is not None:
+                vllm_config = replace(vllm_config, quant_config=None)
+                logger.info("Voxtral TTS FP8 routing: stage-1 audio_tokenizer=bf16")
             self.audio_tokenizer = init_vllm_registered_model(
                 vllm_config=vllm_config,
                 hf_config=config,

vllm_omni/model_executor/models/voxtral_tts/voxtral_tts_audio_generation.py

@@ -3,7 +3,7 @@
 
 import math
 from collections.abc import Iterable, Mapping, Sequence
-from dataclasses import dataclass, fields, is_dataclass
+from dataclasses import dataclass, fields, is_dataclass, replace
 from enum import Enum
 from functools import cached_property
 from math import ceil
@@ -58,6 +58,8 @@ from vllm.sequence import IntermediateTensors
 from vllm.tokenizers import cached_tokenizer_from_config
 from vllm.tokenizers.mistral import MistralTokenizer
 
+from vllm_omni.quantization.component_config import ComponentQuantizationConfig
+
 weight_norm = torch.nn.utils.parametrizations.weight_norm
 
 
@@ -197,6 +199,10 @@ def from_nested_dict(cls, d):
     return cls(**kwargs)
 
 
+def _is_fp8_quant_config(quant_config) -> bool:
+    return quant_config is not None and quant_config.get_name() == "fp8"
+
+
 class FeedForward(nn.Module):
     def __init__(
         self,
@@ -916,6 +922,28 @@ class VoxtralTTSAudioGenerationForConditionalGeneration(nn.Module, SupportsMulti
 
         config = vllm_config.model_config.hf_config
         self.config = config
+        quant_config = vllm_config.quant_config
+        if isinstance(quant_config, ComponentQuantizationConfig):
+            component_configs = dict(quant_config.component_configs)
+            component_configs[maybe_prefix(prefix, "audio_tokenizer")] = None
+            component_configs[maybe_prefix(prefix, "acoustic_transformer")] = None
+            quant_config = ComponentQuantizationConfig(
+                component_configs=component_configs,
+                default_config=None,
+            )
+            vllm_config = replace(vllm_config, quant_config=quant_config)
+        elif _is_fp8_quant_config(quant_config):
+            quant_config = ComponentQuantizationConfig(
+                component_configs={
+                    maybe_prefix(prefix, "language_model"): quant_config,
+                    maybe_prefix(prefix, "audio_tokenizer"): None,
+                    maybe_prefix(prefix, "acoustic_transformer"): None,
+                },
+                default_config=None,
+            )
+            vllm_config = replace(vllm_config, quant_config=quant_config)
+            logger.info("Voxtral TTS FP8 routing: language_model=fp8, acoustic_transformer=bf16, audio_tokenizer=bf16")
+
         self.language_model = init_vllm_registered_model(
             vllm_config=vllm_config,
             hf_config=config.text_config,