diff --git a/vllm_omni/core/sched/omni_generation_scheduler.py b/vllm_omni/core/sched/omni_generation_scheduler.py index 96f9f2402..b85229a5d 100644 --- a/vllm_omni/core/sched/omni_generation_scheduler.py +++ b/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: diff --git a/vllm_omni/engine/output_processor.py b/vllm_omni/engine/output_processor.py index c2520f168..697622f12 100644 --- a/vllm_omni/engine/output_processor.py +++ b/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: diff --git a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code_predictor_vllm.py b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code_predictor_vllm.py index ac05d23c8..4820aa2a7 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code_predictor_vllm.py +++ b/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] 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 1bdcd566e..e80f1785f 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 @@ -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 diff --git a/vllm_omni/model_executor/models/qwen3_tts/tokenizer_12hz/modeling_qwen3_tts_tokenizer_v2.py b/vllm_omni/model_executor/models/qwen3_tts/tokenizer_12hz/modeling_qwen3_tts_tokenizer_v2.py index bbec6a525..5719d1024 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/tokenizer_12hz/modeling_qwen3_tts_tokenizer_v2.py +++ b/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) diff --git a/vllm_omni/worker/gpu_model_runner.py b/vllm_omni/worker/gpu_model_runner.py index af2344ce7..058cba444 100644 --- a/vllm_omni/worker/gpu_model_runner.py +++ b/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: