# YuE2-3B deploy: single-stage native-AR text-to-music (~3.6B bf16 + VAE).
# Four-slot CUDA configuration; validate memory on the target card/workload.
#
# The architectures pin is required: the checkpoint's config.json names
# "YuE2ForCausalLM" with an auto_map that only makes sense under
# trust_remote_code, so the deploy pins vLLM to this package's model class.
# model_type stays "yue2" so the registered Yue2Config is used; no
# trust_remote_code is needed (the registered config/model win over the
# checkpoint's auto_map).
pipeline: yue2
async_chunk: false
dtype: bfloat16

stages:
  - stage_id: 0
    # A finishing song keeps its slot during side-stream synthesis,
    # emitting HOLD tokens until its audio is ready.
    max_num_seqs: 4
    # Reserve synthesis headroom: NAR K/V, acoustic graph buffers and VAE
    # activations are allocated beyond the engine budget. The earlier 4090
    # memory measurement predates async/compiled synthesis; 0.70 is a starting
    # point, not a guarantee for 9000 frames or long ABC prefixes. Lower it if
    # the target workload needs more synthesis memory (see the model recipe).
    gpu_memory_utilization: 0.70
    # Keep AR capture off by default. Sampler/NAR graphs still need
    # separate synthesis headroom.
    enforce_eager: true
    # Prefix caching stays off: the omni prefix-cache save snapshots the
    # step's mm payload before the model-owned sampler runs, but YuE2
    # appends the decoded waveform inside sample() (the terminal NAR/VAE
    # pass), so the cached payload would ship empty audio. Re-enable once
    # the runner saves post-sample.
    enable_prefix_caching: false
    max_num_batched_tokens: 24576
    max_model_len: 24576
    devices: "0"
    # The checkpoint ships no HF tokenizer files (only qwen.tiktoken, which
    # the offline driver reads itself); every prompt arrives as token ids,
    # so the engine builds no tokenizer. Placeholder tokenizer dir follows
    # the indextts2 precedent.
    skip_tokenizer_init: true
    tokenizer: gpt2
    hf_overrides:
      architectures: ["Yue2ForCausalLM"]
    default_sampling_params:
      # Defaults for callers that omit model-owned sampling options. The
      # opt-in prepare_runner_inputs hook reads each request's extra_args;
      # it does not depend on the forward's has_sampling_extra_args flag.
      extra_args:
        yue2_phase: "semantic"
        yue2_seed: 831001
        yue2_max_audio_frames: 200
        yue2_max_hold_steps: 4096
      detokenize: false
      stop_token_ids: [151848, 151852]  # ABC_END, MUSIC_END
      # 200 frames + 4096 synthesis holds + the final end token. Callers
      # changing the frame budget must also reserve this token headroom.
      max_tokens: 4297
      temperature: 1.0