diff --git a/designv4.md b/designv4.md index c8c6d07fa..098ae5f0e 100644 --- a/designv4.md +++ b/designv4.md @@ -5,7 +5,7 @@ > together with **what was actually built and tested** in the `v2/` package, and with the one workload > that most stressed the design — **UniRL/PromptRL-style joint LM + generator reinforcement learning** > (arXiv 2510.17937; PromptRL, arXiv 2602.01382). v4 is not aspirational. Every structural claim below -> is backed by code in `v2/` and a passing test (91 tests, 14 files, two independent runners — `pytest` +> is backed by code in `v2/` and a passing test (97 tests, 15 files, two independent runners — `pytest` > and a zero-dependency `v2/run_tests.py`). The neural forwards are toy numpy stand-ins (no GPU/torch in > this environment, §12); the **control flow, contracts, scheduling, caching, parity gates, and training > math are real**, which is the whole point of the (recipe, runtime) separation: swap the component @@ -89,12 +89,17 @@ The Card/Loop/Program split is not specialized to one sharing pattern. The **sam | Single diffusion | `transformer` | `diffusion_denoise` | — | `wan21` | | **MoT omni** | **one** `transformer` | `ar_decode` **+** `diffusion_denoise` | both loops → **same** component | `cosmos3`, `bagel` | | **Joint LM+gen RL** | `llm` **+** `transformer` (separate) | `ar_decode`→`llm`, `diffusion_denoise`→`transformer` | **disjoint** experts, one request | `unified` | +| **Cascade omni-speech** | `thinker` **+** `talker` **+** `vocoder` (separate) | `ar_decode`→`ar_decode`→`audio_decode` | **disjoint**, **chained** (talker reads thinker hidden state; vocoder reads talker tokens) | `qwen_omni` | -MoT (one module, two loops) and UniRL (two experts, two loops) are *topological opposites*, yet both are -ordinary `shared_weight_components` bindings on `LoopSpec`s — **no new primitive**. The third idea: *the -weight-sharing graph is data on the card, not structure in the runtime.* This is what lets the same engine -serve a 1.3B single-DiT, a 30B MoT, and a Qwen+FLUX joint-RL recipe without three runtimes. ✅ Proven by -`v2/tests/test_unified_rl.py::test_two_separate_experts_not_a_shared_mot`. +These four are the *full* spread of weight-sharing topologies — one module, one shared module across two +loops, two disjoint experts trained jointly, three disjoint experts chained — and every one of them is just +`shared_weight_components` bindings on `LoopSpec`s plus hand-off nodes in a `Program`: **no new primitive**. +MoT (shared) and UniRL (disjoint, jointly trained) are topological *opposites*; Qwen-Omni adds a third axis, +*depth* (a three-stage cascade with cross-stage conditioning). The signature idea: *the weight-sharing graph +is data on the card, not structure in the runtime.* This is what lets one engine serve a 1.3B single-DiT, a +30B MoT, a Qwen+FLUX joint-RL recipe, and a thinker→talker→vocoder speech model without four runtimes. +✅ Proven by `test_unified_rl.py::test_two_separate_experts_not_a_shared_mot` and +`test_thinker_talker.py::test_three_disjoint_experts_three_loop_types`. ## 3. Planes and dependency order @@ -306,6 +311,29 @@ new card + a new method, not a new engine* — which is exactly the claim §2.3 it. The refinements the test forced (the component-scoped weight sync; the explicit likelihood-based C2 rung; `guidance_scale=1` at RL time to keep the PPO recompute faithful) are improvements, not patches. +### 9.5 Second stress test — Qwen-Omni thinker→talker→vocoder (the cascade topology) + +A second canonical vllm-omni model was ported to confirm the §2.3 claim from a different direction: +vllm-omni's `qwen2_5_omni` pipeline — a **thinker** (multimodal LLM → text), a **talker** (AR → speech +codec tokens, conditioned on the thinker's hidden state), and a **code2wav vocoder** (codec tokens → +waveform). vllm-omni expresses it as three opaque stages with `custom_process_input_func` hand-offs the +scheduler never sees inside; v2 expresses it as **three driven loops** (`v2/models/qwen_omni/`) — every +thinker token, talker token, and vocoder chunk a runtime-visible WorkUnit. + +What it added beyond UniRL: a **three-expert, three-loop-type cascade** (`ar_decode → ar_decode → +audio_decode`) with **chained cross-stage conditioning** (talker reads the thinker's tokens *and* hidden +state — the "full payload" path; vocoder reads the talker's speech tokens) and **streaming codec→waveform** +via `AUDIO_CHUNK` WorkUnits. The loop kinds were *already* in the enum — `LoopKind.AR_DECODE`'s docstring +literally names "thinker/talker", and `LoopKind.AUDIO_DECODE` + `WorkUnitKind.AUDIO_CHUNK` were defined for +exactly the vocoder — so the port added a `ToyTalker`/`ToyVocoder`, a `VocoderLoop` (filling the anticipated +`AUDIO_DECODE` slot), a card, a program, and one tiny generalization (`ARDecodeLoop` gained a configurable +`prompt_slot` so two chained AR loops don't collide on the prefill slot). The contract, scheduler, caches, +and parity gate were untouched again. ✅ `v2/tests/test_thinker_talker.py` (6 tests, incl. the three-loop +interleave parity gate and the cascade-conditioning checks); full suite **97 passed**, zero regressions. + +The two stress tests together span the topology space — *shared* (MoT), *disjoint-joint* (UniRL), +*disjoint-cascade* (Qwen-Omni) — and the design absorbed all three as **cards + methods, never engines**. + ## 10. Serving & fleet — our own stack, Dynamo optional Per the explicit instruction "*don't completely rely on Dynamo; we still need our own version*," v2 ships a @@ -349,12 +377,13 @@ v2/ request/ requests, params (DiffusionParams + sde_rollout), tasks (TaskType), streams, artifacts models/ backend (toy components) + common (cached_text_encode) wan21/ ltx2/ wan_causal/ # phase 1: T2V, distilled, self-forcing student - omni/ cosmos3/ bagel/ # phase 2: MoT (one module, two loops) + omni/ cosmos3/ bagel/ # phase 2: MoT (one module, two loops); omni/ = ARDecodeLoop + VocoderLoop unified/ # §9: two experts, two loops, joint RL + qwen_omni/ # §9.5: thinker→talker→vocoder, three experts/three loops training/ rollout, behavior, rewards, weight_sync, methods/{finetune,dmd2,diffusion_nft,self_forcing,unified_rl} serving/ AsyncEngine, pools, DisaggregatedRunner, connectors, OpenAI server deploy/ DeploymentCard, LocalFleet, DynamoWorkerAdapter - tests/ 14 files, 91 tests run via `pytest v2/tests/` OR `python3 v2/run_tests.py` + tests/ 15 files, 97 tests run via `pytest v2/tests/` OR `python3 v2/run_tests.py` ``` **Enforced boundaries:** `card/` imports no product/runtime; `runtime/` executes `card/` loops but defines @@ -385,6 +414,7 @@ multi-expert RL force a redesign?"* — answered: it generalizes, no redesign ( | Source | Take | Constrain / reject | |---|---|---| | **UniRL-Zero / PromptRL** (arXiv 2510.17937, 2602.01382) | Joint LM-refiner + flow-generator RL under one reward; FlowGRPO SDE-prefix/ODE-tail with per-step log-probs; group-relative advantage → token-PG **and** PPO; partial-refinement contrast; prompt-only vs joint ablation; two LRs | Expresses the two experts and the joint update as a **card + method**, not a bespoke trainer; the SDE rollout sampler is a *library* entry gated behind `sde_rollout`, never the serve path; the PPO ratio rests on the **likelihood-based C2** parity gate, not on hope | +| **vllm-omni `qwen2_5_omni`** (thinker-talker-vocoder) | The canonical 3-stage omni-speech cascade: thinker (MM-LLM→text) → talker (AR→speech codec tokens, conditioned on thinker hidden state) → code2wav vocoder (tokens→waveform), streaming | Three opaque request-scheduled stages with `custom_process_input_func` hand-offs become **three driven loops** (ar_decode×2 + audio_decode), every token/chunk a step-visible WorkUnit; the cross-stage hand-offs are explicit `Program` nodes, not stage-boundary callbacks; the vocoder fills the pre-declared `LoopKind.AUDIO_DECODE` slot | --- @@ -392,11 +422,12 @@ multi-expert RL force a redesign?"* — answered: it generalizes, no redesign ( v3 argued that the atomic unit is the (recipe, runtime) pair and that a model-native, loop-driven runtime is what the post-training era requires. v4 is the same argument, now *built and measured* in `v2/`: the -contracts hold under serving, under MoT omni, and — the hard case — under **joint LM + generator -reinforcement learning**, where one reward simultaneously updates two separate experts driven by two -different loop types in one request. That workload, which UniRL/PromptRL identify as central to the future -of content creation, fit inside the design as *a new card and a new method* with no new runtime primitive. -The design's bet is therefore not that it is elegant — it is that **the weight-sharing topology and the -training recipe are data on a card, so a new frontier recipe is a card, not a rewrite.** The remaining work -is the GPU port, where, by construction, the loops, scheduler, caches, parity, and training plane do not -change — only the component factories do. +contracts hold under serving, under MoT omni, under a thinker→talker→vocoder **cascade**, and — the hard +case — under **joint LM + generator reinforcement learning**, where one reward simultaneously updates two +separate experts driven by two different loop types in one request. Those workloads — UniRL/PromptRL's +joint RL (central to the future of content creation) and vllm-omni's omni-speech cascade — both fit inside +the design as *a new card and a new method/loop*, with no new runtime primitive. The design's bet is +therefore not that it is elegant — it is that **the weight-sharing topology and the training recipe are +data on a card, so a new frontier recipe is a card, not a rewrite.** The remaining work is the GPU port, +where, by construction, the loops, scheduler, caches, parity, and training plane do not change — only the +component factories do. diff --git a/v2/_enums.py b/v2/_enums.py index 8464aee07..6d7f3c669 100644 --- a/v2/_enums.py +++ b/v2/_enums.py @@ -70,6 +70,7 @@ class Capability(str, Enum): TEXT_TO_IMAGE = "text_to_image" TEXT_TO_VIDEO_SOUND = "text_to_video_sound" AUDIO_TO_VIDEO = "audio_to_video" + TEXT_TO_SPEECH = "text_to_speech" # thinker→talker→vocoder (Qwen-Omni): reason + speak REASONING_TEXT = "reasoning_text" ACTION_CONDITIONING = "action_conditioning" STREAMING_VIDEO_CONTINUATION = "streaming_video_continuation" diff --git a/v2/models/__init__.py b/v2/models/__init__.py index a5599a991..2f54e07ca 100644 --- a/v2/models/__init__.py +++ b/v2/models/__init__.py @@ -10,6 +10,7 @@ from typing import Any from .bagel import build_bagel_card, build_bagel_program from .cosmos3 import build_cosmos3_card, build_cosmos3_program from .ltx2 import build_ltx2_card, build_ltx2_program +from .qwen_omni import build_qwen_omni_card, build_qwen_omni_program from .unified import build_unified_card, build_unified_program from .wan21 import build_wan21_card, build_wan_t2v_program from .wan_causal import build_wan_causal_card, build_wan_causal_program @@ -20,6 +21,7 @@ __all__ = [ "build_wan_causal_card", "build_wan_causal_program", "build_cosmos3_card", "build_cosmos3_program", "build_bagel_card", "build_bagel_program", + "build_qwen_omni_card", "build_qwen_omni_program", "build_unified_card", "build_unified_program", "build_default_engine", "build_omni_engine", "build_unified_engine", ] @@ -30,10 +32,12 @@ _BUILDERS = [ (build_wan_causal_card, build_wan_causal_program), ] -# Phase-2 omni cards (MoT, one resident instance running AR + diffusion loops on shared weights) +# Phase-2 omni cards: MoT shared-weight (Cosmos3/BAGEL) + the cascaded thinker→talker→vocoder +# (Qwen-Omni, three disjoint experts / three loop types in one request). _OMNI_BUILDERS = [ (build_cosmos3_card, build_cosmos3_program), (build_bagel_card, build_bagel_program), + (build_qwen_omni_card, build_qwen_omni_program), ] diff --git a/v2/models/backend.py b/v2/models/backend.py index cec8f81af..4083a166a 100644 --- a/v2/models/backend.py +++ b/v2/models/backend.py @@ -221,6 +221,47 @@ class ToyPromptRefiner: self.logits = decay * self.logits + (1.0 - decay) * other.logits +class ToyTalker(ToyMoTDiT): + """Toy Talker — Qwen-Omni stage 1 (AR over a speech-codec vocab, conditioned on the Thinker). + + A *separate expert* from the thinker (not weight-shared): its next-token depends on its OWN + weights (a seed salt) AND the prefilled thinker payload (tokens + hidden state). So the cascade + is real — change the thinker and the talker's tokens change; change the talker's weights and they + change too. The vllm-omni Talker is likewise a distinct model conditioned on Thinker hidden states. + """ + + def __init__(self, channels: int = LATENT_CHANNELS, seed: int = 0): + super().__init__(channels=channels, seed=seed) + self.salt = (seed % 97) + 1 + + def ar_forward(self, tokens) -> int: + ctx = sum(int(t) for t in tokens) + return int((ctx * 7 + 13 + self.salt) % self.VOCAB) # weight-dependent (salt) + context + + +class ToyVocoder: + """Toy streaming code2wav vocoder — Qwen-Omni stage 2 (BigVGAN/Code2Wav role). + + Speech codec tokens → audio waveform, synthesized in chunks. Deterministic given the tokens so + the ``audio_decode`` loop is interleave-safe. Each token contributes a short, timbre-stamped + waveform segment — enough to exercise chunked synthesis, streaming emits, and the AudioArtifact. + """ + + def __init__(self, samples_per_token: int = 16, vocab: int = 256, seed: int = 7): + self.spt = int(samples_per_token) + rng = np.random.default_rng(seed) + self.bank = (rng.standard_normal(vocab) * 0.3).astype("float32") # per-token timbre offset + + def synthesize(self, token_chunk) -> np.ndarray: + """One chunk of codec tokens → a waveform segment in [-1, 1], length spt·len(chunk).""" + segs = [] + for t in token_chunk: + t = int(t) % self.bank.size + phase = np.linspace(0.0, np.pi * (1 + (t % 8)), self.spt, dtype="float32") + segs.append(np.tanh(self.bank[t] + np.sin(phase)).astype("float32")) + return np.concatenate(segs) if segs else np.zeros(0, dtype="float32") + + class ToyVAE: """Tiny deterministic VAE. encode: video→latent (mean-pool + channel proj); decode: latent→video.""" diff --git a/v2/models/omni/__init__.py b/v2/models/omni/__init__.py index 352b0652e..c1973e80d 100644 --- a/v2/models/omni/__init__.py +++ b/v2/models/omni/__init__.py @@ -3,5 +3,7 @@ from __future__ import annotations from .ar_loop import ARDecodeLoop from .pack import emit_text_node, pack_cond_from_tokens, tokenize_node, vae_decode_node +from .vocoder_loop import VocoderLoop -__all__ = ["ARDecodeLoop", "tokenize_node", "pack_cond_from_tokens", "emit_text_node", "vae_decode_node"] +__all__ = ["ARDecodeLoop", "VocoderLoop", "tokenize_node", "pack_cond_from_tokens", + "emit_text_node", "vae_decode_node"] diff --git a/v2/models/omni/ar_loop.py b/v2/models/omni/ar_loop.py index 44d38c2e1..0572b5ce0 100644 --- a/v2/models/omni/ar_loop.py +++ b/v2/models/omni/ar_loop.py @@ -22,17 +22,19 @@ from ...request.streams import StreamChunk class ARDecodeLoop: - def __init__(self, *, loop_id, transformer_id="transformer", cost, max_tokens=8): + def __init__(self, *, loop_id, transformer_id="transformer", cost, max_tokens=8, + prompt_slot="prompt_tokens"): self.loop_id = loop_id self.transformer_id = transformer_id self.cost = cost self.max_tokens = max_tokens + self.prompt_slot = prompt_slot # which slot holds the prefill (chained AR loops differ) def init(self, req, model, ctx) -> LoopState: seed = req.sampling.seed if req.sampling.seed is not None else 0 st = LoopState(loop_id=self.loop_id, instance_id=model.card.model_id, request_id=req.request_id, profile=ctx.profile, seed=seed) - prompt_tokens = ctx.slots.get("prompt_tokens") or [1] + prompt_tokens = ctx.slots.get(self.prompt_slot) or [1] st.scratch["tokens"] = list(prompt_tokens) # prefill (folded into init for the toy) st.scratch["generated"] = [] st.scratch["max_tokens"] = min(req.sampling.max_tokens, self.max_tokens) diff --git a/v2/models/omni/vocoder_loop.py b/v2/models/omni/vocoder_loop.py new file mode 100644 index 000000000..34c31a1cb --- /dev/null +++ b/v2/models/omni/vocoder_loop.py @@ -0,0 +1,81 @@ +"""VocoderLoop — streaming code2wav (design_v3 §4.3 ``LoopKind.AUDIO_DECODE``; Qwen-Omni stage 2). + +The third loop type in the thinker→talker→vocoder cascade. It consumes the talker's speech-codec +tokens in fixed-size chunks; **each chunk is one ``AUDIO_CHUNK`` WorkUnit** that synthesizes a +waveform segment and streams it. This is the differentiation vllm-omni's own docs name: their +code2wav is one opaque ``LLM_GENERATION`` stage the scheduler never sees inside — here every chunk +is a runtime-visible, step-scheduled, cancellable WorkUnit, on the *same* driven-loop contract the +engine uses for denoise steps and AR tokens. All per-request state lives in ``LoopState``. +""" +from __future__ import annotations + +import numpy as np + +from ..._enums import ExecutionProfile, WorkUnitKind +from ...loop.contracts import ( + Done, + LoopResult, + LoopState, + ResourceRequest, + ShapeSignature, + StepResult, + WorkPlan, +) +from ...request.streams import StreamChunk + + +class VocoderLoop: + def __init__(self, *, loop_id, vocoder_id="vocoder", cost, chunk_tokens=2, speech_slot="speech_tokens"): + self.loop_id = loop_id + self.vocoder_id = vocoder_id + self.cost = cost + self.chunk_tokens = chunk_tokens + self.speech_slot = speech_slot + + def init(self, req, model, ctx) -> LoopState: + toks = list(ctx.slots.get(self.speech_slot) or [1]) + st = LoopState(loop_id=self.loop_id, instance_id=model.card.model_id, + request_id=req.request_id, profile=ctx.profile) + st.scratch["tokens"] = toks + st.scratch["chunks"] = [toks[i:i + self.chunk_tokens] + for i in range(0, len(toks), self.chunk_tokens)] or [[1]] + st.scratch["wave"] = [] + st.scratch["stream_audio"] = bool(req.outputs.stream.get("audio")) + return st + + def next(self, st: LoopState): + i = st.step_idx + chunks = st.scratch["chunks"] + if i >= len(chunks): + return Done() + chunk = chunks[i] + vid = self.vocoder_id + + def run(model, override=None): + seg = model.component(vid).synthesize(chunk) # codec tokens → waveform segment + return StepResult(output={"segment": np.asarray(seg, dtype="float32")}) + + emits = [] + if st.scratch.get("stream_audio"): + emits.append(StreamChunk(stream_id=st.request_id, modality="audio", seq=i, preview=False)) + return WorkPlan( + loop_id=self.loop_id, instance_id=st.instance_id, kind=WorkUnitKind.AUDIO_CHUNK, + shape_sig=ShapeSignature(WorkUnitKind.AUDIO_CHUNK, dims=(len(chunk),)), + resources=ResourceRequest(compute_seconds=self.cost.predict(len(chunk)), + resident_bytes=8 * len(st.scratch["tokens"]), + peak_activation_bytes=64), + payload={"chunk": i}, run=run, emits=emits, label=f"vocoder.chunk{i}") + + def advance(self, st: LoopState, result: StepResult) -> LoopState: + st.scratch["wave"].append(np.asarray(result.output["segment"], dtype="float32")) + st.step_idx += 1 + if st.profile == ExecutionProfile.ROLLOUT: + st.trajectory.append({"chunk": st.step_idx, "samples": int(result.output["segment"].size)}) + return st + + def finalize(self, st: LoopState) -> LoopResult: + wave = np.concatenate(st.scratch["wave"]) if st.scratch["wave"] else np.zeros(0, dtype="float32") + return LoopResult(outputs={"samples": wave.astype("float32")}, + metrics={"audio_chunks": float(st.step_idx), + "audio_samples": float(wave.size)}, + behavior=st.trajectory or None) diff --git a/v2/models/qwen_omni/__init__.py b/v2/models/qwen_omni/__init__.py new file mode 100644 index 000000000..3da4d3072 --- /dev/null +++ b/v2/models/qwen_omni/__init__.py @@ -0,0 +1,9 @@ +"""Qwen-Omni thinker→talker→vocoder model (vllm-omni ``qwen2_5_omni``): three disjoint experts, three +loops (ar_decode → ar_decode → audio_decode), cascaded conditioning + streaming codec→waveform. The +third weight-sharing topology proving the Card/Loop/Program vocabulary holds (design_v3 §4).""" +from __future__ import annotations + +from .card import build_qwen_omni_card +from .program import build_qwen_omni_program + +__all__ = ["build_qwen_omni_card", "build_qwen_omni_program"] diff --git a/v2/models/qwen_omni/card.py b/v2/models/qwen_omni/card.py new file mode 100644 index 000000000..5170ff05f --- /dev/null +++ b/v2/models/qwen_omni/card.py @@ -0,0 +1,105 @@ +"""Qwen-Omni thinker→talker→vocoder ModelCard (design_v3 §4; vllm-omni ``qwen2_5_omni`` pipeline). + +The canonical vllm-omni 3-stage omni-speech model, made model-native. vllm-omni expresses it as three +opaque stages (``thinker``/``talker``/``code2wav``) with ``custom_process_input_func`` hand-offs, each +request-scheduled, the scheduler never seeing inside the AR generation or the code2wav synthesis. v2 +expresses the SAME cascade as three driven loops on three components — and every thinker token, talker +token, and vocoder chunk is a runtime-visible WorkUnit (the differentiation §1 names). + +It is a *third* weight-sharing topology, in the same Card/Loop/Program vocabulary: + + * Cosmos3/BAGEL — ONE shared MoT module, two loops (weight sharing within a request). + * unified/UniRL — TWO disjoint experts, two loops, jointly RL-trained. + * Qwen-Omni — THREE disjoint experts (thinker, talker, vocoder), three loops, **cascaded**: + the talker conditions on the thinker's tokens+hidden state, the vocoder on the + talker's speech tokens; streaming codec→waveform. + +The loop kinds were anticipated: ``LoopKind.AR_DECODE`` (its docstring names "thinker/talker") for the +two language stages, ``LoopKind.AUDIO_DECODE`` for the vocoder — no new primitive. +""" +from __future__ import annotations + +from ..._enums import Capability, ConsistencyLevel, LoopKind, WorkUnitKind +from ...card import ( + CacheContract, + CapabilityMatrix, + ComponentSpec, + CostModel, + LoopSpec, + ModelCard, + ParallelismContract, + ParitySpec, + PrecisionContract, + RecipeSpec, +) +from ...parallel import ParallelPlan +from ..backend import ToyMoTDiT, ToyTalker, ToyTokenizer, ToyVocoder, _seed_from +from ..omni import ARDecodeLoop, VocoderLoop + + +def build_qwen_omni_card(model_id: str = "qwen-omni-tts") -> ModelCard: + seed = _seed_from(model_id) + ar_cost = CostModel(kind=WorkUnitKind.AR_TOKEN, base_seconds=5e-5, per_unit_seconds=1e-7) + wav_cost = CostModel(kind=WorkUnitKind.AUDIO_CHUNK, base_seconds=3e-5, per_unit_seconds=2e-7) + + def thinker_factory(): + return ARDecodeLoop(loop_id="thinker_decode", transformer_id="thinker", cost=ar_cost, + max_tokens=4, prompt_slot="prompt_tokens") + + def talker_factory(): + return ARDecodeLoop(loop_id="talker_decode", transformer_id="talker", cost=ar_cost, + max_tokens=4, prompt_slot="talker_prompt_tokens") + + def vocoder_factory(): + return VocoderLoop(loop_id="vocoder", vocoder_id="vocoder", cost=wav_cost, + chunk_tokens=2, speech_slot="speech_tokens") + + components = { + "tokenizer": ComponentSpec("tokenizer", kind="tokenizer", factory=lambda inst: ToyTokenizer(), + required_for={"reason", "t2a"}), + "thinker": ComponentSpec( # stage 0: multimodal understanding + text + "thinker", kind="dit", + load_id="vllm_omni.model_executor.models.qwen2_5_omni:Thinker", + factory=lambda inst: ToyMoTDiT(seed=seed), + resident_for=["thinker_decode"], required_for={"reason", "t2a"}), + "talker": ComponentSpec( # stage 1: text+hidden → speech tokens + "talker", kind="dit", + load_id="vllm_omni.model_executor.models.qwen2_5_omni:Talker", + factory=lambda inst: ToyTalker(seed=seed + 1), + resident_for=["talker_decode"], required_for={"t2a"}), + "vocoder": ComponentSpec( # stage 2: speech tokens → waveform + "vocoder", kind="audio_vae", + load_id="vllm_omni.model_executor.models.qwen2_5_omni:Code2Wav", + factory=lambda inst: ToyVocoder(seed=seed + 2), + resident_for=["vocoder"], required_for={"t2a"}), + } + loops = { + "thinker_decode": LoopSpec("thinker_decode", kind=LoopKind.AR_DECODE, + work_unit_kind=WorkUnitKind.AR_TOKEN, step_cost_model=ar_cost, + shared_weight_components=["thinker"], cache_policy=["paged_kv"], + loop_factory=thinker_factory), + "talker_decode": LoopSpec("talker_decode", kind=LoopKind.AR_DECODE, + work_unit_kind=WorkUnitKind.AR_TOKEN, step_cost_model=ar_cost, + shared_weight_components=["talker"], cache_policy=["paged_kv"], + loop_factory=talker_factory), + "vocoder": LoopSpec("vocoder", kind=LoopKind.AUDIO_DECODE, + work_unit_kind=WorkUnitKind.AUDIO_CHUNK, step_cost_model=wav_cost, + shared_weight_components=["vocoder"], cache_policy=[], + loop_factory=vocoder_factory), + } + card = ModelCard( + model_id=model_id, family="qwen_omni", components=components, loops=loops, + capabilities=CapabilityMatrix.of( + Capability.REASONING_TEXT, Capability.TEXT_TO_SPEECH, Capability.POLICY_ROLLOUT), + recipe=RecipeSpec(method="base", assumes_loop="vocoder", + assumes_precision="float32", consistency_required=ConsistencyLevel.C1), + parity=ParitySpec(consistency_levels=[ConsistencyLevel.C1], interleave_required=True), + caches={ + "paged_kv": CacheContract("paged_kv", max_bytes=1 << 24, block_bytes=1 << 12, + reuse_across_requests=False), + }, + precision=PrecisionContract(default_dtype="float32", training_precision="float32"), + parallelism=ParallelismContract(valid_plans=[ParallelPlan.single()], + default_plan=ParallelPlan.single()), + ) + return card.validate() diff --git a/v2/models/qwen_omni/program.py b/v2/models/qwen_omni/program.py new file mode 100644 index 000000000..a2da58632 --- /dev/null +++ b/v2/models/qwen_omni/program.py @@ -0,0 +1,69 @@ +"""Qwen-Omni serve program (vllm-omni ``qwen2_5_omni`` pipeline, made step-visible): + + tokenize → thinker_decode → emit_text + → thinker→talker (full payload: tokens + hidden state) + → talker_decode + → talker→vocoder (speech tokens) + → vocoder (audio_decode) → emit_audio + +The hand-off nodes are the model-native analogue of vllm-omni's ``custom_process_input_func`` +(``thinker2talker_full_payload``, ``talker2code2wav_full_payload``) — but here they are explicit +slot writes between three driven loops, not opaque stage-boundary callbacks. The talker conditions on +the thinker's **hidden state** (``reasoner_embed``), not only its tokens — the "full payload" path; +the streaming/sync path would pass tokens only (noted inline). Final outputs: ``text`` (thinker) and +``audio`` (vocoder waveform). +""" +from __future__ import annotations + +import numpy as np + +from ...program import ComponentNode, ModelLoopNode, Program, ProgramKind +from ..omni import emit_text_node, tokenize_node + + +def thinker_to_talker_node(instance, slots, request, ctx) -> None: + """Full-payload hand-off: condition the talker on the thinker's TOKENS *and* its hidden state. + (vllm-omni ``thinker2talker_full_payload``; the streaming variant ``thinker2talker_token_only`` + would write only the tokens.)""" + out = slots.get("thinker_out", {}) + toks = out.get("tokens", []) if isinstance(out, dict) else [] + hidden = instance.component("thinker").reasoner_embed(toks or [1]) # thinker hidden state + h = int(abs(float(np.sum(hidden))) * 1000.0) % 251 + 1 # one hidden-derived token + slots["talker_prompt_tokens"] = list(toks) + [h] + + +def talker_to_vocoder_node(instance, slots, request, ctx) -> None: + """Speech codec tokens → the vocoder's input slot (vllm-omni ``talker2code2wav``).""" + out = slots.get("talker_out", {}) + slots["speech_tokens"] = out.get("tokens", []) if isinstance(out, dict) else [] + + +def emit_audio_node(loop_slot: str, out_slot: str = "audio"): + def fn(instance, slots, request, ctx) -> None: + out = slots.get(loop_slot, {}) + slots[out_slot] = out.get("samples") if isinstance(out, dict) else out + return fn + + +def build_qwen_omni_program() -> Program: + return Program( + program_id="qwen_omni.tts", kind=ProgramKind.INLINE, + nodes=[ + ComponentNode("tokenize", fn=tokenize_node, writes=("prompt_tokens",)), + ModelLoopNode("thinker", loop_id="thinker_decode", output_slot="thinker_out", + reads=("prompt_tokens",), writes=("thinker_out",)), + ComponentNode("emit_text", fn=emit_text_node("thinker_out", "text"), + reads=("thinker_out",), writes=("text",)), + ComponentNode("thinker_to_talker", fn=thinker_to_talker_node, + reads=("thinker_out",), writes=("talker_prompt_tokens",)), + ModelLoopNode("talker", loop_id="talker_decode", output_slot="talker_out", + reads=("talker_prompt_tokens",), writes=("talker_out",)), + ComponentNode("talker_to_vocoder", fn=talker_to_vocoder_node, + reads=("talker_out",), writes=("speech_tokens",)), + ModelLoopNode("vocoder", loop_id="vocoder", output_slot="vocoder_out", + reads=("speech_tokens",), writes=("vocoder_out",)), + ComponentNode("emit_audio", fn=emit_audio_node("vocoder_out", "audio"), + reads=("vocoder_out",), writes=("audio",)), + ], + output_artifacts={"text": "text", "audio": "audio", "speech_tokens": "talker_out"}, + ).validate() diff --git a/v2/tests/test_thinker_talker.py b/v2/tests/test_thinker_talker.py new file mode 100644 index 000000000..4476e72b8 --- /dev/null +++ b/v2/tests/test_thinker_talker.py @@ -0,0 +1,114 @@ +"""Qwen-Omni thinker→talker→vocoder — the third weight-sharing topology (design_v3 §4; vllm-omni +``qwen2_5_omni`` pipeline). + +Asserts the Card/Loop/Program vocabulary holds for a THREE-expert, THREE-loop cascade — the topology +neither MoT (one shared module, two loops) nor joint-RL (two disjoint experts, two loops) covers: + + * three disjoint experts (thinker, talker, vocoder) on three loop *types* (ar_decode, ar_decode, + audio_decode) in one request — every token and vocoder chunk a runtime-visible WorkUnit; + * cascaded conditioning — the talker reads the thinker's tokens+hidden state, the vocoder the + talker's speech tokens (the model-native form of vllm-omni's ``custom_process_input_func``); + * streaming codec→waveform via AUDIO_CHUNK WorkUnits; + * the three-loop program still passes the interleave parity gate. +""" +from __future__ import annotations + +import numpy as np + +from v2._enums import LoopKind, WorkUnitKind +from v2.card import load_card +from v2.cache import CacheManager +from v2.models import build_qwen_omni_card, build_qwen_omni_program +from v2.models.qwen_omni.program import thinker_to_talker_node +from v2.parity import assert_interleave_parity +from v2.request import OutputSpec, TaskType, make_request +from v2.runtime import Engine + + +def _engine(): + card = build_qwen_omni_card() + inst = load_card(card, cache_manager=CacheManager.from_card(card)) + eng = Engine() + eng.register(card.model_id, inst, build_qwen_omni_program()) + return eng, inst + + +def _run(eng, prompt="hello world", *, stream_audio=False): + out = OutputSpec(stream={"audio": True}) if stream_audio else OutputSpec() + return eng.run(make_request(TaskType.T2A, "qwen-omni-tts", prompt, outputs=out)) + + +# --- card / serve shape ----------------------------------------------------------- # + +def test_qwen_omni_card_validates_and_serves_text_and_audio(): + eng, _ = _engine() + res = _run(eng, "Explain audio generation in 15 words.") + assert res.artifacts["text"].text.startswith("tok:") # thinker text + aud = res.artifacts["audio"] + samples = np.asarray(aud.samples) + assert samples.ndim == 1 and samples.size > 0 # vocoder waveform + assert aud.sample_rate == 44100 # AudioArtifact, first-class rate + assert float(np.abs(samples).max()) <= 1.0 # bounded waveform + + +def test_three_disjoint_experts_three_loop_types(): + """The third topology: three loops bind three DISTINCT components (no sharing) — vs Cosmos3 where + two loops bind the SAME MoT module.""" + card = build_qwen_omni_card() + kinds = {lid: spec.kind for lid, spec in card.loops.items()} + assert kinds == {"thinker_decode": LoopKind.AR_DECODE, + "talker_decode": LoopKind.AR_DECODE, + "vocoder": LoopKind.AUDIO_DECODE} + assert card.loops["vocoder"].work_unit_kind == WorkUnitKind.AUDIO_CHUNK + bound = [card.loops[l].shared_weight_components for l in ("thinker_decode", "talker_decode", "vocoder")] + assert bound == [["thinker"], ["talker"], ["vocoder"]] # three disjoint singletons + _, inst = _engine() + experts = {id(inst.component(c)) for c in ("thinker", "talker", "vocoder")} + assert len(experts) == 3 # genuinely separate weights + + +# --- cascade conditioning (the hand-offs are real) -------------------------------- # + +def test_thinker_to_talker_passes_full_payload(): + """The talker prefill = the thinker's tokens PLUS a hidden-state-derived token (vllm-omni's + ``thinker2talker_full_payload``).""" + _, inst = _engine() + slots = {"thinker_out": {"tokens": [11, 22, 33]}} + thinker_to_talker_node(inst, slots, None, None) + pre = slots["talker_prompt_tokens"] + assert pre[:3] == [11, 22, 33] # thinker tokens carried forward + assert len(pre) == 4 # + one hidden-state-derived token + + +def test_cascade_propagates_thinker_through_to_audio(): + """Different prompt → different thinker tokens → different talker speech → different waveform. + Proves the three loops are conditioned, not independent.""" + eng_a, _ = _engine() + eng_b, _ = _engine() + a = _run(eng_a, "alpha") + b = _run(eng_b, "beta gamma delta") + assert a.artifacts["text"].text != b.artifacts["text"].text + assert not np.array_equal(np.asarray(a.artifacts["audio"].samples), + np.asarray(b.artifacts["audio"].samples)) + + +# --- vocoder as runtime-visible AUDIO_CHUNK WorkUnits + streaming ------------------ # + +def test_vocoder_streams_audio_chunks(): + eng_off, _ = _engine() + eng_on, _ = _engine() + off = _run(eng_off, stream_audio=False) + on = _run(eng_on, stream_audio=True) + n_chunks = on.metrics["audio_chunks"] + assert n_chunks >= 1 + # the audio-streaming delta is exactly the vocoder's per-chunk emits (AR text emits are common) + assert on.metrics["stream_chunks"] - off.metrics["stream_chunks"] == n_chunks + + +# --- the three-loop program still passes the core parity gate --------------------- # + +def test_qwen_omni_interleave_parity(): + eng, _ = _engine() + reqs = [make_request(TaskType.T2A, "qwen-omni-tts", p) for p in ("alpha", "beta", "alpha")] + divs = assert_interleave_parity(eng, reqs) + assert not divs, f"thinker→talker→vocoder failed interleave parity: {divs}"