Revert "[BugFix] Fix 'GPUARWorker' attribute error (#5478)"

This reverts commit f4640d9a2e.

tests/core/sched/test_omni_ar_scheduler_logprobs.py

@@ -76,7 +76,6 @@ class _Request:
         self.status = RequestStatus.RUNNING
         self.sampling_params = SimpleNamespace(num_logprobs=1)
         self.num_computed_tokens = 0
-        self.num_in_flight_tokens = 0
         self.num_output_placeholders = 0
         self.has_encoder_inputs = False
         self.pooling_params = None
@@ -105,6 +104,7 @@ class _RequestQueue(list):
                 self.remove(request)
 
 
+@pytest.mark.skip(reason="#issue 5484")
 def test_invalid_logprobs_finish_only_the_affected_scheduler_request() -> None:
     """A bad runner response must not bring down unrelated batch requests."""
     bad = _Request("bad")
@@ -138,8 +138,7 @@ def test_invalid_logprobs_finish_only_the_affected_scheduler_request() -> None:
         scheduler.requests.pop(request.request_id, None)
         scheduler.finished_req_ids.add(request.request_id)
         scheduler.finished_req_ids_dict[request.client_index].add(request.request_id)
-        # vLLM 0.26 contract: (kv_xfer_params, ec_xfer_params)
-        return None, None
+        return None
 
     scheduler._update_request_with_output = update_request
     scheduler._process_kv_transfer_trigger = lambda _request, _tokens: False

vllm_omni/config/model.py

@@ -119,7 +119,6 @@ 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"

vllm_omni/config/stage_config.py

@@ -223,9 +223,6 @@ 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
@@ -825,7 +822,6 @@ 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:

vllm_omni/core/sched/omni_generation_scheduler.py

