SolitaryThinker
32a7a6b87b
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).
81 lines
3.6 KiB
Python
81 lines
3.6 KiB
Python
"""VAETileLoop — tiled VAE decode (design_v3 §4.3 ``LoopKind.VAE_TILE``; §17 falsifier probe).
|
|
|
|
Decodes a latent in spatial row-tiles, **each tile a ``VAE_TILE`` WorkUnit** that flows through the
|
|
exact same admission + step-scheduling + interleave machinery as a denoise step. This is the §17
|
|
falsifier made testable: *does the general WorkUnit scheduler earn its generality on non-step kinds, or
|
|
should tiles collapse to an in-loop call?* At the mechanism level it composes cleanly — VAE_TILE units
|
|
interleave bit-identically with DIFFUSION_STEP units across requests (the §9.3 gate over mixed kinds).
|
|
Whether it *pays* on a real GPU duty-cycle is the empirical half of the falsifier (deferred to the port).
|
|
|
|
The toy VAE is spatially local (channel-mix + repeat-upsample), so tiled decode is **exactly** equal to
|
|
one-shot decode — a C0 component parity the tests assert.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import numpy as np
|
|
|
|
from ..._enums import WorkUnitKind
|
|
from ...loop.contracts import (
|
|
Done,
|
|
LoopResult,
|
|
LoopState,
|
|
ResourceRequest,
|
|
ShapeSignature,
|
|
StepResult,
|
|
WorkPlan,
|
|
)
|
|
|
|
|
|
class VAETileLoop:
|
|
def __init__(self, *, loop_id, vae_id="vae", cost, tile_rows=1, latent_slot="denoise_out"):
|
|
self.loop_id = loop_id
|
|
self.vae_id = vae_id
|
|
self.cost = cost
|
|
self.tile_rows = tile_rows
|
|
self.latent_slot = latent_slot
|
|
|
|
def init(self, req, model, ctx) -> LoopState:
|
|
out = ctx.slots.get(self.latent_slot)
|
|
latent = out["latents"] if isinstance(out, dict) else out
|
|
latent = np.asarray(latent, dtype="float32")
|
|
st = LoopState(loop_id=self.loop_id, instance_id=model.card.model_id,
|
|
request_id=req.request_id, profile=ctx.profile)
|
|
h = latent.shape[2]
|
|
st.scratch["latent"] = latent
|
|
st.scratch["tiles"] = [(r, min(r + self.tile_rows, h)) for r in range(0, h, self.tile_rows)] or [(0, h)]
|
|
st.scratch["decoded"] = [None] * len(st.scratch["tiles"])
|
|
return st
|
|
|
|
def next(self, st: LoopState):
|
|
i = st.step_idx
|
|
tiles = st.scratch["tiles"]
|
|
if i >= len(tiles):
|
|
return Done()
|
|
r0, r1 = tiles[i]
|
|
latent, vid = st.scratch["latent"], self.vae_id
|
|
|
|
def run(model, override=None):
|
|
tile = latent[:, :, r0:r1, :] # a spatial row-band of the latent
|
|
dec = model.component(vid).decode(tile)
|
|
return StepResult(output={"tile": np.asarray(dec, dtype="float32"), "i": i})
|
|
|
|
rows = r1 - r0
|
|
res = ResourceRequest(
|
|
compute_seconds=self.cost.predict(int(rows * latent.shape[1] * latent.shape[3])),
|
|
resident_bytes=int(latent.nbytes), peak_activation_bytes=int(latent[:, :, r0:r1, :].nbytes))
|
|
return WorkPlan(
|
|
loop_id=self.loop_id, instance_id=st.instance_id, kind=WorkUnitKind.VAE_TILE,
|
|
shape_sig=ShapeSignature(WorkUnitKind.VAE_TILE, dims=(rows, latent.shape[3])),
|
|
resources=res, payload={"tile": i}, run=run, label=f"vae.tile{i}")
|
|
|
|
def advance(self, st: LoopState, result: StepResult) -> LoopState:
|
|
st.scratch["decoded"][int(result.output["i"])] = result.output["tile"]
|
|
st.step_idx += 1
|
|
return st
|
|
|
|
def finalize(self, st: LoopState) -> LoopResult:
|
|
tiles = [t for t in st.scratch["decoded"] if t is not None]
|
|
video = np.concatenate(tiles, axis=2) if tiles else None # stitch row-bands back together
|
|
return LoopResult(outputs={"video": video, "latents": st.scratch["latent"]},
|
|
metrics={"vae_tiles": float(st.step_idx)})
|