[Bugfix] Fix transformers 5.x compat issues in online TTS serving (#1536)

Signed-off-by: linyueqian <linyueqian@outlook.com>

vllm_omni/core/sched/omni_generation_scheduler.py

@@ -21,6 +21,8 @@ from vllm_omni.distributed.omni_connectors.transfer_adapter.chunk_transfer_adapt
 )
 from vllm_omni.outputs import OmniModelRunnerOutput
 
+logger = init_logger(__name__)
+
 
 class OmniGenerationScheduler(VLLMScheduler):
     def __init__(self, *args, **kwargs):
@@ -293,7 +295,7 @@ class OmniGenerationScheduler(VLLMScheduler):
 
         except Exception:
             # If anything goes wrong, leave the original output unchanged
-            init_logger(__name__).exception("Failed to wrap scheduled_new_reqs with OmniNewRequestData")
+            logger.exception("Failed to wrap scheduled_new_reqs with OmniNewRequestData")
 
         return scheduler_output
 
@@ -448,6 +450,10 @@ class OmniGenerationScheduler(VLLMScheduler):
             prompt_logprobs_tensors = prompt_logprobs_dict.get(req_id)
             if new_token_ids or pooler_output is not None or kv_transfer_params or stopped:
                 # Add EngineCoreOutput for this Request.
