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.
81 lines
3.6 KiB
Python
81 lines
3.6 KiB
Python
"""Shared vocabulary enums (design_v3 §6.1, §9.2, §9.4).
|
||
|
||
These live in a leaf module so both ``card/`` (which references loop kinds and
|
||
consistency levels in its specs) and ``loop/``/``runtime/`` can import them
|
||
without a circular dependency. design_v3 §18 boundary: ``card/`` imports no
|
||
runtime; this is pure vocabulary.
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
from enum import Enum
|
||
|
||
|
||
class LoopKind(str, Enum):
|
||
"""The kind of iterative computation a LoopSpec describes (design_v3 §4.3)."""
|
||
DIFFUSION_DENOISE = "diffusion_denoise" # N solver steps over full clip (Wan, LTX-2)
|
||
CHUNK_ROLLOUT = "chunk_rollout" # causal chunk loop × inner denoise (self-forcing, world models)
|
||
AR_DECODE = "ar_decode" # token loop until EOS (reasoner/thinker/talker — phase 2)
|
||
VAE_TILE = "vae_tile" # tiled VAE encode/decode
|
||
ENCODER = "encoder" # one-shot encoder (text/vision) — degenerate single-step loop
|
||
AUDIO_DECODE = "audio_decode" # vocoder / codec decode
|
||
TRAIN_FORWARD = "train_forward" # grad-enabled forward for a training method
|
||
|
||
|
||
class WorkUnitKind(str, Enum):
|
||
"""The smallest schedulable action — the scheduler's currency unit (design_v3 §6.1).
|
||
|
||
Tokens are *one kind*, not the scheduler — the generalization of vLLM's token
|
||
scheduler that diffusion forces.
|
||
"""
|
||
AR_PREFILL = "ar_prefill"
|
||
AR_TOKEN = "ar_token"
|
||
DIFFUSION_STEP = "diffusion_step"
|
||
DIFFUSION_WINDOW = "diffusion_window"
|
||
CHUNK_STEP = "chunk_step"
|
||
ENCODER_CHUNK = "encoder_chunk"
|
||
VAE_TILE = "vae_tile"
|
||
AUDIO_CHUNK = "audio_chunk"
|
||
REWARD_BATCH = "reward_batch" # RL
|
||
LOGPROB_BATCH = "logprob_batch" # RL (likelihood-based)
|
||
TRANSFER = "transfer"
|
||
CACHE_IO = "cache_io"
|
||
GRAPH_CAPTURE = "graph_capture"
|
||
|
||
|
||
class ConsistencyLevel(str, Enum):
|
||
"""The consistency ladder (design_v3 §9.2). RL methods declare their required rung."""
|
||
C0 = "C0" # component parity (VAE/encoder/block/scheduler-step in isolation)
|
||
C1 = "C1" # loop parity (full denoise trajectory / AR logits, fixed seed)
|
||
C2 = "C2" # behavioral identity (train-forward ≡ serve-forward on the RL objective's quantity)
|
||
C3 = "C3" # distribution parity (rollout distribution under allowed nondeterminism)
|
||
C4 = "C4" # artifact quality (SSIM / reward agreement / human preference)
|
||
|
||
@property
|
||
def rank(self) -> int:
|
||
return {"C0": 0, "C1": 1, "C2": 2, "C3": 3, "C4": 4}[self.value]
|
||
|
||
|
||
class ExecutionProfile(str, Enum):
|
||
"""Three forwards, one loop definition — differ only in grad mode + capture (design_v3 §9.4)."""
|
||
SERVE = "serve" # no-grad, graphed, cached, possibly quantized
|
||
ROLLOUT = "rollout" # serve profile + behavior capture
|
||
TRAIN = "train" # grad, checkpointed, FSDP-gathered
|
||
|
||
|
||
class Capability(str, Enum):
|
||
"""CapabilityMatrix entries (designv2 model plane)."""
|
||
TEXT_TO_VIDEO = "text_to_video"
|
||
IMAGE_TO_VIDEO = "image_to_video"
|
||
VIDEO_TO_VIDEO = "video_to_video"
|
||
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"
|
||
VAE_ENCODE = "vae_encode"
|
||
VAE_DECODE = "vae_decode"
|
||
POLICY_ROLLOUT = "policy_rollout" # can serve as an RL rollout engine
|
||
LOGPROB_RECOMPUTE = "logprob_recompute"
|