SolitaryThinker
c7e0a8e894
Ports vllm-omni's canonical qwen2_5_omni omni-speech cascade as a v2 card: a third weight-sharing topology — three disjoint experts (thinker, talker, vocoder) on three loop types (ar_decode → ar_decode → audio_decode) in one request, with chained cross-stage conditioning and streaming codec→waveform. vllm-omni runs these as three opaque request-scheduled stages; v2 makes every thinker token, talker token, and vocoder chunk a runtime-visible WorkUnit. - models/backend.py: ToyTalker (a genuinely distinct AR expert, weight-salted) + ToyVocoder (streaming code2wav: codec tokens → waveform chunks). - models/omni/vocoder_loop.py: VocoderLoop filling the pre-declared LoopKind.AUDIO_DECODE / WorkUnitKind.AUDIO_CHUNK slot. - models/omni/ar_loop.py: ARDecodeLoop gains a configurable prompt_slot so two chained AR loops don't collide on the prefill slot (thinker vs talker). - models/qwen_omni/: card (3 experts/3 loops) + program (tokenize → thinker → emit_text → thinker→talker full-payload hand-off → talker → talker→vocoder → vocoder → emit_audio). Cross-stage hand-offs are explicit Program nodes, the model-native form of vllm-omni's custom_process_input_func. - _enums.py: Capability.TEXT_TO_SPEECH. - tests/test_thinker_talker.py: 6 tests incl. three-loop interleave parity, cascade conditioning, AUDIO_CHUNK streaming. Full suite 97 passed. designv4.md: §2.3 topology table extended to four topologies; new §9.5 on the cascade; reference-synthesis + package layout updated.
79 lines
3.7 KiB
Python
79 lines
3.7 KiB
Python
"""ARDecodeLoop — token decode until EOS/max_tokens (design_v3 §4.3, §5; §15b Cosmos3 reasoner).
|
|
|
|
The autoregressive pathway: ``init`` prefills the prompt tokens (and allocates paged text-KV);
|
|
``next`` emits one ``ar_token`` WorkUnit (kernel-free); ``advance`` appends the sampled token and
|
|
streams it; ``done`` on EOS/max_tokens. It binds the SAME ``transformer`` component as the omni
|
|
model's ``diffusion_denoise`` loop — one resident MoT instance running both loop types (the §16
|
|
claim). All per-request state (the token list, paged-KV namespace) lives in ``LoopState``.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
from ..._enums import ExecutionProfile, WorkUnitKind
|
|
from ...loop.contracts import (
|
|
Done,
|
|
LoopResult,
|
|
LoopState,
|
|
ResourceRequest,
|
|
ShapeSignature,
|
|
StepResult,
|
|
WorkPlan,
|
|
)
|
|
from ...request.streams import StreamChunk
|
|
|
|
|
|
class ARDecodeLoop:
|
|
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(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)
|
|
caches = getattr(model, "caches", None)
|
|
if caches is not None and caches.has("paged_kv"): # reasoner paged text-KV (minority case)
|
|
caches.pool("paged_kv").allocate(req.request_id, n_blocks=1)
|
|
return st
|
|
|
|
def next(self, st: LoopState):
|
|
if len(st.scratch["generated"]) >= st.scratch["max_tokens"]:
|
|
return Done()
|
|
tokens = list(st.scratch["tokens"])
|
|
tid = self.transformer_id
|
|
|
|
def run(model, override=None):
|
|
nxt = model.component(tid).ar_forward(tokens) # und pathway of the shared MoT weights
|
|
return StepResult(output={"token": int(nxt)})
|
|
|
|
pos = len(st.scratch["generated"])
|
|
emits = [StreamChunk(stream_id=st.request_id, modality="text", seq=pos, preview=True)]
|
|
return WorkPlan(
|
|
loop_id=self.loop_id, instance_id=st.instance_id, kind=WorkUnitKind.AR_TOKEN,
|
|
shape_sig=ShapeSignature(WorkUnitKind.AR_TOKEN, dims=(len(tokens),)),
|
|
resources=ResourceRequest(compute_seconds=self.cost.predict(len(tokens)),
|
|
resident_bytes=8 * len(tokens), peak_activation_bytes=8),
|
|
payload={"branch": "ar", "pos": pos}, run=run, emits=emits, label=f"ar.tok{pos}")
|
|
|
|
def advance(self, st: LoopState, result: StepResult) -> LoopState:
|
|
tok = int(result.output["token"])
|
|
st.scratch["tokens"].append(tok)
|
|
st.scratch["generated"].append(tok)
|
|
st.step_idx += 1
|
|
if st.profile == ExecutionProfile.ROLLOUT:
|
|
st.trajectory.append({"token": tok, "pos": st.step_idx})
|
|
return st
|
|
|
|
def finalize(self, st: LoopState) -> LoopResult:
|
|
toks = st.scratch["generated"]
|
|
return LoopResult(outputs={"tokens": toks, "text": "tok:" + ",".join(map(str, toks))},
|
|
metrics={"tokens": float(len(toks))},
|
|
behavior=st.trajectory or None)
|