vLLM-Omni 单阶段 AR 模式实践：以 MOSS-TTS-Nano 为例的端到端流式 TTS 接入指南

vLLM-Omni 单阶段 AR 模式实践：以 MOSS-TTS-Nano 为例的端到端流式 TTS 接入指南

【免费下载链接】vllm-omniA framework for efficient model inference with omni-modality models项目地址: https://gitcode.com/GitHub_Trending/vl/vllm-omni

导读

本文讲解 vLLM-Omni 中的Single-Stage AR（单阶段自回归）接入模式：当上游模型无法干净地拆分为"AR 语言模型 + 独立解码器"时（典型如 MOSS-TTS-Nano，以及所有通过inference_stream()生成器把 AR 与 codec 捆绑在一起的模型），如何把整条流水线放进一个 AR worker 中、按请求逐块产出音频。读完本文，你将掌握这种模式的模型实现骨架、pipeline.py固定拓扑声明、deploy YAML 的运行时配置要点，以及流式生命周期与 lint 纪律的完整实战方案。本文以仓库内.claude/skills/add-tts-model/references/single-stage-ar.md为骨架，结合 MOSS-TTS-Nano 的真实落地代码展开。

为什么需要单阶段 AR 模式

vLLM-Omni 的编排模型通常鼓励"多阶段"拓扑：AR 阶段做文本/自回归推理，扩散或 VAE 阶段做音频/图像解码。但并非所有上游模型都适合这种切分：

• 无法干净拆分：部分模型的inference_stream()生成器把 AR LM 与音频 codec 绑定在同一个循环里，逐帧产出音频事件；
• 强共享状态：远程模型与音频 tokenizer 共享进程级 RNG 与流式解码状态，跨请求并发会导致音频损坏（这直接决定了后文max_num_seqs=1的限制）；
• 体积小、无需多卡：如 MOSS-TTS-Nano 仅 0.1B，AR LM 与 MOSS-Audio-Tokenizer-Nano 合起来约 2 GiB，单卡即可容纳。

此时正确的做法是：把整条流水线运行在单个 AR worker 内，模型forward()每次被调用时从 per-request 生成器中弹出一个音频块，由 AR 调度器维持请求存活直到最后一个块产出。

需要区分的是，VoxCPM2 虽然也是单阶段运行，但它在基础语言模型上使用 vLLM 原生 PagedAttention，并在 vLLM 之外做扩散/VAE 旁路计算，属于另一变体（见plan/voxcpm2_native_ar_design.md的原始设计说明），不在本文范围内。

模型实现骨架

参考文档给出的实现分为三步，每一步在 MOSS-TTS-Nano 的真实代码中都有对应落点。

1. 单一模型文件，同时加载 AR LM 与 codec

MOSS-TTS-Nano 的 AR 语言模型与音频 tokenizer 都加载在同一个MossTTSNanoForGeneration内（见 modeling_moss_tts_nano.py）：

class MossTTSNanoForGeneration(nn.Module): requires_raw_input_tokens = True have_multimodal_outputs = True has_preprocess = False has_postprocess = False enable_update_additional_information = True def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: super().__init__() ... self._lm: nn.Module = lm # AR LM，0.1B self._audio_tokenizer: nn.Module = audio_tokenizer # MOSS-Audio-Tokenizer-Nano self._stream_gens: dict[str, Any] = {} # request_key → generator self._ar_last_chunk_flags: list[bool] = [] # 与最近一次 forward batch 对齐的 EOS 掩码

模型类声明了have_multimodal_outputs = True，并以OmniOutput作为返回类型，通过multimodal_outputs={"model_outputs": ..., "sr": ...}携带音频波形与采样率。

2. 权重加载的时机：__init__与load_weights()的取舍

参考文档强调"在load_weights()中加载权重，而不是__init__()"，理由是 vLLM 会在任何 CUDA 分配之前初始化分布式状态。但 MOSS-TTS-Nano 的真实实现注释给出了一个有意偏离：它在__init__中通过AutoModelForCausalLM.from_pretrained(...)急切构造模型，目的是：

