Merge origin/main into dev/rebase-v0170
- Resolved conflict in vllm_omni/engine/output_processor.py: use tracing_enabled for vLLM v0.17.0 API - Resolved conflict in vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code_predictor_vllm.py: keep correct _CodePredictorDecoderLayer.forward() - Update Dockerfile.ci vLLM wheel commit to b31e9326a (releases/v0.17.0) Made-with: Cursor Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>
This commit is contained in:
Taichang Zhou committed
commit
5511bc14c7
37 files changed
+3617
-535
No files matched your search
@@ -62,6 +62,27 @@ steps:
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
- label: ":full_moon: Qwen3-TTS Non-Async-Chunk E2E Test"
|
||||
timeout_in_minutes: 30
|
||||
depends_on: upload-nightly-pipeline
|
||||
if: build.env("NIGHTLY") == "1"
|
||||
commands:
|
||||
- |
|
||||
export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
pytest -s -v tests/e2e/online_serving/test_qwen3_tts.py::TestQwen3TTSNoAsyncChunk
|
||||
agents:
|
||||
queue: "gpu_4_queue"
|
||||
plugins:
|
||||
- docker#v5.2.0:
|
||||
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
always-pull: true
|
||||
propagate-environment: true
|
||||
shm-size: "8gb"
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
- label: ":full_moon: Omni Model Perf Test with H100"
|
||||
key: nightly-performance
|
||||
timeout_in_minutes: 180
|
||||
|
||||
@@ -0,0 +1,301 @@
|
||||
"""Benchmark comparing async_chunk on vs off for Qwen3-TTS.
|
||||
|
||||
Measures TTFP (Time-to-First-Packet), E2E latency, and RTF across
|
||||
concurrency levels for both async_chunk modes. Saves results as JSON.
|
||||
|
||||
Usage:
|
||||
# Run against a server already serving with a given config:
|
||||
python bench_async_chunk.py \
|
||||
--host 127.0.0.1 --port 8000 \
|
||||
--config-name async_chunk_on \
|
||||
--num-prompts 50 \
|
||||
--max-concurrency 1 10 \
|
||||
--result-dir results/
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import json
|
||||
import time
|
||||
from dataclasses import asdict, dataclass, field
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
|
||||
import aiohttp
|
||||
import numpy as np
|
||||
from tqdm.asyncio import tqdm
|
||||
|
||||
PROMPTS = [
|
||||
"Hello, welcome to the voice synthesis benchmark test.",
|
||||
"She said she would be here by noon, but nobody showed up.",
|
||||
"The quick brown fox jumps over the lazy dog near the riverbank.",
|
||||
"I can't believe how beautiful the sunset looks from up here on the mountain.",
|
||||
"Please remember to bring your identification documents to the appointment tomorrow morning.",
|
||||
"Have you ever wondered what it would be like to travel through time and visit ancient civilizations?",
|
||||
"The restaurant on the corner serves the best pasta I have ever tasted in my entire life.",
|
||||
"After the meeting, we should discuss the quarterly results and plan for the next phase.",
|
||||
"Learning a new language takes patience, practice, and a genuine curiosity about other cultures.",
|
||||
"The train leaves at half past seven, so we need to arrive at the station before then.",
|
||||
"Could you please turn down the music a little bit, I'm trying to concentrate on my work.",
|
||||
"It was a dark and stormy night when the old lighthouse keeper heard a knock at the door.",
|
||||
]
|
||||
|
||||
|
||||
@dataclass
|
||||
class RequestResult:
|
||||
success: bool = False
|
||||
ttfp: float = 0.0
|
||||
e2e: float = 0.0
|
||||
audio_bytes: int = 0
|
||||
audio_duration: float = 0.0
|
||||
rtf: float = 0.0
|
||||
prompt: str = ""
|
||||
error: str = ""
|
||||
|
||||
|
||||
@dataclass
|
||||
class BenchmarkResult:
|
||||
config_name: str = ""
|
||||
concurrency: int = 0
|
||||
num_prompts: int = 0
|
||||
completed: int = 0
|
||||
failed: int = 0
|
||||
duration_s: float = 0.0
|
||||
mean_ttfp_ms: float = 0.0
|
||||
median_ttfp_ms: float = 0.0
|
||||
std_ttfp_ms: float = 0.0
|
||||
p90_ttfp_ms: float = 0.0
|
||||
p95_ttfp_ms: float = 0.0
|
||||
p99_ttfp_ms: float = 0.0
|
||||
mean_e2e_ms: float = 0.0
|
||||
median_e2e_ms: float = 0.0
|
||||
std_e2e_ms: float = 0.0
|
||||
p90_e2e_ms: float = 0.0
|
||||
p95_e2e_ms: float = 0.0
|
||||
p99_e2e_ms: float = 0.0
|
||||
mean_rtf: float = 0.0
|
||||
median_rtf: float = 0.0
|
||||
std_rtf: float = 0.0
|
||||
mean_audio_duration_s: float = 0.0
|
||||
total_audio_duration_s: float = 0.0
|
||||
audio_throughput: float = 0.0
|
||||
request_throughput: float = 0.0
|
||||
per_request: list = field(default_factory=list)
|
||||
|
||||
|
||||
def pcm_bytes_to_duration(num_bytes: int, sample_rate: int = 24000, sample_width: int = 2) -> float:
|
||||
return num_bytes / sample_width / sample_rate
|
||||
|
||||
|
||||
async def send_tts_request(
|
||||
session: aiohttp.ClientSession,
|
||||
api_url: str,
|
||||
prompt: str,
|
||||
voice: str = "vivian",
|
||||
language: str = "English",
|
||||
stream: bool = True,
|
||||
pbar: tqdm | None = None,
|
||||
) -> RequestResult:
|
||||
payload = {
|
||||
"input": prompt,
|
||||
"voice": voice,
|
||||
"language": language,
|
||||
"stream": stream,
|
||||
"response_format": "pcm",
|
||||
}
|
||||
|
||||
result = RequestResult(prompt=prompt)
|
||||
st = time.perf_counter()
|
||||
|
||||
try:
|
||||
async with session.post(api_url, json=payload) as response:
|
||||
if response.status != 200:
|
||||
result.error = f"HTTP {response.status}: {await response.text()}"
|
||||
return result
|
||||
|
||||
first_chunk = True
|
||||
total_bytes = 0
|
||||
|
||||
async for chunk in response.content.iter_any():
|
||||
if first_chunk and len(chunk) > 0:
|
||||
result.ttfp = time.perf_counter() - st
|
||||
first_chunk = False
|
||||
total_bytes += len(chunk)
|
||||
|
||||
result.e2e = time.perf_counter() - st
|
||||
result.audio_bytes = total_bytes
|
||||
result.audio_duration = pcm_bytes_to_duration(total_bytes)
|
||||
if result.audio_duration > 0:
|
||||
result.rtf = result.e2e / result.audio_duration
|
||||
result.success = True
|
||||
|
||||
except Exception as e:
|
||||
result.error = str(e)
|
||||
result.e2e = time.perf_counter() - st
|
||||
|
||||
if pbar:
|
||||
pbar.update(1)
|
||||
return result
|
||||
|
||||
|
||||
async def run_benchmark(
|
||||
host: str,
|
||||
port: int,
|
||||
num_prompts: int,
|
||||
max_concurrency: int,
|
||||
num_warmups: int = 3,
|
||||
voice: str = "vivian",
|
||||
language: str = "English",
|
||||
stream: bool = True,
|
||||
) -> BenchmarkResult:
|
||||
api_url = f"http://{host}:{port}/v1/audio/speech"
|
||||
|
||||
connector = aiohttp.TCPConnector(limit=max_concurrency, limit_per_host=max_concurrency, keepalive_timeout=60)
|
||||
session = aiohttp.ClientSession(connector=connector, timeout=aiohttp.ClientTimeout(total=600))
|
||||
|
||||
if num_warmups > 0:
|
||||
print(f" Warming up with {num_warmups} requests...")
|
||||
warmup_tasks = [
|
||||
send_tts_request(session, api_url, PROMPTS[i % len(PROMPTS)], voice, language, stream)
|
||||
for i in range(num_warmups)
|
||||
]
|
||||
await asyncio.gather(*warmup_tasks)
|
||||
print(" Warmup done.")
|
||||
|
||||
request_prompts = [PROMPTS[i % len(PROMPTS)] for i in range(num_prompts)]
|
||||
|
||||
print(f" Running {num_prompts} requests with concurrency={max_concurrency}...")
|
||||
semaphore = asyncio.Semaphore(max_concurrency)
|
||||
pbar = tqdm(total=num_prompts, desc=f" concurrency={max_concurrency}")
|
||||
|
||||
async def limited_request(prompt):
|
||||
async with semaphore:
|
||||
return await send_tts_request(session, api_url, prompt, voice, language, stream, pbar)
|
||||
|
||||
start_time = time.perf_counter()
|
||||
tasks = [asyncio.create_task(limited_request(p)) for p in request_prompts]
|
||||
results: list[RequestResult] = await asyncio.gather(*tasks)
|
||||
duration = time.perf_counter() - start_time
|
||||
pbar.close()
|
||||
|
||||
await session.close()
|
||||
|
||||
successful = [r for r in results if r.success]
|
||||
failed = [r for r in results if not r.success]
|
||||
|
||||
bench = BenchmarkResult(
|
||||
concurrency=max_concurrency,
|
||||
num_prompts=num_prompts,
|
||||
completed=len(successful),
|
||||
failed=len(failed),
|
||||
duration_s=duration,
|
||||
)
|
||||
|
||||
if successful:
|
||||
ttfps = [r.ttfp * 1000 for r in successful]
|
||||
e2es = [r.e2e * 1000 for r in successful]
|
||||
rtfs = [r.rtf for r in successful]
|
||||
audio_durs = [r.audio_duration for r in successful]
|
||||
|
||||
bench.mean_ttfp_ms = float(np.mean(ttfps))
|
||||
bench.median_ttfp_ms = float(np.median(ttfps))
|
||||
bench.std_ttfp_ms = float(np.std(ttfps))
|
||||
bench.p90_ttfp_ms = float(np.percentile(ttfps, 90))
|
||||
bench.p95_ttfp_ms = float(np.percentile(ttfps, 95))
|
||||
bench.p99_ttfp_ms = float(np.percentile(ttfps, 99))
|
||||
|
||||
bench.mean_e2e_ms = float(np.mean(e2es))
|
||||
bench.median_e2e_ms = float(np.median(e2es))
|
||||
bench.std_e2e_ms = float(np.std(e2es))
|
||||
bench.p90_e2e_ms = float(np.percentile(e2es, 90))
|
||||
bench.p95_e2e_ms = float(np.percentile(e2es, 95))
|
||||
bench.p99_e2e_ms = float(np.percentile(e2es, 99))
|
||||
|
||||
bench.mean_rtf = float(np.mean(rtfs))
|
||||
bench.median_rtf = float(np.median(rtfs))
|
||||
bench.std_rtf = float(np.std(rtfs))
|
||||
|
||||
bench.mean_audio_duration_s = float(np.mean(audio_durs))
|
||||
bench.total_audio_duration_s = float(np.sum(audio_durs))
|
||||
bench.audio_throughput = bench.total_audio_duration_s / duration
|
||||
bench.request_throughput = len(successful) / duration
|
||||
|
||||
bench.per_request = [
|
||||
{
|
||||
"ttfp_ms": r.ttfp * 1000,
|
||||
"e2e_ms": r.e2e * 1000,
|
||||
"rtf": r.rtf,
|
||||
"audio_duration_s": r.audio_duration,
|
||||
"prompt": r.prompt,
|
||||
}
|
||||
for r in successful
|
||||
]
|
||||
|
||||
print(f"\n{'=' * 60}")
|
||||
print(f" Concurrency: {max_concurrency} | Completed: {bench.completed} | Failed: {bench.failed}")
|
||||
print(f" Duration: {duration:.2f}s | Throughput: {bench.request_throughput:.2f} req/s")
|
||||
print(
|
||||
f" TTFP (ms): mean={bench.mean_ttfp_ms:.1f} median={bench.median_ttfp_ms:.1f}"
|
||||
f" p90={bench.p90_ttfp_ms:.1f} p99={bench.p99_ttfp_ms:.1f}"
|
||||
)
|
||||
print(
|
||||
f" E2E (ms): mean={bench.mean_e2e_ms:.1f} median={bench.median_e2e_ms:.1f}"
|
||||
f" p90={bench.p90_e2e_ms:.1f} p99={bench.p99_e2e_ms:.1f}"
|
||||
)
|
||||
print(f" RTF: mean={bench.mean_rtf:.3f} median={bench.median_rtf:.3f}")
|
||||
print(f" Throughput: {bench.audio_throughput:.2f} audio-sec/wall-sec")
|
||||
print(f"{'=' * 60}\n")
|
||||
|
||||
if failed:
|
||||
for r in failed[:3]:
|
||||
print(f" [ERROR] {r.error[:200]}")
|
||||
|
||||
return bench
|
||||
|
||||
|
||||
async def main(args):
|
||||
all_results = []
|
||||
|
||||
for concurrency in args.max_concurrency:
|
||||
result = await run_benchmark(
|
||||
host=args.host,
|
||||
port=args.port,
|
||||
num_prompts=args.num_prompts,
|
||||
max_concurrency=concurrency,
|
||||
num_warmups=args.num_warmups,
|
||||
voice=args.voice,
|
||||
language=args.language,
|
||||
stream=args.stream,
|
||||
)
|
||||
result.config_name = args.config_name
|
||||
all_results.append(asdict(result))
|
||||
|
||||
result_dir = Path(args.result_dir)
|
||||
result_dir.mkdir(parents=True, exist_ok=True)
|
||||
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
|
||||
result_file = result_dir / f"bench_{args.config_name}_{timestamp}.json"
|
||||
|
||||
with open(result_file, "w") as f:
|
||||
json.dump(all_results, f, indent=2)
|
||||
print(f"Results saved to {result_file}")
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(description="Qwen3-TTS async_chunk benchmark client")
|
||||
parser.add_argument("--host", type=str, default="127.0.0.1")
|
||||
parser.add_argument("--port", type=int, default=8000)
|
||||
parser.add_argument("--num-prompts", type=int, default=50)
|
||||
parser.add_argument("--max-concurrency", type=int, nargs="+", default=[1, 10])
|
||||
parser.add_argument("--num-warmups", type=int, default=3)
|
||||
parser.add_argument("--voice", type=str, default="vivian")
|
||||
parser.add_argument("--language", type=str, default="English")
|
||||
parser.add_argument("--stream", action="store_true", default=True)
|
||||
parser.add_argument("--no-stream", dest="stream", action="store_false")
|
||||
parser.add_argument("--config-name", type=str, default="async_chunk_on")
|
||||
parser.add_argument("--result-dir", type=str, default="results")
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = parse_args()
|
||||
asyncio.run(main(args))
|
||||
@@ -0,0 +1,249 @@
|
||||
"""Plot TTFP comparison: async_chunk off vs on.
|
||||
|
||||
Generates a bar chart with improvement arrows, matching the Qwen3-Omni
|
||||
async_chunk benchmark figure style.
|
||||
|
||||
Usage:
|
||||
python plot_async_chunk.py \
|
||||
--off results/bench_async_chunk_off_*.json \
|
||||
--on results/bench_async_chunk_on_*.json \
|
||||
--output results/qwen3_tts_async_chunk_ttfp.png
|
||||
|
||||
# Also supports E2E and RTF metrics:
|
||||
python plot_async_chunk.py \
|
||||
--off results/bench_async_chunk_off_*.json \
|
||||
--on results/bench_async_chunk_on_*.json \
|
||||
--metric e2e \
|
||||
--output results/qwen3_tts_async_chunk_e2e.png
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
import matplotlib.pyplot as plt
|
||||
import numpy as np
|
||||
|
||||
METRIC_CONFIG = {
|
||||
"ttfp": {
|
||||
"key": "mean_ttfp_ms",
|
||||
"ylabel": "TTFP (s)",
|
||||
"title": "TTFP (Time to First Audio Packet) - Qwen3-TTS, by concurrency",
|
||||
"to_seconds": True,
|
||||
},
|
||||
"e2e": {
|
||||
"key": "mean_e2e_ms",
|
||||
"ylabel": "E2E (s)",
|
||||
"title": "E2E Latency - Qwen3-TTS, by concurrency",
|
||||
"to_seconds": True,
|
||||
},
|
||||
"rtf": {
|
||||
"key": "mean_rtf",
|
||||
"ylabel": "RTF",
|
||||
"title": "Real-Time Factor - Qwen3-TTS, by concurrency",
|
||||
"to_seconds": False,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def load_results(path: str) -> list[dict]:
|
||||
with open(path) as f:
|
||||
return json.load(f)
|
||||
|
||||
|
||||
def plot_ttfp_comparison(
|
||||
off_results: list[dict],
|
||||
on_results: list[dict],
|
||||
metric: str,
|
||||
output_path: str,
|
||||
title_override: str | None = None,
|
||||
):
|
||||
cfg = METRIC_CONFIG[metric]
|
||||
key = cfg["key"]
|
||||
to_seconds = cfg["to_seconds"]
|
||||
|
||||
off_map = {r["concurrency"]: r for r in off_results}
|
||||
on_map = {r["concurrency"]: r for r in on_results}
|
||||
concurrencies = sorted(set(off_map.keys()) & set(on_map.keys()))
|
||||
|
||||
off_vals = []
|
||||
on_vals = []
|
||||
for c in concurrencies:
|
||||
v_off = off_map[c][key]
|
||||
v_on = on_map[c][key]
|
||||
if to_seconds:
|
||||
v_off /= 1000.0
|
||||
v_on /= 1000.0
|
||||
off_vals.append(v_off)
|
||||
on_vals.append(v_on)
|
||||
|
||||
fig, ax = plt.subplots(figsize=(8, 6))
|
||||
|
||||
x = np.arange(len(concurrencies))
|
||||
width = 0.3
|
||||
|
||||
ax.bar(x - width / 2, off_vals, width, label="async_chunk off", color="#87CEEB", edgecolor="none")
|
||||
ax.bar(x + width / 2, on_vals, width, label="async_chunk on", color="#FFF8DC", edgecolor="#DDD8B8")
|
||||
|
||||
# Draw improvement arrows and labels
|
||||
for i in range(len(concurrencies)):
|
||||
v_off = off_vals[i]
|
||||
v_on = on_vals[i]
|
||||
if v_on > 0:
|
||||
improvement = v_off / v_on
|
||||
else:
|
||||
improvement = float("inf")
|
||||
|
||||
# Arrow from top of off-bar to top of on-bar
|
||||
arrow_start_x = x[i] - width / 2
|
||||
arrow_start_y = v_off * 0.95
|
||||
arrow_end_x = x[i] + width / 2
|
||||
arrow_end_y = v_on * 1.05
|
||||
|
||||
ax.annotate(
|
||||
"",
|
||||
xy=(arrow_end_x, arrow_end_y),
|
||||
xytext=(arrow_start_x, arrow_start_y),
|
||||
arrowprops=dict(arrowstyle="->", color="red", lw=1.5),
|
||||
)
|
||||
|
||||
# Improvement label
|
||||
label_x = (arrow_start_x + arrow_end_x) / 2
|
||||
label_y = arrow_start_y + (v_off - v_on) * 0.15
|
||||
ax.text(
|
||||
label_x,
|
||||
label_y,
|
||||
f"{improvement:.1f}x improvement",
|
||||
ha="center",
|
||||
va="bottom",
|
||||
fontsize=10,
|
||||
color="red",
|
||||
fontweight="bold",
|
||||
)
|
||||
|
||||
title = title_override or cfg["title"]
|
||||
ax.set_title(title, fontsize=13, fontweight="bold")
|
||||
ax.set_ylabel(cfg["ylabel"], fontsize=12)
|
||||
ax.set_xlabel("Max concurrency", fontsize=12)
|
||||
ax.set_xticks(x)
|
||||
ax.set_xticklabels([str(c) for c in concurrencies])
|
||||
ax.set_yscale("log")
|
||||
ax.legend(loc="upper left", fontsize=11)
|
||||
ax.grid(axis="y", alpha=0.3, linestyle="--")
|
||||
ax.set_axisbelow(True)
|
||||
|
||||
plt.tight_layout()
|
||||
Path(output_path).parent.mkdir(parents=True, exist_ok=True)
|
||||
plt.savefig(output_path, dpi=150, bbox_inches="tight")
|
||||
print(f"Plot saved to {output_path}")
|
||||
plt.close()
|
||||
|
||||
|
||||
def plot_all_metrics(off_results: list[dict], on_results: list[dict], output_path: str):
|
||||
"""Generate a 1x3 subplot with TTFP, E2E, and RTF comparisons."""
|
||||
off_map = {r["concurrency"]: r for r in off_results}
|
||||
on_map = {r["concurrency"]: r for r in on_results}
|
||||
concurrencies = sorted(set(off_map.keys()) & set(on_map.keys()))
|
||||
|
||||
fig, axes = plt.subplots(1, 3, figsize=(18, 6))
|
||||
fig.suptitle("Qwen3-TTS: async_chunk on vs off", fontsize=15, fontweight="bold")
|
||||
|
||||
for ax, metric in zip(axes, ["ttfp", "e2e", "rtf"]):
|
||||
cfg = METRIC_CONFIG[metric]
|
||||
key = cfg["key"]
|
||||
to_seconds = cfg["to_seconds"]
|
||||
|
||||
off_vals = []
|
||||
on_vals = []
|
||||
for c in concurrencies:
|
||||
v_off = off_map[c][key]
|
||||
v_on = on_map[c][key]
|
||||
if to_seconds:
|
||||
v_off /= 1000.0
|
||||
v_on /= 1000.0
|
||||
off_vals.append(v_off)
|
||||
on_vals.append(v_on)
|
||||
|
||||
x = np.arange(len(concurrencies))
|
||||
width = 0.3
|
||||
ax.bar(x - width / 2, off_vals, width, label="async_chunk off", color="#87CEEB")
|
||||
ax.bar(x + width / 2, on_vals, width, label="async_chunk on", color="#FFF8DC", edgecolor="#DDD8B8")
|
||||
|
||||
for i in range(len(concurrencies)):
|
||||
if on_vals[i] > 0:
|
||||
improvement = off_vals[i] / on_vals[i]
|
||||
ax.annotate(
|
||||
"",
|
||||
xy=(x[i] + width / 2, on_vals[i] * 1.05),
|
||||
xytext=(x[i] - width / 2, off_vals[i] * 0.95),
|
||||
arrowprops=dict(arrowstyle="->", color="red", lw=1.5),
|
||||
)
|
||||
label_y = off_vals[i] * 0.85
|
||||
ax.text(x[i], label_y, f"{improvement:.1f}x", ha="center", fontsize=10, color="red", fontweight="bold")
|
||||
|
||||
ax.set_title(cfg["title"].split(" - ")[0], fontsize=12, fontweight="bold")
|
||||
ax.set_ylabel(cfg["ylabel"], fontsize=11)
|
||||
ax.set_xlabel("Max concurrency", fontsize=11)
|
||||
ax.set_xticks(x)
|
||||
ax.set_xticklabels([str(c) for c in concurrencies])
|
||||
if metric != "rtf":
|
||||
ax.set_yscale("log")
|
||||
ax.legend(fontsize=9)
|
||||
ax.grid(axis="y", alpha=0.3, linestyle="--")
|
||||
ax.set_axisbelow(True)
|
||||
|
||||
plt.tight_layout()
|
||||
Path(output_path).parent.mkdir(parents=True, exist_ok=True)
|
||||
plt.savefig(output_path, dpi=150, bbox_inches="tight")
|
||||
print(f"Plot saved to {output_path}")
|
||||
plt.close()
|
||||
|
||||
|
||||
def print_table(off_results: list[dict], on_results: list[dict]):
|
||||
off_map = {r["concurrency"]: r for r in off_results}
|
||||
on_map = {r["concurrency"]: r for r in on_results}
|
||||
concurrencies = sorted(set(off_map.keys()) & set(on_map.keys()))
|
||||
|
||||
print("\n## Benchmark Results: async_chunk off vs on\n")
|
||||
print("| Metric | Concurrency | async_chunk off | async_chunk on | Improvement |")
|
||||
print("| --- | --- | --- | --- | --- |")
|
||||
|
||||
for name, key, fmt in [
|
||||
("TTFP (ms)", "mean_ttfp_ms", ".1f"),
|
||||
("E2E (ms)", "mean_e2e_ms", ".1f"),
|
||||
("RTF", "mean_rtf", ".3f"),
|
||||
("Throughput", "audio_throughput", ".2f"),
|
||||
]:
|
||||
for c in concurrencies:
|
||||
v_off = off_map[c].get(key, 0)
|
||||
v_on = on_map[c].get(key, 0)
|
||||
if v_on > 0 and key != "audio_throughput":
|
||||
ratio = f"{v_off / v_on:.1f}x"
|
||||
elif v_off > 0 and key == "audio_throughput":
|
||||
ratio = f"{v_on / v_off:.1f}x"
|
||||
else:
|
||||
ratio = "N/A"
|
||||
print(f"| {name} | {c} | {v_off:{fmt}} | {v_on:{fmt}} | {ratio} |")
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(description="Plot async_chunk comparison for Qwen3-TTS")
|
||||
parser.add_argument("--off", type=str, required=True, help="JSON results for async_chunk off")
|
||||
parser.add_argument("--on", type=str, required=True, help="JSON results for async_chunk on")
|
||||
parser.add_argument("--metric", type=str, default="ttfp", choices=["ttfp", "e2e", "rtf", "all"])
|
||||
parser.add_argument("--output", type=str, default="results/qwen3_tts_async_chunk.png")
|
||||
parser.add_argument("--title", type=str, default=None, help="Custom title override")
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = parse_args()
|
||||
off_results = load_results(args.off)
|
||||
on_results = load_results(args.on)
|
||||
|
||||
print_table(off_results, on_results)
|
||||
|
||||
if args.metric == "all":
|
||||
plot_all_metrics(off_results, on_results, args.output)
|
||||
else:
|
||||
plot_ttfp_comparison(off_results, on_results, args.metric, args.output, args.title)
|
||||
+167
@@ -0,0 +1,167 @@
|
||||
#!/bin/bash
|
||||
# Qwen3-TTS async_chunk on vs off Benchmark
|
||||
#
|
||||
# Starts two servers (async_chunk on and off), benchmarks both,
|
||||
# and generates comparison plots.
|
||||
#
|
||||
# Usage:
|
||||
# bash run_async_chunk_benchmark.sh
|
||||
#
|
||||
# Environment variables:
|
||||
# GPU_DEVICE - GPU index (default: 0)
|
||||
# MODEL - Model path (default: Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice)
|
||||
# NUM_PROMPTS - Prompts per concurrency level (default: 50)
|
||||
# CONCURRENCY - Space-separated concurrency levels (default: "1 10")
|
||||
# PORT_ON - Port for async_chunk on server (default: 8000)
|
||||
# PORT_OFF - Port for async_chunk off server (default: 8001)
|
||||
|
||||
set -euo pipefail
|
||||
|
||||
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
||||
PROJECT_ROOT="$(cd "${SCRIPT_DIR}/../../.." && pwd)"
|
||||
cd "$PROJECT_ROOT"
|
||||
|
||||
GPU_DEVICE="${GPU_DEVICE:-0}"
|
||||
MODEL="${MODEL:-Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice}"
|
||||
NUM_PROMPTS="${NUM_PROMPTS:-50}"
|
||||
CONCURRENCY="${CONCURRENCY:-1 10}"
|
||||
NUM_WARMUPS="${NUM_WARMUPS:-3}"
|
||||
PORT_ON="${PORT_ON:-8000}"
|
||||
PORT_OFF="${PORT_OFF:-8001}"
|
||||
RESULT_DIR="${SCRIPT_DIR}/results"
|
||||
TIMESTAMP="$(date +%Y%m%d_%H%M%S)"
|
||||
|
||||
STAGE_CONFIG_ON="vllm_omni/model_executor/stage_configs/qwen3_tts.yaml"
|
||||
STAGE_CONFIG_OFF="vllm_omni/model_executor/stage_configs/qwen3_tts_no_async_chunk.yaml"
|
||||
|
||||
mkdir -p "${RESULT_DIR}"
|
||||
|
||||
echo "============================================================"
|
||||
echo " Qwen3-TTS async_chunk Benchmark"
|
||||
echo "============================================================"
|
||||
echo " GPU: ${GPU_DEVICE}"
|
||||
echo " Model: ${MODEL}"
|
||||
echo " Prompts: ${NUM_PROMPTS}"
|
||||
echo " Concurrency: ${CONCURRENCY}"
|
||||
echo " Port (on/off): ${PORT_ON} / ${PORT_OFF}"
|
||||
echo " Results: ${RESULT_DIR}"
|
||||
echo "============================================================"
|
||||
|
||||
cleanup() {
|
||||
echo "Cleaning up servers..."
|
||||
kill "$PID_ON" 2>/dev/null || true
|
||||
kill "$PID_OFF" 2>/dev/null || true
|
||||
wait "$PID_ON" 2>/dev/null || true
|
||||
wait "$PID_OFF" 2>/dev/null || true
|
||||
}
|
||||
trap cleanup EXIT
|
||||
|
||||
wait_for_server() {
|
||||
local port=$1
|
||||
local name=$2
|
||||
local max_wait=300
|
||||
local elapsed=0
|
||||
echo "Waiting for ${name} server on port ${port}..."
|
||||
while ! curl -s "http://localhost:${port}/health" >/dev/null 2>&1; do
|
||||
sleep 5
|
||||
elapsed=$((elapsed + 5))
|
||||
if [ $elapsed -ge $max_wait ]; then
|
||||
echo "ERROR: ${name} server failed to start within ${max_wait}s"
|
||||
exit 1
|
||||
fi
|
||||
done
|
||||
echo "${name} server ready (${elapsed}s)"
|
||||
}
|
||||
|
||||
# ---- Phase 1: Start async_chunk ON server ----
|
||||
echo ""
|
||||
echo "[Phase 1] Starting async_chunk ON server on port ${PORT_ON}..."
|
||||
CUDA_VISIBLE_DEVICES=${GPU_DEVICE} vllm-omni serve "${MODEL}" \
|
||||
--stage-configs-path "${STAGE_CONFIG_ON}" \
|
||||
--host 0.0.0.0 --port "${PORT_ON}" \
|
||||
--trust-remote-code --enforce-eager --omni \
|
||||
> "${RESULT_DIR}/server_on_${TIMESTAMP}.log" 2>&1 &
|
||||
PID_ON=$!
|
||||
|
||||
wait_for_server "${PORT_ON}" "async_chunk_on"
|
||||
|
||||
echo "[Phase 1] Benchmarking async_chunk ON..."
|
||||
# shellcheck disable=SC2086
|
||||
python "${SCRIPT_DIR}/bench_async_chunk.py" \
|
||||
--host 127.0.0.1 --port "${PORT_ON}" \
|
||||
--config-name "async_chunk_on" \
|
||||
--num-prompts "${NUM_PROMPTS}" \
|
||||
--max-concurrency ${CONCURRENCY} \
|
||||
--num-warmups "${NUM_WARMUPS}" \
|
||||
--result-dir "${RESULT_DIR}"
|
||||
|
||||
echo "[Phase 1] Stopping async_chunk ON server..."
|
||||
kill "$PID_ON" 2>/dev/null || true
|
||||
wait "$PID_ON" 2>/dev/null || true
|
||||
sleep 5
|
||||
|
||||
# ---- Phase 2: Start async_chunk OFF server ----
|
||||
echo ""
|
||||
echo "[Phase 2] Starting async_chunk OFF server on port ${PORT_OFF}..."
|
||||
CUDA_VISIBLE_DEVICES=${GPU_DEVICE} vllm-omni serve "${MODEL}" \
|
||||
--stage-configs-path "${STAGE_CONFIG_OFF}" \
|
||||
--host 0.0.0.0 --port "${PORT_OFF}" \
|
||||
--trust-remote-code --enforce-eager --omni \
|
||||
> "${RESULT_DIR}/server_off_${TIMESTAMP}.log" 2>&1 &
|
||||
PID_OFF=$!
|
||||
|
||||
wait_for_server "${PORT_OFF}" "async_chunk_off"
|
||||
|
||||
echo "[Phase 2] Benchmarking async_chunk OFF (non-streaming)..."
|
||||
# shellcheck disable=SC2086
|
||||
python "${SCRIPT_DIR}/bench_async_chunk.py" \
|
||||
--host 127.0.0.1 --port "${PORT_OFF}" \
|
||||
--config-name "async_chunk_off" \
|
||||
--num-prompts "${NUM_PROMPTS}" \
|
||||
--max-concurrency ${CONCURRENCY} \
|
||||
--num-warmups "${NUM_WARMUPS}" \
|
||||
--no-stream \
|
||||
--result-dir "${RESULT_DIR}"
|
||||
|
||||
echo "[Phase 2] Stopping async_chunk OFF server..."
|
||||
kill "$PID_OFF" 2>/dev/null || true
|
||||
wait "$PID_OFF" 2>/dev/null || true
|
||||
|
||||
# ---- Phase 3: Plot results ----
|
||||
echo ""
|
||||
echo "[Phase 3] Generating plots..."
|
||||
|
||||
# Find the latest result files
|
||||
RESULT_ON=$(ls -t "${RESULT_DIR}"/bench_async_chunk_on_*.json 2>/dev/null | head -1)
|
||||
RESULT_OFF=$(ls -t "${RESULT_DIR}"/bench_async_chunk_off_*.json 2>/dev/null | head -1)
|
||||
|
||||
if [ -z "$RESULT_ON" ] || [ -z "$RESULT_OFF" ]; then
|
||||
echo "ERROR: Could not find result files. Check logs in ${RESULT_DIR}/"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
echo " ON results: ${RESULT_ON}"
|
||||
echo " OFF results: ${RESULT_OFF}"
|
||||
|
||||
# TTFP comparison (main figure)
|
||||
python "${SCRIPT_DIR}/plot_async_chunk.py" \
|
||||
--off "${RESULT_OFF}" \
|
||||
--on "${RESULT_ON}" \
|
||||
--metric ttfp \
|
||||
--output "${RESULT_DIR}/qwen3_tts_async_chunk_ttfp.png"
|
||||
|
||||
# All metrics comparison
|
||||
python "${SCRIPT_DIR}/plot_async_chunk.py" \
|
||||
--off "${RESULT_OFF}" \
|
||||
--on "${RESULT_ON}" \
|
||||
--metric all \
|
||||
--output "${RESULT_DIR}/qwen3_tts_async_chunk_all.png"
|
||||
|
||||
echo ""
|
||||
echo "============================================================"
|
||||
echo " Benchmark complete!"
|
||||
echo " Results: ${RESULT_DIR}/"
|
||||
echo " Plots:"
|
||||
echo " - ${RESULT_DIR}/qwen3_tts_async_chunk_ttfp.png"
|
||||
echo " - ${RESULT_DIR}/qwen3_tts_async_chunk_all.png"
|
||||
echo "============================================================"
|
||||
@@ -14,7 +14,7 @@ RUN apt-get update && \
|
||||
|
||||
# Install vLLM wheel for the specific commit
|
||||
# URL format: https://wheels.vllm.ai/<commit_hash>/
|
||||
# Current vLLM wheel commit: 097eb544e9a22810c9b7a59e586b61627b308362 (main)
|
||||
# Current vLLM wheel commit: b31e9326a7d9394aab8c767f8ebe225c65594b60 (releases/v0.17.0)
|
||||
# Note: --prerelease=allow is required because nightly/commit wheels use
|
||||
# pre-release version numbers (e.g. 2.10.0.dev...) which uv ignores by default,
|
||||
# causing it to fall back to the older stable release from PyPI.
|
||||
|
||||
@@ -66,6 +66,7 @@ nav:
|
||||
- Model Implementation:
|
||||
- contributing/model/README.md
|
||||
- contributing/model/adding_omni_model.md
|
||||
- contributing/model/adding_tts_model.md
|
||||
- contributing/model/adding_diffusion_model.md
|
||||
- CI: contributing/ci
|
||||
- Design Documents:
|
||||
|
||||
@@ -69,8 +69,20 @@ With `--log-stats` enabled, the server will output detailed metrics logs after e
|
||||
| rx_decode_time_ms | 111.865 | 31.706 |
|
||||
| in_flight_time_ms | 2.015 | 2.819 |
|
||||
|
||||
### DiffusionStats
|
||||
|
||||
These logs include:
|
||||
| Field | Value |
|
||||
|---------------------------------|------------|
|
||||
| batch_id | 1 |
|
||||
| batch_size | 1 |
|
||||
| diffusion_engine_exec_time_ms | 16,363.419 |
|
||||
| diffusion_engine_total_time_ms | 16,234.849 |
|
||||
| image_num | 1.000 |
|
||||
| num_inference_steps | 50.000 |
|
||||
| postprocess_time_ms | 67.685 |
|
||||
| preprocess_time_ms | 60.106 |
|
||||
| resolution | 640.000 |
|
||||
| stage_gen_time_ms | 16,363.714 |
|
||||
|
||||
- **Overall summary**: total requests, wall time, average tokens/sec, etc.
|
||||
|
||||
@@ -80,6 +92,9 @@ These logs include:
|
||||
|
||||
- **Transfer table**: data transfer and timing for each edge.
|
||||
|
||||
- **DiffusionStats**:preprocessing time, model execution time, postprocessing time, generated images, resolution, and inference steps.
|
||||
|
||||
|
||||
You can use these logs to monitor system health, debug performance, and analyze request-level metrics as described above.
|
||||
|
||||
|
||||
@@ -130,6 +145,11 @@ For **online inference** (serving mode), the summary is always per-request. `e2e
|
||||
| `image_num` | Number of images generated (for diffusion/image stages). |
|
||||
| `resolution` | Image resolution (for diffusion/image stages). |
|
||||
| `postprocess_time_ms` | Diffusion/image: post-processing time in ms. |
|
||||
| `preprocess_time_ms` | Diffusion/image: pre-processing time in ms (before model execution). |
|
||||
| `diffusion_engine_exec_time_ms` | Diffusion/image: model execution time in ms, excluding pre-processing and post-processing. |
|
||||
| `diffusion_engine_total_time_ms` | Diffusion/image: total diffusion engine time in ms, including pre-processing, model execution, and post-processing. |
|
||||
| `num_inference_steps` | Number of inference steps for diffusion generation. |
|
||||
|
||||
|
||||
---
|
||||
|
||||
@@ -143,6 +163,8 @@ For **online inference** (serving mode), the summary is always per-request. `e2e
|
||||
| `in_flight_time_ms` | In-flight time in ms. |
|
||||
|
||||
|
||||
|
||||
|
||||
### Expectation of the Numbers (Verification)
|
||||
|
||||
**Formulas:**
|
||||
@@ -171,3 +193,14 @@ For each edge:
|
||||
- 1->2: tx_time_ms (**18.790**) + rx_decode_time_ms (**31.706**) + in_flight_time_ms (**2.819**) = **53.315**
|
||||
|
||||
192.581 + 53.315 = **245.896** = transfers_total_time_ms, which matches the calculation (difference is due to rounding)
|
||||
|
||||
**diffusion_engine_total_time_ms**
|
||||
|
||||
- preprocess_time_ms: Time spent on input preprocessing
|
||||
- diffusion_engine_exec_time_ms: Pure model execution time
|
||||
- postprocess_time_ms: Time spent on output postprocessing
|
||||
- diffusion_engine_total_time_ms: Total time spent in the diffusion engine
|
||||
|
||||
so `diffusion_engine_total_time_ms` = `diffusion_engine_exec_time_ms` + `preprocess_time_ms` + `postprocess_time_ms`
|
||||
= **16,363.419** + **60.106** + **67.685**
|
||||
= **16,491.210**
|
||||
@@ -0,0 +1,584 @@
|
||||
# Adding a TTS Model
|
||||
|
||||
This guide walks through adding a new TTS model to vLLM-Omni, using **Qwen3-TTS**
|
||||
as a reference. Qwen3-TTS demonstrates the standard two-stage TTS pipeline and the
|
||||
key optimizations all TTS models in this repo should follow.
|
||||
|
||||
## Table of Contents
|
||||
|
||||
1. [Overview](#overview)
|
||||
2. [Directory Structure](#directory-structure)
|
||||
3. [Step-by-Step Implementation](#step-by-step-implementation)
|
||||
4. [Key Components](#key-components)
|
||||
5. [Model Registration](#model-registration)
|
||||
6. [Stage Configuration](#stage-configuration)
|
||||
7. [Stage Input Processors](#stage-input-processors)
|
||||
8. [Testing](#testing)
|
||||
9. [Summary](#summary)
|
||||
|
||||
## Overview
|
||||
|
||||
vLLM-Omni supports TTS models as multi-stage pipelines where each stage runs independently
|
||||
and can be placed on different devices. Qwen3-TTS has two stages:
|
||||
|
||||
| Stage | Name | Input | Output |
|
||||
|-------|------|-------|--------|
|
||||
| 0 | Code Predictor (AR) | Text tokens | Discrete RVQ codec codes |
|
||||
| 1 | Code2Wav (Decoder) | RVQ codec codes | Audio waveform |
|
||||
|
||||
Each stage is a separate model class configured independently via YAML. The two stages
|
||||
are connected by the `async_chunk` framework, which enables inter-stage streaming for
|
||||
low first-packet latency (see [Async Chunk Design](../../design/feature/async_chunk_design.md)).
|
||||
|
||||
### Without async_chunk (batch mode)
|
||||
|
||||
Stage 0 runs to completion before Stage 1 starts, resulting in long first-packet latency:
|
||||
|
||||
```mermaid
|
||||
flowchart TB
|
||||
subgraph stage0["Stage 0: AR Code Predictor"]
|
||||
direction LR
|
||||
P[Prefill] --> D1[Decode 1]
|
||||
D1 --> D2[Decode 2]
|
||||
D2 --> Dots1["..."]
|
||||
Dots1 --> DN[Decode N]
|
||||
end
|
||||
|
||||
subgraph stage1["Stage 1: Code2Wav"]
|
||||
direction LR
|
||||
DEC[Decode all codes at once]
|
||||
end
|
||||
|
||||
stage0 -- "all N codes" --> stage1
|
||||
stage1 --> FPL["First Packet Latency = Stage 0 + Stage 1"]
|
||||
|
||||
style stage0 fill:#dae8fc,stroke:#6c8ebf
|
||||
style stage1 fill:#f8d7c8,stroke:#d4856a
|
||||
style FPL fill:#e8f0fe,stroke:#3366CC,stroke-width:2px
|
||||
```
|
||||
|
||||
### With async_chunk (streaming mode)
|
||||
|
||||
Stage 0 sends codec codes to Stage 1 every `chunk_size=25` tokens. Stage 1 begins decoding
|
||||
immediately, reducing first-packet latency from the full AR time to just the first chunk:
|
||||
|
||||
```mermaid
|
||||
flowchart TB
|
||||
subgraph stage0["Stage 0: Code Predictor (AR)"]
|
||||
direction LR
|
||||
P[Prefill] --> D1["Decode 1-25"]
|
||||
D1 --> D2["Decode 26-50"]
|
||||
D2 --> Dots1["..."]
|
||||
Dots1 --> DN["Decode N"]
|
||||
end
|
||||
|
||||
subgraph stage1["Stage 1: Code2Wav"]
|
||||
direction LR
|
||||
C1["Chunk 1\n(25 frames)"] --> C2["Chunk 2\n(context + 25)"]
|
||||
C2 --> Dots2["..."]
|
||||
Dots2 --> CN["Final chunk"]
|
||||
end
|
||||
|
||||
D1 -. "chunk 1 (25 codes)" .-> C1
|
||||
D2 -. "chunk 2 (context + 25)" .-> C2
|
||||
DN -. "final" .-> CN
|
||||
|
||||
stage0 --> FPL["⏱ First Packet Latency = Prefill + 25 decode steps only"]
|
||||
|
||||
style stage0 fill:#dae8fc,stroke:#6c8ebf
|
||||
style stage1 fill:#e8d4f8,stroke:#8a6cad
|
||||
style FPL fill:#e8f0fe,stroke:#3366CC,stroke-width:2px
|
||||
```
|
||||
|
||||
Key parameters: `chunk_size=25`, `left_context_size=25` (validated defaults from Qwen3-TTS
|
||||
and Qwen3-Omni).
|
||||
|
||||
## Directory Structure
|
||||
|
||||
When adding a new TTS model, create the following structure:
|
||||
|
||||
```
|
||||
vllm_omni/model_executor/models/
|
||||
your_model_name/
|
||||
__init__.py
|
||||
your_model.py # Unified class (stage dispatch)
|
||||
your_model_ar_stage.py # Stage 0: AR stage
|
||||
your_model_decoder.py # Stage 1: audio decoder
|
||||
|
||||
vllm_omni/model_executor/stage_input_processors/
|
||||
your_model_name.py # Stage 0 -> Stage 1 transition
|
||||
|
||||
vllm_omni/model_executor/stage_configs/
|
||||
your_model_name.yaml # Batch mode config
|
||||
your_model_name_async_chunk.yaml # Streaming mode config
|
||||
```
|
||||
|
||||
**Qwen3-TTS reference files:**
|
||||
|
||||
| File | Purpose |
|
||||
|------|---------|
|
||||
| `models/qwen3_tts/qwen3_tts.py` | Unified model class |
|
||||
| `models/qwen3_tts/qwen3_tts_code_predictor_vllm.py` | Stage 0 - optimized AR |
|
||||
| `models/qwen3_tts/qwen3_tts_code2wav.py` | Stage 1 - decoder |
|
||||
| `stage_configs/qwen3_tts.yaml` | Stage config (async_chunk enabled) |
|
||||
| `stage_configs/qwen3_tts_batch.yaml` | Batch mode config |
|
||||
| `stage_input_processors/qwen3_tts.py` | Stage transition processors |
|
||||
|
||||
## Step-by-Step Implementation
|
||||
|
||||
### Step 1: Implement Stage 0 - AR Stage
|
||||
|
||||
Stage 0 is the autoregressive stage that generates intermediate audio representations.
|
||||
**It must use vLLM's native decoder layers with fused ops and PagedAttention** for the LLM
|
||||
backbone - this is the primary source of speedup over HuggingFace inference.
|
||||
|
||||
#### 1.1 Use vLLM Decoder Layers Directly
|
||||
|
||||
Build your transformer layers from the corresponding vLLM decoder layer class (e.g.
|
||||
`Qwen3DecoderLayer` for Qwen3-based backbones, or the equivalent for LLaMA, Qwen2, etc.).
|
||||
Do not wrap the HuggingFace model directly - that bypasses PagedAttention and fused kernels.
|
||||
|
||||
```python
|
||||
# your_model_ar_stage.py
|
||||
|
||||
from vllm.model_executor.models.qwen3 import Qwen3DecoderLayer
|
||||
|
||||
class YourTTSARStage(nn.Module):
|
||||
|
||||
def __init__(self, config, vllm_config, prefix):
|
||||
self.layers = nn.ModuleList([
|
||||
Qwen3DecoderLayer(
|
||||
config, vllm_config=vllm_config, prefix=f"{prefix}.layers.{i}"
|
||||
)
|
||||
for i in range(config.num_hidden_layers)
|
||||
])
|
||||
self.lm_head = ParallelLMHead(config.codec_size, config.hidden_size)
|
||||
```
|
||||
|
||||
See `qwen3_tts_code_predictor_vllm.py` for the full implementation.
|
||||
|
||||
#### 1.2 Forward Pass
|
||||
|
||||
Implement `forward()` to return an `OmniOutput` with intermediate data for Stage 1:
|
||||
|
||||
```python
|
||||
def forward(self, input_ids, positions, intermediate_tensors=None,
|
||||
inputs_embeds=None, **kwargs) -> OmniOutput:
|
||||
hidden_states = self.run_layers(input_ids, positions, intermediate_tensors, inputs_embeds)
|
||||
logits = self.lm_head(hidden_states)
|
||||
|
||||
return OmniOutput(
|
||||
text_hidden_states=hidden_states,
|
||||
multimodal_outputs={
|
||||
"audio_codes": self.extract_codes(logits),
|
||||
},
|
||||
)
|
||||
```
|
||||
|
||||
The keys in `multimodal_outputs` are what your stage input processor will read to build
|
||||
Stage 1 inputs.
|
||||
|
||||
#### 1.3 Weight Loading with Fused QKV
|
||||
|
||||
When using vLLM's fused `QKVParallelLinear`, pack the HF `q_proj`/`k_proj`/`v_proj` weights
|
||||
into `qkv_proj` using `stacked_params_mapping`. See the `load_weights()` method in
|
||||
`qwen3_tts_code_predictor_vllm.py` for the standard pattern - it can be reused as-is
|
||||
for any Qwen-family backbone.
|
||||
|
||||
#### 1.4 Custom Stop Condition (if needed)
|
||||
|
||||
Some TTS models use a learned stop head rather than an EOS token. If your model does this,
|
||||
implement it inside `sample()`:
|
||||
|
||||
```python
|
||||
def sample(self, logits, sampling_metadata) -> SamplerOutput | None:
|
||||
output = self.sampler(logits, sampling_metadata)
|
||||
if self._stop_head_fired():
|
||||
output = mark_as_finished(output)
|
||||
return output
|
||||
```
|
||||
|
||||
### Step 2: Implement Stage 1 - Decoder
|
||||
|
||||
Stage 1 decodes Stage 0 output into audio. It runs outside the scheduler (no PagedAttention
|
||||
needed). Implement `chunked_decode_streaming()` to support async_chunk streaming:
|
||||
|
||||
```python
|
||||
# your_model_decoder.py
|
||||
|
||||
class YourTTSDecoder(nn.Module):
|
||||
|
||||
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
|
||||
super().__init__()
|
||||
# Initialize your audio decoder (SpeechTokenizer, HiFiGAN, etc.)
|
||||
|
||||
def forward(self, codes: torch.Tensor, **kwargs) -> torch.Tensor:
|
||||
return self.decoder(codes)
|
||||
|
||||
def chunked_decode_streaming(self, codes, chunk_size=25,
|
||||
left_context_size=25) -> torch.Tensor:
|
||||
"""Decode with a sliding context window for smooth chunk boundaries."""
|
||||
end_index = codes.shape[-1]
|
||||
context_size = 0 if end_index <= chunk_size else left_context_size
|
||||
wav_chunk = self(codes)
|
||||
# Trim left context to avoid duplicate audio
|
||||
return wav_chunk[..., context_size * self.total_upsample:]
|
||||
```
|
||||
|
||||
### Step 3: Implement the Unified Model Class
|
||||
|
||||
The unified class dispatches to the correct stage based on `model_stage` in the config:
|
||||
|
||||
```python
|
||||
# your_model.py
|
||||
|
||||
class YourTTSModelForConditionalGeneration(nn.Module, SupportsPP):
|
||||
|
||||
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
|
||||
super().__init__()
|
||||
self.model_stage = vllm_config.model_config.model_stage
|
||||
|
||||
if self.model_stage == "ar_stage":
|
||||
ar_vllm_config = vllm_config.with_hf_config(
|
||||
vllm_config.model_config.hf_config.ar_config,
|
||||
architectures=["YourTTSARStageForConditionalGeneration"],
|
||||
)
|
||||
self.ar_stage = init_vllm_registered_model(
|
||||
vllm_config=ar_vllm_config,
|
||||
prefix=maybe_prefix(prefix, "ar"),
|
||||
hf_config=ar_vllm_config.model_config.hf_config,
|
||||
architectures=["YourTTSARStageForConditionalGeneration"],
|
||||
)
|
||||
self.model = self.ar_stage
|
||||
|
||||
elif self.model_stage == "decoder":
|
||||
self.decoder = YourTTSDecoder(vllm_config=vllm_config, prefix=prefix)
|
||||
self.model = self.decoder
|
||||
```
|
||||
|
||||
### Step 4: Create `__init__.py`
|
||||
|
||||
```python
|
||||
# vllm_omni/model_executor/models/your_model_name/__init__.py
|
||||
from .your_model import YourTTSModelForConditionalGeneration
|
||||
|
||||
__all__ = ["YourTTSModelForConditionalGeneration"]
|
||||
```
|
||||
|
||||
## Key Components
|
||||
|
||||
### Model Interfaces
|
||||
|
||||
Your unified model class should implement the appropriate interfaces:
|
||||
|
||||
- **`SupportsPP`**: Required for pipeline parallelism support (all models should implement this)
|
||||
- **`SupportsMultiModal`**: Only if your model accepts multimodal inputs (e.g. reference audio for voice cloning)
|
||||
|
||||
### Output Format
|
||||
|
||||
Use `OmniOutput` so the orchestrator can route intermediate data between stages:
|
||||
|
||||
```python
|
||||
from vllm_omni.model_executor.models.output_templates import OmniOutput
|
||||
|
||||
return OmniOutput(
|
||||
text_hidden_states=hidden_states,
|
||||
multimodal_outputs={
|
||||
"audio_codes": codec_codes,
|
||||
},
|
||||
)
|
||||
```
|
||||
|
||||
### Weight Loading from a Single Checkpoint
|
||||
|
||||
If both stages load from one checkpoint, separate them by prefix in the unified class:
|
||||
|
||||
```python
|
||||
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
|
||||
ar_weights, decoder_weights = [], []
|
||||
for name, tensor in weights:
|
||||
if name.startswith("decoder."):
|
||||
decoder_weights.append((name, tensor))
|
||||
else:
|
||||
ar_weights.append((name, tensor))
|
||||
|
||||
if self.model_stage == "ar_stage":
|
||||
return self.ar_stage.load_weights(ar_weights)
|
||||
elif self.model_stage == "decoder":
|
||||
return self.decoder.load_weights(decoder_weights)
|
||||
```
|
||||
|
||||
## Model Registration
|
||||
|
||||
Register all stage classes in `vllm_omni/model_executor/models/registry.py`:
|
||||
|
||||
```python
|
||||
_OMNI_MODELS = {
|
||||
# (package_name, module_name, class_name)
|
||||
"YourTTSModelForConditionalGeneration": (
|
||||
"your_model_name", "your_model",
|
||||
"YourTTSModelForConditionalGeneration",
|
||||
),
|
||||
"YourTTSARStageForConditionalGeneration": (
|
||||
"your_model_name", "your_model_ar_stage",
|
||||
"YourTTSARStageForConditionalGeneration",
|
||||
),
|
||||
"YourTTSDecoder": (
|
||||
"your_model_name", "your_model_decoder",
|
||||
"YourTTSDecoder",
|
||||
),
|
||||
}
|
||||
```
|
||||
|
||||
The registry uses lazy loading - model classes are only imported when needed.
|
||||
|
||||
## Stage Configuration
|
||||
|
||||
Each stage has a `worker_type` that determines how it is scheduled:
|
||||
|
||||
- `worker_type: ar` - autoregressive stage, uses `OmniARScheduler` with PagedAttention
|
||||
- `worker_type: generation` - non-AR stage (e.g. decoder), uses `OmniGenerationScheduler`
|
||||
|
||||
Key configuration fields:
|
||||
|
||||
| Field | Description |
|
||||
|-------|-------------|
|
||||
| `model_stage` | Which stage to initialize (`ar_stage`, `decoder`, etc.) |
|
||||
| `model_arch` | Architecture name, must match `registry.py` |
|
||||
| `engine_input_source` | List of upstream stage IDs that provide input (e.g. `[0]`) |
|
||||
| `engine_output_type` | Output type: `latent` for intermediate, `audio` for final |
|
||||
| `custom_process_next_stage_input_func` | Async chunk processor function path (streaming only) |
|
||||
| `final_output` | Whether this stage produces the final user-facing output |
|
||||
| `final_output_type` | Type of final output (`audio`, `text`, etc.) |
|
||||
|
||||
### Batch mode
|
||||
|
||||
```yaml
|
||||
# stage_configs/your_model_name.yaml
|
||||
|
||||
stage_args:
|
||||
- stage_id: 0
|
||||
stage_type: llm
|
||||
runtime:
|
||||
devices: "0"
|
||||
max_batch_size: 64
|
||||
engine_args:
|
||||
model_stage: ar_stage
|
||||
model_arch: YourTTSModelForConditionalGeneration
|
||||
worker_type: ar
|
||||
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
|
||||
engine_output_type: latent
|
||||
default_sampling_params:
|
||||
temperature: 0.9
|
||||
top_k: 50
|
||||
max_tokens: 2048
|
||||
|
||||
- stage_id: 1
|
||||
stage_type: llm
|
||||
runtime:
|
||||
devices: "0"
|
||||
engine_args:
|
||||
model_stage: decoder
|
||||
model_arch: YourTTSModelForConditionalGeneration
|
||||
worker_type: generation
|
||||
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
|
||||
engine_output_type: audio
|
||||
engine_input_source: [0]
|
||||
final_output: true
|
||||
final_output_type: audio
|
||||
```
|
||||
|
||||
### Streaming mode (async_chunk)
|
||||
|
||||
Add `async_chunk: true` at the top level and specify `custom_process_next_stage_input_func`
|
||||
on Stage 0 to define how intermediate outputs are chunked and forwarded:
|
||||
|
||||
```yaml
|
||||
# stage_configs/your_model_name_async_chunk.yaml
|
||||
|
||||
async_chunk: true
|
||||
|
||||
stage_args:
|
||||
- stage_id: 0
|
||||
stage_type: llm
|
||||
runtime:
|
||||
devices: "0"
|
||||
max_batch_size: 64
|
||||
engine_args:
|
||||
model_stage: ar_stage
|
||||
model_arch: YourTTSModelForConditionalGeneration
|
||||
worker_type: ar
|
||||
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
|
||||
engine_output_type: latent
|
||||
custom_process_next_stage_input_func: >
|
||||
vllm_omni.model_executor.stage_input_processors.your_model_name.ar2decoder_async_chunk
|
||||
default_sampling_params:
|
||||
temperature: 0.9
|
||||
top_k: 50
|
||||
max_tokens: 2048
|
||||
|
||||
- stage_id: 1
|
||||
stage_type: llm
|
||||
runtime:
|
||||
devices: "0"
|
||||
engine_args:
|
||||
model_stage: decoder
|
||||
model_arch: YourTTSModelForConditionalGeneration
|
||||
worker_type: generation
|
||||
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
|
||||
engine_output_type: audio
|
||||
engine_input_source: [0]
|
||||
final_output: true
|
||||
final_output_type: audio
|
||||
```
|
||||
|
||||
## Stage Input Processors
|
||||
|
||||
Stage input processors convert Stage 0 outputs into Stage 1 inputs. Create yours in
|
||||
`vllm_omni/model_executor/stage_input_processors/your_model_name.py`.
|
||||
|
||||
See `stage_input_processors/qwen3_tts.py` for the full reference implementation.
|
||||
|
||||
### Data structures
|
||||
|
||||
Understanding what's available in stage outputs:
|
||||
|
||||
- `stage_list[source_id].engine_outputs` - list of `EngineCoreOutput` objects
|
||||
- Each `EngineCoreOutput` has `outputs` - list of `RequestOutput` objects
|
||||
- Each `RequestOutput` has:
|
||||
- `token_ids` - generated token IDs
|
||||
- `multimodal_output` - dict with keys matching your model's `OmniOutput.multimodal_outputs`
|
||||
- `prompt_token_ids` - original prompt token IDs
|
||||
|
||||
### Batch mode (non-streaming)
|
||||
|
||||
Collects all Stage 0 outputs and forwards them to Stage 1 in one shot:
|
||||
|
||||
```python
|
||||
def ar2decoder(
|
||||
stage_list: list[Any],
|
||||
engine_input_source: list[int],
|
||||
prompt: OmniTokensPrompt | TextPrompt | None = None,
|
||||
requires_multimodal_data: bool = False,
|
||||
) -> list[OmniTokensPrompt]:
|
||||
source_id = engine_input_source[0]
|
||||
decoder_inputs = []
|
||||
|
||||
for output in stage_list[source_id].engine_outputs:
|
||||
result = output.outputs[0]
|
||||
codes = result.multimodal_output["audio_codes"].cpu()
|
||||
decoder_inputs.append(
|
||||
OmniTokensPrompt(prompt_token_ids=codes.reshape(-1).tolist())
|
||||
)
|
||||
|
||||
return decoder_inputs
|
||||
```
|
||||
|
||||
### Streaming mode (async_chunk)
|
||||
|
||||
Buffers Stage 0 outputs and forwards a chunk to Stage 1 once `chunk_size` frames
|
||||
have accumulated. The function signature follows the `OmniChunkTransferAdapter` protocol:
|
||||
|
||||
```python
|
||||
def ar2decoder_async_chunk(
|
||||
transfer_manager: Any,
|
||||
pooling_output: dict[str, Any] | None,
|
||||
request: Any,
|
||||
is_finished: bool = False,
|
||||
) -> dict[str, Any] | None:
|
||||
"""Forward chunks of AR output to the decoder stage."""
|
||||
request_id = request.external_req_id
|
||||
finished = bool(is_finished or request.is_finished())
|
||||
|
||||
# Extract and buffer the latest frame
|
||||
if isinstance(pooling_output, dict):
|
||||
frame = extract_frame(pooling_output)
|
||||
if frame is not None:
|
||||
transfer_manager.code_prompt_token_ids[request_id].append(
|
||||
frame.cpu().tolist()
|
||||
)
|
||||
elif not finished:
|
||||
return None
|
||||
|
||||
# Read chunk config from connector
|
||||
chunk_size = 25
|
||||
left_context_size = 25
|
||||
|
||||
length = len(transfer_manager.code_prompt_token_ids[request_id])
|
||||
if length <= 0:
|
||||
if finished:
|
||||
return {"codes": [], "finished": torch.tensor(True, dtype=torch.bool)}
|
||||
return None
|
||||
|
||||
# Wait until a full chunk is ready (or request is finished)
|
||||
chunk_length = length % chunk_size
|
||||
if chunk_length != 0 and not finished:
|
||||
return None
|
||||
|
||||
# Build context window: left_context + chunk
|
||||
context_length = chunk_length if chunk_length != 0 else chunk_size
|
||||
end_index = min(length, left_context_size + context_length)
|
||||
window = transfer_manager.code_prompt_token_ids[request_id][-end_index:]
|
||||
|
||||
return {
|
||||
"codes": torch.tensor(window).transpose(0, 1).reshape(-1).tolist(),
|
||||
"left_context_size": max(0, int(end_index - context_length)),
|
||||
"finished": torch.tensor(finished, dtype=torch.bool),
|
||||
}
|
||||
```
|
||||
|
||||
Key points:
|
||||
- `transfer_manager` is the `OmniChunkTransferAdapter` that owns the chunk lifecycle
|
||||
- Each call appends one AR decode step's output; a chunk is emitted every `chunk_size` steps
|
||||
- The final (possibly partial) chunk is flushed when `is_finished` is true
|
||||
- `left_context_size` frames of overlap are included for smooth audio boundaries
|
||||
|
||||
## Testing
|
||||
|
||||
For general testing conventions, see [tests_style.md](../ci/tests_style.md).
|
||||
|
||||
Recommended test cases for a new TTS model:
|
||||
|
||||
1. **Single request** - verify waveform output shape and sample rate
|
||||
2. **Batched requests** - verify each request in the batch finishes independently
|
||||
3. **async_chunk streaming** - verify audio chunks arrive incrementally and decode correctly
|
||||
4. **Speaker conditioning** (if applicable) - verify different speaker inputs produce different outputs
|
||||
|
||||
Reference test: `tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py`
|
||||
|
||||
## Adding a Model Recipe
|
||||
|
||||
After implementing and testing your model, add a model recipe to the
|
||||
[vllm-project/recipes](https://github.com/vllm-project/recipes) repository so users can
|
||||
get started quickly. See [Adding an Omni-Modality Model](./adding_omni_model.md#adding-a-model-recipe)
|
||||
for the expected format.
|
||||
|
||||
## Summary
|
||||
|
||||
Adding a TTS model to vLLM-Omni involves:
|
||||
|
||||
1. **Create model directory** with AR stage, decoder stage, and unified class
|
||||
2. **AR stage** - use vLLM's native decoder layers with fused QKV; do not wrap HF directly
|
||||
3. **Decoder stage** - thin wrapper around your audio decoder; implement `chunked_decode_streaming()`
|
||||
4. **Unified class** - dispatches on `model_stage`; same structure as `Qwen3TTSModelForGeneration`
|
||||
5. **Register** all stage classes in `registry.py`
|
||||
6. **YAML configs** - provide both batch and `async_chunk` variants
|
||||
7. **Stage input processor** - buffer Stage 0 outputs and forward in chunks of 25
|
||||
8. **Tests** - cover single request, batching, and async_chunk streaming
|
||||
9. **Model recipe** - add to [vllm-project/recipes](https://github.com/vllm-project/recipes)
|
||||
|
||||
### Qwen3-TTS Reference Files
|
||||
|
||||
| File | Purpose |
|
||||
|------|---------|
|
||||
| `models/qwen3_tts/qwen3_tts.py` | Unified model class |
|
||||
| `models/qwen3_tts/qwen3_tts_code_predictor_vllm.py` | AR stage with vLLM fused ops |
|
||||
| `models/qwen3_tts/qwen3_tts_code2wav.py` | Decoder stage with `chunked_decode_streaming()` |
|
||||
| `stage_configs/qwen3_tts.yaml` | Stage configuration |
|
||||
| `stage_input_processors/qwen3_tts.py` | Stage transition processors |
|
||||
|
||||
For more information, see:
|
||||
|
||||
- [Architecture Overview](../../design/architecture_overview.md)
|
||||
- [Async Chunk Design](../../design/feature/async_chunk_design.md)
|
||||
- [Stage Configuration Guide](../../configuration/stage_configs.md)
|
||||
@@ -0,0 +1,196 @@
|
||||
# RFC-001: Fully Automated Nightly Rebase via Multi-Agent System
|
||||
|
||||
| Field | Value |
|
||||
|-------------|--------------------------------------|
|
||||
| **Status** | Draft |
|
||||
| **Authors** | vLLM-Omni Team |
|
||||
| **Created** | 2026-03-01 |
|
||||
|
||||
## 1. Summary
|
||||
|
||||
This RFC proposes building a fully automated agent system that performs a nightly rebase of vllm-omni onto the latest upstream vLLM commit. The primary challenge is that the codebase is large enough to exceed any single LLM context window, necessitating a **multi-agent architecture** that decomposes the rebase workflow into context-bounded subtasks.
|
||||
|
||||
## 2. Motivation
|
||||
|
||||
### Current State
|
||||
|
||||
vllm-omni is an independent package that extends upstream vLLM (`vllm.v1.*`) with multimodal capabilities (text, image, video, audio, and diffusion models). Keeping vllm-omni in sync with upstream vLLM is critical because:
|
||||
|
||||
- **API drift**: vLLM frequently changes internal APIs (e.g., `GPUModelRunner`, `OutputProcessor`, `SchedulerOutput`). Each day of lag compounds merge conflict complexity.
|
||||
- **Manual burden**: Past rebase efforts (e.g., `75770c93 - Rebase to vllm v0.16.0`) required significant manual work, including fixing downstream references across workers, model runners, and platform adapters.
|
||||
|
||||
### Why Automation Is Difficult
|
||||
|
||||
| Challenge | Detail |
|
||||
|-----------|--------|
|
||||
| **Context window limits** | vllm-omni has 1,200+ lines in `gpu_model_runner.py`, 1,500+ in `omni_stage.py`, 800+ in `gpu_generation_model_runner.py`, and dozens of model-specific files. A single-agent approach cannot hold enough context to resolve all conflicts. |
|
||||
| **Semantic conflicts** | Many conflicts are not textual but semantic — upstream may rename a parameter, change a method signature, or restructure a class hierarchy. Blind merge tools fail silently. |
|
||||
| **Cross-file dependencies** | A change in `vllm.v1.worker.gpu_model_runner.GPUModelRunner` can cascade into `OmniGPUModelRunner`, `GPUARModelRunner`, `GPUGenerationModelRunner`, and all platform variants (NPU, XPU, ROCm). |
|
||||
| **Test validation** | After resolving conflicts, the result must be validated against CI (L1–L5 test levels, see `docs/contributing/ci/CI_5levels.md`). |
|
||||
|
||||
## 3. Proposed Design
|
||||
|
||||
### 3.1 Architecture Overview
|
||||
|
||||
```
|
||||
┌─────────────────────────────────────────────────────────────────────┐
|
||||
│ Orchestrator Agent │
|
||||
│ (Nightly cron trigger, git operations, conflict decomposition) │
|
||||
└───────┬──────────────┬──────────────┬──────────────┬────────────────┘
|
||||
│ │ │ │
|
||||
▼ ▼ ▼ ▼
|
||||
┌──────────────┐ ┌──────────────┐ ┌──────────────┐ ┌──────────────┐
|
||||
│ Conflict │ │ Conflict │ │ Conflict │ │ Validation │
|
||||
│ Resolver │ │ Resolver │ │ Resolver │ │ Agent │
|
||||
│ Agent #1 │ │ Agent #2 │ │ Agent #N │ │ │
|
||||
│ (worker/) │ │ (engine/) │ │ (platform/) │ │ (CI runner) │
|
||||
└──────────────┘ └──────────────┘ └──────────────┘ └──────────────┘
|
||||
```
|
||||
|
||||
### 3.2 Agent Roles
|
||||
|
||||
#### Orchestrator Agent
|
||||
|
||||
- **Trigger**: GitHub Actions cron schedule (e.g., `0 2 * * *` UTC)
|
||||
- **Responsibilities**:
|
||||
1. Fetch the latest upstream vLLM commit (`pip install vllm` or track a pinned ref)
|
||||
2. Attempt `git rebase` or dependency version bump on a fresh branch `auto-rebase/YYYY-MM-DD`
|
||||
3. If conflicts arise, decompose them by file/module into independent work units
|
||||
4. Dispatch work units to Conflict Resolver Agents
|
||||
5. Collect results, apply patches, and trigger Validation Agent
|
||||
6. If validation passes, open a PR; if not, open an issue with diagnostics
|
||||
|
||||
#### Conflict Resolver Agent (per-module)
|
||||
|
||||
- **Context boundary**: Each agent receives only the conflicting file(s) plus their direct upstream counterpart(s) and a module-level dependency graph
|
||||
- **Strategy**:
|
||||
1. Receive: `(omni_file, upstream_diff, conflict_markers, api_change_summary)`
|
||||
2. Analyze the upstream change (renamed method? new parameter? restructured class?)
|
||||
3. Apply the minimal adaptation to vllm-omni code
|
||||
4. Output: patch file + confidence score + human-review flags
|
||||
- **Module partitions** (based on current codebase structure):
|
||||
- `worker/` — `OmniGPUModelRunner`, `GPUARModelRunner`, `GPUGenerationModelRunner`
|
||||
- `engine/` — `OmniInputProcessor`, `MultimodalOutputProcessor`, `OmniEngineArgs`
|
||||
- `entrypoints/` — `omni_stage.py`, `omni_llm.py`, `async_omni_llm.py`
|
||||
- `model_executor/` — Model-specific code (Qwen3, MiMo, GLM, etc.)
|
||||
- `platforms/` — NPU, XPU, ROCm adapters
|
||||
- `diffusion/` — DiffusionModelRunner and related code
|
||||
- `core/sched/` — Scheduler adaptations
|
||||
|
||||
#### Validation Agent
|
||||
|
||||
- **Responsibilities**:
|
||||
1. Run lint checks (`ruff`, `pre-commit`)
|
||||
2. Run import-level smoke tests
|
||||
3. Trigger CI L1/L2 tests (unit tests, basic e2e)
|
||||
4. Report pass/fail with logs
|
||||
- **Escalation**: If validation fails, feed error context back to the relevant Conflict Resolver Agent for a retry (up to N attempts)
|
||||
|
||||
### 3.3 Context Management Strategy
|
||||
|
||||
The key innovation is **bounded context windows** per agent:
|
||||
|
||||
| Strategy | Description |
|
||||
|----------|-------------|
|
||||
| **File-level isolation** | Each Conflict Resolver Agent sees at most 2-3 files (~2-4K lines total) |
|
||||
| **API summary injection** | Instead of full upstream source, agents receive a structured summary of API changes (method signatures, parameter changes, class hierarchy diffs) |
|
||||
| **Dependency graph pruning** | Only relevant imports and call sites are included in context |
|
||||
| **Iterative refinement** | If a patch introduces cross-module issues, the Orchestrator re-dispatches with expanded context |
|
||||
|
||||
### 3.4 Workflow
|
||||
|
||||
```
|
||||
1. [Cron] Trigger nightly at 02:00 UTC
|
||||
2. [Orchestrator] git fetch upstream; attempt rebase/version bump
|
||||
3. [Orchestrator] If clean → skip to step 7
|
||||
4. [Orchestrator] Parse conflicts → group by module → create work units
|
||||
5. [Resolver Agents] Resolve conflicts in parallel (bounded context per agent)
|
||||
6. [Orchestrator] Collect patches → apply sequentially → check for cross-module issues
|
||||
7. [Validation Agent] Run lint + L1/L2 tests
|
||||
8. [Orchestrator] If pass → open PR "Auto-rebase to vLLM <commit>"
|
||||
9. [Orchestrator] If fail → retry up to 3 times with error feedback
|
||||
10. [Orchestrator] If still failing → open issue with diagnostics for human review
|
||||
```
|
||||
|
||||
### 3.5 GitHub Actions Integration
|
||||
|
||||
```yaml
|
||||
name: Nightly Auto-Rebase
|
||||
on:
|
||||
schedule:
|
||||
- cron: '0 2 * * *'
|
||||
workflow_dispatch:
|
||||
|
||||
jobs:
|
||||
auto-rebase:
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
with:
|
||||
fetch-depth: 0
|
||||
- name: Run Orchestrator Agent
|
||||
run: python scripts/auto_rebase/orchestrator.py
|
||||
env:
|
||||
AGENT_API_KEY: ${{ secrets.AGENT_API_KEY }}
|
||||
- name: Open PR or Issue
|
||||
run: python scripts/auto_rebase/finalize.py
|
||||
```
|
||||
|
||||
## 4. Key Technical Decisions
|
||||
|
||||
### 4.1 Rebase vs. Dependency Bump
|
||||
|
||||
vllm-omni imports vLLM as a dependency rather than forking it. The "rebase" is actually:
|
||||
|
||||
1. Bumping the pinned vLLM version in `requirements/common.txt` or `setup.py`
|
||||
2. Adapting vllm-omni code to API changes in the new vLLM version
|
||||
3. Updating Dockerfiles (e.g., `Dockerfile.ci` base image `vllm/vllm-openai:v0.16.0`, `Dockerfile.rocm` notes)
|
||||
|
||||
### 4.2 Agent Backend
|
||||
|
||||
- Recommend using a code-specialized LLM with 128K+ context window
|
||||
- Each Conflict Resolver Agent should receive structured prompts with:
|
||||
- The specific vllm-omni file(s) being adapted
|
||||
- A diff of the upstream API change
|
||||
- The inheritance/import chain relevant to the file
|
||||
- Test files that cover the changed code
|
||||
|
||||
### 4.3 Rollback Safety
|
||||
|
||||
- All changes happen on a dedicated branch (`auto-rebase/YYYY-MM-DD`)
|
||||
- The PR requires at least one human approval before merge
|
||||
- Failed rebases are tracked in issues with full diagnostic context
|
||||
|
||||
## 5. Risks and Mitigations
|
||||
|
||||
| Risk | Mitigation |
|
||||
|------|------------|
|
||||
| Agent produces semantically wrong patches | Confidence scoring + human review for low-confidence patches |
|
||||
| Cross-module inconsistencies | Orchestrator performs a global consistency check after collecting all patches |
|
||||
| Upstream breaking changes too large | Fall back to opening a detailed issue rather than a broken PR |
|
||||
| API key / cost management | Rate-limit agent calls; cache upstream API summaries across runs |
|
||||
| Flaky CI tests | Distinguish test failures caused by rebase vs. pre-existing flakiness |
|
||||
|
||||
## 6. Success Metrics
|
||||
|
||||
- **Automation rate**: % of nightly rebases that produce a clean PR without human intervention
|
||||
- **Time to rebase**: Wall-clock time from cron trigger to PR creation (target: < 30 minutes)
|
||||
- **Conflict resolution accuracy**: % of agent-resolved conflicts that pass human review
|
||||
- **CI pass rate**: % of auto-rebase PRs that pass L1/L2 CI on first attempt
|
||||
|
||||
## 7. Implementation Plan
|
||||
|
||||
| Phase | Scope | Timeline |
|
||||
|-------|-------|----------|
|
||||
| Phase 1 | Orchestrator skeleton: git operations, conflict detection, branch management | 2 weeks |
|
||||
| Phase 2 | Conflict Resolver Agents: per-module resolution with bounded context | 3 weeks |
|
||||
| Phase 3 | Validation Agent: lint + smoke test integration | 1 week |
|
||||
| Phase 4 | GitHub Actions integration + PR/issue creation | 1 week |
|
||||
| Phase 5 | Monitoring, confidence tuning, feedback loops | Ongoing |
|
||||
|
||||
## 8. Open Questions
|
||||
|
||||
1. Should the system attempt to rebase daily or only when upstream has significant changes?
|
||||
2. What is the maximum acceptable cost per rebase run?
|
||||
3. Should platform-specific adapters (NPU, XPU, ROCm) be included in the nightly rebase or handled separately?
|
||||
4. How should we handle upstream changes that require new test coverage in vllm-omni?
|
||||
@@ -0,0 +1,465 @@
|
||||
## 1. Summary
|
||||
|
||||
This RFC proposes a comprehensive refactoring of the output processing pipeline in vllm-omni. The current `MultimodalOutputProcessor` and `OmniRequestState` classes have grown organically to handle multiple output modalities (text, image, audio, latents) but lack a clear, extensible architecture. A significant portion of the current code is **dead or unreachable**. We propose first simplifying the existing processor by removing dead code, then introducing a **Modality Registry**, a **Router abstraction**, and **composable output handlers** to support emerging multi-output models (e.g., DeepSeek with text+video+image).
|
||||
|
||||
## 2. Motivation
|
||||
|
||||
### Current Architecture
|
||||
|
||||
The output processing pipeline is concentrated in `vllm_omni/engine/output_processor.py` (469 lines) with two main classes:
|
||||
|
||||
- **`MultimodalOutputProcessor`** (extends `vllm.v1.engine.output_processor.OutputProcessor`): Routes `EngineCoreOutput` objects by `output_type` string to modality-specific handler methods.
|
||||
- **`OmniRequestState`** (extends `RequestState`): Accumulates multimodal tensors per-request during generation, with deferred concatenation.
|
||||
|
||||
### Problems
|
||||
|
||||
#### P0: Significant Dead / Unreachable Code
|
||||
|
||||
A thorough audit of the output processor reveals that **a large portion of the code is never executed** in any real production path. This dead code adds confusion, increases maintenance burden, and obscures the actual data flow. See Section 3 for the full dead code inventory.
|
||||
|
||||
#### P1: Rigid String-Based Routing
|
||||
|
||||
The `_route_and_normalize` method uses hardcoded string matching:
|
||||
|
||||
```python
|
||||
if output_type == "image":
|
||||
self._process_image_output(eco)
|
||||
elif output_type in ("text+image", "text,image", "image+text"):
|
||||
self._process_text_image_output(eco)
|
||||
elif output_type in ("latents", "latent"):
|
||||
self._process_latents_output(eco)
|
||||
elif output_type in ("audio", "speech"):
|
||||
self._process_audio_output(eco)
|
||||
elif output_type == "text":
|
||||
self._process_text_output(eco)
|
||||
```
|
||||
|
||||
Adding a new modality (e.g., `video`, `text+video`, `text+video+image`) requires modifying this central dispatch chain.
|
||||
|
||||
#### P2: No Composable Multi-Output Support
|
||||
|
||||
Models like DeepSeek may produce **text + video + image** simultaneously. The current architecture handles at most two modalities (e.g., `text+image`) via a single combined handler. There is no composition mechanism for arbitrary output type combinations.
|
||||
|
||||
#### P3: Inconsistent Key Extraction
|
||||
|
||||
Each `_process_*_output` method has its own hardcoded key lists:
|
||||
|
||||
- Image: `("image", "images", "pixel_values", "pixels")`
|
||||
- Audio: `("audio", "audios", "wav", "waveform", "audio_pcm", "pcm")`
|
||||
- Latents: `("latent", "latents", "z", "posterior")`
|
||||
|
||||
There is no standard contract between model output format and the output processor.
|
||||
|
||||
#### P4: Monolithic Accumulation Logic
|
||||
|
||||
`OmniRequestState.add_multimodal_tensor` handles all modalities in a single method with complex branching for dict/tensor/list inputs and nested dict merging (~65 lines). The special-case for audio tensor concatenation in `_consolidate_multimodal_tensors` (skipping `torch.cat` due to shape inconsistency) further indicates the need for per-modality accumulation strategies.
|
||||
|
||||
#### P5: Missing Modalities
|
||||
|
||||
- **Video output**: No dedicated `"video"` routing branch exists. Video generation would currently fall through to the heuristic fallback.
|
||||
- **Text+Video, Text+Video+Image**: No handling exists for these combinations, which future models (e.g., DeepSeek multimodal) will require.
|
||||
|
||||
## 3. Dead Code Audit and Simplification Plan
|
||||
|
||||
A line-by-line analysis of `output_processor.py` cross-referenced with all callers and stage configs reveals the following dead/unused code:
|
||||
|
||||
### 3.1 Dead Code Inventory
|
||||
|
||||
#### D1: `register_handler()` — Never Called (lines 281-292)
|
||||
|
||||
The `register_handler(modality, handler)` method is defined on `MultimodalOutputProcessor` and initializes `self.output_handlers: dict[str, Callable]`. However, **no code in the entire codebase calls `register_handler()`**. The `output_handlers` dict is always empty, making the custom handler dispatch path in `_route_and_normalize` (lines 373-379) unreachable.
|
||||
|
||||
**Evidence**: `rg "register_handler" vllm_omni/` only finds the definition and the docstring reference. Zero callers.
|
||||
|
||||
**Action**: Remove `register_handler()`, `self.output_handlers`, and the custom handler dispatch block.
|
||||
|
||||
#### D2: `_reqid_to_mm_type` — Written but Never Read (line 278, 346-350)
|
||||
|
||||
The `_reqid_to_mm_type: dict[str, str]` is populated in `process_outputs()`:
|
||||
|
||||
```python
|
||||
self._reqid_to_mm_type.clear()
|
||||
for eco in engine_core_outputs:
|
||||
mm_type = (self.engine_core_output_type or "").lower()
|
||||
if mm_type:
|
||||
self._reqid_to_mm_type[eco.request_id] = mm_type
|
||||
```
|
||||
|
||||
But this dict is **never read anywhere** — not in this class, not in any subclass, not in any caller. It is populated and then cleared on the next call.
|
||||
|
||||
**Action**: Remove `_reqid_to_mm_type` and the population code.
|
||||
|
||||
#### D3: `_process_text_image_output` — No Stage Config Uses "text+image" (lines 383-384, 406-424)
|
||||
|
||||
The routing branch `elif output_type in ("text+image", "text,image", "image+text")` and its handler `_process_text_image_output` are never triggered because:
|
||||
|
||||
- No stage config in the entire project uses `engine_output_type: text+image`, `text,image`, or `image+text`
|
||||
- `OmniEngineCoreOutput` does not have an `output_type` attribute, so `getattr(eco, "output_type", ...)` always falls back to `engine_core_output_type`
|
||||
- No code path ever sets `engine_core_output_type` to any of these composite strings
|
||||
|
||||
**Actual stage config values** (exhaustive search across all YAML files):
|
||||
- `latent` — Qwen3-TTS, Qwen3-Omni, Qwen2.5-Omni, MiMo-Audio
|
||||
- `audio` — Qwen3-TTS, Qwen3-Omni, Qwen2.5-Omni, MiMo-Audio
|
||||
- `text` — Bagel
|
||||
- `image` — Bagel
|
||||
- `token_ids` — GLM-Image
|
||||
|
||||
**Action**: Remove `_process_text_image_output` and its routing branch.
|
||||
|
||||
#### D4: `"speech"` Alias — Never Used (line 387)
|
||||
|
||||
The routing matches `("audio", "speech")`, but no stage config ever uses `engine_output_type: speech`. Only `"audio"` is used in practice.
|
||||
|
||||
**Action**: Remove `"speech"` from the match (minor, but reduces confusion).
|
||||
|
||||
#### D5: `_extract_from_multimodal_outputs` — Always Returns None (lines 457-469)
|
||||
|
||||
This method tries to read `eco.multimodal_outputs`:
|
||||
|
||||
```python
|
||||
mm = getattr(eco, "multimodal_outputs", None)
|
||||
```
|
||||
|
||||
However, `OmniEngineCoreOutput` (defined in `vllm_omni/engine/__init__.py:76-77`) only has:
|
||||
|
||||
```python
|
||||
class OmniEngineCoreOutput(EngineCoreOutput):
|
||||
pooling_output: dict[str, torch.Tensor] | None = None
|
||||
```
|
||||
|
||||
There is **no `multimodal_outputs` attribute** on `EngineCoreOutput` or `OmniEngineCoreOutput`. The `getattr` always returns `None`, making this method always return `None`.
|
||||
|
||||
This means the "extract from multimodal_outputs" fallback inside `_process_image_output`, `_process_latents_output`, `_process_audio_output`, and `_process_text_image_output` is **dead code**. These methods only execute when `eco.pooling_output is None`, and in that case they try `_extract_from_multimodal_outputs` which always returns `None` — so they effectively do nothing.
|
||||
|
||||
**Action**: Remove `_extract_from_multimodal_outputs` and the extraction fallback paths in all `_process_*_output` methods.
|
||||
|
||||
#### D6: `_process_pooling_output` — Broken for Actual Input Type (lines 446-455)
|
||||
|
||||
This method is hit as a fallback when `output_type` is unrecognized (e.g., `"token_ids"` for GLM-Image). It tries:
|
||||
|
||||
```python
|
||||
if not isinstance(eco.pooling_output, torch.Tensor):
|
||||
eco.pooling_output = torch.as_tensor(eco.pooling_output)
|
||||
```
|
||||
|
||||
But `OmniEngineCoreOutput.pooling_output` is typed as `dict[str, torch.Tensor] | None`. Calling `torch.as_tensor()` on a dict will raise a `TypeError`. This method is either dead or broken.
|
||||
|
||||
**Action**: Remove or fix `_process_pooling_output`. For `"token_ids"`, the data already flows correctly through `pooling_output` → `add_multimodal_tensor` in `process_outputs`.
|
||||
|
||||
#### D7: `_process_text_output` — No-Op (lines 442-444)
|
||||
|
||||
```python
|
||||
def _process_text_output(self, eco: EngineCoreOutput) -> None:
|
||||
"""No-op; base processor will detokenize new_token_ids → text."""
|
||||
return
|
||||
```
|
||||
|
||||
This is a no-op method. It adds no value and can be replaced by a direct `pass` or `return` in the routing.
|
||||
|
||||
**Action**: Remove `_process_text_output` and replace the routing branch with a simple return/pass.
|
||||
|
||||
#### D8: Key Remapping Logic — Obscure and Partially Redundant (lines 64-69)
|
||||
|
||||
In `add_multimodal_tensor`, the code remaps keys `"model_outputs"` and `"hidden"` to `target_key` (which is `mm_type`):
|
||||
|
||||
```python
|
||||
if k == "model_outputs":
|
||||
k = target_key
|
||||
elif k == "hidden" and target_key != "hidden":
|
||||
k = target_key
|
||||
```
|
||||
|
||||
While this renaming does execute (AR runners produce `{"hidden": ...}` and generation runners produce `{"model_outputs": ...}`), downstream consumers like `serving_speech.py` still check for the **original key** `"model_outputs"`:
|
||||
|
||||
```python
|
||||
key = "audio" if "audio" in mm else ("model_outputs" if "model_outputs" in mm else None)
|
||||
```
|
||||
|
||||
This means the renaming creates an inconsistency: the output processor renames `"model_outputs"` → `"audio"`, but the serving layer expects `"model_outputs"`. The system works only because the serving layer has a fallback, but the renaming adds confusion.
|
||||
|
||||
**Action**: Audit all downstream consumers and either commit to the renaming (update consumers) or remove it (keep original keys).
|
||||
|
||||
### 3.2 Effectively Unused Code Summary
|
||||
|
||||
| Code | Lines | Status | Reason |
|
||||
|------|-------|--------|--------|
|
||||
| `register_handler()` + `output_handlers` dict | 281-292, 277, 373-379 | **Dead** | Zero callers in entire codebase |
|
||||
| `_reqid_to_mm_type` dict | 278, 346-350 | **Dead** | Populated but never read |
|
||||
| `_process_text_image_output` | 383-384, 406-424 | **Dead** | No stage config uses "text+image"/"text,image"/"image+text" |
|
||||
| `"speech"` alias in routing | 387 | **Dead** | No config uses "speech", only "audio" |
|
||||
| `_extract_from_multimodal_outputs` | 457-469 | **Dead** | `eco.multimodal_outputs` attribute doesn't exist on `OmniEngineCoreOutput` |
|
||||
| Extraction fallbacks in `_process_*_output` | 399-440 (partial) | **Dead** | All call `_extract_from_multimodal_outputs` which always returns None |
|
||||
| `_process_pooling_output` | 446-455 | **Broken** | `torch.as_tensor()` on a dict raises TypeError |
|
||||
| `_process_text_output` | 442-444 | **No-op** | Empty method body |
|
||||
| Key remapping (`"model_outputs"` → mm_type) | 64-69 | **Inconsistent** | Renames keys that consumers still expect under original names |
|
||||
|
||||
**Estimated dead/broken code**: ~90 lines out of 469 (19% of the file).
|
||||
|
||||
### 3.3 What Actually Executes
|
||||
|
||||
After removing dead code, the **actual live data flow** is remarkably simple:
|
||||
|
||||
```
|
||||
1. process_outputs() is called with a list of EngineCoreOutput
|
||||
2. For each eco:
|
||||
a. engine_core_output_type determines mm_type (from stage config: "latent", "audio", "text", "image", "token_ids")
|
||||
b. _route_and_normalize() is called but effectively does nothing useful
|
||||
(all _process_*_output methods only act when pooling_output is None,
|
||||
then try _extract_from_multimodal_outputs which always returns None)
|
||||
c. If eco.pooling_output is not None AND req_state has a detokenizer:
|
||||
- pooling_output (a dict) is passed to req_state.add_multimodal_tensor()
|
||||
- eco.pooling_output is set to None (forcing text path in base processor)
|
||||
3. Base vLLM OutputProcessor.process_outputs() handles text detokenization
|
||||
4. On finish, _consolidate_multimodal_tensors() concatenates accumulated tensors
|
||||
5. _new_completion_output() attaches mm_accumulated to CompletionOutput
|
||||
```
|
||||
|
||||
The entire routing layer (`_route_and_normalize` and all `_process_*_output` methods) is effectively **a no-op** in production — the real work happens in the `pooling_output` capture at lines 355-361 and the accumulation in `OmniRequestState`.
|
||||
|
||||
## 4. Proposed Design
|
||||
|
||||
### 4.1 Core Abstractions
|
||||
|
||||
#### ModalityHandler (Protocol)
|
||||
|
||||
```python
|
||||
class ModalityHandler(Protocol):
|
||||
"""Handles extraction, normalization, and accumulation for a single output modality."""
|
||||
|
||||
modality: str # e.g., "image", "audio", "video", "latent"
|
||||
|
||||
def extract(self, eco: EngineCoreOutput) -> torch.Tensor | dict | None:
|
||||
"""Extract this modality's data from an EngineCoreOutput."""
|
||||
...
|
||||
|
||||
def accumulate(self, existing: Any, incoming: Any) -> Any:
|
||||
"""Accumulate incremental outputs (e.g., streaming audio chunks)."""
|
||||
...
|
||||
|
||||
def consolidate(self, accumulated: Any) -> Any:
|
||||
"""Finalize accumulated data (e.g., torch.cat for tensors)."""
|
||||
...
|
||||
|
||||
def validate(self, output: Any) -> bool:
|
||||
"""Optional validation of the final output."""
|
||||
...
|
||||
```
|
||||
|
||||
#### ModalityRegistry
|
||||
|
||||
```python
|
||||
class ModalityRegistry:
|
||||
"""Central registry mapping modality names to their handlers."""
|
||||
|
||||
def register(self, handler: ModalityHandler) -> None: ...
|
||||
def get(self, modality: str) -> ModalityHandler | None: ...
|
||||
def list_modalities(self) -> list[str]: ...
|
||||
```
|
||||
|
||||
#### OutputRouter
|
||||
|
||||
```python
|
||||
class OutputRouter:
|
||||
"""Routes EngineCoreOutput to one or more ModalityHandlers."""
|
||||
|
||||
def __init__(self, registry: ModalityRegistry): ...
|
||||
|
||||
def route(self, eco: EngineCoreOutput, output_type: str) -> list[tuple[str, ModalityHandler]]:
|
||||
"""Parse output_type and return list of (modality, handler) pairs.
|
||||
|
||||
Supports composite types like "text+video+image" by splitting
|
||||
on '+' or ',' and looking up each component in the registry.
|
||||
"""
|
||||
...
|
||||
```
|
||||
|
||||
### 4.2 Architecture
|
||||
|
||||
```
|
||||
┌────────────────────────────────────────────────────────────────────────┐
|
||||
│ MultimodalOutputProcessor (refactored) │
|
||||
│ │
|
||||
│ ┌──────────────┐ ┌──────────────┐ ┌───────────────────────┐ │
|
||||
│ │ OutputRouter │────▶│ ModalityReg │────▶│ ModalityHandler[] │ │
|
||||
│ │ │ │ │ │ - TextHandler │ │
|
||||
│ │ "text+video" │ │ "text" → H1 │ │ - ImageHandler │ │
|
||||
│ │ ↓ │ │ "video" → H2 │ │ - VideoHandler │ │
|
||||
│ │ ["text", │ │ "image" → H3 │ │ - AudioHandler │ │
|
||||
│ │ "video"] │ │ "audio" → H4 │ │ - LatentHandler │ │
|
||||
│ └──────────────┘ │ "video" → H5 │ │ - (custom handlers) │ │
|
||||
│ │ "latent"→ H6 │ └───────────────────────┘ │
|
||||
│ └──────────────┘ │
|
||||
│ │
|
||||
│ ┌─────────────────────────────────────────────────────────────────┐ │
|
||||
│ │ OmniRequestState (refactored) │ │
|
||||
│ │ mm_accumulators: dict[str, ModalityAccumulator] │ │
|
||||
│ │ - Each accumulator uses its handler's accumulate/consolidate │ │
|
||||
│ └─────────────────────────────────────────────────────────────────┘ │
|
||||
└────────────────────────────────────────────────────────────────────────┘
|
||||
```
|
||||
|
||||
### 4.3 Built-in Handlers
|
||||
|
||||
| Handler | Modality | Extract Keys | Accumulation | Consolidation |
|
||||
|---------|----------|-------------|--------------|---------------|
|
||||
| `TextHandler` | `text` | (none — handled by base vLLM detokenizer) | N/A | N/A |
|
||||
| `ImageHandler` | `image` | `image`, `images`, `pixel_values`, `pixels` | List append | `torch.cat(dim=0)` |
|
||||
| `VideoHandler` | `video` | `video`, `videos`, `frames`, `video_frames` | List append | `torch.cat(dim=0)` or `torch.stack` |
|
||||
| `AudioHandler` | `audio` | `audio`, `audios`, `wav`, `waveform`, `audio_pcm`, `pcm` | List append | `torch.cat(dim=-1)` (last dim for variable-length) |
|
||||
| `LatentHandler` | `latent` | `latent`, `latents`, `z`, `posterior` | List append | `torch.cat(dim=0)` |
|
||||
|
||||
### 4.4 Composite Output Types
|
||||
|
||||
The `OutputRouter` supports composite types by splitting on delimiters:
|
||||
|
||||
```python
|
||||
def route(self, eco, output_type):
|
||||
modalities = re.split(r'[+,]', output_type)
|
||||
handlers = []
|
||||
for m in modalities:
|
||||
m = m.strip().lower()
|
||||
handler = self.registry.get(m)
|
||||
if handler:
|
||||
handlers.append((m, handler))
|
||||
else:
|
||||
logger.warning("No handler for modality: %s", m)
|
||||
return handlers
|
||||
```
|
||||
|
||||
This naturally supports:
|
||||
- `"text+video+image"` → `[TextHandler, VideoHandler, ImageHandler]`
|
||||
- `"text+audio"` → `[TextHandler, AudioHandler]`
|
||||
- `"image"` → `[ImageHandler]`
|
||||
|
||||
### 4.5 Refactored OmniRequestState
|
||||
|
||||
```python
|
||||
class OmniRequestState(RequestState):
|
||||
def __init__(self, *args, registry: ModalityRegistry, **kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
self._registry = registry
|
||||
self._accumulators: dict[str, Any] = {}
|
||||
|
||||
def add_output(self, modality: str, payload: Any) -> None:
|
||||
handler = self._registry.get(modality)
|
||||
if handler is None:
|
||||
return
|
||||
existing = self._accumulators.get(modality)
|
||||
self._accumulators[modality] = handler.accumulate(existing, payload)
|
||||
|
||||
def consolidate(self) -> dict[str, Any]:
|
||||
result = {}
|
||||
for modality, data in self._accumulators.items():
|
||||
handler = self._registry.get(modality)
|
||||
if handler:
|
||||
result[modality] = handler.consolidate(data)
|
||||
else:
|
||||
result[modality] = data
|
||||
return result
|
||||
```
|
||||
|
||||
### 4.6 Model Output Contract
|
||||
|
||||
To standardize the interface between models and the output processor, we propose that models producing multimodal outputs should return an `OmniOutput` (already defined in `model_executor/models/output_templates.py`) with a well-defined `multimodal_outputs` dict:
|
||||
|
||||
```python
|
||||
@dataclass
|
||||
class OmniOutput:
|
||||
text_hidden_states: torch.Tensor
|
||||
multimodal_outputs: dict[str, torch.Tensor | dict]
|
||||
# Keys should match registered modality names: "image", "video", "audio", etc.
|
||||
```
|
||||
|
||||
## 5. Migration Plan
|
||||
|
||||
### Phase 0: Simplify — Remove Dead Code and Clarify Live Paths
|
||||
|
||||
This phase is a prerequisite for any further refactoring. It reduces the file from 469 lines to ~300 and makes the actual data flow visible.
|
||||
|
||||
**Step 0.1**: Remove dead code (all items from Section 3.2):
|
||||
|
||||
- Delete `register_handler()`, `self.output_handlers`, and the custom handler dispatch block in `_route_and_normalize`
|
||||
- Delete `self._reqid_to_mm_type` and all code that populates it
|
||||
- Delete `_process_text_image_output` and its routing branch
|
||||
- Delete `_extract_from_multimodal_outputs` and all extraction fallback paths
|
||||
- Delete `_process_pooling_output`
|
||||
- Delete `_process_text_output`
|
||||
- Remove `"speech"` from the audio routing match
|
||||
|
||||
**Step 0.2**: Simplify `_route_and_normalize` to reflect what actually executes:
|
||||
|
||||
```python
|
||||
def _route_and_normalize(self, eco: EngineCoreOutput) -> None:
|
||||
# Currently a no-op for all production paths.
|
||||
# Routing is actually handled by the pooling_output capture in process_outputs().
|
||||
# Kept as an extension point for future modality-specific normalization.
|
||||
pass
|
||||
```
|
||||
|
||||
Or, if we want to keep it as a dispatch point for future use, collapse into:
|
||||
|
||||
```python
|
||||
def _route_and_normalize(self, eco: EngineCoreOutput) -> None:
|
||||
output_type = (self.engine_core_output_type or "").lower()
|
||||
if not output_type or output_type in ("text", "latent", "audio", "image", "token_ids"):
|
||||
return # All handled by pooling_output capture in process_outputs()
|
||||
```
|
||||
|
||||
**Step 0.3**: Audit and resolve the key remapping inconsistency (D8):
|
||||
|
||||
- Option A: Remove renaming, keep original keys (`"model_outputs"`, `"hidden"`), update the proposed ModalityHandler to standardize keys at extraction time
|
||||
- Option B: Commit to renaming, update `serving_speech.py` and other consumers to use mm_type keys
|
||||
|
||||
**Step 0.4**: Simplify `make_request_output` — the current override largely duplicates base class logic for streaming/DELTA handling. If the only additions are `_consolidate_multimodal_tensors()` and `_new_completion_output()`, consider hooking these without copying the full method.
|
||||
|
||||
**Estimated result**: ~300 lines (from 469), with clear separation between "what the code does" and "extension points for the future".
|
||||
|
||||
### Phase 1: Introduce Abstractions (Non-Breaking)
|
||||
|
||||
1. Create `ModalityHandler` protocol and `ModalityRegistry` class
|
||||
2. Implement built-in handlers (`TextHandler`, `ImageHandler`, `AudioHandler`, `LatentHandler`)
|
||||
3. Create `OutputRouter` with delimiter-based composite type parsing
|
||||
4. Place these in a new module: `vllm_omni/engine/output_handlers.py`
|
||||
|
||||
### Phase 2: Refactor MultimodalOutputProcessor
|
||||
|
||||
1. Replace the simplified `_route_and_normalize` with `OutputRouter.route()`
|
||||
2. Introduce handler-based dispatch for `pooling_output` extraction and normalization
|
||||
3. The existing `register_handler()` API concept is naturally replaced by `ModalityRegistry.register()`
|
||||
|
||||
### Phase 3: Refactor OmniRequestState
|
||||
|
||||
1. Replace `mm_accumulated` dict with per-modality accumulators
|
||||
2. Replace `add_multimodal_tensor` with `add_output(modality, payload)`
|
||||
3. Replace `_consolidate_multimodal_tensors` with handler-driven `consolidate()`
|
||||
|
||||
### Phase 4: Add New Modalities
|
||||
|
||||
1. Add `VideoHandler` for video output support
|
||||
2. Test composite types: `text+video`, `text+image+video`, etc.
|
||||
3. Integrate with upcoming DeepSeek model support
|
||||
|
||||
## 6. Backward Compatibility
|
||||
|
||||
- `register_handler()` was never called externally — removing it has zero backward compatibility impact
|
||||
- Existing stage configs with `engine_output_type` strings (e.g., `"audio"`, `"latent"`, `"image"`) continue to work unchanged
|
||||
- The `output_type` field on `EngineCoreOutput` is unchanged
|
||||
- Phase 0 (dead code removal) is a pure deletion with no behavioral change, since none of the removed paths were reachable
|
||||
|
||||
## 7. Testing Strategy
|
||||
|
||||
| Test Type | Phase | Scope |
|
||||
|-----------|-------|-------|
|
||||
| Regression (existing) | Phase 0 | Run `test_outputs.py` and all E2E tests after dead code removal — must pass unchanged |
|
||||
| Smoke test | Phase 0 | Verify Qwen3-TTS (audio), GLM-Image (token_ids), Bagel (text+image) pipelines still work |
|
||||
| Unit tests | Phase 1 | Each `ModalityHandler` (extract, accumulate, consolidate) |
|
||||
| Integration tests | Phase 2 | `OutputRouter` with composite types |
|
||||
| E2E tests | Phase 2-3 | Full pipeline with existing models (Qwen3-TTS, GLM-Image, etc.) |
|
||||
| New modality tests | Phase 4 | Video output handling, `text+video+image` combinations |
|
||||
|
||||
## 8. Open Questions
|
||||
|
||||
1. Should we support dynamic handler priority (e.g., if both `image` and `latent` are present, which takes precedence for `pooling_output`)?
|
||||
2. Should the modality registry be global (singleton) or per-engine-instance?
|
||||
3. How should error handling work when one handler in a composite route fails but others succeed?
|
||||
4. Should we support async accumulation for streaming scenarios?
|
||||
@@ -0,0 +1,280 @@
|
||||
# RFC-003: Model Runner Clearance — Extracting Model-Specific Logic
|
||||
|
||||
| Field | Value |
|
||||
|-------------|--------------------------------------|
|
||||
| **Status** | Draft |
|
||||
| **Authors** | vLLM-Omni Team |
|
||||
| **Created** | 2026-03-01 |
|
||||
|
||||
## 1. Summary
|
||||
|
||||
This RFC proposes extracting model-specific logic currently embedded in the model runner classes (`OmniGPUModelRunner`, `GPUARModelRunner`, `GPUGenerationModelRunner`) into well-defined extension points and per-model adapter classes. The goal is to make the model runner a **model-agnostic execution framework** that delegates model-specific behavior through a clean plugin/adapter interface.
|
||||
|
||||
## 2. Motivation
|
||||
|
||||
### Current State
|
||||
|
||||
`OmniGPUModelRunner` (1,292 lines in `vllm_omni/worker/gpu_model_runner.py`) is the core execution class that inherits from `vllm.v1.worker.gpu_model_runner.GPUModelRunner`. Over time, it has accumulated significant model-specific logic that should not live in a generic runner:
|
||||
|
||||
#### M-RoPE / Positional Encoding (Qwen2-VL, HunYuan-VL, GLM-Image)
|
||||
|
||||
- `_init_mrope_positions()` (lines 99-157): Extracts `image_grid_thw`, `video_grid_thw`, `second_per_grid_ts`, `audio_feature_lengths`, `use_audio_in_video` — all Qwen-specific metadata.
|
||||
- `_calc_mrope_positions()` (line 159): Delegates to upstream, then calls `_fixup_precomputed_mrope_decode_positions` for GLM-Image.
|
||||
- `_fixup_precomputed_mrope_decode_positions()` (lines 180-222): Entirely GLM-Image-specific — handles 2D spatial decode positions for image generation models with `precomputed_mrope_decode = True`.
|
||||
- XD-RoPE handling via `uses_xdrope_dim` for HunYuan-VL.
|
||||
|
||||
#### Talker MTP (Qwen3-TTS, Qwen3-Omni)
|
||||
|
||||
- `load_model()` (lines 71-97): Contains TTS-specific buffer allocation (`talker_mtp_input_ids`, `talker_mtp_inputs_embeds`, `last_talker_hidden`, `text_step`) with a `TODO` comment: "move this model specific logic to a separate class".
|
||||
- `_talker_mtp_forward()` (lines 1215-1246): Entirely Qwen3-TTS/Omni-specific — runs the talker MTP forward pass with CUDA graph wrapping.
|
||||
- `_preprocess()` (lines 1039-1213): Contains ~50 lines of talker MTP path (lines 1154-1204) interleaved with generic preprocessing.
|
||||
|
||||
#### MiMo-Audio
|
||||
|
||||
- `_maybe_attach_mimo_audio_req_infos()` (lines 1016-1037): Checks `self.model.__class__.__name__ == "MiMoAudioForConditionalGeneration"` — a direct class name check that couples the runner to a specific model.
|
||||
|
||||
#### Additional Information / Postprocess
|
||||
|
||||
- `_process_additional_information_updates()` (lines 937-969): Checks `self.model.has_postprocess` and calls `self.model.postprocess()` — a generic hook that's only used by specific models.
|
||||
- `_preprocess()` similarly checks `self.model.has_preprocess` and calls `self.model.preprocess()`.
|
||||
|
||||
### Problems
|
||||
|
||||
| Problem | Impact |
|
||||
|---------|--------|
|
||||
| **Tight coupling** | Adding a new model with unique positional encoding or pre/post-processing requires modifying the runner itself |
|
||||
| **Code readability** | The 1,292-line runner mixes framework logic with model-specific branches, making it hard to understand the core execution flow |
|
||||
| **Testing difficulty** | Model-specific logic cannot be unit-tested in isolation from the full runner |
|
||||
| **Rebase fragility** | Model-specific code in the runner is a frequent source of conflicts during upstream vLLM rebases |
|
||||
| **Platform duplication** | NPU, XPU, and ROCm runners must replicate or work around model-specific logic from the GPU runner |
|
||||
|
||||
## 3. Proposed Design
|
||||
|
||||
### 3.1 Model Runner Adapter Protocol
|
||||
|
||||
```python
|
||||
class ModelRunnerAdapter(Protocol):
|
||||
"""Extension point for model-specific behavior in the runner."""
|
||||
|
||||
def on_model_loaded(self, runner: "OmniGPUModelRunner", model: nn.Module) -> None:
|
||||
"""Called after model is loaded. Allocate model-specific buffers."""
|
||||
...
|
||||
|
||||
def init_request_positions(
|
||||
self, runner: "OmniGPUModelRunner", req_state: CachedRequestState
|
||||
) -> None:
|
||||
"""Initialize custom positional encoding for a new request (e.g., M-RoPE, XD-RoPE)."""
|
||||
...
|
||||
|
||||
def fixup_decode_positions(
|
||||
self, runner: "OmniGPUModelRunner", scheduler_output: SchedulerOutput
|
||||
) -> None:
|
||||
"""Apply model-specific position fixups for decode (e.g., GLM-Image 2D positions)."""
|
||||
...
|
||||
|
||||
def preprocess(
|
||||
self,
|
||||
runner: "OmniGPUModelRunner",
|
||||
input_ids: torch.Tensor,
|
||||
inputs_embeds: torch.Tensor | None,
|
||||
scheduler_output: SchedulerOutput,
|
||||
num_scheduled_tokens_np: np.ndarray,
|
||||
) -> tuple[torch.Tensor, torch.Tensor | None, dict]:
|
||||
"""Model-specific preprocessing before forward pass."""
|
||||
...
|
||||
|
||||
def postprocess(
|
||||
self,
|
||||
runner: "OmniGPUModelRunner",
|
||||
hidden_states: torch.Tensor,
|
||||
multimodal_outputs: Any,
|
||||
scheduler_output: SchedulerOutput,
|
||||
) -> None:
|
||||
"""Model-specific postprocessing after forward pass."""
|
||||
...
|
||||
|
||||
def extract_outputs(
|
||||
self,
|
||||
runner: "OmniGPUModelRunner",
|
||||
hidden_states: torch.Tensor | OmniOutput,
|
||||
) -> tuple[torch.Tensor, dict]:
|
||||
"""Extract text hidden states and multimodal outputs."""
|
||||
...
|
||||
|
||||
def enrich_request_info(
|
||||
self,
|
||||
runner: "OmniGPUModelRunner",
|
||||
req_id: str,
|
||||
req_state: CachedRequestState,
|
||||
req_infos: dict | None,
|
||||
) -> dict | None:
|
||||
"""Enrich per-request info before forward (e.g., MiMo-Audio mm_features)."""
|
||||
...
|
||||
```
|
||||
|
||||
### 3.2 Adapter Registry
|
||||
|
||||
```python
|
||||
class ModelAdapterRegistry:
|
||||
"""Maps model architecture names to their runner adapters."""
|
||||
|
||||
_adapters: dict[str, type[ModelRunnerAdapter]] = {}
|
||||
|
||||
@classmethod
|
||||
def register(cls, model_cls_name: str):
|
||||
def decorator(adapter_cls):
|
||||
cls._adapters[model_cls_name] = adapter_cls
|
||||
return adapter_cls
|
||||
return decorator
|
||||
|
||||
@classmethod
|
||||
def get_adapter(cls, model: nn.Module) -> ModelRunnerAdapter:
|
||||
name = model.__class__.__name__
|
||||
adapter_cls = cls._adapters.get(name, DefaultModelRunnerAdapter)
|
||||
return adapter_cls()
|
||||
```
|
||||
|
||||
### 3.3 Concrete Adapters
|
||||
|
||||
#### Default Adapter
|
||||
|
||||
```python
|
||||
class DefaultModelRunnerAdapter:
|
||||
"""No-op adapter for models without special requirements."""
|
||||
def on_model_loaded(self, runner, model): pass
|
||||
def init_request_positions(self, runner, req_state): pass
|
||||
def fixup_decode_positions(self, runner, scheduler_output): pass
|
||||
def preprocess(self, runner, input_ids, inputs_embeds, so, nst):
|
||||
return input_ids, inputs_embeds, {}
|
||||
def postprocess(self, runner, hs, mm, so): pass
|
||||
def extract_outputs(self, runner, hs):
|
||||
return (hs, {}) if isinstance(hs, torch.Tensor) else (hs[0], {})
|
||||
def enrich_request_info(self, runner, req_id, req_state, req_infos):
|
||||
return req_infos
|
||||
```
|
||||
|
||||
#### Qwen3-Omni / TTS Adapter
|
||||
|
||||
```python
|
||||
@ModelAdapterRegistry.register("Qwen3OmniMoeForConditionalGeneration")
|
||||
@ModelAdapterRegistry.register("Qwen3TTSForConditionalGeneration")
|
||||
class Qwen3TalkerAdapter(DefaultModelRunnerAdapter):
|
||||
def on_model_loaded(self, runner, model):
|
||||
# Allocate talker_mtp buffers, CUDAGraphWrapper, etc.
|
||||
...
|
||||
|
||||
def init_request_positions(self, runner, req_state):
|
||||
# M-RoPE with image_grid_thw, video_grid_thw, audio_feature_lengths
|
||||
...
|
||||
|
||||
def preprocess(self, runner, input_ids, inputs_embeds, so, nst):
|
||||
# Talker MTP forward for decode requests
|
||||
...
|
||||
```
|
||||
|
||||
#### GLM-Image Adapter
|
||||
|
||||
```python
|
||||
@ModelAdapterRegistry.register("GLMImageForConditionalGeneration")
|
||||
class GLMImageAdapter(DefaultModelRunnerAdapter):
|
||||
def fixup_decode_positions(self, runner, scheduler_output):
|
||||
# 2D spatial decode position fixup for precomputed_mrope_decode
|
||||
...
|
||||
```
|
||||
|
||||
#### MiMo-Audio Adapter
|
||||
|
||||
```python
|
||||
@ModelAdapterRegistry.register("MiMoAudioForConditionalGeneration")
|
||||
class MiMoAudioAdapter(DefaultModelRunnerAdapter):
|
||||
def enrich_request_info(self, runner, req_id, req_state, req_infos):
|
||||
# Attach mm_features and req_id
|
||||
...
|
||||
```
|
||||
|
||||
### 3.4 Refactored OmniGPUModelRunner
|
||||
|
||||
After extraction, the runner becomes much simpler:
|
||||
|
||||
```python
|
||||
class OmniGPUModelRunner(GPUModelRunner):
|
||||
def load_model(self, *args, **kwargs):
|
||||
super().load_model(*args, **kwargs)
|
||||
self._adapter = ModelAdapterRegistry.get_adapter(self.model)
|
||||
self._adapter.on_model_loaded(self, self.model)
|
||||
|
||||
def _update_states(self, scheduler_output):
|
||||
super()._update_states(scheduler_output)
|
||||
for req_id, req_state in ...:
|
||||
self._adapter.init_request_positions(self, req_state)
|
||||
|
||||
def _calc_mrope_positions(self, scheduler_output):
|
||||
super()._calc_mrope_positions(scheduler_output)
|
||||
self._adapter.fixup_decode_positions(self, scheduler_output)
|
||||
|
||||
def _preprocess(self, scheduler_output, num_input_tokens, ...):
|
||||
# Generic preprocessing
|
||||
...
|
||||
input_ids, inputs_embeds, extra = self._adapter.preprocess(...)
|
||||
return ...
|
||||
|
||||
def extract_multimodal_outputs(self, hidden_states):
|
||||
return self._adapter.extract_outputs(self, hidden_states)
|
||||
```
|
||||
|
||||
### 3.5 File Structure
|
||||
|
||||
```
|
||||
vllm_omni/worker/
|
||||
├── gpu_model_runner.py # Cleaned OmniGPUModelRunner (~600 lines)
|
||||
├── gpu_ar_model_runner.py # Unchanged
|
||||
├── gpu_generation_model_runner.py # Unchanged
|
||||
├── adapters/ # NEW: Model-specific adapters
|
||||
│ ├── __init__.py # ModelAdapterRegistry
|
||||
│ ├── base.py # ModelRunnerAdapter protocol + DefaultAdapter
|
||||
│ ├── qwen3_talker.py # Qwen3-Omni / TTS adapter
|
||||
│ ├── glm_image.py # GLM-Image adapter
|
||||
│ ├── mimo_audio.py # MiMo-Audio adapter
|
||||
│ └── hunyuan_vl.py # HunYuan-VL (XD-RoPE) adapter
|
||||
```
|
||||
|
||||
## 4. Migration Plan
|
||||
|
||||
| Phase | Scope | Risk |
|
||||
|-------|-------|------|
|
||||
| Phase 1 | Define `ModelRunnerAdapter` protocol and `ModelAdapterRegistry` | Low — additive only |
|
||||
| Phase 2 | Extract M-RoPE and positional encoding into adapters | Medium — touches core init path |
|
||||
| Phase 3 | Extract talker MTP logic into `Qwen3TalkerAdapter` | Medium — performance-critical path |
|
||||
| Phase 4 | Extract MiMo-Audio logic | Low — isolated code |
|
||||
| Phase 5 | Extract GLM-Image precomputed position fixup | Low — isolated code |
|
||||
| Phase 6 | Clean up `_preprocess` to use adapter dispatch | High — central hot path |
|
||||
| Phase 7 | Port platform adapters (NPU, XPU) to use the same adapter system | Medium |
|
||||
|
||||
Each phase should be a separate PR with targeted tests.
|
||||
|
||||
## 5. Testing Strategy
|
||||
|
||||
| Test | Description |
|
||||
|------|-------------|
|
||||
| Adapter unit tests | Test each adapter's methods independently with mock runner |
|
||||
| Registry tests | Verify correct adapter resolution by model class name |
|
||||
| Integration tests | Run existing E2E tests (Qwen3-TTS, GLM-Image, MiMo-Audio) |
|
||||
| Performance regression | Benchmark to ensure adapter indirection doesn't add measurable overhead |
|
||||
| Platform compatibility | Verify NPU/XPU runners work with the adapter system |
|
||||
|
||||
## 6. Benefits
|
||||
|
||||
| Benefit | Impact |
|
||||
|---------|--------|
|
||||
| **Code clarity** | Runner reduced from ~1,300 to ~600 lines, with model logic isolated |
|
||||
| **Extensibility** | New models only need an adapter class, not runner modifications |
|
||||
| **Testability** | Model-specific logic can be unit-tested independently |
|
||||
| **Rebase safety** | Upstream changes to `GPUModelRunner` are less likely to conflict with model-specific adapters |
|
||||
| **Platform parity** | NPU/XPU runners can share adapters or provide platform-specific overrides |
|
||||
|
||||
## 7. Open Questions
|
||||
|
||||
1. Should adapters be auto-discovered from the model class or explicitly configured in stage configs?
|
||||
2. Should the adapter have access to the full runner state or a restricted view (for safety)?
|
||||
3. How do we handle models that need multiple adapter behaviors composed together?
|
||||
4. Should `GPUARModelRunner` and `GPUGenerationModelRunner` also use the adapter pattern?
|
||||
@@ -0,0 +1,275 @@
|
||||
# RFC-004: Long Function Decomposition for Improved Readability
|
||||
|
||||
| Field | Value |
|
||||
|-------------|--------------------------------------|
|
||||
| **Status** | Draft |
|
||||
| **Authors** | vLLM-Omni Team |
|
||||
| **Created** | 2026-03-01 |
|
||||
|
||||
## 1. Summary
|
||||
|
||||
This RFC proposes systematically decomposing the long functions in vllm-omni's engine and worker modules. Several functions span hundreds of lines, making them difficult to read, test, and maintain. We identify the worst offenders, propose decomposition strategies, and define guidelines to prevent future regression.
|
||||
|
||||
## 2. Motivation
|
||||
|
||||
### Identified Long Functions
|
||||
|
||||
A survey of the vllm-omni codebase reveals the following functions that significantly exceed reasonable length limits:
|
||||
|
||||
| Function | File | Lines | Description |
|
||||
|----------|------|-------|-------------|
|
||||
| `_stage_worker_async` | `entrypoints/omni_stage.py` | ~400 | Async stage worker: device setup, engine init, main loop, profiler handling, output processing |
|
||||
| `_stage_worker` | `entrypoints/omni_stage.py` | ~400 | Sync stage worker: near-identical structure to async version |
|
||||
| `handle_profiler_task_local` | `entrypoints/omni_stage.py` (nested) | ~270 | Profiler start/stop/dump handling for LLM and diffusion engines |
|
||||
| `_dummy_run` | `worker/gpu_model_runner.py` | ~320 | CUDA graph warm-up/capture with multiple execution paths |
|
||||
| `_preprocess` | `worker/gpu_model_runner.py` | ~175 | Input preparation with multiple branches for mm/embeds/text/preprocess |
|
||||
| `process_inputs` | `engine/input_processor.py` | ~193 | Multimodal input processing with tokenization, embeddings, LoRA |
|
||||
| `create_model_config` | `engine/arg_utils.py` | ~146 | Config construction (duplicated for sync/async) |
|
||||
| `_sequential_init_lock` | `entrypoints/omni_stage.py` | ~146 | Lock management for sequential GPU initialization |
|
||||
| `process_engine_inputs` | `entrypoints/omni_stage.py` (class method) | ~50 | Input wiring between stages |
|
||||
|
||||
### Problems Caused by Long Functions
|
||||
|
||||
| Problem | Impact |
|
||||
|---------|--------|
|
||||
| **Cognitive load** | Developers must hold 200-400 lines of context in their head to understand control flow |
|
||||
| **Testing difficulty** | Cannot unit-test internal logic phases; must test the entire function end-to-end |
|
||||
| **Merge conflicts** | Long functions increase the probability of merge conflicts during rebases |
|
||||
| **Code duplication** | `_stage_worker` and `_stage_worker_async` are near-duplicates (~800 lines combined) with minor async/await differences |
|
||||
| **Debugging difficulty** | Stack traces point to a single long function; hard to isolate which phase failed |
|
||||
|
||||
## 3. Decomposition Proposals
|
||||
|
||||
### 3.1 `_stage_worker` and `_stage_worker_async` (omni_stage.py)
|
||||
|
||||
These two functions share ~90% of their logic. The proposed decomposition:
|
||||
|
||||
#### Phase Extraction
|
||||
|
||||
```python
|
||||
# New helper functions (shared between sync and async):
|
||||
|
||||
def _init_stage_environment(stage_payload: dict) -> StageContext:
|
||||
"""Phase 1: Environment setup — ZMQ, device mapping, plugins."""
|
||||
...
|
||||
|
||||
def _init_stage_engine(
|
||||
model: str, stage_payload: dict, stage_ctx: StageContext
|
||||
) -> LLMEngine | OmniDiffusion:
|
||||
"""Phase 2: Engine initialization with sequential lock."""
|
||||
...
|
||||
|
||||
def _init_stage_connectors(
|
||||
stage_id: str, connectors_config: dict
|
||||
) -> dict[tuple[str, str], OmniConnectorBase]:
|
||||
"""Phase 3: Build OmniConnectors."""
|
||||
...
|
||||
|
||||
def _handle_profiler_task(
|
||||
task_type: OmniStageTaskType, engine: Any, stage_id: str, stage_type: str
|
||||
) -> dict:
|
||||
"""Phase 4: Profiler start/stop/dump (extracted from nested function)."""
|
||||
...
|
||||
|
||||
def _process_single_batch(
|
||||
batch: list[dict], engine: Any, stage_ctx: StageContext
|
||||
) -> list[dict]:
|
||||
"""Phase 5: Execute a single batch through the engine."""
|
||||
...
|
||||
|
||||
def _format_and_send_outputs(
|
||||
outputs: list, out_q: Queue, stage_ctx: StageContext
|
||||
) -> None:
|
||||
"""Phase 6: Format outputs and send via IPC."""
|
||||
...
|
||||
```
|
||||
|
||||
#### Sync/Async Unification
|
||||
|
||||
```python
|
||||
def _stage_worker(model, stage_payload, in_q, out_q, ...):
|
||||
ctx = _init_stage_environment(stage_payload)
|
||||
engine = _init_stage_engine(model, stage_payload, ctx)
|
||||
connectors = _init_stage_connectors(ctx.stage_id, ctx.connectors_config)
|
||||
|
||||
while True:
|
||||
batch = _collect_batch(in_q, ctx.max_batch_size, ctx.batch_timeout)
|
||||
if _is_shutdown(batch):
|
||||
break
|
||||
if _is_profiler_task(batch):
|
||||
result = _handle_profiler_task(batch.task_type, engine, ...)
|
||||
out_q.put(result)
|
||||
continue
|
||||
outputs = _process_single_batch(batch, engine, ctx)
|
||||
_format_and_send_outputs(outputs, out_q, ctx)
|
||||
|
||||
async def _stage_worker_async(model, stage_payload, in_q, out_q, ...):
|
||||
ctx = _init_stage_environment(stage_payload)
|
||||
engine = _init_stage_engine(model, stage_payload, ctx) # async variant
|
||||
connectors = _init_stage_connectors(ctx.stage_id, ctx.connectors_config)
|
||||
|
||||
while True:
|
||||
batch = await _collect_batch_async(in_q, ...)
|
||||
# ... same structure, async versions of batch processing
|
||||
```
|
||||
|
||||
### 3.2 `_dummy_run` (gpu_model_runner.py)
|
||||
|
||||
This function handles CUDA graph capture, warm-up, and profiling. Decompose into:
|
||||
|
||||
```python
|
||||
def _dummy_run(self, num_tokens, ...):
|
||||
num_scheduled_tokens, attn_metadata = self._prepare_dummy_batch(num_tokens, ...)
|
||||
hidden_states = self._execute_dummy_forward(
|
||||
num_scheduled_tokens, attn_metadata, ...
|
||||
)
|
||||
self._run_dummy_drafter(num_tokens, is_graph_capturing, ...)
|
||||
self._finalize_dummy_run(hidden_states, num_scheduled_tokens, ...)
|
||||
|
||||
def _prepare_dummy_batch(self, num_tokens, ...):
|
||||
"""Set up fake batch data, attention metadata, LoRA if needed."""
|
||||
...
|
||||
|
||||
def _execute_dummy_forward(self, num_scheduled_tokens, attn_metadata, ...):
|
||||
"""Run the model forward pass for CUDA graph capture or warm-up."""
|
||||
...
|
||||
|
||||
def _run_dummy_drafter(self, num_tokens, is_graph_capturing, ...):
|
||||
"""Run speculative decoding drafter dummy pass if configured."""
|
||||
...
|
||||
|
||||
def _finalize_dummy_run(self, hidden_states, num_scheduled_tokens, ...):
|
||||
"""Register NVTX hooks, run EPLB step, compute logit indices."""
|
||||
...
|
||||
```
|
||||
|
||||
### 3.3 `_preprocess` (gpu_model_runner.py)
|
||||
|
||||
The current function has 4 major branches. Decompose into:
|
||||
|
||||
```python
|
||||
def _preprocess(self, scheduler_output, num_input_tokens, intermediate_tensors=None):
|
||||
input_ids, inputs_embeds, model_kwargs, ec_connector_output = (
|
||||
self._prepare_model_inputs(scheduler_output, num_input_tokens)
|
||||
)
|
||||
positions = self._compute_positions(num_input_tokens)
|
||||
intermediate_tensors = self._prepare_intermediate_tensors(
|
||||
num_input_tokens, intermediate_tensors
|
||||
)
|
||||
self._apply_model_specific_preprocessing(
|
||||
scheduler_output, input_ids, inputs_embeds, model_kwargs
|
||||
)
|
||||
return input_ids, inputs_embeds, positions, intermediate_tensors, model_kwargs, ec_connector_output
|
||||
|
||||
def _prepare_model_inputs(self, scheduler_output, num_input_tokens):
|
||||
"""Branch: mm inputs / prompt embeds / has_preprocess / text-only."""
|
||||
...
|
||||
|
||||
def _compute_positions(self, num_input_tokens):
|
||||
"""Select M-RoPE / XD-RoPE / standard positions."""
|
||||
...
|
||||
|
||||
def _apply_model_specific_preprocessing(self, ...):
|
||||
"""Overlay prompt_embeds, run has_preprocess, talker MTP."""
|
||||
...
|
||||
```
|
||||
|
||||
### 3.4 `process_inputs` (input_processor.py)
|
||||
|
||||
Decompose into phases:
|
||||
|
||||
```python
|
||||
def process_inputs(self, request_id, prompt, params):
|
||||
preprocessed = self._preprocess_prompt(request_id, prompt, params)
|
||||
tokenized = self._tokenize_and_validate(preprocessed, params)
|
||||
mm_data = self._extract_multimodal_data(tokenized, params)
|
||||
prompt_embeds = self._extract_prompt_embeds(tokenized, params)
|
||||
additional_info = self._extract_additional_information(tokenized, params)
|
||||
return self._build_engine_core_request(
|
||||
request_id, tokenized, mm_data, prompt_embeds, additional_info, params
|
||||
)
|
||||
```
|
||||
|
||||
### 3.5 `create_model_config` (arg_utils.py)
|
||||
|
||||
This function is duplicated between `OmniEngineArgs` and `AsyncOmniEngineArgs` (~260 lines combined). Decompose and deduplicate:
|
||||
|
||||
```python
|
||||
class OmniEngineArgs:
|
||||
def create_model_config(self):
|
||||
return _create_model_config_common(self, is_async=False)
|
||||
|
||||
class AsyncOmniEngineArgs:
|
||||
def create_model_config(self):
|
||||
return _create_model_config_common(self, is_async=True)
|
||||
|
||||
def _create_model_config_common(args, is_async: bool):
|
||||
"""Shared config construction logic."""
|
||||
base_config = _build_base_config(args)
|
||||
_apply_omni_overrides(base_config, args)
|
||||
_validate_config(base_config)
|
||||
return base_config
|
||||
```
|
||||
|
||||
## 4. Guidelines for Future Code
|
||||
|
||||
To prevent regression, we propose the following guidelines:
|
||||
|
||||
### 4.1 Function Length Limits
|
||||
|
||||
| Severity | Line Count | Action |
|
||||
|----------|-----------|--------|
|
||||
| OK | < 50 lines | No action needed |
|
||||
| Warning | 50-100 lines | Consider decomposition |
|
||||
| Required | > 100 lines | Must decompose before merge |
|
||||
|
||||
### 4.2 Lint Rule
|
||||
|
||||
Add a custom lint check (or use `ruff` with `pylint.max-statements`) to flag functions exceeding the threshold:
|
||||
|
||||
```toml
|
||||
# pyproject.toml
|
||||
[tool.ruff.pylint]
|
||||
max-statements = 50
|
||||
```
|
||||
|
||||
### 4.3 Decomposition Principles
|
||||
|
||||
1. **Extract by phase**: Each phase of a multi-step function becomes its own method
|
||||
2. **Extract by branch**: If a function has N major branches, each branch becomes a helper
|
||||
3. **Extract by concern**: Separate I/O, computation, and error handling
|
||||
4. **Preserve hot paths**: Do not add unnecessary function call overhead in performance-critical paths (profile after decomposition)
|
||||
5. **Keep related logic together**: Extracted functions should be in the same module, not scattered
|
||||
|
||||
## 5. Implementation Plan
|
||||
|
||||
| Phase | Target | Functions | Est. Impact |
|
||||
|-------|--------|-----------|-------------|
|
||||
| Phase 1 | `omni_stage.py` | `_stage_worker`, `_stage_worker_async`, `handle_profiler_task_local` | ~800 lines → ~200 + shared helpers |
|
||||
| Phase 2 | `gpu_model_runner.py` | `_dummy_run`, `_preprocess` | ~500 lines → ~200 + helpers |
|
||||
| Phase 3 | `input_processor.py` | `process_inputs` | ~193 lines → ~80 + helpers |
|
||||
| Phase 4 | `arg_utils.py` | `create_model_config` (dedup) | ~260 lines → ~130 shared |
|
||||
| Phase 5 | Add lint rules | Configure ruff/pylint | Prevents regression |
|
||||
|
||||
## 6. Risk Mitigation
|
||||
|
||||
| Risk | Mitigation |
|
||||
|------|------------|
|
||||
| **Performance regression** | Benchmark critical paths (model runner `_preprocess`, `_dummy_run`) before and after |
|
||||
| **Behavioral changes** | Each decomposition should be a pure refactor — no logic changes |
|
||||
| **Merge conflicts** | Coordinate with ongoing PRs; complete in small incremental PRs |
|
||||
| **Variable scoping issues** | Use `dataclass` or `NamedTuple` for passing multi-value state between phases rather than long parameter lists |
|
||||
|
||||
## 7. Success Metrics
|
||||
|
||||
- No function in the codebase exceeds 100 lines (excluding comments/docstrings)
|
||||
- `_stage_worker` and `_stage_worker_async` share > 80% of logic via shared helpers
|
||||
- No performance regression in model runner benchmarks
|
||||
- Improved code review velocity for PRs touching refactored files
|
||||
|
||||
## 8. Open Questions
|
||||
|
||||
1. Should we introduce a `StageContext` dataclass to pass state between phases in `_stage_worker`, or use simpler parameter passing?
|
||||
2. Is the 100-line limit too aggressive for GPU kernel-related code paths where inlining matters?
|
||||
3. Should we use `ruff` max-statements or a custom pre-commit hook for enforcement?
|
||||
4. How do we handle the sync/async duplication — use a code generation approach or manual unification?
|
||||
@@ -0,0 +1,277 @@
|
||||
# RFC-005: Model Runner V2 Integration Roadmap
|
||||
|
||||
| Field | Value |
|
||||
|-------------|--------------------------------------|
|
||||
| **Status** | Draft |
|
||||
| **Authors** | vLLM-Omni Team |
|
||||
| **Created** | 2026-03-01 |
|
||||
|
||||
## 1. Summary
|
||||
|
||||
Upstream vLLM is developing a V2 model runner to replace the current V1 architecture (`vllm.v1.worker.gpu_model_runner.GPUModelRunner`). vllm-omni currently **explicitly disables** V2 model runner support with a fallback warning. This RFC outlines the adaptation work required to integrate the V2 model runner into vllm-omni, including identifying the "omni hooks" that must be ported, defining a phased migration plan, and establishing compatibility testing.
|
||||
|
||||
## 2. Motivation
|
||||
|
||||
### Current State
|
||||
|
||||
Both `GPUARWorker` and `GPUGenerationWorker` contain identical guard code:
|
||||
|
||||
```python
|
||||
# vllm_omni/worker/gpu_ar_worker.py:96-99
|
||||
if self.use_v2_model_runner:
|
||||
logger.warning("OMNI GPUARWorker forces v1 model runner for omni hooks.")
|
||||
self.use_v2_model_runner = False
|
||||
|
||||
# vllm_omni/worker/gpu_generation_worker.py:96-99
|
||||
if self.use_v2_model_runner:
|
||||
logger.warning("OMNI GPUGenerationWorker forces v1 model runner for omni hooks.")
|
||||
self.use_v2_model_runner = False
|
||||
```
|
||||
|
||||
The scheduler already has partial V2 awareness:
|
||||
|
||||
```python
|
||||
# vllm_omni/core/sched/output.py:44
|
||||
prefill_token_ids: Optional # prefill token IDs for v2 model runner
|
||||
```
|
||||
|
||||
`OmniNewRequestData.from_request()` conditionally populates `prefill_token_ids` based on `use_v2_model_runner`.
|
||||
|
||||
### Why V2 Matters
|
||||
|
||||
| Reason | Detail |
|
||||
|--------|--------|
|
||||
| **Upstream convergence** | vLLM is moving toward V2 as the default; V1 may eventually be deprecated |
|
||||
| **Performance** | V2 model runner introduces optimizations for attention metadata, micro-batching, and CUDA graph management |
|
||||
| **Feature access** | New upstream features may only be available in V2 |
|
||||
| **Maintenance burden** | Supporting V1 indefinitely means maintaining a growing delta from upstream |
|
||||
|
||||
### What Are "Omni Hooks"?
|
||||
|
||||
The warning message references "omni hooks" — the vllm-omni extensions to the model runner that V2 does not yet support. Based on code analysis, these are:
|
||||
|
||||
| Hook | Location | Description |
|
||||
|------|----------|-------------|
|
||||
| `additional_information` storage | `OmniGPUModelRunner.__init__` | Per-request `additional_information_cpu` attribute on `CachedRequestState` |
|
||||
| `_decode_and_store_request_payloads` | `OmniGPUModelRunner` | Decodes prompt_embeds and additional_information from scheduler output for new requests |
|
||||
| `_collect_additional_information_for_prefill` | `OmniGPUModelRunner` | Overlays prompt_embeds on inputs_embeds during prefill |
|
||||
| `_build_model_kwargs_extra` | `OmniGPUModelRunner` | Injects `runtime_additional_information` into model forward kwargs |
|
||||
| `_process_additional_information_updates` | `OmniGPUModelRunner` | Merges model postprocess outputs back into request state |
|
||||
| `_merge_additional_information_update` | `OmniGPUModelRunner` | CPU-side merge of update dicts with tensor offloading |
|
||||
| `extract_multimodal_outputs` | `OmniGPUModelRunner` | Splits `OmniOutput` into text hidden states + multimodal outputs |
|
||||
| `_model_forward` override | `OmniGPUModelRunner` | Injects extra kwargs and caches model output for multimodal |
|
||||
| M-RoPE extensions | `OmniGPUModelRunner` | Extended M-RoPE with audio, video, spatial metadata |
|
||||
| XD-RoPE | `OmniGPUModelRunner` | HunYuan-VL dimensional RoPE |
|
||||
| Talker MTP | `OmniGPUModelRunner` | Qwen3-TTS/Omni code predictor forward pass |
|
||||
| `has_preprocess` / `has_postprocess` | `OmniGPUModelRunner._preprocess` | Per-request model preprocessing/postprocessing |
|
||||
| `OmniModelRunnerOutput` | `outputs.py` | Extended `ModelRunnerOutput` with `multimodal_outputs` and `kv_extracted_req_ids` |
|
||||
|
||||
## 3. Analysis of V2 Model Runner Architecture
|
||||
|
||||
### 3.1 Key V2 Changes (Based on Upstream vLLM)
|
||||
|
||||
The V2 model runner refactors several areas that affect omni hooks:
|
||||
|
||||
| Area | V1 Behavior | V2 Behavior | Omni Impact |
|
||||
|------|------------|-------------|-------------|
|
||||
| **Input batch** | `GPUInputBatch` with direct attribute access | Restructured batch with new token layout | `_collect_additional_information_for_prefill` needs adaptation |
|
||||
| **Forward context** | `set_forward_context()` with CUDA graph mode | Enhanced forward context with batch descriptors | `_model_forward` override must align |
|
||||
| **Attention metadata** | Per-layer metadata builders | Restructured metadata pipeline | M-RoPE position injection path changes |
|
||||
| **Output handling** | `ModelRunnerOutput` returned from `execute_model` | Potentially different output structure | `OmniModelRunnerOutput` needs remapping |
|
||||
| **CUDA graph capture** | `_dummy_run` with manual graph management | V2 CUDA graph system | `_dummy_run` override may need rewrite |
|
||||
| **Scheduler output** | `SchedulerOutput` with `scheduled_new_reqs` | V2 scheduler output format | Request payload decoding path changes |
|
||||
|
||||
### 3.2 Hook Compatibility Assessment
|
||||
|
||||
| Hook | V2 Compatibility | Effort |
|
||||
|------|-----------------|--------|
|
||||
| `additional_information` storage | **Needs adaptation** — V2 may change `CachedRequestState` | Medium |
|
||||
| `_decode_and_store_request_payloads` | **Needs adaptation** — scheduler output format may differ | Medium |
|
||||
| `_collect_additional_information_for_prefill` | **Needs adaptation** — input batch layout changes | High |
|
||||
| `_build_model_kwargs_extra` | **Likely compatible** — model forward kwargs injection | Low |
|
||||
| `_model_forward` override | **Needs adaptation** — forward context API changes | Medium |
|
||||
| `extract_multimodal_outputs` | **Likely compatible** — output extraction is model-level | Low |
|
||||
| M-RoPE extensions | **Needs adaptation** — position computation pipeline changes | High |
|
||||
| Talker MTP | **Needs adaptation** — CUDA graph management changes | High |
|
||||
| `OmniModelRunnerOutput` | **Needs adaptation** — output class hierarchy may change | Medium |
|
||||
|
||||
## 4. Proposed Integration Plan
|
||||
|
||||
### 4.1 Phase 0: Upstream Tracking (Ongoing)
|
||||
|
||||
- Monitor upstream V2 model runner PRs and API stabilization
|
||||
- Maintain a mapping document of V1 → V2 API changes relevant to omni hooks
|
||||
- Track which omni hooks have upstream equivalents in V2
|
||||
|
||||
### 4.2 Phase 1: Abstract the Hook Interface
|
||||
|
||||
Before integrating V2, abstract the hook points so that both V1 and V2 can be supported:
|
||||
|
||||
```python
|
||||
class OmniModelRunnerHooks(Protocol):
|
||||
"""Interface for omni extensions to the model runner."""
|
||||
|
||||
def on_new_requests(self, scheduler_output) -> None:
|
||||
"""Decode and store payloads for new requests."""
|
||||
...
|
||||
|
||||
def prepare_prefill_embeds(self, num_scheduled_tokens_np) -> None:
|
||||
"""Overlay prompt_embeds for prefill."""
|
||||
...
|
||||
|
||||
def build_extra_kwargs(self) -> dict:
|
||||
"""Build extra model forward kwargs."""
|
||||
...
|
||||
|
||||
def process_model_output(self, hidden_states, mm_outputs, scheduler_output) -> None:
|
||||
"""Post-forward processing of model outputs."""
|
||||
...
|
||||
|
||||
def get_model_runner_output(self, base_output) -> ModelRunnerOutput:
|
||||
"""Wrap base output with omni-specific fields."""
|
||||
...
|
||||
```
|
||||
|
||||
This abstraction aligns with RFC-003 (Model Runner Clearance) and can be implemented concurrently.
|
||||
|
||||
### 4.3 Phase 2: V2 Runner Skeleton
|
||||
|
||||
Create a V2-compatible runner class:
|
||||
|
||||
```python
|
||||
class OmniGPUModelRunnerV2(GPUModelRunnerV2):
|
||||
"""V2 model runner with omni hooks."""
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
self._hooks = OmniModelRunnerHooksImpl(self)
|
||||
|
||||
def execute_model(self, scheduler_output, ...):
|
||||
self._hooks.on_new_requests(scheduler_output)
|
||||
self._hooks.prepare_prefill_embeds(...)
|
||||
result = super().execute_model(scheduler_output, ...)
|
||||
self._hooks.process_model_output(result, ...)
|
||||
return self._hooks.get_model_runner_output(result)
|
||||
```
|
||||
|
||||
### 4.4 Phase 3: Port Individual Hooks
|
||||
|
||||
Port each hook to V2, prioritized by complexity and importance:
|
||||
|
||||
| Priority | Hook | Rationale |
|
||||
|----------|------|-----------|
|
||||
| P0 | `_model_forward` override | Core execution path |
|
||||
| P0 | `extract_multimodal_outputs` | Required for any multimodal output |
|
||||
| P0 | `OmniModelRunnerOutput` | Required for output pipeline |
|
||||
| P1 | `additional_information` lifecycle | Required for multi-stage pipelines |
|
||||
| P1 | M-RoPE extensions | Required for Qwen2-VL, Qwen3-Omni |
|
||||
| P2 | Talker MTP | Required for Qwen3-TTS |
|
||||
| P2 | `has_preprocess` / `has_postprocess` | Required for custom models |
|
||||
| P3 | XD-RoPE | Only needed for HunYuan-VL |
|
||||
|
||||
### 4.5 Phase 4: Worker Updates
|
||||
|
||||
Update workers to conditionally use V2:
|
||||
|
||||
```python
|
||||
class GPUARWorker(OmniWorkerMixin, OmniGPUWorkerBase):
|
||||
def init_device(self):
|
||||
...
|
||||
if self.use_v2_model_runner:
|
||||
self.model_runner = GPUARModelRunnerV2(self.vllm_config, self.device)
|
||||
else:
|
||||
self.model_runner = GPUARModelRunner(self.vllm_config, self.device)
|
||||
```
|
||||
|
||||
### 4.6 Phase 5: Validation and Cutover
|
||||
|
||||
1. Run full test suite with `use_v2_model_runner=True`
|
||||
2. Benchmark V2 against V1 for key models (Qwen3-Omni, Qwen3-TTS, GLM-Image)
|
||||
3. Enable V2 as opt-in via environment variable
|
||||
4. Eventually make V2 the default when all hooks are ported
|
||||
|
||||
## 5. Compatibility Strategy
|
||||
|
||||
### 5.1 Dual Runner Support
|
||||
|
||||
During the transition, both V1 and V2 runners coexist:
|
||||
|
||||
```
|
||||
vllm_omni/worker/
|
||||
├── gpu_model_runner.py # V1: OmniGPUModelRunner (existing)
|
||||
├── gpu_model_runner_v2.py # V2: OmniGPUModelRunnerV2 (new)
|
||||
├── gpu_ar_model_runner.py # V1 AR runner
|
||||
├── gpu_ar_model_runner_v2.py # V2 AR runner (new)
|
||||
├── gpu_generation_model_runner.py # V1 generation runner
|
||||
├── gpu_generation_model_runner_v2.py # V2 generation runner (new)
|
||||
```
|
||||
|
||||
### 5.2 Feature Parity Matrix
|
||||
|
||||
Track feature parity with a compatibility matrix:
|
||||
|
||||
| Feature | V1 | V2 | Notes |
|
||||
|---------|----|----|-------|
|
||||
| Text generation | Yes | Phase 2 | |
|
||||
| Multimodal input (image/video/audio) | Yes | Phase 2 | |
|
||||
| Multimodal output | Yes | Phase 3 | |
|
||||
| Multi-stage pipeline | Yes | Phase 3 | |
|
||||
| M-RoPE | Yes | Phase 3 | |
|
||||
| Talker MTP | Yes | Phase 3 | |
|
||||
| CUDA graph capture | Yes | Phase 2 | |
|
||||
| Prompt embeds | Yes | Phase 3 | |
|
||||
| Diffusion models | Yes | Phase 4 | Separate runner |
|
||||
| NPU support | Yes | Phase 5 | Depends on `vllm_ascend` V2 |
|
||||
| XPU support | Yes | Phase 5 | |
|
||||
|
||||
## 6. Platform Considerations
|
||||
|
||||
### 6.1 NPU (Ascend)
|
||||
|
||||
NPU currently uses `vllm_ascend.worker.model_runner_v1.NPUModelRunner` as its base. The V2 integration depends on:
|
||||
- Whether `vllm_ascend` provides a V2 model runner
|
||||
- Whether NPU-specific changes can be layered on top of `OmniGPUModelRunnerV2`
|
||||
|
||||
### 6.2 XPU (Intel)
|
||||
|
||||
XPU runners (`XPUARModelRunner`, `XPUGenerationModelRunner`) inherit from the GPU runners. They should automatically benefit from V2 if the GPU V2 runner is correct, but may need platform-specific overrides for XPU memory management.
|
||||
|
||||
## 7. Testing Plan
|
||||
|
||||
| Level | Description |
|
||||
|-------|-------------|
|
||||
| Unit | Test each ported hook independently with mock V2 runner |
|
||||
| Integration | Run E2E tests for each model with `use_v2_model_runner=True` |
|
||||
| Benchmark | Compare V1 vs V2 throughput/latency for Qwen3-Omni, Qwen3-TTS |
|
||||
| CI gate | Add V2 runner to CI matrix (opt-in initially, required later) |
|
||||
| Platform | Test NPU/XPU with V2 runner (dependent on platform V2 support) |
|
||||
|
||||
## 8. Timeline
|
||||
|
||||
| Phase | Scope | Dependencies | Est. Duration |
|
||||
|-------|-------|-------------|---------------|
|
||||
| Phase 0 | Upstream tracking | None | Ongoing |
|
||||
| Phase 1 | Abstract hook interface | RFC-003 | 2 weeks |
|
||||
| Phase 2 | V2 runner skeleton + core hooks | Phase 1 + upstream V2 stabilization | 3 weeks |
|
||||
| Phase 3 | Port all hooks | Phase 2 | 4 weeks |
|
||||
| Phase 4 | Worker updates + opt-in flag | Phase 3 | 1 week |
|
||||
| Phase 5 | Validation + platform ports | Phase 4 | 3 weeks |
|
||||
| Cutover | V2 as default | Phase 5 + confidence from CI | TBD |
|
||||
|
||||
## 9. Risks
|
||||
|
||||
| Risk | Mitigation |
|
||||
|------|------------|
|
||||
| Upstream V2 API instability | Phase 0 tracking; abstract hook interface insulates from changes |
|
||||
| Performance regression in V2 | Benchmark gating before cutover |
|
||||
| Feature incompleteness | Feature parity matrix; V1 fallback always available |
|
||||
| Platform support gaps | NPU/XPU V2 treated as separate Phase 5 work |
|
||||
| Large merge conflicts during V2 port | Coordinate with rebase automation (RFC-001) |
|
||||
|
||||
## 10. Open Questions
|
||||
|
||||
1. Is upstream V2 model runner stable enough to begin Phase 2, or should we wait for an official release?
|
||||
2. Should we attempt to contribute omni hooks upstream to V2, or keep them in vllm-omni?
|
||||
3. How does V2 affect the `DiffusionModelRunner` (which has its own independent runner implementation)?
|
||||
4. Should the V2 integration be gated behind a configuration flag or environment variable during the transition?
|
||||
5. What is the minimum set of hooks needed for a useful V2 MVP (i.e., which models can work with a partial port)?
|
||||
@@ -68,6 +68,7 @@ The following table shows which models are currently supported by each accelerat
|
||||
| **Bagel** | `ByteDance-Seed/BAGEL-7B-MoT` | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ |
|
||||
| **FLUX.1-dev** | `black-forest-labs/FLUX.1-dev` | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ |
|
||||
| **NextStep-1.1** | `stepfun-ai/NextStep-1.1` | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ |
|
||||
| **FLUX.2-klein** | `black-forest-labs/FLUX.2-klein-4B` | ❌ | ✅ | ❌ | ❌ | ❌ | ✅ | ❌ |
|
||||
|
||||
### VideoGen
|
||||
|
||||
|
||||
@@ -23,6 +23,29 @@ Or use the startup script:
|
||||
bash run_server.sh
|
||||
```
|
||||
|
||||
### Start with Parallelism Acceleration
|
||||
|
||||
Enable Tensor Parallelism and VAE Patch Parallelism for faster inference:
|
||||
|
||||
```bash
|
||||
# With Tensor Parallelism (requires >= 2 GPUs)
|
||||
vllm serve Qwen/Qwen-Image --omni --port 8091 --tensor-parallel-size 2
|
||||
|
||||
# With Tensor Parallelism and VAE Patch Parallelism (requires >= 2 GPUs)
|
||||
vllm serve Qwen/Qwen-Image --omni --port 8091 --tensor-parallel-size 2 --vae-patch-parallel-size 2 --vae-use-tiling
|
||||
|
||||
# With Sequence Parallelism (Ulysses-SP, requires >= 2 GPUs)
|
||||
vllm serve Qwen/Qwen-Image --omni --port 8091 --usp 2
|
||||
|
||||
# With Ring-Attention (requires >= 2 GPUs)
|
||||
vllm serve Qwen/Qwen-Image --omni --port 8091 --ring 2
|
||||
|
||||
# Combined: Ulysses + Ring (requires >= 4 GPUs)
|
||||
vllm serve Qwen/Qwen-Image --omni --port 8091 --usp 2 --ring 2
|
||||
```
|
||||
|
||||
For more details on parallelism acceleration, see the [Parallelism Acceleration Guide](../../diffusion/parallelism_acceleration.md).
|
||||
|
||||
## API Calls
|
||||
|
||||
### Method 1: Using curl
|
||||
|
||||
@@ -99,6 +99,7 @@ def parse_args() -> argparse.Namespace:
|
||||
parser.add_argument("--vae_use_slicing", action="store_true")
|
||||
parser.add_argument("--vae_use_tiling", action="store_true")
|
||||
parser.add_argument("--enable-cpu-offload", action="store_true")
|
||||
parser.add_argument("--log-stats", "--log_stats", dest="log_stats", action="store_true", default=False)
|
||||
|
||||
return parser.parse_args()
|
||||
|
||||
@@ -155,6 +156,7 @@ async def main():
|
||||
cache_config=cache_config,
|
||||
parallel_config=parallel_config,
|
||||
enforce_eager=args.enforce_eager,
|
||||
log_stats=args.log_stats,
|
||||
enable_cpu_offload=args.enable_cpu_offload,
|
||||
diffusion_load_format="dummy",
|
||||
custom_pipeline_args={"pipeline_class": "custom_pipeline.CustomPipeline"},
|
||||
|
||||
@@ -320,6 +320,11 @@ def parse_args() -> argparse.Namespace:
|
||||
action="store_true",
|
||||
help="Enable layerwise (blockwise) offloading on DiT modules.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--log-stats",
|
||||
action="store_true",
|
||||
help="Enable vLLM-Omni statistics logging.",
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
@@ -383,6 +388,7 @@ def main():
|
||||
parallel_config=parallel_config,
|
||||
enforce_eager=args.enforce_eager,
|
||||
enable_cpu_offload=args.enable_cpu_offload,
|
||||
log_stats=args.log_stats,
|
||||
)
|
||||
print("Pipeline loaded")
|
||||
|
||||
|
||||
@@ -176,6 +176,11 @@ def parse_args() -> argparse.Namespace:
|
||||
"Default 1 means pure sharding (no replication). "
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--log-stats",
|
||||
action="store_true",
|
||||
help="Enable vLLM-Omni statistics logging.",
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
@@ -278,6 +283,7 @@ def main():
|
||||
enable_cpu_offload=args.enable_cpu_offload,
|
||||
parallel_config=parallel_config,
|
||||
enforce_eager=args.enforce_eager,
|
||||
log_stats=args.log_stats,
|
||||
model_class_name=model_class_name,
|
||||
cache_backend=args.cache_backend,
|
||||
cache_config=cache_config,
|
||||
|
||||
@@ -94,6 +94,11 @@ def parse_args() -> argparse.Namespace:
|
||||
action="store_true",
|
||||
help="Enable cache-dit summary logging after diffusion forward passes.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--log-stats",
|
||||
action="store_true",
|
||||
help="Enable vLLM-Omni statistics logging.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--ulysses-degree",
|
||||
type=int,
|
||||
@@ -299,6 +304,7 @@ def main():
|
||||
"parallel_config": parallel_config,
|
||||
"enforce_eager": args.enforce_eager,
|
||||
"enable_cpu_offload": args.enable_cpu_offload,
|
||||
"log_stats": args.log_stats,
|
||||
**lora_args,
|
||||
**quant_kwargs,
|
||||
}
|
||||
|
||||
@@ -127,6 +127,11 @@ def parse_args() -> argparse.Namespace:
|
||||
default=1,
|
||||
help="Number of GPUs used for VAE patch/tile parallelism (decode).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--log-stats",
|
||||
action="store_true",
|
||||
help="Enable vLLM-Omni statistics logging.",
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
@@ -176,6 +181,7 @@ def main():
|
||||
enable_cpu_offload=args.enable_cpu_offload,
|
||||
parallel_config=parallel_config,
|
||||
enforce_eager=args.enforce_eager,
|
||||
log_stats=args.log_stats,
|
||||
cache_backend=args.cache_backend,
|
||||
cache_config=cache_config,
|
||||
)
|
||||
|
||||
@@ -23,11 +23,9 @@ from tests.utils import hardware_test
|
||||
MODEL = "Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice"
|
||||
|
||||
|
||||
def get_stage_config():
|
||||
def get_stage_config(name: str = "qwen3_tts.yaml"):
|
||||
"""Get the stage config path for Qwen3-TTS."""
|
||||
return str(
|
||||
Path(__file__).parent.parent.parent.parent / "vllm_omni" / "model_executor" / "stage_configs" / "qwen3_tts.yaml"
|
||||
)
|
||||
return str(Path(__file__).parent.parent.parent.parent / "vllm_omni" / "model_executor" / "stage_configs" / name)
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
@@ -225,3 +223,63 @@ class TestQwen3TTSAPIEndpoints:
|
||||
data = response.json()
|
||||
assert "data" in data
|
||||
assert len(data["data"]) > 0
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def omni_server_no_async_chunk():
|
||||
"""Start vLLM-Omni server with non-async-chunk config."""
|
||||
stage_config_path = get_stage_config("qwen3_tts_no_async_chunk.yaml")
|
||||
|
||||
with OmniServer(
|
||||
MODEL,
|
||||
[
|
||||
"--stage-configs-path",
|
||||
stage_config_path,
|
||||
"--stage-init-timeout",
|
||||
"120",
|
||||
"--trust-remote-code",
|
||||
"--enforce-eager",
|
||||
"--disable-log-stats",
|
||||
],
|
||||
) as server:
|
||||
yield server
|
||||
|
||||
|
||||
class TestQwen3TTSNoAsyncChunk:
|
||||
"""E2E tests for Qwen3-TTS in non-async-chunk (full decode) mode."""
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.omni
|
||||
@hardware_test(res={"cuda": "L4"}, num_cards=4)
|
||||
def test_speech_english(self, omni_server_no_async_chunk) -> None:
|
||||
"""Test English TTS with non-async-chunk pipeline."""
|
||||
response = make_speech_request(
|
||||
host=omni_server_no_async_chunk.host,
|
||||
port=omni_server_no_async_chunk.port,
|
||||
text="Hello, how are you?",
|
||||
voice="vivian",
|
||||
language="English",
|
||||
)
|
||||
|
||||
assert response.status_code == 200, f"Request failed: {response.text}"
|
||||
assert response.headers.get("content-type") == "audio/wav"
|
||||
assert verify_wav_audio(response.content), "Response is not valid WAV audio"
|
||||
assert len(response.content) > MIN_AUDIO_BYTES
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.omni
|
||||
@hardware_test(res={"cuda": "L4"}, num_cards=4)
|
||||
def test_speech_chinese(self, omni_server_no_async_chunk) -> None:
|
||||
"""Test Chinese TTS with non-async-chunk pipeline."""
|
||||
response = make_speech_request(
|
||||
host=omni_server_no_async_chunk.host,
|
||||
port=omni_server_no_async_chunk.port,
|
||||
text="你好，我是通义千问",
|
||||
voice="vivian",
|
||||
language="Chinese",
|
||||
)
|
||||
|
||||
assert response.status_code == 200, f"Request failed: {response.text}"
|
||||
assert response.headers.get("content-type") == "audio/wav"
|
||||
assert verify_wav_audio(response.content), "Response is not valid WAV audio"
|
||||
assert len(response.content) > MIN_AUDIO_BYTES
|
||||
+79
@@ -346,6 +346,84 @@ def enable_cache_for_flux(pipeline: Any, cache_config: Any) -> Callable[[int], N
|
||||
return refresh_cache_context
|
||||
|
||||
|
||||
def enable_cache_for_flux2_klein(pipeline: Any, cache_config: Any) -> Callable[[int], None]:
|
||||
"""Enable cache-dit for FLUX.2-klein-4B pipeline.
|
||||
|
||||
Args:
|
||||
pipeline: The FLUX.2-klein-4B pipeline instance.
|
||||
cache_config: DiffusionCacheConfig instance with cache configuration.
|
||||
Returns:
|
||||
A refresh function that can be called with a new ``num_inference_steps``
|
||||
to update the cache context for the pipeline.
|
||||
"""
|
||||
# Build DBCacheConfig for transformer
|
||||
db_cache_config = _build_db_cache_config(cache_config)
|
||||
# The Fn_compute_blocks = 2 override is the most important decision here,
|
||||
# and the rationale (quality degradation at Fn=1) only lives in flux2_klein.
|
||||
db_cache_config.Fn_compute_blocks = 2
|
||||
|
||||
calibrator = None
|
||||
if cache_config.enable_taylorseer:
|
||||
taylorseer_order = cache_config.taylorseer_order
|
||||
calibrator = TaylorSeerCalibratorConfig(taylorseer_order=taylorseer_order)
|
||||
logger.info(f"TaylorSeer enabled with order={taylorseer_order}")
|
||||
|
||||
# Build ParamsModifier for transformer
|
||||
modifier = ParamsModifier(
|
||||
cache_config=db_cache_config,
|
||||
calibrator_config=calibrator,
|
||||
)
|
||||
|
||||
logger.info(
|
||||
f"Enabling cache-dit on Flux2-Klein transformer with BlockAdapter: "
|
||||
f"Fn={db_cache_config.Fn_compute_blocks}, "
|
||||
f"Bn={db_cache_config.Bn_compute_blocks}, "
|
||||
f"W={db_cache_config.max_warmup_steps}, "
|
||||
)
|
||||
|
||||
# Enable cache-dit using BlockAdapter for transformer
|
||||
cache_dit.enable_cache(
|
||||
(
|
||||
BlockAdapter(
|
||||
transformer=pipeline.transformer,
|
||||
blocks=[
|
||||
pipeline.transformer.transformer_blocks,
|
||||
pipeline.transformer.single_transformer_blocks,
|
||||
],
|
||||
forward_pattern=[ForwardPattern.Pattern_1, ForwardPattern.Pattern_2],
|
||||
params_modifiers=[modifier],
|
||||
)
|
||||
),
|
||||
cache_config=db_cache_config,
|
||||
)
|
||||
|
||||
def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None:
|
||||
"""Refresh cache context for the transformer with new num_inference_steps.
|
||||
|
||||
Args:
|
||||
pipeline: The FLUX.2-klein-4B pipeline instance.
|
||||
num_inference_steps: New number of inference steps.
|
||||
"""
|
||||
if cache_config.scm_steps_mask_policy is None:
|
||||
cache_dit.refresh_context(pipeline.transformer, num_inference_steps=num_inference_steps, verbose=verbose)
|
||||
else:
|
||||
cache_dit.refresh_context(
|
||||
pipeline.transformer,
|
||||
cache_config=DBCacheConfig().reset(
|
||||
num_inference_steps=num_inference_steps,
|
||||
steps_computation_mask=cache_dit.steps_mask(
|
||||
mask_policy=cache_config.scm_steps_mask_policy,
|
||||
total_steps=num_inference_steps,
|
||||
),
|
||||
Fn_compute_blocks=db_cache_config.Fn_compute_blocks,
|
||||
steps_computation_policy=cache_config.scm_steps_policy,
|
||||
),
|
||||
verbose=verbose,
|
||||
)
|
||||
|
||||
return refresh_cache_context
|
||||
|
||||
|
||||
def enable_cache_for_sd3(pipeline: Any, cache_config: Any) -> Callable[[int], None]:
|
||||
"""Enable cache-dit for StableDiffusion3Pipeline.
|
||||
|
||||
@@ -913,6 +991,7 @@ CUSTOM_DIT_ENABLERS.update(
|
||||
"Wan22I2VPipeline": enable_cache_for_wan22,
|
||||
"Wan22TI2VPipeline": enable_cache_for_wan22,
|
||||
"FluxPipeline": enable_cache_for_flux,
|
||||
"Flux2KleinPipeline": enable_cache_for_flux2_klein,
|
||||
"LongCatImagePipeline": enable_cache_for_longcat_image,
|
||||
"LongCatImageEditPipeline": enable_cache_for_longcat_image,
|
||||
"StableDiffusion3Pipeline": enable_cache_for_sd3,
|
||||
|
||||
@@ -67,14 +67,19 @@ class DiffusionEngine:
|
||||
raise e
|
||||
|
||||
def step(self, request: OmniDiffusionRequest) -> list[OmniRequestOutput]:
|
||||
diffusion_engine_start_time = time.time()
|
||||
# Apply pre-processing if available
|
||||
preprocess_time = 0.0
|
||||
if self.pre_process_func is not None:
|
||||
preprocess_start_time = time.time()
|
||||
request = self.pre_process_func(request)
|
||||
preprocess_time = time.time() - preprocess_start_time
|
||||
logger.info(f"Pre-processing completed in {preprocess_time:.4f} seconds")
|
||||
|
||||
exec_start_time = time.time()
|
||||
output = self.add_req_and_wait_for_response(request)
|
||||
exec_total_time = time.time() - exec_start_time
|
||||
|
||||
if output.error:
|
||||
raise Exception(f"{output.error}")
|
||||
logger.info("Generation completed successfully.")
|
||||
@@ -107,14 +112,20 @@ class DiffusionEngine:
|
||||
outputs = [outputs] if outputs is not None else []
|
||||
|
||||
metrics = {
|
||||
"preprocess_time_ms": preprocess_time * 1000,
|
||||
"diffusion_engine_exec_time_ms": (time.time() - diffusion_engine_start_time) * 1000,
|
||||
"diffusion_engine_total_time_ms": exec_total_time * 1000,
|
||||
"image_num": int(request.sampling_params.num_outputs_per_prompt),
|
||||
"resolution": int(request.sampling_params.resolution),
|
||||
"postprocess_time_ms": postprocess_time * 1000,
|
||||
}
|
||||
|
||||
if self.pre_process_func is not None:
|
||||
metrics["preprocessing_time_ms"] = preprocess_time * 1000
|
||||
|
||||
# Handle single request or multiple requests
|
||||
metrics["postprocess_time_ms"] = postprocess_time * 1000
|
||||
metrics["num_inference_steps"] = int(request.sampling_params.num_inference_steps)
|
||||
|
||||
if len(request.prompts) == 1:
|
||||
# Single request: return single OmniRequestOutput
|
||||
prompt = request.prompts[0]
|
||||
@@ -201,7 +212,7 @@ class DiffusionEngine:
|
||||
),
|
||||
)
|
||||
|
||||
return results
|
||||
return results
|
||||
|
||||
@staticmethod
|
||||
def make_engine(config: OmniDiffusionConfig) -> "DiffusionEngine":
|
||||
|
||||
@@ -1,4 +1,3 @@
|
||||
from collections.abc import Callable
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
@@ -60,9 +59,11 @@ class OmniRequestState(RequestState):
|
||||
incoming: dict[str, Any] = {}
|
||||
target_key = self.mm_type or "hidden"
|
||||
|
||||
# Iterate directly without unnecessary dict copy
|
||||
for k, v in payload.items():
|
||||
# Optional remap: if producer used "model_outputs" or "hidden", rename to mm_type
|
||||
# Normalize producer keys to the modality name.
|
||||
# AR runners produce {"hidden": ...} and generation
|
||||
# runners produce {"model_outputs": ...}; remap both
|
||||
# to the semantic modality key (e.g. "audio", "latent").
|
||||
if k == "model_outputs":
|
||||
k = target_key
|
||||
elif k == "hidden" and target_key != "hidden":
|
||||
@@ -247,16 +248,17 @@ class OmniRequestState(RequestState):
|
||||
|
||||
|
||||
class MultimodalOutputProcessor(VLLMOutputProcessor):
|
||||
"""Handles multimodal output processing by normalizing EngineCoreOutput
|
||||
before delegating to the base vLLM OutputProcessor.
|
||||
"""Handles multimodal output processing by capturing pooling_output
|
||||
from EngineCoreOutput and accumulating it as multimodal tensors,
|
||||
before delegating to the base vLLM OutputProcessor for text handling.
|
||||
|
||||
Strategy:
|
||||
- Route by EngineCoreOutput.output_type when present
|
||||
("image", "text+image", "latents", "text").
|
||||
- Fallback to pooling/text heuristics when output_type is absent.
|
||||
- Mutate EngineCoreOutput in-place to ensure vLLM's base processor can
|
||||
produce the correct RequestOutput/PoolingRequestOutput.
|
||||
- Allow custom per-modality handlers via register_handler().
|
||||
The actual data flow is:
|
||||
1. For each EngineCoreOutput with pooling_output and a detokenizer:
|
||||
- Capture pooling_output into OmniRequestState.add_multimodal_tensor()
|
||||
- Clear eco.pooling_output to force text path in base processor
|
||||
2. Base vLLM OutputProcessor handles text detokenization
|
||||
3. On finish, _consolidate_multimodal_tensors() concatenates accumulated tensors
|
||||
4. _new_completion_output() attaches mm_accumulated to CompletionOutput
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
@@ -275,8 +277,8 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
|
||||
log_stats: Whether to log statistics
|
||||
stream_interval: Stream interval for output generation
|
||||
engine_core_output_type: Optional output type specification
|
||||
(e.g., "image", "audio", "latents"). Used to route outputs
|
||||
to appropriate processors. If None, output type is inferred.
|
||||
(e.g., "image", "audio", "latent"). Used to tag multimodal
|
||||
outputs with the correct modality key.
|
||||
"""
|
||||
super().__init__(
|
||||
tokenizer=tokenizer,
|
||||
@@ -284,23 +286,8 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
|
||||
stream_interval=stream_interval,
|
||||
tracing_enabled=tracing_enabled,
|
||||
)
|
||||
self.output_handlers: dict[str, Callable[[EngineCoreOutput], None]] = {}
|
||||
self._reqid_to_mm_type: dict[str, str] = {}
|
||||
self.engine_core_output_type = engine_core_output_type
|
||||
|
||||
def register_handler(self, modality: str, handler: Callable[[EngineCoreOutput], None]) -> None:
|
||||
"""Register a custom handler for a specific modality.
|
||||
|
||||
Allows custom processing logic for specific output modalities.
|
||||
The handler is called before default processing for outputs
|
||||
matching the specified modality.
|
||||
|
||||
Args:
|
||||
modality: Modality name (e.g., "image", "audio", "latents")
|
||||
handler: Callable that takes an EngineCoreOutput and processes it
|
||||
"""
|
||||
self.output_handlers[modality.lower()] = handler
|
||||
|
||||
def add_request(
|
||||
self,
|
||||
request: EngineCoreRequest,
|
||||
@@ -351,20 +338,13 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
|
||||
engine_core_timestamp: float | None = None,
|
||||
iteration_stats: IterationStats | None = None,
|
||||
) -> OutputProcessorOutput:
|
||||
self._reqid_to_mm_type.clear()
|
||||
for eco in engine_core_outputs:
|
||||
mm_type = (self.engine_core_output_type or "").lower()
|
||||
if mm_type:
|
||||
self._reqid_to_mm_type[eco.request_id] = mm_type
|
||||
self._route_and_normalize(eco)
|
||||
req_state = self.request_states.get(eco.request_id)
|
||||
if req_state is None or not isinstance(req_state, OmniRequestState):
|
||||
continue
|
||||
if eco.pooling_output is not None and req_state.detokenizer is not None:
|
||||
req_state.add_multimodal_tensor(
|
||||
eco.pooling_output,
|
||||
(getattr(eco, "output_type", self.engine_core_output_type) or "").lower(),
|
||||
)
|
||||
mm_type = (getattr(eco, "output_type", self.engine_core_output_type) or "").lower()
|
||||
req_state.add_multimodal_tensor(eco.pooling_output, mm_type)
|
||||
# Force text path in base processor for multimodal outputs.
|
||||
eco.pooling_output = None
|
||||
|
||||
@@ -373,105 +353,3 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
|
||||
engine_core_timestamp=engine_core_timestamp,
|
||||
iteration_stats=iteration_stats,
|
||||
)
|
||||
|
||||
# ---- routing helpers ----
|
||||
def _route_and_normalize(self, eco: EngineCoreOutput) -> None:
|
||||
output_type = (getattr(eco, "output_type", self.engine_core_output_type) or "").lower()
|
||||
|
||||
# Custom handler first (if registered)
|
||||
if output_type in self.output_handlers:
|
||||
try:
|
||||
self.output_handlers[output_type](eco)
|
||||
# Fall through to default fixups in case the handler left gaps
|
||||
except Exception:
|
||||
logger.exception("Error in custom output handler for %s", output_type)
|
||||
|
||||
if output_type == "image":
|
||||
self._process_image_output(eco)
|
||||
elif output_type in ("text+image", "text,image", "image+text"):
|
||||
self._process_text_image_output(eco)
|
||||
elif output_type in ("latents", "latent"):
|
||||
self._process_latents_output(eco)
|
||||
elif output_type in ("audio", "speech"):
|
||||
self._process_audio_output(eco)
|
||||
elif output_type == "text":
|
||||
self._process_text_output(eco)
|
||||
else:
|
||||
# Fallback heuristic
|
||||
if eco.pooling_output is not None:
|
||||
self._process_pooling_output(eco)
|
||||
else:
|
||||
self._process_text_output(eco)
|
||||
|
||||
# ---- modality processors ----
|
||||
def _process_image_output(self, eco: EngineCoreOutput) -> None:
|
||||
"""Ensure image tensors are surfaced via pooling_output for vLLM."""
|
||||
if eco.pooling_output is None:
|
||||
tensor = self._extract_from_multimodal_outputs(eco, keys=("image", "images", "pixel_values", "pixels"))
|
||||
if tensor is not None:
|
||||
eco.pooling_output = tensor
|
||||
|
||||
def _process_text_image_output(self, eco: EngineCoreOutput) -> None:
|
||||
"""Allow text+image outputs. Text path stays as new_token_ids;
|
||||
image/latents route via pooling_output."""
|
||||
# Preserve text tokens as-is; ensure pooling_output carries image/latents
|
||||
if eco.pooling_output is None:
|
||||
tensor = self._extract_from_multimodal_outputs(
|
||||
eco,
|
||||
keys=(
|
||||
"image",
|
||||
"images",
|
||||
"pixel_values",
|
||||
"pixels",
|
||||
"latent",
|
||||
"latents",
|
||||
"z",
|
||||
),
|
||||
)
|
||||
if tensor is not None:
|
||||
eco.pooling_output = tensor
|
||||
|
||||
def _process_latents_output(self, eco: EngineCoreOutput) -> None:
|
||||
"""Ensure latent tensors are surfaced via pooling_output."""
|
||||
if eco.pooling_output is None:
|
||||
tensor = self._extract_from_multimodal_outputs(eco, keys=("latent", "latents", "z", "posterior"))
|
||||
if tensor is not None:
|
||||
eco.pooling_output = tensor
|
||||
|
||||
def _process_audio_output(self, eco: EngineCoreOutput) -> None:
|
||||
"""Ensure audio tensors are surfaced via pooling_output."""
|
||||
if eco.pooling_output is None:
|
||||
tensor = self._extract_from_multimodal_outputs(
|
||||
eco, keys=("audio", "audios", "wav", "waveform", "audio_pcm", "pcm")
|
||||
)
|
||||
if tensor is not None:
|
||||
eco.pooling_output = tensor
|
||||
|
||||
def _process_text_output(self, eco: EngineCoreOutput) -> None:
|
||||
"""No-op; base processor will detokenize new_token_ids → text."""
|
||||
return
|
||||
|
||||
def _process_pooling_output(self, eco: EngineCoreOutput) -> None:
|
||||
"""Optional sanity checks for pooling tensor."""
|
||||
if eco.pooling_output is None:
|
||||
return
|
||||
if not isinstance(eco.pooling_output, torch.Tensor):
|
||||
# Best-effort: convert to tensor if it's a list/ndarray-like
|
||||
try:
|
||||
eco.pooling_output = torch.as_tensor(eco.pooling_output)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def _extract_from_multimodal_outputs(self, eco: EngineCoreOutput, keys: tuple[str, ...]) -> torch.Tensor | None:
|
||||
mm = getattr(eco, "multimodal_outputs", None)
|
||||
if not isinstance(mm, dict):
|
||||
return None
|
||||
for k in keys:
|
||||
v = mm.get(k)
|
||||
if isinstance(v, torch.Tensor):
|
||||
return v
|
||||
# Try the first tensor in the dict as a fallback
|
||||
for v in mm.values():
|
||||
if isinstance(v, torch.Tensor):
|
||||
return v
|
||||
return None
|
||||
@@ -141,6 +141,7 @@ class AsyncOmni(OmniBase):
|
||||
sequence_parallel_size = kwargs.get("sequence_parallel_size")
|
||||
tensor_parallel_size = kwargs.get("tensor_parallel_size") or 1
|
||||
cfg_parallel_size = kwargs.get("cfg_parallel_size") or 1
|
||||
vae_patch_parallel_size = kwargs.get("vae_patch_parallel_size") or 1
|
||||
use_hsdp = kwargs.get("use_hsdp", False)
|
||||
hsdp_shard_size = kwargs.get("hsdp_shard_size", -1)
|
||||
hsdp_replicate_size = kwargs.get("hsdp_replicate_size", 1)
|
||||
@@ -165,6 +166,7 @@ class AsyncOmni(OmniBase):
|
||||
ulysses_degree=ulysses_degree,
|
||||
ring_degree=ring_degree,
|
||||
cfg_parallel_size=cfg_parallel_size,
|
||||
vae_patch_parallel_size=vae_patch_parallel_size,
|
||||
use_hsdp=use_hsdp,
|
||||
hsdp_shard_size=hsdp_shard_size,
|
||||
hsdp_replicate_size=hsdp_replicate_size,
|
||||
|
||||
@@ -219,6 +219,7 @@ class OmniServeCommand(CLISubcommand):
|
||||
)
|
||||
omni_config_group.add_argument(
|
||||
"--ring",
|
||||
"--ring-degree",
|
||||
dest="ring_degree",
|
||||
type=int,
|
||||
default=None,
|
||||
@@ -335,6 +336,14 @@ class OmniServeCommand(CLISubcommand):
|
||||
help="Number of devices for CFG parallel computation for diffusion models. "
|
||||
"Equivalent to setting DiffusionParallelConfig.cfg_parallel_size.",
|
||||
)
|
||||
omni_config_group.add_argument(
|
||||
"--vae-patch-parallel-size",
|
||||
type=int,
|
||||
default=1,
|
||||
help="VAE Patch Parallelism degree for diffusion models. "
|
||||
"Distributes VAE decode workload across multiple ranks by splitting the latent spatially. "
|
||||
"Equivalent to setting DiffusionParallelConfig.vae_patch_parallel_size.",
|
||||
)
|
||||
|
||||
# Default sampling parameters
|
||||
omni_config_group.add_argument(
|
||||
|
||||
@@ -1135,6 +1135,15 @@ class Omni(OmniBase):
|
||||
final_output_type=stage.final_output_type, # type: ignore[attr-defined]
|
||||
request_output=engine_outputs,
|
||||
)
|
||||
try:
|
||||
if stage.final_output_type == "text" or metrics.log_stats:
|
||||
output_to_yield.metrics = metrics.build_output_metrics(stage_id, req_id)
|
||||
except Exception as e:
|
||||
# Make metrics contract explicit on failure.
|
||||
output_to_yield.metrics = {}
|
||||
logger.exception(
|
||||
f"[{self._name}] Failed to attach output metrics for req {req_id} at stage {stage_id}: {e}",
|
||||
)
|
||||
|
||||
# Record audio generated frames (only when finished)
|
||||
try:
|
||||
|
||||
@@ -418,7 +418,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
|
||||
mm = getattr(ro, "multimodal_output", None) if ro else None
|
||||
if not mm:
|
||||
return None, None
|
||||
key = "audio" if "audio" in mm else ("model_outputs" if "model_outputs" in mm else None)
|
||||
key = "audio" if "audio" in mm else None
|
||||
return mm, key
|
||||
|
||||
def _build_tts_params(self, request: OpenAICreateSpeechRequest) -> dict[str, Any]:
|
||||
|
||||
+50
-14
@@ -39,7 +39,7 @@ class StageRequestStats:
|
||||
final_output_type: str | None = None
|
||||
request_id: str | None = None
|
||||
postprocess_time_ms: float = 0.0
|
||||
diffusion_metrics: dict[str, int] = None
|
||||
diffusion_metrics: dict[str, float] = None
|
||||
audio_generated_frames: int = 0
|
||||
|
||||
@property
|
||||
@@ -293,23 +293,15 @@ class OrchestratorAggregator:
|
||||
return
|
||||
|
||||
# 4. Finished with output: assign text metrics if available
|
||||
output_to_yield.metrics = {}
|
||||
stage_event = next(
|
||||
(evt for evt in reversed(self.stage_events.get(req_id, [])) if evt.stage_id == stage_id),
|
||||
None,
|
||||
)
|
||||
if stage_event is not None and stage_event.final_output_type == "text":
|
||||
output_to_yield.metrics = {
|
||||
"num_tokens_in": stage_event.num_tokens_in,
|
||||
"num_tokens_out": stage_event.num_tokens_out,
|
||||
"stage_id": stage_event.stage_id,
|
||||
"final_output_type": stage_event.final_output_type,
|
||||
}
|
||||
output_to_yield.metrics = self.build_output_metrics(stage_id, req_id)
|
||||
|
||||
# 5. Finished: record audio generated frames
|
||||
self.record_audio_generated_frames(output_to_yield, stage_id, req_id)
|
||||
|
||||
except Exception:
|
||||
if output_to_yield is not None:
|
||||
# Make metrics contract explicit on failure.
|
||||
output_to_yield.metrics = {}
|
||||
logger.exception(
|
||||
"Failed to process metrics for stage %s, req %s",
|
||||
stage_id,
|
||||
@@ -329,12 +321,56 @@ class OrchestratorAggregator:
|
||||
stats.request_id = req_id
|
||||
stats.final_output_type = final_output_type
|
||||
stats.diffusion_metrics = (
|
||||
{k: int(v) for k, v in self.diffusion_metrics.pop(req_id, {}).items()}
|
||||
{k: float(v) for k, v in self.diffusion_metrics.pop(req_id, {}).items()}
|
||||
if req_id in self.diffusion_metrics
|
||||
else None
|
||||
)
|
||||
return stats
|
||||
|
||||
def _get_stage_event(self, stage_id: int, req_id: Any) -> StageRequestStats | None:
|
||||
rid_key = str(req_id)
|
||||
for evt in reversed(self.stage_events.get(rid_key, [])):
|
||||
if evt.stage_id == stage_id:
|
||||
return evt
|
||||
return None
|
||||
|
||||
def _collect_diffusion_metrics(self, req_id: Any) -> dict[str, float]:
|
||||
"""Aggregate diffusion metrics across all stages for a request."""
|
||||
rid_key = str(req_id)
|
||||
merged: dict[str, float] = {}
|
||||
for evt in self.stage_events.get(rid_key, []):
|
||||
if not evt.diffusion_metrics:
|
||||
continue
|
||||
for key, value in evt.diffusion_metrics.items():
|
||||
merged[key] = merged.get(key, 0.0) + float(value)
|
||||
return merged
|
||||
|
||||
def build_output_metrics(self, stage_id: int, req_id: Any) -> dict[str, Any]:
|
||||
stage_event = self._get_stage_event(stage_id, req_id)
|
||||
if stage_event is None:
|
||||
return {}
|
||||
|
||||
merged: dict[str, Any] = {}
|
||||
|
||||
if stage_event.final_output_type == "text":
|
||||
merged.update(
|
||||
{
|
||||
"num_tokens_in": stage_event.num_tokens_in,
|
||||
"num_tokens_out": stage_event.num_tokens_out,
|
||||
}
|
||||
)
|
||||
|
||||
if self.log_stats:
|
||||
diffusion_metrics = self._collect_diffusion_metrics(req_id)
|
||||
if diffusion_metrics:
|
||||
merged.update(diffusion_metrics)
|
||||
|
||||
if merged:
|
||||
merged["stage_id"] = stage_event.stage_id
|
||||
merged["final_output_type"] = stage_event.final_output_type
|
||||
|
||||
return merged
|
||||
|
||||
def on_stage_metrics(
|
||||
self,
|
||||
stage_id: int,
|
||||
|
||||
@@ -556,12 +556,14 @@ class MiMoAudioLLMForConditionalGeneration(nn.Module, SupportsMultiModal, Suppor
|
||||
config.hidden_size,
|
||||
bias=False,
|
||||
return_bias=False,
|
||||
gather_output=True,
|
||||
)
|
||||
self.hidden_states_downcast = ColumnParallelLinear(
|
||||
config.hidden_size,
|
||||
self.local_config.hidden_size,
|
||||
bias=False,
|
||||
return_bias=False,
|
||||
gather_output=True,
|
||||
)
|
||||
|
||||
self.lm_head = ColumnParallelLinear(
|
||||
@@ -569,6 +571,7 @@ class MiMoAudioLLMForConditionalGeneration(nn.Module, SupportsMultiModal, Suppor
|
||||
config.vocab_size,
|
||||
bias=False,
|
||||
return_bias=False,
|
||||
gather_output=True,
|
||||
)
|
||||
|
||||
# Re-encode the sum of multi-layer RVQ embeddings to obtain true Audio Code Embeddings
|
||||
|
||||
@@ -4,7 +4,6 @@ import os
|
||||
from collections.abc import Iterable
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from transformers.utils.hub import cached_file
|
||||
@@ -21,7 +20,8 @@ logger = init_logger(__name__)
|
||||
|
||||
class Qwen3TTSCode2Wav(nn.Module):
|
||||
"""Stage-1 code2wav model for Qwen3-TTS (GenerationModelRunner).
|
||||
Consumes frame-aligned codec tokens from input_ids and decodes waveform via SpeechTokenizer."""
|
||||
Consumes frame-aligned codec tokens from input_ids and decodes waveform
|
||||
via the SpeechTokenizer decoder directly (bypassing HF wrapper overhead)."""
|
||||
|
||||
input_modalities = "audio"
|
||||
|
||||
@@ -34,13 +34,13 @@ class Qwen3TTSCode2Wav(nn.Module):
|
||||
self.has_preprocess = False
|
||||
self.has_postprocess = False
|
||||
self.enable_update_additional_information = True
|
||||
# Generation-only stage (no logits / sampling).
|
||||
self.requires_raw_input_tokens = True
|
||||
|
||||
self._speech_tokenizer: Qwen3TTSTokenizer | None = None
|
||||
self._decoder: nn.Module | None = None
|
||||
self._num_quantizers: int | None = None
|
||||
self._decode_upsample_rate: int | None = None
|
||||
self._output_sample_rate: int | None = None
|
||||
self._total_upsample: int | None = None
|
||||
self._logged_codec_stats = False
|
||||
|
||||
@staticmethod
|
||||
@@ -52,18 +52,15 @@ class Qwen3TTSCode2Wav(nn.Module):
|
||||
return buf.device
|
||||
return torch.device("cpu")
|
||||
|
||||
def _ensure_speech_tokenizer_loaded(self) -> Qwen3TTSTokenizer:
|
||||
if self._speech_tokenizer is not None:
|
||||
return self._speech_tokenizer
|
||||
def _ensure_speech_tokenizer_loaded(self) -> None:
|
||||
if self._decoder is not None:
|
||||
return
|
||||
|
||||
# Locate speech_tokenizer dir from HF cache (or local path).
|
||||
cfg_path = cached_file(self.model_path, "speech_tokenizer/config.json")
|
||||
if cfg_path is None:
|
||||
raise ValueError(f"{self.model_path}/speech_tokenizer/config.json not found")
|
||||
speech_tokenizer_dir = os.path.dirname(cfg_path)
|
||||
|
||||
# Stage-1 only needs decode; skip HF feature extractor to avoid heavy optional deps.
|
||||
# Still require preprocessor_config.json (use cached_file so online runs can fetch it).
|
||||
prep_cfg = cached_file(self.model_path, "speech_tokenizer/preprocessor_config.json")
|
||||
if prep_cfg is None:
|
||||
raise ValueError(
|
||||
@@ -77,12 +74,10 @@ class Qwen3TTSCode2Wav(nn.Module):
|
||||
load_feature_extractor=False,
|
||||
)
|
||||
|
||||
# Align device with vLLM worker, then read back from module.
|
||||
if tok.model is not None:
|
||||
tok.model.to(device=self.vllm_config.device_config.device)
|
||||
tok.device = self._module_device(tok.model)
|
||||
|
||||
# Derive codec group count and rates from tokenizer config.
|
||||
dec_cfg = getattr(tok.model.config, "decoder_config", None)
|
||||
num_q = getattr(dec_cfg, "num_quantizers", None) if dec_cfg is not None else None
|
||||
if num_q is None:
|
||||
@@ -103,11 +98,39 @@ class Qwen3TTSCode2Wav(nn.Module):
|
||||
except Exception as e:
|
||||
raise ValueError(f"Failed to get output sample rate: {e}") from e
|
||||
|
||||
decoder = tok.model.decoder
|
||||
decoder.eval()
|
||||
|
||||
self._speech_tokenizer = tok
|
||||
self._decoder = decoder
|
||||
self._num_quantizers = num_q
|
||||
self._decode_upsample_rate = upsample
|
||||
self._output_sample_rate = out_sr
|
||||
return tok
|
||||
self._total_upsample = int(decoder.total_upsample)
|
||||
|
||||
if hasattr(decoder, "enable_cudagraph"):
|
||||
device = self._module_device(decoder)
|
||||
if device.type == "cuda":
|
||||
try:
|
||||
capture_sizes = None
|
||||
model_cfg = getattr(self.vllm_config, "model_config", None)
|
||||
connector_cfg = getattr(model_cfg, "stage_connector_config", None)
|
||||
extra_cfg = (
|
||||
connector_cfg.get("extra", connector_cfg)
|
||||
if isinstance(connector_cfg, dict)
|
||||
else getattr(connector_cfg, "extra", None)
|
||||
)
|
||||
if isinstance(extra_cfg, dict):
|
||||
chunk_frames = int(extra_cfg.get("codec_chunk_frames") or 0)
|
||||
left_frames = int(extra_cfg.get("codec_left_context_frames") or 0)
|
||||
if chunk_frames > 0 and left_frames >= 0:
|
||||
from .cuda_graph_decoder_wrapper import CUDAGraphDecoderWrapper
|
||||
|
||||
steady_window = left_frames + chunk_frames
|
||||
capture_sizes = sorted({*CUDAGraphDecoderWrapper.DEFAULT_CAPTURE_SIZES, steady_window})
|
||||
decoder.enable_cudagraph(capture_sizes=capture_sizes, device=device)
|
||||
logger.info("Code2Wav decoder CUDA Graph enabled")
|
||||
except Exception:
|
||||
logger.warning("Failed to enable CUDA Graph for Code2Wav decoder", exc_info=True)
|
||||
|
||||
def embed_input_ids(self, input_ids: torch.Tensor, **_: Any) -> torch.Tensor:
|
||||
# This stage ignores token embeddings. Keep a stable dummy embedding for vLLM runner.
|
||||
@@ -155,17 +178,19 @@ class Qwen3TTSCode2Wav(nn.Module):
|
||||
input_ids layout per request: [codec_context_frames, *flat_codes]
|
||||
where flat_codes is codebook-major [q*F].
|
||||
|
||||
When batched, uses forward context ubatch_slices to split the
|
||||
concatenated input_ids and decode via a single batched forward pass.
|
||||
Bypasses the HF Qwen3TTSTokenizer.decode() wrapper and calls the
|
||||
decoder.chunked_decode() directly to avoid GPU->CPU->GPU round-trips.
|
||||
Length management is done here instead of relying on HF's padding=-1
|
||||
sentinel logic.
|
||||
"""
|
||||
self._ensure_speech_tokenizer_loaded()
|
||||
assert self._decoder is not None
|
||||
assert self._num_quantizers is not None
|
||||
assert self._decode_upsample_rate is not None
|
||||
assert self._output_sample_rate is not None
|
||||
assert self._total_upsample is not None
|
||||
|
||||
tok = self._speech_tokenizer
|
||||
decoder = self._decoder
|
||||
q = int(self._num_quantizers)
|
||||
upsample = int(self._decode_upsample_rate)
|
||||
upsample = int(self._total_upsample)
|
||||
sr_val = int(self._output_sample_rate)
|
||||
sr_tensor = torch.tensor(sr_val, dtype=torch.int32)
|
||||
empty = torch.zeros((0,), dtype=torch.float32)
|
||||
@@ -179,13 +204,9 @@ class Qwen3TTSCode2Wav(nn.Module):
|
||||
ids = input_ids.reshape(-1).to(dtype=torch.long)
|
||||
request_ids_list = self._split_request_ids(ids, kwargs.get("seq_token_counts"))
|
||||
|
||||
# Parse each request: extract left_context_size, validate, reshape codes.
|
||||
# input_ids layout per request: [*flat_codes]
|
||||
# where flat_codes is codebook-major [q*F].
|
||||
parsed = [] # (left_context_size, actual_frames)
|
||||
valid_codes = []
|
||||
valid_indices = []
|
||||
# Get left_context_size from runtime_additional_information (passed via kwargs)
|
||||
parsed: list[tuple[int, int]] = []
|
||||
valid_codes_qf: list[torch.Tensor] = []
|
||||
valid_indices: list[int] = []
|
||||
left_context_size = [0] * len(request_ids_list)
|
||||
if runtime_additional_information is not None:
|
||||
for i, info in enumerate(runtime_additional_information):
|
||||
@@ -193,8 +214,6 @@ class Qwen3TTSCode2Wav(nn.Module):
|
||||
break
|
||||
if "left_context_size" in info:
|
||||
left_context_size[i] = info["left_context_size"]
|
||||
else:
|
||||
logger.debug("No additional_information provided to code2wav stage.")
|
||||
for i, req_ids in enumerate(request_ids_list):
|
||||
if req_ids.numel() < 1:
|
||||
parsed.append((0, 0))
|
||||
@@ -202,7 +221,6 @@ class Qwen3TTSCode2Wav(nn.Module):
|
||||
ctx_frames = left_context_size[i]
|
||||
flat = req_ids
|
||||
n = flat.numel()
|
||||
# Warmup / dummy_run: not divisible by num_quantizers.
|
||||
if n == 0 or n % q != 0:
|
||||
if n > 0:
|
||||
logger.warning(
|
||||
@@ -214,14 +232,14 @@ class Qwen3TTSCode2Wav(nn.Module):
|
||||
parsed.append((0, 0))
|
||||
continue
|
||||
frames = n // q
|
||||
# Reshape codebook-major flat [q*F] -> [q, F] -> [F, q] for SpeechTokenizer.
|
||||
codes_fq = flat.reshape(q, frames).transpose(0, 1).contiguous()
|
||||
# [q*F] -> [Q, F] for direct decoder call (decoder expects [B, Q, F])
|
||||
codes_qf = flat.reshape(q, frames)
|
||||
parsed.append((ctx_frames, frames))
|
||||
valid_codes.append({"audio_codes": codes_fq})
|
||||
valid_codes_qf.append(codes_qf)
|
||||
valid_indices.append(i)
|
||||
|
||||
num_req = len(request_ids_list)
|
||||
if not valid_codes:
|
||||
if not valid_codes_qf:
|
||||
return OmniOutput(
|
||||
text_hidden_states=None,
|
||||
multimodal_outputs={
|
||||
@@ -233,50 +251,55 @@ class Qwen3TTSCode2Wav(nn.Module):
|
||||
if not self._logged_codec_stats:
|
||||
self._logged_codec_stats = True
|
||||
try:
|
||||
c = valid_codes[0]["audio_codes"]
|
||||
c = valid_codes_qf[0]
|
||||
logger.info(
|
||||
"Code2Wav codec: frames=%d q=%d uniq=%d range=[%d,%d] head=%s batch=%d",
|
||||
c.shape[0],
|
||||
"Code2Wav codec: frames=%d q=%d uniq=%d range=[%d,%d] batch=%d",
|
||||
c.shape[1],
|
||||
q,
|
||||
int(torch.unique(c).numel()),
|
||||
int(c.min().item()),
|
||||
int(c.max().item()),
|
||||
c[: min(2, c.shape[0]), : min(8, q)].cpu().tolist(),
|
||||
len(valid_codes),
|
||||
len(valid_codes_qf),
|
||||
)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Batched decode: single forward pass through SpeechTokenizer.
|
||||
wavs, _ = tok.decode(valid_codes)
|
||||
if len(wavs) != len(valid_codes):
|
||||
raise RuntimeError(f"Code2Wav returned {len(wavs)} waveforms for {len(valid_codes)} requests")
|
||||
# Decode directly via decoder.chunked_decode(), staying entirely on GPU.
|
||||
# For single request: no padding needed, fast path.
|
||||
# For multiple requests: decode each individually to avoid padding overhead.
|
||||
wav_tensors: list[torch.Tensor] = []
|
||||
if len(valid_codes_qf) == 1:
|
||||
codes_bqf = valid_codes_qf[0].unsqueeze(0) # [1, Q, F]
|
||||
wav = decoder.chunked_decode(codes_bqf) # [1, 1, wav_len]
|
||||
wav_tensors.append(wav.squeeze(0).squeeze(0)) # [wav_len]
|
||||
else:
|
||||
for codes_qf in valid_codes_qf:
|
||||
codes_bqf = codes_qf.unsqueeze(0) # [1, Q, F]
|
||||
wav = decoder.chunked_decode(codes_bqf)
|
||||
wav_tensors.append(wav.squeeze(0).squeeze(0))
|
||||
|
||||
# Build per-request outputs, trimming padding and left-context.
|
||||
audios = [empty] * num_req
|
||||
audios: list[torch.Tensor] = [empty] * num_req
|
||||
srs = [sr_tensor] * num_req
|
||||
|
||||
for j, idx in enumerate(valid_indices):
|
||||
ctx_frames, actual_frames = parsed[idx]
|
||||
audio_np = wavs[j].astype(np.float32, copy=False)
|
||||
# Trim decoder padding (output may be longer due to batch padding).
|
||||
wav = wav_tensors[j]
|
||||
expected_len = actual_frames * upsample
|
||||
if audio_np.shape[0] > expected_len:
|
||||
audio_np = audio_np[:expected_len]
|
||||
# Trim left-context waveform samples (streaming sliding window).
|
||||
if wav.shape[0] > expected_len:
|
||||
wav = wav[:expected_len]
|
||||
if ctx_frames > 0:
|
||||
cut = ctx_frames * upsample
|
||||
if cut < audio_np.shape[0]:
|
||||
audio_np = audio_np[cut:]
|
||||
if cut < wav.shape[0]:
|
||||
wav = wav[cut:]
|
||||
else:
|
||||
logger.warning(
|
||||
"Context trim %d >= decoded length %d; returning empty audio.",
|
||||
cut,
|
||||
audio_np.shape[0],
|
||||
wav.shape[0],
|
||||
)
|
||||
continue
|
||||
if audio_np.shape[0] > 0:
|
||||
audios[idx] = torch.from_numpy(audio_np).to(dtype=torch.float32).reshape(-1)
|
||||
if wav.shape[0] > 0:
|
||||
audios[idx] = wav.to(dtype=torch.float32).reshape(-1)
|
||||
|
||||
return OmniOutput(
|
||||
text_hidden_states=None,
|
||||
|
||||
@@ -1,228 +1,217 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Iterable
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from vllm.config import VllmConfig
|
||||
from vllm.config.vllm import set_current_vllm_config
|
||||
from vllm.forward_context import set_forward_context
|
||||
from vllm.logger import init_logger
|
||||
from vllm.model_executor.layers.layernorm import RMSNorm
|
||||
from vllm.model_executor.layers.linear import (
|
||||
MergedColumnParallelLinear,
|
||||
QKVParallelLinear,
|
||||
RowParallelLinear,
|
||||
)
|
||||
from vllm.model_executor.layers.rotary_embedding import get_rope
|
||||
from vllm.model_executor.model_loader.weight_utils import (
|
||||
default_weight_loader,
|
||||
maybe_remap_kv_scale_name,
|
||||
)
|
||||
from vllm.model_executor.models.qwen3 import Qwen3DecoderLayer
|
||||
from vllm.model_executor.models.utils import is_pp_missing_parameter
|
||||
from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheGroupSpec, KVCacheSpec, KVCacheTensor
|
||||
from vllm.v1.worker.gpu import attn_utils
|
||||
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
from .configuration_qwen3_tts import Qwen3TTSTalkerCodePredictorConfig, Qwen3TTSTalkerConfig
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
class _LocalPredictorKVCache:
|
||||
"""Minimal local KV cache + attention metadata for running
|
||||
code_predictor inside one worker (independent of engine KV)."""
|
||||
|
||||
# ===================================================================
|
||||
# Standalone Code Predictor Layers (no vLLM paged attention)
|
||||
# ===================================================================
|
||||
#
|
||||
# These replace vLLM's Qwen3DecoderLayer for the code predictor.
|
||||
# Input is batch-major [B, seq_len, H], attention uses F.scaled_dot_product_attention.
|
||||
# Weight names match the checkpoint (self_attn.qkv_proj, mlp.gate_up_proj, etc.)
|
||||
# so load_weights works unchanged.
|
||||
|
||||
|
||||
class _CodePredictorAttention(nn.Module):
|
||||
"""Standalone multi-head attention for code predictor.
|
||||
|
||||
Uses F.scaled_dot_product_attention (SDPA) instead of vLLM's paged Attention.
|
||||
Supports fused QKV, RoPE, q/k normalization, and native GQA via enable_gqa.
|
||||
Input: [B, seq_len, hidden_size], output: [B, seq_len, hidden_size].
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
config: Qwen3TTSTalkerCodePredictorConfig,
|
||||
*,
|
||||
vllm_config: VllmConfig,
|
||||
max_seq_len: int,
|
||||
max_batch_size: int,
|
||||
device: torch.device,
|
||||
prefix: str = "",
|
||||
) -> None:
|
||||
self.vllm_config = vllm_config
|
||||
self.device = device
|
||||
|
||||
# Collect attention layers registered in this vllm_config.
|
||||
kv_cache_spec_by_layer = attn_utils.get_kv_cache_spec(vllm_config)
|
||||
if not kv_cache_spec_by_layer:
|
||||
raise RuntimeError("Local predictor KVCache requires vLLM Attention layers to be registered.")
|
||||
|
||||
# We only need enough blocks for a tiny per-frame sequence (<= max_seq_len).
|
||||
any_spec = next(iter(kv_cache_spec_by_layer.values()))
|
||||
block_size = int(any_spec.block_size)
|
||||
blocks_per_seq = (int(max_seq_len) + block_size - 1) // block_size
|
||||
num_blocks = max(1, int(max_batch_size) * int(blocks_per_seq))
|
||||
|
||||
# Allocate per-layer KV caches (small, independent).
|
||||
kv_cache_tensors: list[KVCacheTensor] = []
|
||||
for layer_name, spec in kv_cache_spec_by_layer.items():
|
||||
kv_cache_tensors.append(KVCacheTensor(size=int(spec.page_size_bytes) * num_blocks, shared_by=[layer_name]))
|
||||
|
||||
merged_spec: KVCacheSpec = KVCacheSpec.merge(list(kv_cache_spec_by_layer.values()))
|
||||
self.kv_cache_config = KVCacheConfig(
|
||||
num_blocks=num_blocks,
|
||||
kv_cache_tensors=kv_cache_tensors,
|
||||
kv_cache_groups=[
|
||||
KVCacheGroupSpec(layer_names=list(kv_cache_spec_by_layer.keys()), kv_cache_spec=merged_spec)
|
||||
],
|
||||
super().__init__()
|
||||
self.hidden_size = config.hidden_size
|
||||
self.num_heads = config.num_attention_heads
|
||||
self.num_kv_heads = config.num_key_value_heads
|
||||
self.head_dim = getattr(
|
||||
config,
|
||||
"head_dim",
|
||||
config.hidden_size // config.num_attention_heads,
|
||||
)
|
||||
self.q_size = self.num_heads * self.head_dim
|
||||
self.kv_size = self.num_kv_heads * self.head_dim
|
||||
self.scaling = self.head_dim**-0.5
|
||||
self._use_gqa = self.num_kv_heads != self.num_heads
|
||||
|
||||
# Init backend + bind KV cache tensors to attention modules.
|
||||
self.attn_backends, self.attn_metadata_builders = attn_utils.init_attn_backend(
|
||||
self.kv_cache_config, vllm_config, device
|
||||
self.qkv_proj = QKVParallelLinear(
|
||||
hidden_size=self.hidden_size,
|
||||
head_size=self.head_dim,
|
||||
total_num_heads=self.num_heads,
|
||||
total_num_kv_heads=self.num_kv_heads,
|
||||
bias=getattr(config, "attention_bias", False),
|
||||
prefix=f"{prefix}.qkv_proj",
|
||||
disable_tp=True,
|
||||
)
|
||||
self.runner_kv_caches: list[torch.Tensor] = []
|
||||
attn_utils.init_kv_cache(
|
||||
self.runner_kv_caches,
|
||||
vllm_config.compilation_config.static_forward_context,
|
||||
self.kv_cache_config,
|
||||
self.attn_backends,
|
||||
device,
|
||||
self.o_proj = RowParallelLinear(
|
||||
input_size=self.num_heads * self.head_dim,
|
||||
output_size=self.hidden_size,
|
||||
bias=False,
|
||||
prefix=f"{prefix}.o_proj",
|
||||
disable_tp=True,
|
||||
)
|
||||
self.rotary_emb = get_rope(
|
||||
self.head_dim,
|
||||
max_position=config.max_position_embeddings,
|
||||
rope_parameters=getattr(config, "rope_parameters", None),
|
||||
dual_chunk_attention_config=None,
|
||||
)
|
||||
self.q_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)
|
||||
self.k_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)
|
||||
|
||||
# Precompute a fixed block table mapping for the maximum batch.
|
||||
self.block_size = block_size
|
||||
self.blocks_per_seq = blocks_per_seq
|
||||
self.max_batch_size = int(max_batch_size)
|
||||
|
||||
bt = torch.full((self.max_batch_size, self.blocks_per_seq), -1, dtype=torch.int32, device=device)
|
||||
for i in range(self.max_batch_size):
|
||||
for j in range(self.blocks_per_seq):
|
||||
bt[i, j] = i * self.blocks_per_seq + j
|
||||
self._block_table = bt
|
||||
|
||||
def build_attn_metadata(
|
||||
def forward(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
position_ids: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
bsz, seq_len, _ = hidden_states.shape
|
||||
|
||||
qkv, _ = self.qkv_proj(hidden_states.reshape(bsz * seq_len, -1))
|
||||
q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
|
||||
|
||||
q = self.q_norm(q.view(-1, self.num_heads, self.head_dim)).view(q.shape)
|
||||
k = self.k_norm(k.view(-1, self.num_kv_heads, self.head_dim)).view(k.shape)
|
||||
|
||||
q, k = self.rotary_emb(position_ids, q, k)
|
||||
|
||||
q = q.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
|
||||
k = k.view(bsz, seq_len, self.num_kv_heads, self.head_dim).transpose(1, 2)
|
||||
v = v.view(bsz, seq_len, self.num_kv_heads, self.head_dim).transpose(1, 2)
|
||||
|
||||
attn_out = F.scaled_dot_product_attention(
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
scale=self.scaling,
|
||||
is_causal=True,
|
||||
enable_gqa=self._use_gqa,
|
||||
)
|
||||
|
||||
attn_out = attn_out.transpose(1, 2).reshape(bsz * seq_len, -1)
|
||||
output, _ = self.o_proj(attn_out)
|
||||
return output.view(bsz, seq_len, -1)
|
||||
|
||||
|
||||
class _CodePredictorMLP(nn.Module):
|
||||
"""SiLU-gated MLP for code predictor, matching Qwen3MLP structure."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
config: Qwen3TTSTalkerCodePredictorConfig,
|
||||
*,
|
||||
num_reqs: int,
|
||||
query_lens: torch.Tensor, # (num_reqs,) int32 on cpu
|
||||
seq_lens: torch.Tensor, # (num_reqs,) int32 on cpu
|
||||
) -> tuple[dict[str, Any], torch.Tensor, dict[str, torch.Tensor]]:
|
||||
"""Build attention metadata, positions, and slot_mapping dict.
|
||||
prefix: str = "",
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.gate_up_proj = MergedColumnParallelLinear(
|
||||
input_size=config.hidden_size,
|
||||
output_sizes=[config.intermediate_size] * 2,
|
||||
bias=False,
|
||||
prefix=f"{prefix}.gate_up_proj",
|
||||
disable_tp=True,
|
||||
)
|
||||
self.down_proj = RowParallelLinear(
|
||||
input_size=config.intermediate_size,
|
||||
output_size=config.hidden_size,
|
||||
bias=False,
|
||||
prefix=f"{prefix}.down_proj",
|
||||
disable_tp=True,
|
||||
)
|
||||
|
||||
Returns:
|
||||
(attn_metadata, positions, slot_mappings_by_layer)
|
||||
- attn_metadata: per-layer attention metadata for attn backends.
|
||||
- positions: (num_tokens,) position IDs on device.
|
||||
- slot_mappings_by_layer: {layer_name: slot_mapping_tensor} for
|
||||
set_forward_context so that unified_kv_cache_update can write
|
||||
the KV cache correctly.
|
||||
"""
|
||||
num_reqs = int(num_reqs)
|
||||
if num_reqs <= 0:
|
||||
return {}, torch.empty((0,), dtype=torch.int64, device=self.device), {}
|
||||
if num_reqs > self.max_batch_size:
|
||||
raise ValueError(f"num_reqs={num_reqs} exceeds local predictor max_batch_size={self.max_batch_size}")
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
gate_up, _ = self.gate_up_proj(x)
|
||||
gate, up = gate_up.chunk(2, dim=-1)
|
||||
x = F.silu(gate) * up
|
||||
x, _ = self.down_proj(x)
|
||||
return x
|
||||
|
||||
query_lens_i32 = query_lens.to(dtype=torch.int32, device="cpu")
|
||||
seq_lens_i32 = seq_lens.to(dtype=torch.int32, device="cpu")
|
||||
|
||||
# query_start_loc: prefix sums of query_lens.
|
||||
qsl = torch.zeros((num_reqs + 1,), dtype=torch.int32, device="cpu")
|
||||
qsl[1:] = torch.cumsum(query_lens_i32, dim=0)
|
||||
num_tokens = int(qsl[-1].item())
|
||||
if num_tokens <= 0:
|
||||
return {}, torch.empty((0,), dtype=torch.int64, device=self.device), {}
|
||||
class _CodePredictorDecoderLayer(nn.Module):
|
||||
"""Transformer decoder layer for code predictor (SDPA, no KV cache)."""
|
||||
|
||||
# positions: for each request i, emit positions [seq_len-query_len .. seq_len-1]
|
||||
pos_list: list[torch.Tensor] = []
|
||||
for i in range(num_reqs):
|
||||
ql = int(query_lens_i32[i].item())
|
||||
sl = int(seq_lens_i32[i].item())
|
||||
start = sl - ql
|
||||
pos_list.append(torch.arange(start, sl, dtype=torch.int64))
|
||||
positions_cpu = torch.cat(pos_list, dim=0)
|
||||
def __init__(
|
||||
self,
|
||||
config: Qwen3TTSTalkerCodePredictorConfig,
|
||||
*,
|
||||
prefix: str = "",
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.self_attn = _CodePredictorAttention(config, prefix=f"{prefix}.self_attn")
|
||||
self.mlp = _CodePredictorMLP(config, prefix=f"{prefix}.mlp")
|
||||
self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
|
||||
self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
|
||||
|
||||
# slot_mapping: map each query token to a physical slot in the paged KV cache.
|
||||
# We allocate per-request contiguous blocks; slot = base + position.
|
||||
slot_mapping = torch.empty((num_tokens,), dtype=torch.int64, device="cpu")
|
||||
cursor = 0
|
||||
for i in range(num_reqs):
|
||||
ql = int(query_lens_i32[i].item())
|
||||
sl = int(seq_lens_i32[i].item())
|
||||
start = sl - ql
|
||||
for p in range(start, sl):
|
||||
block_idx = p // self.block_size
|
||||
offset = p % self.block_size
|
||||
block_id = int(self._block_table[i, block_idx].item())
|
||||
slot_mapping[cursor] = block_id * self.block_size + offset
|
||||
cursor += 1
|
||||
def forward(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
position_ids: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
residual = hidden_states
|
||||
hidden_states = self.input_layernorm(hidden_states)
|
||||
hidden_states = self.self_attn(hidden_states, position_ids)
|
||||
hidden_states = residual + hidden_states
|
||||
|
||||
max_seq_len = int(seq_lens_i32[:num_reqs].max().item())
|
||||
max_query_len = int(query_lens_i32[:num_reqs].max().item())
|
||||
query_start_loc_gpu = qsl.to(device=self.device)
|
||||
seq_lens_gpu = seq_lens_i32.to(device=self.device)
|
||||
block_table = self._block_table[:num_reqs].contiguous()
|
||||
slot_mapping_gpu = slot_mapping.to(device=self.device)
|
||||
residual = hidden_states
|
||||
hidden_states = self.post_attention_layernorm(hidden_states)
|
||||
hidden_states = self.mlp(hidden_states)
|
||||
hidden_states = residual + hidden_states
|
||||
return hidden_states
|
||||
|
||||
if current_omni_platform.is_npu():
|
||||
from vllm_ascend.worker.v2 import attn_utils as attn_utils_npu
|
||||
|
||||
max_query_len = int(query_lens_i32[:num_reqs].max().item())
|
||||
slot_mappings_tensor = slot_mapping_gpu.unsqueeze(0)
|
||||
attn_metadata = attn_utils_npu.build_attn_metadata(
|
||||
attn_metadata_builders=self.attn_metadata_builders,
|
||||
num_reqs=num_reqs,
|
||||
num_tokens=num_tokens,
|
||||
query_start_loc_gpu=query_start_loc_gpu,
|
||||
query_start_loc_cpu=qsl,
|
||||
max_query_len=max_query_len,
|
||||
seq_lens=seq_lens_gpu,
|
||||
max_seq_len=max_seq_len,
|
||||
block_tables=[block_table],
|
||||
slot_mappings=slot_mappings_tensor,
|
||||
kv_cache_config=self.kv_cache_config,
|
||||
)
|
||||
else:
|
||||
attn_metadata = attn_utils.build_attn_metadata(
|
||||
self.attn_metadata_builders,
|
||||
num_reqs=num_reqs,
|
||||
num_tokens=num_tokens,
|
||||
query_start_loc_gpu=query_start_loc_gpu,
|
||||
query_start_loc_cpu=qsl,
|
||||
max_query_len=max_query_len,
|
||||
seq_lens=seq_lens_gpu,
|
||||
max_seq_len=max_seq_len,
|
||||
block_tables=[block_table],
|
||||
slot_mappings=[slot_mapping_gpu],
|
||||
kv_cache_config=self.kv_cache_config,
|
||||
)
|
||||
|
||||
# Build slot_mappings_by_layer for set_forward_context.
|
||||
# Fix for vllm 0.15.0
|
||||
slot_mappings_by_layer: dict[str, torch.Tensor] = {}
|
||||
for kv_cache_group in self.kv_cache_config.kv_cache_groups:
|
||||
for layer_name in kv_cache_group.layer_names:
|
||||
slot_mappings_by_layer[layer_name] = slot_mapping_gpu
|
||||
|
||||
return attn_metadata, positions_cpu.to(device=self.device), slot_mappings_by_layer
|
||||
# ===================================================================
|
||||
# Code Predictor Transformer Model
|
||||
# ===================================================================
|
||||
|
||||
|
||||
class Qwen3TTSTalkerCodePredictorModelVLLM(nn.Module):
|
||||
"""Transformer model for the code predictor (re-prefill, no KV cache)."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
config: Qwen3TTSTalkerCodePredictorConfig,
|
||||
*,
|
||||
talker_hidden_size: int | None = None,
|
||||
cache_config=None,
|
||||
quant_config=None,
|
||||
prefix: str = "",
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.config = config
|
||||
self.quant_config = quant_config
|
||||
|
||||
self.layers = nn.ModuleList(
|
||||
[
|
||||
Qwen3DecoderLayer(
|
||||
config, cache_config=cache_config, quant_config=quant_config, prefix=f"{prefix}.layers.{i}"
|
||||
)
|
||||
for i in range(config.num_hidden_layers)
|
||||
]
|
||||
[_CodePredictorDecoderLayer(config, prefix=f"{prefix}.layers.{i}") for i in range(config.num_hidden_layers)]
|
||||
)
|
||||
self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
|
||||
|
||||
# Official code_predictor uses one embedding table per residual group.
|
||||
# Some Qwen3-TTS checkpoints store codec embeddings in the talker hidden
|
||||
# space, even when `code_predictor_config.hidden_size` is smaller.
|
||||
# We keep the embedding dim aligned with the checkpoint and project down
|
||||
# via `small_to_mtp_projection` in the wrapper module.
|
||||
# Codec embeddings: one per residual group. Stored in talker hidden dim
|
||||
# (some checkpoints use talker_hidden_size != code_predictor hidden_size).
|
||||
emb_dim = int(talker_hidden_size) if talker_hidden_size is not None else int(config.hidden_size)
|
||||
self.codec_embedding = nn.ModuleList(
|
||||
[nn.Embedding(config.vocab_size, emb_dim) for _ in range(config.num_code_groups - 1)]
|
||||
@@ -231,18 +220,18 @@ class Qwen3TTSTalkerCodePredictorModelVLLM(nn.Module):
|
||||
def get_input_embeddings(self) -> nn.ModuleList:
|
||||
return self.codec_embedding
|
||||
|
||||
def forward(self, positions: torch.Tensor, inputs_embeds: torch.Tensor) -> torch.Tensor:
|
||||
# Token-major: [num_tokens, hidden]
|
||||
def forward(
|
||||
self,
|
||||
inputs_embeds: torch.Tensor,
|
||||
position_ids: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
hidden_states = inputs_embeds
|
||||
residual = None
|
||||
for layer in self.layers:
|
||||
hidden_states, residual = layer(positions, hidden_states, residual)
|
||||
hidden_states, _ = self.norm(hidden_states, residual)
|
||||
hidden_states = layer(hidden_states, position_ids)
|
||||
hidden_states = self.norm(hidden_states)
|
||||
return hidden_states
|
||||
|
||||
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
|
||||
# Match vLLM Qwen2/Qwen3 packing conventions: q_proj/k_proj/v_proj -> qkv_proj,
|
||||
# gate_proj/up_proj -> gate_up_proj.
|
||||
stacked_params_mapping = [
|
||||
("qkv_proj", "q_proj", "q"),
|
||||
("qkv_proj", "k_proj", "k"),
|
||||
@@ -255,13 +244,6 @@ class Qwen3TTSTalkerCodePredictorModelVLLM(nn.Module):
|
||||
for name, loaded_weight in weights:
|
||||
if "rotary_emb.inv_freq" in name:
|
||||
continue
|
||||
if self.quant_config is not None and (scale_name := self.quant_config.get_cache_scale(name)):
|
||||
param = params_dict[scale_name]
|
||||
weight_loader = getattr(param, "weight_loader", default_weight_loader)
|
||||
loaded_weight = loaded_weight if loaded_weight.dim() == 0 else loaded_weight[0]
|
||||
weight_loader(param, loaded_weight)
|
||||
loaded_params.add(scale_name)
|
||||
continue
|
||||
|
||||
for param_name, weight_name, shard_id in stacked_params_mapping:
|
||||
if weight_name not in name:
|
||||
@@ -302,8 +284,31 @@ class Qwen3TTSTalkerCodePredictorModelVLLM(nn.Module):
|
||||
return loaded_params
|
||||
|
||||
|
||||
# ===================================================================
|
||||
# Code Predictor Wrapper (optimized re-prefill + torch.compile)
|
||||
# ===================================================================
|
||||
|
||||
|
||||
class Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM(nn.Module):
|
||||
"""vLLM-native code_predictor used by the AR talker (residual codebooks)."""
|
||||
"""vLLM-native code_predictor for the AR talker (residual codebooks).
|
||||
|
||||
Re-prefill approach: each AR step forwards the full growing sequence
|
||||
through the 5-layer transformer. No KV cache needed. This trades
|
||||
~O(T^2) extra attention FLOPs (negligible for T=16, 5 layers) for
|
||||
zero KV cache management overhead and a simpler execution model.
|
||||
|
||||
Optimizations over baseline:
|
||||
1. torch.compile on model forward -- fuses 60+ small kernel launches per step
|
||||
into fewer fused kernels (4x speedup on model_fwd, ~75% of total time).
|
||||
2. Pre-allocated embedding buffer [B, max_seq, H] -- no torch.cat per step.
|
||||
3. Projection caching -- each token projected once and cached, avoids O(T^2)
|
||||
redundant projections.
|
||||
4. Pre-allocated position_ids [max_seq] -- no torch.arange per step.
|
||||
5. Inline sampling -- no custom op / forward_context overhead.
|
||||
6. No context managers in forward().
|
||||
7. Cached module references -- bypass nn.Module.__call__ and ModuleList indexing.
|
||||
8. Pre-allocated output tensor.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
@@ -318,16 +323,12 @@ class Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM(nn.Module):
|
||||
self.config = config
|
||||
self.talker_config = talker_config
|
||||
|
||||
# Keep module/weight names aligned with official checkpoint (talker.code_predictor.model.*).
|
||||
self.model = Qwen3TTSTalkerCodePredictorModelVLLM(
|
||||
config,
|
||||
talker_hidden_size=int(talker_config.hidden_size),
|
||||
cache_config=vllm_config.cache_config,
|
||||
quant_config=vllm_config.quant_config,
|
||||
prefix=f"{prefix}.model",
|
||||
)
|
||||
|
||||
# One head per residual group.
|
||||
self.lm_head = nn.ModuleList(
|
||||
[nn.Linear(config.hidden_size, config.vocab_size, bias=False) for _ in range(config.num_code_groups - 1)]
|
||||
)
|
||||
@@ -337,14 +338,21 @@ class Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM(nn.Module):
|
||||
else:
|
||||
self.small_to_mtp_projection = nn.Identity()
|
||||
|
||||
self._kv_cache: _LocalPredictorKVCache | None = None
|
||||
self._num_groups = int(config.num_code_groups)
|
||||
self._talker_hidden = int(talker_config.hidden_size)
|
||||
self._cp_hidden = int(config.hidden_size)
|
||||
|
||||
# Pre-allocated buffers (lazily initialized on first forward).
|
||||
self._proj_buf: torch.Tensor | None = None
|
||||
self._pos_ids: torch.Tensor | None = None
|
||||
|
||||
# torch.compile: fuse small kernels in the 5-layer transformer.
|
||||
self._compiled_model_fwd: object | None = None
|
||||
|
||||
def get_input_embeddings(self) -> nn.ModuleList:
|
||||
return self.model.get_input_embeddings()
|
||||
|
||||
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
|
||||
# Ensure all vLLM custom layers consult the predictor vllm_config
|
||||
# (esp. for Attention static_forward_context).
|
||||
with set_current_vllm_config(self._vllm_config):
|
||||
loaded: set[str] = set()
|
||||
model_weights: list[tuple[str, torch.Tensor]] = []
|
||||
@@ -366,95 +374,52 @@ class Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM(nn.Module):
|
||||
loaded.add(name)
|
||||
return loaded
|
||||
|
||||
def _maybe_init_kv_cache(self, device: torch.device) -> None:
|
||||
if self._kv_cache is not None:
|
||||
# ------------------------------------------------------------------
|
||||
# Pre-allocated buffer management
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _ensure_buffers(self, bsz: int, device: torch.device, dtype: torch.dtype) -> None:
|
||||
max_seq = self._num_groups + 1
|
||||
if (
|
||||
self._proj_buf is not None
|
||||
and self._proj_buf.shape[0] >= bsz
|
||||
and self._proj_buf.device == device
|
||||
and self._proj_buf.dtype == dtype
|
||||
):
|
||||
return
|
||||
with set_current_vllm_config(self._vllm_config):
|
||||
max_seq_len = int(getattr(self.config, "num_code_groups", 16) or 16)
|
||||
# Upper bound on batch size: vLLM scheduler max_num_seqs (fallback 8).
|
||||
max_batch = int(getattr(self._vllm_config.scheduler_config, "max_num_seqs", 8) or 8)
|
||||
max_batch = max(1, max_batch)
|
||||
self._kv_cache = _LocalPredictorKVCache(
|
||||
vllm_config=self._vllm_config,
|
||||
max_seq_len=max_seq_len,
|
||||
max_batch_size=max_batch,
|
||||
device=device,
|
||||
)
|
||||
|
||||
@torch.inference_mode()
|
||||
def reset_cache(self) -> None:
|
||||
# We reuse a fixed kv cache buffer and overwrite starting at slot 0.
|
||||
# No action required here (seq_lens controls what is read).
|
||||
return
|
||||
|
||||
@torch.inference_mode()
|
||||
def prefill_logits(self, inputs_embeds: torch.Tensor) -> torch.Tensor:
|
||||
"""Prefill with 2 tokens: [past_hidden, layer0_embed]. Returns logits for residual group 0."""
|
||||
self._maybe_init_kv_cache(inputs_embeds.device)
|
||||
assert self._kv_cache is not None
|
||||
|
||||
bsz = int(inputs_embeds.shape[0])
|
||||
qlen = 2
|
||||
# Flatten to token-major.
|
||||
model_dtype = next(self.model.parameters()).dtype
|
||||
hs = inputs_embeds.to(dtype=model_dtype).reshape(bsz * qlen, -1)
|
||||
hs = self.small_to_mtp_projection(hs)
|
||||
|
||||
query_lens = torch.full((bsz,), qlen, dtype=torch.int32)
|
||||
seq_lens = query_lens.clone()
|
||||
attn_metadata, positions, slot_mappings = self._kv_cache.build_attn_metadata(
|
||||
num_reqs=bsz, query_lens=query_lens, seq_lens=seq_lens
|
||||
self._proj_buf = torch.zeros(
|
||||
bsz,
|
||||
max_seq,
|
||||
self._cp_hidden,
|
||||
dtype=dtype,
|
||||
device=device,
|
||||
)
|
||||
self._pos_ids = torch.arange(
|
||||
max_seq,
|
||||
dtype=torch.long,
|
||||
device=device,
|
||||
)
|
||||
|
||||
with (
|
||||
set_current_vllm_config(self._vllm_config),
|
||||
set_forward_context(
|
||||
attn_metadata,
|
||||
self._vllm_config,
|
||||
num_tokens=int(hs.shape[0]),
|
||||
slot_mapping=slot_mappings,
|
||||
),
|
||||
):
|
||||
out = self.model(positions=positions, inputs_embeds=hs)
|
||||
def _setup_compile(self) -> None:
|
||||
"""Lazily set up torch.compiled model forward for kernel fusion.
|
||||
|
||||
# Gather last token per request.
|
||||
last_idx = torch.arange(qlen - 1, bsz * qlen, step=qlen, device=out.device, dtype=torch.long)
|
||||
last_h = out.index_select(0, last_idx)
|
||||
logits = self.lm_head[0](last_h)
|
||||
return logits
|
||||
|
||||
@torch.inference_mode()
|
||||
def decode_logits(self, input_ids: torch.Tensor, *, generation_step: int, past_seq_len: int) -> torch.Tensor:
|
||||
"""Decode one new token for residual group `generation_step` (1..Q-1)."""
|
||||
self._maybe_init_kv_cache(input_ids.device)
|
||||
assert self._kv_cache is not None
|
||||
bsz = int(input_ids.shape[0])
|
||||
if generation_step <= 0:
|
||||
raise ValueError("generation_step must be >= 1 for decode_logits")
|
||||
|
||||
embed_idx = generation_step - 1
|
||||
hs = self.model.get_input_embeddings()[embed_idx](input_ids.to(dtype=torch.long).reshape(bsz, 1))
|
||||
hs = self.small_to_mtp_projection(hs.reshape(bsz, -1))
|
||||
|
||||
query_lens = torch.ones((bsz,), dtype=torch.int32)
|
||||
seq_lens = torch.full((bsz,), int(past_seq_len) + 1, dtype=torch.int32)
|
||||
attn_metadata, positions, slot_mappings = self._kv_cache.build_attn_metadata(
|
||||
num_reqs=bsz, query_lens=query_lens, seq_lens=seq_lens
|
||||
Uses ``mode="default"`` so Inductor performs operator fusion without
|
||||
capturing its own CUDA graphs. This avoids conflicts with vLLM's
|
||||
``CUDAGraphWrapper`` which manages CUDA graphs for the main Talker
|
||||
model on the default stream.
|
||||
"""
|
||||
if self._compiled_model_fwd is not None:
|
||||
return
|
||||
self._compiled_model_fwd = torch.compile(
|
||||
self.model.forward,
|
||||
mode="default",
|
||||
dynamic=True,
|
||||
)
|
||||
logger.info("code_predictor: torch.compile enabled (mode=default)")
|
||||
|
||||
with (
|
||||
set_current_vllm_config(self._vllm_config),
|
||||
set_forward_context(
|
||||
attn_metadata,
|
||||
self._vllm_config,
|
||||
num_tokens=int(hs.shape[0]),
|
||||
slot_mapping=slot_mappings,
|
||||
),
|
||||
):
|
||||
out = self.model(positions=positions, inputs_embeds=hs)
|
||||
|
||||
logits = self.lm_head[generation_step](out)
|
||||
return logits
|
||||
# ------------------------------------------------------------------
|
||||
# Optimized forward: re-prefill + torch.compile + projection cache
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
@torch.inference_mode()
|
||||
def forward(
|
||||
@@ -463,62 +428,71 @@ class Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM(nn.Module):
|
||||
layer0_embed: torch.Tensor,
|
||||
last_talker_hidden: torch.Tensor,
|
||||
do_sample: bool = True,
|
||||
temperature: float = 1.0,
|
||||
temperature: float = 0.9,
|
||||
top_k: int = 50,
|
||||
top_p: float = 1.0,
|
||||
) -> torch.Tensor:
|
||||
"""Full autoregressive prediction of residual codebooks 1..Q-1.
|
||||
"""Predict residual codebooks 1..Q-1 autoregressively via re-prefill.
|
||||
|
||||
Args:
|
||||
layer0_code: [B, 1] first-layer codec token ids.
|
||||
layer0_embed: [B, 1, H] embedding of layer0_code.
|
||||
last_talker_hidden: [B, 1, H] hidden state from the talker.
|
||||
do_sample: whether to sample or take argmax.
|
||||
temperature: sampling temperature.
|
||||
top_k: top-k filtering.
|
||||
top_p: top-p (nucleus) filtering.
|
||||
torch.compile fuses the ~60 small kernel launches per step into fewer
|
||||
fused kernels, reducing kernel launch overhead by ~75%.
|
||||
|
||||
Returns:
|
||||
audio_codes: [B, Q] all codebook tokens (layer0 + residuals).
|
||||
Projection caching: each token is projected once via small_to_mtp_projection
|
||||
and cached in _proj_buf, avoiding redundant re-projection of past tokens.
|
||||
"""
|
||||
bsz = int(layer0_code.shape[0])
|
||||
num_groups = int(self.config.num_code_groups)
|
||||
max_steps = num_groups - 1
|
||||
num_groups = self._num_groups
|
||||
device = layer0_code.device
|
||||
dtype = layer0_embed.dtype
|
||||
|
||||
# Reset KV cache for a fresh sequence.
|
||||
self.reset_cache()
|
||||
all_codes = torch.empty(bsz, num_groups, dtype=torch.long, device=device)
|
||||
all_codes[:, 0] = layer0_code.reshape(bsz)
|
||||
|
||||
# Prefill: feed [last_talker_hidden, layer0_embed] → logits for group 1.
|
||||
prefill_input = torch.cat([last_talker_hidden, layer0_embed], dim=1) # [B, 2, H]
|
||||
logits = self.prefill_logits(prefill_input) # [B, vocab]
|
||||
self._ensure_buffers(bsz, device, dtype)
|
||||
self._setup_compile()
|
||||
|
||||
all_codes = [layer0_code.reshape(bsz, 1)]
|
||||
past_seq_len = 2
|
||||
proj_buf = self._proj_buf
|
||||
pos_ids = self._pos_ids
|
||||
|
||||
projection = self.small_to_mtp_projection
|
||||
model_fwd = self._compiled_model_fwd
|
||||
lm_heads = list(self.lm_head)
|
||||
codec_embeds = list(self.model.codec_embedding)
|
||||
|
||||
proj_buf[:bsz, 0, :] = projection(last_talker_hidden.reshape(bsz, 1, -1)).reshape(bsz, -1)
|
||||
proj_buf[:bsz, 1, :] = projection(layer0_embed.reshape(bsz, 1, -1)).reshape(bsz, -1)
|
||||
|
||||
use_sampling = do_sample and temperature > 0
|
||||
inv_temperature = 1.0 / max(temperature, 1e-6) if use_sampling else 0.0
|
||||
if use_sampling and top_p != 1.0:
|
||||
raise NotImplementedError(
|
||||
"top_p sampling is not implemented for the vLLM-native code predictor; please set top_p=1.0."
|
||||
)
|
||||
|
||||
for step in range(1, num_groups):
|
||||
# Sample or argmax from logits.
|
||||
if do_sample and temperature > 0:
|
||||
scaled = logits.float() / temperature
|
||||
seq_len = step + 1
|
||||
|
||||
projected = proj_buf[:bsz, :seq_len, :]
|
||||
step_pos_ids = pos_ids[:seq_len] if bsz == 1 else pos_ids[:seq_len].repeat(bsz)
|
||||
|
||||
hidden_out = model_fwd(projected, step_pos_ids)
|
||||
|
||||
logits = lm_heads[step - 1](hidden_out[:, -1, :])
|
||||
|
||||
if use_sampling:
|
||||
scaled = logits * inv_temperature
|
||||
if top_k > 0:
|
||||
topk_vals, _ = scaled.topk(top_k, dim=-1)
|
||||
scaled = scaled.masked_fill(scaled < topk_vals[:, -1:], float("-inf"))
|
||||
probs = torch.softmax(scaled, dim=-1)
|
||||
# Guard against NaN/Inf from lower-precision dtypes.
|
||||
probs = probs.clamp(min=0.0)
|
||||
row_sums = probs.sum(dim=-1, keepdim=True)
|
||||
probs = torch.where(row_sums > 0, probs / row_sums, torch.ones_like(probs) / probs.shape[-1])
|
||||
next_ids = torch.multinomial(probs, num_samples=1) # [B, 1]
|
||||
probs = F.softmax(scaled, dim=-1)
|
||||
next_ids = torch.multinomial(probs, num_samples=1)
|
||||
else:
|
||||
next_ids = logits.argmax(dim=-1, keepdim=True) # [B, 1]
|
||||
all_codes.append(next_ids)
|
||||
next_ids = logits.argmax(dim=-1, keepdim=True)
|
||||
|
||||
# If not the last step, decode one more token.
|
||||
if step < max_steps:
|
||||
logits = self.decode_logits(
|
||||
next_ids.reshape(bsz),
|
||||
generation_step=step,
|
||||
past_seq_len=past_seq_len,
|
||||
)
|
||||
past_seq_len += 1
|
||||
all_codes[:, step] = next_ids.reshape(bsz)
|
||||
|
||||
return torch.cat(all_codes, dim=1) # [B, Q]
|
||||
if step < num_groups - 1:
|
||||
new_embed = codec_embeds[step - 1](next_ids)
|
||||
proj_buf[:bsz, step + 1, :] = projection(new_embed.reshape(bsz, 1, -1)).reshape(bsz, -1)
|
||||
|
||||
return all_codes
|
||||
@@ -15,6 +15,7 @@ stage_args:
|
||||
model_arch: MiMoAudioForConditionalGeneration
|
||||
worker_type: ar
|
||||
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
|
||||
tensor_parallel_size: 1 # Change to desired TP size for multi-GPU inference (e.g., 4 for 4 GPUs)
|
||||
gpu_memory_utilization: 0.3
|
||||
enforce_eager: true # need to discuss
|
||||
trust_remote_code: true
|
||||
@@ -45,6 +46,7 @@ stage_args:
|
||||
model_arch: MiMoAudioForConditionalGeneration
|
||||
worker_type: generation
|
||||
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
|
||||
tensor_parallel_size: 1 # Change to desired TP size for multi-GPU inference (e.g., 4 for 4 GPUs)
|
||||
gpu_memory_utilization: 0.2
|
||||
enforce_eager: true
|
||||
trust_remote_code: true
|
||||
|
||||
@@ -15,6 +15,7 @@ stage_args:
|
||||
model_arch: MiMoAudioForConditionalGeneration
|
||||
worker_type: ar
|
||||
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
|
||||
tensor_parallel_size: 1 # Change to desired TP size for multi-GPU inference (e.g., 4 for 4 GPUs)
|
||||
gpu_memory_utilization: 0.3
|
||||
enforce_eager: true # need to discuss
|
||||
trust_remote_code: true
|
||||
@@ -46,6 +47,7 @@ stage_args:
|
||||
model_arch: MiMoAudioForConditionalGeneration
|
||||
worker_type: generation
|
||||
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
|
||||
tensor_parallel_size: 1 # Change to desired TP size for multi-GPU inference (e.g., 4 for 4 GPUs)
|
||||
gpu_memory_utilization: 0.2
|
||||
enforce_eager: true
|
||||
trust_remote_code: true
|
||||
|
||||
@@ -0,0 +1,67 @@
|
||||
async_chunk: false
|
||||
stage_args:
|
||||
- stage_id: 0
|
||||
stage_type: llm
|
||||
runtime:
|
||||
devices: "0"
|
||||
max_batch_size: 1
|
||||
engine_args:
|
||||
model_stage: qwen3_tts
|
||||
model_arch: Qwen3TTSTalkerForConditionalGeneration
|
||||
hf_overrides:
|
||||
architectures: [Qwen3TTSTalkerForConditionalGeneration]
|
||||
worker_type: ar
|
||||
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
|
||||
enforce_eager: false
|
||||
trust_remote_code: true
|
||||
async_scheduling: false
|
||||
enable_prefix_caching: false
|
||||
engine_output_type: latent
|
||||
gpu_memory_utilization: 0.3
|
||||
distributed_executor_backend: "mp"
|
||||
max_num_batched_tokens: 512
|
||||
max_model_len: 4096
|
||||
default_sampling_params:
|
||||
temperature: 0.9
|
||||
top_k: 50
|
||||
max_tokens: 4096
|
||||
seed: 42
|
||||
detokenize: false
|
||||
repetition_penalty: 1.05
|
||||
stop_token_ids: [2150]
|
||||
|
||||
- stage_id: 1
|
||||
stage_type: llm
|
||||
runtime:
|
||||
devices: "0"
|
||||
max_batch_size: 1
|
||||
engine_args:
|
||||
model_stage: code2wav
|
||||
model_arch: Qwen3TTSCode2Wav
|
||||
hf_overrides:
|
||||
architectures: [Qwen3TTSCode2Wav]
|
||||
worker_type: generation
|
||||
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
|
||||
enforce_eager: true
|
||||
trust_remote_code: true
|
||||
async_scheduling: false
|
||||
enable_prefix_caching: false
|
||||
engine_output_type: audio
|
||||
gpu_memory_utilization: 0.2
|
||||
distributed_executor_backend: "mp"
|
||||
max_num_batched_tokens: 8192
|
||||
max_model_len: 32768
|
||||
engine_input_source: [0]
|
||||
custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_tts.talker2code2wav
|
||||
final_output: true
|
||||
final_output_type: audio
|
||||
tts_args:
|
||||
max_instructions_length: 500
|
||||
default_sampling_params:
|
||||
temperature: 0.0
|
||||
top_p: 1.0
|
||||
top_k: -1
|
||||
max_tokens: 65536
|
||||
seed: 42
|
||||
detokenize: true
|
||||
repetition_penalty: 1.0
|
||||
@@ -8,6 +8,37 @@ from vllm.logger import init_logger
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
def talker2code2wav(
|
||||
stage_list: list[Any],
|
||||
engine_input_source: list[int],
|
||||
prompt: Any = None,
|
||||
requires_multimodal_data: bool = False,
|
||||
) -> list[Any]:
|
||||
"""Non-async processor: wait for talker to finish, then pass all codes to code2wav at once."""
|
||||
from vllm_omni.inputs.data import OmniTokensPrompt
|
||||
from vllm_omni.model_executor.stage_input_processors.qwen3_omni import _validate_stage_inputs
|
||||
|
||||
talker_outputs = _validate_stage_inputs(stage_list, engine_input_source)
|
||||
code2wav_inputs: list[OmniTokensPrompt] = []
|
||||
for talker_output in talker_outputs:
|
||||
output = talker_output.outputs[0]
|
||||
# audio_codes shape: [num_frames, Q] where Q=num_quantizers (16)
|
||||
audio_codes = output.multimodal_output["audio_codes"].to(torch.long)
|
||||
# Filter zero-padded frames (EOS/invalid steps), matching _extract_last_frame behavior
|
||||
valid_mask = audio_codes.any(dim=1)
|
||||
audio_codes = audio_codes[valid_mask]
|
||||
# Code2Wav expects codebook-major flat: [Q*num_frames]
|
||||
codec_codes = audio_codes.transpose(0, 1).cpu().reshape(-1).tolist()
|
||||
code2wav_inputs.append(
|
||||
OmniTokensPrompt(
|
||||
prompt_token_ids=codec_codes,
|
||||
multi_modal_data=None,
|
||||
mm_processor_kwargs=None,
|
||||
)
|
||||
)
|
||||
return code2wav_inputs
|
||||
|
||||
|
||||
def _extract_last_frame(pooling_output: dict[str, Any]) -> torch.Tensor | None:
|
||||
audio_codes = pooling_output.get("audio_codes")
|
||||
if not isinstance(audio_codes, torch.Tensor) or audio_codes.numel() == 0:
|
||||
|
||||
@@ -248,7 +248,6 @@ class OmniRequestOutput:
|
||||
f"images={images_repr}",
|
||||
f"prompt={self.prompt!r}",
|
||||
f"latents={self.latents}",
|
||||
f"metrics={self.metrics}",
|
||||
f"multimodal_output={self._multimodal_output}",
|
||||
]
|
||||
|
||||
|
||||
Reference in new issue
Block a user