75 lines
3.4 KiB
YAML
75 lines
3.4 KiB
YAML
# 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
|