1. 让load_format: dummy可用（DummyModelLoader会跳过load_weights，但急切构造保证参数已就位）；
2. 避免 vLLM init 后 KV-cache profiling 带来的 OOM；
3. 与 qwen3_tts（PR #3117）的做法保持一致。

对应的load_weights()只做"排空迭代器 + 上报全部参数"：

def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: for _ in weights: pass return {name for name, _ in self.named_parameters()}

从源码结构看，这属于"从__init__急切加载 +load_weights空实现"的变体——如果你的模型通过 vLLM 的load_format机制加载权重，仍应遵循文档的通用原则，在load_weights()中加载；选择哪种取决于是否需要dummy加载与 KV-cache profiling 的兼容性。

此外，真实代码还包含两个值得注意的加载细节：

• transformers 兼容垫片：_hf_load_without_tp_warmup()上下文管理器临时把transformers.modeling_utils._is_torch_distributed_initialized置为False，规避 transformers 5.8.x 在torch.distributed已初始化时对_tp_plan=None的 remote-code 模型做 warmup 导致的TypeError；transformers_keys_to_ignore_compat()则处理 transformers 5.9 的 keys list-vs-set 变更；
• RoPE 权重修复：reinit_rotary_inv_freq(lm, base=10000.0)重新初始化以persistent=False注册的inv_freq缓冲区，否则trust_remote_code的自定义 RoPE 类会在首次 forward 时产生 NaN logits。

3. 通过 per-request 生成器流式输出

参考文档给出的核心模式是：forward()根据runtime_additional_information（每个请求一个 dict）拿到 request_key，首次见到该 key 时创建生成器存入self._stream_gens，之后每次next()弹出一个(waveform, is_last)元组；is_last为真时删除生成器。MOSS-TTS-Nano 的实现进一步明确了 request_key 的取值：

request_key = str(info.get("global_request_id") or info.get("_omni_req_id") or id(info))

代码注释说明：global_request_id由引擎设置（info 键集合为['text', 'mode', 'prompt_audio_array', 'global_request_id', 'omni_final_stage_id', 'generated_len']），而_omni_req_id是永不会被当前引擎设置的遗留回退键——如果回退到常量，所有请求会塌缩到同一个生成器，请求 N 的残留块会重放给请求 N+1，表现为"后一个请求的音频与前一请求输入一致"的串扰。这一点对文档中"request ID 由 vLLM 设置"的说明做了重要补充：新引擎下应优先使用global_request_id。

生成器主体遍历上游inference_stream()事件流，同时处理两种事件类型：

def _create_stream_gen(self, info: dict[str, Any]): ... for event in self._lm.inference_stream( text=text, output_audio_path=output_path, mode=mode, prompt_text=prompt_text, prompt_audio_path=prompt_audio_path, text_tokenizer_path=self.model_path, audio_tokenizer=self._audio_tokenizer, device=device, nq=None, max_new_frames=max_new_frames, do_sample=True, use_kv_cache=True, **sampling, ): event_type = str(event.get("type", "")) if event_type == "audio": waveform = event.get("waveform") if waveform is not None: chunk = _to_mono_1d(waveform) audio_chunks.append(chunk) yield chunk, False # 增量块，is_last=False elif event_type == "result": if not audio_chunks: # 回退：无增量块时使用最终合并结果 waveform = event.get("waveform") if waveform is not None: yield _to_mono_1d(waveform), True return yield torch.zeros((0,), dtype=torch.float32), True # 结束哨兵

这里包含三个工程细节：

• 事件双路径："audio"（增量）与"result"（最终合并）两种事件都要处理，短句在某些后端上可能只发result；
• 声道混合：_to_mono_1d()对(channels, samples)张量按 channel 求均值降为单声道——MOSS 音频 tokenizer 配置为双声道，若用.T.reshape(-1)直接展平会把 L/R 交错成 2 倍长度流，以 1× 采样率回放导致播放速度慢一倍；
• RNG 状态保存/恢复：seed非空时快照并恢复 CPU/GPU RNG 状态，避免上游依赖全局 RNG 的采样逻辑污染引擎其他组件。

