diff --git a/tests/test_config_factory.py b/tests/test_config_factory.py index de403df2d..6cf8dcd40 100644 --- a/tests/test_config_factory.py +++ b/tests/test_config_factory.py @@ -940,6 +940,36 @@ class TestQwen3TTSPipeline: # Stage 1 uses its per-stage override assert stages[1].yaml_engine_args["model_arch"] == "Qwen3TTSCode2Wav" + def test_subtalker_sampling_params_deep_merge_preserves_base_keys(self): + """Verify subtalker sampling params participate in stage deep-merge.""" + from vllm_omni.config.stage_config import _deep_merge_stage + + base = { + "stage_id": 0, + "subtalker_sampling_params": { + "do_sample": True, + "temperature": 0.9, + "top_k": 50, + "top_p": 1.0, + }, + } + overlay = { + "stage_id": 0, + "subtalker_sampling_params": { + "temperature": 0.7, + "top_k": 32, + }, + } + + merged = _deep_merge_stage(base, overlay) + + assert merged["subtalker_sampling_params"] == { + "do_sample": True, + "temperature": 0.7, + "top_k": 32, + "top_p": 1.0, + } + class TestBaseConfigInheritance: """Test deploy YAML base_config inheritance.""" diff --git a/tests/worker/test_omni_gpu_model_runner.py b/tests/worker/test_omni_gpu_model_runner.py index b2d619315..a74c9ffc2 100644 --- a/tests/worker/test_omni_gpu_model_runner.py +++ b/tests/worker/test_omni_gpu_model_runner.py @@ -41,7 +41,17 @@ class MiMoAudioForConditionalGeneration(torch.nn.Module): class DummyTalkerMTP(torch.nn.Module): """A fake talker_mtp module for deterministic CPU testing.""" - def forward(self, req_input_ids, req_embeds, last_talker_hidden, text_step): + def forward( + self, + req_input_ids, + req_embeds, + last_talker_hidden, + text_step, + do_sample=None, + temperature=None, + top_k=None, + top_p=None, + ): # Deterministic behavior: # - output embeds = input embeds + 1 # - output codes = [[0], [1], ...] @@ -51,6 +61,36 @@ class DummyTalkerMTP(torch.nn.Module): return new_embeds, codes +class CaptureTalkerMTP(torch.nn.Module): + """A fake talker_mtp module that records sampling kwargs.""" + + def __init__(self): + super().__init__() + self.calls = [] + + def forward( + self, + req_input_ids, + req_embeds, + last_talker_hidden, + text_step, + do_sample=None, + temperature=None, + top_k=None, + top_p=None, + ): + self.calls.append( + { + "do_sample": do_sample, + "temperature": temperature, + "top_k": top_k, + "top_p": top_p, + } + ) + codes = torch.zeros((req_embeds.shape[0], 1), dtype=torch.int64) + return req_embeds, codes + + @contextmanager def _noop_forward_context(*args, **kwargs): """A no-op context manager to replace vLLM forward context in CPU tests.""" @@ -80,7 +120,7 @@ def _make_runner(req_ids=("r1", "r2"), hidden_size=4): runner.talker_mtp = DummyTalkerMTP() runner.model = SimpleNamespace(talker_mtp_output_key="code_predictor_codes") - runner.vllm_config = object() + runner.vllm_config = SimpleNamespace(model_config=SimpleNamespace()) # Provide a minimal implementation that returns the expected 4-tuple. def _determine_batch_execution_and_padding(**kwargs): @@ -168,6 +208,43 @@ def test_talker_mtp_forward_cpu_empty_batch_noop(monkeypatch): assert torch.allclose(inputs_embeds, before) +def test_talker_mtp_forward_passes_qwen3_tts_subtalker_sampling_params_to_talker(monkeypatch): + import vllm_omni.worker.gpu_model_runner as mod + + monkeypatch.setattr(mod, "set_forward_context", _noop_forward_context) + + runner = _make_runner(req_ids=("r1",), hidden_size=4) + runner.talker_mtp = CaptureTalkerMTP() + runner.vllm_config = SimpleNamespace( + model_config=SimpleNamespace( + subtalker_sampling_params={ + "do_sample": False, + "temperature": 0.2, + "top_k": 9, + "top_p": 0.55, + } + ) + ) + + def fake_determine(self, num_tokens, num_reqs, num_scheduled_tokens_np, max_num_scheduled_tokens, use_cascade_attn): + batch_desc = SimpleNamespace(num_tokens=int(num_tokens)) + return (False, batch_desc, None, None, None) + + monkeypatch.setattr(runner, "_determine_batch_execution_and_padding", fake_determine.__get__(runner, type(runner))) + + inputs_embeds = torch.zeros((2, 4), dtype=torch.float32) + OmniGPUModelRunner._talker_mtp_forward(runner, ["r1"], inputs_embeds) + + assert runner.talker_mtp.calls == [ + { + "do_sample": False, + "temperature": 0.2, + "top_k": 9, + "top_p": 0.55, + } + ] + + def test_update_intermediate_buffer_writes_to_buffer_and_setattr(monkeypatch): """Validate that _update_intermediate_buffer writes to model_intermediate_buffer (forward path) and mirrors to additional_information_cpu setattr (backward compat).""" diff --git a/vllm_omni/config/model.py b/vllm_omni/config/model.py index 5e38d3065..96a34a8d7 100644 --- a/vllm_omni/config/model.py +++ b/vllm_omni/config/model.py @@ -109,6 +109,7 @@ class OmniModelConfig(ModelConfig): "extra": {}, } ) + subtalker_sampling_params: dict[str, Any] | None = None omni_kv_config: dict | None = None codec_frame_rate_hz: float | None = None task_type: str | None = None diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index 392a550be..9b41660a5 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -393,6 +393,7 @@ class StageDeployConfig: output_connectors: dict[str, str] | None = None input_connectors: dict[str, str] | None = None default_sampling_params: dict[str, Any] | None = None + subtalker_sampling_params: dict[str, Any] | None = None engine_extras: dict[str, Any] = field(default_factory=dict) @@ -463,7 +464,7 @@ def _parse_stage_deploy(stage_data: dict[str, Any]) -> StageDeployConfig: return StageDeployConfig(**kwargs) -_DEEP_MERGE_KEYS = frozenset({"default_sampling_params", "engine_extras", "engine_args"}) +_DEEP_MERGE_KEYS = frozenset({"default_sampling_params", "subtalker_sampling_params", "engine_extras", "engine_args"}) def _deep_merge_stage(base: dict, overlay: dict) -> dict: diff --git a/vllm_omni/deploy/qwen3_tts.yaml b/vllm_omni/deploy/qwen3_tts.yaml index 32dceebd8..5948007a9 100644 --- a/vllm_omni/deploy/qwen3_tts.yaml +++ b/vllm_omni/deploy/qwen3_tts.yaml @@ -43,6 +43,11 @@ stages: max_tokens: 4096 seed: 42 repetition_penalty: 1.05 + subtalker_sampling_params: + do_sample: true + temperature: 0.9 + top_k: 50 + top_p: 1.0 - stage_id: 1 max_num_seqs: 1 diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index 33dee4880..c93bd32c2 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -140,6 +140,7 @@ class OmniEngineArgs(EngineArgs): hf_config_name: str | None = None custom_process_next_stage_input_func: str | None = None stage_connector_spec: dict[str, Any] = field(default_factory=dict) + subtalker_sampling_params: dict[str, Any] | None = None async_chunk: bool = False omni_kv_config: dict | None = None quantization_config: Any | None = None @@ -315,6 +316,7 @@ class OmniEngineArgs(EngineArgs): hf_config_name=self.hf_config_name, custom_process_next_stage_input_func=self.custom_process_next_stage_input_func, stage_connector_config=stage_connector_config, + subtalker_sampling_params=self.subtalker_sampling_params, omni_kv_config=self.omni_kv_config, task_type=self.task_type, ) diff --git a/vllm_omni/model_executor/models/fish_speech/fish_speech_slow_ar.py b/vllm_omni/model_executor/models/fish_speech/fish_speech_slow_ar.py index 62776cbb3..45f7afc69 100644 --- a/vllm_omni/model_executor/models/fish_speech/fish_speech_slow_ar.py +++ b/vllm_omni/model_executor/models/fish_speech/fish_speech_slow_ar.py @@ -649,6 +649,7 @@ class FishSpeechSlowARForConditionalGeneration(nn.Module): input_embeds: torch.Tensor, last_talker_hidden: torch.Tensor, text_step: torch.Tensor, + **kwargs: Any, ) -> tuple[torch.Tensor, torch.Tensor]: """GPU fast-path: run Fast AR to predict residual codebook codes. diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py index f06ecf41d..ce3360375 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py @@ -689,6 +689,7 @@ class Qwen3OmniMoeForConditionalGeneration( input_embeds: torch.Tensor, last_talker_hidden: torch.Tensor, text_step: torch.Tensor, + **kwargs: Any, ): # TODO(Peiqi): not support intermediate_tensors now input_ids = safe_tensor_reshape(input_ids, (input_ids.shape[0], -1)) diff --git a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py index d9cbcf7d4..31fd80622 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py +++ b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py @@ -415,6 +415,10 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module): # In-memory LRU cache for voice extraction artifacts (Base voice clone). self._voice_cache = VoiceEmbeddingCache() + raw_subtalker_sampling = getattr(vllm_config.model_config, "subtalker_sampling_params", None) + self._subtalker_sampling_params: dict[str, Any] = ( + dict(raw_subtalker_sampling) if isinstance(raw_subtalker_sampling, Mapping) else {} + ) # -------------------- vLLM required hooks -------------------- @@ -1638,6 +1642,10 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module): input_embeds: torch.Tensor, last_talker_hidden: torch.Tensor, text_step: torch.Tensor, + do_sample: bool | None = None, + temperature: float | None = None, + top_k: int | None = None, + top_p: float | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: """GPU fast-path used by OmniGPUModelRunner to predict residual codebooks (1..Q-1). Returns (inputs_embeds, audio_codes) for the current step.""" @@ -1656,15 +1664,24 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module): audio_codes = input_ids.reshape(bsz, 1) return (last_id_hidden + text_step).reshape(bsz, -1), audio_codes - # Predict residual codes (1..Q-1) with HF reference sampling params. + subtalker_params = self._subtalker_sampling_params + if do_sample is None: + do_sample = bool(subtalker_params.get("do_sample", True)) + if temperature is None: + temperature = float(subtalker_params.get("temperature", 0.9)) + if top_k is None: + top_k = int(subtalker_params.get("top_k", 50)) + if top_p is None: + top_p = float(subtalker_params.get("top_p", 1.0)) + audio_codes = self.code_predictor( layer0_code=input_ids.reshape(bsz, 1), layer0_embed=last_id_hidden, last_talker_hidden=past_hidden, - do_sample=True, - temperature=0.9, - top_k=50, - top_p=1.0, + do_sample=do_sample, + temperature=temperature, + top_k=top_k, + top_p=top_p, ) # [B, Q] # Map invalid layer-0 ids (e.g. EOS) to PAD=0 so SpeechTokenizer sees only real codes. diff --git a/vllm_omni/platforms/npu/worker/npu_model_runner.py b/vllm_omni/platforms/npu/worker/npu_model_runner.py index 8ef39adfa..310d09311 100644 --- a/vllm_omni/platforms/npu/worker/npu_model_runner.py +++ b/vllm_omni/platforms/npu/worker/npu_model_runner.py @@ -417,10 +417,22 @@ class OmniNPUModelRunner(OmniGPUModelRunner, NPUModelRunner): req_embeds = self.talker_mtp_inputs_embeds.gpu[:num_tokens_padded] last_talker_hidden = self.last_talker_hidden.gpu[:num_tokens_padded] text_step = self.text_step.gpu[:num_tokens_padded] + subtalker_params = getattr(self.vllm_config.model_config, "subtalker_sampling_params", None) + if not isinstance(subtalker_params, dict): + subtalker_params = {} with set_ascend_forward_context( None, self.vllm_config, aclgraph_runtime_mode=_cudagraph_mode, batch_descriptor=batch_desc ): - req_embeds, code_predictor_codes = self.talker_mtp(req_input_ids, req_embeds, last_talker_hidden, text_step) + req_embeds, code_predictor_codes = self.talker_mtp( + req_input_ids, + req_embeds, + last_talker_hidden, + text_step, + do_sample=subtalker_params.get("do_sample"), + temperature=subtalker_params.get("temperature"), + top_k=subtalker_params.get("top_k"), + top_p=subtalker_params.get("top_p"), + ) # code_predictor_codes stays on GPU here; _update_intermediate_buffer # keeps it device-resident when the key is in gpu_resident_buffer_keys. # D2H is deferred to sample_tokens where hidden_states.to("cpu") already diff --git a/vllm_omni/worker/gpu_model_runner.py b/vllm_omni/worker/gpu_model_runner.py index 77f487725..d3ccbaaf3 100644 --- a/vllm_omni/worker/gpu_model_runner.py +++ b/vllm_omni/worker/gpu_model_runner.py @@ -1355,10 +1355,22 @@ class OmniGPUModelRunner(GPUModelRunner): req_embeds = self.talker_mtp_inputs_embeds.gpu[:num_tokens_padded] last_talker_hidden = self.last_talker_hidden.gpu[:num_tokens_padded] text_step = self.text_step.gpu[:num_tokens_padded] + subtalker_params = getattr(self.vllm_config.model_config, "subtalker_sampling_params", None) + if not isinstance(subtalker_params, dict): + subtalker_params = {} with set_forward_context( None, self.vllm_config, cudagraph_runtime_mode=_cudagraph_mode, batch_descriptor=batch_desc ): - req_embeds, code_predictor_codes = self.talker_mtp(req_input_ids, req_embeds, last_talker_hidden, text_step) + req_embeds, code_predictor_codes = self.talker_mtp( + req_input_ids, + req_embeds, + last_talker_hidden, + text_step, + do_sample=subtalker_params.get("do_sample"), + temperature=subtalker_params.get("temperature"), + top_k=subtalker_params.get("top_k"), + top_p=subtalker_params.get("top_p"), + ) # code_predictor_codes stays on GPU here; _update_intermediate_buffer # keeps it device-resident when the key is in gpu_resident_buffer_keys. # D2H is deferred to sample_tokens where hidden_states.to("cpu") already