fix(models): remove invalid sampling_metadata parameter from logits_processor (#4456)

Signed-off-by: akshatvishu <akshatnayak197@gmail.com>

tests/model_executor/models/moss_tts/test_talker_logits_processor.py

@@ -0,0 +1,43 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Regression tests for MOSS-TTS talker logits processing."""
+
+from __future__ import annotations
+
+import pytest
+import torch
+import torch.nn as nn
+
+pytest.importorskip("vllm")
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.tts]
+
+
+class _RecordingLogitsProcessor:
+    def __init__(self) -> None:
+        self.args: tuple[object, ...] | None = None
+        self.kwargs: dict[str, object] | None = None
+
+    def __call__(self, *args: object, **kwargs: object) -> torch.Tensor:
+        self.args = args
+        self.kwargs = kwargs
+        return torch.zeros((1, 4))
+
+
+def test_moss_tts_delay_compute_logits_does_not_forward_sampling_metadata() -> None:
+    from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_talker import (
+        MossTTSDelayTalkerForGeneration,
+    )
+
+    model = MossTTSDelayTalkerForGeneration.__new__(MossTTSDelayTalkerForGeneration)
+    nn.Module.__init__(model)
+    model.text_lm_head = object()
+    model.logits_processor = _RecordingLogitsProcessor()
+    model._batch_state = None
+    hidden_states = torch.randn(1, 4)
+    sampling_metadata = object()
+
+    logits = model.compute_logits(hidden_states, sampling_metadata=sampling_metadata)
+
+    assert logits.shape == (1, 4)
+    assert model.logits_processor.args == (model.text_lm_head, hidden_states)
+    assert model.logits_processor.kwargs == {}

vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_v2_talker.py

@@ -887,7 +887,7 @@ class HiggsAudioV2TalkerForConditionalGeneration(nn.Module):
         # single step at a time; cleared in :meth:`sample` after consumption
         # so a missed call doesn't leak stale tensors.
         self._last_logits_hidden = hidden_states
-        return self.logits_processor(self.lm_head, hidden_states, sampling_metadata)
+        return self.logits_processor(self.lm_head, hidden_states)
 
     def audio_codebook_logits(
         self,

vllm_omni/model_executor/models/higgs_audio_v3/higgs_audio_v3_talker.py

@@ -531,7 +531,7 @@ class HiggsAudioV3TalkerForConditionalGeneration(nn.Module):
 
     def compute_logits(self, hidden_states: torch.Tensor, sampling_metadata: Any = None) -> torch.Tensor:
         self._last_logits_hidden = hidden_states
-        return self.logits_processor(self.lm_head, hidden_states, sampling_metadata)
+        return self.logits_processor(self.lm_head, hidden_states)
 
     def embed_input_ids(self, input_ids: torch.Tensor, **_: Any) -> torch.Tensor:
         safe_ids = torch.where(input_ids < 0, torch.zeros_like(input_ids), input_ids)

vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_talker.py

@@ -194,7 +194,7 @@ class MossTTSDelayTalkerForGeneration(nn.Module):
             hidden_states = hidden_states.text_hidden_states
         if hidden_states is None:
             return None
-        logits = self.logits_processor(self.text_lm_head, hidden_states, sampling_metadata)
+        logits = self.logits_processor(self.text_lm_head, hidden_states)
         if logits is None or self._batch_state is None:
             return logits