@@ -47,7 +47,6 @@ 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):
@@ -146,17 +145,12 @@ 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 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
-
+        while (
+            self.waiting
+            and token_budget > 0
+            and len(self.running) < self.max_num_running_reqs
+            and self._pause_state == PauseState.UNPAUSED
+        ):
             request = self.waiting.peek_request()
             # OMNI: Skip requests that are not in self.requests
             if request.request_id not in self.requests or (

vllm_omni/deploy/moss_tts_local.yaml

@@ -9,25 +9,24 @@ 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: 64
+    max_num_seqs: 8
     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:
@@ -39,22 +38,17 @@ stages:
       repetition_penalty: 1.1
 
   - stage_id: 1
-    max_num_seqs: 64
-    gpu_memory_utilization: 0.30
+    max_num_seqs: 8
+    gpu_memory_utilization: 0.15
     enforce_eager: false
     enable_prefix_caching: false
     async_scheduling: true
     compilation_config:
-      # 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_capture_sizes: [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15]
       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: "1"
+    devices: "0"
     input_connectors:
       from_stage_0: shm
     default_sampling_params:

vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py

@@ -550,11 +550,6 @@ 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

vllm_omni/engine/arg_utils.py

@@ -159,7 +159,6 @@ 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.
@@ -342,7 +341,6 @@ 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,

vllm_omni/model_executor/models/moss_tts/audio_tokenizer_v2.py

@@ -102,22 +102,12 @@ 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:
-        if self._owns_exec_mask:
-            self.exec_mask.masked_fill_(reset_mask, True)
+        self.exec_mask[:] = torch.where(reset_mask, torch.ones_like(self.exec_mask), self.exec_mask)
 
     def __enter__(self):
         # ExitStack expects a context manager; returning self is conventional and useful for debugging.
@@ -127,28 +117,6 @@ 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."""
 
@@ -479,57 +447,11 @@ 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 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:
+    def complete(self, k: torch.Tensor, v: torch.Tensor, exec_mask: torch.Tensor) -> 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
@@ -557,8 +479,6 @@ 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)
@@ -587,14 +507,6 @@ 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,
@@ -688,8 +600,7 @@ 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)
-        weight_dtype = cast(torch.dtype, in_proj.weight.dtype)
-        dtype = torch.bfloat16 if device.type == "cuda" else weight_dtype
+        dtype = cast(torch.dtype, in_proj.weight.dtype)
 
         dim_per_head = self.embed_dim // self.num_heads
         if self.context is None:
@@ -714,37 +625,21 @@ class MossAudioTokenizerMultiheadAttention(StreamingModule):
             offset_cpu=0,
         )
 
-    def _complete_kv(
-        self,
-        k: torch.Tensor,
-        v: torch.Tensor,
-        execution_context: StreamingExecutionContext | None = None,
-    ) -> KVCacheResult:
+    def _complete_kv(self, k, v) -> 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, execution_context)
+        return state.kv_cache.complete(k, v, state.exec_mask)
 
-    def forward(
-        self,
-        query: torch.Tensor,
-        key: torch.Tensor,
-        value: torch.Tensor,
-        execution_context: StreamingExecutionContext | None = None,
-    ):
+    def forward(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor):
         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
@@ -757,7 +652,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, execution_context)
+        k, v, pos_k = self._complete_kv(k, v)
         pos_k = pos_k[:, None]
 
         if self.causal:
@@ -775,12 +670,8 @@ class MossAudioTokenizerMultiheadAttention(StreamingModule):
         x = apply_weights_per_step(self.out_projs, self.weights_per_step_schedule, x, offset_cpu)
 
         if state is not None:
-            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)
+            state.offset[:] = torch.where(state.exec_mask, state.offset + T, state.offset)
+            state.offset_cpu += T
         return x
 
 
@@ -797,11 +688,6 @@ 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."""
@@ -877,14 +763,8 @@ class MossAudioTokenizerTransformerLayer(StreamingModule):
         device = next(iter(self.parameters())).device
         return LayerState(batch_size, device, offset_cpu=0)
 
-    def _ff_block(
-        self,
-        x: torch.Tensor,
-        execution_context: StreamingExecutionContext | None = None,
-    ) -> torch.Tensor:
+    def _ff_block(self, x: torch.Tensor) -> 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
@@ -902,25 +782,17 @@ class MossAudioTokenizerTransformerLayer(StreamingModule):
                 update = self.gating(x)
         return x_orig.to(update) + self.layer_scale_2(update)
 
-    def _sa_block(
-        self,
-        x: torch.Tensor,
-        execution_context: StreamingExecutionContext | None = None,
-    ):
+    def _sa_block(self, x: torch.Tensor):
         x_orig = x
         x = self.norm1(x)
-        update = self.self_attn(x, x, x, execution_context=execution_context)
+        update = self.self_attn(x, x, x)
         return x_orig.to(update) + self.layer_scale_1(update)
 
-    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)
+    def forward(self, x: torch.Tensor):
+        x = self._sa_block(x)
+        x = self._ff_block(x)
         state = self._streaming_state
-        if state is not None and execution_context is None:
+        if state is not None:
             assert isinstance(state, LayerState)
             state.offset_cpu += x.shape[1]
         return x
@@ -939,11 +811,6 @@ 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."""
@@ -1002,17 +869,15 @@ class MossAudioTokenizerTransformer(StreamingModule):
     def forward(self, x: torch.Tensor, *args, **kwargs):
         B, T, C = x.shape
         state = self._streaming_state
-        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)
+        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)
+            )
+        )
 
         if self.positional_embedding in {"sin", "sin_rope"}:
             positions = torch.arange(T, device=x.device).view(1, -1, 1)
@@ -1025,11 +890,7 @@ class MossAudioTokenizerTransformer(StreamingModule):
 
         if state is not None:
             assert isinstance(state, TransformerState)
-            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)
+            state.offsets[:] = torch.where(state.exec_mask, state.offsets + T, state.offsets)
         return x
 
 
@@ -1374,76 +1235,6 @@ 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(
@@ -1485,12 +1276,6 @@ 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)
@@ -1604,114 +1389,34 @@ 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, *, 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_streaming(self, batch_size: int):
+        """Start streaming mode for all modules."""
 
         def _start(module):
             if isinstance(module, StreamingModule):
-                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)
+                module._streaming_state = module._init_streaming_state(batch_size)
 
-        root.apply(_start)
-        self._streaming_modules = streaming_modules
-        self._streaming_exec_mask = shared_exec_mask
+        self.apply(_start)
 
-    def _stop_streaming(self) -> None:
+    def _stop_streaming(self):
         """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 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
