Unskip test for invalid logprobs handling

Remove skip marker from test for invalid logprobs.

Signed-off-by: Gao Han <hgaoaf@connect.ust.hk>

tests/core/sched/test_omni_ar_scheduler_logprobs.py

@@ -76,6 +76,7 @@ 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
@@ -104,7 +105,6 @@ 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,7 +138,8 @@ 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)
-        return None
+        # vLLM 0.26 contract: (kv_xfer_params, ec_xfer_params)
+        return None, None
 
     scheduler._update_request_with_output = update_request
     scheduler._process_kv_transfer_trigger = lambda _request, _tokens: False

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"

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:

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 (

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:

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

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,

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)

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"]

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
 

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},
         ),

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()),

vllm_omni/worker/base.py

@@ -120,7 +120,6 @@ class OmniGPUWorkerBase(GPUWorker):
             self.model_runner.profile_run()
 
         self.non_torch_memory = profile_result.non_torch_increase
-        self.total_consumed = profile_result.total_consumed
         self.peak_activation_memory = profile_result.torch_peak_increase
 
         process_memory = (