# SPDX-License-Identifier: Apache-2.0 """Qwen3-TTS model runner for the OmniScheduler AR stage.""" from __future__ import annotations from typing import Any import torch from sglang.srt.managers.scheduler import GenerationBatchResult from sglang_omni.model_runner.base import ModelRunner from sglang_omni.models.qwen3_omni.talker_model_runner import QwenTalkerModelRunner from sglang_omni.scheduling.types import RequestOutput class Qwen3TTSModelRunner(ModelRunner): """Runs Qwen3-TTS AR steps and stores generated codec frames per request.""" def __init__(self, tp_worker: Any, output_processor: Any): super().__init__(tp_worker, output_processor) self._has_pending_code_step = False def before_prefill( self, forward_batch: Any, schedule_batch: Any, requests: list, ) -> None: del forward_batch, schedule_batch self.model.prepare_decode_buffers(requests) def custom_prefill_forward( self, forward_batch: Any, schedule_batch: Any, requests: list, ) -> GenerationBatchResult | None: del schedule_batch input_embeds = self._build_prefill_input_embeds(forward_batch, requests) return self._forward_with_input_embeds( forward_batch, input_embeds, ) def before_decode( self, forward_batch: Any, schedule_batch: Any, requests: list, *, is_lookahead: bool = False, ) -> None: del is_lookahead del schedule_batch self.model.prepare_decode_buffers(requests) self._write_feedback_buffers(forward_batch, requests) def post_prefill( self, result: Any, forward_batch: Any, schedule_batch: Any, requests: list, ) -> None: self._collect_codes(result, forward_batch, schedule_batch, requests) def post_decode( self, result: Any, forward_batch: Any, schedule_batch: Any, requests: list, ) -> None: self._collect_codes(result, forward_batch, schedule_batch, requests) def sample_before_post_prefill( self, forward_batch: Any, schedule_batch: Any, requests: list ) -> bool: del forward_batch, schedule_batch, requests return True def sample_before_post_decode( self, forward_batch: Any, schedule_batch: Any, requests: list ) -> bool: del forward_batch, schedule_batch, requests return True def _sample_next_token_ids( self, logits_output: Any, forward_batch: Any, schedule_batch: Any, requests: list, ) -> Any: self._install_semantic_sampling_seeds(forward_batch, requests) return super()._sample_next_token_ids( logits_output, forward_batch, schedule_batch, requests, ) def _install_semantic_sampling_seeds( self, forward_batch: Any, requests: list, ) -> None: batch_size = len(requests) forward_batch.sampling_info.sampling_seed = ( self.model._semantic_sampling_seed_tensor[:batch_size] ) def _collect_codes( self, result: Any, forward_batch: Any, schedule_batch: Any, requests: list, ) -> None: self._has_pending_code_step = False if result.next_token_ids is None: return layer0_codes = result.next_token_ids if layer0_codes.ndim == 1: layer0_codes = layer0_codes.unsqueeze(1) hidden = result.logits_output.hidden_states if isinstance(hidden, torch.Tensor) and hidden.ndim == 2: hidden = hidden.unsqueeze(1) semantic_positions = self._sample_positions(forward_batch, layer0_codes.device) self.model.code_predictor_forward( layer0_codes, hidden, semantic_positions=semantic_positions, ) schedule_batch.output_ids = result.next_token_ids self._has_pending_code_step = True def post_process_outputs( self, result: Any, scheduler_output: Any, outputs: dict[str, RequestOutput], ) -> None: del result if not self._has_pending_code_step: return self._has_pending_code_step = False eos_id = int(self.model.config.codec_eos_token_id) for row_idx, sched_req in enumerate(scheduler_output.requests): req_output = outputs[sched_req.request_id] if req_output.data is None or int(req_output.data) == eos_id: continue code_chunk = self.model._output_codes[row_idx].detach().clone() feedback = self.model._output_embeds[row_idx].detach().clone() sched_req.data.output_codes.append(code_chunk) sched_req.data.pending_feedback_queue.append(feedback) def _sample_positions( self, forward_batch: Any, device: torch.device ) -> torch.Tensor: forward_mode = getattr(forward_batch, "forward_mode", None) is_decode = ( forward_mode is not None and hasattr(forward_mode, "is_decode") and bool(forward_mode.is_decode()) ) if is_decode: positions = getattr(forward_batch, "positions", None) if positions is not None: return positions.to(device=device, dtype=torch.long) seq_lens = getattr(forward_batch, "seq_lens", None) if seq_lens is not None: return (seq_lens.to(device=device, dtype=torch.long) - 1).clamp_min(0) positions = getattr(forward_batch, "positions", None) if positions is not None: return positions.to(device=device, dtype=torch.long) raise RuntimeError("Qwen3-TTS subtalker sampling requires semantic positions") def _write_feedback_buffers(self, forward_batch: Any, requests: list) -> None: batch_size = len(requests) if batch_size == 0: return decode_feedback_embedding = self.model._decode_feedback_embedding input_ids = forward_batch.input_ids if input_ids.numel() < batch_size: raise RuntimeError( "Qwen3-TTS decode input_ids must contain one row id per request" ) if batch_size > decode_feedback_embedding.num_embeddings: raise RuntimeError( "Qwen3-TTS decode batch exceeds staged feedback embedding rows" ) row_ids = torch.arange( batch_size, device=input_ids.device, dtype=input_ids.dtype, ) rows = [] for row_idx, sched_req in enumerate(requests): combined = QwenTalkerModelRunner._take_next_decode_input_embed( sched_req=sched_req, device=decode_feedback_embedding.weight.device, dtype=decode_feedback_embedding.weight.dtype, ) if combined is None: token_id = input_ids[row_idx : row_idx + 1].to( device=decode_feedback_embedding.weight.device ) combined = self.model.get_input_embeddings()(token_id).reshape(-1) rows.append(combined) stacked = torch.stack(rows, dim=0).to( device=decode_feedback_embedding.weight.device, dtype=decode_feedback_embedding.weight.dtype, ) with torch.no_grad(): decode_feedback_embedding.weight[:batch_size].copy_(stacked) # During graph decode, input_ids carries staged embedding row ids. input_ids[:batch_size].copy_(row_ids) def _build_prefill_input_embeds( self, forward_batch: Any, requests: list, ) -> torch.Tensor: pieces = [] for sched_req in requests: data = sched_req.data req = data.req req_len = int(req.extend_input_len) prefix_len = len(req.prefix_indices) prompt_embeds = data.prompt_input_embeds if prompt_embeds is None: raise RuntimeError("Qwen3-TTS prefill requires prompt_input_embeds") pieces.append(prompt_embeds[prefix_len : prefix_len + req_len]) return torch.cat(pieces, dim=0).to( device=forward_batch.input_ids.device, dtype=next(self.model.parameters()).dtype, ) def _forward_with_input_embeds( self, forward_batch: Any, input_embeds: torch.Tensor, ) -> GenerationBatchResult: model_runner = self.tp_worker.model_runner model_dtype = next(self.model.parameters()).dtype model_runner.attn_backend.init_forward_metadata(forward_batch) positions = forward_batch.positions if forward_batch.mrope_positions is not None: positions = forward_batch.mrope_positions input_embeds = input_embeds.to( device=forward_batch.input_ids.device, dtype=model_dtype, ) logits_output = self.model( input_ids=forward_batch.input_ids, positions=positions, forward_batch=forward_batch, input_embeds=input_embeds, input_embeds_are_projected=True, ) return GenerationBatchResult( logits_output=logits_output, can_run_cuda_graph=False, )