+        def _stop(module):
+            if isinstance(module, StreamingModule):
+                module._streaming_state = None
 
-    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)
+        self.apply(_stop)
 
     def _set_streaming_exec_mask(self, exec_mask: torch.Tensor) -> None:
-        """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)
+        """Set the active slot mask on every live streaming module."""
 
-    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)
+        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)
 
     @contextmanager
     def streaming(self, batch_size: int = 1):
@@ -1722,15 +1427,6 @@ 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.")
@@ -1929,7 +1625,6 @@ 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
@@ -1938,21 +1633,13 @@ class MossAudioTokenizerModel(MossAudioTokenizerPreTrainedModel):
         if codes_lengths is None:
             codes_lengths = torch.full((B,), T, device=device, dtype=torch.long)
 
-        # 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)
+        # Decode from codes
+        quantizer = cast(MossAudioTokenizerResidualVQ | MossAudioTokenizerResidualLFQ, self.quantizer)
+        zq = quantizer.decode_codes(codes)
 
-            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 = zq, codes_lengths
+        for decoder_module in self.decoder:
+            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)

vllm_omni/model_executor/models/moss_tts/cuda_graph_streaming_decoder_wrapper.py

@@ -1,11 +1,12 @@
 # Copyright 2026 OpenMOSS and the vLLM-Omni team. All rights reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License").
-"""CUDA graphs for compact MOSS-TTS streaming codec batches."""
+"""CUDA graph wrapper for MOSS-TTS Local streaming codec decode."""
 
 from __future__ import annotations
 
 import time
+from collections.abc import Callable
 from dataclasses import dataclass
 
 import torch
@@ -22,35 +23,31 @@ 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 streaming decode graphs keyed by ``(B_bucket, exact_T)``.
+    """Replay fixed-size streaming ``_decode_frame`` steps with CUDA graphs.
 
