vllm-omni/docs/design/feature/async_chunk.md

# Async Chunk

## Table of Contents

1. [Overview](#overview)
2. [Performance](#performance)
3. [Architecture](#architecture)
4. [Configuration](#configuration)
5. [Related Files](#related-files)

## Overview

The `async_chunk` feature enables asynchronous, chunked processing of data across multiple stages in a multi-stage pipeline (e.g., Qwen3-Omni with Thinker → Talker → Code2Wav stages). Instead of waiting for a complete stage output before forwarding to the next stage, this feature allows stages to process and forward data in chunks as it becomes available, significantly reducing latency and improving throughput.

**Chunk Size Definition**

- **Prefill Phase**: `chunk_size = num_scheduled_tokens` for chunked prefill processing
- **Decode Phase**: `chunk_size = num_scheduled_tokens = 1 ` for per-token streaming

For qwen3-omni:
- **Thinker → Talker**: Per decode step (typically chunk_size=1)
- **Talker → Code2Wav**: Accumulated to `codec_chunk_frames` (default=25) before sending. During the initial phase, a dynamic initial chunk size (IC) is automatically selected based on server load to reduce TTFP. Use the per-request `initial_codec_chunk_frames` API field to override.
- **Code2Wav**: Streaming decode with code2wav chunk_size

With `async_chunk`:
- Stages can start processing as soon as chunks are available
- Overlapping execution across stages
- Reduced latency and improved throughput
- Better resource utilization
- Async scheduling: Chunk IO (get/put) overlaps with compute via background threads so the scheduler is not blocked waiting for chunks

## Performance
1. **Reduced Latency**: Next stage can start processing immediately
2. **Streaming Support**: Enables streaming for audio generation
3. **IO-Compute Overlap**: Chunk retrieval happens asynchronously while other requests compute
4. **Non-blocking Scheduler**: Requests waiting for chunks don't block the entire scheduler
5. **Code2Wav Batch Inference**: Supports batched processing in code2wav stage

| Input     | Output           | Async_chunk enabled | Code2Wav batch size | Max_Concurrency | Prompts | Mean E2E   | Mean TTFT  | Mean TPOT | Mean TTFP   | Mean RTF | Mean ITL |
|-----------|------------------|---------------------|---------------------|----------------|---------|------------|------------|-----------|-------------|----------|----------|
| text 100  | text 100+audio   | False               | 1                   | 1              | 50      | 6581.80    | 43.22      | 8.31      | 6459.34     | 0.24     | 8.22     |
| text 100  | text 100+audio   | False               | 1                   | 4              | 50      | 7398.63    | 67.57      | 9.14      | 7285.35     | 0.27     | 9.05     |
| text 100  | text 100+audio   | False               | 1                   | 10             | 50      | 13522.99   | 131.82     | 12.72     | 13410.44    | 0.49     | 12.60    |
| text 100  | text 100+audio   | False               | 64                  | 1              | 50      | 6505.13    | 43.14      | 8.52      | 6395.40     | 0.24     | 8.44     |
| text 100  | text 100+audio   | False               | 64                  | 4              | 50      | 7668.15    | 51.15      | 9.36      | 7562.37     | 0.28     | 9.27     |
| text 100  | text 100+audio   | False               | 64                  | 10             | 50      | 9516.18    | 138.06     | 14.75     | 9409.26     | 0.34     | 14.60    |
| text 100  | text 100+audio   | True                | 1                   | 1              | 50      | 6179.79    | 44.58      | 8.69      | 522.99      | 0.22     | 8.60     |
| text 100  | text 100+audio   | True                | 1                   | 4              | 50      | 7692.69    | 103.96     | 10.22     | 785.85      | 0.29     | 10.12    |
| text 100  | text 100+audio   | True                | 1                   | 10             | 50      | 11152.71   | 685.60     | 17.64     | 1628.88     | 0.41     | 17.62    |


Performance data collected on H800 GPUs through comprehensive benchmarking with cudagraph enabled. text input uses random dataset.

Enabling **async_chunk** (False→True) sharply reduces time-to-first-audio (TTFP)—e.g. ~92% at concurrency 1 (6.5s→0.52s)—and improves E2E latency (e.g. ~6% at conc 1, ~17% at conc 10). RTF (Real Time Factor) also improves with async_chunk on (e.g. ~8% at conc 1: 0.24→0.22, ~16% at conc 10: 0.49→0.41). Enabling **Code2Wav batch size 64** (vs 1) improves E2E and TTFP at higher concurrency when async_chunk is off (e.g. ~30% at conc 10: 13.5s→9.5s E2E, 13.4s→9.4s TTFP).

<p align="center">
  <picture>
    <source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/performance/qwen3-omni_ttfp_performance.png">
    <img alt="TTFP Performance Data Comparison" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/performance/qwen3-omni_ttfp_performance.png" width=100%>
  </picture>
</p>

<p align="center">
  <picture>
    <source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/performance/qwen3-omni_rtf_performance.png">
    <img alt="RTF Performance Data Comparison" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/performance/qwen3-omni_rtf_performance.png" width=100%>
  </picture>
</p>

<p align="center">
  <picture>
    <source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/performance/qwen3-omni_e2e_performance.png">
    <img alt="E2E Performance Data Comparison" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/performance/qwen3-omni_e2e_performance.png" width=100%>
  </picture>
</p>

## Architecture

### Async Chunk Pipeline Overview

The following diagram illustrates the **Async Chunk Architecture** for multi-stage models (e.g., Qwen3-Omni with Thinker → Talker → Code2Wav), showing how data flows through the 4-stage pipeline with parallel processing and dual-stream output:

<p align="center">
  <picture>
    <source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/qwen3-omni-async-chunk.png">
    <img alt="Async Chunk Pipeline Architecture" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/qwen3-omni-async-chunk.png" width=100%>
  </picture>
</p>

**Diagram Legend:**

| Step | Stage Type | Description |
|------|-----------|------------|
| `prefill` | Initialization | Context processing, KV cache initialization |
| `decode` | Autoregressive | Token-by-token generation in AR stages |
| `codes` | Audio Encoding | RVQ codec codes from Talker stage |
| `output` | Final Output | Text chunks or audio waveforms |

### Data Flow

#### Stage 0: Thinker (Multimodal Understanding + Text Generation)
- **Prefill**: Processes multimodal input (text/image/audio/video), initializes KV cache
- **Decode Loop**: Generates text tokens autoregressively
- **Chunk Triggers**: Each decode step (typically `chunk_size=1`) can trigger downstream processing
- **Dual Output**:
  - **Text Stream**: `text_0`, `text_1`, `text_2`... `text_n` streamed to output
  - **Hidden States**: Passed to Talker stage for audio synthesis

#### Stage 1: Talker (Text → RVQ Audio Codes)
- **Prefill**: Receives hidden states from Thinker as semantic condition
- **Decode Loop**: Generates RVQ codec codes autoregressively
- **Accumulation**: Codes accumulate to `codec_chunk_frames` (default=25) before forwarding
- **Dynamic IC**: Initial chunk size auto-selected based on server load to optimize TTFP
- **Output**: `codes` blocks (chunk 0, 1, ... n) sent to Code2Wav

#### Stage 2: Code2Wav (Vocoder Decoder)
- **Non-Autoregressive**: Processes RVQ codes in parallel batches
- **Streaming Decode**: Converts codes to audio waveforms chunk-by-chunk
- **Batching**: Supports batched inference for multiple concurrent requests
- **Output**: Audio segments `audio_0`, `audio_1`, ... `audio_n`

#### Stage 3: Output (Dual Stream)
- **Text Streaming**: `text_0` → `text_1` → `text_2` → ... (user sees response in real-time)
- **Audio Streaming**: `audio_0` → `audio_1` → ... (user hears audio progressively)

### Execution Timeline

```
Timeline: Parallel vs Sequential

Sequential (async_chunk=false):
[Thinker: ████████████████████]  (2.0s)
                            [Talker: ████████████████████]  (3.0s)
                                                        [Code2Wav: ████]  (1.0s)
Total: 6.0s, TTFP: 6.0s

Async Chunk (async_chunk=true):
[Thinker: ████░░░░████░░░░████]  (2.0s, streaming)
     [Talker: ░░████░░░░████░░]  (3.0s, parallel)
         [Code2Wav: ░░░░████░░]  (1.0s, batched)
Total: ~3.5s, TTFP: ~0.5s

█ = Active computation  ░ = Waiting/idle
```

#### Sequential Flow (for comparison)
<p align="center">
  <picture>
    <source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/qwen3-omni-non-async-chunk.png">
    <img alt="Sequential Data Flow" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/qwen3-omni-non-async-chunk.png" width=100%>
  </picture>
</p>

In sequential mode, each stage must wait for the previous stage to complete entirely before starting.

### Async Chunk System Architecture
<p align="center">
  <picture>
    <source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/async-chunk-architecture.png">
    <img alt="Async Chunk Architecture" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/async-chunk-architecture.png" width=100%>
  </picture>
</p>


### Key Components

1. **OmniConnector**: Inter-stage data transport only
   - Shared memory or other IPC mechanisms
   - **Transport-only API**: `put(from_stage, to_stage, put_key, data)` and `get(from_stage, to_stage, get_key)` (optionally with timeout)
   - **No request-specific state**: Connector does not track put_requests, get_requests, request_payload, finished_requests, or other request-bound metadata; it only performs put/get operations
   - Chunk keys and request/chunk lifecycle are managed by **OmniChunkTransferAdapter**

2. **Transfer Adapter Layer**: Extensible abstraction for managing data transfer via connectors
   - **OmniTransferAdapterBase**: Base class with background **recv_loop** and **save_loop** threads;
   - **OmniChunkTransferAdapter**: Chunk-specific implementation that owns the full chunk lifecycle when async_chunk is enabled
     - **Chunk ID and key construction**: Builds keys like `{req_id}_{stage_id}_{chunk_id}` for put/get
     - **Async get**: `load_async(request)` enqueues the request; background **recv_loop** polls the connector (non-blocking); when data is available, updates the request and marks it in `_finished_load_reqs`; scheduler calls `get_finished_requests()` to learn which requests have chunks ready
     - **Async put**: `save_async(pooling_output, request)` invokes `custom_process_next_stage_input_func` in the main thread to build the payload, then enqueues a save task; background **save_loop** performs `connector.put()`; payload processing and chunk accumulation (e.g. code2wav chunk_size) remain in the main thread

3. **Stage Input Processors**: Custom functions that process stage outputs into chunks for different models
   - Receive **transfer_manager** (OmniChunkTransferAdapter)
   - Qwen3-omni reference: `thinker2talker_async_chunk`, `talker2code2wav_async_chunk`

4. **Schedulers**: Modified to handle chunk-based scheduling with async IO-compute overlap
   - `OmniARScheduler`: For autoregressive stages
   - `OmniGenerationScheduler`: For generation stages
   - Both schedulers use **OmniChunkTransferAdapter** and **before/after** hooks around `super().schedule()`:
     - **Before** `super().schedule()`: `process_pending_chunks(waiting, running)` moves requests waiting for chunks to `WAITING_FOR_CHUNK`, enqueues load tasks for background polling
     - **After** `super().schedule()`: `restore_queues(waiting, running)` restores requests with ready chunks back to waiting/running, `postprocess_scheduler_output(scheduler_output)` attaches cached additional_information, clears chunk-ready flags
   - **put_chunk** `save_async(pooler_output, request)`; **get_chunk** / **get_chunk_for_generation** `load_async(request)`

5. **Model Runners**: Handle chunk processing
   - `OmniGPUModelRunner`: Processes chunks in AR stages
   - `GPUGenerationModelRunner`: Processes chunks in generation stages
     - Uses `ubatch_slices` from `get_forward_context()` to track per-request sequence lengths in batched inference
     - Reuses `ubatch_slices_padded` for code2wav batching to properly split batch outputs
     - Handles list-type multimodal outputs: iterates through requests and assigns corresponding tensor to each
     - Improved request state management: removes unscheduled and finished requests from input batch

6. **Model Implementation**: Model-specific chunk handling
   - `Qwen3OmniMoeForConditionalGeneration`: Main model with async_chunk support
     - **Code2Wav stage batching**: Uses `ubatch_slices` to construct batched codec codes tensor `[batch_size, 16, max_seq_len]`
     - **Batch output handling**: `generate_audio()` returns `list[torch.Tensor]`, one audio tensor per request
     - **Multimodal outputs**: Returns list of audio tensors for batch processing instead of single concatenated tensor
   - `Qwen3OmniCode2WavDecoder`: Audio generation model
     - `chunked_decode()` and `chunked_decode_streaming()`: Return `list[torch.Tensor]` (one per request)
     - Uses `ubatch_slices` to split batched waveform output into per-request audio chunks
     - Each request gets correctly sized audio based on its code sequence length: `waveform[:, :, :code_seq_len * total_upsample]`

7. **Request status**: `RequestStatus.WAITING_FOR_CHUNK` is added via patch (e.g. in `vllm_omni/patch.py`) so requests waiting for a chunk are not scheduled by the base vLLM scheduler until the chunk is ready.

## Configuration

Enable async_chunk in stage configuration YAML:

```yaml
async_chunk: true
stage_args:
  - stage_id: 0
    engine_args:
      custom_process_next_stage_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker_async_chunk
  - stage_id: 1
    engine_args:
      custom_process_next_stage_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav_async_chunk
```

### Stage Configuration

- `async_chunk: bool`: Enable/disable async chunk mode
- `custom_process_next_stage_input_func: str`: Path to custom chunk processing function; receives `(transfer_manager, pooling_output, request)`. For qwen3-omni: `thinker2talker_async_chunk`, `talker2code2wav_async_chunk`
- `stage_connector_config: dict`: Connector configuration
- `worker_type: str`: Model type, e.g. `"ar"` or `"generation"` (used by OmniChunkTransferAdapter for mode-specific payload handling)
- `max_num_seqs: int`: Maximum number of sequences for concurrent processing in the stage


### Connector Configuration

```yaml
connectors:
  - from_stage: 0
    to_stage: 1
    spec:
      name: SharedMemoryConnector
      extra:
        stage_id: 0
```

### Code2Wav Batch Configuration

For optimal performance with async_chunk, the code2wav stage should be configured with batching:

```yaml
stage_args:
  - stage_id: 2  # code2wav stage
    runtime:
      devices: "1"
    engine_args:
      model_stage: code2wav
      max_num_seqs: 64  # Enables batched audio generation
```

## Related Files

- `vllm_omni/model_executor/stage_input_processors/qwen3_omni.py`: Chunk processing functions (receive `transfer_manager` as first param)
- `vllm_omni/distributed/omni_connectors/transfer_adapter/base.py`: OmniTransferAdapterBase (recv_loop, save_loop, load_async, save_async)
- `vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py`: OmniChunkTransferAdapter (process_pending_chunks, restore_queues, postprocess_scheduler_output)
- `vllm_omni/distributed/omni_connectors/connectors/shm_connector.py`: SharedMemoryConnector (transport-only put/get)
- `vllm_omni/core/sched/omni_ar_scheduler.py`: AR scheduler with chunk_transfer_adapter
- `vllm_omni/core/sched/omni_generation_scheduler.py`: Generation scheduler with same async chunk pattern
- `vllm_omni/worker/gpu_model_runner.py`: Model runner with chunk handling
- `vllm_omni/worker/gpu_generation_model_runner.py`: Generation model runner with batch output handling and ubatch_slices support
- `vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py`: Model implementation with code2wav batching
- `vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_code2wav.py`: Code2wav decoder with batch support
- `vllm_omni/engine/arg_utils.py`: Configuration definitions (async_chunk, worker_type)
- `vllm_omni/config/model.py`: Model config with async_chunk field