From d8a72be20601ee9dc3a31dc398ec958591a8095e Mon Sep 17 00:00:00 2001
From: akshatvishu <33392262+akshatvishu@users.noreply.github.com>
Date: Tue, 16 Jun 2026 12:11:40 +0530
Subject: [PATCH] fix(models): remove invalid sampling_metadata parameter from
 logits_processor (#4456)

Signed-off-by: akshatvishu <akshatnayak197@gmail.com>
---
 .../moss_tts/test_talker_logits_processor.py  | 43 +++++++++++++++++++
 .../higgs_audio_v2/higgs_audio_v2_talker.py   |  2 +-
 .../higgs_audio_v3/higgs_audio_v3_talker.py   |  2 +-
 .../moss_tts/modeling_moss_tts_talker.py      |  2 +-
 4 files changed, 46 insertions(+), 3 deletions(-)
 create mode 100644 tests/model_executor/models/moss_tts/test_talker_logits_processor.py

diff --git a/tests/model_executor/models/moss_tts/test_talker_logits_processor.py b/tests/model_executor/models/moss_tts/test_talker_logits_processor.py
new file mode 100644
index 000000000..35c309979
--- /dev/null
+++ b/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 == {}
diff --git a/vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_v2_talker.py b/vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_v2_talker.py
index 5e3883fd2..df636a4bb 100644
--- a/vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_v2_talker.py
+++ b/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,
diff --git a/vllm_omni/model_executor/models/higgs_audio_v3/higgs_audio_v3_talker.py b/vllm_omni/model_executor/models/higgs_audio_v3/higgs_audio_v3_talker.py
index 9d4e05356..79c612dc1 100644
--- a/vllm_omni/model_executor/models/higgs_audio_v3/higgs_audio_v3_talker.py
+++ b/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)
diff --git a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_talker.py b/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_talker.py
index e5fcdb234..e7e100a18 100644
--- a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_talker.py
+++ b/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