-    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.
+    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.
     """
 
     def __init__(
         self,
         codec: nn.Module,
         *,
-        state_capacity: int,
-        batch_sizes: list[int],
-        frame_sizes: list[int],
+        batch_size: int,
         num_quantizers: int,
+        reset_streaming_state: Callable[[], None] | None = None,
     ) -> None:
         self.codec = codec
-        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.batch_size = int(batch_size)
         self.num_quantizers = int(num_quantizers)
-        self.graphs: dict[tuple[int, int], _CapturedStreamingDecodeGraph] = {}
+        self.reset_streaming_state = reset_streaming_state
+        self.graphs: dict[int, _CapturedStreamingDecodeGraph] = {}
         self._pool = None
         self._warmed_up = False
 
@@ -59,135 +56,109 @@ class CUDAGraphStreamingDecoderWrapper:
         return bool(self.graphs)
 
     @property
-    def capture_sizes(self) -> list[tuple[int, int]]:
+    def capture_sizes(self) -> list[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) -> None:
+    def warmup(self, device: torch.device, capture_sizes: list[int]) -> None:
         if self._warmed_up:
             return
         self._warmed_up = True
-        if device.type != "cuda" or not self.batch_sizes or not self.frame_sizes:
+        if device.type != "cuda":
+            return
+        sizes = sorted({int(size) for size in capture_sizes if int(size) > 0}, reverse=True)
+        if not 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: n_vq=%d (B,T)=%s",
+            "MOSS-TTS streaming decoder CUDA graph warmup: batch_size=%d n_vq=%d capture_sizes=%s",
+            self.batch_size,
             self.num_quantizers,
-            list(reversed(capture_keys)),
+            list(reversed(sizes)),
         )
         start_s = time.perf_counter()
         with torch.cuda.device(device):
-            for batch_size, frame_size in capture_keys:
-                key = (batch_size, frame_size)
+            for size in sizes:
                 try:
-                    self._capture(batch_size, frame_size, device)
-                    logger.info("  Captured MOSS-TTS streaming decoder CUDA graph for (B,T)=%s", key)
+                    self._capture(size, device)
+                    logger.info("  Captured MOSS-TTS streaming decoder CUDA graph for size=%d", size)
                 except Exception:
-                    self.graphs.pop(key, None)
+                    self.graphs.pop(size, None)
                     logger.warning(
-                        "  Failed to capture MOSS-TTS streaming decoder CUDA graph for (B,T)=%s; using eager",
-                        key,
+                        "  Failed to capture MOSS-TTS streaming decoder CUDA graph for size=%d; "
+                        "this size will use eager decode",
+                        size,
                         exc_info=True,
                     )
         logger.info(
             "MOSS-TTS streaming decoder CUDA graph warmup complete: %d/%d captured in %.1f ms",
             len(self.graphs),
-            len(capture_keys),
+            len(sizes),
             (time.perf_counter() - start_s) * 1000.0,
         )
 
     @torch.no_grad()
-    def _capture(self, batch_size: int, frame_size: int, device: torch.device) -> None:
+    def _capture(self, size: int, device: torch.device) -> None:
         codes = torch.zeros(
             self.num_quantizers,
-            batch_size,
-            frame_size,
+            self.batch_size,
+            size,
             dtype=torch.long,
             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)
+        lengths = torch.full((self.batch_size,), size, dtype=torch.long, device=device)
+        exec_mask = torch.ones(self.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_streaming_batch(codes, lengths, scratch_slots, valid_rows)
+            _ = self.codec._decode_frame(codes, lengths)
         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_streaming_batch(codes, lengths, scratch_slots, valid_rows)
+            output = self.codec._decode_frame(codes, lengths)
 
         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[(batch_size, frame_size)] = _CapturedStreamingDecodeGraph(
+
+        self.graphs[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: 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():
+        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():
             return None
-        n_vq, actual_batch_size, frame_size = codes.shape
-        if int(n_vq) != self.num_quantizers or state_slot_ids.shape != (actual_batch_size,):
+        n_vq, batch_size, step_t = codes_step.shape
+        if int(n_vq) != self.num_quantizers or int(batch_size) != self.batch_size:
             return None
-        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))
+        # 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))
         if entry is None:
             return None
 
-        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)
+        self.codec._set_streaming_exec_mask(exec_mask)
+        entry.static_codes.copy_(codes_step)
         entry.graph.replay()
-        return entry.static_audio, entry.static_audio_lengths, int(actual_batch_size)
+        return entry.static_audio, entry.static_audio_lengths
 
 
 __all__ = ["CUDAGraphStreamingDecoderWrapper"]

vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_codec.py

@@ -5,6 +5,7 @@
 
 from __future__ import annotations
 
+import contextlib
 from collections.abc import Iterable
 from typing import Any
 
@@ -37,136 +38,104 @@ logger = init_logger(__name__)
 
 
 class _MossCodecStreamSession:
-    """Persistent state pool with compact per-step execution batches."""
+    """Persistent streaming decode session for vendored MOSS-Audio-Tokenizer-v2."""
 
     def __init__(
         self,
         codec: nn.Module,
         *,
-        state_capacity: int,
+        stream_slots: int,
         n_vq: int,
-        graph_batch_sizes: list[int] | None = None,
-        graph_frame_sizes: list[int] | None = None,
+        cudagraph_capture_sizes: list[int] | None = None,
     ) -> None:
         self._codec = codec
-        self._state_capacity = int(state_capacity)
+        self._stream_slots = int(stream_slots)
+        self._batch_size = self._stream_slots
         self._n_vq = int(n_vq)
         self._device = next(codec.parameters()).device
-        self._free_stream_slots = list(reversed(range(self._state_capacity)))
-        self._leased_slots: set[int] = set()
+        self._free_stream_slots = list(range(self._stream_slots))
+        self._exit_stack = contextlib.ExitStack()
         self._closed = False
-        self._cudagraph_wrapper: CUDAGraphStreamingDecoderWrapper | None = None
-        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._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":
             self._cudagraph_wrapper = CUDAGraphStreamingDecoderWrapper(
                 codec,
-                state_capacity=self._state_capacity,
-                batch_sizes=batch_sizes,
-                frame_sizes=frame_sizes,
+                batch_size=self._batch_size,
                 num_quantizers=self._n_vq,
+                reset_streaming_state=lambda: self.reset_slots(list(range(self._batch_size))),
             )
-            self._cudagraph_wrapper.warmup(self._device)
-            self.reset_slots(list(range(self._state_capacity + scratch_capacity)))
+            self._cudagraph_wrapper.warmup(self._device, capture_sizes)
+            self.reset_slots(list(range(self._batch_size)))
             if not self._cudagraph_wrapper.is_ready:
                 self._cudagraph_wrapper = None
 
     def acquire(self) -> int | None:
         if not self._free_stream_slots:
             return None
-        slot = self._free_stream_slots.pop()
-        self._leased_slots.add(slot)
-        return slot
+        return self._free_stream_slots.pop()
 
     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_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.")
+        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)
 
     def close(self) -> None:
         if self._closed:
             return
-        close_state_pool = getattr(self._codec, "close_decoder_state_pool", None)
         with torch.no_grad():
-            if callable(close_state_pool):
-                close_state_pool()
+            self._exit_stack.close()
         self._closed = True
 
     @torch.no_grad()
-    def step(
-        self,
-        slot_codes: dict[int, torch.Tensor],
-        *,
-        terminal_slots: set[int] | None = None,
-    ) -> dict[int, torch.Tensor]:
+    def step(self, slot_codes: dict[int, torch.Tensor]) -> 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
-        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_step = torch.zeros(
+            self._n_vq,
+            self._batch_size,
+            step_t,
+            dtype=torch.long,
+            device=self._device,
         )
-        state_slot_ids = torch.tensor(slots, device=self._device, dtype=torch.long)
+        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
 
-        graph_output: tuple[torch.Tensor, torch.Tensor, int] | None = None
+        graph_output: tuple[torch.Tensor, torch.Tensor] | None = None
         if self._cudagraph_wrapper is not None:
-            # 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),
-            )
+            graph_output = self._cudagraph_wrapper.decode(codes_step, exec_mask)
 
         used_cudagraph = graph_output is not None
         if used_cudagraph:
-            audio_tensor, lengths_tensor, actual_batch_size = graph_output
-            audio_tensor = audio_tensor[:actual_batch_size]
-            lengths_tensor = lengths_tensor[:actual_batch_size]
+            audio_tensor, lengths_tensor = graph_output
         else:
-            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,
-            )
+            self._codec._set_streaming_exec_mask(exec_mask)
+            result = self._codec._decode_frame(codes_step, codes_lengths)
             if result.audio is None:
                 return {}
             audio_tensor = result.audio
@@ -175,10 +144,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 row, slot in enumerate(slots):
-            wav = audio[row]
+        for slot in slot_codes:
+            wav = audio[slot]
             if lengths is not None:
-                wav = wav[..., : int(lengths[row].item())]
+                wav = wav[..., : int(lengths[slot].item())]
             out[slot] = wav.contiguous()
         return out
 
@@ -232,16 +201,17 @@ 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
-        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_slots: int = self._connector_int("codec_stream_slots", 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._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}
+        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
         )
 
     # ------------------------------------------------------------------
@@ -305,7 +275,7 @@ class MossTTSCodecDecoder(nn.Module):
                 },
             )
 
-        if self._async_chunk and runtime_additional_information is None:
+        if self._codec_streaming and runtime_additional_information is None:
             return OmniOutput(
                 text_hidden_states=None,
                 multimodal_outputs={
@@ -320,6 +290,8 @@ 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},
@@ -369,7 +341,12 @@ 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 = self._async_chunk
+            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
             if seg.numel() % self._n_vq != 0:
                 logger.warning(
                     "MossTTS codec input length %d not divisible by n_vq %d; skipping.",
@@ -437,6 +414,40 @@ 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:
@@ -475,12 +486,15 @@ 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,
-            state_capacity=self._stream_state_capacity,
+            stream_slots=max(1, slots),
             n_vq=self._n_vq,
-            graph_batch_sizes=self._streaming_graph_batch_sizes,
-            graph_frame_sizes=self._streaming_graph_frame_sizes,
+            cudagraph_capture_sizes=self._streaming_cudagraph_capture_sizes,
         )
         return self._stream_session
 
@@ -497,23 +511,40 @@ 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:
-                    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._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
                 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,
-                    terminal=finished,
-                )
+                wav = self._decode_stream_slot_sequence(session, slot, codes_nq_t)
                 if wav is not None:
                     outputs[output_index] = wav
                 if finished:
@@ -526,8 +557,7 @@ class MossTTSCodecDecoder(nn.Module):
 
         for group in grouped.values():
             plan = {slot: codes_nq_t for _, _, slot, codes_nq_t, _ in group}
-            terminal_slots = {slot for _, _, slot, _, finished in group if finished}
-            decoded = session.step(plan, terminal_slots=terminal_slots)
+            decoded = session.step(plan)
             for output_index, request_id, slot, _, finished in group:
                 wav = decoded.get(slot)
                 if wav is not None:
@@ -537,13 +567,22 @@ 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
@@ -551,11 +590,7 @@ 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]
-            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,
-            )
+            decoded = session.step({slot: chunk})
             wav = decoded.get(slot)
             if wav is not None:
                 parts.append(wav)
@@ -572,6 +607,8 @@ 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.
@@ -586,12 +623,17 @@ class MossTTSCodecDecoder(nn.Module):
         for req_id in finished_req_ids:
             request_id = str(req_id)
             slot = self._stream_req_slots.get(request_id)
-            if slot is None:
+            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:
                 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)
@@ -604,14 +646,29 @@ class MossTTSCodecDecoder(nn.Module):
             return int(extra_cfg[name])
         return default
 
-    def _streaming_graph_batch_sizes_from_compilation_config(self) -> list[int]:
+    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]:
         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:
-            return sorted({int(size) for size in capture_sizes if 0 < int(size) <= self._stream_state_capacity})
-        return []
+            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)
 
     # ------------------------------------------------------------------
     # Weight loading
@@ -713,17 +770,6 @@ 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(
@@ -744,7 +790,7 @@ class MossTTSCodecDecoder(nn.Module):
         )
 
         self._configure_decoder_cudagraph(device)
-        if self._async_chunk and self._streaming_graph_batch_sizes and self._streaming_graph_frame_sizes:
+        if self._codec_streaming and self._streaming_cudagraph_capture_sizes:
             self._ensure_stream_session()
 
         # vLLM's track_weights_loading() compares the returned set against
@@ -772,19 +818,18 @@ class MossTTSCodecDecoder(nn.Module):
         """Select the codec CUDA Graph path.
 
         ``enforce_eager`` is the single graph on/off switch. If graphing is
-        enabled, ``async_chunk`` decides whether decode uses the persistent
+        enabled, ``codec_streaming`` 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_graph_batch_sizes = []
+            self._streaming_cudagraph_capture_sizes = []
             return
         if self._codec is None:
             return
-        if self._async_chunk:
+        if self._codec_streaming:
             logger.info(
-                "MOSS-TTS codec CUDA Graph selected streaming wrapper: B=%s exact_T=%s",
-                self._streaming_graph_batch_sizes,
-                self._streaming_graph_frame_sizes,
+                "MOSS-TTS codec CUDA Graph selected streaming wrapper: capture_sizes=%s",
+                self._streaming_cudagraph_capture_sizes,
             )
             return
 

vllm_omni/model_executor/models/moss_tts/pipeline.py

@@ -46,7 +46,6 @@ 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},
         ),
@@ -77,7 +76,6 @@ 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},
         ),
@@ -111,7 +109,6 @@ 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},
         ),

vllm_omni/model_executor/stage_input_processors/moss_tts.py

@@ -294,6 +294,7 @@ 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),
@@ -325,6 +326,7 @@ 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()),