[refactor]: remove training semantics from v2 inference contracts

This commit is contained in:
Peiyuan Zhang
2026-07-04 19:31:35 +00:00
parent a4d8978c37
commit 30cfbfc4f6
33 changed files with 123 additions and 260 deletions
+6 -7
View File
@@ -1,15 +1,14 @@
"""v2 — a scoped, CPU-testable realization of the model-native runtime (see v2/README.md).
"""v2 - the FastVideo inference runtime (see v2/README.md).
> A model card is a (recipe, runtime) pair with a parity obligation.
> The model owns loop semantics; the runtime owns loop lifecycle.
> One resident instance runs many loops; one scheduler runs their steps in one currency.
> Caches are correct by key; parity is correct by test; the interleave gate is non-negotiable.
> Training records behavior on the same loops it serves.
> Caches are correct by key; parity is correct by test.
Phase 1 supports Wan2.1-1.3B (T2V) and LTX2.3 (2-stage distilled), plus four training
methods on Wan2.1-1.3B (finetuning, DMD2, DiffusionNFT, self-forcing). The spine is
omni-ready (multi-loop ModelInstance, ar_decode/chunk_rollout loop kinds) for the phase-2
Cosmos3 + vllm-omni omni ports.
v2 is inference-only. Training, finetuning, distillation, RL, and optimizer
loops belong to ``fastvideo/train`` or ``fastvideo/training``. v2 only records
checkpoint provenance in recipe metadata so inference can bind weights to the
right loop and precision policy.
The core is numpy-only and CPU-testable; heavy Wan/LTX neural forwards become lazy torch
adapters (see ``v2/platform/backends/``) that are off the test path.
+2 -5
View File
@@ -21,8 +21,6 @@ if TYPE_CHECKING: # avoid runtime card -> parallel coupling
pass
# --------------------------------------------------------------------------- #
# Precision, parallelism, parity, cache, checkpoint contracts #
# --------------------------------------------------------------------------- #
@@ -31,7 +29,6 @@ class PrecisionContract:
default_dtype: str = "float32"
component_overrides: dict[str, str] = field(default_factory=dict)
quantization_scheme: str | None = None # "nvfp4" | "int8" | None
training_precision: str = "float32" # distinct from serving precision
def dtype_for(self, component_id: str) -> str:
return self.component_overrides.get(component_id, self.default_dtype)
@@ -52,7 +49,7 @@ class CacheContract:
eviction: str = "lru" # "lru" | "fifo" | "none"
reuse_across_requests: bool = True # paged/feature=True; slab depends on mode
per_component: dict[str, int] = field(default_factory=dict)
training_mode_disables_recycle: bool = False # chunk-KV training mode
disable_recycle: bool = False # keep full causal/world context instead of sliding-window recycle
@dataclass
@@ -133,7 +130,7 @@ class LoopSpec:
loop_id: str
kind: LoopKind
work_unit_kind: WorkUnitKind
behavior_schema: type | None = None # what to capture for RL (None if not training-relevant)
behavior_schema: type | None = None # optional inference trajectory/debug payload schema
extension_schema: type | None = None # per-model LoopState extension (Cosmos3PackedSeq, etc.)
cache_policy: list[str] = field(default_factory=list) # cache class names this loop draws from
graph_capture: str = "eager" # eager | breakable_cudagraph
+4 -10
View File
@@ -17,7 +17,6 @@ class LoopKind(str, Enum):
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):
@@ -34,18 +33,16 @@ class WorkUnitKind(str, Enum):
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. RL methods declare their required rung."""
"""The inference parity ladder."""
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)
C2 = "C2" # behavioral identity for inference-visible trajectories
C3 = "C3" # distribution parity (rollout distribution under allowed nondeterminism)
C4 = "C4" # artifact quality (SSIM / reward agreement / human preference)
@@ -55,10 +52,9 @@ class ConsistencyLevel(str, Enum):
class ExecutionProfile(str, Enum):
"""Three forwards, one loop definition — differ only in grad mode + capture."""
"""Inference execution profiles."""
SERVE = "serve" # no-grad, graphed, cached, possibly quantized
ROLLOUT = "rollout" # serve profile + behavior capture
TRAIN = "train" # grad, checkpointed, FSDP-gathered
ROLLOUT = "rollout" # serve profile + trajectory capture / stochastic rollout
class Capability(str, Enum):
@@ -75,5 +71,3 @@ class Capability(str, Enum):
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"
+4 -6
View File
@@ -47,13 +47,11 @@ def flow_sde_step_with_logprob(x_t,
noise=None,
prev_sample=None,
noise_scale: float = 0.7):
"""FlowGRPO-style stochastic step + Gaussian log-prob.
"""Stochastic diffusion step + Gaussian log-prob metadata.
The RL *rollout* sampler (vs the deterministic ODE ``flow_match_euler_step`` used at serve time).
Injects noise for exploration and returns the log-prob of the realized sample under the step's
Gaussian — the quantity the FlowGRPO PPO ratio uses. ``prev_sample=None`` samples a new step;
passing ``prev_sample`` (the rollout sample) recomputes its log-prob under the *current* velocity
(the ratio's numerator at update time). Returns ``(prev_sample, log_prob, mean, eff_std)``.
This is the stochastic counterpart to deterministic ``flow_match_euler_step``. ``prev_sample=None``
samples a new step; passing ``prev_sample`` recomputes its log-prob under the current velocity.
Returns ``(prev_sample, log_prob, mean, eff_std)``.
"""
x_t = np.asarray(x_t, dtype=np.float64)
velocity = np.asarray(velocity, dtype=np.float64)
+3 -4
View File
@@ -1,10 +1,9 @@
"""The consistency ladder + numeric comparison helpers.
C0 component · C1 loop · C2 behavioral · C3 distribution · C4 artifact.
C0 component · C1 loop · C2 trajectory behavior · C3 distribution · C4 artifact.
The C2 split is load-bearing: likelihood-based methods compare per-step log-probs;
likelihood-free methods (DiffusionNFT) compare seeded final-sample + prediction-space
identity (old_deviate / ref-MSE) — there are NO log-probs to match.
The ladder is for inference parity. C0/C1 are component and loop parity, C2 can compare
captured trajectory records, C3 covers stochastic distribution parity, and C4 covers artifact quality.
"""
from __future__ import annotations
+1 -2
View File
@@ -132,8 +132,7 @@ class ParallelWorkflow(Workflow):
class BestOfNWorkflow(Workflow):
"""Inference-time scaling / rejection sampling: generate N candidates (varying seed), score each with
a reward scorer, return the best. A feedback loop across models — generator + reward (e.g. the served
REWARD_BATCH card) — the other non-linear shape."""
a scorer, return the best. This is a feedback loop across servable models, not a training loop."""
def __init__(self, workflow_id, generator_stage, *, scorer, n: int = 4, score_key: str = "latents"):
super().__init__(workflow_id, [generator_stage])
+3 -3
View File
@@ -35,8 +35,8 @@ class DiffusionParams:
num_frames: int = 81
fps: int = 16
seed: int | None = None
# FlowGRPO-style RL rollout: use the SDE sampler (with per-step log-prob capture) instead of the
# deterministic ODE serve sampler. Default False ⇒ the serve path is unchanged.
# Stochastic sampling: use the SDE sampler instead of the deterministic ODE serve sampler.
# Default False means the normal serve path is unchanged.
sde_rollout: bool = False
sde_noise_scale: float = 0.7
# Active adapters for this request (LoRA / ControlNet ids on the base).
@@ -47,7 +47,7 @@ class DiffusionParams:
class CaptureMode(str, Enum):
"""What the loop captures while running."""
NONE = "none" # serve forward: no-grad, no capture
BEHAVIOR = "behavior" # rollout forward: emit BehaviorRecord slices for RL/distill
BEHAVIOR = "behavior" # trajectory/debug capture for inference diagnostics
@dataclass(frozen=True)
+1 -1
View File
@@ -1,6 +1,6 @@
"""Shared lightweight type aliases for v2.
The core (card/loop/runtime/cache/parity/program/request/training) is numpy-only
The core (card/loop/runtime/cache/parity/program/request) is numpy-only
and CPU-testable, so tensors are typed structurally as ``TensorLike``: a
``numpy.ndarray`` on the CPU test path, a ``torch.Tensor`` on GPU. The core never
imports torch — only model-component adapters do, lazily (see
+2 -15
View File
@@ -11,7 +11,6 @@ from typing import Any
from v2.recipes import build_default_engine, build_omni_engine
from v2.recipes.wan21 import build_wan21_card
from v2.core.request import DiffusionParams, OutputSpec, Request, SamplingParams, TaskType, make_request
from v2.training import build_diffusion_nft
def _t2v(mid: str, prompt: str, seed: int, steps: int = 4, **kw: Any) -> Request:
@@ -40,17 +39,6 @@ def example_c_causal_streaming(eng) -> None:
f"streamed_chunks={out.metrics.get('stream_chunks', 0)}")
def example_d_rl_rollout() -> None:
print("\n(d) RL rollout (DiffusionNFT): the SAME denoise loop + behavior capture (train ≡ serve)")
nft = build_diffusion_nft(build_wan21_card(), num_video_per_prompt=4, num_inner_timesteps=2)
loss, m = nft.managed_train_step({"prompts": ["a red car", "a blue boat"], "seeds": [1, 2]}, 0)
fc = nft.old.caches.stats()["feature"]
print(f" policy_loss={loss['policy_loss']:.3f} kl={loss['kl_div_loss']:.5f} "
f"reward_mean={m['reward_mean']:.3f} consistency={nft.consistency_level().value} (likelihood-free)")
print(f" shared-prompt feature-cache reuse: {fc['hits']} hits / {fc['misses']} misses "
f"(K samples encode the prompt once — the 24× reduction)")
def example_g_omni_mot() -> None:
print("\n(g) Omni / MoT (§16): ONE resident instance runs AR + diffusion loops on shared weights")
eng = build_omni_engine()
@@ -78,7 +66,7 @@ async def _serving_demo() -> None:
import asyncio
from v2.serving.deploy import DynamoWorkerAdapter, FakeDynamoRuntime, LocalFleet, build_deployment_card
from v2.recipes.wan21 import build_wan21_card, build_wan_t2v_program
from v2.recipes.wan21 import build_wan_t2v_program
from v2.runtime import AsyncEngine, PoolSet, wan_t2v_disaggregated
from v2.serving import OmniOpenAIServer
@@ -143,12 +131,11 @@ def main() -> None:
example_a_text_to_video(eng)
example_b_ltx2_two_stage(eng)
example_c_causal_streaming(eng)
example_d_rl_rollout()
example_g_omni_mot()
example_h_serving_and_fleet()
print("\n" + "=" * 78)
print("All examples ran on CPU with numpy toy components. The architecture (cards, driven")
print("loops, scheduler, caches, parity, training-on-shared-loops) is real; the neural")
print("loops, scheduler, caches, parity, and serving) is real; the neural")
print("forwards are toys. On a GPU box, swap ComponentSpec.factory for the torch adapters.")
print("=" * 78)
+1 -1
View File
@@ -3,7 +3,7 @@
Two tuple-keyed registries (``COMPONENTS``, ``KERNELS``) + a ``Platform`` that detects ``(device,
arch)`` and resolves both, with the numpy reference as the terminal fallback rung and parity oracle.
Lets CPU, GPU, and other backends coexist: a model's loops, policies, scheduler, caches, and
training code never name a device; they call through the component/kernel seams and the resolved
inference code never names a device; it calls through the component/kernel seams and the resolved
``Platform`` decides the implementation.
Importing this package is cheap and cycle-free (registry + platform classes only); the concrete
+1 -1
View File
@@ -45,7 +45,7 @@ class AccelDiT:
return self._inner(*args, **kwargs)
def __getattr__(self, name: str):
# delegate the trainable surface (clone/copy_from/blend_from/mse_grad_step) to the inner toy
# delegate the weight/version surface (clone/copy_from/blend_from) to the inner toy
return getattr(self._inner, name)
+6 -109
View File
@@ -2,8 +2,8 @@
There is no GPU/torch/weights in this environment, so the heavy Wan/LTX neural forwards are
represented by small, deterministic numpy components. They exercise the *real* loop control
flow, CFG/flow-shift policies, scheduler steps, cache reuse, parity gates, and training-method
math with real numbers — just with a toy network instead of a 1.3B DiT.
flow, CFG/flow-shift policies, scheduler steps, cache reuse, and parity gates with real numbers -
just with a toy network instead of a 1.3B DiT.
The contract these toys honor is what matters:
* deterministic given (weights, inputs) → bit-reproducible (the interleave/parity gates rely on this);
@@ -12,8 +12,8 @@ The contract these toys honor is what matters:
On a GPU box, ``ComponentSpec.factory`` swaps these for lazy torch adapters wrapping the real
``fastvideo.models`` modules + weights (see each model's ``components.py``); the loops, policies,
scheduler, caches, parity, and training code are unchanged. That is the whole point of the
(recipe, runtime) separation.
scheduler, caches, and parity are unchanged. That is the whole point of the (recipe, runtime)
separation.
"""
from __future__ import annotations
@@ -105,7 +105,7 @@ class ToyDiT:
audio = np.tanh(0.9 * au + 0.1 * float(sigma) + 0.3 * a_cond).astype("float32")
return video, audio
# --- minimal trainable surface (so training methods do real optimizer steps) --------- #
# --- weight/version surface for serving hot-reload tests ----------------------------- #
def clone(self) -> ToyDiT:
c = ToyDiT.__new__(ToyDiT)
c.C, c.s_text, c.s_ctx = self.C, self.s_text, self.s_ctx
@@ -113,40 +113,13 @@ class ToyDiT:
return c
def blend_from(self, other: ToyDiT, decay: float) -> None:
"""EMA / decay-blended-old-policy update: self ← decay·self + (1-decay)·other."""
"""Blend resident weights toward another component version."""
self.w_x = (decay * self.w_x + (1.0 - decay) * other.w_x).astype("float32")
self.w_t = (decay * self.w_t + (1.0 - decay) * other.w_t).astype("float32")
def copy_from(self, other: ToyDiT) -> None:
self.w_x, self.w_t = other.w_x.copy(), other.w_t.copy()
def mse_grad_step(self,
latent: np.ndarray,
text_embed,
sigma: float,
target: np.ndarray,
lr: float,
context=None,
weight: float = 1.0) -> tuple[float, float]:
"""One exact SGD step minimizing ``weight·MSE(forward(latent), target)`` w.r.t. ``w_x``.
Returns (loss, grad_norm). The grad_norm is the #1396-style per-method regression signal.
This is a *real* gradient (chain rule through tanh), so the toy student actually learns —
the structural stand-in for the FSDP/optimizer step a GPU trainer would run."""
latent = np.asarray(latent, dtype=np.float32)
target = np.asarray(target, dtype=np.float32)
z = self._pre_tanh(latent, text_embed, sigma, context)
pred = np.tanh(z)
err = pred - target
loss = float(weight * np.mean(err**2))
n = err.size
g = weight * (2.0 * err / n) * (1.0 - pred**2)
axes = (list(range(1, g.ndim)), list(range(1, latent.ndim)))
grad_wx = np.tensordot(g, latent, axes=axes) # [C_out, C_in]
grad_norm = float(np.linalg.norm(grad_wx))
self.w_x = (self.w_x - lr * grad_wx).astype("float32")
return loss, grad_norm
class ToyTokenizer:
"""Deterministic byte-ish tokenizer for the omni AR pathway (phase 2)."""
@@ -185,82 +158,6 @@ class ToyMoTDiT(ToyDiT):
return (rng.standard_normal((4, 8)) * 0.1).astype("float32")
class ToyPromptRefiner:
"""Toy LM prompt-refiner — UniRL/PromptRL's Qwen role (a *separate expert*, not MoT-shared).
A categorical policy over a small set of refinement "actions" (stand-in for the LM emitting a
refined prompt). Sampling an action picks an embedding offset that shifts the diffusion
conditioning; the realized reward then drives the policy by REINFORCE. This is the *second
trainable expert* in the unified-RL stress test: one reward → a token-policy-gradient update
here AND a FlowGRPO update on the DiT, under one advantage (PromptRL §6).
Real gradient: loss = −A·logπ(a); ∇_logits = −A·(onehot(a) − softmax). So the toy LM actually
learns to prefer the reward-favored action — the LM-side analogue of ToyDiT.mse_grad_step.
"""
def __init__(self, n_actions: int = 8, dim: int = TEXT_DIM, seq: int = TEXT_SEQ, seed: int = 0):
self.n = int(n_actions)
self.dim, self.seq = int(dim), int(seq)
self.logits: np.ndarray = np.zeros(self.n, dtype="float64") # the trainable policy params
rng = np.random.default_rng(seed)
# each action ⇒ a fixed conditioning offset (the "content" of that refined prompt)
self._emb = (rng.standard_normal((self.n, self.seq, self.dim)) * 0.1).astype("float32")
def _probs(self) -> np.ndarray:
z = self.logits - self.logits.max()
e = np.exp(z)
return e / e.sum()
def sample_refinement(self, rng) -> tuple[int, float]:
"""Sample an action (exploration) and return (action, logπ(action))."""
p = self._probs()
a = int(rng.choice(self.n, p=p))
return a, float(np.log(p[a] + 1e-12))
def argmax_refinement(self) -> tuple[int, float]:
a = int(np.argmax(self.logits))
return a, float(np.log(self._probs()[a] + 1e-12))
def action_logprob(self, action: int) -> float:
return float(np.log(self._probs()[int(action)] + 1e-12))
def ar_forward(self, tokens) -> int:
"""und-pathway hook so the declared ``ar_decode`` loop is runnable (greedy ⇒ interleave-safe)."""
return int(np.argmax(self.logits))
def refined_embed(self, base_embed, action: int) -> np.ndarray:
"""Apply the chosen refinement to the base text embedding → the diffusion conditioning."""
base = np.zeros((self.seq, self.dim), dtype="float32") if base_embed is None \
else np.asarray(base_embed, dtype="float32")
return (base + self._emb[int(action)]).astype("float32")
def pg_step(self, action: int, advantage: float, lr: float) -> float:
"""REINFORCE on the categorical policy. Returns grad_norm (the per-method regression signal)."""
p = self._probs()
onehot: np.ndarray = np.zeros(self.n, dtype="float64")
onehot[int(action)] = 1.0
grad = -float(advantage) * (onehot - p) # ∇_logits (−A·logπ(a))
self.logits = self.logits - float(lr) * grad # descent on −A·logπ ⇒ ascent on A·logπ
return float(np.linalg.norm(grad))
def kl_to(self, other: ToyPromptRefiner) -> float:
"""KL(π_self ‖ π_other) over the categorical — the LM-PG reference-KL penalty term."""
p, q = self._probs(), other._probs()
return float(np.sum(p * (np.log(p + 1e-12) - np.log(q + 1e-12))))
def clone(self) -> ToyPromptRefiner:
c = ToyPromptRefiner.__new__(ToyPromptRefiner)
c.n, c.dim, c.seq = self.n, self.dim, self.seq
c.logits, c._emb = self.logits.copy(), self._emb.copy()
return c
def copy_from(self, other: ToyPromptRefiner) -> None:
self.logits = other.logits.copy()
def blend_from(self, other: ToyPromptRefiner, decay: float) -> None:
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).
+1 -1
View File
@@ -40,7 +40,7 @@ from collections.abc import Callable, Iterable
# Op names — the stable identifiers a loop dispatches on (not free functions). #
# --------------------------------------------------------------------------- #
FLOW_MATCH_STEP = "flow_match_step" # deterministic flow-match Euler solver step (ODE serve)
FLOW_SDE_STEP = "flow_sde_step" # FlowGRPO stochastic step + log-prob (RL rollout)
FLOW_SDE_STEP = "flow_sde_step" # stochastic diffusion step + log-prob metadata
def _always_available() -> bool:
+8 -7
View File
@@ -38,12 +38,13 @@ def build_bagel_card(model_id: str = "bagel-mot") -> ModelCard:
return ARDecodeLoop(loop_id="generate_text", transformer_id="transformer", max_tokens=6)
def image_factory():
return WanDenoiseLoop(loop_id="generate_image",
cfg=cfg,
flow_shift=flow,
precision=precision,
expert=expert,
)
return WanDenoiseLoop(
loop_id="generate_image",
cfg=cfg,
flow_shift=flow,
precision=precision,
expert=expert,
)
components = {
"tokenizer":
@@ -93,7 +94,7 @@ def build_bagel_card(model_id: str = "bagel-mot") -> ModelCard:
"feature": CacheContract("feature", max_bytes=1 << 24, reuse_across_requests=True),
"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"),
precision=PrecisionContract(default_dtype="float32"),
parallelism=ParallelismContract(valid_plans=[ParallelPlan.single()], default_plan=ParallelPlan.single()),
)
return card.validate()
+1 -2
View File
@@ -1,7 +1,6 @@
"""Shared component-node helpers (text encode with feature cache, etc.).
The content-hash feature cache lets a K-sample RL group encode its shared prompt once instead of
K times.
The content-hash feature cache lets repeated requests reuse prompt embeddings safely.
"""
from __future__ import annotations
+9 -9
View File
@@ -37,12 +37,13 @@ def build_cosmos3_card(model_id: str = "cosmos3-vfm") -> ModelCard:
return ARDecodeLoop(loop_id="ar_decode", transformer_id="transformer", max_tokens=6)
def dn_factory():
return WanDenoiseLoop(loop_id="diffusion_denoise",
cfg=cfg,
flow_shift=flow,
precision=precision,
expert=expert,
)
return WanDenoiseLoop(
loop_id="diffusion_denoise",
cfg=cfg,
flow_shift=flow,
precision=precision,
expert=expert,
)
components = {
"tokenizer":
@@ -89,8 +90,7 @@ def build_cosmos3_card(model_id: str = "cosmos3-vfm") -> ModelCard:
components=components,
loops=loops,
capabilities=CapabilityMatrix.of(Capability.TEXT_TO_VIDEO, Capability.REASONING_TEXT,
Capability.TEXT_TO_VIDEO_SOUND, Capability.VAE_DECODE,
Capability.POLICY_ROLLOUT),
Capability.TEXT_TO_VIDEO_SOUND, Capability.VAE_DECODE),
recipe=RecipeSpec(method="base",
assumes_loop="diffusion_denoise",
assumes_precision="float32",
@@ -100,7 +100,7 @@ def build_cosmos3_card(model_id: str = "cosmos3-vfm") -> ModelCard:
"feature": CacheContract("feature", max_bytes=1 << 24, reuse_across_requests=True),
"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"),
precision=PrecisionContract(default_dtype="float32"),
parallelism=ParallelismContract(valid_plans=[ParallelPlan.single()], default_plan=ParallelPlan.single()),
)
return card.validate()
+1 -1
View File
@@ -107,7 +107,7 @@ def _build_flux2_card(model_id: str, *, text_encoder_load_id: str, text_encoder_
parity=ParitySpec(consistency_levels=[ConsistencyLevel.C1],
tests=[ParityTestSpec(name="denoise_trajectory", level=ConsistencyLevel.C1, tap="latents")]),
caches={"feature": CacheContract(cache_class="feature", max_bytes=1 << 24, reuse_across_requests=True)},
precision=PrecisionContract(default_dtype="float32", training_precision="float32"),
precision=PrecisionContract(default_dtype="float32"),
parallelism=ParallelismContract(valid_plans=[ParallelPlan.single()], default_plan=ParallelPlan.single()),
sampling_defaults=sampling_defaults,
)
+3 -3
View File
@@ -135,7 +135,7 @@ def build_ltx2_card(model_id: str = "ltx2-2stage-distilled") -> ModelCard:
negative_prompt=""),
parity=ParitySpec(consistency_levels=[ConsistencyLevel.C1]),
caches={"feature": CacheContract(cache_class="feature", max_bytes=1 << 24, reuse_across_requests=True)},
precision=PrecisionContract(default_dtype="float32", training_precision="float32"),
precision=PrecisionContract(default_dtype="float32"),
parallelism=ParallelismContract(valid_plans=[ParallelPlan.single()], default_plan=ParallelPlan.single()),
)
return card.validate()
@@ -211,7 +211,7 @@ def build_ltx2_base_card(model_id: str = "ltx2-single-stage") -> ModelCard:
negative_prompt=LTX2_NEG),
parity=ParitySpec(consistency_levels=[ConsistencyLevel.C1]),
caches={"feature": CacheContract(cache_class="feature", max_bytes=1 << 24, reuse_across_requests=True)},
precision=PrecisionContract(default_dtype="float32", training_precision="float32"),
precision=PrecisionContract(default_dtype="float32"),
parallelism=ParallelismContract(valid_plans=[ParallelPlan.single()], default_plan=ParallelPlan.single()),
)
return card.validate()
@@ -301,7 +301,7 @@ def build_ltx2_3_card(model_id: str = "ltx2.3-distilled") -> ModelCard:
}),
parity=ParitySpec(consistency_levels=[ConsistencyLevel.C1]),
caches={"feature": CacheContract(cache_class="feature", max_bytes=1 << 24, reuse_across_requests=True)},
precision=PrecisionContract(default_dtype="float32", training_precision="float32"),
precision=PrecisionContract(default_dtype="float32"),
parallelism=ParallelismContract(valid_plans=[ParallelPlan.single()], default_plan=ParallelPlan.single()),
)
return card.validate()
+1 -1
View File
@@ -127,7 +127,7 @@ def build_matrixgame2_card(model_id: str = "matrix-game-2.0-base-distilled",
parity=ParitySpec(consistency_levels=[ConsistencyLevel.C1],
tests=[ParityTestSpec(name="denoise_trajectory", level=ConsistencyLevel.C1, tap="latents")]),
caches={"feature": CacheContract(cache_class="feature", max_bytes=1 << 24, reuse_across_requests=True)},
precision=PrecisionContract(default_dtype="float32", training_precision="float32"),
precision=PrecisionContract(default_dtype="float32"),
parallelism=ParallelismContract(valid_plans=[ParallelPlan.single()], default_plan=ParallelPlan.single()),
sampling_defaults=sampling_defaults or SamplingDefaults(num_steps=3,
guidance_scale=1.0,
+4 -8
View File
@@ -8,7 +8,7 @@ talker token, and vocoder chunk is a runtime-visible WorkUnit.
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.
* multi-expert inference — TWO disjoint experts, two loops, one request.
* 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.
@@ -51,10 +51,7 @@ def build_qwen_omni_card(model_id: str = "qwen-omni-tts") -> ModelCard:
prompt_slot="talker_prompt_tokens")
def vocoder_factory():
return VocoderLoop(loop_id="vocoder",
vocoder_id="vocoder",
chunk_tokens=2,
speech_slot="speech_tokens")
return VocoderLoop(loop_id="vocoder", vocoder_id="vocoder", chunk_tokens=2, speech_slot="speech_tokens")
components = {
"tokenizer":
@@ -115,8 +112,7 @@ def build_qwen_omni_card(model_id: str = "qwen-omni-tts") -> ModelCard:
family="qwen_omni",
components=components,
loops=loops,
capabilities=CapabilityMatrix.of(Capability.REASONING_TEXT, Capability.TEXT_TO_SPEECH,
Capability.POLICY_ROLLOUT),
capabilities=CapabilityMatrix.of(Capability.REASONING_TEXT, Capability.TEXT_TO_SPEECH),
recipe=RecipeSpec(method="base",
assumes_loop="vocoder",
assumes_precision="float32",
@@ -125,7 +121,7 @@ def build_qwen_omni_card(model_id: str = "qwen-omni-tts") -> ModelCard:
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"),
precision=PrecisionContract(default_dtype="float32"),
parallelism=ParallelismContract(valid_plans=[ParallelPlan.single()], default_plan=ParallelPlan.single()),
)
return card.validate()
+16 -14
View File
@@ -53,12 +53,13 @@ def build_wan21_i2v_card(model_id: str = "wan2.1-i2v-1.3b",
precision, expert = PrecisionPolicy(compute_dtype="float32", scheduler_step_in_fp32=True), NoRouting("transformer")
def loop_factory():
return WanDenoiseLoop(loop_id="i2v_denoise",
cfg=cfg,
flow_shift=flow,
precision=precision,
expert=expert,
)
return WanDenoiseLoop(
loop_id="i2v_denoise",
cfg=cfg,
flow_shift=flow,
precision=precision,
expert=expert,
)
components = {
"text_encoder":
@@ -109,7 +110,7 @@ def build_wan21_i2v_card(model_id: str = "wan2.1-i2v-1.3b",
consistency_required=ConsistencyLevel.C1),
parity=ParitySpec(consistency_levels=[ConsistencyLevel.C1]),
caches={"feature": CacheContract("feature", max_bytes=1 << 24, reuse_across_requests=True)},
precision=PrecisionContract(default_dtype="float32", training_precision="float32"),
precision=PrecisionContract(default_dtype="float32"),
parallelism=ParallelismContract(valid_plans=[ParallelPlan.single()], default_plan=ParallelPlan.single()),
sampling_defaults=SamplingDefaults(num_steps=40,
guidance_scale=5.0,
@@ -138,12 +139,13 @@ def build_wan22_i2v_a14b_card(model_id: str = "wan2.2-i2v-a14b",
expert = BoundaryTimestepRouting(high_noise="transformer", low_noise="transformer_2", boundary=boundary)
def loop_factory():
return WanDenoiseLoop(loop_id="i2v_denoise",
cfg=cfg,
flow_shift=flow,
precision=precision,
expert=expert,
)
return WanDenoiseLoop(
loop_id="i2v_denoise",
cfg=cfg,
flow_shift=flow,
precision=precision,
expert=expert,
)
def _dit(cid: str) -> ComponentSpec:
return ComponentSpec(cid,
@@ -198,7 +200,7 @@ def build_wan22_i2v_a14b_card(model_id: str = "wan2.2-i2v-a14b",
consistency_required=ConsistencyLevel.C1),
parity=ParitySpec(consistency_levels=[ConsistencyLevel.C1]),
caches={"feature": CacheContract("feature", max_bytes=1 << 24, reuse_across_requests=True)},
precision=PrecisionContract(default_dtype="float32", training_precision="float32"),
precision=PrecisionContract(default_dtype="float32"),
parallelism=ParallelismContract(valid_plans=[ParallelPlan.single()], default_plan=ParallelPlan.single()),
sampling_defaults=SamplingDefaults(num_steps=40,
guidance_scale=3.5,
+3 -3
View File
@@ -1,7 +1,7 @@
"""Wan2.1 T2V program: text_encode → diffusion_denoise → vae_decode.
"""Wan2.1 T2V program: text_encode -> diffusion_denoise -> vae_decode.
The text-encode node uses the content-hash feature cache: the same prompt across requests (e.g. a
K=24 RL group) encodes once and reuses.
The text-encode node uses the content-hash feature cache: the same prompt across requests encodes
once and reuses.
"""
from __future__ import annotations
+14 -13
View File
@@ -32,7 +32,7 @@ def build_wan_causal_card(model_id: str = "wan-causal-sf-1.3b",
num_chunks: int = 7,
chunk_size: int = 3,
steps_per_chunk: int = 4,
training_mode: bool = False) -> ModelCard:
preserve_full_context: bool = False) -> ModelCard:
# Self-forcing DMD reference (fastvideo SelfForcingWanT2V480PConfig + causal_denoising.py): the
# distilled student is CFG-FREE (guidance 1.0, single forward/step) and denoises with the 4-step DMD
# schedule (dmd_denoising_steps=[1000,750,500,250], warped — reproduced by FlowShiftPolicy(5.0) over 4
@@ -44,14 +44,15 @@ def build_wan_causal_card(model_id: str = "wan-causal-sf-1.3b",
precision = PrecisionPolicy(compute_dtype="float32", scheduler_step_in_fp32=True)
def loop_factory():
return ChunkRolloutLoop(loop_id="chunk_rollout",
num_chunks=num_chunks,
chunk_size=chunk_size,
steps_per_chunk=steps_per_chunk,
cfg=cfg,
flow_shift=flow,
precision=precision,
)
return ChunkRolloutLoop(
loop_id="chunk_rollout",
num_chunks=num_chunks,
chunk_size=chunk_size,
steps_per_chunk=steps_per_chunk,
cfg=cfg,
flow_shift=flow,
precision=precision,
)
components = {
"text_encoder":
@@ -89,7 +90,7 @@ def build_wan_causal_card(model_id: str = "wan-causal-sf-1.3b",
components=components,
loops=loops,
capabilities=CapabilityMatrix.of(Capability.TEXT_TO_VIDEO, Capability.STREAMING_VIDEO_CONTINUATION,
Capability.VAE_DECODE, Capability.POLICY_ROLLOUT),
Capability.VAE_DECODE),
recipe=RecipeSpec(method="self_forcing",
parents=["wan2.1-1.3b"],
assumes_loop="chunk_rollout",
@@ -103,10 +104,10 @@ def build_wan_causal_card(model_id: str = "wan-causal-sf-1.3b",
CacheContract(cache_class="slab_kv",
max_bytes=1 << 26,
reuse_across_requests=False,
per_component={"window": (num_chunks if training_mode else 2)},
training_mode_disables_recycle=training_mode),
per_component={"window": (num_chunks if preserve_full_context else 2)},
disable_recycle=preserve_full_context),
},
precision=PrecisionContract(default_dtype="float32", training_precision="float32"),
precision=PrecisionContract(default_dtype="float32"),
parallelism=ParallelismContract(valid_plans=[ParallelPlan.single()], default_plan=ParallelPlan.single()),
sampling_defaults=SamplingDefaults(
num_steps=4, guidance_scale=1.0, height=480, width=832, num_frames=81, fps=16,
+5 -7
View File
@@ -1,9 +1,9 @@
"""ChunkRolloutLoop — causal/AR video.
Loop shape: outer loop over latent-frame chunks × inner denoise per chunk, carrying a per-chunk
KV/context via the slab-KV cache (sink/local window in inference; full history in self-forcing
training mode). Each completed chunk is streamable (``StepResult.emit``), and the rollout profile
captures per-chunk behavior for the self-forcing distillation method.
Loop shape: outer loop over latent-frame chunks x inner denoise per chunk, carrying a per-chunk
KV/context via the slab-KV cache. The cache can use a sink/local window or preserve full request
history. Each completed chunk is streamable (``StepResult.emit``), and the rollout profile can
capture per-chunk inference trajectory records.
State is entirely in ``LoopState`` (the slab list + chunk position), so interleaving causal
rollouts of two sessions cannot smear context — and the slab-KV pool is namespaced per request.
@@ -151,9 +151,7 @@ class ChunkRolloutLoop:
seq=st.scratch["chunk_idx"],
data=x,
preview=False))
res = ResourceRequest(
resident_bytes=int(x.nbytes),
peak_activation_bytes=int(x.nbytes * len(branches)))
res = ResourceRequest(resident_bytes=int(x.nbytes), peak_activation_bytes=int(x.nbytes * len(branches)))
return WorkPlan(loop_id=self.loop_id,
instance_id=st.instance_id,
kind=WorkUnitKind.CHUNK_STEP,
+6 -6
View File
@@ -4,7 +4,7 @@ No single unified block pool: cache classes differ by 150-500x in natural granul
each gets its own statically budgeted pool behind one ``CacheHandle``. Four classes:
* ``FeatureCache`` — content-hash keyed, partitioned by adapter+weights (text/vision encoders)
* ``ResidualCache`` — cache-dit residuals, scoped per request AND per CFG branch
* ``SlabKVCache`` — chunk-KV slabs (self-forcing / world models); training mode disables recycle
* ``SlabKVCache`` — chunk-KV slabs for causal/world models; can disable recycle
* ``PagedKVCache`` — paged text-KV stub for ar_decode (phase-2 omni); minority case, lazy
KV is the minority case — a pure bidirectional deployment (Wan/LTX T2V) allocates none of it.
@@ -71,7 +71,7 @@ class FeatureCache(_Pool):
self.used_bytes -= nb
def invalidate_weights(self, version: str) -> None:
"""RL update_weights bumps weight epoch → drop entries from older epochs (wholesale)."""
"""Drop entries from older resident-weight epochs."""
drop = [h for h, (_v, _nb, k) in self._store.items() if k.weights_version != version]
for h in drop:
_v, nb, _k = self._store.pop(h)
@@ -124,21 +124,21 @@ class Slab:
class SlabKVCache(_Pool):
"""Chunk-KV slabs for causal/world-model rollout.
``training_mode`` disables mid-rollout recycling so activation-checkpoint recompute
doesn't double-advance the cache (self-forcing).
``disable_recycle`` preserves the full request-local context. Otherwise the
cache behaves as a sliding window with the declared ``window`` size.
"""
def __init__(self, policy: CachePolicy):
super().__init__(policy)
self._store: dict[str, list[Slab]] = {}
self.window = max(1, policy.per_component.get("window", 1 << 30))
self.training_mode = policy.training_mode_disables_recycle
self.disable_recycle = policy.disable_recycle
def append(self, namespace: str, slab: Slab) -> None:
slabs = self._store.setdefault(namespace, [])
slabs.append(slab)
self.used_bytes += _nbytes(slab.k) + _nbytes(slab.v)
if not self.training_mode and len(slabs) > self.window:
if not self.disable_recycle and len(slabs) > self.window:
dropped = slabs.pop(0) # sliding-window recycle (inference only)
self.used_bytes -= _nbytes(dropped.k) + _nbytes(dropped.v)
+3 -4
View File
@@ -2,8 +2,7 @@
If a field can change output semantics, it is in the key (incorrect reuse is worse than no reuse).
The serving hazard this kills: a request that shares a prompt but differs in te-LoRA stack must not
serve stale embeddings — so the key is *partitioned* by ``adapter_versions``, not flushed. An RL
``update_weights`` bumps ``weights_version`` and invalidates wholesale.
serve stale embeddings, so the key is *partitioned* by ``adapter_versions``.
"""
from __future__ import annotations
@@ -15,7 +14,7 @@ from typing import Any
def content_hash(obj: Any) -> str:
"""Stable content hash for feature-cache keys (text/vision embeddings).
Lets K rollout samples of one prompt reuse a single text encode.
Lets repeated requests for one prompt reuse a single text encode.
"""
h = hashlib.sha256()
if isinstance(obj, str):
@@ -78,4 +77,4 @@ class CachePolicy:
eviction: str = "lru" # "lru" | "fifo" | "none"
reuse_across_requests: bool = True
per_component: dict[str, int] = field(default_factory=dict)
training_mode_disables_recycle: bool = False
disable_recycle: bool = False
+2 -2
View File
@@ -2,7 +2,7 @@
Static partitioning makes cross-class fragmentation impossible (jumbo slab traffic cannot
strand text-KV pages and vice versa). Each class gets a budget carved at init from the card's
CacheContracts. ``invalidate_weights`` implements the wholesale RL weight-epoch bump.
CacheContracts.
"""
from __future__ import annotations
@@ -32,7 +32,7 @@ class CacheManager:
eviction=cc.eviction,
reuse_across_requests=cc.reuse_across_requests,
per_component=dict(cc.per_component),
training_mode_disables_recycle=cc.training_mode_disables_recycle,
disable_recycle=cc.disable_recycle,
))
return cls(policies)
+2 -2
View File
@@ -19,7 +19,7 @@ correctness-critical and bug-prone regardless of device:
* the **eager-break** — a step that is model-declared non-capturable (host RNG / data-dependent
branch, e.g. the FlowGRPO SDE step) or that an interceptor overrode this step never touches the
graph cache (vLLM's ``@eager_break_during_capture``).
* **invalidation** — a resident-weights version bump (RL weight sync) or a shape change changes the
* **invalidation** — a resident-weights version bump or a shape change changes the
key → recapture, so a graph never outlives the weights/shape it was captured for. This is
AUTOMATIC: the resident-weight version is *in* the key, so a synced component's old graphs simply
become unreachable and the next step recaptures.
@@ -159,7 +159,7 @@ class GraphCapturer:
return plan.graph_fn(instance, g.workspace)
def invalidate(self, component_ids: Any = None) -> int:
"""Evict captured graphs whose resident weights just changed (RL weight sync), mirroring
"""Evict captured graphs whose resident weights just changed, mirroring
``CacheManager.invalidate_components``. Without this, a synced component's old graphs become
unreachable (the version is in the key) but linger — a real GPU memory leak across FlowGRPO's
per-iteration syncs. With ``component_ids=None``, clears all. Returns the count dropped."""
+1 -1
View File
@@ -3,7 +3,7 @@
Reservation must be pre-flight: a scheduler that admits a diffusion step and then discovers it
cannot allocate the VAE tile is wrong.
Sleep/wake is component-granular (tags are component names) for RL: drop DiT + caches,
Sleep/wake is component-granular (tags are component names): drop DiT + caches,
keep VAE/text-encoder resident (CuMem-style).
"""
from __future__ import annotations
+5 -6
View File
@@ -1,4 +1,4 @@
"""Per-class cache pools: feature partition/invalidation, slab-KV window/training."""
"""Per-class cache pools: feature partition/invalidation and slab-KV windows."""
from __future__ import annotations
import numpy as np
@@ -23,7 +23,7 @@ def test_feature_cache_weight_invalidation_is_wholesale():
fc = FeatureCache(CachePolicy("feature", max_bytes=1 << 20))
k = CacheKey(model_id="m", component_id="te", weights_version="v0")
fc.put(k, np.ones((2,), dtype="float32"))
fc.invalidate_weights("v1") # RL update_weights bump
fc.invalidate_weights("v1")
assert fc.get(k) is None
@@ -35,17 +35,16 @@ def test_feature_cache_budget_eviction():
assert fc.used_bytes <= 8
def test_slab_kv_inference_window_vs_training_mode():
def test_slab_kv_window_vs_disable_recycle():
inf = SlabKVCache(CachePolicy("slab_kv", reuse_across_requests=False, per_component={"window": 2}))
for i in range(4):
inf.append("r", Slab(i, np.zeros((3,), dtype="float32"), None))
assert len(inf.get("r")) == 2 # sliding window recycles (inference)
tr = SlabKVCache(CachePolicy("slab_kv", training_mode_disables_recycle=True,
per_component={"window": 2}))
tr = SlabKVCache(CachePolicy("slab_kv", disable_recycle=True, per_component={"window": 2}))
for i in range(4):
tr.append("r", Slab(i, np.zeros((3,), dtype="float32"), None))
assert len(tr.get("r")) == 4 # no mid-rollout recycle in training mode (self-forcing)
assert len(tr.get("r")) == 4 # no mid-rollout recycle when full context is requested
assert tr.used_bytes > 0
tr.clear_namespace("r")
assert tr.used_bytes == 0
+1 -1
View File
@@ -89,7 +89,7 @@ def test_weight_version_bump_recaptures_and_evicts_stale():
inst = _wan()
_drive(inst, seed=1, steps=4)
assert inst.graphs.stats["captures"] == 1 and inst.graphs.captured == 1
inst.set_weights_version("v1", ["transformer"]) # RL weight sync: bumps version + evicts graph
inst.set_weights_version("v1", ["transformer"]) # version bump evicts stale graphs
assert inst.graphs.captured == 0 # stale graph dropped (no leak)
_drive(inst, seed=1, steps=4)
assert inst.graphs.stats["captures"] == 2 # recaptured under the new version...
+2 -2
View File
@@ -43,12 +43,12 @@ def test_feature_key_partitions_by_adapter_stack():
assert inst.caches.stats()["feature"]["misses"] == m + 1
def test_transformer_weight_sync_keeps_text_encoder_cache():
def test_transformer_weight_reload_keeps_text_encoder_cache():
card = build_wan21_card()
inst = load_card(card, cache_manager=CacheManager.from_card(card))
cached_text_encode(inst, "p")
assert inst.caches.stats()["feature"]["used_bytes"] > 0
inst.set_weights_version("w1", components=["transformer"]) # RL-style transformer-only sync
inst.set_weights_version("w1", components=["transformer"]) # transformer-only reload
h = inst.caches.stats()["feature"]["hits"]
cached_text_encode(inst, "p") # text-encoder embed survived ⇒ hit (not flushed)
assert inst.caches.stats()["feature"]["hits"] == h + 1
+1 -3
View File
@@ -1,8 +1,7 @@
"""Qwen-Omni thinker→talker→vocoder — the third weight-sharing topology (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:
Asserts the Card/Loop/Program vocabulary holds for a THREE-expert, THREE-loop cascade:
* 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;
@@ -102,4 +101,3 @@ def test_vocoder_streams_audio_chunks():
# 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