forward()对 dummy/profiling 调用有专门分支：当runtime_additional_information为空或全部_is_dummy时，直接返回空输出并置self._ar_last_chunk_flags = [True] * len(infos)立即结束请求（dummy 信息由get_dummy_runtime_additional_information()提供）。正常路径下，每次next(generator)弹出一个块并更新逐行 EOS 掩码。

4. 用compute_logits()控制 AR 调度器的请求生命周期

单阶段 AR 模式下forward()本身不产出真正的 logits，请求何时结束由compute_logits()决定。MOSS-TTS-Nano 的实现按行构造 logits：

def compute_logits(self, hidden_states, sampling_metadata=None): ... for row in range(num_rows): is_last = flags[row] if row < len(flags) else True # 失同步时保守视为结束 if is_last: logits[row, eos_id] = 1.0e6 # EOS 占优 → 调度器结束该请求 else: logits[row, eos_id] = -1.0e9 logits[row, safe_id] = 1.0e6 # 非 EOS → 请求保持存活，等待下一块 return logits

eos_id = 2（vocab_size > 2时），与 pipeline 配置中sampling_constraints={"stop_token_ids": [2]}的硬性兜底对应：即使compute_logits()逻辑被绕过，调度器侧仍有 stop token 兜底。

5. 异常终止时的生成器清理

forward()只在正常完成（is_last或StopIteration）时弹出生成器。取消、超时、抢占等异常终止会泄漏生成器并跳过其finally块，导致临时 WAV 文件残留。因此模型实现了on_requests_finished(finished_req_ids)：

def on_requests_finished(self, finished_req_ids): for req_id in finished_req_ids: gen = self._stream_gens.pop(str(req_id), None) if gen is not None: gen.close() # 触发 GeneratorExit，让 finally 清理块运行

关键点速查

参考文档归纳的三条 key points，在真实代码中均可验证：

• runtime_additional_information是正确的参数名（不是**kwargs），它按批内每个请求携带一个 dict——forward()签名中显式声明了runtime_additional_information: list[dict[str, Any]] | None；
• 请求 ID 的语义是info.get("_omni_req_id")由 vLLM 设置，不是用户代码设置——实际引擎中该键已废弃，应使用global_request_id；
• 必须同时处理上游模型的"audio"（增量）与"result"（最终合并）两种事件类型。

Pipeline 与 deploy 配置

固定拓扑声明：pipeline.py

单阶段 AR 的拓扑是固定不可变的，在pipeline.py中声明。MOSS-TTS-Nano 的真实声明位于 pipeline.py：

