diff --git a/vllm_omni/config/model.py b/vllm_omni/config/model.py index 3da7de7d4..683c6bd52 100644 --- a/vllm_omni/config/model.py +++ b/vllm_omni/config/model.py @@ -119,6 +119,7 @@ class OmniModelConfig(ModelConfig): stage_id: int = 0 async_chunk: bool = False + retains_state_across_chunks: bool = False # Stage-1 active stream slots; 0 keeps legacy chunk-level round-robin. active_stream_window: int = 0 model_stage: str = "thinker" diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index 595003ec7..43944084c 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -223,6 +223,9 @@ class StagePipelineConfig: hf_config_name: str | None = None engine_output_type: str | None = None model_arch: str | None = None + # The model keeps per-request execution state while awaiting the next + # async chunk, so the parked request continues to consume model capacity. + retains_state_across_chunks: bool = False sampling_constraints: dict[str, Any] = field(default_factory=dict) custom_process_input_func: str | None = None custom_process_next_stage_input_func: str | None = None @@ -822,6 +825,7 @@ def _build_engine_args( ``engine_extras`` can still carry a stage-specific ``dtype``). """ engine_args: dict[str, Any] = {"model_arch": ps.model_arch or pipeline.model_arch or None} + engine_args["retains_state_across_chunks"] = ps.retains_state_across_chunks if ps.execution_type == StageExecutionType.DIFFUSION and ps.model_arch: engine_args.setdefault("model_class_name", ps.model_arch) if ps.engine_output_type: diff --git a/vllm_omni/core/sched/omni_generation_scheduler.py b/vllm_omni/core/sched/omni_generation_scheduler.py index 06082b0dc..26e2cfb83 100644 --- a/vllm_omni/core/sched/omni_generation_scheduler.py +++ b/vllm_omni/core/sched/omni_generation_scheduler.py @@ -47,6 +47,7 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): self.chunk_transfer_adapter = None if getattr(model_config, "async_chunk", False): self.chunk_transfer_adapter = OmniChunkTransferAdapter(self.vllm_config) + self._retains_state_across_chunks = bool(getattr(model_config, "retains_state_across_chunks", False)) self._pending_finish_reqs: list[Request] = [] self.input_coordinator: OmniSchedulingCoordinator | None = None if uses_full_payload_input_coordinator(model_config): @@ -145,12 +146,17 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): # Fast path selection and scheduling for one-shot generation requests, # independent of pooling_params. - while ( - self.waiting - and token_budget > 0 - and len(self.running) < self.max_num_running_reqs - and self._pause_state == PauseState.UNPAUSED - ): + while self.waiting and token_budget > 0 and self._pause_state == PauseState.UNPAUSED: + # Requests waiting for their next chunk are temporarily absent + # from `running`, but stateful models still retain their model + # runner slot. Mirror vLLM's treatment of + # `num_waiting_for_streaming_input` when enforcing max_num_seqs. + num_running = len(self.running) + if self._retains_state_across_chunks and self.chunk_transfer_adapter is not None: + num_running += self.chunk_transfer_adapter.num_running_waiting_for_chunk + if num_running >= self.max_num_running_reqs: + break + request = self.waiting.peek_request() # OMNI: Skip requests that are not in self.requests if request.request_id not in self.requests or ( diff --git a/vllm_omni/deploy/moss_tts_local.yaml b/vllm_omni/deploy/moss_tts_local.yaml index deffe4a17..7d3ae8601 100644 --- a/vllm_omni/deploy/moss_tts_local.yaml +++ b/vllm_omni/deploy/moss_tts_local.yaml @@ -9,24 +9,25 @@ connectors: shm: name: SharedMemoryConnector extra: - codec_streaming: true connector_get_sleep_s: 0.005 connector_get_max_wait_first_chunk: 1000 connector_get_max_wait: 300 initial_codec_chunk_frames: 1 # 1 frame * 80 ms/frame = 80 ms first audio chunk codec_chunk_frames: 15 # 15 frames * 80 ms/frame = 1.2 s steady chunks codec_left_context_frames: 0 - codec_stream_slots: 8 stages: - stage_id: 0 - max_num_seqs: 8 + max_num_seqs: 64 gpu_memory_utilization: 0.60 enable_prefix_caching: false async_scheduling: true max_num_batched_tokens: 65536 max_model_len: 4096 devices: "0" + # profiler_config: + # profiler: torch + # torch_profiler_dir: /root/vllm-omni-workspace/vllm-omni/results/moss_local_stage0/batch_8 output_connectors: to_stage_1: shm default_sampling_params: @@ -38,17 +39,22 @@ stages: repetition_penalty: 1.1 - stage_id: 1 - max_num_seqs: 8 - gpu_memory_utilization: 0.15 + max_num_seqs: 64 + gpu_memory_utilization: 0.30 enforce_eager: false enable_prefix_caching: false async_scheduling: true compilation_config: - cudagraph_capture_sizes: [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15] + # Reused by the codec as execution-batch buckets; T is derived from + # initial_codec_chunk_frames and codec_chunk_frames. + cudagraph_capture_sizes: [1, 2, 4, 8, 16, 32, 64] cudagraph_num_of_warmups: 1 + # profiler_config: + # profiler: torch + # torch_profiler_dir: /root/vllm-omni-workspace/vllm-omni/results/moss_local_stage1/batch_8 max_num_batched_tokens: 65536 max_model_len: 65536 - devices: "0" + devices: "1" input_connectors: from_stage_0: shm default_sampling_params: diff --git a/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py b/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py index 83def2dc0..54d1477a9 100644 --- a/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py +++ b/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py @@ -550,6 +550,11 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): self._active_streams[request_id] = request return True + @property + def num_running_waiting_for_chunk(self) -> int: + """Count running requests temporarily removed while awaiting a chunk.""" + return len(self.waiting_for_chunk_running_requests) + def _preempt_non_active_running(self, waiting_queue: Any, running_queue: list[Request]) -> None: # Hold non-active running requests in a private deque rather than # routing them back through waiting_queue. Routing through the diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index 4d7622b91..924f41c3f 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -159,6 +159,7 @@ class OmniEngineArgs(EngineArgs): stage_connector_spec: dict[str, Any] = field(default_factory=dict) subtalker_sampling_params: dict[str, Any] | None = None async_chunk: bool = False + retains_state_across_chunks: bool = False # WS-A: Stage-1 active stream slots. 0 = legacy preempt-everything. # Must be declared here so engine_args dict propagation does not silently # drop the value when constructing OmniEngineArgs from kwargs. @@ -341,6 +342,7 @@ class OmniEngineArgs(EngineArgs): # All kwargs below are Omni specific stage_id=self.stage_id, async_chunk=self.async_chunk, + retains_state_across_chunks=self.retains_state_across_chunks, active_stream_window=self.active_stream_window, model_stage=self.model_stage, model_arch=self.model_arch, diff --git a/vllm_omni/model_executor/models/moss_tts/audio_tokenizer_v2.py b/vllm_omni/model_executor/models/moss_tts/audio_tokenizer_v2.py index f3ba18c36..26ea9c567 100644 --- a/vllm_omni/model_executor/models/moss_tts/audio_tokenizer_v2.py +++ b/vllm_omni/model_executor/models/moss_tts/audio_tokenizer_v2.py @@ -102,12 +102,22 @@ class StreamingState: def __post_init__(self): self.exec_mask = torch.ones(self.batch_size, dtype=torch.bool, device=self.device) + self._owns_exec_mask = True + + def bind_exec_mask(self, exec_mask: torch.Tensor) -> None: + if exec_mask.shape != (self.batch_size,): + raise ValueError(f"Expected exec_mask shape ({self.batch_size},), got {tuple(exec_mask.shape)}") + if exec_mask.device != self.device: + raise ValueError(f"Expected exec_mask on {self.device}, got {exec_mask.device}") + self.exec_mask = exec_mask + self._owns_exec_mask = False def set_exec_mask(self, exec_mask: torch.Tensor): self.exec_mask[:] = exec_mask def reset(self, reset_mask: torch.Tensor) -> None: - self.exec_mask[:] = torch.where(reset_mask, torch.ones_like(self.exec_mask), self.exec_mask) + if self._owns_exec_mask: + self.exec_mask.masked_fill_(reset_mask, True) def __enter__(self): # ExitStack expects a context manager; returning self is conventional and useful for debugging. @@ -117,6 +127,28 @@ class StreamingState: pass +@dataclass(frozen=True) +class StreamingExecutionContext: + """Map compact execution rows to persistent streaming-state slots. + + ``state_slot_ids`` and ``valid_rows`` both have shape ``(B_execution,)``. + CUDA Graph padding rows use dedicated scratch slots and ``valid_rows=False``. + """ + + state_slot_ids: torch.Tensor + valid_rows: torch.Tensor + + def validate(self, *, batch_size: int, state_capacity: int, device: torch.device) -> None: + if self.state_slot_ids.shape != (batch_size,): + raise ValueError(f"Expected state_slot_ids shape ({batch_size},), got {tuple(self.state_slot_ids.shape)}") + if self.valid_rows.shape != (batch_size,): + raise ValueError(f"Expected valid_rows shape ({batch_size},), got {tuple(self.valid_rows.shape)}") + if self.state_slot_ids.device != device or self.valid_rows.device != device: + raise ValueError("Streaming execution metadata must be on the same device as decoder inputs.") + if self.state_slot_ids.dtype != torch.long or self.valid_rows.dtype != torch.bool: + raise TypeError("state_slot_ids must be int64 and valid_rows must be bool.") + + class StreamingModule(nn.Module): """Base class for streaming components.""" @@ -447,11 +479,57 @@ class RingKVCache: def reset(self, reset_mask: torch.Tensor) -> None: self.end_offset[:] = torch.where(reset_mask, torch.zeros_like(self.end_offset), self.end_offset) - def complete(self, k: torch.Tensor, v: torch.Tensor, exec_mask: torch.Tensor) -> KVCacheResult: + def reset_slots(self, state_slot_ids: torch.Tensor) -> None: + if state_slot_ids.numel() == 0: + return + slots = state_slot_ids.to(device=self.end_offset.device, dtype=torch.long) + self.end_offset.index_fill_(0, slots, 0) + self.cache.index_fill_(1, slots, 0) + + def complete( + self, + k: torch.Tensor, + v: torch.Tensor, + exec_mask: torch.Tensor | None = None, + execution_context: StreamingExecutionContext | None = None, + ) -> KVCacheResult: B, H, T, D = k.shape if T <= 0: raise ValueError(f"Expected T > 0, got T={T}") + if execution_context is not None: + if not self.respect_exec_mask: + raise RuntimeError("Dynamic state-slot execution does not support weights_per_step attention.") + slots = execution_context.state_slot_ids + valid_rows = execution_context.valid_rows + end_offset = self.end_offset.index_select(0, slots) + row_cache = self.cache.index_select(1, slots) + + indexes = torch.arange(T, device=end_offset.device, dtype=end_offset.dtype) + indexes = (indexes + end_offset.view(-1, 1)) % self.capacity + scatter_indexes = indexes.view(B, 1, T, 1).expand(-1, H, T, D) + row_cache[0].scatter_(2, scatter_indexes, k) + row_cache[1].scatter_(2, scatter_indexes, v) + # Live and graph-padding rows always map to distinct slots. The + # latter map only to scratch state, so this write cannot corrupt a + # request even though dense graph operators still execute it. + self.cache.index_copy_(1, slots, row_cache) + + cache_indexes = torch.arange(self.capacity, device=end_offset.device, dtype=torch.long) + last_offset = end_offset.view(-1, 1) + T - 1 + end_index = last_offset % self.capacity + delta = cache_indexes - end_index + positions = torch.where( + delta <= 0, + last_offset + delta, + last_offset + delta - self.capacity, + ) + next_offset = torch.where(valid_rows, end_offset + T, end_offset) + self.end_offset.index_copy_(0, slots, next_offset) + invalid = cache_indexes >= next_offset.view(-1, 1) + positions = torch.where(invalid, torch.full_like(positions, -1), positions) + return KVCacheResult(row_cache[0], row_cache[1], positions) + indexes = torch.arange(T, device=self.end_offset.device, dtype=self.end_offset.dtype) indexes = indexes + self.end_offset.view(-1, 1) indexes = indexes % self.capacity @@ -479,6 +557,8 @@ class RingKVCache: ) if self.respect_exec_mask: + if exec_mask is None: + raise ValueError("exec_mask is required for fixed-width streaming execution.") self.end_offset[:] = torch.where(exec_mask, self.end_offset + T, self.end_offset) else: self.end_offset.add_(T) @@ -507,6 +587,14 @@ class MHAState(StreamingState): self.kv_cache.reset(reset_mask) self.offset_cpu = 0 + def reset_slots(self, state_slot_ids: torch.Tensor) -> None: + if state_slot_ids.numel() == 0: + return + slots = state_slot_ids.to(device=self.offset.device, dtype=torch.long) + self.offset.index_fill_(0, slots, 0) + if self.kv_cache is not None: + self.kv_cache.reset_slots(slots) + def apply_weights_per_step( modules: nn.ModuleList, @@ -600,7 +688,8 @@ class MossAudioTokenizerMultiheadAttention(StreamingModule): def _init_streaming_state(self, batch_size: int) -> MHAState: in_proj = cast(nn.Linear, self.in_projs[0]) device = cast(torch.device, in_proj.weight.device) - dtype = cast(torch.dtype, in_proj.weight.dtype) + weight_dtype = cast(torch.dtype, in_proj.weight.dtype) + dtype = torch.bfloat16 if device.type == "cuda" else weight_dtype dim_per_head = self.embed_dim // self.num_heads if self.context is None: @@ -625,21 +714,37 @@ class MossAudioTokenizerMultiheadAttention(StreamingModule): offset_cpu=0, ) - def _complete_kv(self, k, v) -> KVCacheResult: + def _complete_kv( + self, + k: torch.Tensor, + v: torch.Tensor, + execution_context: StreamingExecutionContext | None = None, + ) -> KVCacheResult: state = cast(MHAState | None, self._streaming_state) if state is None: return KVCacheResult.from_kv(k, v) if state.kv_cache is None: return KVCacheResult.from_kv(k, v) - return state.kv_cache.complete(k, v, state.exec_mask) + return state.kv_cache.complete(k, v, state.exec_mask, execution_context) - def forward(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor): + def forward( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + execution_context: StreamingExecutionContext | None = None, + ): state = cast(MHAState | None, self._streaming_state) B, T = query.shape[:2] if state is None: offset = torch.zeros(B, device=query.device, dtype=torch.long) offset_cpu = 0 + elif execution_context is not None: + if self.weights_per_step: + raise RuntimeError("Dynamic codec state slots do not support weights_per_step attention.") + offset = state.offset.index_select(0, execution_context.state_slot_ids) + offset_cpu = 0 else: offset = state.offset offset_cpu = state.offset_cpu @@ -652,7 +757,7 @@ class MossAudioTokenizerMultiheadAttention(StreamingModule): if self.rope: q, k = self.rope(q, k, offset, time_before_heads=False) - k, v, pos_k = self._complete_kv(k, v) + k, v, pos_k = self._complete_kv(k, v, execution_context) pos_k = pos_k[:, None] if self.causal: @@ -670,8 +775,12 @@ class MossAudioTokenizerMultiheadAttention(StreamingModule): x = apply_weights_per_step(self.out_projs, self.weights_per_step_schedule, x, offset_cpu) if state is not None: - state.offset[:] = torch.where(state.exec_mask, state.offset + T, state.offset) - state.offset_cpu += T + if execution_context is None: + state.offset[:] = torch.where(state.exec_mask, state.offset + T, state.offset) + state.offset_cpu += T + else: + next_offset = torch.where(execution_context.valid_rows, offset + T, offset) + state.offset.index_copy_(0, execution_context.state_slot_ids, next_offset) return x @@ -688,6 +797,11 @@ class LayerState(StreamingState): super().reset(reset_mask) self.offset_cpu = 0 + def reset_slots(self, state_slot_ids: torch.Tensor) -> None: + # This state is only meaningful for weights_per_step models, which are + # rejected by the dynamic state-slot path. + self.offset_cpu = 0 + class MossAudioTokenizerTransformerLayer(StreamingModule): """Transformer layer with streaming support.""" @@ -763,8 +877,14 @@ class MossAudioTokenizerTransformerLayer(StreamingModule): device = next(iter(self.parameters())).device return LayerState(batch_size, device, offset_cpu=0) - def _ff_block(self, x: torch.Tensor) -> torch.Tensor: + def _ff_block( + self, + x: torch.Tensor, + execution_context: StreamingExecutionContext | None = None, + ) -> torch.Tensor: state = self._streaming_state + if execution_context is not None and self.weights_per_step: + raise RuntimeError("Dynamic codec state slots do not support weights_per_step transformer layers.") offset = state.offset_cpu if isinstance(state, LayerState) else 0 x_orig = x @@ -782,17 +902,25 @@ class MossAudioTokenizerTransformerLayer(StreamingModule): update = self.gating(x) return x_orig.to(update) + self.layer_scale_2(update) - def _sa_block(self, x: torch.Tensor): + def _sa_block( + self, + x: torch.Tensor, + execution_context: StreamingExecutionContext | None = None, + ): x_orig = x x = self.norm1(x) - update = self.self_attn(x, x, x) + update = self.self_attn(x, x, x, execution_context=execution_context) return x_orig.to(update) + self.layer_scale_1(update) - def forward(self, x: torch.Tensor): - x = self._sa_block(x) - x = self._ff_block(x) + def forward( + self, + x: torch.Tensor, + execution_context: StreamingExecutionContext | None = None, + ): + x = self._sa_block(x, execution_context) + x = self._ff_block(x, execution_context) state = self._streaming_state - if state is not None: + if state is not None and execution_context is None: assert isinstance(state, LayerState) state.offset_cpu += x.shape[1] return x @@ -811,6 +939,11 @@ class TransformerState(StreamingState): super().reset(reset_mask) self.offsets[:] = torch.where(reset_mask, torch.zeros_like(self.offsets), self.offsets) + def reset_slots(self, state_slot_ids: torch.Tensor) -> None: + if state_slot_ids.numel() == 0: + return + self.offsets.index_fill_(0, state_slot_ids.to(device=self.offsets.device, dtype=torch.long), 0) + class MossAudioTokenizerTransformer(StreamingModule): """Transformer with streaming/causal support.""" @@ -869,15 +1002,17 @@ class MossAudioTokenizerTransformer(StreamingModule): def forward(self, x: torch.Tensor, *args, **kwargs): B, T, C = x.shape state = self._streaming_state - offsets = ( - torch.zeros(1, dtype=torch.long, device=x.device) - if state is None - else ( - state.offsets - if isinstance(state, TransformerState) - else torch.zeros(1, dtype=torch.long, device=x.device) - ) - ) + execution_context = kwargs.get("execution_context") + if execution_context is not None and not isinstance(execution_context, StreamingExecutionContext): + raise TypeError("execution_context must be a StreamingExecutionContext.") + if state is None: + offsets = torch.zeros(1, dtype=torch.long, device=x.device) + elif not isinstance(state, TransformerState): + raise TypeError(f"Unexpected transformer state: {type(state).__name__}") + elif execution_context is None: + offsets = state.offsets + else: + offsets = state.offsets.index_select(0, execution_context.state_slot_ids) if self.positional_embedding in {"sin", "sin_rope"}: positions = torch.arange(T, device=x.device).view(1, -1, 1) @@ -890,7 +1025,11 @@ class MossAudioTokenizerTransformer(StreamingModule): if state is not None: assert isinstance(state, TransformerState) - state.offsets[:] = torch.where(state.exec_mask, state.offsets + T, state.offsets) + if execution_context is None: + state.offsets[:] = torch.where(state.exec_mask, state.offsets + T, state.offsets) + else: + next_offsets = torch.where(execution_context.valid_rows, offsets + T, offsets) + state.offsets.index_copy_(0, execution_context.state_slot_ids, next_offsets) return x @@ -1235,6 +1374,76 @@ class MossAudioTokenizerResidualLFQ(nn.Module): for _ in range(num_quantizers) ] ) + self.register_buffer("_decode_lut", None, persistent=False) + self.register_buffer("_decode_lut_bias", None, persistent=False) + self.register_buffer("_decode_lut_quantizer_ids", None, persistent=False) + + @torch.no_grad() + def build_decode_lut( + self, + num_quantizers: int, + *, + dtype: torch.dtype, + ) -> None: + """Fold the inference-only LFQ decode projections into code tables.""" + if not 0 < num_quantizers <= len(self.quantizers): + raise ValueError(f"num_quantizers must be in [1, {len(self.quantizers)}], got {num_quantizers}") + + output_proj = self.output_proj + if isinstance(output_proj, nn.Identity): + output_weight = None + output_bias = torch.zeros( + self.output_dim, + device=self.quantizers[0].codebook.weight.device, + dtype=torch.float32, + ) + elif isinstance(output_proj, nn.Conv1d) and output_proj.kernel_size == (1,): + output_weight = output_proj.weight.detach().squeeze(-1).float() + output_bias = ( + output_proj.bias.detach().float() + if output_proj.bias is not None + else torch.zeros( + self.output_dim, + device=output_weight.device, + dtype=torch.float32, + ) + ) + else: + raise TypeError( + "LFQ decoded LUT requires output_proj to be Identity or Conv1d(kernel_size=1), " + f"got {type(output_proj).__name__}" + ) + + decoded_tables: list[torch.Tensor] = [] + for quantizer in self.quantizers[:num_quantizers]: + quantizer = cast(MossAudioTokenizerLFQ, quantizer) + inner_proj = quantizer.out_proj + codebook = quantizer.codebook.weight.detach().float() + if isinstance(inner_proj, nn.Identity): + decoded = codebook + elif isinstance(inner_proj, nn.Conv1d) and inner_proj.kernel_size == (1,): + inner_weight = inner_proj.weight.detach().squeeze(-1).float() + inner_bias = inner_proj.bias.detach().float() if inner_proj.bias is not None else None + decoded = F.linear(codebook, inner_weight, inner_bias) + else: + raise TypeError( + "LFQ decoded LUT requires quantizer out_proj to be Identity or Conv1d(kernel_size=1), " + f"got {type(inner_proj).__name__}" + ) + + if output_weight is not None: + # Apply the shared outer bias once after summing all codebooks. + decoded = F.linear(decoded, output_weight, bias=None) + decoded_tables.append(decoded) + + lut = torch.stack(decoded_tables).to(dtype=dtype) + self._decode_lut = lut + self._decode_lut_bias = output_bias.to(dtype=dtype) + self._decode_lut_quantizer_ids = torch.arange( + num_quantizers, + device=lut.device, + dtype=torch.long, + ).view(-1, 1, 1) @torch.no_grad() def forward( @@ -1276,6 +1485,12 @@ class MossAudioTokenizerResidualLFQ(nn.Module): def decode_codes(self, codes: torch.Tensor) -> torch.Tensor: nq, B, T = codes.shape + if self._decode_lut is not None and nq <= self._decode_lut.shape[0]: + quantizer_ids = self._decode_lut_quantizer_ids[:nq] + decoded = self._decode_lut[quantizer_ids, codes].sum(dim=0) + decoded = decoded + self._decode_lut_bias + return decoded.transpose(1, 2) + emb = torch.zeros(B, self.rvq_dim, T, device=codes.device, dtype=torch.float32) for i, quantizer in enumerate(self.quantizers[:nq]): quantizer = cast(MossAudioTokenizerLFQ, quantizer) @@ -1389,34 +1604,114 @@ class MossAudioTokenizerModel(MossAudioTokenizerPreTrainedModel): f"got current_frame_rate={current_frame_rate}, expected={expected_output_frame_rate}." ) + self._streaming_modules: list[StreamingModule] = [] + self._streaming_exec_mask: torch.Tensor | None = None + self._decoder_state_capacity = 0 self.post_init() - def _start_streaming(self, batch_size: int): - """Start streaming mode for all modules.""" + def _start_streaming(self, batch_size: int, *, decoder_only: bool = False) -> None: + """Start streaming with one shared, address-stable execution mask.""" + if self._streaming_modules: + raise RuntimeError("MOSS Audio Tokenizer is already streaming.") + + root: nn.Module = self.decoder if decoder_only else self + shared_exec_mask = torch.ones( + batch_size, + dtype=torch.bool, + device=next(self.parameters()).device, + ) + streaming_modules: list[StreamingModule] = [] def _start(module): if isinstance(module, StreamingModule): - module._streaming_state = module._init_streaming_state(batch_size) + if module._streaming_state is not None: + raise RuntimeError("A MOSS Audio Tokenizer submodule is already streaming.") + state = module._init_streaming_state(batch_size) + state.bind_exec_mask(shared_exec_mask) + module._streaming_state = state + streaming_modules.append(module) - self.apply(_start) + root.apply(_start) + self._streaming_modules = streaming_modules + self._streaming_exec_mask = shared_exec_mask - def _stop_streaming(self): + def _stop_streaming(self) -> None: """Stop streaming mode for all modules.""" + for module in self._streaming_modules: + module._streaming_state = None + self._streaming_modules = [] + self._streaming_exec_mask = None + self._decoder_state_capacity = 0 - def _stop(module): - if isinstance(module, StreamingModule): - module._streaming_state = None + def initialize_decoder_state_pool(self, state_capacity: int, scratch_capacity: int = 0) -> None: + """Allocate persistent decoder state independently of execution B.""" + if state_capacity <= 0 or scratch_capacity < 0: + raise ValueError(f"Invalid decoder state capacities: state={state_capacity}, scratch={scratch_capacity}.") + self._start_streaming(state_capacity + scratch_capacity, decoder_only=True) + self._decoder_state_capacity = state_capacity + scratch_capacity - self.apply(_stop) + def close_decoder_state_pool(self) -> None: + self._stop_streaming() + + def reset_decoder_state_slots(self, state_slot_ids: torch.Tensor) -> None: + if not self._streaming_modules: + raise RuntimeError("MOSS Audio Tokenizer decoder state pool is not initialized.") + if state_slot_ids.numel() == 0: + return + slots = state_slot_ids.to(device=next(self.parameters()).device, dtype=torch.long) + if not torch.cuda.is_current_stream_capturing(): + slots_cpu = slots.detach().to("cpu") + if int(slots_cpu.min()) < 0 or int(slots_cpu.max()) >= self._decoder_state_capacity: + raise ValueError( + f"Decoder state slots must be in [0, {self._decoder_state_capacity}), got {slots_cpu.tolist()}" + ) + for module in self._streaming_modules: + state = module._streaming_state + reset_slots = getattr(state, "reset_slots", None) + if callable(reset_slots): + reset_slots(slots) + + def decode_streaming_batch( + self, + codes: torch.Tensor, + codes_lengths: torch.Tensor, + state_slot_ids: torch.Tensor, + valid_rows: torch.Tensor, + ) -> MossAudioTokenizerDecoderOutput: + if not self._streaming_modules: + raise RuntimeError("MOSS Audio Tokenizer decoder state pool is not initialized.") + execution_context = StreamingExecutionContext(state_slot_ids=state_slot_ids, valid_rows=valid_rows) + execution_context.validate( + batch_size=int(codes.shape[1]), + state_capacity=self._decoder_state_capacity, + device=codes.device, + ) + return self._decode_frame(codes, codes_lengths, execution_context=execution_context) def _set_streaming_exec_mask(self, exec_mask: torch.Tensor) -> None: - """Set the active slot mask on every live streaming module.""" + """Update the shared active-slot mask with one device copy.""" + shared_exec_mask = self._streaming_exec_mask + if shared_exec_mask is None: + raise RuntimeError("MOSS Audio Tokenizer is not streaming.") + if exec_mask.shape != shared_exec_mask.shape: + raise ValueError(f"Expected exec_mask shape {tuple(shared_exec_mask.shape)}, got {tuple(exec_mask.shape)}") + shared_exec_mask.copy_(exec_mask.to(device=shared_exec_mask.device, dtype=torch.bool), non_blocking=True) - def _set(module): - if isinstance(module, StreamingModule) and module._streaming_state is not None: - module._streaming_state.set_exec_mask(exec_mask.to(module._streaming_state.device)) - - self.apply(_set) + def _reset_streaming_slots(self, reset_mask: torch.Tensor) -> None: + """Reset cached decoder state without traversing the full module tree.""" + shared_exec_mask = self._streaming_exec_mask + if shared_exec_mask is None: + raise RuntimeError("MOSS Audio Tokenizer is not streaming.") + if reset_mask.shape != shared_exec_mask.shape: + raise ValueError( + f"Expected reset_mask shape {tuple(shared_exec_mask.shape)}, got {tuple(reset_mask.shape)}" + ) + reset_mask = reset_mask.to(device=shared_exec_mask.device, dtype=torch.bool) + shared_exec_mask.masked_fill_(reset_mask, True) + for module in self._streaming_modules: + state = module._streaming_state + if state is not None: + state.reset(reset_mask) @contextmanager def streaming(self, batch_size: int = 1): @@ -1427,6 +1722,15 @@ class MossAudioTokenizerModel(MossAudioTokenizerPreTrainedModel): finally: self._stop_streaming() + @contextmanager + def decoder_streaming(self, batch_size: int = 1): + """Stream decoder-only inference without allocating encoder state.""" + self._start_streaming(batch_size, decoder_only=True) + try: + yield + finally: + self._stop_streaming() + def _prepare_waveform_batch(self, wav_list: list[torch.Tensor]) -> tuple[torch.Tensor, torch.Tensor]: if len(wav_list) == 0: raise ValueError("`wav_list` must contain at least one waveform.") @@ -1625,6 +1929,7 @@ class MossAudioTokenizerModel(MossAudioTokenizerPreTrainedModel): self, codes: torch.Tensor, codes_lengths: torch.Tensor | None = None, + execution_context: StreamingExecutionContext | None = None, ) -> MossAudioTokenizerDecoderOutput: """Detokenize discrete tokens into audio waveform.""" nq, B, T = codes.shape @@ -1633,13 +1938,21 @@ class MossAudioTokenizerModel(MossAudioTokenizerPreTrainedModel): if codes_lengths is None: codes_lengths = torch.full((B,), T, device=device, dtype=torch.long) - # Decode from codes - quantizer = cast(MossAudioTokenizerResidualVQ | MossAudioTokenizerResidualLFQ, self.quantizer) - zq = quantizer.decode_codes(codes) + # Keep eager execution and CUDA Graph capture on the same BF16 path. + with torch.autocast(device_type=device.type, dtype=torch.bfloat16, enabled=device.type == "cuda"): + quantizer = cast(MossAudioTokenizerResidualVQ | MossAudioTokenizerResidualLFQ, self.quantizer) + zq = quantizer.decode_codes(codes) - d, d_lengths = zq, codes_lengths - for decoder_module in self.decoder: - d, d_lengths = decoder_module(d, d_lengths) + d, d_lengths = zq, codes_lengths + for decoder_module in self.decoder: + if isinstance(decoder_module, StreamingModule): + d, d_lengths = decoder_module( + d, + d_lengths, + execution_context=execution_context, + ) + else: + d, d_lengths = decoder_module(d, d_lengths) d, d_lengths = self._restore_channels_from_codec(d, d_lengths) return MossAudioTokenizerDecoderOutput(audio=d, audio_lengths=d_lengths) diff --git a/vllm_omni/model_executor/models/moss_tts/cuda_graph_streaming_decoder_wrapper.py b/vllm_omni/model_executor/models/moss_tts/cuda_graph_streaming_decoder_wrapper.py index 1daa5081e..4aeb17785 100644 --- a/vllm_omni/model_executor/models/moss_tts/cuda_graph_streaming_decoder_wrapper.py +++ b/vllm_omni/model_executor/models/moss_tts/cuda_graph_streaming_decoder_wrapper.py @@ -1,12 +1,11 @@ # Copyright 2026 OpenMOSS and the vLLM-Omni team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"). -"""CUDA graph wrapper for MOSS-TTS Local streaming codec decode.""" +"""CUDA graphs for compact MOSS-TTS streaming codec batches.""" from __future__ import annotations import time -from collections.abc import Callable from dataclasses import dataclass import torch @@ -23,31 +22,35 @@ class _CapturedStreamingDecodeGraph: graph: CUDAGraph static_codes: torch.Tensor static_lengths: torch.Tensor + static_state_slot_ids: torch.Tensor + static_valid_rows: torch.Tensor static_audio: torch.Tensor static_audio_lengths: torch.Tensor class CUDAGraphStreamingDecoderWrapper: - """Replay fixed-size streaming ``_decode_frame`` steps with CUDA graphs. + """Replay streaming decode graphs keyed by ``(B_bucket, exact_T)``. - The wrapped codec must already be inside ``codec.streaming(batch_size)``. - Graphs are keyed by the streaming step frame count ``T``. The batch width is - fixed to the codec stream slot count. + Persistent codec state is addressed through ``state_slot_ids``. Graph + padding rows use non-leaseable scratch slots, so graph batch width is no + longer tied to the number of live stream states. """ def __init__( self, codec: nn.Module, *, - batch_size: int, + state_capacity: int, + batch_sizes: list[int], + frame_sizes: list[int], num_quantizers: int, - reset_streaming_state: Callable[[], None] | None = None, ) -> None: self.codec = codec - self.batch_size = int(batch_size) + self.state_capacity = int(state_capacity) + self.batch_sizes = sorted({int(size) for size in batch_sizes if 0 < int(size) <= state_capacity}) + self.frame_sizes = sorted({int(size) for size in frame_sizes if int(size) > 0}) self.num_quantizers = int(num_quantizers) - self.reset_streaming_state = reset_streaming_state - self.graphs: dict[int, _CapturedStreamingDecodeGraph] = {} + self.graphs: dict[tuple[int, int], _CapturedStreamingDecodeGraph] = {} self._pool = None self._warmed_up = False @@ -56,109 +59,135 @@ class CUDAGraphStreamingDecoderWrapper: return bool(self.graphs) @property - def capture_sizes(self) -> list[int]: + def capture_sizes(self) -> list[tuple[int, int]]: return sorted(self.graphs) + @property + def scratch_capacity(self) -> int: + return max(self.batch_sizes, default=0) + @torch.no_grad() - def warmup(self, device: torch.device, capture_sizes: list[int]) -> None: + def warmup(self, device: torch.device) -> None: if self._warmed_up: return self._warmed_up = True - if device.type != "cuda": - return - sizes = sorted({int(size) for size in capture_sizes if int(size) > 0}, reverse=True) - if not sizes: + if device.type != "cuda" or not self.batch_sizes or not self.frame_sizes: return + capture_keys = sorted( + ((batch_size, frame_size) for batch_size in self.batch_sizes for frame_size in self.frame_sizes), + reverse=True, + ) logger.info( - "MOSS-TTS streaming decoder CUDA graph warmup: batch_size=%d n_vq=%d capture_sizes=%s", - self.batch_size, + "MOSS-TTS streaming decoder CUDA graph warmup: n_vq=%d (B,T)=%s", self.num_quantizers, - list(reversed(sizes)), + list(reversed(capture_keys)), ) start_s = time.perf_counter() with torch.cuda.device(device): - for size in sizes: + for batch_size, frame_size in capture_keys: + key = (batch_size, frame_size) try: - self._capture(size, device) - logger.info(" Captured MOSS-TTS streaming decoder CUDA graph for size=%d", size) + self._capture(batch_size, frame_size, device) + logger.info(" Captured MOSS-TTS streaming decoder CUDA graph for (B,T)=%s", key) except Exception: - self.graphs.pop(size, None) + self.graphs.pop(key, None) logger.warning( - " Failed to capture MOSS-TTS streaming decoder CUDA graph for size=%d; " - "this size will use eager decode", - size, + " Failed to capture MOSS-TTS streaming decoder CUDA graph for (B,T)=%s; using eager", + key, exc_info=True, ) logger.info( "MOSS-TTS streaming decoder CUDA graph warmup complete: %d/%d captured in %.1f ms", len(self.graphs), - len(sizes), + len(capture_keys), (time.perf_counter() - start_s) * 1000.0, ) @torch.no_grad() - def _capture(self, size: int, device: torch.device) -> None: + def _capture(self, batch_size: int, frame_size: int, device: torch.device) -> None: codes = torch.zeros( self.num_quantizers, - self.batch_size, - size, + batch_size, + frame_size, dtype=torch.long, device=device, ) - lengths = torch.full((self.batch_size,), size, dtype=torch.long, device=device) - exec_mask = torch.ones(self.batch_size, dtype=torch.bool, device=device) + lengths = torch.zeros(batch_size, dtype=torch.long, device=device) + scratch_slots = self.state_capacity + torch.arange(batch_size, dtype=torch.long, device=device) + valid_rows = torch.zeros(batch_size, dtype=torch.bool, device=device) - self.codec._set_streaming_exec_mask(exec_mask) stream = torch.cuda.Stream() stream.wait_stream(torch.cuda.current_stream()) with torch.cuda.stream(stream): - _ = self.codec._decode_frame(codes, lengths) + _ = self.codec.decode_streaming_batch(codes, lengths, scratch_slots, valid_rows) torch.cuda.current_stream().wait_stream(stream) torch.accelerator.synchronize(device) + self.codec.reset_decoder_state_slots(scratch_slots) - if self.reset_streaming_state is not None: - self.reset_streaming_state() - self.codec._set_streaming_exec_mask(exec_mask) if self._pool is None: self._pool = current_platform.get_global_graph_pool() graph = CUDAGraph() with torch.cuda.graph(graph, pool=self._pool, capture_error_mode="thread_local"): - output = self.codec._decode_frame(codes, lengths) + output = self.codec.decode_streaming_batch(codes, lengths, scratch_slots, valid_rows) if output.audio is None or output.audio_lengths is None: raise RuntimeError("MOSS-TTS streaming decoder CUDA graph capture produced empty audio.") - - self.graphs[size] = _CapturedStreamingDecodeGraph( + self.graphs[(batch_size, frame_size)] = _CapturedStreamingDecodeGraph( graph=graph, static_codes=codes, static_lengths=lengths, + static_state_slot_ids=scratch_slots, + static_valid_rows=valid_rows, static_audio=output.audio, static_audio_lengths=output.audio_lengths, ) + def _select_batch_size(self, actual_batch_size: int) -> int | None: + return next((size for size in self.batch_sizes if size >= actual_batch_size), None) + + def _select_frame_size(self, actual_frame_size: int, allow_padding: bool) -> int | None: + if actual_frame_size in self.frame_sizes: + return actual_frame_size + if not allow_padding: + return None + return next((size for size in self.frame_sizes if size >= actual_frame_size), None) + @torch.no_grad() def decode( self, - codes_step: torch.Tensor, - exec_mask: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor] | None: - if not codes_step.is_cuda or torch.cuda.is_current_stream_capturing(): + codes: torch.Tensor, + state_slot_ids: torch.Tensor, + *, + allow_frame_padding: bool = False, + ) -> tuple[torch.Tensor, torch.Tensor, int] | None: + if not codes.is_cuda or torch.cuda.is_current_stream_capturing(): return None - n_vq, batch_size, step_t = codes_step.shape - if int(n_vq) != self.num_quantizers or int(batch_size) != self.batch_size: + n_vq, actual_batch_size, frame_size = codes.shape + if int(n_vq) != self.num_quantizers or state_slot_ids.shape != (actual_batch_size,): return None - # The streaming codec advances its internal KV offsets by tensor T, so - # bucket-padding a smaller step into a larger captured graph would write - # padded frames into state. Capture exact step sizes instead. - entry = self.graphs.get(int(step_t)) + batch_size = self._select_batch_size(int(actual_batch_size)) + if batch_size is None: + return None + graph_frame_size = self._select_frame_size(int(frame_size), allow_frame_padding) + if graph_frame_size is None: + return None + entry = self.graphs.get((batch_size, graph_frame_size)) if entry is None: return None - self.codec._set_streaming_exec_mask(exec_mask) - entry.static_codes.copy_(codes_step) + entry.static_codes.zero_() + entry.static_codes[:, :actual_batch_size, :frame_size].copy_(codes, non_blocking=True) + entry.static_lengths.zero_() + entry.static_lengths[:actual_batch_size].fill_(int(frame_size)) + entry.static_state_slot_ids.copy_( + self.state_capacity + torch.arange(batch_size, dtype=torch.long, device=entry.static_state_slot_ids.device) + ) + entry.static_state_slot_ids[:actual_batch_size].copy_(state_slot_ids, non_blocking=True) + entry.static_valid_rows.zero_() + entry.static_valid_rows[:actual_batch_size].fill_(True) entry.graph.replay() - return entry.static_audio, entry.static_audio_lengths + return entry.static_audio, entry.static_audio_lengths, int(actual_batch_size) __all__ = ["CUDAGraphStreamingDecoderWrapper"] diff --git a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_codec.py b/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_codec.py index d0f9ea1bc..d6bc525ac 100644 --- a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_codec.py +++ b/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_codec.py @@ -5,7 +5,6 @@ from __future__ import annotations -import contextlib from collections.abc import Iterable from typing import Any @@ -38,104 +37,136 @@ logger = init_logger(__name__) class _MossCodecStreamSession: - """Persistent streaming decode session for vendored MOSS-Audio-Tokenizer-v2.""" + """Persistent state pool with compact per-step execution batches.""" def __init__( self, codec: nn.Module, *, - stream_slots: int, + state_capacity: int, n_vq: int, - cudagraph_capture_sizes: list[int] | None = None, + graph_batch_sizes: list[int] | None = None, + graph_frame_sizes: list[int] | None = None, ) -> None: self._codec = codec - self._stream_slots = int(stream_slots) - self._batch_size = self._stream_slots + self._state_capacity = int(state_capacity) self._n_vq = int(n_vq) self._device = next(codec.parameters()).device - self._free_stream_slots = list(range(self._stream_slots)) - self._exit_stack = contextlib.ExitStack() + self._free_stream_slots = list(reversed(range(self._state_capacity))) + self._leased_slots: set[int] = set() self._closed = False - with torch.no_grad(): - self._exit_stack.enter_context(codec.streaming(self._batch_size)) self._cudagraph_wrapper: CUDAGraphStreamingDecoderWrapper | None = None - capture_sizes = sorted({int(size) for size in (cudagraph_capture_sizes or []) if int(size) > 0}) - if capture_sizes and self._device.type == "cuda": + batch_sizes = sorted({int(size) for size in (graph_batch_sizes or []) if 0 < int(size) <= self._state_capacity}) + frame_sizes = sorted({int(size) for size in (graph_frame_sizes or []) if int(size) > 0}) + scratch_capacity = max(batch_sizes, default=0) if self._device.type == "cuda" else 0 + initialize_state_pool = getattr(codec, "initialize_decoder_state_pool", None) + if not callable(initialize_state_pool): + raise RuntimeError("The streaming codec does not implement a decoder state pool.") + with torch.no_grad(): + initialize_state_pool(self._state_capacity, scratch_capacity) + if batch_sizes and frame_sizes and self._device.type == "cuda": self._cudagraph_wrapper = CUDAGraphStreamingDecoderWrapper( codec, - batch_size=self._batch_size, + state_capacity=self._state_capacity, + batch_sizes=batch_sizes, + frame_sizes=frame_sizes, num_quantizers=self._n_vq, - reset_streaming_state=lambda: self.reset_slots(list(range(self._batch_size))), ) - self._cudagraph_wrapper.warmup(self._device, capture_sizes) - self.reset_slots(list(range(self._batch_size))) + self._cudagraph_wrapper.warmup(self._device) + self.reset_slots(list(range(self._state_capacity + scratch_capacity))) if not self._cudagraph_wrapper.is_ready: self._cudagraph_wrapper = None def acquire(self) -> int | None: if not self._free_stream_slots: return None - return self._free_stream_slots.pop() + slot = self._free_stream_slots.pop() + self._leased_slots.add(slot) + return slot def release(self, slot: int) -> None: if self._closed: return + if slot not in self._leased_slots: + return self.reset_slots([slot]) + self._leased_slots.remove(slot) self._free_stream_slots.append(slot) def reset_slots(self, slots: list[int]) -> None: if not slots: return - reset_mask = torch.zeros(self._batch_size, dtype=torch.bool, device=self._device) - reset_mask[slots] = True - - def _reset(module: nn.Module) -> None: - state = getattr(module, "_streaming_state", None) - if state is not None: - state.reset(reset_mask.to(state.device)) - - with torch.no_grad(): - self._codec.apply(_reset) + reset_streaming_slots = getattr(self._codec, "reset_decoder_state_slots", None) + if callable(reset_streaming_slots): + with torch.no_grad(): + reset_streaming_slots(torch.tensor(slots, device=self._device, dtype=torch.long)) + return + raise RuntimeError("The streaming codec does not implement state-slot reset.") def close(self) -> None: if self._closed: return + close_state_pool = getattr(self._codec, "close_decoder_state_pool", None) with torch.no_grad(): - self._exit_stack.close() + if callable(close_state_pool): + close_state_pool() self._closed = True @torch.no_grad() - def step(self, slot_codes: dict[int, torch.Tensor]) -> dict[int, torch.Tensor]: + def step( + self, + slot_codes: dict[int, torch.Tensor], + *, + terminal_slots: set[int] | None = None, + ) -> dict[int, torch.Tensor]: if not slot_codes: return {} + terminal_slots = terminal_slots or set() + if not terminal_slots.issubset(slot_codes): + raise ValueError("terminal_slots must be a subset of the decode batch slots.") step_lengths = {int(codes.shape[1]) for codes in slot_codes.values()} if len(step_lengths) != 1: raise ValueError(f"MOSS codec streaming step needs uniform T, got {sorted(step_lengths)}") (step_t,) = step_lengths - codes_step = torch.zeros( - self._n_vq, - self._batch_size, - step_t, - dtype=torch.long, - device=self._device, + slots = list(slot_codes) + if any(slot not in self._leased_slots for slot in slots): + raise RuntimeError(f"Streaming decode references an unleased state slot: {slots}") + codes_step = torch.stack( + [slot_codes[slot].to(device=self._device, dtype=torch.long) for slot in slots], + dim=1, ) - codes_lengths = torch.zeros(self._batch_size, dtype=torch.long, device=self._device) - exec_mask = torch.zeros(self._batch_size, dtype=torch.bool, device=self._device) - for slot, codes in slot_codes.items(): - codes_step[:, slot, :] = codes.to(device=self._device, dtype=torch.long) - codes_lengths[slot] = int(codes.shape[1]) - exec_mask[slot] = True + state_slot_ids = torch.tensor(slots, device=self._device, dtype=torch.long) - graph_output: tuple[torch.Tensor, torch.Tensor] | None = None + graph_output: tuple[torch.Tensor, torch.Tensor, int] | None = None if self._cudagraph_wrapper is not None: - graph_output = self._cudagraph_wrapper.decode(codes_step, exec_mask) + # A terminal tail can safely pad T to a larger graph: decoder + # attention is causal, audio_lengths trims the padding, and every + # affected state slot is reset immediately after this step. + graph_output = self._cudagraph_wrapper.decode( + codes_step, + state_slot_ids, + allow_frame_padding=len(terminal_slots) == len(slots), + ) used_cudagraph = graph_output is not None if used_cudagraph: - audio_tensor, lengths_tensor = graph_output + audio_tensor, lengths_tensor, actual_batch_size = graph_output + audio_tensor = audio_tensor[:actual_batch_size] + lengths_tensor = lengths_tensor[:actual_batch_size] else: - self._codec._set_streaming_exec_mask(exec_mask) - result = self._codec._decode_frame(codes_step, codes_lengths) + codes_lengths = torch.full( + (len(slots),), + step_t, + dtype=torch.long, + device=self._device, + ) + valid_rows = torch.ones(len(slots), dtype=torch.bool, device=self._device) + result = self._codec.decode_streaming_batch( + codes_step, + codes_lengths, + state_slot_ids, + valid_rows, + ) if result.audio is None: return {} audio_tensor = result.audio @@ -144,10 +175,10 @@ class _MossCodecStreamSession: audio = audio_tensor.detach().to("cpu", torch.float32) lengths = lengths_tensor.detach().to("cpu") if lengths_tensor is not None else None out: dict[int, torch.Tensor] = {} - for slot in slot_codes: - wav = audio[slot] + for row, slot in enumerate(slots): + wav = audio[row] if lengths is not None: - wav = wav[..., : int(lengths[slot].item())] + wav = wav[..., : int(lengths[row].item())] out[slot] = wav.contiguous() return out @@ -201,17 +232,16 @@ class MossTTSCodecDecoder(nn.Module): self._n_channels: int = 1 self._sr_tensor = torch.tensor(self._OUTPUT_SAMPLE_RATE, dtype=torch.int32) self._stream_session: _MossCodecStreamSession | None = None - self._stream_slots: int = self._connector_int("codec_stream_slots", default=0) + scheduler_cfg = getattr(self.vllm_config, "scheduler_config", None) + self._stream_state_capacity = max(1, int(getattr(scheduler_cfg, "max_num_seqs", 1) or 1)) + self._initial_stream_chunk_frames: int = self._connector_int("initial_codec_chunk_frames", default=0) self._stream_chunk_frames: int = self._connector_int("codec_chunk_frames", default=0) self._stream_max_step_frames: int = self._stream_chunk_frames or 100 self._stream_req_slots: dict[str, int] = {} - self._stream_pending_codes: dict[str, list[torch.Tensor]] = {} - self._stream_starved_reqs: set[str] = set() - self._codec_streaming: bool = self._connector_bool("codec_streaming", default=False) - default_graph_capture_max = self._stream_chunk_frames or min(self._stream_max_step_frames, 32) - default_graph_capture_sizes = list(range(1, max(1, default_graph_capture_max) + 1)) - self._streaming_cudagraph_capture_sizes = self._streaming_cudagraph_capture_sizes_from_compilation_config( - default_graph_capture_sizes + self._async_chunk = bool(getattr(self.vllm_config.model_config, "async_chunk", False)) + self._streaming_graph_batch_sizes = self._streaming_graph_batch_sizes_from_compilation_config() + self._streaming_graph_frame_sizes = sorted( + {frames for frames in (self._initial_stream_chunk_frames, self._stream_chunk_frames) if frames > 0} ) # ------------------------------------------------------------------ @@ -275,7 +305,7 @@ class MossTTSCodecDecoder(nn.Module): }, ) - if self._codec_streaming and runtime_additional_information is None: + if self._async_chunk and runtime_additional_information is None: return OmniOutput( text_hidden_states=None, multimodal_outputs={ @@ -290,8 +320,6 @@ class MossTTSCodecDecoder(nn.Module): streaming_work: list[tuple[int, str, torch.Tensor, bool]] = [] if input_ids is None or input_ids.numel() == 0: - for i, wav in self._finish_empty_streaming_requests(info_list).items(): - audios[i] = wav.reshape(-1) if wav.ndim == 1 or int(wav.shape[0]) == 1 else wav return OmniOutput( text_hidden_states=None, multimodal_outputs={"model_outputs": audios, "sr": srs}, @@ -341,12 +369,7 @@ class MossTTSCodecDecoder(nn.Module): continue meta = (info.get("meta", {}) if isinstance(info, dict) else {}) or {} finished = bool(meta.get("stream_finished", meta.get("finished", False))) - streaming_enabled = bool(meta.get("codec_streaming", self._codec_streaming)) - code_flat_numel = meta.get("code_flat_numel") - if streaming_enabled and finished and code_flat_numel is not None and int(code_flat_numel) == 0: - for _, wav in self._finish_empty_streaming_requests([info]).items(): - audios[i] = wav.reshape(-1) if wav.ndim == 1 or int(wav.shape[0]) == 1 else wav - continue + streaming_enabled = self._async_chunk if seg.numel() % self._n_vq != 0: logger.warning( "MossTTS codec input length %d not divisible by n_vq %d; skipping.", @@ -414,40 +437,6 @@ class MossTTSCodecDecoder(nn.Module): multimodal_outputs={"model_outputs": audios, "sr": srs}, ) - def _finish_empty_streaming_requests(self, info_list: list[dict[str, Any]]) -> dict[int, torch.Tensor]: - """Release codec stream state for empty finish sentinels. - - Stage-0 can finish on a step that emits no new audio frame. The stage - input processor forwards that as an empty payload with finished=true. - If a request never acquired a stream slot, do not offline-decode its - buffered codes here: streaming requests must only emit client deltas. - """ - session = self._stream_session - outputs: dict[int, torch.Tensor] = {} - if session is None: - return outputs - for i, info in enumerate(info_list): - if not isinstance(info, dict): - continue - meta = (info.get("meta", {}) or {}) if isinstance(info.get("meta", {}), dict) else {} - if not bool(meta.get("codec_streaming", self._codec_streaming)): - continue - finished = bool(meta.get("stream_finished", meta.get("finished", False))) - if not finished: - continue - req_key = self._runtime_request_key(info, meta, i) - slot = self._stream_req_slots.get(req_key) - pending = req_key in self._stream_pending_codes - if slot is not None or pending: - if pending and slot is None: - logger.warning( - "MOSS codec stream request %s finished before a codec stream slot became available; " - "dropping buffered codes instead of offline-decoding a non-delta waveform.", - req_key, - ) - self._finish_stream_request(req_key, session, slot) - return outputs - @staticmethod def _normalize_seq_token_counts(value: Any) -> list[int] | None: if value is None: @@ -486,15 +475,12 @@ class MossTTSCodecDecoder(nn.Module): return None if self._stream_session is not None: return self._stream_session - slots = self._stream_slots - if slots <= 0: - scheduler_cfg = getattr(self.vllm_config, "scheduler_config", None) - slots = int(getattr(scheduler_cfg, "max_num_seqs", 1) or 1) self._stream_session = _MossCodecStreamSession( self._codec, - stream_slots=max(1, slots), + state_capacity=self._stream_state_capacity, n_vq=self._n_vq, - cudagraph_capture_sizes=self._streaming_cudagraph_capture_sizes, + graph_batch_sizes=self._streaming_graph_batch_sizes, + graph_frame_sizes=self._streaming_graph_frame_sizes, ) return self._stream_session @@ -511,40 +497,23 @@ class MossTTSCodecDecoder(nn.Module): max_step_frames = max(1, int(self._stream_max_step_frames)) for output_index, request_id, codes_nq_t, finished in items: - pending = self._stream_pending_codes.get(request_id) slot = self._stream_req_slots.get(request_id) if slot is None: slot = session.acquire() if slot is None: - self._append_stream_pending(request_id, codes_nq_t) - if request_id not in self._stream_starved_reqs: - logger.warning( - "MOSS codec streaming slots exhausted; buffering %s until a stream slot is available.", - request_id, - ) - self._stream_starved_reqs.add(request_id) - if finished: - logger.warning( - "MOSS codec stream request %s finished before a codec stream slot became available; " - "dropping buffered codes instead of offline-decoding a non-delta waveform.", - request_id, - ) - self._finish_stream_request(request_id, session, None) - continue + raise RuntimeError( + "MOSS codec state capacity exhausted. Stage-1 admission must count idle live streams " + f"against max_num_seqs={self._stream_state_capacity}; refusing to drop audio codes." + ) self._stream_req_slots[request_id] = slot - if pending: - self._append_stream_pending(request_id, codes_nq_t) - replay_codes = self._pop_stream_pending(request_id) - wav = self._decode_stream_slot_sequence(session, slot, replay_codes) - if wav is not None: - outputs[output_index] = wav - if finished: - self._finish_stream_request(request_id, session, slot) - continue - if int(codes_nq_t.shape[1]) > max_step_frames: - wav = self._decode_stream_slot_sequence(session, slot, codes_nq_t) + wav = self._decode_stream_slot_sequence( + session, + slot, + codes_nq_t, + terminal=finished, + ) if wav is not None: outputs[output_index] = wav if finished: @@ -557,7 +526,8 @@ class MossTTSCodecDecoder(nn.Module): for group in grouped.values(): plan = {slot: codes_nq_t for _, _, slot, codes_nq_t, _ in group} - decoded = session.step(plan) + terminal_slots = {slot for _, _, slot, _, finished in group if finished} + decoded = session.step(plan, terminal_slots=terminal_slots) for output_index, request_id, slot, _, finished in group: wav = decoded.get(slot) if wav is not None: @@ -567,22 +537,13 @@ class MossTTSCodecDecoder(nn.Module): return outputs - def _append_stream_pending(self, request_id: str, codes_nq_t: torch.Tensor) -> None: - self._stream_pending_codes.setdefault(request_id, []).append( - codes_nq_t.detach().to("cpu", torch.long).contiguous() - ) - - def _pop_stream_pending(self, request_id: str) -> torch.Tensor: - pending = self._stream_pending_codes.pop(request_id, []) - if not pending: - return torch.empty((self._n_vq, 0), dtype=torch.long) - return torch.cat(pending, dim=1).contiguous() - def _decode_stream_slot_sequence( self, session: _MossCodecStreamSession, slot: int, codes_nq_t: torch.Tensor, + *, + terminal: bool = False, ) -> torch.Tensor | None: if codes_nq_t.numel() == 0: return None @@ -590,7 +551,11 @@ class MossTTSCodecDecoder(nn.Module): parts: list[torch.Tensor] = [] for start in range(0, int(codes_nq_t.shape[1]), max_step_frames): chunk = codes_nq_t[:, start : start + max_step_frames] - decoded = session.step({slot: chunk}) + is_terminal_chunk = terminal and start + max_step_frames >= int(codes_nq_t.shape[1]) + decoded = session.step( + {slot: chunk}, + terminal_slots={slot} if is_terminal_chunk else None, + ) wav = decoded.get(slot) if wav is not None: parts.append(wav) @@ -607,8 +572,6 @@ class MossTTSCodecDecoder(nn.Module): if slot is not None: session.release(slot) self._stream_req_slots.pop(request_id, None) - self._stream_pending_codes.pop(request_id, None) - self._stream_starved_reqs.discard(request_id) def on_requests_finished(self, finished_req_ids: set[str] | list[str]) -> None: """Release codec streaming slots when requests finish outside payload flow. @@ -623,17 +586,12 @@ class MossTTSCodecDecoder(nn.Module): for req_id in finished_req_ids: request_id = str(req_id) slot = self._stream_req_slots.get(request_id) - has_state = ( - slot is not None or request_id in self._stream_pending_codes or request_id in self._stream_starved_reqs - ) - if not has_state: + if slot is None: continue if session is not None: self._finish_stream_request(request_id, session, slot) else: self._stream_req_slots.pop(request_id, None) - self._stream_pending_codes.pop(request_id, None) - self._stream_starved_reqs.discard(request_id) def _connector_int(self, name: str, default: int = 0) -> int: model_cfg = getattr(self.vllm_config, "model_config", None) @@ -646,29 +604,14 @@ class MossTTSCodecDecoder(nn.Module): return int(extra_cfg[name]) return default - def _connector_bool(self, name: str, default: bool = False) -> bool: - model_cfg = getattr(self.vllm_config, "model_config", None) - connector_cfg = getattr(model_cfg, "stage_connector_config", None) - if isinstance(connector_cfg, dict): - extra_cfg: dict | None = connector_cfg.get("extra", connector_cfg) - else: - extra_cfg = getattr(connector_cfg, "extra", None) - if isinstance(extra_cfg, dict) and name in extra_cfg: - return bool(extra_cfg[name]) - return default - - def _streaming_cudagraph_capture_sizes_from_compilation_config( - self, - default: list[int], - ) -> list[int]: + def _streaming_graph_batch_sizes_from_compilation_config(self) -> list[int]: if getattr(self.vllm_config.model_config, "enforce_eager", True): return [] compilation_config = getattr(self.vllm_config, "compilation_config", None) capture_sizes = getattr(compilation_config, "cudagraph_capture_sizes", None) if capture_sizes: - max_step_frames = max(1, int(self._stream_max_step_frames)) - return sorted({int(size) for size in capture_sizes if 0 < int(size) <= max_step_frames}) - return list(default) + return sorted({int(size) for size in capture_sizes if 0 < int(size) <= self._stream_state_capacity}) + return [] # ------------------------------------------------------------------ # Weight loading @@ -770,6 +713,17 @@ class MossTTSCodecDecoder(nn.Module): device = self.vllm_config.device_config.device codec.to(device=device, dtype=torch.float32) codec.eval() + build_decode_lut = getattr(codec.quantizer, "build_decode_lut", None) + if callable(build_decode_lut): + lut_dtype = torch.bfloat16 if device.type == "cuda" else torch.float32 + build_decode_lut(self._n_vq, dtype=lut_dtype) + lut = codec.quantizer._decode_lut + logger.info( + "MOSS Audio Tokenizer LFQ decoded LUT: shape=%s dtype=%s size=%.1f MiB", + tuple(lut.shape), + lut.dtype, + lut.numel() * lut.element_size() / (1024**2), + ) self._codec = codec inferred_channels = 2 if "v2" in codec_path.lower() else 1 self._n_channels = int( @@ -790,7 +744,7 @@ class MossTTSCodecDecoder(nn.Module): ) self._configure_decoder_cudagraph(device) - if self._codec_streaming and self._streaming_cudagraph_capture_sizes: + if self._async_chunk and self._streaming_graph_batch_sizes and self._streaming_graph_frame_sizes: self._ensure_stream_session() # vLLM's track_weights_loading() compares the returned set against @@ -818,18 +772,19 @@ class MossTTSCodecDecoder(nn.Module): """Select the codec CUDA Graph path. ``enforce_eager`` is the single graph on/off switch. If graphing is - enabled, ``codec_streaming`` decides whether decode uses the persistent + enabled, ``async_chunk`` decides whether decode uses the persistent streaming-state wrapper or the offline full-chunk wrapper. """ if getattr(self.vllm_config.model_config, "enforce_eager", True): - self._streaming_cudagraph_capture_sizes = [] + self._streaming_graph_batch_sizes = [] return if self._codec is None: return - if self._codec_streaming: + if self._async_chunk: logger.info( - "MOSS-TTS codec CUDA Graph selected streaming wrapper: capture_sizes=%s", - self._streaming_cudagraph_capture_sizes, + "MOSS-TTS codec CUDA Graph selected streaming wrapper: B=%s exact_T=%s", + self._streaming_graph_batch_sizes, + self._streaming_graph_frame_sizes, ) return diff --git a/vllm_omni/model_executor/models/moss_tts/pipeline.py b/vllm_omni/model_executor/models/moss_tts/pipeline.py index 254bc18fe..f499654a4 100644 --- a/vllm_omni/model_executor/models/moss_tts/pipeline.py +++ b/vllm_omni/model_executor/models/moss_tts/pipeline.py @@ -46,6 +46,7 @@ MOSS_TTS_PIPELINE = PipelineConfig( final_output_type="audio", engine_output_type="audio", model_arch="MossTTSCodecDecoder", + retains_state_across_chunks=True, sync_process_input_func=f"{_PROC}.talker2codec", sampling_constraints={"detokenize": True}, ), @@ -76,6 +77,7 @@ MOSS_TTS_REALTIME_PIPELINE = PipelineConfig( final_output_type="audio", engine_output_type="audio", model_arch="MossTTSCodecDecoder", + retains_state_across_chunks=True, sync_process_input_func=f"{_PROC}.talker2codec", sampling_constraints={"detokenize": True}, ), @@ -109,6 +111,7 @@ MOSS_TTS_LOCAL_PIPELINE = PipelineConfig( final_output_type="audio", engine_output_type="audio", model_arch="MossTTSCodecDecoder", + retains_state_across_chunks=True, sync_process_input_func=f"{_PROC}.talker2codec", sampling_constraints={"detokenize": True}, ), diff --git a/vllm_omni/model_executor/stage_input_processors/moss_tts.py b/vllm_omni/model_executor/stage_input_processors/moss_tts.py index 2a7da4a4d..3ba32271e 100644 --- a/vllm_omni/model_executor/stage_input_processors/moss_tts.py +++ b/vllm_omni/model_executor/stage_input_processors/moss_tts.py @@ -294,7 +294,6 @@ def talker2codec_raw_async_chunk( meta=MetaStruct( req_id=[req_id], left_context_size=0, - codec_streaming=True, codec_chunk_frames=0, codec_left_context_frames=0, stream_finished=torch.tensor(True, dtype=torch.bool), @@ -326,7 +325,6 @@ def talker2codec_raw_async_chunk( meta=MetaStruct( req_id=[req_id], left_context_size=0, - codec_streaming=True, codec_chunk_frames=int(chunk_codes.shape[0]), codec_left_context_frames=0, code_flat_numel=int(codec_flat.numel()),