+                num_cached = request.num_cached_tokens
+                if num_cached < 0:
+                    logger.warning("Negative num_cached_tokens (%d) for request %s, clamping to 0", num_cached, req_id)
+                    num_cached = 0
                 outputs[request.client_index].append(
                     EngineCoreOutput(
                         request_id=req_id,
@@ -460,7 +466,7 @@ class OmniGenerationScheduler(VLLMScheduler):
                         events=request.take_events(),
                         kv_transfer_params=kv_transfer_params,
                         trace_headers=request.trace_headers,
-                        num_cached_tokens=request.num_cached_tokens,
+                        num_cached_tokens=num_cached,
                         num_external_computed_tokens=request.num_external_computed_tokens,
                         routed_experts=routed_experts,
                         num_nans_in_logits=request.num_nans_in_logits,
@@ -482,6 +488,12 @@ class OmniGenerationScheduler(VLLMScheduler):
             requests = [self.requests[req_id] for req_id in failed_kv_load_req_ids]
             self.finish_requests(failed_kv_load_req_ids, RequestStatus.FINISHED_ERROR)
             for request in requests:
+                num_cached = request.num_cached_tokens
+                if num_cached < 0:
+                    logger.warning(
+                        "Negative num_cached_tokens (%d) for request %s, clamping to 0", num_cached, request.request_id
+                    )
+                    num_cached = 0
                 outputs[request.client_index].append(
                     EngineCoreOutput(
                         request_id=request.request_id,
@@ -489,7 +501,7 @@ class OmniGenerationScheduler(VLLMScheduler):
                         finish_reason=request.get_finished_reason(),
                         events=request.take_events(),
                         trace_headers=request.trace_headers,
-                        num_cached_tokens=request.num_cached_tokens,
+                        num_cached_tokens=num_cached,
                     )
                 )
                 if self.chunk_transfer_adapter is not None:

vllm_omni/engine/output_processor.py

@@ -120,6 +120,9 @@ class OmniRequestState(RequestState):
                             # When the audio tensor shape is inconsistent, torch.cat will fail.
                             # We need to use torch.cat in -1 dimension.
                             continue
+                        elif k == "sr":
+                            # Sample rate is a constant scalar, keep last value.
+                            self.mm_accumulated[k] = v[-1]
                         else:
                             self.mm_accumulated[k] = torch.cat(v, dim=0)
                     except Exception:

vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code_predictor_vllm.py

@@ -397,7 +397,8 @@ class Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM(nn.Module):
         bsz = int(inputs_embeds.shape[0])
         qlen = 2
         # Flatten to token-major.
-        hs = inputs_embeds.to(dtype=torch.bfloat16).reshape(bsz * qlen, -1)
+        model_dtype = next(self.model.parameters()).dtype
+        hs = inputs_embeds.to(dtype=model_dtype).reshape(bsz * qlen, -1)
         hs = self.small_to_mtp_projection(hs)
 
         query_lens = torch.full((bsz,), qlen, dtype=torch.int32)
@@ -498,11 +499,15 @@ class Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM(nn.Module):
         for step in range(1, num_groups):
             # Sample or argmax from logits.
             if do_sample and temperature > 0:
-                scaled = logits / temperature
+                scaled = logits.float() / temperature
                 if top_k > 0:
                     topk_vals, _ = scaled.topk(top_k, dim=-1)
                     scaled = scaled.masked_fill(scaled < topk_vals[:, -1:], float("-inf"))
                 probs = torch.softmax(scaled, dim=-1)
+                # Guard against NaN/Inf from lower-precision dtypes.
+                probs = probs.clamp(min=0.0)
+                row_sums = probs.sum(dim=-1, keepdim=True)
+                probs = torch.where(row_sums > 0, probs / row_sums, torch.ones_like(probs) / probs.shape[-1])
                 next_ids = torch.multinomial(probs, num_samples=1)  # [B, 1]
             else:
                 next_ids = logits.argmax(dim=-1, keepdim=True)  # [B, 1]

vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py

@@ -638,12 +638,12 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
 
     def _get_tokenizer(self):
         if self._tokenizer is None:
-            self._tokenizer = AutoTokenizer.from_pretrained(
-                self.model_path,
-                trust_remote_code=True,
-                fix_mistral_regex=True,
-                use_fast=True,
-            )
+            import transformers
+
+            kwargs = dict(trust_remote_code=True, use_fast=True)
+            if transformers.__version__ < "5":
+                kwargs["fix_mistral_regex"] = True
+            self._tokenizer = AutoTokenizer.from_pretrained(self.model_path, **kwargs)
             self._tokenizer.padding_side = "left"
         return self._tokenizer
 

vllm_omni/model_executor/models/qwen3_tts/tokenizer_12hz/modeling_qwen3_tts_tokenizer_v2.py

@@ -48,6 +48,22 @@ from .configuration_qwen3_tts_tokenizer_v2 import (
 logger = logging.get_logger(__name__)
 
 
+def _default_rope_init(config, device=None, seq_len=None, layer_type=None):
+    """Vanilla sinusoidal RoPE (no scaling).
+
+    transformers>=5.0 removed the 'default' entry from
+    ``ROPE_INIT_FUNCTIONS`` (see ``modeling_rope_utils.py`` in
+    huggingface/transformers), but the speech tokenizer config still
+    declares ``rope_type="default"``. This reimplements the original
+    behaviour.
+    """
+    head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
+    inv_freq = 1.0 / (
+        config.rope_theta ** (torch.arange(0, head_dim, 2, dtype=torch.float32, device=device) / head_dim)
+    )
+    return inv_freq, 1.0
+
+
 @dataclass
 @auto_docstring
 class Qwen3TTSTokenizerV2EncoderOutput(ModelOutput):
@@ -256,7 +272,15 @@ class Qwen3TTSTokenizerV2DecoderRotatoryEmbedding(nn.Module):
         self.original_max_seq_len = config.max_position_embeddings
 
         self.config = config
-        self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
+
+        if self.rope_type in ROPE_INIT_FUNCTIONS:
+            self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
+        elif self.rope_type == "default":
+            self.rope_init_fn = _default_rope_init
+        else:
+            raise ValueError(
+                f"Unsupported rope_type '{self.rope_type}'. Expected one of {list(ROPE_INIT_FUNCTIONS)} or 'default'."
+            )
 
         inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device)
         self.register_buffer("inv_freq", inv_freq, persistent=False)

vllm_omni/worker/gpu_model_runner.py

@@ -7,7 +7,6 @@ from vllm.config import CUDAGraphMode
 from vllm.distributed.parallel_state import get_pp_group
 from vllm.forward_context import set_forward_context
 from vllm.logger import init_logger
-from vllm.model_executor.layers.rotary_embedding import MRotaryEmbedding
 from vllm.model_executor.models.interfaces import supports_mrope
 from vllm.model_executor.models.interfaces_base import VllmModelForPooling
 from vllm.sampling_params import SamplingType
@@ -19,6 +18,7 @@ from vllm.v1.worker.gpu_input_batch import CachedRequestState
 from vllm.v1.worker.gpu_model_runner import GPUModelRunner, IntermediateTensors, PerLayerAttnMetadata
 from vllm.v1.worker.ubatch_utils import maybe_create_ubatch_slices
 
+from vllm_omni.model_executor.layers.rotary_embedding.mrope import OmniMRotaryEmbedding as MRotaryEmbedding
 from vllm_omni.model_executor.models.output_templates import OmniOutput
 
 if TYPE_CHECKING: