# AuK deploy: thinker encoder (stage 0) + rectified-flow DiT (stage 1).
# Topology declared in vllm_omni/model_executor/models/auk/pipeline.py.
#
# Fields omitted from a stage fall back to StageDeployConfig dataclass
# defaults (see vllm_omni/config/stage_config.py). trust_remote_code,
# enable_prefix_caching and enable_chunked_prefill are pipeline-wide
# DeployConfig fields and apply to both stages.
#
# The ODE runs over the whole target segment, so there is nothing to stream
# between the stages: async_chunk stays off.
async_chunk: false
trust_remote_code: true
# Prefix caching must stay off: a cache hit skips the prefix positions, and
# the fused [tokens, 2048] condition needs every position of the prompt.
# Chunked prefill works (the runner concatenates the per-chunk hidden states)
# but measured no wall-time gain, since the encoder is a few percent of a
# request, and it costs bit-exact reproducibility across scheduling.
enable_prefix_caching: false
enable_chunked_prefill: false

stages:
  # Stage 0: Qwen2.5-Omni thinker as a frozen text encoder.
  # Prefill only: one sampled token, discarded downstream. The fused-hidden
  # forward walks the decoder layers itself, so torch.compile has nothing to
  # hook (mode 0); the whole prefill is instead captured as FULL CUDA graphs
  # padded to the sizes below, which removes most of its launch overhead
  # (about 20 ms per request on H200). Longer prompts run eager.
  - stage_id: 0
    devices: "0"
    max_num_seqs: 1
    gpu_memory_utilization: 0.45
    enforce_eager: false
    compilation_config:
      mode: 0
      cudagraph_mode: FULL
      cudagraph_capture_sizes: [64, 128, 192, 256, 384, 512]
    max_num_batched_tokens: 32768
    # AuK takes text plus at most one clip; no image or video tower is built.
    limit_mm_per_prompt: {audio: 1}
    default_sampling_params:
      temperature: 0.0
      max_tokens: 1
      detokenize: false
      seed: 0

  # Stage 1: AuKPipeline (DiT + BigVGAN-flow VAE encoder and decoder).
  # Shares GPU 0 with the encoder; 32 Euler steps at CFG 2.0 is the Base
  # recipe, and AuK-Flash overrides both with its distilled 4-step grid.
  # Not eager: the codec decode is compiled into bucketed CUDA graphs at
  # startup (about a minute more on a cold Inductor cache, ~30 s warm, and
  # about 1 GB of graph pools); clips longer than the largest bucket are
  # decoded in overlapping tiles of it. Each DiT denoise step replays a
  # per-shape CUDA graph of the regionally compiled blocks; the common 3/6/12 s
  # shapes are compiled and captured at startup (auk_dit_warmup_frames), other
  # shapes on first use. Buckets, tile size and the
  # plain-graph cache are tunable through model_config:
  #   model_config: {auk_vae_compile_shapes: [...], auk_vae_tile_frames: <frames>, auk_vae_max_graphs: 32}
  - stage_id: 1
    devices: "0"
    max_num_seqs: 1
    gpu_memory_utilization: 0.35
    enforce_eager: false
    model_config:
      # Codec decode buckets in latent frames (3.2 / 6.4 / 12.8 s). Tighter
      # than the 128/256/512 defaults for the common 3-12 s requests: a
      # 3 s clip decodes in the 160 bucket instead of padding to 256.
      auk_vae_compile_shapes: [160, 320, 640]
      # Maximum number of cached DiT graphs. When the cache reaches this limit,
      # capturing a new shape retires all existing graphs before capture.
      max_dit_graphs: 32
    default_sampling_params:
      seed: 0
      num_inference_steps: 32
      guidance_scale: 2.0