from vllm_omni.config.stage_config import ( PipelineConfig, StageExecutionType, StagePipelineConfig, ) MOSS_TTS_NANO_PIPELINE = PipelineConfig( model_type="moss_tts_nano", default_deploy_config_name="moss_tts_nano.yaml", model_arch="MossTTSNanoForCausalLM", stages=( StagePipelineConfig( stage_id=0, model_stage="moss_tts_nano", execution_type=StageExecutionType.LLM_AR, input_sources=(), final_output=True, final_output_type="audio", owns_tokenizer=True, engine_output_type="audio", sampling_constraints={ "detokenize": False, "stop_token_ids": [2], # compute_logits() 的硬性兜底 }, ), ), )

各字段语义（见 stage_config.py 中PipelineConfig/StagePipelineConfig定义）：

• model_type/model_arch：模型唯一标识与 HF 架构名，StageConfigFactory据此路由到正确的 pipeline；HF 架构名撞车时可用hf_architectures元组消歧；
• execution_type=StageExecutionType.LLM_AR：声明为 AR 阶段。调度器由此解析：async_chunk为 false 时用OmniARScheduler，否则用OmniARAsyncScheduler（_resolve_scheduler的对应逻辑）；
• input_sources=()：单阶段无上游，输入源为空；
• owns_tokenizer=True：本阶段拥有 tokenizer；
• engine_output_type="audio"与final_output=True/final_output_type="audio"：引擎侧输出与最终对外输出均为音频。

注册与部署配置

在 pipeline_registry.py 中注册（第 102 行导入MOSS_TTS_NANO_PIPELINE，第 191 行挂到"moss_tts_nano"键），并在 registry.py 中注册模型架构"MossTTSNanoForCausalLM"到moss_tts_nano.modeling_moss_tts_nano.MossTTSNanoForGeneration。

运行时放置与尺寸由 deploy YAML 决定，参考 moss_tts_nano.yaml：

async_chunk: false trust_remote_code: true stages: - stage_id: 0 max_num_seqs: 1 gpu_memory_utilization: 0.3 enforce_eager: true enable_prefix_caching: false max_num_batched_tokens: 4096 max_model_len: 4096 devices: "0" skip_mm_profiling: true default_sampling_params: temperature: 1.0 top_p: 1.0 top_k: 50 max_tokens: 4096 seed: 42

每个字段的选择依据（文件头注释 + 源码共同支撑）：

• async_chunk: false：同步分块模式，与_resolve_scheduler中非异步路径的OmniARScheduler对应；
• trust_remote_code: true：上游模型与 codec 均需远程代码执行；
• enforce_eager: true：模型走inference_stream()+ trust_remote_code，未接 CUDA graph 捕获；
• max_num_seqs: 1：硬性约束。_validate_max_num_seqs()会在__init__中直接raise ValueError拒绝大于 1 的值，因为远程模型与音频 tokenizer 共享进程级 RNG 与流式解码状态，并发请求会互相污染音频；
• gpu_memory_utilization: 0.3：AR LM + codec 合计约 2 GiB，0.3 足够；
• skip_mm_profiling: true：MOSS-TTS-Nano 无多模态输入预处理，profile 阶段跳过 dummy-MM pass；
• default_sampling_params：与上游 demo 默认值一致（_DEFAULT_TEXT_TEMPERATURE=1.0、_DEFAULT_TEXT_TOP_P=1.0、_DEFAULT_TEXT_TOP_K=50）。

该配置在 1× L4 24GB 上验证通过，且对 H20 / A100 / 3090 同样足够。

拓扑与部署的职责边界：execution_type、输出归属、tokenizer 归属、模型架构都属于 pipeline 拓扑（frozen 的StagePipelineConfig），绝不能挪进 deploy YAML；deploy YAML 只负责设备、显存、批量尺寸等运行时放置。

Lint 纪律

参考文档提醒：只从additional_information中提取实际会转发给模型调用的变量，未使用的提取会触发 pre-commit 的ruff F841（unused variable）。在 MOSS-TTS-Nano 的_create_stream_gen中可以看到规范的提取模式——每个键经_pick(info, key, default)提取后都被用于inference_stream(...)的具名参数或分支逻辑，无悬空变量。

参考实现清单

• 拓扑声明：vllm_omni/model_executor/models/moss_tts_nano/pipeline.py
• 部署配置：vllm_omni/deploy/moss_tts_nano.yaml
• 模型实现：vllm_omni/model_executor/models/moss_tts_nano/modeling_moss_tts_nano.py
• 流水线注册：vllm_omni/config/pipeline_registry.py
• 架构注册：vllm_omni/model_executor/models/registry.py
• 拓扑/执行类型定义：vllm_omni/config/stage_config.py
• 端到端测试：tests/e2e/offline_inference/test_moss_tts_nano_expansion.py、tests/e2e/online_serving/test_moss_tts_nano_expansion.py、tests/model_executor/models/moss_tts_nano/test_npu_compat.py

小结

Single-Stage AR 模式是 vLLM-Omni 为"AR 与 codec 深度捆绑"类模型提供的标准接入范式：单模型文件承载全部计算，per-request 生成器驱动逐块流式输出，compute_logits()按行决定请求生命周期，固定拓扑留在pipeline.py、运行时放置留在 deploy YAML。MOSS-TTS-Nano 作为仓库内参考实现，完整演示了从_stream_gens管理、事件双路径、声道混合、RNG 状态恢复到on_requests_finished清理的全部细节——接入同类模型时，可以直接以其为模板照此结构落地。

【免费下载链接】vllm-omniA framework for efficient model inference with omni-modality models项目地址: https://gitcode.com/GitHub_Trending/vl/vllm-omni