Official FastVideo/LTX2-Distilled-Diffusers. New torch_ltx2.py adapters (build_torch_* dispatch on
class):
- TorchLTX2DiT: patchify-internal; per-token timestep ones(B,tok,1)*sigma (sigma direct) + per-sample
video_sigma; DiT predicts x0 so the adapter returns velocity=(x_t-x0)/sigma for the v2 flow-match step.
- TorchLTX2VAE: CausalVideoAutoencoder decode (internal per-channel un_normalize).
- TorchGemma: LTX2GemmaTextEncoderModel (Gemma + feature-extractor + connectors) -> last_hidden_state.
- ltx2 loop: real 128-ch latent geometry on cuda (32x spatial / 8x temporal; half-res base, 2x upsample).
e2e two-stage (8+3 steps) -> coherent, high-quality video (surfers at sunset), (3,9,512,768). NOTE:
still uses the v2 program's np.repeat upsampler between stages (the refine regenerates from noise so
output is faithful-quality); real LTX2LatentUpsampler swap-in is a follow-up.
Official SF checkpoint wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers (reuses TorchWanVAE + TorchT5Encoder).
- TorchWanDiT detects the causal transformer: ignores the chunk_rollout loop's latent `context` (the
real model conditions across chunks via an internal kv_cache, not a forward arg) -> dispatches to
full-attention _forward_train, and passes a per-latent-frame timestep [B, num_frames] (the causal
block asserts a per-frame temb), uniform per chunk.
- chunk_rollout real geometry on cuda (16ch; chunk_size latent frames; 8x spatial).
e2e: chunk_rollout over the SF student -> coherent video (cat in a garden), (3,21,480,832). Fidelity
gap (artifacts): the v2 loop's per-chunk few-step sampling != the official kv-cache streaming + SF
schedule (a follow-up).
The cuda-availability probe imports torch by design; the no-torch-import invariant is only
verifiable when torch is absent. skipif torch installed -> green on GPU box (202 passed, 2 skipped),
still enforced in torchless CI.
- latent_shape(req, model): real Wan geometry (16ch; (T-1)//4+1, H/8, W/8) on the cuda backend;
toy stand-in stays on accel/cpu.
- stamp_wan21_checkpoints(card, model_root): map a root (local dir or HF id) onto the 3 components'
ComponentSpec.checkpoint; build_wan21_card(checkpoint_root=...) optional.
Take the written-not-run torch adapters to runs-and-generates on 1x H100 (aarch64):
- A: FastVideoArgs.from_kwargs(model_path=root) builds the real pipeline_config; single-GPU dist
init; load each component from its subfolder; tokenizer from the sibling <root>/tokenizer.
- B/C: DiT forward wrapped in set_forward_context(attn_metadata=None) (SDPA dense path);
timestep=sigma*1000 + bare-velocity output confirmed.
- D: VAE decode denormalizes z*std+mean; removed the double-mean (it re-added
shift_factor==latents_mean) that washed out the video.
- E: UMT5 from config; text embeds zero-padded to text_len (Wan t5_postprocess_text) - the fix
that took output from a dark blur to a coherent prompt-matching scene.
- Components run at native precision (DiT bf16, VAE/text fp32); checkpoint check before dist init.
Orientation + process doc for an agent on a GPU branch: the 6 commits already
landed, the files to touch, the gating tasks (Risk A FastVideoArgs/checkpoint),
the verification bar (CPU suite stays 204; GPU generation matches a reference via
the SSIM harness), commit/push rules (no Claude co-author; don't rewrite history;
wandb token referenced not embedded), and the gotchas. Points to
GPU_BRINGUP.md for the detailed checklist + risk table.
Adversarial cross-check of the written-not-run torch adapters against the real
fastvideo source confirmed the interface contracts (DiT returns bare velocity
tensor; timestep=sigma*1000; encode().mode() + bare decode; .last_hidden_state;
no fused solver kernel) but caught a wrong construction layer. Fixed in code:
- Construction: WanTransformer3DModel / AutoencoderKLWan have NO from_pretrained.
Replace it with the real FastVideo loaders (TransformerLoader / VAELoader /
TextEncoderLoader + TokenizerLoader, each load(model_path, fastvideo_args)).
The loader resolves the class from the checkpoint config — so UMT5-vs-T5 is
chosen correctly instead of hardcoded (was BLOCKER #1/#4/#5).
- Text encoder: wrap the forward in set_forward_context(...) — the (U)MT5
attention reads global state via get_forward_context(); a bare call mis-encodes
(was BLOCKER #2). Drop the wrong padding="max_length".
- VAE: apply latent normalization the DiT expects — (z-mean)*inv_std on encode,
inverse on decode, with latents_std stored as its reciprocal; shift_factor
before decode (was BLOCKER #3). Skipping it yields washed-out video, not error.
The remaining unknowns are genuinely box-dependent (FastVideoArgs fields,
shift_factor placement, exact tokenizer kwargs, FSDP) — GPU_BRINGUP.md reconciled
to mark what's now fixed-in-code vs what still needs the box. 204 CPU tests pass.
Implement the GPU backend the substrate was built for: Platform.detect() ->
cuda resolves real torch adapters + torch solver ops instead of the numpy
rungs, with the existing loops/policies/scheduler/training unchanged.
WRITTEN-NOT-RUN: this environment has no GPU/torch, so the torch code is
grounded in the verbatim real fastvideo APIs (DiT forward signature confirmed
from source) but cannot be executed/verified here. It is gated available=False
(CPU mini stays green; importing the backends never imports torch), with every
on-box confirm point marked `# BRINGUP` and an ordered checklist in
platform/backends/GPU_BRINGUP.md.
- torch_adapters.py: TorchWanDiT / TorchWanVAE / TorchT5Encoder wrap the real
module named by each card's load_id and bridge it to the mini's duck-typed
surface (numpy<->torch at the boundary; loop math stays numpy fp32). DiT
weight-surface (copy_from/blend_from/clone) for serving sync; mse_grad_step
raises (GPU training is a separate workstream).
- torch_kernels.py: flow_match_step / flow_sde_step as plain torch elementwise.
Grounded conclusion from the kernel audit: fastvideo-kernel ships NO fused
solver kernel (only attention/norm/quant primitives), so the cuda solver is
torch, registered at arch generic with an honest source string.
- torch_cuda.py: rewritten as lazy trampolines (torch imported only inside
builder/kernel bodies). Adds the missing vae + text_encoder cuda components
(they'd otherwise silently fall back to the toy) and corrects the dishonest
"fastvideo-kernel:flow_*" labels.
- ComponentSpec.checkpoint: the weights source for the torch adapter (risk A;
the one field the cards didn't carry). Empty on the CPU toys.
- 7 CPU-verifiable wiring tests: honest registration/sources, torch-free import,
cuda-resolves-real-cells-not-toy, build-fails-loudly-without-torch.
204 tests pass (CPU). The torch path needs a GPU box to verify (GPU_BRINGUP.md).
Close the loudest deferred gap from the cudagraph audit: the capturable step
now binds its I/O to address-stable static buffers (modeling real CUDA static
I/O buffers), instead of allocating fresh arrays per call.
- StaticWorkspace: address-stable buffers allocated once per capture key; bind()
copies the current step's inputs in place via np.copyto, which RAISES on a
shape/dtype mismatch — turning the weak peak_activation_bytes proxy into a real
key-soundness backstop (a step whose shape doesn't fit can't replay an
incompatible graph). Output written into a static buffer too.
- WorkPlan.graph_fn / graph_inputs: a capturable step exposes its deterministic
op-structure as graph_fn(model, workspace) reading EVERY per-step input (latent,
sigmas, conditioning, scale) from the workspace — never from closure over
per-step data — plus the dict of current values. Loops without both stay on the
eager path. wan21 factors a shared _velocity() so graph_fn and the eager run
stay bit-identical.
- Capturer dispatch captures/replays via graph_fn against the keyed workspace;
the workspace is shared per key on the instance. Correct under the engine's
synchronous step execution (bind+graph_fn atomic per dispatch, output returned
as a copy) — proven by the batch-of-N interleave gate running two same-key
requests through the shared workspace bit-identically. A concurrent/multi-stream
executor would need a per-stream pool (documented).
- 2 new tests (no-static-form eager-break, static-buffer shape-mismatch raises);
workspace collision test split into bytes-proxy vs shape-backstop.
197 tests pass.
Wire the capture/replay lifecycle into the driven-loop step boundary — the
other half of Path A (hand-fused kernels behind the registry + piecewise
cudagraphs, no compiler). Models and tests the correctness-critical control
logic; replay re-runs the current step thunk (CPU models the lifecycle, not the
GPU speedup).
- GraphCapturer on the instance (cross-request cache), wired into
RuntimeLoopContext.execute and gated by LoopSpec.graph_capture ==
"breakable_cudagraph". Capture key = (device, arch, loop, shape_sig,
resident-weight-versions, graph_key). Eager-break for non-capturable (SDE) and
interceptor-overridden steps. Never executes a stored thunk, so interleaved
requests can't smear state.
- WorkPlan.capturable / graph_key. wan21 sets capturable=not sde and folds
compute dtype (shape_sig.dtype) + CFG branch set + expert + scheduler-precision
into the key — closing a cross-precision key-collision corruption path a
non-fp32 build would otherwise hit on a real GPU (audit finding).
- Version-in-key auto-invalidation + real eviction: set_weights_version evicts
the synced component's graphs (duck-typed, so card/ imports no runtime),
preventing a GPU graph leak across FlowGRPO's per-iteration syncs.
- register_kernel gains the workspace_bytes capture-safety contract (declared in
the matrix; cuda cells declare real scratch, numpy reference is 0).
- 11 tests: capture-once/replay-many, eager-break (SDE + override), capture ≡
pure-eager bit-identical, recapture on shape + weight-version change with
eviction, eager-loop gating, accel-backend capture, and capturer unit tests
(key discrimination, eager-break, workspace-collision safety net, eviction).
Honestly deferred (bite a real GPU, not the CPU tests): the static-buffer
refactor of the step body, admission budgeting of capture cost (GRAPH_CAPTURE),
and per-card opt-in beyond wan21 — all documented in cudagraph.py + README.
195 tests pass.
Finish the kernel seam across the board so the platform's KernelTable is the
universal solver-dispatch path, not just wan21.
- Loops: ltx2, wan_causal, adapters, adaptive now resolve the flow-match solver
via model.platform.kernels.get(FLOW_MATCH_STEP) instead of importing the numpy
sampler directly (wan21 already did). On CPU this is bit-identical (the cpu
kernel IS the old function); a GPU/accel backend now overrides every loop.
- RL: the FlowGRPO log-prob recompute in unified_rl / joint_multi_rl /
workflow_rl dispatches FLOW_SDE_STEP through the platform, pinned to the SAME
kernel the rollout used (C2 kernel-pinning — otherwise the PPO ratio biases on
a real GPU where rollout and recompute kernels could differ).
- accel backend: add a kind-generic AccelComponent wrapper and override the vae
component too (text_encoder left unregistered to keep the device->cpu fallback
demonstrated), closing the "only dit is overridden" gap.
- tests: vae override assertion; a second-loop (wan-causal chunk rollout) parity
oracle proving accel == cpu bit-identical beyond wan21. README scope updated.
184 tests pass.
Add v2/platform/: the (recipe, runtime) backend membrane that lets CPU, GPU,
and other devices coexist behind one dispatch substrate.
- Two tuple-keyed registries: COMPONENTS(kind, device, variant) for
weight-bearing components and KERNELS(op, device, arch, variant) for
stateless primitives, each with an availability predicate and an enumerable
manifest (declared-but-unavailable cells listed without importing torch).
- Platform: detected (device, arch) owning the device + arch fallback chains
and a per-platform cached KernelTable; detect() is honest (CPU/numpy unless
torch+CUDA are actually present). Arch fallback is monotonic — only degrades
to older/portable archs, never a newer binary-incompatible one.
- Seams wired: ModelInstance.component() -> platform.build_component(spec,self)
with spec.factory as the numpy terminal rung (existing cards untouched); the
wan21 denoise thunk dispatches solver ops through model.platform.kernels.
- Three backends: cpu (numpy terminal + parity oracle), accel (pure-python
stand-in proving cross-device resolution, arch fallback, and the oracle),
torch_cuda (declared-but-unavailable; no faked GPU).
- test_platform.py: 16 tests — detection, terminal rung, arch-fallback walk +
monotonicity, device precedence, variant fallback, component override +
per-kind device fallback, the parity oracle (accel == cpu, bit_identical via
the C1 ladder), and matrix enumeration without importing torch.
Scope is honestly bounded in the README/docstrings: only wan21-denoise routes
through the kernel table and only the dit kind is overridden today (each a
one-line adoption); the torch/CUDA path and a cudagraph workspace-safety
contract are declared/deferred, not implemented. 183 tests pass.
Three more capabilities on distinct untested surfaces (167 tests pass, no new runtime primitive).
A — Adapter plane (§9.19): one base + swappable LoRA/ControlNet adapters, selected per request
(DiffusionParams.adapters); AdapterDenoiseLoop applies each active adapter's velocity delta. Per-request
selection changes output, multi-LoRA composes, ControlNet conditions on a control image, mixed-adapter
requests interleave without smearing, hot-swap changes generation, cache key partitions by adapter stack
(the adapter_versions field, previously declared-only). ToyLoRA/ToyControlNet; models/adapters/. 6 tests.
B — Non-linear workflows (§9.17): ParallelWorkflow (fan-out: one input → N models → merged) and
BestOfNWorkflow (generate N → score with the served reward card → return best; inference-time scaling).
The shapes a linear chain can't express. 4 tests.
D — RL→distill flywheel (§9.18): run_flywheel RL-improves the base (NFT), then distills FROM the RL'd model
(DMD2 teacher = RL'd policy) into a faster card, recording the base→rl→distilled provenance chain in
RecipeSpec.parents. The distilled student is measurably closer to the RL'd teacher than the base; the
distilled card serves few-step. training/flywheel.py. 4 tests.
designv4 §9.17–§9.19 + layout/closing/counts (167 tests, 29 files).
The last audit stress test. A cheap draft model proposes K tokens, the target verifies
them in one batched step, and SpeculativeARLoop accepts the matching prefix + one target
correction — a variable accepted-length per round (a ragged AR loop the model owns).
- Exactness: the emitted sequence equals the target's OWN greedy decode for any draft
quality (every accepted token is one the target would produce; the correction is the
target's token) — the speedup is free.
- Speedup scales with accept rate: draft-agree 0.3→1x, 0.7→3x, 1.0→4x=K tokens/round
(fewer verify_rounds, the expensive model's latency steps, for the same output).
- Two components (draft + target) co-scheduled on one resident instance; each round an
AR_TOKEN WorkUnit.
models/speculative/ (loop+card+program); backend ToyTargetModel/ToyDraftModel with a
shared length-dependent target formula (no degenerate fixed point). 5 tests; full suite
153 passed. designv4 §9.16 + counts (153 tests, 26 files).
The remaining design_v3 probes (all except 8b speculative decoding). All fit with no new runtime
primitive (148 tests pass).
#6 LTX-2 joint audio+video (§9.11) — LTX-2 declared an audio_vae required_for t2vs but never used
it; now a single 2-stage denoise carries a synchronized audio latent (conditioned on video),
applies per-modality CFG (guidance_per_modality), and decodes via video VAE + audio VAE → video +
audio. Gated on requesting audio, so the T2V path is byte-identical (existing tests untouched).
ToyAudioVAE; build_ltx2_av_program. 5 tests.
#4 Live weight-sync under in-flight serving (§9.14) — WeightSyncController makes the freeze → drain →
transfer → bump version + invalidate → resume lifecycle explicit. Tests: a mid-flight swap corrupts
(the hazard); draining first leaves the in-flight request bit-identical to baseline while a
post-sync request reflects new weights; transformer-only sync, so the frozen text-encoder cache
survives. The RL flywheel's hardest correctness. 3 tests.
#5 Reward-model-as-a-served-card (§9.15) — a reward model is a card (scorer + a score loop emitting
REWARD_BATCH units); ServedRewardScorer drop-in-replaces the numpy scorer so any RL method becomes
RLHF/RLAIF with no method change. ToyRewardModel; models/reward/. 4 tests.
#7 Content-adaptive control flow (§9.12) — CacheDiTDenoiseLoop (isolated WanDenoiseLoop subclass)
reuses the cached velocity when predictions barely change (cache-dit skip) and early-exits on
convergence — variable step count; interleave parity holds across ragged loops. models/adaptive/. 4 tests.
#8a Nested workflows (§9.13) — a workflow stage can invoke another workflow (engine.run routes ids);
requires/validate recurse; cycles caught at registration + a run-time guard (engine._wf_running).
build_t2i_i2v_extend_workflow. 5 tests.
designv4 §9.11–§9.15 + falsifier/layout/closing updates (148 tests, 25 files). Also removed a
pre-existing unused import in ltx2/loop.py.
Three more example folders alongside inference/, all CPU/numpy, self-contained
(sys.path bootstrap), public API only, every script verified to run green.
training/ (7) — one per method, all via the uniform method.train_step seam:
01 finetune · 02 dmd2 distillation · 03 diffusion_nft (likelihood-free RL,
samples from the old policy, feature-cache reuse) · 04 self-forcing (causal
chunk_rollout) · 05 joint LM+generator RL (UniRL; joint + prompt-only) ·
06 N-way joint RL (per_expert vs shared credit) · 07 end-to-end workflow RL
(T2I+I2V from one final-video reward).
omni/ (4) — 01 Cosmos3 (reason→joint denoise, shared MoT) · 02 BAGEL
(text→image, shared MoT; scheduler prices both WorkUnit kinds) · 03 Qwen-Omni
(thinker→talker→vocoder, three separate experts, text+audio) · 04 interleave
parity across AR + diffusion loop types.
workflows/ (2) — 01 cross-model T2I→I2V workflow (image provably conditions the
video) · 02 workflow as a first-class servable (requires/validate, address by
id, register_workflows catalog, WorkflowRegistry).
Each folder has a README indexing its scripts.
Five self-contained, runnable scripts (CPU/numpy) for the Wan2.1-1.3B card on the
v2 runtime, each bootstrapping the repo onto sys.path so they run from anywhere:
- 01_basic_t2v.py minimal path: build engine → T2V request → run → video
- 02_params_and_reproducibility.py DiffusionParams knobs + seeded bit-identical reproducibility
- 03_streaming.py per-denoise-step preview chunks (OutputSpec stream)
- 04_concurrent_interleaved.py step-interleaved batching + interleave parity gate + cache reuse
- 05_async_serving.py AsyncEngine: concurrent generate, event stream, step-boundary cancel
+ README.md indexing them. All five run green; use only the public API.
- v2/README.md: point to designv4.md as the unified design (design_v3 as
north star); refresh the scope table (joint/N-way/workflow RL, Qwen-Omni
cascade, cross-model Workflow, tiled VAE co-scheduling, WorldModelSession),
package layout (program/Workflow, runtime/session, the 7 methods, new model
dirs), the demonstrated-stress-tests list (§9.3–§9.10), and counts (49→127,
20 files). Sessions moved out of "out of scope"; WebRTC wire stays out.
- designv4.md: drop two intermediate absolute suite totals (milestone "91/97
passed") in favor of "zero regressions" so the only absolute count is the
current 127 (intro + layout) — no stale numbers.
Targets the three design_v3 claims that were most load-bearing AND least
exercised (sessions/realtime, training-plane boundary, the WorkUnit-generality
falsifier). All fit with no new runtime primitive (127 tests pass).
1. Interactive world-model session (runtime/session.py, §9.8) — the Session
plane had ZERO coverage. WorldModelSession drives the causal chunk_rollout
loop as a long-lived session: persistent cross-request world state on the
Session.kv_handle, frame streaming, transactional step-boundary cancellation
(a cancelled act leaves the world resumable), no cross-session smearing.
Added only a continuation seam to the chunk loop (init seeds context from a
world_context slot; default empty = unchanged one-shot path). 5 tests.
2. End-to-end RL over a cross-model workflow (training/methods/workflow_rl.py,
§9.9) — trains BOTH flux-t2i and wan-i2v from ONE final-video reward. Rolls
out the whole workflow with SDE capture in both instances; the same final
advantage drives FlowGRPO PPO on each stage's transformer; two WeightSyncPlans
on two instances. The earlier model (T2I) is trained by a reward on the final
video — end-to-end credit across a model boundary — proven causal by a control
(constant reward => zero advantage => nothing moves). 4 tests.
3. Heterogeneous WorkUnit co-scheduling (models/tiled/, §9.10) — the §17
falsifier. VAETileLoop makes VAE decode a loop of VAE_TILE units; tiling is
exact (== one-shot, C0), and VAE_TILE + DIFFUSION_STEP pipelines interleave
bit-identically and co-run in one batch. Validates the mechanism; the
economic half (does it pay) stays a port-time measurement. 4 tests.
designv4 §9.8–§9.10 + falsifier/layout/closing updates (127 tests, 20 files).
Answers "what's the right way to name/register custom pipelines like T2I→I2V":
treat a Workflow like a card — a stable namespaced id in the same servable
namespace, declared dependencies, and a two-level registry. No new concepts;
mirrors how cards are registered (and vllm-omni's pipeline_registry).
- program/workflow.py: Workflow gains `requires` (the cards it composes, derived
from stages) and `validate(engine)` (fail-fast if a required card is absent,
P7). New WorkflowRegistry: declarative workflow_id -> builder catalog for
out-of-tree/ad hoc use.
- runtime/engine.py: `_workflows` registry + register_workflow (validates deps,
rejects id collision with a model_id) + serves(); engine.run routes a request
whose model_id is a workflow to workflow.run — addressable exactly like a model.
Single-model hot path untouched.
- runtime/async_engine.py + serving/server.py: serves() and /models include
workflows (discoverable as servables).
- models/__init__.py: declarative `_WORKFLOWS` catalog (the cross-model analog of
_BUILDERS) + register_workflows() helper; build_image_video_engine now registers
the workflow too. Adding a custom pipeline = one catalog line.
- Naming convention: dotted/namespaced workflow_id (`image_video.t2i_i2v`),
distinct from kebab model ids, collision-checked. Renamed from `t2i_then_i2v`.
- tests (+5, 12 total in the file): addressable by id, requires/validate, id
collision, registry catalog, register_workflows helper. Full suite 114 passed.
- designv4 §9.6: the naming & registration convention documented.
Two more pipelines stress-testing the design, plus a BAGEL-placement note in
designv4. Both fit with no new runtime primitive (109 tests pass).
Pipeline 1 — cross-model T2I→I2V (program/workflow.py, models/image_video/):
- Realizes ProgramKind.WORKFLOW as a thin multi-instance orchestrator ABOVE the
engine (the hot path stays single-instance). A Program composes one model's
loops; a Workflow chains full engine.run calls across distinct cards, threading
artifacts. (LTX-2 already covers same-card multi-stage; cross-model — FLUX→Wan
— is the new capability the single-instance runner can't express.)
- flux-t2i (text→image) and wan-i2v (text+image→video) cards; the I2V program
folds the conditioning image into text_embeds so WanDenoiseLoop is unchanged.
- Each model keeps its own interleave-parity guarantee (crossing instances is a
Workflow boundary, not a loop step). 7 tests incl. video-depends-on-image.
Pipeline 2 — N-way joint RL (training/methods/joint_multi_rl.py, models/multi_expert/):
- JointMultiExpertRL generalizes UnifiedRLMethod (N=2) to N refiner LMs + a
generator: one reward → one group advantage → N token-PG updates + 1 FlowGRPO
PPO update, N+1 independent WeightSyncPlans. Proves the substrate was already
N-ready (card holds N components/loops; per-component weight-sync; dict grad
targets) — only the method body looped over two; now it loops over a list.
- credit="per_expert" learns all N cleanly; credit="shared" (faithful to UniRL)
works but is noisier — the honest multi-agent credit-assignment result, a
reward-shaping choice, not a substrate limit. 6 tests (N=1,3,4; prompt-only).
Fix — flow_sde_ml_velocity (loop/sampler.py): the toy FlowGRPO generator update
targeted the velocity the model already produced (a no-op once guidance_scale=1
was set for the C2 identity; the unified generator moved only on ~1e-7 noise).
The correct PG surrogate targets the max-likelihood velocity of the realized
sample — nonzero at ratio==1. Both UniRL and N-way generators now learn for real;
the C2 ratio==1 identity still holds (measured before the update).
BAGEL: MoT/shared-weight (one transformer on both loops), same row as Cosmos3;
real BAGEL's co-resident experts are expressible via the expert-routing policy
(partial sharing) — captured in designv4 §2.3.
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.
Stress-tests the v2 Card/Loop/Program design with a UniRL/PromptRL-style
joint RL recipe: a prompt-refiner LM expert and a flow generator expert,
two separate experts driven by two loop types in one request, both updated
simultaneously from a single RL reward.
- loop/sampler.py: flow_sde_step_with_logprob — FlowGRPO SDE rollout sampler
(per-step Gaussian log-prob), distinct from the deterministic ODE serve step.
- request/params.py: gated sde_rollout/sde_noise_scale on DiffusionParams so
the serve path stays byte-identical (default ODE).
- models/wan21/loop.py: gated SDE-rollout capture in WanDenoiseLoop.advance.
- models/backend.py: ToyPromptRefiner — a real REINFORCE categorical policy
(the Qwen role), separate weights from the generator.
- models/unified/: the unified card+program — two disjoint experts (llm +
transformer) on ar_decode + diffusion_denoise; the topological opposite of
the Cosmos3 MoT card, same vocabulary.
- training/methods/unified_rl.py: joint GRPO — one reward -> group advantage
-> LM token policy gradient + DiT FlowGRPO PPO; two LRs; prompt-only/joint
flag; reuses the shared diffusion loop for rollout.
- training/weight_sync.py: WeightSyncPlan gains a component scope so the two
experts version + cache-invalidate independently (LM sync never flushes the
frozen text-encoder feature cache).
- tests/test_unified_rl.py: 9 tests incl. likelihood-based C2 identity, the
two-loop interleave parity gate, joint vs prompt-only. Full suite 91 passed.
designv4.md: unified design doc reflecting v2 as built+tested, with the joint
RL stress test as the validating case study (the design held — new card +
new method, no new runtime primitive).
Directory rename (git mv, history preserved) plus rewrite of all references: absolute imports in
tests, the zero-dep runner, docstrings, comments, and the README. No behavior change.
Run: python3 -m pytest v2/tests/ -q ; python3 v2/run_tests.py ; python3 -m v2.examples
Review confirmed the core bets (concurrent disaggregation is bit-identical, design conformance holds,
Dynamo genuinely optional, engine stays step-scheduled). Fixes for the untested failure paths:
- HIGH: pool capacity (RolePool.in_flight) no longer leaks when a disaggregated request is cancelled
or errors mid-occupancy — DisaggregatedRunner.close() releases the occupied pool and AsyncEngine._run
calls it in a finally (a cancel on a capacity-1 denoiser no longer bricks the pool).
- credit flow-control: cross-pool transfer wraps acquire/release in try/finally (no credit leak on a
failing transfer); slot is re-homed only on a successful fetch.
- AsyncEngine: duplicate in-flight request_id is rejected (was a deadlock); submit() cancels the driver
task when the consumer abandons the stream (client disconnect → no orphaned compute); bounded
per-request history (no unbounded _events/_states/_results/_runners growth).
- cancellation is common-path on the OFFLINE path too (cancel check at the top of every runner.tick()).
- HTTP server: read timeout (slowloris guard → 408), body-size cap (→ 413), invalid Content-Length
(→ 400), explicit StreamReader limit, and aclose() of the SSE generator on client disconnect.
- build_deployment_card no longer aliases one mutable CostModel across replica cards (dataclasses.replace),
so online calibration of one worker's cost doesn't mutate another's.
- video-job tasks tracked (not fire-and-forget); server.close() cancels/drains them; jobs dict bounded;
fleet affinity map bounded.
6 regression tests added for these paths. 82 tests pass (pytest + zero-dep runner).
Builds the full serving layer the design files specify, instead of deferring it to Dynamo:
- transport/ (§7.3): pluggable Connectors (in-proc zero-copy / SHM-fake copy) with chunk_ready
readiness (vllm-omni) AND credit-based flow control (sglang-omni Relay); KVConnector protocol shape;
TransferManifest.
- runtime/ (§6, §13; plan M3/M4): AsyncEngine — request queue, lifecycle state machine
(waiting→running→completed/cancelled/failed), live AsyncIterator[OmniEvent] streaming, common-path
cancellation, step-level concurrency. RolePool + DisaggregatedRunner (encoder→denoiser→decoder,
capacity-aware dispatch, cross-pool transfers via connectors); disaggregated output is bit-identical
to inline. No-progress detection (no busy-spin).
- deploy/ (§14, §6.3.5-6): DeploymentCard; OUR OWN LocalFleet (discovery, health/drain, least-loaded
/ cost-model / sticky-affinity routing) so we never rely on Dynamo; DynamoWorkerAdapter +
FakeDynamoRuntime export the SAME card + cost model so Dynamo CAN front us — one object, two consumers.
- serving/ (§6.3.5, §12): framework-free stdlib-asyncio OpenAI server (our own version of the
vllm-omni pattern): /v1/chat/completions (SSE), /v1/images/generations, /v1/videos (async job+poll)
+ /v1/videos/sync, /v1/models, /health, /metrics. A thin shim over the STEP-scheduled engine — the
runtime-visible loop scheduler vllm-omni's request-scheduled opaque DIFFUSION stage lacks.
15 serving tests (real-socket HTTP+SSE via stdlib asyncio, disagg==inline, fleet routing, Dynamo
contract, cancellation); 76 tests pass total (pytest + zero-dep runner). ~7900 LOC.
One resident MoT instance runs BOTH an ar_decode loop and a diffusion_denoise loop on shared weights
(the §16 claim no DAG-of-engines can express), with both loops runtime-visible: the scheduler prices
ar_token AND diffusion_step WorkUnits — unlike vllm-omni's opaque DIFFUSION stage the scheduler never
sees inside.
- ARDecodeLoop: token decode until EOS/max_tokens, paged text-KV — the omni AR pathway (loop/§5).
- ToyMoTDiT: one module exposing an und pathway (ar_forward) AND a gen pathway (denoise __call__);
ToyTokenizer. Binding both loops to one instance = shared weights, no duplication.
- models/cosmos3/: tokenize → reason(ar_decode) → pack(tokens→conditioning) → diffusion_denoise →
vae_decode; sound_vae declared optional_for non-t2vs (the lazy-component P8 fix, not an env-var hack).
- models/bagel/: the canonical vllm-omni model — generate_text(ar_decode) → generate_image(diffusion),
text+image outputs, both loops step-scheduled.
- The diffusion loop is WanDenoiseLoop reused (one loop definition bound to the MoT module).
- build_omni_engine() + an omni worked example; 7 omni tests (shared-instance, both-kinds-scheduled,
interleave parity across loop types, lazy sound_vae). 61 tests pass (pytest + zero-dep runner).
- Admission fails fast with AdmissionInfeasible on infeasible/deadlocked reservations instead of a
10M-iteration busy-spin: no-progress detection in run_to_completion/run_interleaved via a real
progress token, plus feasibility pre-checks (need > pool capacity).
- Compute budget is now a refundable concurrency gate (release() refunds spent), not a
never-refunded lifetime cap that silently deadlocks.
- §7.1: the text-encoder feature CacheKey carries adapter_versions + precision (no stale serve across
te-LoRA stacks); per-component weight versions mean a transformer-only RL weight sync no longer
flushes the frozen text-encoder cache (component-scoped invalidation, not wholesale).
- Interleave gate flags symmetric-empty output instead of passing it vacuously.
- skipped_steps counted only when the override is actually consumed; BatchScheduler wired for round
batch-accounting (metric renamed stepped_units); ResidualCache.get cleanup; stream chunks carry a
latent preview payload; dead progress-vars removed.
- 5 regression tests added for the previously-untested paths. 54 tests pass (pytest + zero-dep runner).
A scoped, CPU-testable realization of design_v3.md — the architecture where the atomic unit
is a typed (recipe, runtime) ModelCard, the model owns loop semantics while the runtime owns
loop lifecycle, one resident instance runs many loops, and training records behavior on the
same loops it serves.
Implements:
- card/ loop/ runtime/ cache/ memory/ parallel/ parity/ extend/ program/ request/ training/
spanning design_v3 §4-§13: ModelCard + validate(); driven loops (init/next/advance/finalize);
step-interleaving Engine with reservation-before-admission + per-class caches keyed by CacheKey;
the C0-C4 consistency ladder + the non-negotiable batch-of-N interleave parity gate.
- Inference: Wan2.1-1.3B (T2V), LTX2.3 (two-stage distilled, shared transformer), Wan-causal
(chunk rollout + slab-KV streaming).
- Training (Wan2.1-1.3B), each driving the SAME loops the engine serves: finetune (flow-match),
DMD2 (teacher/critic distribution matching), DiffusionNFT (likelihood-free C2, samples from the
decay-blended old policy, group-relative advantages, shared-prompt cache reuse), self-forcing
(causal chunk loop). The engine never imports training (the §10 dependency rule, grep-verified).
numpy-only core (no torch/GPU here); heavy Wan/LTX forwards are deterministic toy stand-ins with
lazy torch-adapter seams (ComponentSpec.load_id/factory) for a GPU box. 49 tests pass via pytest
and a zero-dependency runner; interleave parity verified (and a buggy module-global interceptor
provably breaks it). Omni-ready spine (ar_decode/chunk_step loop kinds, multi-loop instances,
LoopState.extension) for the phase-2 Cosmos3 + vllm-omni omni ports.
Run: python3 -m pytest mini_fastvideo/tests/ -q ; python3 -m mini_fastvideo.examples
SolitaryThinker