Compare commits
29
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
82b6d2dbf0 | ||
|
|
c6468e8def | ||
|
|
2c0d9a0828 | ||
|
|
b9876ee06d | ||
|
|
a568406b8c | ||
|
|
39cc075452 | ||
|
|
88f50fa7d6 | ||
|
|
4b25087bd8 | ||
|
|
a314489b78 | ||
|
|
a1b8a78e9e | ||
|
|
df55a31824 | ||
|
|
d022cf00e3 | ||
|
|
7dc3166827 | ||
|
|
ce4c4f37c7 | ||
|
|
48dff4be69 | ||
|
|
d9736d5e5d | ||
|
|
7396024f95 | ||
|
|
978896f534 | ||
|
|
560810c592 | ||
|
|
e3a7450a05 | ||
|
|
6aab7f3832 | ||
|
|
9cd53fe5f8 | ||
|
|
6a32cf3a5e | ||
|
|
98be9b3da2 | ||
|
|
c53e85b767 | ||
|
|
40a8bd2d3b | ||
|
|
31aa115611 | ||
|
|
98ac10a528 | ||
|
|
a5a6d171e5 |
@@ -1,207 +0,0 @@
|
||||
# v2 ← M\*: Architecture Gap-Analysis & Improvement Roadmap
|
||||
|
||||
**Status:** exploration, flagged for review. **Date:** 2026-06-19.
|
||||
**Source paper:** *M\*: A Modular, Extensible, Serving System for Multimodal Models* (arXiv 2606.12688,
|
||||
Stanford/UW/CMU; Jha, Sagan, Kamahori, …, Kasikci, S. Wang). It is a universal serving runtime for composite
|
||||
multimodal models built on the **Walk Graph** abstraction (a model is a dataflow graph `G`; a request is a
|
||||
*Walk* — a labeled subgraph — and the runtime executes walks). It beats vLLM-Omni (~20% lower T2I latency on
|
||||
**BAGEL**, up to 2.64× on I2I), SGLang-Omni (2.7× TTS throughput on **Qwen3-Omni**), and native V-JEPA2
|
||||
rollout (12.5×). It explicitly names **FastVideo's own** sparse/sliding-tile attention, xDiT/PipeFusion/USP,
|
||||
Inferix, and FlashDrive as techniques integratable into the graph runtime.
|
||||
|
||||
**Method:** a 28-agent workflow — 6 parallel v2-subsystem maps → 10 M\*-dimension analyses, each
|
||||
*adversarially verified against the actual v2 code* → synthesis + a completeness critic. The critic's
|
||||
corrections and three P0 claims were then **spot-verified by hand** (file:line below). This doc folds those
|
||||
corrections in; it is the corrected, authoritative synthesis.
|
||||
|
||||
---
|
||||
|
||||
## 1. Executive summary
|
||||
|
||||
v2 already implements the **harder half** of M\*'s thesis and in several axes **exceeds** it:
|
||||
|
||||
- v2's `Program` *is* M\*'s graph `G` (typed `ComponentNode`/`ModelLoopNode` + edges).
|
||||
- v2's `shared_weight_components` *is* M\*'s cross-Walk node sharing — BAGEL/Cosmos3/LTX2 each bind two
|
||||
`ModelLoopNode`s to **one resident transformer** (`instance.component()` returns the same live object). This
|
||||
is the exact MoT serving property the omni cards in this repo already express.
|
||||
- v2 adds three things M\* (serving-only) has **no equivalent for**: a required+validated per-loop **cost
|
||||
model**, a non-negotiable **interleave bit-parity gate**, and an **integrated training plane** (RL→distill
|
||||
flywheel driving the *same* serving Loop).
|
||||
- The `extend/` plugin seam (interceptors/observers/registry with capability negotiation) is precisely the
|
||||
hook M\*'s "extensible / integrate FastVideo-STA, xDiT, Inferix, FlashDrive" call-out asks for — **v2
|
||||
already has the seam M\* only gestures at.**
|
||||
|
||||
What v2 lacks is M\*'s **declarative authoring layer above the substrate**, and — the key insight — *much of
|
||||
that substrate is already authored but inert*: v2 has declared the metadata for "minimum components per
|
||||
request" (`required_for`/`optional_for` on every omni card) and "branch as a cache axis" (`guidance_sig`,
|
||||
`CacheKey`) but **never wired it to an executor**. The substrate is ~80% built and switched off.
|
||||
|
||||
**Highest-leverage cluster:** three small, parity-safe wires that turn on inert substrate and unblock the
|
||||
BAGEL/Qwen-Omni/Cosmos3 latency wins M\* measured **on the exact models this repo already runs** — plus one
|
||||
P1 that aligns v2 with the paper's headline "extensible" claim using a seam v2 already has.
|
||||
|
||||
### Verified P0 correctness findings (spot-checked by hand)
|
||||
1. **Runner divergence (real bug).** `v2/runtime/engine.py:88` → `nodes = self.program.nodes`;
|
||||
`v2/runtime/disaggregated.py:96` → `nodes = self.program.active_nodes(self.request)`. The inline and
|
||||
disaggregated runners execute *different node sets*. ✅ confirmed.
|
||||
2. **EOS is faked.** `v2/recipes/omni/ar_loop.py` docstring says "done on EOS/max_tokens"; `next()` (`:46-48`)
|
||||
checks **only** `max_tokens`. M\*'s marquee `DynamicLoop` use case (EOS) is unimplemented in the loop that
|
||||
serves the Qwen-Omni Thinker/Talker and Cosmos3 reasoner. ✅ confirmed.
|
||||
3. **`required_for`/`optional_for` have zero runtime consumers** (grep outside `specs.py`/recipes/tests is
|
||||
empty). The min-components metadata is declared on every card and never read. ✅ confirmed.
|
||||
|
||||
---
|
||||
|
||||
## 2. Dimension table (corrected)
|
||||
|
||||
| # | Dimension | v2 status | Gap | Priority | Effort | Payoff | Action |
|
||||
|---|---|---|---|---|---|---|---|
|
||||
| 1 | Min-components per request (`required_for` + `when_task`) | substrate built, **inert** | real, cheap | **P0** | S | Consume `required_for` in `active_nodes`; unify `engine.py:88` onto `active_nodes`; deliver via registry/card builder so all ~40 cards inherit it |
|
||||
| 2 | Real EOS + declarative `DynamicLoop` | early-exit emergent; **EOS faked** | real | **P0** | S | `ARDecodeLoop` honors `eos_id` + `req.sampling.stop`; add `LoopSpec.dynamic_stop` + `register_loop_stop`. **Training-enabling** (world-model rollout horizon) |
|
||||
| 3 | CFG/branch as label over one paged KV pool | absent (`PagedKVCache` is a counter) | real | **P1** | L | `(namespace,label)` paged store w/ one budget; reuse `guidance_sig` for hash (NOT `partition_field`); by-ref via existing `InProcKVConnector`. AR path only (diffusion has no KV) |
|
||||
| 4 | `extend/` plugin seam → integrate FastVideo-STA / Inferix | **seam exists, unused for attn** | real (paper headline) | **P1** | M | Expose FastVideo sparse/sliding-tile attention + Inferix block-diffusion as `Interceptor`/`EngineKind` plugins — the paper's named integration targets, on this repo's own code |
|
||||
| 5 | `ParitySpec.output_determinism` (C3 distributional) | C3 rung defined, **0 users** | real, dormant | **P1** | S | Add field; `compare_outputs` consults it. **Training-enabling** (SDE/FlowGRPO stochastic rollouts) |
|
||||
| 6 | Registry-driven delivery of #1 | present, not leveraged | integration | **P1** | S | Express `when_task`/min-components through `WorkflowRegistry`/card builders, not 3 bespoke recipe patches |
|
||||
| 7 | Serving conductor + pluggable data plane | conductor exists (`serving/http.py`); **single-process transport** | real | **P2** | L | v2 already has the step-scheduled worker surface; gap is ZeroMQ/Mooncake + direct worker→worker tensor routing (today `InProcKVConnector` only) |
|
||||
| 8 | Fleet/Dynamo placement + replicas | **live** (`deploy/fleet.py`,`dynamo.py`) | partial | **P2** | M | Fleet-level placement/affinity/replica is real & ≥M\*; missing piece is only the intra-engine `(node,Walk)→rank` map decoupled from model code |
|
||||
| 9 | Per-node TP / SP + cross-rank transport | axis vocab **exists** (`sp` incl.); not wired to runtime | partial | **P2** | XL | Wire declarative degrees into runtime; Wan/LTX are **SP-native** (TP is a no-op there); populate `parallel_plan_hash` on the serving cache path |
|
||||
| 10 | Named Walks + per-model state machine | `Program`=G, sharing real; no Walk/SM | real | **P2** | M | Defer until a *re-entrant* phase graph (Thinker↔Talker, rollout) needs it; #1 captures the min-components win without it |
|
||||
| 11 | Declarative `Parallel/Sequential/Loop` IR | imperative loop classes | real (authoring) | **P2** | M | Thin Section IR lowering to flat `Program`; scope to one AR recipe |
|
||||
| 12 | Streaming `ChunkPolicy` + `StreamBuffer` | causal-chunk emit **already ships** (`wan_causal`); `EdgeKind.STREAM` inert | real | **P2** | L | Declarative `ChunkPolicy` vocab over the existing chunk mechanism; needs concurrent producer/consumer runner (= pipelined scheduling). Inferix integration point |
|
||||
| 13 | Speculative deferred-termination; loop-spanning CUDA graphs; N+1 prefetch; attn double-buffer | absent / per-step capture (14 cards) | real | **P3** | L | Gate behind a real GPU executor; unobservable on CPU-toy CI; loop-span needs an `allows_interleaving=False` carve-out |
|
||||
| — | Cost model + interleave/consistency parity | **exceeds M\*** | none | **guard** | — | Do not regress; keep `step_cost_model` mandatory + `bit_identical` default |
|
||||
| — | Integrated training plane (flywheel, weight-sync) | **exceeds M\*** | none | **guard** | — | Protect train==serve loop identity with a toy fixture |
|
||||
|
||||
---
|
||||
|
||||
## 3. P0/P1 deep-dives (sequenced)
|
||||
|
||||
```
|
||||
PR-1 (P0) min-components ──┐
|
||||
PR-2 (P0) real EOS ─┼─► prereqs for honest "DynamicLoop" + min-component claims; both training-enabling
|
||||
PR-3 (P1) output_determinism (independent)
|
||||
PR-5 (P1) extend/ plugin: FastVideo-STA / Inferix as Interceptors (independent; highest paper-alignment)
|
||||
PR-4 (P1) CFG-as-label paged pool ──► depends on PR-2 (AR loop is the only KV consumer)
|
||||
```
|
||||
PR-1, PR-2, PR-3, PR-5 are mutually independent; PR-4 depends on PR-2.
|
||||
|
||||
### PR-1 (P0) — Turn on the inert min-components substrate + fix runner divergence
|
||||
- **Change.** Extend `Program.active_nodes(request)` (`v2/program/specs.py`) to also drop any node whose bound
|
||||
`ComponentSpec.required_for` (`v2/card/specs.py:144`) excludes `request.task` (and isn't in `optional_for`).
|
||||
**Fix the bug:** change `v2/runtime/engine.py:88` to `nodes = self.program.active_nodes(self.request)` so the
|
||||
inline `ProgramRunner` matches `DisaggregatedRunner` (`disaggregated.py:96`). Deliver the `when_task` gating
|
||||
through the **registry/card builder** (`recipes/__init__.py`, `program/workflow.py:WorkflowRegistry`) so all
|
||||
~40 cards inherit it uniformly — not three bespoke `program.py` patches.
|
||||
- **Why (this repo's models).** BAGEL T2I currently steps the AR-text loop and Cosmos3 t2v materializes the
|
||||
reasoner even though the cards declare `transformer required_for={'reason','t2i'}`, `vae required_for={'t2i'}`.
|
||||
On the GPU backend that is wasted resident-weight load + wasted steps on every single-modality request —
|
||||
exactly M\*'s "execute the MINIMUM components per request," delivered by consuming existing metadata.
|
||||
- **Risk/invariant.** Validate in `ModelCard.validate()` that every active node's `reads` are produced by an
|
||||
active node for each declared `TaskType` (avoid dropping a producer). Pure node-id filtering ⇒ serial and
|
||||
interleaved still walk the same filtered list ⇒ §9.3 interleave bit-parity holds by construction. CPU-toy clean.
|
||||
|
||||
### PR-2 (P0) — Real EOS + declarative `dynamic_stop` *(also training-enabling)*
|
||||
- **Change.** In `v2/recipes/omni/ar_loop.py`, `advance()` reads the emitted token; if it equals the model
|
||||
`eos_id` (toy backend exposes `EOS=0`) or matches `req.sampling.stop` (`params.py:21`, currently dead),
|
||||
register termination; `next()` returns `Done()` on stop OR `max_tokens`. Add `StopRegistry` to `LoopState` +
|
||||
`register_loop_stop(name)` to the `LoopContext` protocol (`contracts.py:204`) and to
|
||||
`DisaggregatedRunner`'s `RuntimeLoopContext`. Add `LoopSpec.dynamic_stop: bool=False`, opt the AR cards in.
|
||||
- **Why.** The docstring-vs-code lie sits in the loop serving Qwen-Omni Thinker/Talker and the Cosmos3 reasoner;
|
||||
M\*'s second named `DynamicLoop` use case (world-model **rollout horizon**) is exactly what `self_forcing` RL
|
||||
needs — so this is both a serving-credibility fix and a training enabler (raise its payoff accordingly).
|
||||
- **Risk/invariant.** `dynamic_stop=False` is byte-identical back-compat. Must pass **all three** parity gates:
|
||||
serial==interleaved AND disaggregated==inline. **Not** in this PR: speculative deferred-termination (unobservable
|
||||
on CPU-toy, fights the interleave invariant — P3, gated on GPU executor).
|
||||
|
||||
### PR-3 (P1) — `ParitySpec.output_determinism` (close the dormant C3 hole) *(training-enabling)*
|
||||
- **Change.** Add `output_determinism: str = "bit_identical"` to `ParitySpec` (`card/specs.py:88`); make
|
||||
`compare_outputs` (`parity/interleave_gate.py:54`) consult it (`bit_identical` → today's exact check;
|
||||
`distributional` → a moment/tolerance check — land a simple moment match first; a real KS test is new code).
|
||||
- **Why.** `ConsistencyLevel.C3` is defined and used by zero recipes; an SDE/FlowGRPO stochastic rollout cannot
|
||||
honestly declare its parity contract and would falsely fail the bit-identical gate. Additive; default unchanged.
|
||||
|
||||
### PR-5 (P1) — Expose FastVideo's own attention + Inferix as `extend/` plugins *(highest paper-alignment)*
|
||||
- **Change.** Use the existing `extend/{interceptors,observers,registry}.py` seam (capability-negotiated, with
|
||||
per-(request,branch) `plugin_state` that already passes the interleave gate) to register FastVideo's
|
||||
sparse/sliding-tile attention and Inferix-style block-diffusion as `Interceptor`s / an `EngineKind` plugin.
|
||||
- **Why.** M\*'s title is "Modular, **Extensible**" and it explicitly lists FastVideo-STA, xDiT/PipeFusion/USP,
|
||||
Inferix, FlashDrive as integratable. v2 already has the seam M\* only describes — this is where v2 most
|
||||
directly answers the paper, using this repo's own attention code. Low risk (the seam + capability negotiation
|
||||
already exist and are tested).
|
||||
|
||||
### PR-4 (P1) — CFG/branch as a LABEL over one paged KV pool
|
||||
- **Change.** Rewrite `PagedKVCache` (`cache/classes.py:155-172`) from a block *counter* into a real
|
||||
`(namespace,label)->[block-handle]` store with **one shared `total_blocks` budget** (M\*'s single-pool
|
||||
property). Reuse the existing-but-unpopulated `CacheKey.guidance_sig` (`keys.py:53`) for the hash. Thread the
|
||||
label through `ar_loop.py` (alloc/append/get per `(request_id, branch)`; prefill once per shared-prefix label;
|
||||
combine via `CFGPolicy.combine`). Wire `ResourceRequest.cache_blocks` (`contracts.py:64`, zero consumers) into
|
||||
admission per (class,label).
|
||||
- **Why.** The dossier-identified driver of M\*'s BAGEL win (3 CFG contexts as 3 labels over ONE pool vs dense
|
||||
per-context). Targets AR_DECODE (BAGEL `generate_text`, omni Thinker); **correctly excludes diffusion**
|
||||
(Wan/LTX are bidirectional, no KV — their CFG stays dense-but-batched).
|
||||
- **Corrections to bake in.** Do **NOT** add `branch_label` to `CacheKey.partition_field()` (CFG branches share
|
||||
embeddings; partitioning by branch is a semantic bug). Do **NOT** add a new by-ref type — reuse
|
||||
`InProcKVConnector` + `TransferManifest.cache_key`. Wiring `cache_blocks` admission is greenfield ⇒ effort **L**.
|
||||
CPU version proves label/sharing semantics; the real latency win needs a FlashInfer paged kernel (out of scope)
|
||||
— **merge** with a future "real KVCacheEngine" effort rather than landing isolated.
|
||||
|
||||
---
|
||||
|
||||
## 4. What v2 already does ≥ M\* — do NOT regress
|
||||
1. **Required+validated cost model** on every `LoopSpec` (13-kind `WorkUnitKind`) — typed, pre-GPU-validated.
|
||||
2. **Interleave bit-parity as a hard gate** (`parity.interleave_required=True` on 40+ cards). M\* has no such
|
||||
gate (its speculative scheduling deliberately wastes steps). Load-bearing invariant; every new primitive
|
||||
must pass it.
|
||||
3. **C0–C4 consistency ladder** wired into RL methods, with first-divergence tap reporting. No M\* equivalent.
|
||||
4. **Integrated training plane** — DiffusionNFT/DMD2/self_forcing, RL→distill flywheel, `WeightSyncController`
|
||||
hot weight-sync with drain-to-boundary + scoped cache invalidation, driving the **same** serving Loop.
|
||||
M\* is serving-only. Protect with a toy fixture asserting `rollout_loop` drives the served Loop object.
|
||||
5. **CPU-toy parity for the whole stack** — loops/CFG/caches/parity/RL run in CI without a GPU. Every new
|
||||
primitive must ship a toy exercise (this is what makes all PRs above testable without H100s).
|
||||
6. **Partition-not-flush cache invalidation** + four independent per-class pools.
|
||||
7. **`extend/` plugin seam** with capability negotiation (a 4-step distilled card *rejects* a residual-skip
|
||||
interceptor) — M\* describes extensibility; v2 has the mechanism.
|
||||
8. **Dynamo citizenship** (`deploy/dynamo.py`: one `DeploymentCard`+cost model, two consumers) — beyond M\*'s
|
||||
self-contained runtime.
|
||||
|
||||
---
|
||||
|
||||
## 5. Dropped / merged / deferred (and why)
|
||||
- **DROP declarative `Parallel` as a CFG-execution win.** The runner walks nodes linearly (ignores
|
||||
`Program.edges`), so `Parallel` lowers to sequential sugar and the CFG 3-pass braid is already one
|
||||
co-scheduled `WorkPlan.run`; splitting it risks the interleave gate. Salvage only the no-op refactor
|
||||
extracting `branch_forward` from `WanDenoiseLoop._velocity`. Reassign `Parallel` to the placement workstream.
|
||||
- **MERGE the full Walk/state-machine layer** into "defer until a re-entrant phase graph needs it" (PR-1 gets the
|
||||
min-components win with ~20 lines, no new abstraction). If built: the validator must check a walk's node-id
|
||||
order is a *subsequence* of `program.nodes` (not just membership) or the runner can reorder and break parity.
|
||||
- **MERGE `StreamBuffer`/`ChunkPolicy` into pipelined-scheduling.** Causal-chunk emit *already ships*
|
||||
(`wan_causal/loop.py` per-chunk `StepResult.emit` + slab-KV); the gap is the declarative `ChunkPolicy` vocab
|
||||
+ a concurrent producer/consumer runner. If built: keep all policies pure (per-request `StreamBuffer` history,
|
||||
not shared edge state) and restrict the bit-identical claim to the token-only handoff.
|
||||
- **MERGE CFG-fan-out exec + cross-rank transport + PD loop-splitting into a multi-GPU-runtime program.** These
|
||||
need real collectives (`v2/distributed/` is a stub) and KV-by-reference (KV lives in `CacheManager`, not the
|
||||
transferable `slots`). **Keep cheaply now:** the *declarative* halves — per-component degree, `(node,Walk)`
|
||||
placement key with node-only fallback, `ReplicaSet` under `LocalFleet`, and populate `parallel_plan_hash` on
|
||||
the **serving** cache path (it is already populated in `training/behavior.py:40` — the gap is serving-only).
|
||||
- **DEFER** speculative deferred-termination, loop-spanning CUDA graphs, N+1 prefetch, attention-plan
|
||||
double-buffer — all gated on a real GPU executor; benefit unobservable on CPU-toy CI. Keep the cheap
|
||||
`EngineKind` tag (`STATELESS|KV_CACHE|DIFFUSION`) now. Correct the stale `cudagraph.py:51-52` docstring
|
||||
(per-step capture ships in 14 cards, not just wan21).
|
||||
- **RESCOPE per-node TP.** Wan/LTX use `ReplicatedLinear` + **sequence parallelism** (`sp`), not TP; the `sp`
|
||||
axis already exists in `parallel/plan.py:AXIS_NAMES`. The work is wiring degrees into the runtime, not
|
||||
inventing vocabulary; a `tp_size=2` "one-line activation" is a no-op for the shipped models.
|
||||
|
||||
---
|
||||
|
||||
## 6. The first integration test, if/when multi-GPU placement work starts
|
||||
The **live Qwen-Omni 2-GPU bring-up** (Thinker on rank 0, Talker+Code2Wav on rank 1; see
|
||||
`v2_debug_videos/vlm.md` Session 4) is the natural first validation target for any `(node,Walk)→rank`
|
||||
placement work — it is the one place this repo already has real multi-rank composite-model execution.
|
||||
|
||||
---
|
||||
|
||||
## Anchor files for P0/P1
|
||||
`v2/program/specs.py`, `v2/runtime/engine.py` (**line 88 fix**), `v2/runtime/disaggregated.py`,
|
||||
`v2/recipes/omni/ar_loop.py`, `v2/loop/contracts.py`, `v2/card/specs.py`, `v2/cache/{classes.py,keys.py}`,
|
||||
`v2/parity/interleave_gate.py`, `v2/extend/{interceptors,registry}.py`, `recipes/__init__.py` +
|
||||
`v2/program/workflow.py` (registry-driven delivery).
|
||||
@@ -1,5 +1,9 @@
|
||||
{
|
||||
"benchmark_id": "wan-t2v-1.3b-2gpu",
|
||||
"config_schema_version": 2,
|
||||
"workload_id": "wan-t2v-1.3b",
|
||||
"variant_id": "canonical",
|
||||
"benchmark_version": 1,
|
||||
"description": "Wan2.1 T2V 1.3B inference performance",
|
||||
"model": {
|
||||
"model_path": "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
|
||||
|
||||
@@ -114,6 +114,17 @@ steps:
|
||||
limit: 2
|
||||
agents:
|
||||
queue: "default"
|
||||
- label: ":test_tube: LoRA Extraction Tests"
|
||||
if: build.env("TEST_SCOPE") == "direct" && build.env("TEST_TYPE") == "lora_extraction"
|
||||
command: "timeout 90m .buildkite/scripts/pr_test.sh"
|
||||
retry:
|
||||
automatic:
|
||||
- exit_status: 128
|
||||
limit: 3
|
||||
- exit_status: -1
|
||||
limit: 2
|
||||
agents:
|
||||
queue: "default"
|
||||
- label: ":test_tube: Training Tests"
|
||||
if: build.env("TEST_SCOPE") == "direct" && build.env("TEST_TYPE") == "training"
|
||||
command: "timeout 90m .buildkite/scripts/pr_test.sh"
|
||||
@@ -371,6 +382,21 @@ steps:
|
||||
- TEST_TYPE=inference_lora
|
||||
agents:
|
||||
queue: "default"
|
||||
- path:
|
||||
- "scripts/lora_extraction/**"
|
||||
- "fastvideo/tests/lora_extraction/**"
|
||||
- "fastvideo/models/loader/**"
|
||||
- "fastvideo/training/training_utils.py"
|
||||
- "fastvideo/layers/lora/**"
|
||||
- "pyproject.toml"
|
||||
- "docker/Dockerfile"
|
||||
config:
|
||||
command: "timeout 90m .buildkite/scripts/pr_test.sh"
|
||||
label: ":test_tube: LoRA Extraction Tests"
|
||||
env:
|
||||
- TEST_TYPE=lora_extraction
|
||||
agents:
|
||||
queue: "default"
|
||||
- path:
|
||||
- "fastvideo/**"
|
||||
- "pyproject.toml"
|
||||
|
||||
@@ -125,7 +125,7 @@ jobs:
|
||||
set -euo pipefail
|
||||
TEST_NAME=$(echo "$COMMENT" | grep -oP '(?<=/test\s)\S+' | head -1 || true)
|
||||
|
||||
VALID="encoder vae transformer kernel unit dreamverse ssim training lora-inference lora-training distillation self-forcing vsa vmoba performance api train-framework eval full fastcheck pre-commit"
|
||||
VALID="encoder vae transformer kernel unit dreamverse ssim training lora-inference lora-training lora-extraction distillation self-forcing vsa vmoba performance api train-framework eval full fastcheck pre-commit"
|
||||
if [ -z "$TEST_NAME" ] || ! echo "$VALID" | grep -qw "$TEST_NAME"; then
|
||||
echo "Unknown test: '$TEST_NAME'. Valid: $VALID"
|
||||
exit 1
|
||||
@@ -136,6 +136,7 @@ jobs:
|
||||
[kernel]=kernel_tests [unit]=unit_test [dreamverse]=dreamverse_app
|
||||
[ssim]=ssim [training]=training
|
||||
[lora-inference]=inference_lora [lora-training]=training_lora
|
||||
[lora-extraction]=lora_extraction
|
||||
[distillation]=distillation_dmd [self-forcing]=self_forcing
|
||||
[vsa]=training_vsa [vmoba]=inference_vmoba
|
||||
[performance]=performance [api]=api_server
|
||||
|
||||
@@ -13,12 +13,33 @@ on:
|
||||
required: false
|
||||
default: false
|
||||
type: boolean
|
||||
# Auto-rebuild the CUDA images when their Dockerfile changes on main. The CUDA
|
||||
# matrix is the only lane that builds from docker/Dockerfile, so a path-scoped
|
||||
# push trigger is a sufficient change detector on its own -- no separate
|
||||
# detect-changes/paths-filter job is needed now that there is a single
|
||||
# in-scope Dockerfile. Dreamverse (apps/dreamverse/docker/Dockerfile) and the
|
||||
# rocm Dockerfile stay manual-dispatch only.
|
||||
push:
|
||||
branches: [main]
|
||||
paths:
|
||||
- 'docker/Dockerfile'
|
||||
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
packages: write
|
||||
|
||||
# One static group, no cancellation: every run of this workflow writes the same
|
||||
# mutable registry tags (latest, py3.12-latest, ...), so runs must serialize —
|
||||
# concurrent push/dispatch runs would race on those tags, and cancelling a run
|
||||
# mid-publish can strand the cu126/cu130 tag families at different commits. An
|
||||
# in-flight superseded build wastes its runner time, but its tags are then
|
||||
# overwritten by the newer queued run. GitHub keeps a single pending run per
|
||||
# group: the newest queued run replaces any older queued one.
|
||||
concurrency:
|
||||
group: infra-build-image
|
||||
cancel-in-progress: false
|
||||
|
||||
jobs:
|
||||
# CUDA matrix: Python 3.12 x {12.6.3, 13.0.0} x {amd64, arm64}. Each architecture
|
||||
# builds natively and pushes only by digest; publish-cuda-manifests is the sole
|
||||
@@ -28,7 +49,11 @@ jobs:
|
||||
# aliases; 13.0.0/cu130 is published under explicit versioned tags. Flash-attn
|
||||
# 2.8.3 comes from the architecture-specific prebuilt releases.
|
||||
build-cuda-images:
|
||||
if: ${{ github.event.inputs.build_cuda_matrix == 'true' }}
|
||||
# Runs on a manual dispatch when build_cuda_matrix is set, or automatically
|
||||
# on a push that changed docker/Dockerfile (inputs are null on push). The
|
||||
# repository guard keeps fork syncs from auto-building; manual dispatch
|
||||
# still works in forks.
|
||||
if: ${{ (github.event_name == 'push' && github.repository == 'hao-ai-lab/FastVideo') || github.event.inputs.build_cuda_matrix == 'true' }}
|
||||
strategy:
|
||||
fail-fast: false
|
||||
matrix:
|
||||
@@ -75,10 +100,10 @@ jobs:
|
||||
secrets: inherit
|
||||
|
||||
publish-cuda-manifests:
|
||||
# !cancelled(): a failed sibling build leg must not skip the manifests for a
|
||||
# CUDA lane whose own digests all exist; the digest-count check below fails
|
||||
# the incomplete lane loudly instead.
|
||||
if: ${{ !cancelled() && github.event.inputs.build_cuda_matrix == 'true' }}
|
||||
# !cancelled(): publish lanes whose digests exist even if a sibling build
|
||||
# leg failed (the digest-count check fails incomplete lanes); it also
|
||||
# bypasses skipped-needs propagation, hence the explicit skipped check.
|
||||
if: ${{ !cancelled() && needs.build-cuda-images.result != 'skipped' }}
|
||||
needs: build-cuda-images
|
||||
runs-on: ubuntu-latest
|
||||
permissions:
|
||||
|
||||
@@ -10,8 +10,6 @@ exclude: |
|
||||
scripts/.*|
|
||||
fastvideo/dataset/.*|
|
||||
fastvideo/models/.*|
|
||||
v2/(layers|attention|platforms|configs|distributed|models|logging_utils|third_party|hooks|api)/.*|
|
||||
v2/(envs|logger|utils|version|forward_context|fastvideo_args)\.py|
|
||||
^apps/dreamverse/web/.*|
|
||||
examples/.*|
|
||||
\.agents/.*|
|
||||
|
||||
@@ -84,9 +84,12 @@ RUN source /opt/venv/bin/activate \
|
||||
FFMPEG_NATIVE_CXX=/usr/bin/g++ \
|
||||
bash /opt/FastVideo/apps/dreamverse/scripts/install_native_ffmpeg.sh
|
||||
|
||||
# FASTVIDEO_FA4: FA4 (flash_attn.cute) is opt-in; this image installs it via
|
||||
# the dreamverse extra and is validated with it, so enable it here.
|
||||
ENV FASTVIDEO_DREAMVERSE_HOME=/var/lib/dreamverse \
|
||||
STREAM_MODE=av_fmp4 \
|
||||
FASTVIDEO_ENABLE_PROMPT_SAFETY=0 \
|
||||
FASTVIDEO_FA4=1 \
|
||||
HF_HOME=/root/.cache/huggingface
|
||||
|
||||
RUN mkdir -p /var/lib/dreamverse
|
||||
|
||||
@@ -12,13 +12,17 @@ Defaults:
|
||||
|
||||
- `HF_REPO_ID=FastVideo/performance-tracking`
|
||||
- `PERFORMANCE_TRACKING_ROOT=/tmp/fastvideo-perf-dashboard`
|
||||
- `PERF_MAX_REGRESSION=0.05`
|
||||
|
||||
Records can include source metadata:
|
||||
Records can include source metadata and rolling-baseline policy context:
|
||||
|
||||
- `run_source`: `pr`, `local`, `scheduled_main`, or `unknown`
|
||||
- `baseline_eligible`: only successful scheduled-main records should be true
|
||||
- Buildkite metadata such as branch, PR number, build URL, build ID, and job ID
|
||||
- `regression_thresholds`: per-metric rolling-baseline percent and absolute
|
||||
floors used for recomputed status context
|
||||
|
||||
Dashboard/API metric payloads expose `threshold_exceeded` for raw threshold
|
||||
crossings; `regressed` remains the gated CI-failure signal.
|
||||
|
||||
Set one of `HF_API_KEY`, `HUGGINGFACE_HUB_TOKEN`, or `HF_TOKEN` if the
|
||||
configured dataset repo requires authenticated access:
|
||||
@@ -89,7 +93,8 @@ Trend charts show metric-specific axes and exact point details on hover/focus:
|
||||
- PR number, branch, and Buildkite URL when present
|
||||
|
||||
The latest status table uses the stored JSON `success` value. Recomputed
|
||||
baseline context is shown separately and does not override stored status.
|
||||
baseline context applies each metric's percent and absolute regression floors
|
||||
and does not override stored status.
|
||||
|
||||
## API
|
||||
|
||||
|
||||
@@ -440,6 +440,8 @@ export default function App() {
|
||||
<th>Throughput</th>
|
||||
<th>Memory</th>
|
||||
<th>Worst</th>
|
||||
<th>Exceeded</th>
|
||||
<th>Failing</th>
|
||||
</tr>
|
||||
</thead>
|
||||
<tbody>
|
||||
@@ -465,6 +467,12 @@ export default function App() {
|
||||
<td>{formatNumber(row.metrics.throughput?.current, 3)}</td>
|
||||
<td>{formatNumber(row.metrics.memory?.current, 1)}</td>
|
||||
<td>{formatNumber(row.worst_regression_pct, 1)}%</td>
|
||||
<td>
|
||||
{row.threshold_exceeded_metrics.length
|
||||
? row.threshold_exceeded_metrics.join(", ")
|
||||
: "none"}
|
||||
</td>
|
||||
<td>{row.failing_metrics.length ? row.failing_metrics.join(", ") : "none"}</td>
|
||||
</tr>
|
||||
))}
|
||||
</tbody>
|
||||
|
||||
@@ -2,6 +2,12 @@ export type MetricValue = {
|
||||
current: number | null;
|
||||
baseline: number | null;
|
||||
regression_pct: number | null;
|
||||
absolute_delta: number | null;
|
||||
threshold_percent: number;
|
||||
threshold_absolute: number;
|
||||
gated: boolean;
|
||||
threshold_exceeded: boolean;
|
||||
regressed: boolean;
|
||||
label: string;
|
||||
lower_is_better: boolean;
|
||||
precision: number;
|
||||
@@ -15,7 +21,8 @@ export type SummaryRow = {
|
||||
success: boolean;
|
||||
baseline_n: number;
|
||||
worst_regression_pct: number | null;
|
||||
regression_threshold_pct: number;
|
||||
threshold_exceeded_metrics: string[];
|
||||
failing_metrics: string[];
|
||||
computed_regression_status: "pass" | "fail";
|
||||
status: "pass" | "fail";
|
||||
run_source: RunSource;
|
||||
|
||||
@@ -1,56 +0,0 @@
|
||||
# FastVideo — Design Philosophy
|
||||
|
||||
One page on *why* FastVideo is built the way it is. The full architecture, the as-built status, and the
|
||||
forward roadmap live in **[`v2/README.md`](v2/README.md)** — this is the philosophy beneath it.
|
||||
|
||||
---
|
||||
|
||||
**A deployable model is a post-training artifact.** Unlike an LLM — where inference optimizes frozen weights
|
||||
after the fact — a *usable* video/omni model is *created* by training: step distillation for latency, QAT for
|
||||
precision, distillation + self-forcing for causal/world models. So every inference capability is a
|
||||
**(recipe, runtime) pair**: the weights and the loop that produced-and-assumes them are one versioned object.
|
||||
This is the source of the moat — whoever owns *both* sides of the pair owns the optimization frontier — and it
|
||||
is why training and serving cannot be two systems.
|
||||
|
||||
**The work is loops, not `forward()`.** Denoise timesteps, AR decode, chunked rollout, VAE tiles, encoder
|
||||
chunks, audio tokens, reward batches, optimizer steps, media chunks — video and omni inference is iteration. A
|
||||
runtime that collapses everything to a single `forward` can't schedule, batch, cancel, stream, reserve memory
|
||||
for, or capture the behavior of what actually runs. So loops are first-class, and they are **driven**: the
|
||||
model describes the next step it needs, the runtime decides when and with whom it runs, the model folds the
|
||||
result back. The model keeps content-adaptive control flow; the runtime keeps admission, batching, streaming,
|
||||
and behavior capture. Per-request state lives in typed `LoopState`, never in module globals — so interleaving
|
||||
requests through one model instance cannot smear state, by construction.
|
||||
|
||||
**The model is the center; everything else is a view over it.** A typed `ModelCard` owns components, loops,
|
||||
the recipe, and the parity contract. Programs compose a card's loops into a task; Workflows compose cards into
|
||||
pipelines; the scheduler runs the *steps* of all loops as `WorkUnit`s under one currency (predicted GPU-time,
|
||||
because a bidirectional denoise step and an AR token are ~1000× apart and incommensurable in counts);
|
||||
deployment places and routes; products stream artifacts. None of them define model semantics — they reference
|
||||
the Model Plane. One resident instance can run many loop types on shared weights, which is what makes omni/MoT
|
||||
native rather than a DAG that doubles weights.
|
||||
|
||||
**Correctness is a typed contract, not a hope.** Caches are correct by *key* — if a field can change output
|
||||
semantics it is in the key, so reuse is partitioned, never blindly flushed. Parity between the train-forward
|
||||
and the serve-forward is *measured* on a declared ladder (component → loop → behavioral → distribution →
|
||||
artifact-quality), never assumed. And the non-negotiable gate is **interleave bit-parity**: N requests
|
||||
interleaved at step granularity must be bit-identical to running them serially — the test the whole
|
||||
loop-inversion bet lives or dies on.
|
||||
|
||||
**One substrate for inference, training, and RL.** The rollout forward *is* the serve forward plus capture —
|
||||
same loop, same caches, same batcher, same numerics — so every serving optimization is automatically a rollout
|
||||
optimization, and there is one numerics surface the ladder measures rather than a correction layer papering
|
||||
over it. The engine doubles as the RL rollout engine under a strict rule: `training` consumes the engine; the
|
||||
**engine never imports `training`**.
|
||||
|
||||
**Borrow aggressively; copy nothing as the core.** vLLM/SGLang scheduling, vLLM-Omni/SGLang-Omni omni serving,
|
||||
Dynamo fleet orchestration, diffusers components, xDiT parallelism, TorchTitan mesh discipline,
|
||||
verl-omni/miles RL lessons, ComfyUI workflows, Dreamverse/LiveKit sessions — each contributes a take, none is
|
||||
the center. Deployment orchestration (Dynamo) sits *above* the engine, never inside it. Extensions are
|
||||
versioned hook points, never monkeypatching. New frontier capabilities arrive as a card, a method, a loop, a
|
||||
workflow, or a controller — **not a rewrite**.
|
||||
|
||||
> 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. Deployment places and routes; products stream artifacts; neither defines
|
||||
> the model.
|
||||
+4
-2
@@ -170,12 +170,14 @@ RUN --mount=type=cache,target=/opt/uv/cache \
|
||||
# rmtree clears the wheel's stale cute files first to avoid an install conflict.
|
||||
# Then verify both survive so a broken overlay fails the build instead of shipping
|
||||
# an FA2-less image. x86 only: the FA4 stack (quack-kernels etc.) is unvalidated on
|
||||
# arm64 / GB10 (sm_121), so there we skip the overlay and FA4 falls back to FA2.
|
||||
# arm64 / GB10 (sm_121), so there we skip the overlay; FA4 is opt-in
|
||||
# (FASTVIDEO_FA4=1) and errors if set without the overlay, so leave it unset on
|
||||
# arm64 and the image runs FA3/FA2 as usual.
|
||||
RUN --mount=type=cache,target=/opt/uv/cache \
|
||||
source $HOME/.local/bin/env && \
|
||||
source /opt/venv/bin/activate && \
|
||||
if [ "${TARGETARCH}" = "arm64" ]; then \
|
||||
echo "Skipping FA4 cute overlay on arm64 (FA4 stack unvalidated there; FA4 falls back to FA2)"; \
|
||||
echo "Skipping FA4 cute overlay on arm64 (FA4 stack unvalidated there; do not set FASTVIDEO_FA4)"; \
|
||||
else \
|
||||
python -c "import glob, shutil; [shutil.rmtree(d, ignore_errors=True) for d in glob.glob('/opt/venv/lib/python*/site-packages/flash_attn/cute')]" && \
|
||||
uv pip install "flash-attn-4 @ git+https://github.com/Dao-AILab/flash-attention.git@${FA4_CUTE_REF}#subdirectory=flash_attn/cute" && \
|
||||
|
||||
@@ -103,6 +103,7 @@ can merge a PR.
|
||||
|---|---|---|
|
||||
| SSIM Tests | `ssim` | `fastvideo/**/*.py`, `pyproject.toml`, `docker/Dockerfile` |
|
||||
| LoRA Inference Tests | `inference_lora` | LoRA tests, loader, transformer tests, pipelines, LoRA layers |
|
||||
| LoRA Extraction Tests | `lora_extraction` | LoRA extraction scripts/tests, loader, training utilities, LoRA layers |
|
||||
| Training Tests | `training` | `fastvideo/**`, `pyproject.toml`, `docker/Dockerfile` |
|
||||
| Distillation DMD Tests | `distillation_dmd` | `fastvideo/training/*distillation_pipeline.py` |
|
||||
| Self-Forcing Tests | `self_forcing` | self-forcing distillation pipeline and tests |
|
||||
@@ -144,6 +145,7 @@ Valid direct test names:
|
||||
| `/test training` | `training` |
|
||||
| `/test lora-inference` | `inference_lora` |
|
||||
| `/test lora-training` | `training_lora` |
|
||||
| `/test lora-extraction` | `lora_extraction` |
|
||||
| `/test distillation` | `distillation_dmd` |
|
||||
| `/test self-forcing` | `self_forcing` |
|
||||
| `/test vsa` | `training_vsa` |
|
||||
|
||||
@@ -72,7 +72,10 @@ fastvideo/tests/performance/
|
||||
│ writes Markdown summary + (optionally) uploads new records
|
||||
├── dashboard.py
|
||||
│ └── builds time-series Plotly HTML from HF history
|
||||
└── hf_store.py # shared HF I/O + DataFrame helpers
|
||||
|
||||
fastvideo/performance/
|
||||
├── hf_store.py # shared HF I/O + DataFrame helpers
|
||||
└── metric_policy.py # shared rolling-baseline threshold policy
|
||||
```
|
||||
|
||||
The HF dataset (`FastVideo/performance-tracking` by default) holds one
|
||||
@@ -92,16 +95,18 @@ and recipe changes instead of treating all records for a model as equivalent.
|
||||
|
||||
## Metrics
|
||||
|
||||
Each benchmark records six metrics:
|
||||
Each benchmark records six metrics. The rolling-baseline comparator also has a
|
||||
per-metric policy with direction, percent threshold, absolute threshold, and a
|
||||
`gated` flag.
|
||||
|
||||
| Metric | Raw key | Normalized key | Direction |
|
||||
|---|---|---|---|
|
||||
| End-to-end generation latency | `avg_generation_time_s` | `latency` | Lower is better |
|
||||
| Video throughput | `throughput_fps` | `throughput` | Higher is better |
|
||||
| Peak GPU memory | `max_peak_memory_mb` | `memory` | Lower is better |
|
||||
| Text encoder time | `text_encoder_time_s` | `text_encoder_time_s` | Lower is better |
|
||||
| DiT denoising time | `dit_time_s` | `dit_time_s` | Lower is better |
|
||||
| VAE decode time | `vae_decode_time_s` | `vae_decode_time_s` | Lower is better |
|
||||
| Metric | Raw key | Normalized key | Direction | Default rolling policy |
|
||||
|---|---|---|---|---|
|
||||
| End-to-end generation latency | `avg_generation_time_s` | `latency` | Lower is better | 8% and 0.5 s |
|
||||
| Video throughput | `throughput_fps` | `throughput` | Higher is better | 8% and 0.05 FPS |
|
||||
| Peak GPU memory | `max_peak_memory_mb` | `memory` | Lower is better | 5% and 256 MB |
|
||||
| Text encoder time | `text_encoder_time_s` | `text_encoder_time_s` | Lower is better | 5% and 0.25 s |
|
||||
| DiT denoising time | `dit_time_s` | `dit_time_s` | Lower is better | 5% and 0.25 s |
|
||||
| VAE decode time | `vae_decode_time_s` | `vae_decode_time_s` | Lower is better | 5% and 0.25 s |
|
||||
|
||||
`test_inference_performance.py` temporarily sets `FASTVIDEO_STAGE_LOGGING=1`
|
||||
while it runs so pipeline stage execution times are available in
|
||||
@@ -156,9 +161,22 @@ headroom and almost never need touching.
|
||||
|
||||
`compare_baseline.py` loads the last 5 successful, baseline-eligible records
|
||||
for the same `(model_id, gpu_type)` from the HF dataset, computes the median
|
||||
for each available metric, and fails if the current run regresses by more than
|
||||
`PERF_MAX_REGRESSION` (default 5%). For latency, memory, and component times,
|
||||
higher values are regressions. For throughput, lower values are regressions.
|
||||
for each available metric, and evaluates the current run with the metric's
|
||||
rolling regression policy. For latency, memory, and component times, higher
|
||||
values are regressions. For throughput, lower values are regressions.
|
||||
|
||||
A metric exceeds its rolling threshold when both of these are true:
|
||||
|
||||
```text
|
||||
percent_delta > threshold_percent
|
||||
absolute_delta > threshold_absolute
|
||||
```
|
||||
|
||||
Gated metrics fail CI when that threshold crossing happens. Set `gated: false`
|
||||
for metrics that should remain visible in reports and the dashboard without
|
||||
failing CI. Dashboard/API payloads expose `threshold_exceeded` separately from
|
||||
`regressed`, where `regressed` means a gated CI failure. Missing or `null`
|
||||
metrics are skipped.
|
||||
|
||||
This is the **drift detector** — it catches sub-threshold regressions that
|
||||
slowly add up. Only scheduled-main successful records are baseline eligible.
|
||||
@@ -172,6 +190,45 @@ agent skill to advance the rolling median.
|
||||
|
||||
## Schemas
|
||||
|
||||
### Benchmark config (`.buildkite/performance-benchmarks/tests/*.json`)
|
||||
|
||||
Benchmark configs without `config_schema_version` are treated as legacy v1
|
||||
configs and remain loadable. New or migrated configs should use
|
||||
`config_schema_version: 2` and include explicit comparable identity fields:
|
||||
|
||||
```jsonc
|
||||
{
|
||||
"benchmark_id": "wan-t2v-1.3b-2gpu",
|
||||
"config_schema_version": 2,
|
||||
"workload_id": "wan-t2v-1.3b",
|
||||
"variant_id": "canonical",
|
||||
"benchmark_version": 1
|
||||
}
|
||||
```
|
||||
|
||||
`benchmark_id` is still required in this phase because raw artifact names,
|
||||
generated-video directories, normalized record paths, and the current rolling
|
||||
baseline comparator still depend on it. The v2 identity fields are config
|
||||
metadata that make the measured workload explicit:
|
||||
|
||||
| Field | Purpose |
|
||||
|---|---|
|
||||
| `workload_id` | Stable benchmark family, such as `wan-t2v-1.3b`. |
|
||||
| `variant_id` | Intentional recipe family, such as `canonical`. |
|
||||
| `benchmark_version` | Version of the measurement protocol and comparison policy. |
|
||||
|
||||
If a config declares `config_schema_version: 2`, loading fails clearly when any
|
||||
required v2 identity field is missing. If v2 identity or metadata fields are
|
||||
added without `config_schema_version: 2`, loading also fails so partial
|
||||
migrations do not silently run as v1 configs. Optional v2 metadata fields
|
||||
reserved for follow-up work, such as `recipe`, `metric_threshold_policy`, and
|
||||
`quality_metadata`, must be JSON objects when present.
|
||||
|
||||
Recipe fingerprinting, hardware/software profile IDs, exact-identity
|
||||
comparison, metric-specific threshold policy behavior, promoted baselines, and
|
||||
dashboard regrouping are separate follow-up changes. Until those land, rolling
|
||||
baseline comparison remains keyed by `(model_id, gpu_type)`.
|
||||
|
||||
### Raw record (`results/perf_*.json`)
|
||||
|
||||
Written by `test_inference_performance.py`. One file per benchmark run.
|
||||
@@ -179,6 +236,10 @@ Written by `test_inference_performance.py`. One file per benchmark run.
|
||||
```jsonc
|
||||
{
|
||||
"benchmark_id": "wan-t2v-1.3b-2gpu",
|
||||
"config_schema_version": 2,
|
||||
"workload_id": "wan-t2v-1.3b",
|
||||
"variant_id": "canonical",
|
||||
"benchmark_version": 1,
|
||||
"model_short_name": "Wan2.1-T2V-1.3B-Diffusers",
|
||||
"device": "NVIDIA L40S",
|
||||
"num_gpus": 2,
|
||||
@@ -196,6 +257,13 @@ Written by `test_inference_performance.py`. One file per benchmark run.
|
||||
"max_dit_time_s": 10.0,
|
||||
"max_vae_decode_time_s": 10.0
|
||||
},
|
||||
"regression_thresholds": {
|
||||
"latency": {
|
||||
"threshold_percent": 0.10,
|
||||
"threshold_absolute": 1.0,
|
||||
"gated": true
|
||||
}
|
||||
},
|
||||
"commit": "<full sha>",
|
||||
"pr_number": "1234",
|
||||
"timestamp": "2026-05-08T22:00:00+00:00",
|
||||
@@ -222,6 +290,13 @@ result, used as the rolling-baseline source of truth.
|
||||
"text_encoder_time_s": 2.141,
|
||||
"dit_time_s": 8.437,
|
||||
"vae_decode_time_s": 3.208,
|
||||
"regression_thresholds": {
|
||||
"latency": {
|
||||
"threshold_percent": 0.08,
|
||||
"threshold_absolute": 0.5,
|
||||
"gated": true
|
||||
}
|
||||
},
|
||||
"success": true
|
||||
}
|
||||
```
|
||||
@@ -238,18 +313,17 @@ successful main/full-suite uploads and remain eligible for rolling baselines.
|
||||
|
||||
| Variable | Default | Used by | Purpose |
|
||||
|---|---|---|---|
|
||||
| `PERF_MAX_REGRESSION` | `0.05` | `compare_baseline.py` | Per-metric regression fraction that fails the build. |
|
||||
| `PERFORMANCE_TRACKING_ROOT` | `/tmp/perf-tracking` | `compare_baseline.py`, `dashboard.py` | Local directory the HF dataset is synced to. |
|
||||
| `PERF_REPORTS_DIR` | `/root/data/perf_reports` | `compare_baseline.py`, `dashboard.py` | Where the Markdown summary and Plotly HTML get written for Buildkite to pick up. |
|
||||
| `HF_REPO_ID` | `FastVideo/performance-tracking` | `hf_store.py` | HF dataset repo holding rolling-baseline records. |
|
||||
| `HF_API_KEY`, `HUGGINGFACE_HUB_TOKEN`, `HF_TOKEN` | unset | `hf_store.py` | Required for upload or private dataset reads. |
|
||||
| `HF_REPO_ID` | `FastVideo/performance-tracking` | `fastvideo/performance/hf_store.py` | HF dataset repo holding rolling-baseline records. |
|
||||
| `HF_API_KEY`, `HUGGINGFACE_HUB_TOKEN`, `HF_TOKEN` | unset | `fastvideo/performance/hf_store.py` | Required for upload or private dataset reads. |
|
||||
| `PERF_RUN_SOURCE` | inferred | `compare_baseline.py` | Source metadata for uploaded records: `pr`, `local`, `scheduled_main`, or `unknown`. |
|
||||
| `PERF_UPLOAD_POLICY` | `never` | `compare_baseline.py` | Upload policy: `never`, `pass`, or `always`. |
|
||||
| `PERF_PYTEST_RC` | unset | `compare_baseline.py` | Static-threshold pytest exit code, used so scheduled-main failures can be uploaded with `success=false`. |
|
||||
| `TEST_SCOPE` | unset | `compare_baseline.py` | CI context used to infer scheduled-main runs together with `BUILDKITE_BRANCH=main`. |
|
||||
| `BUILDKITE_BRANCH`, `BUILDKITE_COMMIT`, `BUILDKITE_PULL_REQUEST` | unset | `compare_baseline.py`, `test_inference_performance.py` | CI metadata stamped into records. |
|
||||
| `DASHBOARD_DAYS` | `30` | `dashboard.py` | Lookback window for the Plotly trend pages. |
|
||||
| `PERFORMANCE_TRACKING_SYNC_REUSE_TTL_SECONDS` | `3600` | `hf_store.py` | Freshness window for reusing an existing HF sync when requested by dashboard consumers. |
|
||||
| `PERFORMANCE_TRACKING_SYNC_REUSE_TTL_SECONDS` | `3600` | `fastvideo/performance/hf_store.py` | Freshness window for reusing an existing HF sync when requested by dashboard consumers. |
|
||||
| `FASTVIDEO_STAGE_LOGGING` | set by the pytest test | `test_inference_performance.py` | Enables pipeline stage timing capture for component metrics during benchmark runs. |
|
||||
|
||||
## CI integration
|
||||
@@ -261,9 +335,11 @@ Buildkite artifact upload is in
|
||||
`.buildkite/scripts/pr_test.sh:upload_performance_artifacts`.
|
||||
|
||||
Each performance build runs pytest first. If that fixed-threshold phase fails,
|
||||
`compare_baseline.py` is skipped, so Markdown summaries and normalized JSON
|
||||
artifacts are not emitted. The dashboard still runs best-effort for
|
||||
observability. When pytest passes, the rolling-baseline phase emits:
|
||||
PR/direct runs skip `compare_baseline.py` because they only upload passing
|
||||
records. Scheduled-main runs still execute `compare_baseline.py` with
|
||||
`PERF_PYTEST_RC` set so the failed canonical attempt is visible in normalized
|
||||
JSON and dashboard history. The dashboard runs best-effort for observability.
|
||||
When the rolling-baseline phase runs, it emits:
|
||||
|
||||
* **Markdown summary** — appended to `$GITHUB_STEP_SUMMARY` when that variable
|
||||
is set, and written as `perf_<sha>_<ts>.md` for Buildkite upload. Contains a
|
||||
@@ -279,11 +355,16 @@ observability. When pytest passes, the rolling-baseline phase emits:
|
||||
## Adding a new benchmark
|
||||
|
||||
1. Drop a new JSON config into
|
||||
`.buildkite/performance-benchmarks/tests/<name>.json`. Required keys:
|
||||
`.buildkite/performance-benchmarks/tests/<name>.json`. New configs should
|
||||
use v2 identity fields:
|
||||
|
||||
```json
|
||||
{
|
||||
"benchmark_id": "<unique-id>",
|
||||
"config_schema_version": 2,
|
||||
"workload_id": "<stable-workload-id>",
|
||||
"variant_id": "canonical",
|
||||
"benchmark_version": 1,
|
||||
"model": { "model_path": "...", "model_short_name": "..." },
|
||||
"init_kwargs": { "num_gpus": 1, ... },
|
||||
"generation_kwargs": { "num_frames": 45, ... },
|
||||
@@ -299,9 +380,17 @@ observability. When pytest passes, the rolling-baseline phase emits:
|
||||
"max_vae_decode_time_s": 10.0
|
||||
},
|
||||
"default": { "max_generation_time_s": 120.0, "max_peak_memory_mb": 30000.0 }
|
||||
},
|
||||
"regression_thresholds": {
|
||||
"latency": { "threshold_percent": 0.10, "threshold_absolute": 1.0, "gated": true }
|
||||
}
|
||||
}
|
||||
```
|
||||
}
|
||||
|
||||
```
|
||||
|
||||
Legacy v1 configs without `config_schema_version` still load, but should not
|
||||
gain v2 identity or metadata fields until they are migrated to
|
||||
`config_schema_version: 2`.
|
||||
|
||||
2. The pytest test auto-discovers all configs — no test code needed. CI
|
||||
picks it up on the next `/test performance` run.
|
||||
@@ -320,6 +409,13 @@ observability. When pytest passes, the rolling-baseline phase emits:
|
||||
a useful fixed gate. The rolling baseline will still track component times
|
||||
when static component thresholds are omitted.
|
||||
|
||||
6. Omit `regression_thresholds` to use the default rolling-baseline policy, or
|
||||
include only benchmark-specific deviations. Tune these independently from
|
||||
the fixed thresholds when a metric is noisy or should be informational. The
|
||||
fixed `thresholds` block is an absolute pytest ceiling. The
|
||||
`regression_thresholds` block controls rolling-baseline comparisons against
|
||||
recent scheduled-main records.
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
**"No baseline for ... Initializing"** — first run for this `(model_id,
|
||||
|
||||
@@ -41,6 +41,11 @@ surfaces:
|
||||
disable_autocast: generator.engine.disable_autocast
|
||||
enable_stage_verification: generator.engine.enable_stage_verification
|
||||
prompt_txt: request.inputs.prompt_path
|
||||
batching_mode: generator.engine.batching.mode
|
||||
batching_max_size: generator.engine.batching.max_size
|
||||
batching_delay_ms: generator.engine.batching.delay_ms
|
||||
batching_config: generator.engine.batching.config_path
|
||||
enable_batching_metrics: generator.engine.batching.enable_metrics
|
||||
override_text_encoder_safetensors: generator.pipeline.components.text_encoder_weights
|
||||
override_text_encoder_quant: generator.engine.quantization.text_encoder_quant
|
||||
transformer_quant: generator.engine.quantization.transformer_quant
|
||||
|
||||
@@ -74,6 +74,23 @@ uv pip install ninja
|
||||
python setup.py install
|
||||
```
|
||||
|
||||
### Flash Attention 4 (opt-in)
|
||||
|
||||
FastVideo never auto-selects FlashAttention-4 (`flash_attn.cute`) just because it
|
||||
is installed: its CuTeDSL kernels JIT-compile per shape family and can fail at
|
||||
runtime on some GPU/shape combinations. To use FA4, install the pinned
|
||||
`flash-attn-4` build (see the `flash-attn-4` source in `pyproject.toml`) and set:
|
||||
|
||||
```bash
|
||||
export FASTVIDEO_FA4=1
|
||||
```
|
||||
|
||||
On GPUs below sm90 a capability gate routes to FlashAttention-2 the calls FA4
|
||||
cannot serve there: grad-enabled (training) attention (FA4's backward requires
|
||||
sm90+) and GQA attention (FA4's `pack_gqa` fails to JIT-compile below sm90).
|
||||
On sm90+ both run on FA4. If FA4 is unusable while `FASTVIDEO_FA4=1` is set,
|
||||
FastVideo fails loudly instead of silently falling back.
|
||||
|
||||
### FP4 Flash Attention 4 (Blackwell only)
|
||||
|
||||
**`FLASH_ATTN`** with **`--nvfp4_fa4`**
|
||||
|
||||
@@ -1,94 +0,0 @@
|
||||
# v2 porting status — fastvideo models → the v2 (recipe, runtime) substrate
|
||||
|
||||
Goal: every model in fastvideo's registry resolves through the **v2 `VideoGenerator`** / `Engine`
|
||||
(typed `fastvideo.api` configs + the real torch backend) to a recipe that can construct and run it.
|
||||
|
||||
**Scope: ALL fastvideo models (achieved).** v2 now resolves **63/64** of fastvideo's registered HF ids
|
||||
by exact id (PRIMARY), plus the architecture fallback for local/unregistered checkpoints. The single
|
||||
remaining id — `FastVideo/Wan2.1-VSA-T2V-14B-720P-Diffusers` — is **environment-blocked**: its VSA
|
||||
(Sparse-Linear Attention) kernels require `nvcc` (not built in this bring-up). It arch-resolves to the
|
||||
base Wan card but needs the VSA kernel build to run faithfully.
|
||||
|
||||
Dispatch is **architecture-driven** (`v2/registry.py`): exact HF id → short-name → architecture
|
||||
inference from the checkpoint (pipeline / transformer / VAE class names + `z_dim`, `transformer_2`,
|
||||
`spatial_upsampler`). Adding a model is one `_BUCKET_C` row (HF ids → builders + transformer class).
|
||||
|
||||
## The porting mechanism — self-contained recipe packages
|
||||
Every net-new arch is a **self-contained recipe package** (`v2/recipes/<arch>/` = `card.py` `loop.py`
|
||||
`program.py` [+ `sampler.py`] + an optional `v2/platform/backends/torch_<arch>.py` adapter). The card
|
||||
declares its torch adapter via **`ComponentSpec.adapter="module:Class"`** (the `_explicit_adapter` seam in
|
||||
`torch_backend.py`) instead of editing the shared `_make_dit`/`_make_vae`/`_make_text_encoder` dispatch —
|
||||
so a port adds **only new files**, never touching shared code, and parallel ports never conflict. New
|
||||
samplers/loops live in-package. Registration is one row in `v2/registry.py:_BUCKET_C`.
|
||||
|
||||
## Working today (GPU-verified, real video/audio) — committed on `v2`
|
||||
| Official example(s) | Model | v2 card |
|
||||
|---|---|---|
|
||||
| `basic.py`, `basic_mps.py`, `basic_ray.py` | `Wan-AI/Wan2.1-T2V-1.3B-Diffusers` | wan21 |
|
||||
| `basic_self_forcing_causal.py` | `wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers` | wan_causal |
|
||||
| `basic_ltx2_distilled.py` | `FastVideo/LTX2-Distilled-Diffusers` (2-stage + spatial upsampler) | ltx2 |
|
||||
| `basic_wan2_2_ti2v.py` | `Wan-AI/Wan2.2-TI2V-5B-Diffusers` | wan2.2-ti2v |
|
||||
| `basic_wan2_2.py` | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` (MoE + CPU expert offload) | wan2.2-a14b |
|
||||
| `basic_ltx2.py` | `Davids048/LTX2-Base-Diffusers` | ltx2 base |
|
||||
| `basic_ltx2_3_distilled.py` | `FastVideo/LTX-2.3-Distilled-Diffusers` (joint T2VS, video+audio) | ltx2.3-distilled |
|
||||
|
||||
Plus the **Wan2.1 i2v cluster** (Fun-1.3B-InP GPU-verified; I2V-14B-480P/720P + Wan2.2-I2V-A14B MoE reuse
|
||||
the i2v card) — CLIP image-encoder + first-frame `[mask|cond]` → 36ch DiT.
|
||||
|
||||
## GPU bring-up results (real weights on H100 NVL, single-GPU, TORCH_SDPA)
|
||||
**20 models generate real video/audio on GPU** — the 7 above + **13 of the newly-ported** archs, each run
|
||||
end-to-end through the real `VideoGenerator` (resolve → stamp → CUDA load → generate). The rest are blocked
|
||||
by a **fastvideo-shared-code / missing-kernel / HF-access** wall, NOT a v2 recipe bug (the v2 recipes are
|
||||
faithful — e.g. cosmos25's DiT+VAE produced finite output; only its Qwen2.5-VL encoder hit a library
|
||||
incompat). All ports also resolve + run end-to-end on the CPU toy backend (`test_bucket_c_ports.py`).
|
||||
|
||||
| GPU status | Models |
|
||||
|---|---|
|
||||
| ✅ **Verified** (real GPU output) | stable_audio (audio), matrixgame2, matrixgame3, gen3c, wan_fun_control, lucy_edit, hunyuangamecraft, hunyuan_video, hunyuan_video15, longcat (13.58B), sfwan22 (2×14B MoE, expert offload), lingbotworld (2×14B, offload), fastwan (TI2V-5B-FullAttn DMD) |
|
||||
| 🚫 fastvideo/env-blocked | **cosmos25** (DiT+VAE ran; Qwen2.5-VL encoder → transformers 5.12.1 incompat in fastvideo); **kandinsky5** (fastvideo registry registers a bare `PipelineConfig`); **hyworld** (fastvideo DiT hardcodes `flash_attn`, not built); **turbowan** 1.3B/i2v + **fastwan** VSA-variants (SLA/VSA sparse-attn params + Triton kernels need nvcc) |
|
||||
| 🚫 access-blocked (HF-gated) | cosmos2, flux2, sd35 (no HF token in this env) |
|
||||
|
||||
To unblock the env-blocked: build `fastvideo-kernel` (SLA/VSA Triton, needs nvcc); pin a fastvideo-compatible
|
||||
`transformers` for the Qwen2.5-VL encoder; add a Kandinsky5 `PipelineConfig` + an SDPA fallback in the
|
||||
hyworld DiT (all fastvideo-side / environment, not v2 recipe work).
|
||||
|
||||
## Newly ported (recipe details)
|
||||
Each resolves through the registry AND runs end-to-end on the CPU toy backend via the public `Engine`
|
||||
path (the `v2/tests/test_bucket_c_ports.py` regression guard), emitting the correct modality artifact.
|
||||
|
||||
**15 net-new architectures** (each a new `TorchComponent` adapter + recipe):
|
||||
- **cosmos2** (Cosmos-Predict2-2B-Video2World) — EDM-Karras denoiser; new `CosmosDenoiseLoop` +
|
||||
`build_karras_sigmas` (the reference port). **cosmos25** (Cosmos-Predict2.5 2B/14B) — flow-match,
|
||||
per-frame plain-sigma timestep, Reason1/Qwen2.5-VL encoder. **gen3c** (GEN3C) — EDM + 82ch pose-buffer.
|
||||
- **hunyuan_video** (+FastHunyuan) — reuses WanDenoiseLoop, dual LLaMA+CLIP encoders, Hunyuan VAE.
|
||||
**hunyuan_video15** (480p/720p). **hunyuangamecraft**, **hyworld** — interactive (camera/action).
|
||||
- **longcat** (T2V/I2V/VC). **kandinsky5** (5.0 T2V Lite).
|
||||
- **sd35** (MMDiT, image, triple-encoder). **flux2** (dev/klein, MMDiT image). **stable_audio** (audio).
|
||||
- **lingbotworld** (camera/Plucker), **matrixgame2**, **matrixgame3** — interactive world models.
|
||||
|
||||
**5 Wan-family variants** (reuse the Wan/Causal arch, new in-package sampler/loop/conditioning):
|
||||
- **turbowan** — rCM few-step (faithful RCMScheduler port), 1.3B/14B T2V + I2V-A14B MoE.
|
||||
- **lucy_edit** — v2v editor (video-VAE-encode node → 96ch DiT input). **wan_fun_control** — control input.
|
||||
- **sfwan22** — Self-Forcing Wan2.2-A14B causal + MoE (i2v + t2v). **fastwan** — DMD 3-step (TI2V-5B-FullAttn
|
||||
loadable; VSA-trained variants + non-strict `to_gate_compress` load are BRINGUP).
|
||||
|
||||
BRINGUP scope per port (documented in each package): GPU load/run; for interactive/world-model archs the
|
||||
action/camera/memory conditioning needs a request-API extension (the t2v/degenerate path is what
|
||||
CPU-verifies); video2world/i2v frame-replace conditioning is threaded but inert without conditioning inputs.
|
||||
|
||||
## Environment
|
||||
v2 bring-up runs **single-GPU, resident, on the `TORCH_SDPA` backend** (no fastvideo-kernel / VSA / FP4).
|
||||
The box has been rescheduled across hosts/arches/python versions mid-session; rebuild the venv for the
|
||||
current arch when that happens: `uv venv --python 3.12 .venv`; comment out `fastvideo-kernel` in
|
||||
`pyproject.toml`; `uv pip install -e ".[dev]"`. Source `/home/scratch.willlin_ent/.bringup_env`
|
||||
(`HF_HOME=./.cache` on scratch, `FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA`). v2 CPU mini: 240 passed, 2 skipped.
|
||||
|
||||
## How to add a model to the v2 substrate
|
||||
1. `v2/recipes/<arch>/` — card (declare adapters via `ComponentSpec.adapter`; per-model `SamplingDefaults`),
|
||||
loop (reuse `WanDenoiseLoop`/`chunk_rollout` or a new in-package loop+sampler), program.
|
||||
2. `v2/platform/backends/torch_<arch>.py` — a `TorchComponent` subclass (only the forward semantics) if the
|
||||
arch is genuinely new; reuse `WanDiT`/`LTX2DiT`/`WanVAE`/`T5Encoder` via `load_id` when it isn't.
|
||||
3. One row in `v2/registry.py:_BUCKET_C` (HF ids → builders; `transformer_cls` for the arch fallback, or
|
||||
`""` for explicit-id-only capability variants of an existing arch).
|
||||
4. CPU-verify: it resolves + runs on the toy backend (auto-covered by `test_bucket_c_ports.py`). Then GPU
|
||||
bring-up (`stamp_*_checkpoints` → real weights) per BRINGUP notes.
|
||||
@@ -1,36 +0,0 @@
|
||||
"""v2 port of basic.py — Wan2.1-T2V-1.3B through the v2 VideoGenerator.
|
||||
|
||||
Same convenience API as upstream (from_pretrained + generate_video); only delta is importing
|
||||
VideoGenerator from v2. v2 bring-up: single-GPU, resident, SDPA; modest res/frames for a quick run.
|
||||
"""
|
||||
from v2 import VideoGenerator
|
||||
|
||||
OUTPUT_PATH = "v2_video_samples"
|
||||
|
||||
|
||||
def main() -> None:
|
||||
generator = VideoGenerator.from_pretrained(
|
||||
"Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
|
||||
num_gpus=1,
|
||||
use_fsdp_inference=False,
|
||||
dit_cpu_offload=False,
|
||||
vae_cpu_offload=False,
|
||||
text_encoder_cpu_offload=False,
|
||||
pin_cpu_memory=False,
|
||||
)
|
||||
common = dict(output_path=OUTPUT_PATH, save_video=True,
|
||||
num_frames=25, height=480, width=832, num_inference_steps=30, guidance_scale=5.0)
|
||||
|
||||
prompt = ("A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes wide with "
|
||||
"interest. The playful yet serene atmosphere is complemented by soft natural light "
|
||||
"filtering through the petals. Mid-shot, warm and cheerful tones.")
|
||||
video = generator.generate_video(prompt, output_video_name="wan21_raccoon", **common)
|
||||
|
||||
prompt2 = ("A majestic lion strides across the golden savanna, its powerful frame glistening under "
|
||||
"the warm afternoon sun. Low angle, steady tracking shot, cinematic.")
|
||||
video2 = generator.generate_video(prompt2, output_video_name="wan21_lion", **common)
|
||||
print(f"Outputs: {video.video_path} , {video2.video_path}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,29 +0,0 @@
|
||||
"""v2 port of basic_ltx2.py — LTX-2 base (single-stage) through the v2 VideoGenerator.
|
||||
|
||||
Same convenience API as upstream; only delta is importing VideoGenerator from v2. LTX-2 base is the
|
||||
single-stage (non-distilled) model: the v2 single-stage card (build_ltx2_base_card) runs a request-driven
|
||||
many-step flow-match at FULL latent res (no distilled base/refine split, no spatial upsampler), reusing
|
||||
the LTX-2 DiT/VAE/Gemma adapters. The SAME single-stage card also serves LTX-2.3-Distilled (which is also
|
||||
single-stage) — just pass fewer num_inference_steps for the few-step distilled schedule.
|
||||
|
||||
NOTE: modest res/frames here — upstream defaults to 1088x1920x121, which on an 18.88B base is very slow;
|
||||
raise them for full quality. v2 bring-up: single-GPU, resident, SDPA.
|
||||
"""
|
||||
from v2 import VideoGenerator
|
||||
|
||||
PROMPT = ("A warm sunny backyard, cinematic close-up of two people talking; the camera slowly pans right "
|
||||
"to reveal a grandfather in the garden wearing enormous butterfly wings, flapping his arms like "
|
||||
"he is trying to take off. Deadpan, absurd, quietly tragic.")
|
||||
|
||||
|
||||
def main() -> None:
|
||||
generator = VideoGenerator.from_pretrained("Davids048/LTX2-Base-Diffusers", num_gpus=1)
|
||||
video = generator.generate_video(
|
||||
prompt=PROMPT, output_path="v2_video_samples_ltx2_base", output_video_name="ltx2_base_backyard",
|
||||
save_video=True, num_frames=25, height=512, width=768, num_inference_steps=30)
|
||||
print(f"Output: {video.video_path}")
|
||||
generator.shutdown()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,36 +0,0 @@
|
||||
"""v2 port of basic_ltx2_3_distilled.py — LTX-2.3 Distilled (single-stage, joint A/V) through the v2
|
||||
VideoGenerator.
|
||||
|
||||
Unlike LTX-2.0 distilled (two-stage, video-only), LTX-2.3 is a single-stage *audio+video* model. The
|
||||
shared registry (v2/registry.py) maps ``FastVideo/LTX-2.3-Distilled-Diffusers`` to its OWN card,
|
||||
``build_ltx2_3_card`` — distinct from the LTX-2 base/2-stage cards — which wires the 2.3-specific path:
|
||||
* SEPARATE video + audio text connectors (the Gemma encoder projects the prompt to two embeddings,
|
||||
2048-dim for audio, 4096-dim for video) plus gated attention;
|
||||
* a JOINT DiT forward where video and audio latents cross-attend in a single denoise per step;
|
||||
* a video VAE decode + an AudioDecoder→Vocoder decode → video frames AND a stereo waveform @24kHz.
|
||||
|
||||
Because the model advertises TEXT_TO_VIDEO_SOUND, the VideoGenerator issues a T2VS request by default,
|
||||
so ``generate_video`` returns BOTH modalities: the mp4 plus a sibling ``.wav`` (and ``result.audio`` /
|
||||
``result.audio_sample_rate`` in memory). Being distilled, it wants FEW steps (8). GPU-verified on the
|
||||
rebuilt x86 stack: video (3,33,256,384) + stereo audio (2×61920 @ 24kHz).
|
||||
"""
|
||||
from v2 import VideoGenerator
|
||||
|
||||
PROMPT = "ocean waves crashing on rocks at sunset, seagulls calling in the distance, cinematic, highly detailed"
|
||||
|
||||
|
||||
def main() -> None:
|
||||
generator = VideoGenerator.from_pretrained("FastVideo/LTX-2.3-Distilled-Diffusers", num_gpus=1)
|
||||
# audio=None auto-enables sound for this A/V model (pass audio=False to force video-only).
|
||||
result = generator.generate_video(
|
||||
prompt=PROMPT, output_path="v2_video_samples_ltx2_3", output_video_name="ltx2_3_ocean",
|
||||
save_video=True, num_frames=33, height=512, width=768, num_inference_steps=8, seed=1)
|
||||
print(f"Video: {result.video_path}")
|
||||
audio_path = result.extra.get("audio_path")
|
||||
if audio_path:
|
||||
print(f"Audio: {audio_path} ({result.audio_sample_rate} Hz)")
|
||||
generator.shutdown()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,87 +0,0 @@
|
||||
"""v2 typed-API inference example — mirrors ``basic_dmd_new_api.py`` but drives the **v2
|
||||
(recipe, runtime) substrate + real torch backend** for the three models brought up on GPU
|
||||
(Wan2.1, SF-causal Wan, LTX-2).
|
||||
|
||||
The ONLY delta from the upstream example is importing ``VideoGenerator`` from ``v2`` instead of
|
||||
``fastvideo`` — the typed config classes are the SAME ``fastvideo.api`` dataclasses.
|
||||
|
||||
Run (on a GPU box, with the v2 venv active):
|
||||
python examples/inference/basic/v2_basic_new_api.py
|
||||
|
||||
Notes vs upstream: the v2 bring-up runs single-GPU, resident, on the TORCH_SDPA backend (no
|
||||
fastvideo-kernel / VSA), so resolutions/steps are modest here for a quick runnable demo. LTX-2 loads
|
||||
an 18.88B DiT (slow first load).
|
||||
"""
|
||||
import os
|
||||
import time
|
||||
|
||||
from v2 import VideoGenerator
|
||||
from fastvideo.api import (
|
||||
EngineConfig,
|
||||
GenerationRequest,
|
||||
GeneratorConfig,
|
||||
OffloadConfig,
|
||||
OutputConfig,
|
||||
SamplingConfig,
|
||||
)
|
||||
|
||||
OUTPUT_PATH = "v2_video_samples"
|
||||
|
||||
MODELS = [
|
||||
{
|
||||
"family": "wan21",
|
||||
"model_path": "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
|
||||
"prompt": "a red panda surfing on ocean waves at sunset, cinematic, highly detailed",
|
||||
"sampling": SamplingConfig(num_frames=25, height=480, width=832,
|
||||
num_inference_steps=30, guidance_scale=5.0, seed=1, fps=16),
|
||||
},
|
||||
{
|
||||
"family": "wan_causal",
|
||||
"model_path": "wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers",
|
||||
"prompt": "a cat walking through a sunlit garden, cinematic",
|
||||
"sampling": SamplingConfig(num_frames=25, height=480, width=832,
|
||||
num_inference_steps=4, guidance_scale=5.0, seed=1, fps=16),
|
||||
},
|
||||
{
|
||||
"family": "ltx2",
|
||||
"model_path": "FastVideo/LTX2-Distilled-Diffusers",
|
||||
"prompt": "surfers riding ocean waves at sunset, cinematic, highly detailed",
|
||||
"sampling": SamplingConfig(num_frames=9, height=512, width=768,
|
||||
num_inference_steps=8, guidance_scale=1.0, seed=1, fps=16),
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
def run_one(m: dict) -> None:
|
||||
generator_config = GeneratorConfig(
|
||||
model_path=m["model_path"],
|
||||
engine=EngineConfig(
|
||||
num_gpus=1,
|
||||
use_fsdp_inference=False,
|
||||
offload=OffloadConfig(text_encoder=False, dit=False, vae=False, pin_cpu_memory=False),
|
||||
),
|
||||
)
|
||||
load_start = time.perf_counter()
|
||||
generator = VideoGenerator.from_config(generator_config)
|
||||
load_time = time.perf_counter() - load_start
|
||||
|
||||
request = GenerationRequest(
|
||||
prompt=m["prompt"],
|
||||
sampling=m["sampling"],
|
||||
output=OutputConfig(output_path=OUTPUT_PATH, output_video_name=f"v2_{m['family']}",
|
||||
save_video=True, return_frames=False),
|
||||
)
|
||||
gen_start = time.perf_counter()
|
||||
result = generator.generate(request)
|
||||
gen_time = time.perf_counter() - gen_start
|
||||
|
||||
print(f"[{m['family']:10s}] load={load_time:6.1f}s gen={gen_time:6.1f}s -> {result.video_path}")
|
||||
|
||||
|
||||
def main() -> None:
|
||||
for m in MODELS:
|
||||
run_one(m)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,30 +0,0 @@
|
||||
"""v2 port of basic_self_forcing_causal.py — SF-causal Wan2.1 (CausalWanTransformer3DModel) through
|
||||
the v2 VideoGenerator (chunk_rollout loop).
|
||||
|
||||
Same convenience API as upstream; only delta is importing VideoGenerator from v2. NOTE: the v2 causal
|
||||
loop runs per-chunk few-step (not the upstream kv-cache streaming + SF schedule), so output is coherent
|
||||
but lower-fidelity (a documented gap). num_frames is set by the card's chunk schedule; height/width
|
||||
drive the latent geometry.
|
||||
"""
|
||||
from v2 import VideoGenerator
|
||||
from fastvideo.api.sampling_param import SamplingParam
|
||||
|
||||
OUTPUT_PATH = "v2_video_samples_causal"
|
||||
|
||||
|
||||
def main() -> None:
|
||||
model_name = "wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers"
|
||||
generator = VideoGenerator.from_pretrained(
|
||||
model_name, num_gpus=1, text_encoder_cpu_offload=False, dit_cpu_offload=False)
|
||||
sampling_param = SamplingParam.from_pretrained(model_name)
|
||||
|
||||
prompt = ("A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes wide with "
|
||||
"interest. The playful yet serene atmosphere is complemented by soft natural light "
|
||||
"filtering through the petals. Mid-shot, warm and cheerful tones.")
|
||||
video = generator.generate_video(prompt, output_path=OUTPUT_PATH, output_video_name="causal_raccoon",
|
||||
save_video=True, sampling_param=sampling_param, height=480, width=832)
|
||||
print(f"Output: {video.video_path}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,40 +0,0 @@
|
||||
"""v2 port of basic_wan2_2.py — Wan2.2-T2V-A14B (MoE) through the v2 VideoGenerator.
|
||||
|
||||
Same convenience API as upstream; only delta is importing VideoGenerator from v2. A14B is a 2-expert
|
||||
MoE: WanTransformer3DModel x2 (in_ch=16, Wan2.1 geometry) with a boundary-timestep switch
|
||||
(boundary_ratio 0.875) — ported via build_wan22_a14b_card (BoundaryTimestepRouting: transformer =
|
||||
high-noise expert, transformer_2 = low-noise), reusing the Wan adapters for both experts.
|
||||
|
||||
NOTE: upstream runs A14B with num_gpus=2 + dit_cpu_offload=True ("DiT need to be offloaded for MoE").
|
||||
The v2 bring-up is single-GPU + resident (no offload), so the two 14B experts (~56GB bf16) + UMT5 are
|
||||
near an 80GB GPU's limit — this example uses reduced res/frames to fit. If it OOMs, the A14B card is
|
||||
still correct; it just needs the (not-yet-ported) MoE DiT CPU offload. See V2_PORTING_STATUS.md.
|
||||
"""
|
||||
from v2 import VideoGenerator
|
||||
|
||||
OUTPUT_PATH = "v2_video_samples_wan2_2_14B_t2v"
|
||||
|
||||
|
||||
def main() -> None:
|
||||
generator = VideoGenerator.from_pretrained(
|
||||
"Wan-AI/Wan2.2-T2V-A14B-Diffusers",
|
||||
num_gpus=1,
|
||||
use_fsdp_inference=False,
|
||||
dit_cpu_offload=False,
|
||||
vae_cpu_offload=False,
|
||||
text_encoder_cpu_offload=False,
|
||||
pin_cpu_memory=False,
|
||||
)
|
||||
|
||||
prompt = ("A majestic lion strides across the golden savanna, its powerful frame glistening under "
|
||||
"the warm afternoon sun. The tall grass ripples gently in the breeze. Low angle, steady "
|
||||
"tracking shot, cinematic.")
|
||||
# Reduced res/frames so the two resident 14B experts fit a single 80GB GPU (upstream: 720x1280x81).
|
||||
video = generator.generate_video(prompt, output_path=OUTPUT_PATH, output_video_name="wan22_a14b_lion",
|
||||
save_video=True, num_frames=17, height=480, width=832,
|
||||
num_inference_steps=20, guidance_scale=5.0)
|
||||
print(f"Output: {video.video_path}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,38 +0,0 @@
|
||||
"""v2 port of basic_wan2_2_ti2v.py — Wan2.2-TI2V-5B (T2V mode) through the v2 VideoGenerator.
|
||||
|
||||
Same convenience API as upstream (from_pretrained + generate_video); only delta is importing
|
||||
VideoGenerator from v2. Wan2.2-TI2V-5B reuses the Wan adapter classes (WanTransformer3DModel /
|
||||
AutoencoderKLWan / UMT5) with the higher-compression VAE geometry (z_dim=48, 16x spatial, 4x temporal).
|
||||
|
||||
NOTE: upstream also runs I2V (image_path=...). The v2 program here is T2V-only (image conditioning is
|
||||
not yet ported), so this mirrors the upstream *T2V* branch (prompt2). Modest res/frames for a quick run.
|
||||
"""
|
||||
from v2 import VideoGenerator
|
||||
|
||||
OUTPUT_PATH = "v2_video_samples_wan2_2_5B_ti2v"
|
||||
|
||||
|
||||
def main() -> None:
|
||||
model_name = "Wan-AI/Wan2.2-TI2V-5B-Diffusers"
|
||||
generator = VideoGenerator.from_pretrained(
|
||||
model_name,
|
||||
num_gpus=1,
|
||||
use_fsdp_inference=False,
|
||||
dit_cpu_offload=False,
|
||||
vae_cpu_offload=False,
|
||||
text_encoder_cpu_offload=False,
|
||||
pin_cpu_memory=False,
|
||||
)
|
||||
|
||||
# T2V mode (the v2 program is text-to-video; upstream's image_path I2V branch is not ported yet).
|
||||
prompt = ("A majestic lion strides across the golden savanna, its powerful frame glistening under "
|
||||
"the warm afternoon sun. The tall grass ripples gently in the breeze, enhancing the lion's "
|
||||
"commanding presence. Low angle, steady tracking shot, cinematic.")
|
||||
video = generator.generate_video(prompt, output_path=OUTPUT_PATH, output_video_name="wan22_ti2v_lion",
|
||||
save_video=True, num_frames=25, height=448, width=768,
|
||||
num_inference_steps=20, guidance_scale=5.0)
|
||||
print(f"Output: {video.video_path}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -12,6 +12,9 @@ set(_FASTVIDEO_USER_CUDA_ARCH "${CMAKE_CUDA_ARCHITECTURES}")
|
||||
if(NOT DEFINED GPU_BACKEND AND DEFINED ENV{GPU_BACKEND})
|
||||
set(GPU_BACKEND "$ENV{GPU_BACKEND}")
|
||||
endif()
|
||||
if(NOT GPU_BACKEND)
|
||||
set(GPU_BACKEND "CUDA")
|
||||
endif()
|
||||
|
||||
if(GPU_BACKEND STREQUAL "ROCM")
|
||||
enable_language(HIP)
|
||||
@@ -50,7 +53,16 @@ if(NOT GPU_BACKEND STREQUAL "ROCM")
|
||||
if(_FASTVIDEO_USER_CUDA_ARCH)
|
||||
# Caller pinned -DCMAKE_CUDA_ARCHITECTURES (which torch ignores); translate it
|
||||
# to the TORCH_CUDA_ARCH_LIST spelling: "121" -> "12.1", "90a" -> "9.0a".
|
||||
# Only numeric spellings translate; keywords like "native"/"all" would
|
||||
# otherwise be mangled into nonsense ("nativ.e").
|
||||
foreach(_fv_arch IN LISTS _FASTVIDEO_USER_CUDA_ARCH)
|
||||
if(NOT _fv_arch MATCHES "^[0-9]+[af]?$")
|
||||
message(FATAL_ERROR
|
||||
"fastvideo-kernel: CMAKE_CUDA_ARCHITECTURES='${_fv_arch}' is not "
|
||||
"supported. Use a numeric arch (e.g. 90a, 121), set "
|
||||
"TORCH_CUDA_ARCH_LIST directly (e.g. 9.0a), or unset both to "
|
||||
"auto-detect from the visible GPU.")
|
||||
endif()
|
||||
string(REGEX MATCH "[af]$" _fv_suffix "${_fv_arch}")
|
||||
string(REGEX REPLACE "[af]$" "" _fv_num "${_fv_arch}")
|
||||
string(REGEX REPLACE "(.)$" ".\\1" _fv_num "${_fv_num}") # dot before the last digit
|
||||
@@ -173,6 +185,14 @@ else()
|
||||
endif()
|
||||
endif()
|
||||
|
||||
# ThunderKittens headers don't compile on aarch64 hosts: plain char is unsigned
|
||||
# there, and tk's base_types.cuh brace-initializes signed-char vector members
|
||||
# from char (narrowing error). Skip TK until upstream is aarch64-clean.
|
||||
if(ENABLE_TK_KERNELS AND CMAKE_SYSTEM_PROCESSOR MATCHES "^(aarch64|arm64)$")
|
||||
message(STATUS "ThunderKittens kernels: forced OFF on ${CMAKE_SYSTEM_PROCESSOR} (tk headers are not aarch64-clean)")
|
||||
set(ENABLE_TK_KERNELS OFF)
|
||||
endif()
|
||||
|
||||
if(ENABLE_TK_KERNELS)
|
||||
message(STATUS "ThunderKittens kernels: ENABLED")
|
||||
else()
|
||||
@@ -263,11 +283,6 @@ set(CUDA_FLAGS
|
||||
"--expt-relaxed-constexpr"
|
||||
"-Xcompiler=-fno-strict-aliasing"
|
||||
"-Xcompiler=-fPIC"
|
||||
# ARM/aarch64 defaults `char` to unsigned, but ThunderKittens headers assume the
|
||||
# x86 signed-char behavior (else base_types.cuh hits "narrowing conversion from
|
||||
# char to signed char"). Force signed char so TK compiles on Grace Hopper; this
|
||||
# is a no-op on x86_64, where char is already signed.
|
||||
"-Xcompiler=-fsigned-char"
|
||||
"-DTORCH_COMPILE"
|
||||
"-Xnvlink=--verbose"
|
||||
"-Xptxas=--verbose"
|
||||
@@ -426,3 +441,14 @@ if(ENABLE_ATTN_QAT_INFER)
|
||||
install(TARGETS fp4attn_cuda LIBRARY DESTINATION .)
|
||||
install(TARGETS fp4quant_cuda LIBRARY DESTINATION .)
|
||||
endif()
|
||||
|
||||
# One-look answer to "what is this build producing?" — kept last so it is the
|
||||
# final thing configure prints. The per-kernel matrix lives in README.md.
|
||||
message(STATUS "============== fastvideo-kernel build summary ==============")
|
||||
message(STATUS "host / backend: ${CMAKE_SYSTEM_PROCESSOR} / ${GPU_BACKEND}")
|
||||
message(STATUS "TORCH_CUDA_ARCH_LIST: ${TORCH_CUDA_ARCH_LIST}")
|
||||
message(STATUS "fastvideo_kernel_ops: ON (turbodiffusion int8-gemm/quant/rmsnorm/layernorm, all listed archs)")
|
||||
message(STATUS " + TK sta/block_sparse (sm_90a only): ${ENABLE_TK_KERNELS}")
|
||||
message(STATUS "fp4attn/fp4quant (sm_120a only, CUDA >= 12.8): ${ENABLE_ATTN_QAT_INFER}")
|
||||
message(STATUS "Triton fallbacks ship in python/fastvideo_kernel/triton_kernels regardless.")
|
||||
message(STATUS "============================================================")
|
||||
|
||||
@@ -2,6 +2,42 @@
|
||||
|
||||
CUDA kernels for FastVideo video generation.
|
||||
|
||||
## Kernel inventory
|
||||
|
||||
Compiled CUDA extensions (CMake, see the build summary printed at the end of every configure):
|
||||
|
||||
| Extension | Kernels | Sources | GPU arch | Build gate |
|
||||
|---|---|---|---|---|
|
||||
| `fastvideo_kernel._C.fastvideo_kernel_ops` | TurboDiffusion INT8 GEMM, quant, RMSNorm, LayerNorm | `csrc/turbodiffusion/` | every arch in `TORCH_CUDA_ARCH_LIST` | always built |
|
||||
| same extension, optional part | ThunderKittens sliding-tile attention (`sta_fwd`) and VSA block-sparse (`block_sparse_fwd/bwd`) | `csrc/attention/*_h100.cu` | Hopper `sm_90a` only | `FASTVIDEO_KERNEL_BUILD_TK` (AUTO = ON iff `9.0a` is in the arch list; always OFF on aarch64 hosts — TK headers don't compile there) |
|
||||
| `fp4attn_cuda`, `fp4quant_cuda` | FP4 attention + quantization ("attn_qat_infer", modified SageAttention3) | `attn_qat_infer/` | consumer Blackwell `sm_120a` only, CUDA ≥ 12.8 | `FASTVIDEO_KERNEL_BUILD_ATTN_QAT_INFER` (AUTO = ON iff `12.0a` is in the arch list) |
|
||||
|
||||
Runtime-JIT kernels (no build step, ship in every wheel/image):
|
||||
|
||||
| Kernels | Where | Used when |
|
||||
|---|---|---|
|
||||
| Triton: STA, VSA block-sparse, SLA, fused compress+topk, FP4 QAT training, quant/norm utils | `python/fastvideo_kernel/triton_kernels/` | automatic fallback when the matching C++ op is absent (`ops.py`, `turbodiffusion_ops.py`) |
|
||||
| FA4 CuTe-DSL block-sparse (VSA-256 fastpath on `sm_100`) | `block_sparse_attn_cute_fwd.py` | optional `flash_attn.cute` dependency, see below |
|
||||
| VMoBA `moba_attn_varlen` | `vmoba.py` | wraps flash-attn varlen |
|
||||
|
||||
## What gets built where, and when
|
||||
|
||||
| Surface | Trigger | Leg | `TORCH_CUDA_ARCH_LIST` | TK | FP4 |
|
||||
|---|---|---|---|---|---|
|
||||
| PyPI wheels (`.github/workflows/publish-kernel.yml`) | version bump in `fastvideo-kernel/pyproject.toml` on main, or manual dispatch | x86_64 cu126 | `9.0a` | ON | — (CUDA < 12.8) |
|
||||
| | | x86_64 cu130 | `9.0a;12.0a` | ON | ON |
|
||||
| | | aarch64 cu130 | `10.0a;12.0a` | — | ON |
|
||||
| Docker images `ghcr.io/hao-ai-lab/fastvideo/fastvideo-dev` (`.github/workflows/infra-build-image.yml`) | `docker/Dockerfile` changes on main, or manual dispatch | amd64 cuda12.6.3 + cuda13.0.0 | `9.0a` | ON | — |
|
||||
| | | arm64 cuda12.6.3 (GH200) | `9.0a` | — (aarch64) | — |
|
||||
| | | arm64 cuda13.0.0 (GB10 / DGX Spark) | `12.1` | — | — |
|
||||
| Local `./build.sh` | manual | probes the visible GPU via torch | detected | ON iff sm_90 (non-aarch64 host) | ON iff sm_120 |
|
||||
|
||||
Notes:
|
||||
|
||||
- No Docker image ships the FP4 kernels; only the x86_64/aarch64 cu130 wheels do.
|
||||
- On arm64 images (GH200 included) STA/VSA run on the Triton fallbacks, since TK never builds on aarch64.
|
||||
- Kernel tests run on Buildkite GPU CI for PRs touching `fastvideo-kernel/**` (see `.buildkite/pipeline.yml`).
|
||||
|
||||
## Installation
|
||||
|
||||
### Standard Installation (Local Development)
|
||||
|
||||
@@ -46,6 +46,23 @@ fi
|
||||
if git rev-parse --git-dir >/dev/null 2>&1; then
|
||||
git submodule update --init --recursive include/cutlass include/tk
|
||||
fi
|
||||
# Fail fast with a clear message if the headers are still missing (e.g. a
|
||||
# Docker context that excluded .git AND the submodule contents) instead of
|
||||
# dying later in a wall of nvcc include errors. CUTLASS is consumed by the
|
||||
# always-built turbodiffusion sources, so it is a hard error; ThunderKittens
|
||||
# only feeds the TK-gated Hopper kernels, so a missing tree just warns (the
|
||||
# TK gate resolves later, and non-SM90/ROCm builds never touch it).
|
||||
if [ ! -d include/cutlass/include ]; then
|
||||
echo "ERROR: include/cutlass/include is missing. Outside a git checkout the" >&2
|
||||
echo " CUTLASS sources must already be present (run" >&2
|
||||
echo " 'git submodule update --init --recursive include/cutlass include/tk'" >&2
|
||||
echo " in the source checkout, or include them in the build context)." >&2
|
||||
exit 1
|
||||
fi
|
||||
if [ ! -d include/tk/include ]; then
|
||||
echo "WARNING: include/tk/include is missing; ThunderKittens (Hopper sm_90a)" >&2
|
||||
echo " kernels cannot be built. Fine for non-SM90/ROCm targets." >&2
|
||||
fi
|
||||
|
||||
# Install build dependencies
|
||||
uv pip install scikit-build-core cmake ninja
|
||||
|
||||
@@ -159,6 +159,16 @@ def legacy_from_pretrained_to_config(
|
||||
preset_refine["guidance_scale"] = value
|
||||
elif key in {"enable_stage_verification", "use_fsdp_inference", "disable_autocast"}:
|
||||
engine[key] = value
|
||||
elif key == "batching_mode":
|
||||
engine.setdefault("batching", {})["mode"] = value
|
||||
elif key == "batching_max_size":
|
||||
engine.setdefault("batching", {})["max_size"] = value
|
||||
elif key == "batching_delay_ms":
|
||||
engine.setdefault("batching", {})["delay_ms"] = value
|
||||
elif key == "batching_config":
|
||||
engine.setdefault("batching", {})["config_path"] = value
|
||||
elif key == "enable_batching_metrics":
|
||||
engine.setdefault("batching", {})["enable_metrics"] = value
|
||||
elif key == "override_text_encoder_quant":
|
||||
quantization["text_encoder_quant"] = value
|
||||
elif key == "workload_type":
|
||||
@@ -244,6 +254,11 @@ def generator_config_to_fastvideo_args(config: GeneratorConfig | Mapping[str, An
|
||||
"enable_stage_verification": engine.enable_stage_verification,
|
||||
"use_fsdp_inference": engine.use_fsdp_inference,
|
||||
"disable_autocast": engine.disable_autocast,
|
||||
"batching_mode": engine.batching.mode,
|
||||
"batching_max_size": engine.batching.max_size,
|
||||
"batching_delay_ms": engine.batching.delay_ms,
|
||||
"batching_config": engine.batching.config_path,
|
||||
"enable_batching_metrics": engine.batching.enable_metrics,
|
||||
}
|
||||
if normalized.pipeline.workload_type is not None:
|
||||
kwargs["workload_type"] = normalized.pipeline.workload_type
|
||||
|
||||
@@ -71,6 +71,15 @@ class QuantizationConfig:
|
||||
transformer_quant: str | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class BatchingConfig:
|
||||
mode: Literal["disabled", "dynamic"] = "disabled"
|
||||
max_size: int = 1
|
||||
delay_ms: float = 0.0
|
||||
config_path: str | None = None
|
||||
enable_metrics: bool = False
|
||||
|
||||
|
||||
@dataclass
|
||||
class EngineConfig:
|
||||
num_gpus: int = 1
|
||||
@@ -82,6 +91,7 @@ class EngineConfig:
|
||||
use_fsdp_inference: bool = False
|
||||
disable_autocast: bool = False
|
||||
quantization: QuantizationConfig | None = None
|
||||
batching: BatchingConfig = field(default_factory=BatchingConfig)
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -280,6 +290,7 @@ class ServeConfig:
|
||||
|
||||
|
||||
__all__ = [
|
||||
"BatchingConfig",
|
||||
"CompileConfig",
|
||||
"ComponentConfig",
|
||||
"ContinuationState",
|
||||
|
||||
@@ -1,15 +1,39 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
import importlib.util
|
||||
import os
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from dataclasses import dataclass
|
||||
|
||||
try:
|
||||
from fastvideo.attention.utils.flash_attn_cute import flash_attn_func
|
||||
from fastvideo import envs
|
||||
from fastvideo.attention.backends.abstract import (
|
||||
AttentionBackend,
|
||||
AttentionImpl,
|
||||
AttentionMetadata,
|
||||
AttentionMetadataBuilder,
|
||||
)
|
||||
from fastvideo.logger import init_logger
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
# FA4 (flash_attn.cute) is explicit opt-in via FASTVIDEO_FA4=1, mirroring the
|
||||
# kernel package's FASTVIDEO_VSA_CUTEDSL: its CuTeDSL kernels JIT-compile per
|
||||
# shape family and can fail at runtime on some arch/shape combinations, so it
|
||||
# is never auto-selected just because it is installed. Below sm90 a capability
|
||||
# gate in flash_attn_cute routes to FA2 the calls FA4 cannot serve there:
|
||||
# grad-enabled (its backward asserts sm90+) and GQA (pack_gqa fails CuTeDSL
|
||||
# JIT, observed on sm_89).
|
||||
if envs.FASTVIDEO_FA4:
|
||||
try:
|
||||
from fastvideo.attention.utils.flash_attn_cute import flash_attn_func
|
||||
except ImportError as e:
|
||||
raise RuntimeError(f"FASTVIDEO_FA4=1 but flash_attn.cute (FA4) is not usable ({e}); "
|
||||
"fix the FA4 install (see the flash-attn-4 pin in pyproject.toml) "
|
||||
"or unset FASTVIDEO_FA4.") from e
|
||||
fa_version = "4"
|
||||
except ImportError:
|
||||
else:
|
||||
try:
|
||||
from flash_attn_interface import flash_attn_func as flash_attn_3_func
|
||||
|
||||
@@ -21,6 +45,12 @@ except ImportError:
|
||||
from flash_attn import flash_attn_func as flash_attn_2_func
|
||||
flash_attn_func = flash_attn_2_func
|
||||
fa_version = "2"
|
||||
try:
|
||||
if importlib.util.find_spec("flash_attn.cute") is not None:
|
||||
logger.info("flash_attn.cute (FA4) is installed but not enabled; "
|
||||
"set FASTVIDEO_FA4=1 to use it for inference.")
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
# torch.compile traceability: the FA4/cute path (fa_version=="4") is
|
||||
# already a registered torch.library custom op, so dynamo treats it as a
|
||||
@@ -87,9 +117,10 @@ if fa_version in ("2", "3"):
|
||||
return _fa_default(q, k, v, softmax_scale=softmax_scale, causal=causal)
|
||||
return torch.ops.fastvideo._flash_attn_default_forward(q, k, v, softmax_scale, causal)
|
||||
elif fa_version == "4":
|
||||
# FA4 path: `flash_attn_func` is already a torch.library custom op
|
||||
# (registered in `fastvideo.attention.utils.flash_attn_cute`), so a
|
||||
# passthrough is enough — no extra registration needed.
|
||||
# FA4 path: `flash_attn_func` (from `flash_attn_cute`) goes through a
|
||||
# registered torch.library custom op (with an FA4 backward on sm90+;
|
||||
# grad-enabled and GQA calls below sm90 route to FA2), so a passthrough
|
||||
# is enough — no extra registration needed.
|
||||
def flash_attn_func_compilable(q, k, v, softmax_scale=None, causal=False):
|
||||
return flash_attn_func(q, k, v, softmax_scale=softmax_scale, causal=causal)
|
||||
else:
|
||||
@@ -99,17 +130,6 @@ else:
|
||||
raise RuntimeError(f"Unsupported FlashAttention version: {fa_version!r} — expected "
|
||||
f"'2', '3', or '4' from the import probe above.")
|
||||
|
||||
from fastvideo.attention.backends.abstract import (
|
||||
AttentionBackend,
|
||||
AttentionImpl,
|
||||
AttentionMetadata,
|
||||
AttentionMetadataBuilder,
|
||||
)
|
||||
from fastvideo.logger import init_logger
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
_WARNED_NON_FA_DTYPE = False
|
||||
logger.info("Using FlashAttention-%s backend", fa_version)
|
||||
|
||||
# FP4 FA4 support: quantize Q/K to NVFP4 E2M1 for block-scaled MMA on Blackwell.
|
||||
@@ -271,12 +291,8 @@ class FlashAttentionImpl(AttentionImpl):
|
||||
# SP). Cast through bf16 and restore, matching TORCH_SDPA's tolerance.
|
||||
orig_dtype = query.dtype
|
||||
if orig_dtype not in (torch.float16, torch.bfloat16):
|
||||
global _WARNED_NON_FA_DTYPE
|
||||
if not _WARNED_NON_FA_DTYPE:
|
||||
_WARNED_NON_FA_DTYPE = True
|
||||
logger.warning(
|
||||
"FLASH_ATTN received %s inputs; casting to bfloat16 for the "
|
||||
"kernel and restoring on output.", orig_dtype)
|
||||
logger.warning_once(f"FLASH_ATTN received {orig_dtype} inputs; casting to "
|
||||
f"bfloat16 for the kernel and restoring on output.")
|
||||
query = query.to(torch.bfloat16)
|
||||
key = key.to(torch.bfloat16)
|
||||
value = value.to(torch.bfloat16)
|
||||
|
||||
@@ -4,10 +4,9 @@ import functools
|
||||
from collections.abc import Callable
|
||||
|
||||
import torch
|
||||
from flash_attn import flash_attn_func as _flash_attn_2_func
|
||||
from flash_attn import flash_attn_varlen_func as _flash_attn_2_varlen_func
|
||||
|
||||
from fastvideo.logger import init_logger
|
||||
from fastvideo.platforms import current_platform
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
@@ -16,7 +15,8 @@ if torch.cuda.is_available():
|
||||
from flash_attn.cute.interface import _flash_attn_bwd, _flash_attn_fwd
|
||||
except ImportError:
|
||||
# flash_attn.cute (FA4) is simply not installed -- expected on builds
|
||||
# without it; callers fall back to FA3/FA2 quietly.
|
||||
# without it; callers handle the ImportError (the FASTVIDEO_FA4 gate in
|
||||
# flash_attn.py raises, the FP4 probe treats FA4 as unavailable).
|
||||
raise
|
||||
except Exception as e:
|
||||
# flash_attn.cute IS installed but failed to import -- almost always an
|
||||
@@ -24,23 +24,57 @@ if torch.cuda.is_available():
|
||||
# 'cutlass.cute.core' has no attribute 'ThrMma'" (an AttributeError, not
|
||||
# ImportError). This is fixable by pinning a compatible
|
||||
# nvidia-cutlass-dsl, so warn loudly, then re-raise as ImportError so
|
||||
# callers fall back to FA3/FA2 instead of crashing worker init.
|
||||
# callers can handle it uniformly.
|
||||
logger.warning(
|
||||
"flash_attn.cute (FA4) is installed but failed to import (%r); "
|
||||
"falling back to FA3/FA2. This is usually an nvidia-cutlass-dsl "
|
||||
"version mismatch -- pin a compatible nvidia-cutlass-dsl to "
|
||||
"restore FA4.", e)
|
||||
"flash_attn.cute (FA4) is installed but failed to import (%r). "
|
||||
"This is usually an nvidia-cutlass-dsl version mismatch -- pin a "
|
||||
"compatible nvidia-cutlass-dsl to restore FA4.", e)
|
||||
raise ImportError(f"flash_attn.cute (FA4) import failed: {e!r}") from e
|
||||
else:
|
||||
# This error will be caught in flash_attn.py or flash_attn_no_pad.py
|
||||
raise ImportError("flash_attn.cute is only available on CUDA devices; this error must be handled internally")
|
||||
|
||||
try:
|
||||
# FA2 serves the calls FA4 cute cannot on pre-sm90 GPUs (backward, GQA).
|
||||
# Optional so FA4-only installs can still import this module.
|
||||
from flash_attn import flash_attn_func as _flash_attn_2_func
|
||||
from flash_attn import flash_attn_varlen_func as _flash_attn_2_varlen_func
|
||||
except ImportError:
|
||||
_flash_attn_2_func = None
|
||||
_flash_attn_2_varlen_func = None
|
||||
|
||||
|
||||
def _check_dropout(dropout_p: float) -> None:
|
||||
if dropout_p != 0.0:
|
||||
raise NotImplementedError(f"flash_attn.cute does not support dropout (got dropout_p={dropout_p})")
|
||||
|
||||
|
||||
@functools.cache
|
||||
def _sm90_or_newer() -> bool:
|
||||
return current_platform.has_device_capability(90)
|
||||
|
||||
|
||||
def _use_fa2(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> bool:
|
||||
if _sm90_or_newer():
|
||||
return False
|
||||
# Pre-sm90 FA4 cute limitations, both served by FA2 (deterministic
|
||||
# capability gate, not a runtime fallback):
|
||||
# * the backward asserts sm90+ (L40S/sm_89 dies on its arch check);
|
||||
# * GQA fails CuTeDSL JIT in pack_gqa ("ValueError: Operation creation
|
||||
# failed", observed on sm_89 with HunyuanGameCraft/LTX2).
|
||||
if q.shape[-2] != k.shape[-2]:
|
||||
return True
|
||||
return torch.is_grad_enabled() and any(t.requires_grad for t in (q, k, v))
|
||||
|
||||
|
||||
def _fa2_or_raise(fa2_func: Callable | None) -> Callable:
|
||||
if fa2_func is None:
|
||||
raise RuntimeError("this attention call cannot run on FA4 cute below sm90 (its backward and "
|
||||
"GQA support require sm90+) and flash-attn 2, which serves it there, is "
|
||||
"not installed.")
|
||||
return fa2_func
|
||||
|
||||
|
||||
@torch.library.custom_op(
|
||||
"fastvideo::_flash_attn_cute_forward",
|
||||
mutates_args=(),
|
||||
@@ -243,70 +277,6 @@ torch.library.register_autograd(
|
||||
)
|
||||
|
||||
|
||||
# FA4's CuTeDSL kernels JIT-compile per shape family, and some configurations
|
||||
# fail MLIR op creation at runtime even though the import succeeded (observed:
|
||||
# GQA models on sm_89 dying in pack_gqa with "ValueError: Operation creation
|
||||
# failed"). Degrade to FA2 once, process-wide, instead of crashing inference.
|
||||
class _FA4Policy:
|
||||
"""Per-call gate for the FA4 cute fast path, with FA2 as the fallback.
|
||||
|
||||
FA4 is skipped when:
|
||||
* a previous call failed at runtime -- CuTeDSL JIT compilation is
|
||||
shape-dependent, so the first failure disables FA4 for the rest of
|
||||
the process instead of retrying a broken JIT on every call; or
|
||||
* the call needs autograd -- FA4's backward asserts sm90+ (L40S/sm_89
|
||||
dies on its arch check) and is unvalidated for training in this repo
|
||||
(its lse is not even allocated through our inference-shaped custom
|
||||
op), so training keeps the pre-FA4 behavior: FA2 on every device.
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.broken = False
|
||||
|
||||
def use_fa4(self, *tensors: torch.Tensor) -> bool:
|
||||
if self.broken:
|
||||
return False
|
||||
return not (torch.is_grad_enabled() and any(t.requires_grad for t in tensors))
|
||||
|
||||
def mark_broken(self, error: Exception) -> None:
|
||||
if not self.broken:
|
||||
self.broken = True
|
||||
logger.warning(
|
||||
"flash_attn.cute (FA4) failed at runtime (%r); falling back "
|
||||
"to FA2 for the rest of this process.", error)
|
||||
|
||||
|
||||
_FA4 = _FA4Policy()
|
||||
|
||||
|
||||
def _with_fa2_fallback(fa2_func: Callable) -> Callable:
|
||||
"""Pair an FA4 cute wrapper with its FA2 twin of the same signature.
|
||||
|
||||
The decorated body runs only when ``_FA4`` allows it; otherwise (or after
|
||||
the first FA4 runtime failure) the call is served by ``fa2_func``.
|
||||
``NotImplementedError`` is a contract error (e.g. dropout), not a JIT
|
||||
failure, so it propagates without disabling FA4.
|
||||
"""
|
||||
|
||||
def decorator(fa4_func: Callable) -> Callable:
|
||||
|
||||
@functools.wraps(fa4_func)
|
||||
def wrapper(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, *args, **kwargs) -> torch.Tensor:
|
||||
if _FA4.use_fa4(q, k, v):
|
||||
try:
|
||||
return fa4_func(q, k, v, *args, **kwargs)
|
||||
except NotImplementedError:
|
||||
raise
|
||||
except Exception as e: # CuTeDSL compile errors surface as ValueError
|
||||
_FA4.mark_broken(e)
|
||||
return fa2_func(q, k, v, *args, **kwargs)
|
||||
|
||||
return wrapper
|
||||
|
||||
return decorator
|
||||
|
||||
|
||||
@_with_fa2_fallback(_flash_attn_2_func)
|
||||
def flash_attn_func(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
@@ -317,6 +287,16 @@ def flash_attn_func(
|
||||
deterministic: bool = False,
|
||||
) -> torch.Tensor:
|
||||
"""Only returns the output, not the lse."""
|
||||
if _use_fa2(q, k, v):
|
||||
return _fa2_or_raise(_flash_attn_2_func)(
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
dropout_p=dropout_p,
|
||||
softmax_scale=softmax_scale,
|
||||
causal=causal,
|
||||
deterministic=deterministic,
|
||||
)
|
||||
_check_dropout(dropout_p)
|
||||
out, _ = torch.ops.fastvideo._flash_attn_cute_forward(q, k, v, softmax_scale, causal, deterministic)
|
||||
return out
|
||||
@@ -392,7 +372,6 @@ def flash_attn_fp4_func(
|
||||
return torch.ops.fastvideo._flash_attn_cute_fp4_forward(q, k, v, sfq, sfk, softmax_scale, causal)
|
||||
|
||||
|
||||
@_with_fa2_fallback(_flash_attn_2_varlen_func)
|
||||
def flash_attn_varlen_func(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
@@ -407,6 +386,20 @@ def flash_attn_varlen_func(
|
||||
deterministic: bool = False,
|
||||
) -> torch.Tensor:
|
||||
"""Only returns the output, not the lse."""
|
||||
if _use_fa2(q, k, v):
|
||||
return _fa2_or_raise(_flash_attn_2_varlen_func)(
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
cu_seqlens_q,
|
||||
cu_seqlens_k,
|
||||
max_seqlen_q,
|
||||
max_seqlen_k,
|
||||
dropout_p=dropout_p,
|
||||
softmax_scale=softmax_scale,
|
||||
causal=causal,
|
||||
deterministic=deterministic,
|
||||
)
|
||||
_check_dropout(dropout_p)
|
||||
out, _ = torch.ops.fastvideo._flash_attn_cute_varlen_forward(
|
||||
q,
|
||||
|
||||
@@ -21,24 +21,35 @@ from einops import rearrange
|
||||
from flash_attn import flash_attn_varlen_qkvpacked_func
|
||||
from flash_attn.bert_padding import pad_input, unpad_input
|
||||
|
||||
from fastvideo import envs
|
||||
|
||||
|
||||
def _resolve_flash_attn_varlen_func() -> Any:
|
||||
try:
|
||||
from fastvideo.attention.utils.flash_attn_cute import (
|
||||
flash_attn_varlen_func as flash_attn_varlen_func_cute, )
|
||||
if envs.FASTVIDEO_FA4:
|
||||
# FA4 cute is explicit opt-in (see fastvideo/attention/backends/
|
||||
# flash_attn.py); with FASTVIDEO_FA4=1 an unimportable FA4 build must
|
||||
# fail loudly here rather than fall through to FA3/FA2. RuntimeError,
|
||||
# not ImportError: importers like bsa_attn.py treat ImportError as
|
||||
# "flash-attn not installed" and silently degrade to reference kernels.
|
||||
try:
|
||||
from fastvideo.attention.utils.flash_attn_cute import (
|
||||
flash_attn_varlen_func as flash_attn_varlen_func_cute, )
|
||||
except ImportError as e:
|
||||
raise RuntimeError(f"FASTVIDEO_FA4=1 but flash_attn.cute (FA4) is not usable ({e}); "
|
||||
"fix the FA4 install (see the flash-attn-4 pin in pyproject.toml) "
|
||||
"or unset FASTVIDEO_FA4.") from e
|
||||
|
||||
return flash_attn_varlen_func_cute
|
||||
try:
|
||||
from flash_attn_interface import (
|
||||
flash_attn_varlen_func as flash_attn_varlen_func_interface, )
|
||||
|
||||
return flash_attn_varlen_func_interface
|
||||
except ImportError:
|
||||
try:
|
||||
from flash_attn_interface import (
|
||||
flash_attn_varlen_func as flash_attn_varlen_func_interface, )
|
||||
from flash_attn import (
|
||||
flash_attn_varlen_func as flash_attn_varlen_func_flash, )
|
||||
|
||||
return flash_attn_varlen_func_interface
|
||||
except ImportError:
|
||||
from flash_attn import (
|
||||
flash_attn_varlen_func as flash_attn_varlen_func_flash, )
|
||||
|
||||
return flash_attn_varlen_func_flash
|
||||
return flash_attn_varlen_func_flash
|
||||
|
||||
|
||||
flash_attn_varlen_func_impl = _resolve_flash_attn_varlen_func()
|
||||
|
||||
@@ -0,0 +1,26 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""Dynamic generation batching helpers."""
|
||||
|
||||
from fastvideo.batching.admission import (
|
||||
AdmissionLimit,
|
||||
BatchAdmissionController,
|
||||
BatchingRule,
|
||||
load_batching_config,
|
||||
)
|
||||
from fastvideo.batching.signature import (
|
||||
BatchCompatibility,
|
||||
can_dynamic_batch,
|
||||
dynamic_batch_signature,
|
||||
resolution_key,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"AdmissionLimit",
|
||||
"BatchAdmissionController",
|
||||
"BatchCompatibility",
|
||||
"BatchingRule",
|
||||
"can_dynamic_batch",
|
||||
"dynamic_batch_signature",
|
||||
"load_batching_config",
|
||||
"resolution_key",
|
||||
]
|
||||
@@ -0,0 +1,297 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
from dataclasses import dataclass
|
||||
from difflib import get_close_matches
|
||||
from typing import Any
|
||||
|
||||
from fastvideo.batching.signature import resolution_key
|
||||
from fastvideo.fastvideo_args import FastVideoArgs
|
||||
from fastvideo.logger import init_logger
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
_BYTES_PER_GB = 1024**3
|
||||
_BATCHING_RULE_KEYS = frozenset({
|
||||
"model",
|
||||
"model_contains",
|
||||
"resolution",
|
||||
"device_memory_gb_min",
|
||||
"device_memory_gb_max",
|
||||
"offload",
|
||||
"max_batch_size",
|
||||
"max_cost",
|
||||
"calibration",
|
||||
})
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class AdmissionLimit:
|
||||
max_batch_size: int
|
||||
max_cost: float | None = None
|
||||
cap_reason: str | None = None
|
||||
|
||||
def reject_reason(self, *, batch_size: int, batch_cost: float) -> str | None:
|
||||
if batch_size > self.max_batch_size:
|
||||
return self.cap_reason or f"config_cap:{self.max_batch_size}"
|
||||
if self.max_cost is not None and batch_cost > self.max_cost:
|
||||
return f"cost_budget:{batch_cost:.0f}>{self.max_cost:.0f}"
|
||||
return None
|
||||
|
||||
def stop_reason_for_next_cost(self, next_batch_cost: float) -> str | None:
|
||||
if self.max_cost is not None and next_batch_cost > self.max_cost:
|
||||
return f"cost_budget_next:{next_batch_cost:.0f}>{self.max_cost:.0f}"
|
||||
return None
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class BatchingRule:
|
||||
model: str | None = None
|
||||
model_contains: str | None = None
|
||||
resolution: str | None = None
|
||||
device_memory_gb_min: float | None = None
|
||||
device_memory_gb_max: float | None = None
|
||||
offload: bool | None = None
|
||||
max_batch_size: int = 1
|
||||
max_cost: float | None = None
|
||||
source: str = "user"
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: dict[str, Any], *, source: str) -> BatchingRule:
|
||||
if not isinstance(data, dict):
|
||||
raise ValueError(f"batching config rule from {source} must be an object, got {type(data).__name__}")
|
||||
_validate_rule_keys(data, source=source)
|
||||
if "max_batch_size" not in data:
|
||||
raise ValueError("batching config rule requires max_batch_size")
|
||||
|
||||
rule = cls(
|
||||
model=_optional_str(data.get("model")),
|
||||
model_contains=_optional_str(data.get("model_contains")),
|
||||
resolution=_optional_str(data.get("resolution")),
|
||||
device_memory_gb_min=_optional_float(data.get("device_memory_gb_min")),
|
||||
device_memory_gb_max=_optional_float(data.get("device_memory_gb_max")),
|
||||
offload=_optional_bool(data.get("offload")),
|
||||
max_batch_size=int(data["max_batch_size"]),
|
||||
max_cost=_optional_float(data.get("max_cost")),
|
||||
source=source,
|
||||
)
|
||||
rule.validate()
|
||||
return rule
|
||||
|
||||
def validate(self) -> None:
|
||||
if self.model is not None and self.model_contains is not None:
|
||||
raise ValueError("batching config rule cannot set both model and model_contains")
|
||||
if self.model is None and self.model_contains is None:
|
||||
raise ValueError("batching config rule requires model or model_contains")
|
||||
if self.max_batch_size < 1:
|
||||
raise ValueError("batching config rule max_batch_size must be >= 1")
|
||||
if self.max_cost is not None and self.max_cost <= 0.0:
|
||||
raise ValueError("batching config rule max_cost must be > 0")
|
||||
if (self.device_memory_gb_min is not None and self.device_memory_gb_max is not None
|
||||
and self.device_memory_gb_min > self.device_memory_gb_max):
|
||||
raise ValueError("batching config rule device_memory_gb_min must be <= device_memory_gb_max")
|
||||
|
||||
def matches(
|
||||
self,
|
||||
*,
|
||||
model_path: str,
|
||||
resolution: str | None,
|
||||
device_memory_gb: float | None,
|
||||
offload: bool,
|
||||
) -> bool:
|
||||
if self.model is not None and self.model != model_path:
|
||||
return False
|
||||
if self.model_contains is not None and self.model_contains not in model_path:
|
||||
return False
|
||||
if self.resolution not in (None, "*") and self.resolution != resolution:
|
||||
return False
|
||||
if self.offload is not None and self.offload != offload:
|
||||
return False
|
||||
if device_memory_gb is None:
|
||||
return True
|
||||
if self.device_memory_gb_min is not None and device_memory_gb < self.device_memory_gb_min:
|
||||
return False
|
||||
return not (self.device_memory_gb_max is not None and device_memory_gb > self.device_memory_gb_max)
|
||||
|
||||
|
||||
class BatchAdmissionController:
|
||||
|
||||
def __init__(self, fastvideo_args: FastVideoArgs, *, gpu_id: int = 0):
|
||||
self._mode = fastvideo_args.batching_mode
|
||||
self._user_max_batch_size = max(1, int(fastvideo_args.batching_max_size))
|
||||
self._model_path = fastvideo_args.model_path
|
||||
self._offload = bool(fastvideo_args.dit_cpu_offload or fastvideo_args.dit_layerwise_offload)
|
||||
self._device_memory_gb = self._get_device_memory_gb(gpu_id)
|
||||
self._rules = load_batching_config(fastvideo_args.batching_config)
|
||||
self._pipeline_config = fastvideo_args.pipeline_config
|
||||
|
||||
if self.enabled:
|
||||
logger.info(
|
||||
"Batch admission enabled: user_max=%d, device_memory=%.1fGiB, rules=%d",
|
||||
self._user_max_batch_size,
|
||||
self._device_memory_gb or 0.0,
|
||||
len(self._rules),
|
||||
)
|
||||
|
||||
@property
|
||||
def enabled(self) -> bool:
|
||||
return self._mode == "dynamic" and self._user_max_batch_size > 1
|
||||
|
||||
def reject_reason_for_candidate(self, current_requests: list[Any], candidate_request: Any) -> str | None:
|
||||
if not self.enabled:
|
||||
return None
|
||||
proposed = current_requests + [candidate_request]
|
||||
limit = self.limit_for(proposed[0])
|
||||
return limit.reject_reason(
|
||||
batch_size=len(proposed),
|
||||
batch_cost=self.estimate_batch_cost(proposed),
|
||||
)
|
||||
|
||||
def batch_is_full(self, requests: list[Any]) -> bool:
|
||||
if not self.enabled or not requests:
|
||||
return len(requests) >= self._user_max_batch_size
|
||||
|
||||
limit = self.limit_for(requests[0])
|
||||
if len(requests) >= limit.max_batch_size:
|
||||
return True
|
||||
|
||||
next_cost = self.estimate_batch_cost(requests + [requests[0]])
|
||||
return limit.max_cost is not None and next_cost > limit.max_cost
|
||||
|
||||
def limit_reason_for_batch(self, requests: list[Any]) -> str | None:
|
||||
if not self.enabled or not requests:
|
||||
return None
|
||||
|
||||
limit = self.limit_for(requests[0])
|
||||
if len(requests) >= limit.max_batch_size:
|
||||
return limit.cap_reason or f"config_cap:{limit.max_batch_size}"
|
||||
|
||||
next_cost = self.estimate_batch_cost(requests + [requests[0]])
|
||||
return limit.stop_reason_for_next_cost(next_cost)
|
||||
|
||||
def max_admissible_batch_size(self, request: Any) -> int:
|
||||
return self.limit_for(request).max_batch_size
|
||||
|
||||
def limit_for(self, request: Any) -> AdmissionLimit:
|
||||
rules = self._matching_rules(request)
|
||||
if not rules:
|
||||
return AdmissionLimit(max_batch_size=self._user_max_batch_size)
|
||||
|
||||
config_cap = min(rule.max_batch_size for rule in rules)
|
||||
max_batch_size = min(self._user_max_batch_size, config_cap)
|
||||
cap_reason = f"config_cap:{max_batch_size}" if max_batch_size < self._user_max_batch_size else None
|
||||
costs = [rule.max_cost for rule in rules if rule.max_cost is not None]
|
||||
return AdmissionLimit(
|
||||
max_batch_size=max(1, max_batch_size),
|
||||
max_cost=min(costs) if costs else None,
|
||||
cap_reason=cap_reason,
|
||||
)
|
||||
|
||||
def estimate_batch_cost(self, requests: list[Any]) -> float:
|
||||
return sum(float(self._pipeline_config.estimate_request_cost(request)) for request in requests)
|
||||
|
||||
def _matching_rules(self, request: Any) -> list[BatchingRule]:
|
||||
return [
|
||||
rule for rule in self._rules if rule.matches(
|
||||
model_path=self._model_path,
|
||||
resolution=resolution_key(request),
|
||||
device_memory_gb=self._device_memory_gb,
|
||||
offload=self._offload,
|
||||
)
|
||||
]
|
||||
|
||||
@staticmethod
|
||||
def _get_device_memory_gb(gpu_id: int) -> float | None:
|
||||
try:
|
||||
from fastvideo.platforms import current_platform
|
||||
|
||||
return current_platform.get_device_total_memory(gpu_id) / _BYTES_PER_GB
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
def load_batching_config(path: str | None) -> list[BatchingRule]:
|
||||
if path is None:
|
||||
return []
|
||||
|
||||
with open(path, encoding="utf-8") as f:
|
||||
payload = json.load(f)
|
||||
|
||||
source = os.path.abspath(path)
|
||||
entries = _config_entries(payload)
|
||||
rules = [BatchingRule.from_dict(entry, source=source) for entry in entries]
|
||||
if not rules:
|
||||
raise ValueError(f"batching config {source} does not contain any rules")
|
||||
return rules
|
||||
|
||||
|
||||
def _config_entries(payload: Any) -> list[dict[str, Any]]:
|
||||
if isinstance(payload, dict) and payload.get("schema_version") not in (None, 1):
|
||||
raise ValueError("batching config schema_version must be 1")
|
||||
if isinstance(payload, dict) and isinstance(payload.get("rules"), list):
|
||||
return payload["rules"]
|
||||
if isinstance(payload, list):
|
||||
return payload
|
||||
if isinstance(payload, dict):
|
||||
entries: list[dict[str, Any]] = []
|
||||
for key, value in payload.items():
|
||||
if key == "schema_version" or not isinstance(value, dict):
|
||||
continue
|
||||
model, _sep, resolution = key.partition("|")
|
||||
entry = dict(value)
|
||||
if model:
|
||||
entry.setdefault("model", model)
|
||||
if resolution:
|
||||
entry.setdefault("resolution", resolution)
|
||||
entries.append(entry)
|
||||
return entries
|
||||
raise ValueError("batching config must be a {'schema_version': 1, 'rules': [...]} object, "
|
||||
"a list of rules, or a mapping keyed by model|resolution")
|
||||
|
||||
|
||||
def _validate_rule_keys(data: dict[str, Any], *, source: str) -> None:
|
||||
unknown = sorted(set(data) - _BATCHING_RULE_KEYS)
|
||||
if not unknown:
|
||||
return
|
||||
|
||||
hints = []
|
||||
for key in unknown:
|
||||
matches = get_close_matches(key, _BATCHING_RULE_KEYS, n=1)
|
||||
if matches:
|
||||
hints.append(f"{key!r} (did you mean {matches[0]!r}?)")
|
||||
else:
|
||||
hints.append(repr(key))
|
||||
raise ValueError(f"batching config rule from {source} contains unknown key(s): {', '.join(hints)}")
|
||||
|
||||
|
||||
def _optional_str(value: Any) -> str | None:
|
||||
if value is None:
|
||||
return None
|
||||
return str(value)
|
||||
|
||||
|
||||
def _optional_float(value: Any) -> float | None:
|
||||
if value is None:
|
||||
return None
|
||||
return float(value)
|
||||
|
||||
|
||||
def _optional_bool(value: Any) -> bool | None:
|
||||
if value is None:
|
||||
return None
|
||||
if isinstance(value, bool):
|
||||
return value
|
||||
if isinstance(value, int | float):
|
||||
if value == 1.0:
|
||||
return True
|
||||
if value == 0.0:
|
||||
return False
|
||||
if isinstance(value, str):
|
||||
lowered = value.strip().lower()
|
||||
if lowered in ("1", "true", "yes", "y", "on"):
|
||||
return True
|
||||
if lowered in ("0", "false", "no", "n", "off"):
|
||||
return False
|
||||
raise ValueError(f"cannot parse boolean batching config value: {value!r}")
|
||||
@@ -0,0 +1,178 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
from __future__ import annotations
|
||||
|
||||
import dataclasses
|
||||
from dataclasses import dataclass
|
||||
from enum import Enum
|
||||
from typing import Any
|
||||
|
||||
from fastvideo.api.sampling_param import SamplingParam
|
||||
|
||||
_SIGNATURE_EXCLUDED_FIELDS = frozenset({
|
||||
"prompt",
|
||||
"prompt_path",
|
||||
"output_path",
|
||||
"output_video_name",
|
||||
"seed",
|
||||
"save_video",
|
||||
"return_frames",
|
||||
})
|
||||
|
||||
_UNSUPPORTED_DYNAMIC_BATCH_FIELDS = frozenset({
|
||||
"image_path",
|
||||
"pil_image",
|
||||
"video_path",
|
||||
"mouse_cond",
|
||||
"keyboard_cond",
|
||||
"grid_sizes",
|
||||
"pose",
|
||||
"camera_states",
|
||||
"camera_trajectory",
|
||||
"action_list",
|
||||
"action_speed_list",
|
||||
"gt_latents",
|
||||
"conditioning_mask",
|
||||
"c2ws_plucker_emb",
|
||||
"refine_from",
|
||||
"stage1_video",
|
||||
"trajectory_type",
|
||||
"movement_distance",
|
||||
"camera_rotation",
|
||||
"ltx2_images",
|
||||
"ltx2_conditioning_latent_stage1",
|
||||
"ltx2_conditioning_latent_stage2",
|
||||
"ltx2_video_conditions",
|
||||
"init_audio",
|
||||
"inpaint_audio",
|
||||
"inpaint_mask",
|
||||
"continuation_state",
|
||||
})
|
||||
|
||||
_UNSUPPORTED_EXTRA_KEYS = frozenset({
|
||||
"ltx2_audio_latents",
|
||||
"ltx2_audio_clean_latent",
|
||||
"ltx2_audio_denoise_mask",
|
||||
"audio_num_frames",
|
||||
"video_position_offset_sec",
|
||||
})
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class BatchCompatibility:
|
||||
can_batch: bool
|
||||
reason: str | None = None
|
||||
|
||||
|
||||
def resolution_key(request: Any) -> str:
|
||||
height = _first_scalar(getattr(request, "height", None))
|
||||
width = _first_scalar(getattr(request, "width", None))
|
||||
num_frames = _first_scalar(getattr(request, "num_frames", None))
|
||||
return f"{height}x{width}x{num_frames}"
|
||||
|
||||
|
||||
def dynamic_batch_signature(
|
||||
request: SamplingParam,
|
||||
*,
|
||||
extra: dict[str, Any] | None = None,
|
||||
) -> tuple[tuple[str, Any], ...]:
|
||||
"""Build a hashable compatibility signature for a generation request."""
|
||||
signature_items: list[tuple[str, Any]] = []
|
||||
for field in dataclasses.fields(request):
|
||||
if field.name in _SIGNATURE_EXCLUDED_FIELDS:
|
||||
continue
|
||||
signature_items.append((field.name, _freeze_signature_value(getattr(request, field.name, None))))
|
||||
if extra:
|
||||
signature_items.append(("extra", _freeze_signature_value(extra)))
|
||||
return tuple(signature_items)
|
||||
|
||||
|
||||
def can_dynamic_batch(
|
||||
base: SamplingParam,
|
||||
candidate: SamplingParam,
|
||||
*,
|
||||
base_extra: dict[str, Any] | None = None,
|
||||
candidate_extra: dict[str, Any] | None = None,
|
||||
) -> BatchCompatibility:
|
||||
"""Return whether two FastVideo generation requests can be merged."""
|
||||
base_ready = _request_is_batchable(base, extra=base_extra)
|
||||
if not base_ready.can_batch:
|
||||
return base_ready
|
||||
candidate_ready = _request_is_batchable(candidate, extra=candidate_extra)
|
||||
if not candidate_ready.can_batch:
|
||||
return candidate_ready
|
||||
|
||||
base_sig = dynamic_batch_signature(base, extra=base_extra)
|
||||
candidate_sig = dynamic_batch_signature(candidate, extra=candidate_extra)
|
||||
if base_sig == candidate_sig:
|
||||
return BatchCompatibility(can_batch=True)
|
||||
|
||||
mismatch = _first_mismatch(base_sig, candidate_sig)
|
||||
return BatchCompatibility(can_batch=False, reason=mismatch or "signature_mismatch")
|
||||
|
||||
|
||||
def _request_is_batchable(
|
||||
request: SamplingParam,
|
||||
*,
|
||||
extra: dict[str, Any] | None = None,
|
||||
) -> BatchCompatibility:
|
||||
if not isinstance(request.prompt, str):
|
||||
return BatchCompatibility(can_batch=False, reason="prompt_type")
|
||||
if request.prompt_path is not None:
|
||||
return BatchCompatibility(can_batch=False, reason="prompt_path")
|
||||
if request.num_videos_per_prompt != 1:
|
||||
return BatchCompatibility(can_batch=False, reason="num_videos_per_prompt")
|
||||
if request.return_continuation_state:
|
||||
return BatchCompatibility(can_batch=False, reason="return_continuation_state")
|
||||
|
||||
for name in _UNSUPPORTED_DYNAMIC_BATCH_FIELDS:
|
||||
value = getattr(request, name, None)
|
||||
if _is_present(value):
|
||||
return BatchCompatibility(can_batch=False, reason=name)
|
||||
|
||||
if extra:
|
||||
unsupported = sorted(set(extra) & _UNSUPPORTED_EXTRA_KEYS)
|
||||
if unsupported:
|
||||
return BatchCompatibility(can_batch=False, reason=f"extra.{unsupported[0]}")
|
||||
|
||||
return BatchCompatibility(can_batch=True)
|
||||
|
||||
|
||||
def _freeze_signature_value(value: Any) -> Any:
|
||||
if isinstance(value, str | int | float | bool | type(None)):
|
||||
return value
|
||||
if isinstance(value, Enum):
|
||||
return value.value
|
||||
if isinstance(value, dict):
|
||||
return tuple(
|
||||
(str(key), _freeze_signature_value(item)) for key, item in sorted(value.items(), key=lambda kv: str(kv[0])))
|
||||
if isinstance(value, list | tuple):
|
||||
return tuple(_freeze_signature_value(item) for item in value)
|
||||
return repr(value)
|
||||
|
||||
|
||||
def _is_present(value: Any) -> bool:
|
||||
if value is None:
|
||||
return False
|
||||
if value is False:
|
||||
return False
|
||||
return not (isinstance(value, list | tuple | dict | set) and not value)
|
||||
|
||||
|
||||
def _first_scalar(value: Any) -> Any:
|
||||
if isinstance(value, list | tuple):
|
||||
return value[0] if value else None
|
||||
return value
|
||||
|
||||
|
||||
def _first_mismatch(
|
||||
base_sig: tuple[tuple[str, Any], ...],
|
||||
candidate_sig: tuple[tuple[str, Any], ...],
|
||||
) -> str | None:
|
||||
if len(base_sig) != len(candidate_sig):
|
||||
return "sampling_params"
|
||||
for (name, base_value), (candidate_name, candidate_value) in zip(base_sig, candidate_sig, strict=True):
|
||||
if name != candidate_name:
|
||||
return "sampling_params"
|
||||
if base_value != candidate_value:
|
||||
return f"sampling_params.{name}"
|
||||
return None
|
||||
@@ -276,6 +276,27 @@ class PipelineConfig:
|
||||
f"Length of text postprocess functions ({len(self.postprocess_text_funcs)}) must be equal to length of text preprocessing functions ({len(self.preprocess_text_funcs)})"
|
||||
)
|
||||
|
||||
def estimate_request_cost(self, request: Any) -> float:
|
||||
"""Estimate relative memory/compute cost for batching admission.
|
||||
|
||||
The default is intentionally simple and model-agnostic: pixel count
|
||||
times frame count. Pipeline subclasses can override this when they have
|
||||
calibrated costs.
|
||||
"""
|
||||
height = getattr(request, "height", None)
|
||||
width = getattr(request, "width", None)
|
||||
num_frames = getattr(request, "num_frames", None)
|
||||
if isinstance(height, list):
|
||||
height = height[0] if height else None
|
||||
if isinstance(width, list):
|
||||
width = width[0] if width else None
|
||||
if isinstance(num_frames, list):
|
||||
num_frames = num_frames[0] if num_frames else None
|
||||
height = int(height or 1)
|
||||
width = int(width or 1)
|
||||
num_frames = int(num_frames or 1)
|
||||
return float(max(1, height) * max(1, width) * max(1, num_frames))
|
||||
|
||||
def dump_to_json(self, file_path: str):
|
||||
output_dict = shallow_asdict(self)
|
||||
del_keys = []
|
||||
|
||||
@@ -10,6 +10,7 @@ from fastapi.middleware.cors import CORSMiddleware
|
||||
|
||||
from fastvideo.api.presets import validate_preset_selection
|
||||
from fastvideo.api.schema import GenerationRequest
|
||||
from fastvideo.entrypoints.openai.batching import VideoBatchScheduler
|
||||
from fastvideo.entrypoints.openai.state import (
|
||||
DEFAULT_OUTPUT_DIR,
|
||||
clear_state,
|
||||
@@ -59,11 +60,29 @@ async def lifespan(app: FastAPI) -> AsyncIterator[None]:
|
||||
generator = VideoGenerator.from_fastvideo_args(args)
|
||||
logger.info("Model loaded successfully.")
|
||||
|
||||
set_state(generator, args, output_dir, default_request=default_request)
|
||||
video_batch_scheduler: VideoBatchScheduler | None = None
|
||||
if args.batching_mode == "dynamic" and args.batching_max_size > 1:
|
||||
video_batch_scheduler = VideoBatchScheduler(generator, args)
|
||||
await video_batch_scheduler.start()
|
||||
logger.info(
|
||||
"Started dynamic video batch scheduler: max_size=%d delay_ms=%.2f",
|
||||
args.batching_max_size,
|
||||
args.batching_delay_ms,
|
||||
)
|
||||
|
||||
set_state(
|
||||
generator,
|
||||
args,
|
||||
output_dir,
|
||||
default_request=default_request,
|
||||
video_batch_scheduler=video_batch_scheduler,
|
||||
)
|
||||
|
||||
yield # server is running
|
||||
|
||||
logger.info("Shutting down — releasing model resources ...")
|
||||
if video_batch_scheduler is not None:
|
||||
await video_batch_scheduler.stop()
|
||||
generator.shutdown()
|
||||
clear_state()
|
||||
logger.info("Shutdown complete.")
|
||||
|
||||
@@ -0,0 +1,183 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import time
|
||||
from collections import deque
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
from fastvideo.api.sampling_param import SamplingParam
|
||||
from fastvideo.batching.signature import can_dynamic_batch
|
||||
from fastvideo.fastvideo_args import FastVideoArgs
|
||||
from fastvideo.logger import init_logger
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class _VideoBatchJob:
|
||||
request_id: str
|
||||
kwargs: dict[str, Any]
|
||||
future: asyncio.Future
|
||||
enqueue_time: float
|
||||
|
||||
|
||||
class VideoBatchScheduler:
|
||||
"""Async FIFO scheduler for OpenAI-compatible video generation."""
|
||||
|
||||
def __init__(self, generator: Any, fastvideo_args: FastVideoArgs) -> None:
|
||||
self._generator = generator
|
||||
self._fastvideo_args = fastvideo_args
|
||||
self._queue: asyncio.Queue[_VideoBatchJob | None] = asyncio.Queue()
|
||||
self._pending: deque[_VideoBatchJob] = deque()
|
||||
self._task: asyncio.Task | None = None
|
||||
self._stopped = False
|
||||
|
||||
@property
|
||||
def enabled(self) -> bool:
|
||||
return self._fastvideo_args.batching_mode == "dynamic" and self._fastvideo_args.batching_max_size > 1
|
||||
|
||||
async def start(self) -> None:
|
||||
if self._task is not None:
|
||||
return
|
||||
self._task = asyncio.create_task(self._run(), name="fastvideo-video-batch-scheduler")
|
||||
|
||||
async def stop(self) -> None:
|
||||
self._stopped = True
|
||||
await self._queue.put(None)
|
||||
if self._task is not None:
|
||||
await self._task
|
||||
self._task = None
|
||||
|
||||
async def submit(self, request_id: str, kwargs: dict[str, Any]) -> Any:
|
||||
if self._stopped:
|
||||
raise RuntimeError("Video batch scheduler is stopped; cannot submit new requests")
|
||||
loop = asyncio.get_running_loop()
|
||||
future = loop.create_future()
|
||||
await self._queue.put(
|
||||
_VideoBatchJob(
|
||||
request_id=request_id,
|
||||
kwargs=dict(kwargs),
|
||||
future=future,
|
||||
enqueue_time=time.perf_counter(),
|
||||
))
|
||||
return await future
|
||||
|
||||
async def _run(self) -> None:
|
||||
while not self._stopped:
|
||||
job = await self._get_next_job()
|
||||
if job is None:
|
||||
break
|
||||
batch = await self._collect_batch(job)
|
||||
await self._dispatch(batch)
|
||||
|
||||
while self._pending:
|
||||
pending = self._pending.popleft()
|
||||
if not pending.future.done():
|
||||
pending.future.set_exception(RuntimeError("Video batch scheduler stopped before dispatch"))
|
||||
|
||||
async def _get_next_job(self) -> _VideoBatchJob | None:
|
||||
if self._pending:
|
||||
return self._pending.popleft()
|
||||
return await self._queue.get()
|
||||
|
||||
async def _collect_batch(self, first: _VideoBatchJob) -> list[_VideoBatchJob]:
|
||||
batch = [first]
|
||||
max_size = self._fastvideo_args.batching_max_size
|
||||
delay_s = max(0.0, self._fastvideo_args.batching_delay_ms / 1000.0)
|
||||
deadline = first.enqueue_time + delay_s
|
||||
|
||||
while len(batch) < max_size:
|
||||
if delay_s > 0:
|
||||
timeout = deadline - time.perf_counter()
|
||||
if timeout <= 0:
|
||||
break
|
||||
try:
|
||||
candidate = await asyncio.wait_for(self._get_next_job(), timeout=timeout)
|
||||
except TimeoutError:
|
||||
break
|
||||
else:
|
||||
# delay=0 means "don't wait": greedily drain whatever is
|
||||
# already queued so max_size still coalesces.
|
||||
if self._pending:
|
||||
candidate = self._pending.popleft()
|
||||
else:
|
||||
try:
|
||||
candidate = self._queue.get_nowait()
|
||||
except asyncio.QueueEmpty:
|
||||
break
|
||||
if candidate is None:
|
||||
await self._queue.put(None)
|
||||
break
|
||||
if self._jobs_are_compatible(batch[0], candidate):
|
||||
batch.append(candidate)
|
||||
continue
|
||||
self._pending.appendleft(candidate)
|
||||
break
|
||||
return batch
|
||||
|
||||
async def _dispatch(self, batch: list[_VideoBatchJob]) -> None:
|
||||
loop = asyncio.get_running_loop()
|
||||
request_ids = [job.request_id for job in batch]
|
||||
queue_wait_ms = (time.perf_counter() - min(job.enqueue_time for job in batch)) * 1000.0
|
||||
if self._fastvideo_args.enable_batching_metrics:
|
||||
logger.info(
|
||||
"Dispatching video batch: request_ids=%s size=%d queue_wait_ms=%.2f",
|
||||
request_ids,
|
||||
len(batch),
|
||||
queue_wait_ms,
|
||||
)
|
||||
|
||||
try:
|
||||
results = await loop.run_in_executor(
|
||||
None,
|
||||
lambda: self._generator.generate_video_batch([job.kwargs for job in batch]),
|
||||
)
|
||||
except Exception as exc:
|
||||
for job in batch:
|
||||
if not job.future.done():
|
||||
job.future.set_exception(exc)
|
||||
return
|
||||
|
||||
if len(results) != len(batch):
|
||||
error = RuntimeError(f"Video batch returned {len(results)} results for {len(batch)} requests")
|
||||
for job in batch:
|
||||
if not job.future.done():
|
||||
job.future.set_exception(error)
|
||||
return
|
||||
|
||||
for job, result in zip(batch, results, strict=True):
|
||||
if not job.future.done():
|
||||
job.future.set_result(result)
|
||||
|
||||
def _jobs_are_compatible(self, base: _VideoBatchJob, candidate: _VideoBatchJob) -> bool:
|
||||
try:
|
||||
base_sampling, base_extra = self._sampling_param_from_kwargs(base.kwargs)
|
||||
candidate_sampling, candidate_extra = self._sampling_param_from_kwargs(candidate.kwargs)
|
||||
except Exception:
|
||||
return False
|
||||
return can_dynamic_batch(
|
||||
base_sampling,
|
||||
candidate_sampling,
|
||||
base_extra=base_extra,
|
||||
candidate_extra=candidate_extra,
|
||||
).can_batch
|
||||
|
||||
def _sampling_param_from_kwargs(self, kwargs: dict[str, Any]) -> tuple[SamplingParam, dict[str, Any]]:
|
||||
sampling_param = SamplingParam.from_pretrained(self._fastvideo_args.model_path)
|
||||
updates = dict(kwargs)
|
||||
prompt = updates.pop("prompt", None)
|
||||
extra: dict[str, Any] = {}
|
||||
for key in (
|
||||
"ltx2_audio_latents",
|
||||
"ltx2_audio_clean_latent",
|
||||
"ltx2_audio_denoise_mask",
|
||||
"audio_num_frames",
|
||||
"video_position_offset_sec",
|
||||
):
|
||||
if key in updates:
|
||||
extra[key] = updates.pop(key)
|
||||
sampling_param.update(updates)
|
||||
sampling_param.prompt = prompt
|
||||
return sampling_param, extra
|
||||
@@ -11,6 +11,7 @@ from typing import TYPE_CHECKING
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from fastvideo.api.schema import GenerationRequest
|
||||
from fastvideo.entrypoints.openai.batching import VideoBatchScheduler
|
||||
from fastvideo.entrypoints.video_generator import VideoGenerator
|
||||
from fastvideo.fastvideo_args import FastVideoArgs
|
||||
|
||||
@@ -20,6 +21,7 @@ _generator: VideoGenerator | None = None
|
||||
_fastvideo_args: FastVideoArgs | None = None
|
||||
_output_dir: str = DEFAULT_OUTPUT_DIR
|
||||
_default_request: GenerationRequest | None = None
|
||||
_video_batch_scheduler: VideoBatchScheduler | None = None
|
||||
|
||||
|
||||
def get_generator() -> VideoGenerator:
|
||||
@@ -44,23 +46,31 @@ def get_default_request() -> GenerationRequest | None:
|
||||
return _default_request
|
||||
|
||||
|
||||
def get_video_batch_scheduler() -> VideoBatchScheduler | None:
|
||||
"""Return the video batch scheduler when dynamic batching is enabled."""
|
||||
return _video_batch_scheduler
|
||||
|
||||
|
||||
def set_state(
|
||||
generator: VideoGenerator,
|
||||
fastvideo_args: FastVideoArgs,
|
||||
output_dir: str,
|
||||
default_request: GenerationRequest | None = None,
|
||||
video_batch_scheduler: VideoBatchScheduler | None = None,
|
||||
) -> None:
|
||||
"""Set all server state at once (called from lifespan)."""
|
||||
global _generator, _fastvideo_args, _output_dir, _default_request
|
||||
global _generator, _fastvideo_args, _output_dir, _default_request, _video_batch_scheduler
|
||||
_generator = generator
|
||||
_fastvideo_args = fastvideo_args
|
||||
_output_dir = output_dir
|
||||
_default_request = default_request
|
||||
_video_batch_scheduler = video_batch_scheduler
|
||||
|
||||
|
||||
def clear_state() -> None:
|
||||
"""Clear server state on shutdown."""
|
||||
global _generator, _fastvideo_args, _default_request
|
||||
global _generator, _fastvideo_args, _default_request, _video_batch_scheduler
|
||||
_generator = None
|
||||
_fastvideo_args = None
|
||||
_default_request = None
|
||||
_video_batch_scheduler = None
|
||||
|
||||
@@ -26,6 +26,7 @@ from fastvideo.entrypoints.openai.state import (
|
||||
get_generator,
|
||||
get_output_dir,
|
||||
get_server_args,
|
||||
get_video_batch_scheduler,
|
||||
)
|
||||
from fastvideo.entrypoints.openai.protocol import (
|
||||
VideoGenerationsRequest,
|
||||
@@ -151,15 +152,19 @@ async def _run_generation(request_id: str, kwargs: dict[str, Any]) -> None:
|
||||
is synchronous) and update the store on completion or failure.
|
||||
"""
|
||||
generator = get_generator()
|
||||
scheduler = get_video_batch_scheduler()
|
||||
loop = asyncio.get_running_loop()
|
||||
|
||||
try:
|
||||
start = time.perf_counter()
|
||||
|
||||
result = await loop.run_in_executor(
|
||||
None,
|
||||
lambda: generator.generate_video(**kwargs),
|
||||
)
|
||||
if scheduler is not None and scheduler.enabled:
|
||||
result = await scheduler.submit(request_id, kwargs)
|
||||
else:
|
||||
result = await loop.run_in_executor(
|
||||
None,
|
||||
lambda: generator.generate_video(**kwargs),
|
||||
)
|
||||
|
||||
elapsed = time.perf_counter() - start
|
||||
update: dict[str, Any] = {
|
||||
|
||||
@@ -18,6 +18,7 @@ import warnings
|
||||
from collections.abc import Mapping
|
||||
from contextlib import suppress
|
||||
from copy import deepcopy
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
import imageio
|
||||
@@ -26,6 +27,8 @@ import torch
|
||||
import torchvision
|
||||
from einops import rearrange
|
||||
|
||||
from fastvideo.batching.admission import BatchAdmissionController
|
||||
from fastvideo.batching.signature import can_dynamic_batch
|
||||
from fastvideo.api.compat import (
|
||||
expand_request_prompt_batch,
|
||||
generator_config_to_fastvideo_args,
|
||||
@@ -94,6 +97,11 @@ _FROM_PRETRAINED_CONVENIENCE_KWARGS = frozenset({
|
||||
"pin_cpu_memory",
|
||||
"enable_torch_compile",
|
||||
"torch_compile_kwargs",
|
||||
"batching_mode",
|
||||
"batching_max_size",
|
||||
"batching_delay_ms",
|
||||
"batching_config",
|
||||
"enable_batching_metrics",
|
||||
"output_type",
|
||||
"nvfp4_fa4",
|
||||
})
|
||||
@@ -112,6 +120,17 @@ def _infer_latent_batch_size(batch: ForwardBatch) -> int:
|
||||
return latent_batch_size
|
||||
|
||||
|
||||
@dataclass
|
||||
class _GenerationWorkItem:
|
||||
prompt: str
|
||||
sampling_param: SamplingParam
|
||||
fastvideo_args: FastVideoArgs
|
||||
batch: ForwardBatch
|
||||
output_path: str
|
||||
target_height: int
|
||||
target_width: int
|
||||
|
||||
|
||||
class VideoGenerator:
|
||||
"""
|
||||
A unified class for generating videos using diffusion models.
|
||||
@@ -439,6 +458,70 @@ class VideoGenerator:
|
||||
if log_queue:
|
||||
self.executor.clear_log_queue()
|
||||
|
||||
def generate_video_batch(self, request_kwargs: list[dict[str, Any]]) -> list[dict[str, Any]]:
|
||||
"""Generate multiple legacy video requests, batching compatible items."""
|
||||
work_items: list[_GenerationWorkItem] = []
|
||||
reserved_output_paths: set[str] = set()
|
||||
fastvideo_args_by_pipeline_override: dict[tuple[tuple[str, str], ...], FastVideoArgs] = {
|
||||
(): self.fastvideo_args
|
||||
}
|
||||
for raw_kwargs in request_kwargs:
|
||||
kwargs = dict(raw_kwargs)
|
||||
prompt = kwargs.pop("prompt", None)
|
||||
if prompt is None:
|
||||
raise ValueError("Each batched generation request must include prompt")
|
||||
if not isinstance(prompt, str):
|
||||
raise TypeError(f"`prompt` must be a string, but got {type(prompt)}")
|
||||
|
||||
sampling_param = kwargs.pop("sampling_param", None)
|
||||
if sampling_param is None:
|
||||
sampling_param = SamplingParam.from_pretrained(self.fastvideo_args.model_path)
|
||||
else:
|
||||
sampling_param = deepcopy(sampling_param)
|
||||
|
||||
extra_overrides: dict[str, Any] = {}
|
||||
for _ek in _BATCH_EXTRA_PASSTHROUGH_KEYS:
|
||||
if _ek in kwargs:
|
||||
extra_overrides[_ek] = kwargs.pop(_ek)
|
||||
|
||||
request = legacy_generate_call_to_request(
|
||||
prompt,
|
||||
sampling_param,
|
||||
legacy_kwargs=kwargs,
|
||||
)
|
||||
if not isinstance(request.prompt, str):
|
||||
raise TypeError(f"`prompt` must be a string, but got {type(request.prompt)}")
|
||||
|
||||
fastvideo_args = self.fastvideo_args
|
||||
pipeline_overrides = request_to_pipeline_overrides(request)
|
||||
if pipeline_overrides:
|
||||
override_key = tuple((key, repr(value)) for key, value in sorted(pipeline_overrides.items()))
|
||||
fastvideo_args = fastvideo_args_by_pipeline_override.get(override_key)
|
||||
if fastvideo_args is None:
|
||||
fastvideo_args = deepcopy(self.fastvideo_args)
|
||||
for key, value in pipeline_overrides.items():
|
||||
if not hasattr(fastvideo_args.pipeline_config, key):
|
||||
raise ValueError(f"Request field {key!r} is not supported by pipeline config overrides")
|
||||
setattr(fastvideo_args.pipeline_config, key, deepcopy(value))
|
||||
fastvideo_args_by_pipeline_override[override_key] = fastvideo_args
|
||||
|
||||
resolved_sampling_param = request_to_sampling_param(
|
||||
request,
|
||||
model_path=self.fastvideo_args.model_path,
|
||||
)
|
||||
output_path = self._prepare_output_path(resolved_sampling_param.output_path, request.prompt,
|
||||
reserved_output_paths)
|
||||
work_items.append(
|
||||
self._prepare_generation_work_item(
|
||||
prompt=request.prompt,
|
||||
sampling_param=resolved_sampling_param,
|
||||
fastvideo_args=fastvideo_args,
|
||||
output_path=output_path,
|
||||
_extra_overrides=extra_overrides,
|
||||
))
|
||||
|
||||
return self._generate_prepared_work_items(work_items)
|
||||
|
||||
def _generate_request_impl(
|
||||
self,
|
||||
request: GenerationRequest,
|
||||
@@ -535,6 +618,29 @@ class VideoGenerator:
|
||||
|
||||
logger.info("Found %d prompts in %s", len(prompts), prompt_txt_path)
|
||||
|
||||
if self._dynamic_batching_enabled(fastvideo_args):
|
||||
work_items: list[_GenerationWorkItem] = []
|
||||
reserved_output_paths: set[str] = set()
|
||||
for batch_prompt in prompts:
|
||||
item_kwargs = dict(kwargs)
|
||||
item_kwargs["output_path"] = self._prepare_output_path(sampling_param.output_path, batch_prompt,
|
||||
reserved_output_paths)
|
||||
work_items.append(
|
||||
self._prepare_generation_work_item(
|
||||
prompt=batch_prompt,
|
||||
sampling_param=sampling_param,
|
||||
fastvideo_args=fastvideo_args,
|
||||
**item_kwargs,
|
||||
))
|
||||
|
||||
results = self._generate_prepared_work_items(work_items, tolerate_failures=True)
|
||||
for i, (result, batch_prompt) in enumerate(zip(results, prompts, strict=True)):
|
||||
result["prompt_index"] = i
|
||||
result["prompt"] = batch_prompt
|
||||
logger.info("Completed batch processing. Generated %d videos successfully.",
|
||||
sum(1 for result in results if "error" not in result))
|
||||
return results
|
||||
|
||||
results = []
|
||||
for i, batch_prompt in enumerate(prompts):
|
||||
logger.info("Processing prompt %d/%d: %s...", i + 1, len(prompts), batch_prompt[:100])
|
||||
@@ -588,6 +694,7 @@ class VideoGenerator:
|
||||
self,
|
||||
output_path: str,
|
||||
prompt: str,
|
||||
reserved_paths: set[str] | None = None,
|
||||
) -> str:
|
||||
"""Build a unique, sanitized output file path.
|
||||
|
||||
@@ -602,6 +709,9 @@ class VideoGenerator:
|
||||
- Invalid filename characters are removed; if the name changes, a
|
||||
warning is logged.
|
||||
- If the target path already exists, a numeric suffix is appended.
|
||||
- ``reserved_paths`` lets batch callers resolve every path before any
|
||||
file is written: paths in the set are treated as taken, and the
|
||||
chosen path is added to the set.
|
||||
"""
|
||||
target_ext = ".png" if self._is_image_workload() else ".mp4"
|
||||
|
||||
@@ -646,15 +756,416 @@ class VideoGenerator:
|
||||
if output_dir:
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
|
||||
def _is_taken(path: str) -> bool:
|
||||
return os.path.exists(path) or (reserved_paths is not None and path in reserved_paths)
|
||||
|
||||
new_output_path = os.path.join(output_dir, out_name)
|
||||
counter = 1
|
||||
while os.path.exists(new_output_path):
|
||||
while _is_taken(new_output_path):
|
||||
name_part, ext_part = os.path.splitext(out_name)
|
||||
new_name = f"{name_part}_{counter}{ext_part}"
|
||||
new_output_path = os.path.join(output_dir, new_name)
|
||||
counter += 1
|
||||
if reserved_paths is not None:
|
||||
reserved_paths.add(new_output_path)
|
||||
return new_output_path
|
||||
|
||||
def _dynamic_batching_enabled(self, fastvideo_args: FastVideoArgs) -> bool:
|
||||
batching_mode = getattr(fastvideo_args, "batching_mode", "disabled")
|
||||
batching_max_size = getattr(fastvideo_args, "batching_max_size", 1)
|
||||
return batching_mode == "dynamic" and batching_max_size > 1
|
||||
|
||||
def _prepare_generation_work_item(
|
||||
self,
|
||||
prompt: str | list[str],
|
||||
sampling_param: SamplingParam,
|
||||
fastvideo_args: FastVideoArgs,
|
||||
**kwargs,
|
||||
) -> _GenerationWorkItem:
|
||||
if isinstance(prompt, str):
|
||||
prompt_for_output = prompt.strip()
|
||||
prompt_value: str | list[str] = prompt_for_output
|
||||
elif isinstance(prompt, list) and all(isinstance(item, str) for item in prompt):
|
||||
prompt_value = [item.strip() for item in prompt]
|
||||
prompt_for_output = prompt_value[0] if prompt_value else ""
|
||||
else:
|
||||
raise TypeError(f"`prompt` must be a string or list of strings, but got {type(prompt)}")
|
||||
|
||||
sampling_param = deepcopy(sampling_param)
|
||||
output_path = kwargs["output_path"]
|
||||
sampling_param.prompt = prompt_value
|
||||
if sampling_param.negative_prompt is not None:
|
||||
sampling_param.negative_prompt = sampling_param.negative_prompt.strip()
|
||||
|
||||
if sampling_param.height <= 0 or sampling_param.width <= 0 or sampling_param.num_frames <= 0:
|
||||
raise ValueError(f"Height, width, and num_frames must be positive integers, got "
|
||||
f"height={sampling_param.height}, width={sampling_param.width}, "
|
||||
f"num_frames={sampling_param.num_frames}")
|
||||
|
||||
target_height = align_to(sampling_param.height, 16)
|
||||
target_width = align_to(sampling_param.width, 16)
|
||||
latents_size = [(sampling_param.num_frames - 1) // 4 + 1, sampling_param.height // 8, sampling_param.width // 8]
|
||||
n_tokens = latents_size[0] * latents_size[1] * latents_size[2]
|
||||
|
||||
debug_str = f"""
|
||||
height: {target_height}
|
||||
width: {target_width}
|
||||
video_length: {sampling_param.num_frames}
|
||||
prompt: {sampling_param.prompt}
|
||||
image_path: {sampling_param.image_path}
|
||||
neg_prompt: {sampling_param.negative_prompt}
|
||||
seed: {sampling_param.seed}
|
||||
infer_steps: {sampling_param.num_inference_steps}
|
||||
num_videos_per_prompt: {sampling_param.num_videos_per_prompt}
|
||||
guidance_scale: {sampling_param.guidance_scale}
|
||||
n_tokens: {n_tokens}
|
||||
flow_shift: {fastvideo_args.pipeline_config.flow_shift}
|
||||
embedded_guidance_scale: {fastvideo_args.pipeline_config.embedded_cfg_scale}
|
||||
save_video: {sampling_param.save_video}
|
||||
output_path: {output_path}
|
||||
""" # type: ignore[attr-defined]
|
||||
logger.info(debug_str)
|
||||
|
||||
batch = ForwardBatch(
|
||||
**shallow_asdict(sampling_param),
|
||||
eta=0.0,
|
||||
n_tokens=n_tokens,
|
||||
VSA_sparsity=fastvideo_args.VSA_sparsity,
|
||||
)
|
||||
|
||||
extra_overrides = kwargs.get("_extra_overrides", {})
|
||||
for _ek, _ev in extra_overrides.items():
|
||||
batch.extra[_ek] = _ev
|
||||
|
||||
return _GenerationWorkItem(
|
||||
prompt=prompt_for_output,
|
||||
sampling_param=sampling_param,
|
||||
fastvideo_args=fastvideo_args,
|
||||
batch=batch,
|
||||
output_path=output_path,
|
||||
target_height=target_height,
|
||||
target_width=target_width,
|
||||
)
|
||||
|
||||
def _run_forward_batch(
|
||||
self,
|
||||
batch: ForwardBatch,
|
||||
fastvideo_args: FastVideoArgs,
|
||||
) -> tuple[ForwardBatch, float, float]:
|
||||
start_time = time.perf_counter()
|
||||
result_container = {"output_batch": ForwardBatch(data_type=batch.data_type)}
|
||||
thread_error: dict[str, BaseException | None] = {"error": None}
|
||||
thread_error_traceback: dict[str, str] = {"traceback": ""}
|
||||
|
||||
def execute_forward_thread():
|
||||
import traceback
|
||||
try:
|
||||
result_container["output_batch"] = self.executor.execute_forward(batch, fastvideo_args)
|
||||
except BaseException as error: # noqa: BLE001
|
||||
thread_error["error"] = error
|
||||
thread_error_traceback["traceback"] = traceback.format_exc()
|
||||
|
||||
thread = threading.Thread(target=execute_forward_thread)
|
||||
thread.start()
|
||||
thread.join()
|
||||
|
||||
if thread_error["error"] is not None:
|
||||
raise RuntimeError("Forward execution thread failed.\n"
|
||||
f"{thread_error_traceback['traceback']}") from thread_error["error"]
|
||||
|
||||
output_batch = result_container["output_batch"]
|
||||
if output_batch.output is None:
|
||||
raise RuntimeError("Forward execution returned no output tensor. "
|
||||
"This usually means the executor/pipeline failed earlier.")
|
||||
|
||||
gen_time = time.perf_counter() - start_time
|
||||
logger.info("Generated successfully in %.2f seconds", gen_time)
|
||||
return output_batch, gen_time, start_time
|
||||
|
||||
def _samples_from_output(
|
||||
self,
|
||||
work_item: _GenerationWorkItem,
|
||||
output_batch: ForwardBatch,
|
||||
) -> torch.Tensor:
|
||||
output = output_batch.output
|
||||
if output is None:
|
||||
raise RuntimeError("Forward execution returned no output tensor.")
|
||||
fastvideo_args = work_item.fastvideo_args
|
||||
sampling_param = work_item.sampling_param
|
||||
latent_batch_size = _infer_latent_batch_size(work_item.batch)
|
||||
skip_pixel_prealloc = fastvideo_args.output_type == "latent"
|
||||
expected_shape = (
|
||||
latent_batch_size,
|
||||
3,
|
||||
sampling_param.num_frames,
|
||||
sampling_param.height,
|
||||
sampling_param.width,
|
||||
)
|
||||
if skip_pixel_prealloc:
|
||||
return output.cpu()
|
||||
samples = torch.empty(expected_shape, device="cpu", pin_memory=fastvideo_args.pin_cpu_memory)
|
||||
if output.shape == samples.shape:
|
||||
samples.copy_(output)
|
||||
return samples
|
||||
logger.warning("Output shape %s does not match expected shape %s; use slow path", output.shape, samples.shape)
|
||||
return output.cpu()
|
||||
|
||||
def _postprocess_generation_output(
|
||||
self,
|
||||
work_item: _GenerationWorkItem,
|
||||
output_batch: ForwardBatch,
|
||||
gen_time: float,
|
||||
start_time: float,
|
||||
) -> dict[str, Any]:
|
||||
batch = work_item.batch
|
||||
fastvideo_args = work_item.fastvideo_args
|
||||
output_path = work_item.output_path
|
||||
samples = self._samples_from_output(work_item, output_batch)
|
||||
logging_info = output_batch.logging_info
|
||||
|
||||
is_latent_output = fastvideo_args.output_type == "latent"
|
||||
audio_only = bool(output_batch.extra.get("audio_only"))
|
||||
|
||||
postprocess_start = time.perf_counter()
|
||||
frames: list[np.ndarray] | None
|
||||
if is_latent_output or audio_only:
|
||||
frames = None if is_latent_output else []
|
||||
else:
|
||||
videos = rearrange(samples, "b c t h w -> t b c h w")
|
||||
frames = []
|
||||
for x in videos:
|
||||
x = torchvision.utils.make_grid(x, nrow=6)
|
||||
x = x.permute(1, 2, 0).squeeze(-1)
|
||||
x = (x * 255).to(torch.uint8)
|
||||
frames.append(x.contiguous().cpu().numpy())
|
||||
postprocess_time = time.perf_counter() - postprocess_start
|
||||
logger.info("PostDecodeFrameProcessStage completed in %.3f s", postprocess_time)
|
||||
if logging_info is not None:
|
||||
logging_info.add_stage_execution_time("PostDecodeFrameProcessStage", postprocess_time)
|
||||
|
||||
save_to_disk = batch.save_video and not is_latent_output
|
||||
save_video_time = 0.0
|
||||
audio_mux_time = 0.0
|
||||
if save_to_disk:
|
||||
if audio_only:
|
||||
output_path = self._rewrite_extension(output_path, ".wav")
|
||||
save_start = time.perf_counter()
|
||||
self._write_pcm_wav(
|
||||
output_path,
|
||||
output_batch.extra["audio"],
|
||||
int(output_batch.extra["audio_sample_rate"]),
|
||||
)
|
||||
save_video_time = time.perf_counter() - save_start
|
||||
logger.info("Saved audio to %s", output_path)
|
||||
elif self._is_image_workload():
|
||||
assert frames is not None
|
||||
save_start = time.perf_counter()
|
||||
imageio.imwrite(output_path, frames[0])
|
||||
save_video_time = time.perf_counter() - save_start
|
||||
logger.info("Saved image to %s", output_path)
|
||||
else:
|
||||
assert frames is not None
|
||||
audio = output_batch.extra.get("audio")
|
||||
audio_sample_rate = output_batch.extra.get("audio_sample_rate")
|
||||
if audio is not None and audio_sample_rate is not None:
|
||||
save_start = time.perf_counter()
|
||||
save_ok = self._save_video_with_audio_ffmpeg_pipe(
|
||||
output_path=output_path,
|
||||
frames=frames,
|
||||
fps=batch.fps,
|
||||
audio=audio,
|
||||
sample_rate=int(audio_sample_rate),
|
||||
)
|
||||
if not save_ok:
|
||||
logger.warning("ffmpeg pipe save failed; trying PyAV single-pass save.")
|
||||
save_ok = self._save_video_with_audio_single_pass(
|
||||
output_path=output_path,
|
||||
frames=frames,
|
||||
fps=batch.fps,
|
||||
audio=audio,
|
||||
sample_rate=int(audio_sample_rate),
|
||||
)
|
||||
save_video_time = time.perf_counter() - save_start
|
||||
if save_ok:
|
||||
audio_mux_time = 0.0
|
||||
else:
|
||||
logger.warning("Single-pass save failed; falling back to two-step save/mux.")
|
||||
save_start = time.perf_counter()
|
||||
imageio.mimsave(output_path, frames, fps=batch.fps, format="mp4")
|
||||
save_video_time = time.perf_counter() - save_start
|
||||
mux_start = time.perf_counter()
|
||||
mux_ok = self._mux_audio(output_path, audio, int(audio_sample_rate))
|
||||
audio_mux_time = time.perf_counter() - mux_start
|
||||
if not mux_ok:
|
||||
logger.warning("Audio mux failed; saved video without audio.")
|
||||
else:
|
||||
save_start = time.perf_counter()
|
||||
imageio.mimsave(output_path, frames, fps=batch.fps, format="mp4")
|
||||
save_video_time = time.perf_counter() - save_start
|
||||
audio_mux_time = 0.0
|
||||
logger.info("Saved video to %s", output_path)
|
||||
|
||||
logger.info("VideoSaveStage completed in %.3f s", save_video_time)
|
||||
if logging_info is not None:
|
||||
logging_info.add_stage_execution_time("VideoSaveStage", save_video_time)
|
||||
logger.info("AudioMuxStage completed in %.3f s", audio_mux_time)
|
||||
if logging_info is not None:
|
||||
logging_info.add_stage_execution_time("AudioMuxStage", audio_mux_time)
|
||||
|
||||
e2e_time = time.perf_counter() - start_time
|
||||
logger.info("End-to-end latency: %.2f seconds", e2e_time)
|
||||
|
||||
return {
|
||||
"prompts": work_item.prompt,
|
||||
"samples": samples if batch.return_frames else None,
|
||||
"frames": frames if batch.return_frames else None,
|
||||
"audio": output_batch.extra.get("audio"),
|
||||
"audio_sample_rate": output_batch.extra.get("audio_sample_rate"),
|
||||
"ltx2_audio_latents": output_batch.extra.get("ltx2_audio_latents"),
|
||||
"size": (work_item.target_height, work_item.target_width, batch.num_frames),
|
||||
"generation_time": gen_time,
|
||||
"e2e_latency": e2e_time,
|
||||
"logging_info": logging_info,
|
||||
"trajectory": output_batch.trajectory_latents,
|
||||
"trajectory_timesteps": output_batch.trajectory_timesteps,
|
||||
"trajectory_decoded": output_batch.trajectory_decoded,
|
||||
"video_path": output_path if save_to_disk else None,
|
||||
"peak_memory_mb": output_batch.extra.get("peak_memory_mb"),
|
||||
}
|
||||
|
||||
def _split_output_batch(
|
||||
self,
|
||||
output_batch: ForwardBatch,
|
||||
*,
|
||||
index: int,
|
||||
batch_size: int,
|
||||
) -> ForwardBatch:
|
||||
extra = {}
|
||||
for key, value in (output_batch.extra or {}).items():
|
||||
if torch.is_tensor(value) and value.ndim > 0 and value.shape[0] == batch_size:
|
||||
extra[key] = value[index:index + 1]
|
||||
elif isinstance(value, list) and len(value) == batch_size:
|
||||
extra[key] = value[index]
|
||||
else:
|
||||
extra[key] = value
|
||||
|
||||
result = ForwardBatch(
|
||||
data_type=output_batch.data_type,
|
||||
output=(output_batch.output[index:index + 1] if output_batch.output is not None else None),
|
||||
logging_info=output_batch.logging_info,
|
||||
extra=extra,
|
||||
)
|
||||
if output_batch.trajectory_latents is not None:
|
||||
result.trajectory_latents = output_batch.trajectory_latents[index:index + 1]
|
||||
result.trajectory_timesteps = output_batch.trajectory_timesteps
|
||||
if output_batch.trajectory_decoded is not None:
|
||||
result.trajectory_decoded = [
|
||||
decoded[index:index + 1] if torch.is_tensor(decoded) and decoded.shape[0] == batch_size else decoded
|
||||
for decoded in output_batch.trajectory_decoded
|
||||
]
|
||||
return result
|
||||
|
||||
def _merge_work_items(self, work_items: list[_GenerationWorkItem]) -> _GenerationWorkItem:
|
||||
first = work_items[0]
|
||||
sampling_param = deepcopy(first.sampling_param)
|
||||
prompts = [item.prompt for item in work_items]
|
||||
sampling_param.prompt = prompts
|
||||
sampling_param.seed = work_items[0].sampling_param.seed
|
||||
|
||||
merged = self._prepare_generation_work_item(
|
||||
prompts,
|
||||
sampling_param,
|
||||
first.fastvideo_args,
|
||||
output_path=first.output_path,
|
||||
_extra_overrides=first.batch.extra,
|
||||
)
|
||||
merged.batch.seeds = [int(item.sampling_param.seed) for item in work_items]
|
||||
merged.batch.extra["dynamic_batch_size"] = len(work_items)
|
||||
merged.batch.extra["dynamic_batch_output_paths"] = [item.output_path for item in work_items]
|
||||
return merged
|
||||
|
||||
def _can_merge_work_items(
|
||||
self,
|
||||
base: _GenerationWorkItem,
|
||||
candidate: _GenerationWorkItem,
|
||||
admission: BatchAdmissionController,
|
||||
current_group: list[_GenerationWorkItem],
|
||||
) -> bool:
|
||||
if candidate.fastvideo_args is not base.fastvideo_args:
|
||||
return False
|
||||
compatibility = can_dynamic_batch(
|
||||
base.sampling_param,
|
||||
candidate.sampling_param,
|
||||
base_extra=base.batch.extra,
|
||||
candidate_extra=candidate.batch.extra,
|
||||
)
|
||||
if not compatibility.can_batch:
|
||||
return False
|
||||
current_requests = [item.sampling_param for item in current_group]
|
||||
return admission.reject_reason_for_candidate(current_requests, candidate.sampling_param) is None
|
||||
|
||||
def _run_work_item_group(self, group: list[_GenerationWorkItem]) -> list[dict[str, Any]]:
|
||||
if len(group) == 1:
|
||||
return [self._execute_single_work_item(group[0])]
|
||||
merged = self._merge_work_items(group)
|
||||
output_batch, gen_time, start_time = self._run_forward_batch(merged.batch, merged.fastvideo_args)
|
||||
return [
|
||||
self._postprocess_generation_output(
|
||||
item,
|
||||
self._split_output_batch(output_batch, index=item_index, batch_size=len(group)),
|
||||
gen_time,
|
||||
start_time,
|
||||
) for item_index, item in enumerate(group)
|
||||
]
|
||||
|
||||
def _generate_prepared_work_items(
|
||||
self,
|
||||
work_items: list[_GenerationWorkItem],
|
||||
tolerate_failures: bool = False,
|
||||
) -> list[dict[str, Any]]:
|
||||
"""Execute prepared work items, batching compatible neighbors.
|
||||
|
||||
With ``tolerate_failures`` (prompt-file semantics), a failed group
|
||||
yields one ``{"error": ..., "prompt": ...}`` entry per work item so
|
||||
completed results survive and stay aligned with the inputs; otherwise
|
||||
the exception propagates.
|
||||
"""
|
||||
if not work_items:
|
||||
return []
|
||||
|
||||
def run_group(group: list[_GenerationWorkItem]) -> list[dict[str, Any]]:
|
||||
if not tolerate_failures:
|
||||
return self._run_work_item_group(group)
|
||||
try:
|
||||
return self._run_work_item_group(group)
|
||||
except Exception as e:
|
||||
logger.error("Failed to generate videos for batched prompts %s: %s",
|
||||
[item.prompt[:100] for item in group], e)
|
||||
return [{"error": str(e), "prompt": item.prompt} for item in group]
|
||||
|
||||
fastvideo_args = work_items[0].fastvideo_args
|
||||
if not self._dynamic_batching_enabled(fastvideo_args):
|
||||
return [result for item in work_items for result in run_group([item])]
|
||||
|
||||
admission = BatchAdmissionController(fastvideo_args)
|
||||
results: list[dict[str, Any]] = []
|
||||
index = 0
|
||||
while index < len(work_items):
|
||||
group = [work_items[index]]
|
||||
index += 1
|
||||
while index < len(work_items) and len(group) < fastvideo_args.batching_max_size:
|
||||
candidate = work_items[index]
|
||||
if not self._can_merge_work_items(group[0], candidate, admission, group):
|
||||
break
|
||||
group.append(candidate)
|
||||
index += 1
|
||||
results.extend(run_group(group))
|
||||
return results
|
||||
|
||||
def _execute_single_work_item(self, work_item: _GenerationWorkItem) -> dict[str, Any]:
|
||||
output_batch, gen_time, start_time = self._run_forward_batch(work_item.batch, work_item.fastvideo_args)
|
||||
return self._postprocess_generation_output(work_item, output_batch, gen_time, start_time)
|
||||
|
||||
def _generate_single_video(
|
||||
self,
|
||||
prompt: str,
|
||||
|
||||
@@ -20,6 +20,7 @@ if TYPE_CHECKING:
|
||||
FASTVIDEO_LOGGING_CONFIG_PATH: str | None = None
|
||||
FASTVIDEO_TRACE_FUNCTION: int = 0
|
||||
FASTVIDEO_ATTENTION_BACKEND: str | None = None
|
||||
FASTVIDEO_FA4: bool = False
|
||||
FASTVIDEO_WORKER_MULTIPROC_METHOD: str = "spawn"
|
||||
FASTVIDEO_TARGET_DEVICE: str = "cuda"
|
||||
MAX_JOBS: str | None = None
|
||||
@@ -207,9 +208,18 @@ environment_variables: dict[str, Callable[[], Any]] = {
|
||||
# - "VIDEO_SPARSE_ATTN": use Video Sparse Attention
|
||||
# - "SAGE_ATTN": use Sage Attention
|
||||
# - "SAGE_ATTN_THREE": use Sage Attention 3
|
||||
# FLASH_ATTN uses FlashAttention-3/2; to run FlashAttention-4 set
|
||||
# FASTVIDEO_FA4=1 as well (see below).
|
||||
"FASTVIDEO_ATTENTION_BACKEND":
|
||||
lambda: os.getenv("FASTVIDEO_ATTENTION_BACKEND", None),
|
||||
|
||||
# If set (=1), the FLASH_ATTN backend uses FlashAttention-4
|
||||
# (flash_attn.cute). FA4 is opt-in and never auto-selected just because it
|
||||
# is installed. Below sm90, grad-enabled and GQA calls are routed to FA2
|
||||
# (FA4's backward asserts sm90+ and its pack_gqa fails to JIT there).
|
||||
"FASTVIDEO_FA4":
|
||||
lambda: os.getenv("FASTVIDEO_FA4", "0") != "0",
|
||||
|
||||
# Use dedicated multiprocess context for workers.
|
||||
"FASTVIDEO_WORKER_MULTIPROC_METHOD":
|
||||
lambda: os.getenv("FASTVIDEO_WORKER_MULTIPROC_METHOD", "spawn"),
|
||||
|
||||
@@ -169,6 +169,14 @@ class FastVideoArgs:
|
||||
# Prompt text file for batch processing
|
||||
prompt_txt: str | None = None
|
||||
|
||||
# Dynamic multimodal generation batching. Defaults preserve the historical
|
||||
# one-request-at-a-time execution path.
|
||||
batching_mode: str = "disabled"
|
||||
batching_max_size: int = 1
|
||||
batching_delay_ms: float = 0.0
|
||||
batching_config: str | None = None
|
||||
enable_batching_metrics: bool = False
|
||||
|
||||
# LTX-2 VAE tiling overrides
|
||||
ltx2_vae_tiling: bool | None = None
|
||||
ltx2_vae_spatial_tile_size_in_pixels: int | None = None
|
||||
@@ -446,6 +454,37 @@ class FastVideoArgs:
|
||||
default=FastVideoArgs.prompt_txt,
|
||||
help="Path to a text file containing prompts (one per line) for batch processing",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--batching-mode",
|
||||
type=str,
|
||||
choices=["disabled", "dynamic"],
|
||||
default=FastVideoArgs.batching_mode,
|
||||
help="Request batching mode for inference serving.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--batching-max-size",
|
||||
type=int,
|
||||
default=FastVideoArgs.batching_max_size,
|
||||
help="Maximum number of compatible generation requests to execute as one batch.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--batching-delay-ms",
|
||||
type=float,
|
||||
default=FastVideoArgs.batching_delay_ms,
|
||||
help="Maximum queue delay in milliseconds before dispatching a dynamic batch.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--batching-config",
|
||||
type=str,
|
||||
default=FastVideoArgs.batching_config,
|
||||
help="Optional JSON batching admission rule file.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--enable-batching-metrics",
|
||||
action=StoreBoolean,
|
||||
default=FastVideoArgs.enable_batching_metrics,
|
||||
help="Log dynamic batching utilization and rejection metrics.",
|
||||
)
|
||||
|
||||
# LTX-2 VAE tiling overrides
|
||||
parser.add_argument(
|
||||
@@ -753,6 +792,13 @@ class FastVideoArgs:
|
||||
WorkloadType), f"Workload type must be a WorkloadType enum, got {type(self.workload_type)}"
|
||||
assert self.workload_type in WorkloadType.choices(), f"Invalid workload type: {self.workload_type}"
|
||||
|
||||
if self.batching_mode not in {"disabled", "dynamic"}:
|
||||
raise ValueError(f"batching_mode must be 'disabled' or 'dynamic', got {self.batching_mode!r}")
|
||||
if self.batching_max_size < 1:
|
||||
raise ValueError("batching_max_size must be >= 1")
|
||||
if self.batching_delay_ms < 0:
|
||||
raise ValueError("batching_delay_ms must be >= 0")
|
||||
|
||||
if self.mode in [ExecutionMode.DISTILLATION, ExecutionMode.FINETUNING] and self.inference_mode:
|
||||
logger.warning("Mode is 'training' but inference_mode is True. Setting inference_mode to False.")
|
||||
self.inference_mode = False
|
||||
|
||||
@@ -323,9 +323,9 @@ class CausalWanTransformerBlock(nn.Module):
|
||||
value, _ = self.to_v(norm_hidden_states)
|
||||
|
||||
if self.norm_q is not None:
|
||||
query = self.norm_q.forward_native(query)
|
||||
query = self.norm_q(query)
|
||||
if self.norm_k is not None:
|
||||
key = self.norm_k.forward_native(key)
|
||||
key = self.norm_k(key)
|
||||
|
||||
query = query.squeeze(1).unflatten(2, (self.num_attention_heads, -1))
|
||||
key = key.squeeze(1).unflatten(2, (self.num_attention_heads, -1))
|
||||
|
||||
@@ -537,9 +537,9 @@ class CausalMatrixGame2TransformerBlock(nn.Module):
|
||||
value, _ = self.to_v(norm_hidden_states)
|
||||
|
||||
if self.norm_q is not None:
|
||||
query = self.norm_q.forward_native(query)
|
||||
query = self.norm_q(query)
|
||||
if self.norm_k is not None:
|
||||
key = self.norm_k.forward_native(key)
|
||||
key = self.norm_k(key)
|
||||
|
||||
query = query.squeeze(1).unflatten(2, (self.num_attention_heads, -1))
|
||||
key = key.squeeze(1).unflatten(2, (self.num_attention_heads, -1))
|
||||
|
||||
@@ -0,0 +1,2 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""Performance benchmark and dashboard utilities."""
|
||||
@@ -57,9 +57,7 @@ def is_baseline_eligible_record(record: dict[str, Any]) -> bool:
|
||||
"""
|
||||
if record.get("baseline_eligible") is True:
|
||||
return True
|
||||
if "baseline_eligible" not in record and "run_source" not in record:
|
||||
return True
|
||||
return False
|
||||
return "baseline_eligible" not in record and "run_source" not in record
|
||||
|
||||
|
||||
def resolve_hf_token() -> str | None:
|
||||
@@ -0,0 +1,130 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""Metric policy for rolling performance baseline comparisons."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Mapping
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class MetricPolicy:
|
||||
key: str
|
||||
label: str
|
||||
precision: int
|
||||
lower_is_better: bool
|
||||
threshold_percent: float
|
||||
threshold_absolute: float
|
||||
gated: bool = True
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class MetricDelta:
|
||||
absolute: float
|
||||
percent: float
|
||||
threshold_exceeded: bool
|
||||
regressed: bool
|
||||
|
||||
|
||||
DEFAULT_METRIC_POLICIES: tuple[MetricPolicy, ...] = (
|
||||
MetricPolicy("latency", "Latency", 3, True, 0.08, 0.5),
|
||||
MetricPolicy("throughput", "Throughput", 3, False, 0.08, 0.05),
|
||||
MetricPolicy("memory", "Memory", 1, True, 0.05, 256.0),
|
||||
MetricPolicy("text_encoder_time_s", "Text Enc", 3, True, 0.05, 0.25),
|
||||
MetricPolicy("dit_time_s", "DiT", 3, True, 0.05, 0.25),
|
||||
MetricPolicy("vae_decode_time_s", "VAE Decode", 3, True, 0.05, 0.25),
|
||||
)
|
||||
|
||||
def _optional_float(value: Any) -> float | None:
|
||||
if value is None or isinstance(value, bool):
|
||||
return None
|
||||
try:
|
||||
return float(value)
|
||||
except (TypeError, ValueError):
|
||||
return None
|
||||
|
||||
|
||||
def _optional_bool(value: Any) -> bool | None:
|
||||
if isinstance(value, bool):
|
||||
return value
|
||||
if isinstance(value, str):
|
||||
normalized = value.strip().lower()
|
||||
if normalized in {"1", "true", "yes", "on"}:
|
||||
return True
|
||||
if normalized in {"0", "false", "no", "off"}:
|
||||
return False
|
||||
return None
|
||||
|
||||
|
||||
def resolve_metric_policies(
|
||||
threshold_overrides: Mapping[str, Any] | None,
|
||||
) -> tuple[MetricPolicy, ...]:
|
||||
"""Return default metric policies with optional per-metric overrides."""
|
||||
|
||||
if not isinstance(threshold_overrides, Mapping):
|
||||
threshold_overrides = {}
|
||||
policies: list[MetricPolicy] = []
|
||||
for base_policy in DEFAULT_METRIC_POLICIES:
|
||||
raw_override = threshold_overrides.get(base_policy.key, {})
|
||||
if not isinstance(raw_override, Mapping):
|
||||
raw_override = {}
|
||||
|
||||
threshold_percent = _optional_float(raw_override.get("threshold_percent"))
|
||||
threshold_absolute = _optional_float(raw_override.get("threshold_absolute"))
|
||||
gated = _optional_bool(raw_override.get("gated"))
|
||||
|
||||
policies.append(
|
||||
MetricPolicy(
|
||||
key=base_policy.key,
|
||||
label=base_policy.label,
|
||||
precision=base_policy.precision,
|
||||
lower_is_better=base_policy.lower_is_better,
|
||||
threshold_percent=(
|
||||
base_policy.threshold_percent
|
||||
if threshold_percent is None
|
||||
else threshold_percent
|
||||
),
|
||||
threshold_absolute=(
|
||||
base_policy.threshold_absolute
|
||||
if threshold_absolute is None
|
||||
else threshold_absolute
|
||||
),
|
||||
gated=base_policy.gated if gated is None else gated,
|
||||
)
|
||||
)
|
||||
return tuple(policies)
|
||||
|
||||
|
||||
def serialize_metric_thresholds(
|
||||
policies: tuple[MetricPolicy, ...],
|
||||
) -> dict[str, dict[str, float | bool]]:
|
||||
return {
|
||||
policy.key: {
|
||||
"threshold_percent": policy.threshold_percent,
|
||||
"threshold_absolute": policy.threshold_absolute,
|
||||
"gated": policy.gated,
|
||||
}
|
||||
for policy in policies
|
||||
}
|
||||
|
||||
|
||||
def regression_delta(
|
||||
policy: MetricPolicy,
|
||||
current: float,
|
||||
baseline: float,
|
||||
) -> MetricDelta | None:
|
||||
if baseline <= 0:
|
||||
return None
|
||||
absolute_delta = current - baseline if policy.lower_is_better else baseline - current
|
||||
percent_delta = absolute_delta / baseline
|
||||
threshold_exceeded = (
|
||||
percent_delta > policy.threshold_percent
|
||||
and absolute_delta > policy.threshold_absolute
|
||||
)
|
||||
return MetricDelta(
|
||||
absolute=absolute_delta,
|
||||
percent=percent_delta,
|
||||
threshold_exceeded=threshold_exceeded,
|
||||
regressed=policy.gated and threshold_exceeded,
|
||||
)
|
||||
@@ -13,7 +13,7 @@ from fastapi.middleware.cors import CORSMiddleware
|
||||
from fastapi.responses import FileResponse
|
||||
from fastapi.staticfiles import StaticFiles
|
||||
|
||||
from fastvideo.tests.performance import hf_store
|
||||
from fastvideo.performance import hf_store
|
||||
|
||||
from .service import build_latest_summary, build_trends, filter_records
|
||||
|
||||
@@ -150,7 +150,6 @@ def create_app(store: PerformanceDataStore | None = None) -> FastAPI:
|
||||
filtered = filter_records(loaded, model_id=model_id, gpu_type=gpu_type)
|
||||
rows = build_latest_summary(
|
||||
filtered,
|
||||
max_regression=float(os.environ.get("PERF_MAX_REGRESSION", "0.05")),
|
||||
run_source=run_source,
|
||||
)
|
||||
return {
|
||||
|
||||
@@ -1,26 +1,8 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""Metric definitions shared by the performance dashboard backend."""
|
||||
|
||||
from __future__ import annotations
|
||||
from fastvideo.performance.metric_policy import DEFAULT_METRIC_POLICIES
|
||||
|
||||
from dataclasses import dataclass
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class MetricDefinition:
|
||||
key: str
|
||||
label: str
|
||||
precision: int
|
||||
lower_is_better: bool
|
||||
|
||||
|
||||
METRICS: tuple[MetricDefinition, ...] = (
|
||||
MetricDefinition("latency", "Latency", 3, True),
|
||||
MetricDefinition("throughput", "Throughput", 3, False),
|
||||
MetricDefinition("memory", "Memory", 1, True),
|
||||
MetricDefinition("text_encoder_time_s", "Text Encoder", 3, True),
|
||||
MetricDefinition("dit_time_s", "DiT", 3, True),
|
||||
MetricDefinition("vae_decode_time_s", "VAE Decode", 3, True),
|
||||
)
|
||||
METRICS = DEFAULT_METRIC_POLICIES
|
||||
|
||||
METRIC_BY_KEY = {metric.key: metric for metric in METRICS}
|
||||
|
||||
@@ -13,9 +13,8 @@ from collections import defaultdict
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any
|
||||
|
||||
from fastvideo.tests.performance.hf_store import is_baseline_eligible_record, safe_float
|
||||
|
||||
from .metrics import METRICS
|
||||
from fastvideo.performance.hf_store import is_baseline_eligible_record, safe_float
|
||||
from fastvideo.performance.metric_policy import regression_delta, resolve_metric_policies
|
||||
|
||||
Record = dict[str, Any]
|
||||
|
||||
@@ -95,19 +94,9 @@ def baseline_value(records: list[Record], metric_key: str) -> float | None:
|
||||
return float(statistics.median(values))
|
||||
|
||||
|
||||
def regression_percent(metric_key: str, current: float | None, baseline: float | None) -> float | None:
|
||||
if current is None or baseline is None or baseline <= 0:
|
||||
return None
|
||||
metric = next(metric for metric in METRICS if metric.key == metric_key)
|
||||
if metric.lower_is_better:
|
||||
return (current - baseline) / baseline * 100.0
|
||||
return (baseline - current) / baseline * 100.0
|
||||
|
||||
|
||||
def build_latest_summary(records: list[Record],
|
||||
*,
|
||||
baseline_window: int = 5,
|
||||
max_regression: float = 0.05,
|
||||
run_source: str | None = None) -> list[Record]:
|
||||
rows: list[Record] = []
|
||||
for (model_id, gpu_type), group in group_by_model_gpu(records).items():
|
||||
@@ -123,52 +112,58 @@ def build_latest_summary(records: list[Record],
|
||||
if record is not latest and record.get("success", True) and is_baseline_eligible_record(record)
|
||||
]
|
||||
baseline_records = baseline_pool[-baseline_window:]
|
||||
metric_policies = resolve_metric_policies(latest.get("regression_thresholds"))
|
||||
|
||||
metrics: dict[str, Record] = {}
|
||||
regressions: list[float] = []
|
||||
for metric in METRICS:
|
||||
current = safe_float(latest.get(metric.key))
|
||||
baseline = baseline_value(baseline_records, metric.key)
|
||||
regression = regression_percent(metric.key, current, baseline)
|
||||
metrics[metric.key] = {
|
||||
failing_metrics: list[str] = []
|
||||
threshold_exceeded_metrics: list[str] = []
|
||||
for policy in metric_policies:
|
||||
current = safe_float(latest.get(policy.key))
|
||||
baseline = baseline_value(baseline_records, policy.key)
|
||||
delta = None
|
||||
if current is not None and baseline is not None:
|
||||
delta = regression_delta(policy, current, baseline)
|
||||
regression = None if delta is None else delta.percent * 100.0
|
||||
metrics[policy.key] = {
|
||||
"current": current,
|
||||
"baseline": baseline,
|
||||
"regression_pct": regression,
|
||||
"label": metric.label,
|
||||
"lower_is_better": metric.lower_is_better,
|
||||
"precision": metric.precision,
|
||||
"absolute_delta": None if delta is None else delta.absolute,
|
||||
"threshold_percent": policy.threshold_percent * 100.0,
|
||||
"threshold_absolute": policy.threshold_absolute,
|
||||
"gated": policy.gated,
|
||||
"threshold_exceeded": False if delta is None else delta.threshold_exceeded,
|
||||
"regressed": False if delta is None else delta.regressed,
|
||||
"label": policy.label,
|
||||
"lower_is_better": policy.lower_is_better,
|
||||
"precision": policy.precision,
|
||||
}
|
||||
if regression is not None:
|
||||
regressions.append(regression)
|
||||
if delta is not None and delta.threshold_exceeded:
|
||||
threshold_exceeded_metrics.append(policy.key)
|
||||
if delta is not None and delta.regressed:
|
||||
failing_metrics.append(policy.key)
|
||||
|
||||
worst_regression = max(regressions) if regressions else None
|
||||
success = bool(latest.get("success", True))
|
||||
status = "pass" if success else "fail"
|
||||
|
||||
rows.append({
|
||||
"model_id":
|
||||
model_id,
|
||||
"gpu_type":
|
||||
gpu_type,
|
||||
"timestamp":
|
||||
latest.get("timestamp"),
|
||||
"commit_sha":
|
||||
latest.get("commit_sha"),
|
||||
"model_id": model_id,
|
||||
"gpu_type": gpu_type,
|
||||
"timestamp": latest.get("timestamp"),
|
||||
"commit_sha": latest.get("commit_sha"),
|
||||
**record_metadata(latest),
|
||||
"success":
|
||||
success,
|
||||
"baseline_n":
|
||||
len(baseline_records),
|
||||
"worst_regression_pct":
|
||||
worst_regression,
|
||||
"regression_threshold_pct":
|
||||
max_regression * 100.0,
|
||||
"computed_regression_status":
|
||||
"fail" if worst_regression is not None and worst_regression > max_regression * 100.0 else "pass",
|
||||
"status":
|
||||
status,
|
||||
"metrics":
|
||||
metrics,
|
||||
"success": success,
|
||||
"baseline_n": len(baseline_records),
|
||||
"worst_regression_pct": worst_regression,
|
||||
"threshold_exceeded_metrics": threshold_exceeded_metrics,
|
||||
"failing_metrics": failing_metrics,
|
||||
"computed_regression_status": "fail" if failing_metrics else "pass",
|
||||
"status": status,
|
||||
"metrics": metrics,
|
||||
})
|
||||
|
||||
return sorted(rows, key=lambda row: (row["status"] != "fail", row["model_id"], row["gpu_type"]))
|
||||
@@ -179,14 +174,15 @@ def build_trends(records: list[Record]) -> list[Record]:
|
||||
for (model_id, gpu_type), group in group_by_model_gpu(records).items():
|
||||
points = []
|
||||
for record in group:
|
||||
metric_policies = resolve_metric_policies(record.get("regression_thresholds"))
|
||||
point = {
|
||||
"timestamp": record.get("timestamp"),
|
||||
"commit_sha": record.get("commit_sha"),
|
||||
**record_metadata(record),
|
||||
"success": bool(record.get("success", True)),
|
||||
"metrics": {
|
||||
metric.key: safe_float(record.get(metric.key))
|
||||
for metric in METRICS
|
||||
policy.key: safe_float(record.get(policy.key))
|
||||
for policy in metric_policies
|
||||
},
|
||||
}
|
||||
points.append(point)
|
||||
|
||||
@@ -234,7 +234,6 @@ class DenoisingStage(PipelineStage):
|
||||
|
||||
boundary_timestep = boundary_ratio * self.scheduler.num_train_timesteps if boundary_ratio is not None else None
|
||||
latent_model_input = latents.to(target_dtype)
|
||||
assert latent_model_input.shape[0] == 1, "only support batch size 1"
|
||||
|
||||
if fastvideo_args.pipeline_config.ti2v_task and batch.pil_image is not None:
|
||||
# TI2V directly replaces the first frame of the latent with
|
||||
|
||||
@@ -35,7 +35,11 @@ class InputValidationStage(PipelineStage):
|
||||
num_videos_per_prompt = batch.num_videos_per_prompt
|
||||
|
||||
assert seed is not None
|
||||
seeds = [seed + i for i in range(num_videos_per_prompt)]
|
||||
if batch.seeds is not None:
|
||||
seeds = batch.seeds
|
||||
else:
|
||||
prompt_count = len(batch.prompt) if isinstance(batch.prompt, list) else 1
|
||||
seeds = [seed + i for i in range(prompt_count * num_videos_per_prompt)]
|
||||
batch.seeds = seeds
|
||||
|
||||
# Peiyuan: using GPU seed will cause A100 and H100 to generate different results...
|
||||
|
||||
@@ -65,12 +65,20 @@ class TextEncodingStage(PipelineStage):
|
||||
assert batch.prompt is not None
|
||||
prompt_text: str | list[str] = batch.prompt
|
||||
all_indices: list[int] = list(range(len(self.text_encoders)))
|
||||
prompt_embeds_list, prompt_masks_list = self.encode_text(
|
||||
prompt_text,
|
||||
fastvideo_args,
|
||||
encoder_index=all_indices,
|
||||
return_attention_mask=True,
|
||||
)
|
||||
if isinstance(prompt_text, list):
|
||||
prompt_embeds_list, prompt_masks_list = self._encode_prompt_list_individually(
|
||||
prompt_text,
|
||||
fastvideo_args,
|
||||
encoder_index=all_indices,
|
||||
return_attention_mask=True,
|
||||
)
|
||||
else:
|
||||
prompt_embeds_list, prompt_masks_list = self.encode_text(
|
||||
prompt_text,
|
||||
fastvideo_args,
|
||||
encoder_index=all_indices,
|
||||
return_attention_mask=True,
|
||||
)
|
||||
if self._last_audio_embeds is not None:
|
||||
batch.extra["ltx2_audio_prompt_embeds"] = self._last_audio_embeds
|
||||
|
||||
@@ -82,13 +90,24 @@ class TextEncodingStage(PipelineStage):
|
||||
|
||||
# Encode negative prompt if CFG is enabled
|
||||
if batch.do_classifier_free_guidance:
|
||||
assert isinstance(batch.negative_prompt, str)
|
||||
neg_embeds_list, neg_masks_list = self.encode_text(
|
||||
batch.negative_prompt,
|
||||
fastvideo_args,
|
||||
encoder_index=all_indices,
|
||||
return_attention_mask=True,
|
||||
)
|
||||
assert isinstance(batch.negative_prompt, str | list)
|
||||
negative_prompt: str | list[str] = batch.negative_prompt
|
||||
if isinstance(batch.prompt, list) and isinstance(negative_prompt, str):
|
||||
negative_prompt = [negative_prompt] * len(batch.prompt)
|
||||
if isinstance(negative_prompt, list):
|
||||
neg_embeds_list, neg_masks_list = self._encode_prompt_list_individually(
|
||||
negative_prompt,
|
||||
fastvideo_args,
|
||||
encoder_index=all_indices,
|
||||
return_attention_mask=True,
|
||||
)
|
||||
else:
|
||||
neg_embeds_list, neg_masks_list = self.encode_text(
|
||||
negative_prompt,
|
||||
fastvideo_args,
|
||||
encoder_index=all_indices,
|
||||
return_attention_mask=True,
|
||||
)
|
||||
if self._last_audio_embeds is not None:
|
||||
batch.extra["ltx2_audio_negative_embeds"] = self._last_audio_embeds
|
||||
|
||||
@@ -101,6 +120,81 @@ class TextEncodingStage(PipelineStage):
|
||||
|
||||
return batch
|
||||
|
||||
def _encode_prompt_list_individually(
|
||||
self,
|
||||
texts: list[str],
|
||||
fastvideo_args: FastVideoArgs,
|
||||
*,
|
||||
encoder_index: list[int],
|
||||
return_attention_mask: bool,
|
||||
) -> tuple[list[torch.Tensor], list[torch.Tensor]]:
|
||||
per_prompt_embeds: list[list[torch.Tensor]] = []
|
||||
per_prompt_masks: list[list[torch.Tensor]] = []
|
||||
per_prompt_audio_embeds: list[list[torch.Tensor] | None] = []
|
||||
|
||||
for text in texts:
|
||||
embeds, masks = self.encode_text(
|
||||
text,
|
||||
fastvideo_args,
|
||||
encoder_index=encoder_index,
|
||||
return_attention_mask=return_attention_mask,
|
||||
)
|
||||
per_prompt_embeds.append(embeds)
|
||||
per_prompt_masks.append(masks)
|
||||
per_prompt_audio_embeds.append(self._last_audio_embeds)
|
||||
|
||||
merged_embeds = [
|
||||
self._cat_tensors([prompt_embeds[encoder_pos] for prompt_embeds in per_prompt_embeds])
|
||||
for encoder_pos in range(len(per_prompt_embeds[0]))
|
||||
]
|
||||
merged_masks = [
|
||||
self._cat_attention_masks([prompt_masks[encoder_pos] for prompt_masks in per_prompt_masks])
|
||||
for encoder_pos in range(len(per_prompt_masks[0]))
|
||||
]
|
||||
if per_prompt_audio_embeds and all(audio_embeds is not None for audio_embeds in per_prompt_audio_embeds):
|
||||
audio_embed_lists = [audio_embeds for audio_embeds in per_prompt_audio_embeds if audio_embeds is not None]
|
||||
self._last_audio_embeds = [
|
||||
self._cat_tensors([audio_embeds[encoder_pos] for audio_embeds in audio_embed_lists])
|
||||
for encoder_pos in range(len(audio_embed_lists[0]))
|
||||
]
|
||||
else:
|
||||
self._last_audio_embeds = None
|
||||
return merged_embeds, merged_masks
|
||||
|
||||
@staticmethod
|
||||
def _cat_tensors(tensors: list[torch.Tensor]) -> torch.Tensor:
|
||||
base_shape = tensors[0].shape[1:]
|
||||
if all(tensor.shape[1:] == base_shape for tensor in tensors):
|
||||
return torch.cat(tensors, dim=0)
|
||||
if all(tensor.ndim == 3 for tensor in tensors):
|
||||
base_trailing_shape = tensors[0].shape[2:]
|
||||
if all(tensor.shape[2:] == base_trailing_shape for tensor in tensors):
|
||||
max_length = max(tensor.shape[1] for tensor in tensors)
|
||||
padded_tensors = []
|
||||
for tensor in tensors:
|
||||
pad_width = max_length - tensor.shape[1]
|
||||
if pad_width > 0:
|
||||
tensor = torch.nn.functional.pad(tensor, (0, 0, 0, pad_width), value=0.0)
|
||||
padded_tensors.append(tensor)
|
||||
return torch.cat(padded_tensors, dim=0)
|
||||
raise ValueError(f"Cannot concatenate tensors with shapes: {[list(tensor.shape) for tensor in tensors]}")
|
||||
|
||||
@staticmethod
|
||||
def _cat_attention_masks(masks: list[torch.Tensor]) -> torch.Tensor:
|
||||
base_shape = masks[0].shape[1:]
|
||||
if all(mask.shape[1:] == base_shape for mask in masks):
|
||||
return torch.cat(masks, dim=0)
|
||||
if all(mask.ndim == 2 for mask in masks):
|
||||
max_length = max(mask.shape[1] for mask in masks)
|
||||
padded_masks = []
|
||||
for mask in masks:
|
||||
pad_width = max_length - mask.shape[1]
|
||||
if pad_width > 0:
|
||||
mask = torch.nn.functional.pad(mask, (0, pad_width), value=0)
|
||||
padded_masks.append(mask)
|
||||
return torch.cat(padded_masks, dim=0)
|
||||
raise ValueError(f"Cannot concatenate attention masks with shapes: {[list(mask.shape) for mask in masks]}")
|
||||
|
||||
def verify_input(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult:
|
||||
"""Verify text encoding stage inputs."""
|
||||
result = VerificationResult()
|
||||
@@ -235,6 +329,9 @@ class TextEncodingStage(PipelineStage):
|
||||
attn_masks_list.append(attention_mask)
|
||||
return self.return_embeds(embeds_list, attn_masks_list, return_type, return_attention_mask, indices)
|
||||
|
||||
if len(processed_texts) > 1 and "padding" not in tok_kwargs:
|
||||
tok_kwargs["padding"] = True
|
||||
|
||||
# If tokenizer is a multimodal processor (e.g. Qwen2_5_VLProcessor),
|
||||
# use its inner tokenizer for text-only encoding.
|
||||
tok = getattr(tokenizer, "tokenizer", tokenizer)
|
||||
|
||||
@@ -9,7 +9,7 @@ from fastvideo.api.compat import (
|
||||
generator_config_to_fastvideo_args,
|
||||
legacy_from_pretrained_to_config,
|
||||
)
|
||||
from fastvideo.api.schema import CompileConfig, GeneratorConfig
|
||||
from fastvideo.api.schema import BatchingConfig, CompileConfig, GeneratorConfig
|
||||
|
||||
|
||||
class TestLegacyTorchCompileKwargsTranslation:
|
||||
@@ -200,6 +200,48 @@ class TestLegacyTextEncoderCompileTranslation:
|
||||
assert "enable_torch_compile_text_encoder" not in args.kwargs
|
||||
|
||||
|
||||
class TestBatchingTranslation:
|
||||
|
||||
def test_flat_kwargs_promote_to_engine_batching(self) -> None:
|
||||
config = legacy_from_pretrained_to_config(
|
||||
"/models/wan",
|
||||
{
|
||||
"batching_mode": "dynamic",
|
||||
"batching_max_size": 4,
|
||||
"batching_delay_ms": 25.0,
|
||||
"batching_config": "/tmp/batching.json",
|
||||
"enable_batching_metrics": True,
|
||||
},
|
||||
)
|
||||
|
||||
assert config.engine.batching.mode == "dynamic"
|
||||
assert config.engine.batching.max_size == 4
|
||||
assert config.engine.batching.delay_ms == 25.0
|
||||
assert config.engine.batching.config_path == "/tmp/batching.json"
|
||||
assert config.engine.batching.enable_metrics is True
|
||||
|
||||
def test_typed_batching_emits_fastvideo_args_kwargs(self, monkeypatch) -> None:
|
||||
_stub_fastvideo_args_from_kwargs(monkeypatch)
|
||||
config = GeneratorConfig(
|
||||
model_path="/models/wan",
|
||||
engine=_engine_with_batching(BatchingConfig(
|
||||
mode="dynamic",
|
||||
max_size=3,
|
||||
delay_ms=10.0,
|
||||
config_path="/tmp/batching.json",
|
||||
enable_metrics=True,
|
||||
)),
|
||||
)
|
||||
|
||||
args = generator_config_to_fastvideo_args(config)
|
||||
|
||||
assert args.kwargs["batching_mode"] == "dynamic"
|
||||
assert args.kwargs["batching_max_size"] == 3
|
||||
assert args.kwargs["batching_delay_ms"] == 10.0
|
||||
assert args.kwargs["batching_config"] == "/tmp/batching.json"
|
||||
assert args.kwargs["enable_batching_metrics"] is True
|
||||
|
||||
|
||||
# -------------------------------------------------------------------
|
||||
# Helpers
|
||||
# -------------------------------------------------------------------
|
||||
@@ -213,6 +255,13 @@ def _engine_with_compile(compile_config):
|
||||
return engine
|
||||
|
||||
|
||||
def _engine_with_batching(batching_config):
|
||||
from fastvideo.api.schema import EngineConfig
|
||||
engine = EngineConfig()
|
||||
engine.batching = batching_config
|
||||
return engine
|
||||
|
||||
|
||||
def _stub_fastvideo_args_from_kwargs(monkeypatch):
|
||||
"""Swap ``FastVideoArgs.from_kwargs`` for a capture-only stub so
|
||||
translation tests don't need to construct a valid FastVideoArgs."""
|
||||
|
||||
@@ -130,6 +130,13 @@ def test_load_run_config_supports_yaml_roundtrip(tmp_path) -> None:
|
||||
"use_fsdp_inference": False,
|
||||
"disable_autocast": False,
|
||||
"quantization": None,
|
||||
"batching": {
|
||||
"mode": "disabled",
|
||||
"max_size": 1,
|
||||
"delay_ms": 0.0,
|
||||
"config_path": None,
|
||||
"enable_metrics": False,
|
||||
},
|
||||
},
|
||||
"pipeline": {
|
||||
"workload_type": None,
|
||||
|
||||
@@ -3,15 +3,16 @@
|
||||
|
||||
Landed in PR #1225 slice 5 (Attn-QAT 5/12). The resolver centralises the
|
||||
varlen-flash-attn import-fallback logic that several backends
|
||||
(``bsa_attn.py``, ``video_sparse_attn.py``) used to duplicate. The fallback
|
||||
(``bsa_attn.py``, ``video_sparse_attn.py``) used to duplicate. The resolution
|
||||
order is:
|
||||
|
||||
1. ``fastvideo.attention.utils.flash_attn_cute``
|
||||
1. ``fastvideo.attention.utils.flash_attn_cute`` -- only when
|
||||
``FASTVIDEO_FA4=1`` (explicit opt-in), and then it must import or the
|
||||
resolver raises RuntimeError instead of falling through
|
||||
2. ``flash_attn_interface``
|
||||
3. ``flash_attn``
|
||||
|
||||
These tests verify that the resolver picks the highest-priority impl
|
||||
available and falls through cleanly on ``ImportError``. CPU-only, no
|
||||
These tests verify the opt-in gate and the FA3/FA2 fallthrough. CPU-only, no
|
||||
flash-attn install required.
|
||||
"""
|
||||
|
||||
@@ -35,8 +36,36 @@ def _reload_resolver_module():
|
||||
return importlib.import_module("fastvideo.attention.utils.flash_attn_no_pad")
|
||||
|
||||
|
||||
def test_resolver_falls_back_when_cute_unavailable(monkeypatch) -> None:
|
||||
"""When ``flash_attn_cute`` is unimportable, resolver tries the next impl."""
|
||||
def test_resolver_skips_cute_without_opt_in(monkeypatch) -> None:
|
||||
"""Without ``FASTVIDEO_FA4=1`` the resolver must not even attempt the cute
|
||||
import."""
|
||||
monkeypatch.delenv("FASTVIDEO_FA4", raising=False)
|
||||
attempted: list[str] = []
|
||||
real_import = builtins.__import__
|
||||
|
||||
def spying_import(name, globals=None, locals=None, fromlist=(), level=0):
|
||||
attempted.append(name)
|
||||
return real_import(name, globals, locals, fromlist, level)
|
||||
|
||||
monkeypatch.setattr(builtins, "__import__", spying_import)
|
||||
|
||||
mod = _reload_resolver_module()
|
||||
resolved = mod._resolve_flash_attn_varlen_func()
|
||||
assert resolved is not None
|
||||
assert resolved.__name__ == "flash_attn_varlen_func"
|
||||
assert "fastvideo.attention.utils.flash_attn_cute" not in attempted
|
||||
|
||||
|
||||
def test_resolver_raises_when_opted_in_but_cute_unavailable(monkeypatch) -> None:
|
||||
"""With ``FASTVIDEO_FA4=1`` an unimportable cute build fails loudly instead
|
||||
of silently falling through to FA3/FA2.
|
||||
|
||||
The resolver runs at module import time, so the reload itself must raise.
|
||||
It raises RuntimeError (not ImportError) so importers that treat
|
||||
ImportError as "flash-attn not installed" (``bsa_attn.py``) cannot swallow
|
||||
the opted-in failure.
|
||||
"""
|
||||
monkeypatch.setenv("FASTVIDEO_FA4", "1")
|
||||
real_import = builtins.__import__
|
||||
|
||||
def patched_import(name, globals=None, locals=None, fromlist=(), level=0):
|
||||
@@ -46,21 +75,17 @@ def test_resolver_falls_back_when_cute_unavailable(monkeypatch) -> None:
|
||||
|
||||
monkeypatch.setattr(builtins, "__import__", patched_import)
|
||||
|
||||
mod = _reload_resolver_module()
|
||||
resolved = mod._resolve_flash_attn_varlen_func()
|
||||
assert resolved is not None
|
||||
assert resolved.__name__ == "flash_attn_varlen_func"
|
||||
with pytest.raises(RuntimeError, match="cute disabled for test"):
|
||||
_reload_resolver_module()
|
||||
|
||||
|
||||
def test_resolver_returns_flash_attn_when_cute_and_interface_unavailable(monkeypatch) -> None:
|
||||
def test_resolver_returns_flash_attn_when_interface_unavailable(monkeypatch) -> None:
|
||||
"""The terminal fallback is the plain ``flash_attn`` import."""
|
||||
monkeypatch.delenv("FASTVIDEO_FA4", raising=False)
|
||||
real_import = builtins.__import__
|
||||
|
||||
def patched_import(name, globals=None, locals=None, fromlist=(), level=0):
|
||||
if name in {
|
||||
"fastvideo.attention.utils.flash_attn_cute",
|
||||
"flash_attn_interface",
|
||||
}:
|
||||
if name == "flash_attn_interface":
|
||||
raise ImportError(f"{name} disabled for test")
|
||||
return real_import(name, globals, locals, fromlist, level)
|
||||
|
||||
|
||||
@@ -0,0 +1,198 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
|
||||
from fastvideo import VideoGenerator
|
||||
|
||||
DEFAULT_PROMPTS = (
|
||||
"A small robot sketches a city skyline at sunrise, cinematic lighting.",
|
||||
"A glass teapot steams on a wooden table while rain falls outside.",
|
||||
)
|
||||
|
||||
|
||||
def _build_init_kwargs(args: argparse.Namespace, *, dynamic: bool) -> dict[str, Any]:
|
||||
return {
|
||||
"num_gpus": args.num_gpus,
|
||||
"sp_size": args.sp_size,
|
||||
"tp_size": args.tp_size,
|
||||
"use_fsdp_inference": args.use_fsdp_inference,
|
||||
"dit_cpu_offload": False,
|
||||
"dit_layerwise_offload": False,
|
||||
"flow_shift": args.flow_shift,
|
||||
"text_encoder_precisions": ("fp32",),
|
||||
"output_type": "latent",
|
||||
"batching_mode": "dynamic" if dynamic else "disabled",
|
||||
"batching_max_size": args.batch_size if dynamic else 1,
|
||||
"batching_delay_ms": 0.0,
|
||||
}
|
||||
|
||||
|
||||
def _request_kwargs(args: argparse.Namespace, prompt_index: int) -> dict[str, Any]:
|
||||
return {
|
||||
"prompt": args.prompts[prompt_index],
|
||||
"height": args.height,
|
||||
"width": args.width,
|
||||
"num_frames": args.num_frames,
|
||||
"num_inference_steps": args.num_inference_steps,
|
||||
"guidance_scale": args.guidance_scale,
|
||||
"embedded_cfg_scale": args.embedded_cfg_scale,
|
||||
"seed": args.seed + prompt_index,
|
||||
"fps": 24,
|
||||
"save_video": False,
|
||||
"return_frames": True,
|
||||
"output_path": str(Path(args.output_dir) / f"request_{prompt_index}.mp4"),
|
||||
}
|
||||
|
||||
|
||||
def _sync() -> None:
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.synchronize()
|
||||
|
||||
|
||||
def _run_sequential(generator: VideoGenerator, args: argparse.Namespace) -> tuple[list[dict[str, Any]], float]:
|
||||
_sync()
|
||||
start = time.perf_counter()
|
||||
results = []
|
||||
for index in range(args.batch_size):
|
||||
kwargs = _request_kwargs(args, index)
|
||||
prompt = kwargs.pop("prompt")
|
||||
results.append(generator.generate_video(prompt=prompt, **kwargs))
|
||||
_sync()
|
||||
return results, time.perf_counter() - start
|
||||
|
||||
|
||||
def _run_dynamic(generator: VideoGenerator, args: argparse.Namespace) -> tuple[list[dict[str, Any]], float]:
|
||||
if not hasattr(generator, "generate_video_batch"):
|
||||
raise RuntimeError("VideoGenerator.generate_video_batch is unavailable in this checkout")
|
||||
requests = [_request_kwargs(args, index) for index in range(args.batch_size)]
|
||||
_sync()
|
||||
start = time.perf_counter()
|
||||
results = generator.generate_video_batch(requests)
|
||||
_sync()
|
||||
return results, time.perf_counter() - start
|
||||
|
||||
|
||||
def _tensor_metrics(sequential: list[dict[str, Any]], dynamic: list[dict[str, Any]]) -> dict[str, Any]:
|
||||
per_request = []
|
||||
for index, (seq_result, dyn_result) in enumerate(zip(sequential, dynamic, strict=True)):
|
||||
seq = seq_result["samples"].detach().cpu().to(torch.float32)
|
||||
dyn = dyn_result["samples"].detach().cpu().to(torch.float32)
|
||||
diff = (seq - dyn).abs()
|
||||
per_request.append({
|
||||
"index": index,
|
||||
"shape": list(seq.shape),
|
||||
"max_abs_diff": float(diff.max().item()),
|
||||
"mean_abs_diff": float(diff.mean().item()),
|
||||
"allclose_atol_1e_5": bool(torch.allclose(seq, dyn, atol=1e-5, rtol=1e-5)),
|
||||
"allclose_atol_1e_4": bool(torch.allclose(seq, dyn, atol=1e-4, rtol=1e-4)),
|
||||
})
|
||||
return {
|
||||
"per_request": per_request,
|
||||
"max_abs_diff": max(item["max_abs_diff"] for item in per_request),
|
||||
"mean_abs_diff": sum(item["mean_abs_diff"] for item in per_request) / len(per_request),
|
||||
"allclose_atol_1e_5": all(item["allclose_atol_1e_5"] for item in per_request),
|
||||
"allclose_atol_1e_4": all(item["allclose_atol_1e_4"] for item in per_request),
|
||||
}
|
||||
|
||||
|
||||
def run_parity(args: argparse.Namespace) -> dict[str, Any]:
|
||||
generator = VideoGenerator.from_pretrained(args.model_path, **_build_init_kwargs(args, dynamic=True))
|
||||
try:
|
||||
sequential, sequential_s = _run_sequential(generator, args)
|
||||
dynamic, dynamic_s = _run_dynamic(generator, args)
|
||||
metrics = _tensor_metrics(sequential, dynamic)
|
||||
finally:
|
||||
generator.shutdown()
|
||||
return {
|
||||
"mode": "parity",
|
||||
"model_path": args.model_path,
|
||||
"num_gpus": args.num_gpus,
|
||||
"shape": {
|
||||
"height": args.height,
|
||||
"width": args.width,
|
||||
"num_frames": args.num_frames,
|
||||
"num_inference_steps": args.num_inference_steps,
|
||||
},
|
||||
"batch_size": args.batch_size,
|
||||
"sequential_time_s": sequential_s,
|
||||
"dynamic_time_s": dynamic_s,
|
||||
"speedup": sequential_s / dynamic_s if dynamic_s > 0 else None,
|
||||
"tensor_metrics": metrics,
|
||||
}
|
||||
|
||||
|
||||
def run_benchmark(args: argparse.Namespace, *, dynamic: bool) -> dict[str, Any]:
|
||||
generator = VideoGenerator.from_pretrained(args.model_path, **_build_init_kwargs(args, dynamic=dynamic))
|
||||
run = _run_dynamic if dynamic else _run_sequential
|
||||
try:
|
||||
for _ in range(args.warmup_runs):
|
||||
run(generator, args)
|
||||
times = []
|
||||
for _ in range(args.measurement_runs):
|
||||
_results, elapsed = run(generator, args)
|
||||
times.append(elapsed)
|
||||
finally:
|
||||
generator.shutdown()
|
||||
avg = sum(times) / len(times)
|
||||
return {
|
||||
"mode": "dynamic" if dynamic else "sequential",
|
||||
"model_path": args.model_path,
|
||||
"num_gpus": args.num_gpus,
|
||||
"batch_size": args.batch_size,
|
||||
"measurement_runs": args.measurement_runs,
|
||||
"times_s": times,
|
||||
"avg_time_s": avg,
|
||||
"requests_per_second": args.batch_size / avg if avg > 0 else None,
|
||||
}
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--mode", choices=("parity", "sequential", "dynamic"), required=True)
|
||||
parser.add_argument("--model-path", default="Wan-AI/Wan2.1-T2V-1.3B-Diffusers")
|
||||
parser.add_argument("--num-gpus", type=int, default=1)
|
||||
parser.add_argument("--sp-size", type=int, default=1)
|
||||
parser.add_argument("--tp-size", type=int, default=1)
|
||||
parser.add_argument("--use-fsdp-inference", action="store_true")
|
||||
parser.add_argument("--height", type=int, default=256)
|
||||
parser.add_argument("--width", type=int, default=256)
|
||||
parser.add_argument("--num-frames", type=int, default=9)
|
||||
parser.add_argument("--num-inference-steps", type=int, default=2)
|
||||
parser.add_argument("--guidance-scale", type=float, default=1.0)
|
||||
parser.add_argument("--embedded-cfg-scale", type=float, default=6.0)
|
||||
parser.add_argument("--flow-shift", type=float, default=7.0)
|
||||
parser.add_argument("--seed", type=int, default=1024)
|
||||
parser.add_argument("--batch-size", type=int, default=2)
|
||||
parser.add_argument("--warmup-runs", type=int, default=1)
|
||||
parser.add_argument("--measurement-runs", type=int, default=3)
|
||||
parser.add_argument("--output-dir", default="/tmp/fastvideo_dynamic_batching")
|
||||
parser.add_argument("--output-json", required=True)
|
||||
parser.add_argument("--prompts", nargs="+", default=list(DEFAULT_PROMPTS))
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
if len(args.prompts) < args.batch_size:
|
||||
raise ValueError("--prompts must contain at least --batch-size prompts")
|
||||
os.makedirs(args.output_dir, exist_ok=True)
|
||||
if args.mode == "parity":
|
||||
result = run_parity(args)
|
||||
else:
|
||||
result = run_benchmark(args, dynamic=args.mode == "dynamic")
|
||||
output_path = Path(args.output_json)
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
output_path.write_text(json.dumps(result, indent=2), encoding="utf-8")
|
||||
print(json.dumps(result, indent=2))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,89 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
from __future__ import annotations
|
||||
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
|
||||
from fastvideo.batching.admission import (
|
||||
AdmissionLimit,
|
||||
BatchAdmissionController,
|
||||
BatchingRule,
|
||||
load_batching_config,
|
||||
)
|
||||
from fastvideo.configs.pipelines.base import PipelineConfig
|
||||
|
||||
|
||||
def test_admission_limit_rejects_batch_size_and_cost() -> None:
|
||||
limit = AdmissionLimit(max_batch_size=2, max_cost=10.0)
|
||||
|
||||
assert limit.reject_reason(batch_size=3, batch_cost=1.0) == "config_cap:2"
|
||||
assert limit.reject_reason(batch_size=2, batch_cost=11.0) == "cost_budget:11>10"
|
||||
assert limit.reject_reason(batch_size=2, batch_cost=10.0) is None
|
||||
|
||||
|
||||
def test_batching_rule_validates_unknown_keys() -> None:
|
||||
with pytest.raises(ValueError, match="did you mean 'max_batch_size'"):
|
||||
BatchingRule.from_dict(
|
||||
{
|
||||
"model_contains": "wan",
|
||||
"max_batch_siz": 2,
|
||||
},
|
||||
source="unit",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(("value", "expected"), [(1, True), (0, False), (1.0, True), (0.0, False)])
|
||||
def test_batching_rule_parses_numeric_bool_values(value, expected) -> None:
|
||||
rule = BatchingRule.from_dict(
|
||||
{
|
||||
"model_contains": "wan",
|
||||
"offload": value,
|
||||
"max_batch_size": 2,
|
||||
},
|
||||
source="unit",
|
||||
)
|
||||
|
||||
assert rule.offload is expected
|
||||
|
||||
|
||||
def test_load_batching_config_supports_mapping_form(tmp_path) -> None:
|
||||
path = tmp_path / "batching.json"
|
||||
path.write_text(
|
||||
'{"schema_version": 1, "wan|720x1280x81": {"max_batch_size": 3, "max_cost": 9}}',
|
||||
encoding="utf-8",
|
||||
)
|
||||
|
||||
rules = load_batching_config(str(path))
|
||||
|
||||
assert len(rules) == 1
|
||||
assert rules[0].model == "wan"
|
||||
assert rules[0].resolution == "720x1280x81"
|
||||
assert rules[0].max_batch_size == 3
|
||||
assert rules[0].max_cost == 9.0
|
||||
|
||||
|
||||
def test_admission_controller_applies_user_and_config_caps(tmp_path, monkeypatch) -> None:
|
||||
path = tmp_path / "batching.json"
|
||||
path.write_text(
|
||||
'{"rules": [{"model_contains": "wan", "resolution": "720x1280x81", "max_batch_size": 3}]}',
|
||||
encoding="utf-8",
|
||||
)
|
||||
monkeypatch.setattr(BatchAdmissionController, "_get_device_memory_gb", staticmethod(lambda gpu_id: 48.0))
|
||||
|
||||
args = SimpleNamespace(
|
||||
batching_mode="dynamic",
|
||||
batching_max_size=4,
|
||||
batching_config=str(path),
|
||||
model_path="/models/wan",
|
||||
dit_cpu_offload=False,
|
||||
dit_layerwise_offload=False,
|
||||
pipeline_config=PipelineConfig(),
|
||||
)
|
||||
request = SimpleNamespace(height=720, width=1280, num_frames=81)
|
||||
|
||||
controller = BatchAdmissionController(args)
|
||||
|
||||
assert controller.enabled is True
|
||||
assert controller.max_admissible_batch_size(request) == 3
|
||||
assert controller.batch_is_full([request, request, request]) is True
|
||||
@@ -0,0 +1,57 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
from __future__ import annotations
|
||||
|
||||
from fastvideo.api.sampling_param import SamplingParam
|
||||
from fastvideo.batching.signature import (
|
||||
can_dynamic_batch,
|
||||
dynamic_batch_signature,
|
||||
resolution_key,
|
||||
)
|
||||
|
||||
|
||||
def _request(prompt: str = "a prompt", **overrides) -> SamplingParam:
|
||||
request = SamplingParam(prompt=prompt, height=256, width=384, num_frames=17, num_inference_steps=4)
|
||||
for key, value in overrides.items():
|
||||
setattr(request, key, value)
|
||||
return request
|
||||
|
||||
|
||||
def test_dynamic_batch_signature_excludes_request_local_fields() -> None:
|
||||
first = _request(seed=1, output_path="/tmp/a.mp4", save_video=True, return_frames=False)
|
||||
second = _request(seed=2, output_path="/tmp/b.mp4", save_video=False, return_frames=True)
|
||||
|
||||
assert dynamic_batch_signature(first) == dynamic_batch_signature(second)
|
||||
|
||||
|
||||
def test_can_dynamic_batch_accepts_matching_text_requests() -> None:
|
||||
first = _request("first", seed=1)
|
||||
second = _request("second", seed=2)
|
||||
|
||||
result = can_dynamic_batch(first, second)
|
||||
|
||||
assert result.can_batch is True
|
||||
assert result.reason is None
|
||||
|
||||
|
||||
def test_can_dynamic_batch_rejects_sampling_mismatch() -> None:
|
||||
first = _request(guidance_scale=1.0)
|
||||
second = _request(guidance_scale=3.0)
|
||||
|
||||
result = can_dynamic_batch(first, second)
|
||||
|
||||
assert result.can_batch is False
|
||||
assert result.reason == "sampling_params.guidance_scale"
|
||||
|
||||
|
||||
def test_can_dynamic_batch_rejects_image_conditioning() -> None:
|
||||
first = _request()
|
||||
second = _request(image_path="/tmp/image.png")
|
||||
|
||||
result = can_dynamic_batch(first, second)
|
||||
|
||||
assert result.can_batch is False
|
||||
assert result.reason == "image_path"
|
||||
|
||||
|
||||
def test_resolution_key_uses_generation_shape() -> None:
|
||||
assert resolution_key(_request(height=720, width=1280, num_frames=81)) == "720x1280x81"
|
||||
@@ -1,10 +1,15 @@
|
||||
"""Unit tests for the OpenAI-compatible API server helpers (no GPU needed)."""
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
import time
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
|
||||
from fastvideo.configs.pipelines.base import PipelineConfig
|
||||
from fastvideo.entrypoints.openai.batching import VideoBatchScheduler, _VideoBatchJob
|
||||
from fastvideo.api.parser import parse_config
|
||||
from fastvideo.api.schema import GenerationRequest
|
||||
from fastvideo.entrypoints.openai.protocol import (
|
||||
@@ -21,6 +26,172 @@ from fastvideo.entrypoints.openai.utils import (
|
||||
parse_size,
|
||||
)
|
||||
|
||||
|
||||
class _FakeBatchGenerator:
|
||||
|
||||
def __init__(self):
|
||||
self.calls = []
|
||||
|
||||
def generate_video_batch(self, request_kwargs):
|
||||
self.calls.append([dict(item) for item in request_kwargs])
|
||||
return [{"prompts": item["prompt"], "video_path": item["output_path"]} for item in request_kwargs]
|
||||
|
||||
|
||||
def _make_batch_job(request_id, kwargs):
|
||||
loop = asyncio.get_running_loop()
|
||||
return _VideoBatchJob(
|
||||
request_id=request_id,
|
||||
kwargs=dict(kwargs),
|
||||
future=loop.create_future(),
|
||||
enqueue_time=time.perf_counter(),
|
||||
)
|
||||
|
||||
|
||||
def _batch_scheduler_args(**overrides):
|
||||
defaults = dict(
|
||||
model_path="test-model",
|
||||
batching_mode="dynamic",
|
||||
batching_max_size=2,
|
||||
batching_delay_ms=25.0,
|
||||
enable_batching_metrics=False,
|
||||
pipeline_config=PipelineConfig(),
|
||||
)
|
||||
defaults.update(overrides)
|
||||
return SimpleNamespace(**defaults)
|
||||
|
||||
|
||||
def test_video_batch_scheduler_groups_compatible_requests(tmp_path):
|
||||
async def run():
|
||||
generator = _FakeBatchGenerator()
|
||||
scheduler = VideoBatchScheduler(generator, _batch_scheduler_args())
|
||||
await scheduler.start()
|
||||
try:
|
||||
first = {
|
||||
"prompt": "first",
|
||||
"height": 256,
|
||||
"width": 256,
|
||||
"num_frames": 1,
|
||||
"num_inference_steps": 2,
|
||||
"seed": 1,
|
||||
"output_path": str(tmp_path / "first.mp4"),
|
||||
"save_video": False,
|
||||
}
|
||||
second = {
|
||||
"prompt": "second",
|
||||
"height": 256,
|
||||
"width": 256,
|
||||
"num_frames": 1,
|
||||
"num_inference_steps": 2,
|
||||
"seed": 2,
|
||||
"output_path": str(tmp_path / "second.mp4"),
|
||||
"save_video": False,
|
||||
}
|
||||
results = await asyncio.gather(
|
||||
scheduler.submit("req-1", first),
|
||||
scheduler.submit("req-2", second),
|
||||
)
|
||||
finally:
|
||||
await scheduler.stop()
|
||||
return generator.calls, results
|
||||
|
||||
calls, results = asyncio.run(run())
|
||||
|
||||
assert len(calls) == 1
|
||||
assert [item["prompt"] for item in calls[0]] == ["first", "second"]
|
||||
assert [result["prompts"] for result in results] == ["first", "second"]
|
||||
|
||||
|
||||
def test_video_batch_scheduler_keeps_incompatible_requests_separate(tmp_path):
|
||||
async def run():
|
||||
generator = _FakeBatchGenerator()
|
||||
scheduler = VideoBatchScheduler(generator, _batch_scheduler_args())
|
||||
await scheduler.start()
|
||||
try:
|
||||
text_only = {
|
||||
"prompt": "first",
|
||||
"height": 256,
|
||||
"width": 256,
|
||||
"num_frames": 1,
|
||||
"num_inference_steps": 2,
|
||||
"seed": 1,
|
||||
"output_path": str(tmp_path / "first.mp4"),
|
||||
"save_video": False,
|
||||
}
|
||||
image_conditioned = {
|
||||
"prompt": "second",
|
||||
"height": 256,
|
||||
"width": 256,
|
||||
"num_frames": 1,
|
||||
"num_inference_steps": 2,
|
||||
"seed": 2,
|
||||
"image_path": str(tmp_path / "input.png"),
|
||||
"output_path": str(tmp_path / "second.mp4"),
|
||||
"save_video": False,
|
||||
}
|
||||
results = await asyncio.gather(
|
||||
scheduler.submit("req-1", text_only),
|
||||
scheduler.submit("req-2", image_conditioned),
|
||||
)
|
||||
finally:
|
||||
await scheduler.stop()
|
||||
return generator.calls, results
|
||||
|
||||
calls, results = asyncio.run(run())
|
||||
|
||||
assert len(calls) == 2
|
||||
assert [[item["prompt"] for item in call] for call in calls] == [["first"], ["second"]]
|
||||
assert [result["prompts"] for result in results] == ["first", "second"]
|
||||
|
||||
|
||||
def test_video_batch_scheduler_requeues_incompatible_pending_job_at_front(tmp_path):
|
||||
async def run():
|
||||
generator = _FakeBatchGenerator()
|
||||
scheduler = VideoBatchScheduler(generator, _batch_scheduler_args())
|
||||
first = {
|
||||
"prompt": "first",
|
||||
"height": 256,
|
||||
"width": 256,
|
||||
"num_frames": 1,
|
||||
"num_inference_steps": 2,
|
||||
"seed": 1,
|
||||
"output_path": str(tmp_path / "first.mp4"),
|
||||
"save_video": False,
|
||||
}
|
||||
incompatible = {
|
||||
"prompt": "second",
|
||||
"height": 256,
|
||||
"width": 256,
|
||||
"num_frames": 1,
|
||||
"num_inference_steps": 2,
|
||||
"seed": 2,
|
||||
"image_path": str(tmp_path / "input.png"),
|
||||
"output_path": str(tmp_path / "second.mp4"),
|
||||
"save_video": False,
|
||||
}
|
||||
newer = {
|
||||
"prompt": "third",
|
||||
"height": 256,
|
||||
"width": 256,
|
||||
"num_frames": 1,
|
||||
"num_inference_steps": 2,
|
||||
"seed": 3,
|
||||
"output_path": str(tmp_path / "third.mp4"),
|
||||
"save_video": False,
|
||||
}
|
||||
|
||||
scheduler._pending.extend([
|
||||
_make_batch_job("req-2", incompatible),
|
||||
_make_batch_job("req-3", newer),
|
||||
])
|
||||
batch = await scheduler._collect_batch(_make_batch_job("req-1", first))
|
||||
return [job.request_id for job in batch], [job.request_id for job in scheduler._pending]
|
||||
|
||||
batch_ids, pending_ids = asyncio.run(run())
|
||||
|
||||
assert batch_ids == ["req-1"]
|
||||
assert pending_ids == ["req-2", "req-3"]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# parse_size
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@@ -3,6 +3,7 @@ from types import SimpleNamespace
|
||||
import warnings
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from fastvideo.api import (
|
||||
GenerationRequest,
|
||||
@@ -13,8 +14,10 @@ from fastvideo.api import (
|
||||
load_run_config,
|
||||
)
|
||||
from fastvideo.api.sampling_param import SamplingParam
|
||||
from fastvideo.configs.pipelines.base import PipelineConfig
|
||||
from fastvideo.entrypoints.video_generator import VideoGenerator
|
||||
from fastvideo.fastvideo_args import WorkloadType
|
||||
from fastvideo.pipelines import ForwardBatch
|
||||
|
||||
|
||||
def _new_video_generator() -> VideoGenerator:
|
||||
@@ -37,6 +40,25 @@ def _new_runtime_video_generator() -> VideoGenerator:
|
||||
return generator
|
||||
|
||||
|
||||
def _batching_fastvideo_args(**overrides):
|
||||
defaults = dict(
|
||||
model_path="test-model",
|
||||
prompt_txt=None,
|
||||
workload_type=SimpleNamespace(value="t2v"),
|
||||
batching_mode="dynamic",
|
||||
batching_max_size=4,
|
||||
batching_config=None,
|
||||
dit_cpu_offload=False,
|
||||
dit_layerwise_offload=False,
|
||||
output_type="latent",
|
||||
pin_cpu_memory=False,
|
||||
VSA_sparsity=0.0,
|
||||
pipeline_config=PipelineConfig(),
|
||||
)
|
||||
defaults.update(overrides)
|
||||
return SimpleNamespace(**defaults)
|
||||
|
||||
|
||||
def _patch_from_fastvideo_args(monkeypatch):
|
||||
captured = {}
|
||||
|
||||
@@ -151,6 +173,117 @@ def test_prepare_output_path_empty_prompt_fallback(tmp_path):
|
||||
assert os.path.basename(result) == "output.mp4"
|
||||
|
||||
|
||||
def test_generate_prepared_work_items_merges_compatible_latent_requests(monkeypatch, tmp_path):
|
||||
vg = _new_video_generator()
|
||||
vg.fastvideo_args = _batching_fastvideo_args()
|
||||
calls = []
|
||||
|
||||
def fake_device_memory(gpu_id):
|
||||
return 48.0
|
||||
|
||||
def fake_run_forward(batch, fastvideo_args):
|
||||
calls.append(batch)
|
||||
batch_size = len(batch.prompt) if isinstance(batch.prompt, list) else 1
|
||||
output = torch.arange(batch_size * 4, dtype=torch.float32).reshape(batch_size, 4, 1, 1, 1)
|
||||
return ForwardBatch(data_type=batch.data_type, output=output, extra={"peak_memory_mb": 1.0}), 0.5, 10.0
|
||||
|
||||
monkeypatch.setattr(
|
||||
"fastvideo.batching.admission.BatchAdmissionController._get_device_memory_gb",
|
||||
staticmethod(fake_device_memory),
|
||||
)
|
||||
monkeypatch.setattr(vg, "_run_forward_batch", fake_run_forward)
|
||||
|
||||
first = SamplingParam(prompt="one", height=8, width=8, num_frames=1, seed=11, return_frames=True, save_video=False)
|
||||
second = SamplingParam(prompt="two", height=8, width=8, num_frames=1, seed=22, return_frames=True, save_video=False)
|
||||
work_items = [
|
||||
vg._prepare_generation_work_item("one", first, vg.fastvideo_args, output_path=str(tmp_path / "one.mp4")),
|
||||
vg._prepare_generation_work_item("two", second, vg.fastvideo_args, output_path=str(tmp_path / "two.mp4")),
|
||||
]
|
||||
|
||||
results = vg._generate_prepared_work_items(work_items)
|
||||
|
||||
assert len(calls) == 1
|
||||
assert calls[0].prompt == ["one", "two"]
|
||||
assert calls[0].seeds == [11, 22]
|
||||
assert [result["prompts"] for result in results] == ["one", "two"]
|
||||
assert [result["samples"].shape for result in results] == [(1, 4, 1, 1, 1), (1, 4, 1, 1, 1)]
|
||||
|
||||
|
||||
def test_generate_prepared_work_items_falls_back_for_incompatible_requests(monkeypatch, tmp_path):
|
||||
vg = _new_video_generator()
|
||||
vg.fastvideo_args = _batching_fastvideo_args()
|
||||
calls = []
|
||||
|
||||
def fake_run_forward(batch, fastvideo_args):
|
||||
calls.append(batch)
|
||||
output = torch.zeros((1, 4, 1, 1, 1), dtype=torch.float32)
|
||||
return ForwardBatch(data_type=batch.data_type, output=output), 0.5, 10.0
|
||||
|
||||
monkeypatch.setattr(vg, "_run_forward_batch", fake_run_forward)
|
||||
|
||||
first = SamplingParam(prompt="one", height=8, width=8, num_frames=1, guidance_scale=1.0, save_video=False)
|
||||
second = SamplingParam(prompt="two", height=8, width=8, num_frames=1, guidance_scale=3.0, save_video=False)
|
||||
work_items = [
|
||||
vg._prepare_generation_work_item("one", first, vg.fastvideo_args, output_path=str(tmp_path / "one.mp4")),
|
||||
vg._prepare_generation_work_item("two", second, vg.fastvideo_args, output_path=str(tmp_path / "two.mp4")),
|
||||
]
|
||||
|
||||
results = vg._generate_prepared_work_items(work_items)
|
||||
|
||||
assert len(calls) == 2
|
||||
assert all(isinstance(call.prompt, str) for call in calls)
|
||||
assert [result["prompts"] for result in results] == ["one", "two"]
|
||||
|
||||
|
||||
def test_generate_video_batch_routes_compat_kwargs(monkeypatch, tmp_path):
|
||||
vg = _new_video_generator()
|
||||
vg.fastvideo_args = _batching_fastvideo_args()
|
||||
calls = []
|
||||
|
||||
def fake_device_memory(gpu_id):
|
||||
return 48.0
|
||||
|
||||
def fake_run_forward(batch, fastvideo_args):
|
||||
calls.append((batch, fastvideo_args))
|
||||
output = torch.zeros((len(batch.prompt), 4, 1, 1, 1), dtype=torch.float32)
|
||||
return ForwardBatch(data_type=batch.data_type, output=output), 0.5, 10.0
|
||||
|
||||
monkeypatch.setattr(
|
||||
"fastvideo.batching.admission.BatchAdmissionController._get_device_memory_gb",
|
||||
staticmethod(fake_device_memory),
|
||||
)
|
||||
monkeypatch.setattr(vg, "_run_forward_batch", fake_run_forward)
|
||||
|
||||
results = vg.generate_video_batch([
|
||||
{
|
||||
"prompt": "one",
|
||||
"height": 8,
|
||||
"width": 8,
|
||||
"num_frames": 1,
|
||||
"embedded_cfg_scale": 7.5,
|
||||
"save_video": False,
|
||||
"return_frames": True,
|
||||
"output_path": str(tmp_path / "one.mp4"),
|
||||
},
|
||||
{
|
||||
"prompt": "two",
|
||||
"height": 8,
|
||||
"width": 8,
|
||||
"num_frames": 1,
|
||||
"embedded_cfg_scale": 7.5,
|
||||
"save_video": False,
|
||||
"return_frames": True,
|
||||
"output_path": str(tmp_path / "two.mp4"),
|
||||
},
|
||||
])
|
||||
|
||||
assert len(calls) == 1
|
||||
batch, fastvideo_args = calls[0]
|
||||
assert batch.prompt == ["one", "two"]
|
||||
assert fastvideo_args.pipeline_config.embedded_cfg_scale == 7.5
|
||||
assert [result["prompts"] for result in results] == ["one", "two"]
|
||||
|
||||
|
||||
def test_from_config_normalizes_and_translates(monkeypatch):
|
||||
captured = _patch_from_fastvideo_args(monkeypatch)
|
||||
_patch_fastvideo_args_from_kwargs(monkeypatch)
|
||||
|
||||
@@ -0,0 +1,203 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import subprocess
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
from torch.distributed import init_device_mesh
|
||||
from torch.distributed.fsdp import CPUOffloadPolicy, fully_shard
|
||||
from torch.distributed.tensor import DTensor
|
||||
|
||||
from fastvideo.layers.layernorm import RMSNorm
|
||||
|
||||
WORLD_SIZE = 2
|
||||
HIDDEN_SIZE = 8
|
||||
SEED = 1379
|
||||
REPO_ROOT = Path(__file__).resolve().parents[3]
|
||||
|
||||
|
||||
def _run_torchrun(script_path: Path, mode: str, output_path: Path) -> None:
|
||||
# --standalone binds the rendezvous port atomically, avoiding the
|
||||
# free-port-probe race a hand-picked --master_port would have.
|
||||
cmd = [
|
||||
"torchrun",
|
||||
"--standalone",
|
||||
"--nproc_per_node",
|
||||
str(WORLD_SIZE),
|
||||
str(script_path),
|
||||
"--rmsnorm-fsdp-worker",
|
||||
"--mode",
|
||||
mode,
|
||||
"--output",
|
||||
str(output_path),
|
||||
]
|
||||
env = os.environ.copy()
|
||||
env.setdefault("TORCHDYNAMO_DISABLE", "1")
|
||||
try:
|
||||
process = subprocess.run(
|
||||
cmd,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
env=env,
|
||||
timeout=120,
|
||||
)
|
||||
except subprocess.TimeoutExpired as error:
|
||||
raise RuntimeError(
|
||||
f"{mode} worker timed out after 120 seconds\n"
|
||||
f"STDOUT:\n{error.stdout}\n"
|
||||
f"STDERR:\n{error.stderr}"
|
||||
) from error
|
||||
if process.returncode != 0:
|
||||
raise RuntimeError(
|
||||
f"{mode} worker failed with code {process.returncode}\n"
|
||||
f"STDOUT:\n{process.stdout}\n"
|
||||
f"STDERR:\n{process.stderr}"
|
||||
)
|
||||
|
||||
|
||||
def _summarize_tensor(tensor: torch.Tensor | Any) -> dict[str, Any]:
|
||||
return {
|
||||
"type": type(tensor).__name__,
|
||||
"is_dtensor": isinstance(tensor, DTensor),
|
||||
"shape": list(tensor.shape) if hasattr(tensor, "shape") else None,
|
||||
"device": str(tensor.device) if hasattr(tensor, "device") else None,
|
||||
"dtype": str(tensor.dtype) if hasattr(tensor, "dtype") else None,
|
||||
}
|
||||
|
||||
|
||||
def _run_worker(mode: str, output_path: Path) -> None:
|
||||
if mode not in {
|
||||
"module_no_offload",
|
||||
"direct_no_offload",
|
||||
"module_cpu_offload",
|
||||
"direct_cpu_offload",
|
||||
}:
|
||||
raise ValueError(f"Unsupported mode: {mode}")
|
||||
|
||||
dist.init_process_group("nccl")
|
||||
rank = dist.get_rank()
|
||||
world_size = dist.get_world_size()
|
||||
local_rank = int(os.environ.get("LOCAL_RANK", "0"))
|
||||
device = torch.device(f"cuda:{local_rank}")
|
||||
torch.cuda.set_device(device)
|
||||
torch.manual_seed(SEED + rank)
|
||||
|
||||
try:
|
||||
mesh = init_device_mesh("cuda", (world_size,))
|
||||
norm = RMSNorm(HIDDEN_SIZE, eps=1e-6, has_weight=True).to(device)
|
||||
with torch.no_grad():
|
||||
norm.weight.fill_(1.0)
|
||||
|
||||
fsdp_kwargs: dict[str, Any] = {"mesh": mesh}
|
||||
if mode.endswith("cpu_offload"):
|
||||
fsdp_kwargs["offload_policy"] = CPUOffloadPolicy(pin_memory=False)
|
||||
# fully_shard is applied to the bare RMSNorm to make the hook bypass
|
||||
# observable. Production sharding (fsdp_load.shard_model) only wraps
|
||||
# whole transformer blocks, whose pre-forward all-gather localizes norm
|
||||
# weights before the qk-norm call sites run, so this pins the dispatch
|
||||
# invariant rather than reproducing a production topology.
|
||||
fully_shard(norm, **fsdp_kwargs)
|
||||
|
||||
x = torch.randn(2, 3, HIDDEN_SIZE, device=device, dtype=torch.bfloat16)
|
||||
call_kind = "direct" if mode.startswith("direct") else "module"
|
||||
|
||||
try:
|
||||
if call_kind == "direct":
|
||||
output = norm.forward_native(x)
|
||||
else:
|
||||
output = norm(x)
|
||||
torch.cuda.synchronize(device)
|
||||
result = {
|
||||
"rank": rank,
|
||||
"ok": True,
|
||||
"mode": mode,
|
||||
"weight": _summarize_tensor(norm.weight),
|
||||
"output": _summarize_tensor(output),
|
||||
}
|
||||
except Exception as exc:
|
||||
result = {
|
||||
"rank": rank,
|
||||
"ok": False,
|
||||
"mode": mode,
|
||||
"error_type": type(exc).__name__,
|
||||
"error": str(exc),
|
||||
"weight": _summarize_tensor(norm.weight),
|
||||
}
|
||||
|
||||
gathered = [None for _ in range(world_size)] if rank == 0 else None
|
||||
dist.gather_object(result, object_gather_list=gathered, dst=0)
|
||||
if rank == 0:
|
||||
output_path.write_text(json.dumps(gathered, indent=2), encoding="utf-8")
|
||||
dist.barrier()
|
||||
finally:
|
||||
dist.destroy_process_group()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("mode", "expect_ok"),
|
||||
[
|
||||
("module_no_offload", True),
|
||||
("direct_no_offload", False),
|
||||
("module_cpu_offload", True),
|
||||
("direct_cpu_offload", False),
|
||||
],
|
||||
)
|
||||
def test_rmsnorm_forward_native_bypasses_fsdp_hooks(mode: str, expect_ok: bool, tmp_path: Path) -> None:
|
||||
if not torch.cuda.is_available():
|
||||
pytest.skip("This test requires CUDA.")
|
||||
if torch.cuda.device_count() < WORLD_SIZE:
|
||||
pytest.skip(f"This test requires at least {WORLD_SIZE} CUDA devices.")
|
||||
|
||||
output_path = tmp_path / f"{mode}.json"
|
||||
_run_torchrun(Path(__file__).resolve(), mode, output_path)
|
||||
results = json.loads(output_path.read_text(encoding="utf-8"))
|
||||
print(f"\n{mode} results:\n{json.dumps(results, indent=2)}")
|
||||
|
||||
if expect_ok:
|
||||
failures = [result for result in results if not result["ok"]]
|
||||
assert not failures, json.dumps(results, indent=2)
|
||||
return
|
||||
|
||||
successes = [result for result in results if result["ok"]]
|
||||
assert not successes, json.dumps(results, indent=2)
|
||||
error_text = "\n".join(result.get("error", "") for result in results)
|
||||
# Pin the specific bypassed-hook failure: "got mixed torch.Tensor and
|
||||
# DTensor" ("Tensor" alone is a substring of "DTensor", so it adds nothing).
|
||||
assert "mixed" in error_text and "DTensor" in error_text, json.dumps(results, indent=2)
|
||||
|
||||
|
||||
def test_no_direct_forward_native_calls_in_models() -> None:
|
||||
"""Direct .forward_native(...) calls bypass nn.Module.__call__ and FSDP
|
||||
hooks (issue #1379); model code must use module dispatch instead."""
|
||||
models_dir = REPO_ROOT / "fastvideo" / "models"
|
||||
offenders = [
|
||||
str(path.relative_to(REPO_ROOT))
|
||||
for path in sorted(models_dir.rglob("*.py"))
|
||||
if ".forward_native(" in path.read_text(encoding="utf-8")
|
||||
]
|
||||
assert not offenders, f"Replace .forward_native(...) with module dispatch in: {offenders}"
|
||||
|
||||
|
||||
def _parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--rmsnorm-fsdp-worker", action="store_true")
|
||||
parser.add_argument("--mode", type=str, default=None)
|
||||
parser.add_argument("--output", type=str, default=None)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = _parse_args()
|
||||
if not args.rmsnorm_fsdp_worker:
|
||||
raise SystemExit("This module is intended to be run by pytest.")
|
||||
if args.mode is None or args.output is None:
|
||||
raise SystemExit("--mode and --output are required in worker mode.")
|
||||
_run_worker(mode=args.mode, output_path=Path(args.output))
|
||||
@@ -32,7 +32,8 @@ import modal
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||||
try:
|
||||
from modal_image_utils import resolve_image_ref # noqa: E402
|
||||
from modal_image_utils import ( # noqa: E402
|
||||
resolve_image_ref, resolve_uv_torch_backend)
|
||||
except ModuleNotFoundError:
|
||||
# Remote Modal containers re-import this module but mount only the
|
||||
# entrypoint file; the digest resolution already happened at local
|
||||
@@ -40,6 +41,9 @@ except ModuleNotFoundError:
|
||||
def resolve_image_ref(image_ref: str) -> str:
|
||||
return image_ref
|
||||
|
||||
def resolve_uv_torch_backend(image_tag: str) -> str | None:
|
||||
return os.environ.get("UV_TORCH_BACKEND")
|
||||
|
||||
app = modal.App("fastvideo-gpu-job")
|
||||
|
||||
REPO_DIR = "/FastVideo"
|
||||
@@ -72,12 +76,7 @@ local_secrets = modal.Secret.from_dict({
|
||||
# Mutable tags inherit the registry image's baked backend, including custom
|
||||
# FASTVIDEO_MODAL_IMAGE overrides. Explicit CUDA tags also work with older
|
||||
# images that predate the baked setting, and a caller override always wins.
|
||||
uv_torch_backend_override = os.environ.get("UV_TORCH_BACKEND")
|
||||
if not uv_torch_backend_override:
|
||||
if "cuda13" in IMAGE_TAG.lower():
|
||||
uv_torch_backend_override = "cu130"
|
||||
elif "cuda12.6" in IMAGE_TAG.lower():
|
||||
uv_torch_backend_override = "cu126"
|
||||
uv_torch_backend_override = resolve_uv_torch_backend(IMAGE_TAG)
|
||||
|
||||
image = (
|
||||
modal.Image.from_registry(IMAGE_REF, add_python="3.12")
|
||||
@@ -98,6 +97,9 @@ image = (
|
||||
"TOKENIZERS_PARALLELISM": "false",
|
||||
**({"UV_TORCH_BACKEND": uv_torch_backend_override} if uv_torch_backend_override else {}),
|
||||
"FASTVIDEO_ATTENTION_BACKEND": os.environ.get("FASTVIDEO_ATTENTION_BACKEND", "FLASH_ATTN"),
|
||||
# FA4 is opt-in (FASTVIDEO_FA4); keep CI parity with the seeded
|
||||
# references. Caller override wins.
|
||||
"FASTVIDEO_FA4": os.environ.get("FASTVIDEO_FA4", "1"),
|
||||
})
|
||||
)
|
||||
|
||||
|
||||
@@ -5,6 +5,7 @@ through the ``fastvideo`` package, so they keep working on thin CI hosts that
|
||||
have ``modal`` but not torch.
|
||||
"""
|
||||
import json
|
||||
import os
|
||||
import urllib.request
|
||||
|
||||
_REGISTRY = "ghcr.io"
|
||||
@@ -55,3 +56,23 @@ def resolve_image_ref(image_ref: str) -> str:
|
||||
print(f"WARNING: could not resolve {image_ref} to a digest ({error}); "
|
||||
"Modal may reuse a stale cached image for this tag.")
|
||||
return image_ref
|
||||
|
||||
|
||||
def resolve_uv_torch_backend(image_tag: str) -> str | None:
|
||||
"""UV_TORCH_BACKEND for a launcher image tag.
|
||||
|
||||
A caller-set UV_TORCH_BACKEND always wins. Otherwise sniff the CUDA
|
||||
version from an explicit image tag (cuda13 -> cu130, cuda12.6 -> cu126)
|
||||
so uv resolves torch against the image's toolkit. Mutable tags (e.g.
|
||||
py3.12-latest) return None and inherit the registry image's baked
|
||||
backend, which keeps a latest-tag CUDA transition safe.
|
||||
"""
|
||||
override = os.environ.get("UV_TORCH_BACKEND")
|
||||
if override:
|
||||
return override
|
||||
tag = image_tag.lower()
|
||||
if "cuda13" in tag:
|
||||
return "cu130"
|
||||
if "cuda12.6" in tag:
|
||||
return "cu126"
|
||||
return None
|
||||
|
||||
@@ -5,7 +5,8 @@ import modal
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||||
try:
|
||||
from modal_image_utils import resolve_image_ref # noqa: E402
|
||||
from modal_image_utils import ( # noqa: E402
|
||||
resolve_image_ref, resolve_uv_torch_backend)
|
||||
except ModuleNotFoundError:
|
||||
# Remote Modal containers re-import this module but mount only the
|
||||
# entrypoint file; the digest resolution already happened at local
|
||||
@@ -13,6 +14,9 @@ except ModuleNotFoundError:
|
||||
def resolve_image_ref(image_ref: str) -> str:
|
||||
return image_ref
|
||||
|
||||
def resolve_uv_torch_backend(image_tag: str) -> str | None:
|
||||
return os.environ.get("UV_TORCH_BACKEND")
|
||||
|
||||
app = modal.App()
|
||||
|
||||
model_vol = modal.Volume.from_name("hf-model-weights")
|
||||
@@ -24,12 +28,7 @@ print(f"Using image: {image_ref}")
|
||||
# Mutable tags inherit the registry image's baked backend, keeping a latest-tag
|
||||
# transition safe. Explicit CUDA tags also work with older images that predate
|
||||
# the baked setting, and a caller override always wins.
|
||||
uv_torch_backend_override = os.environ.get("UV_TORCH_BACKEND")
|
||||
if not uv_torch_backend_override:
|
||||
if "cuda13" in image_tag.lower():
|
||||
uv_torch_backend_override = "cu130"
|
||||
elif "cuda12.6" in image_tag.lower():
|
||||
uv_torch_backend_override = "cu126"
|
||||
uv_torch_backend_override = resolve_uv_torch_backend(image_tag)
|
||||
|
||||
image = (modal.Image.from_registry(
|
||||
image_ref, add_python="3.12"
|
||||
@@ -61,6 +60,10 @@ image = (modal.Image.from_registry(
|
||||
**({
|
||||
"UV_TORCH_BACKEND": uv_torch_backend_override
|
||||
} if uv_torch_backend_override else {}),
|
||||
# FA4 is opt-in (FASTVIDEO_FA4); CI lanes keep it enabled to match the
|
||||
# SSIM/perf baselines. Caller override wins.
|
||||
"FASTVIDEO_FA4":
|
||||
os.environ.get("FASTVIDEO_FA4", "1"),
|
||||
"HF_REPO_ID":
|
||||
"FastVideo/performance-tracking",
|
||||
}))
|
||||
@@ -273,7 +276,7 @@ def run_self_forcing_tests():
|
||||
@app.function(gpu="L40S:1", image=image, timeout=900)
|
||||
def run_unit_test():
|
||||
run_test(
|
||||
"pytest ./fastvideo/tests/api/ ./fastvideo/tests/contract/ ./fastvideo/tests/dataset/ ./fastvideo/tests/workflow/ ./fastvideo/tests/entrypoints/ ./fastvideo/tests/train/ --ignore=./fastvideo/tests/entrypoints/test_openai_api_integration.py --ignore=./fastvideo/tests/train/models --ignore=./fastvideo/tests/train/methods -vs"
|
||||
"pytest ./fastvideo/tests/api/ ./fastvideo/tests/batching/ ./fastvideo/tests/contract/ ./fastvideo/tests/dataset/ ./fastvideo/tests/workflow/ ./fastvideo/tests/entrypoints/ ./fastvideo/tests/train/ ./fastvideo/tests/stages/ --ignore=./fastvideo/tests/entrypoints/test_openai_api_integration.py --ignore=./fastvideo/tests/train/models --ignore=./fastvideo/tests/train/methods -vs"
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -13,7 +13,8 @@ import modal
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||||
try:
|
||||
from modal_image_utils import resolve_image_ref # noqa: E402
|
||||
from modal_image_utils import ( # noqa: E402
|
||||
resolve_image_ref, resolve_uv_torch_backend)
|
||||
except ModuleNotFoundError:
|
||||
# Remote Modal containers re-import this module but mount only the
|
||||
# entrypoint file; the digest resolution already happened at local
|
||||
@@ -21,6 +22,9 @@ except ModuleNotFoundError:
|
||||
def resolve_image_ref(image_ref: str) -> str:
|
||||
return image_ref
|
||||
|
||||
def resolve_uv_torch_backend(image_tag: str) -> str | None:
|
||||
return os.environ.get("UV_TORCH_BACKEND")
|
||||
|
||||
app = modal.App()
|
||||
|
||||
model_vol = modal.Volume.from_name("hf-model-weights")
|
||||
@@ -32,12 +36,7 @@ print(f"Using image: {image_ref}")
|
||||
# Mutable tags inherit the registry image's baked backend, keeping a latest-tag
|
||||
# transition safe. Explicit CUDA tags also work with older images that predate
|
||||
# the baked setting, and a caller override always wins.
|
||||
uv_torch_backend_override = os.environ.get("UV_TORCH_BACKEND")
|
||||
if not uv_torch_backend_override:
|
||||
if "cuda13" in image_tag.lower():
|
||||
uv_torch_backend_override = "cu130"
|
||||
elif "cuda12.6" in image_tag.lower():
|
||||
uv_torch_backend_override = "cu126"
|
||||
uv_torch_backend_override = resolve_uv_torch_backend(image_tag)
|
||||
|
||||
image = (
|
||||
modal.Image.from_registry(image_ref, add_python="3.12")
|
||||
@@ -64,6 +63,9 @@ image = (
|
||||
"BUILDKITE_PULL_REQUEST": os.environ.get("BUILDKITE_PULL_REQUEST", ""),
|
||||
"IMAGE_VERSION": image_version,
|
||||
**({"UV_TORCH_BACKEND": uv_torch_backend_override} if uv_torch_backend_override else {}),
|
||||
# FA4 is opt-in (FASTVIDEO_FA4); the SSIM references were seeded
|
||||
# with FA4 inference, so keep it enabled in CI. Caller override wins.
|
||||
"FASTVIDEO_FA4": os.environ.get("FASTVIDEO_FA4", "1"),
|
||||
}
|
||||
)
|
||||
)
|
||||
|
||||
@@ -8,8 +8,8 @@ This script:
|
||||
baseline-eligible successful records (filtered by gpu_type),
|
||||
4) writes normalized records back to the HF dataset repo according to
|
||||
PERF_UPLOAD_POLICY,
|
||||
5) exits non-zero if any metric regresses by more than PERF_MAX_REGRESSION
|
||||
(default 5%).
|
||||
5) exits non-zero if any gated metric exceeds both its percent and absolute
|
||||
regression floors.
|
||||
"""
|
||||
|
||||
import glob
|
||||
@@ -21,21 +21,36 @@ from datetime import datetime, timezone
|
||||
from typing import Any
|
||||
|
||||
try:
|
||||
from .hf_store import (
|
||||
from fastvideo.performance.hf_store import (
|
||||
load_records_for_model,
|
||||
safe_float,
|
||||
sanitize,
|
||||
sync_from_hf,
|
||||
upload_record,
|
||||
)
|
||||
from fastvideo.performance.metric_policy import (
|
||||
MetricPolicy,
|
||||
regression_delta,
|
||||
resolve_metric_policies,
|
||||
serialize_metric_thresholds,
|
||||
)
|
||||
except ImportError:
|
||||
from hf_store import (
|
||||
repo_root = os.path.abspath(os.path.join(os.path.dirname(__file__), "../../.."))
|
||||
if repo_root not in sys.path:
|
||||
sys.path.insert(0, repo_root)
|
||||
from fastvideo.performance.hf_store import (
|
||||
load_records_for_model,
|
||||
safe_float,
|
||||
sanitize,
|
||||
sync_from_hf,
|
||||
upload_record,
|
||||
)
|
||||
from fastvideo.performance.metric_policy import (
|
||||
MetricPolicy,
|
||||
regression_delta,
|
||||
resolve_metric_policies,
|
||||
serialize_metric_thresholds,
|
||||
)
|
||||
|
||||
RESULTS_DIR = os.path.join(
|
||||
os.path.dirname(os.path.abspath(__file__)),
|
||||
@@ -46,25 +61,9 @@ TRACKING_ROOT = os.environ.get(
|
||||
"/tmp/perf-tracking",
|
||||
)
|
||||
PERF_REPORTS_DIR = os.environ.get("PERF_REPORTS_DIR", "/root/data/perf_reports")
|
||||
MAX_REGRESSION = float(os.environ.get("PERF_MAX_REGRESSION", "0.05"))
|
||||
UPLOAD_POLICY = os.environ.get("PERF_UPLOAD_POLICY", "never").strip().lower()
|
||||
VALID_UPLOAD_POLICIES = {"never", "pass", "always"}
|
||||
VALID_RUN_SOURCES = {"pr", "local", "scheduled_main", "unknown"}
|
||||
METRICS = (
|
||||
("latency", "Latency", 3),
|
||||
("throughput", "Throughput", 3),
|
||||
("memory", "Memory", 1),
|
||||
("text_encoder_time_s", "Text Enc", 3),
|
||||
("dit_time_s", "DiT", 3),
|
||||
("vae_decode_time_s", "VAE Decode", 3),
|
||||
)
|
||||
LOWER_IS_BETTER_METRICS = {
|
||||
"latency",
|
||||
"memory",
|
||||
"text_encoder_time_s",
|
||||
"dit_time_s",
|
||||
"vae_decode_time_s",
|
||||
}
|
||||
|
||||
|
||||
def _should_persist_tracking() -> bool:
|
||||
@@ -168,6 +167,7 @@ def normalize_performance_result(result: dict[str, Any]) -> dict[str, Any]:
|
||||
text_encoder_time = safe_float(result.get("text_encoder_time_s"))
|
||||
dit_time = safe_float(result.get("dit_time_s"))
|
||||
vae_decode_time = safe_float(result.get("vae_decode_time_s"))
|
||||
metric_policies = resolve_metric_policies(result.get("regression_thresholds"))
|
||||
|
||||
return {
|
||||
"model_id": model_id,
|
||||
@@ -180,6 +180,7 @@ def normalize_performance_result(result: dict[str, Any]) -> dict[str, Any]:
|
||||
"text_encoder_time_s": text_encoder_time,
|
||||
"dit_time_s": dit_time,
|
||||
"vae_decode_time_s": vae_decode_time,
|
||||
"regression_thresholds": serialize_metric_thresholds(metric_policies),
|
||||
"success": True,
|
||||
**_record_metadata(_detect_run_source(), result),
|
||||
}
|
||||
@@ -229,55 +230,41 @@ def _baseline_metric(records: list[dict[str, Any]], key: str) -> float | None:
|
||||
return statistics.median(values)
|
||||
|
||||
|
||||
def _metric_policy_summary(policy: MetricPolicy) -> str:
|
||||
gated = "gated" if policy.gated else "info"
|
||||
return (
|
||||
f"{gated}, >{policy.threshold_percent * 100:.1f}% "
|
||||
f"and >{policy.threshold_absolute:.{policy.precision}f}"
|
||||
)
|
||||
|
||||
|
||||
def _check_regressions(
|
||||
current: dict[str, Any],
|
||||
baseline_records: list[dict[str, Any]],
|
||||
max_regression: float,
|
||||
metric_policies: tuple[MetricPolicy, ...],
|
||||
) -> list[str]:
|
||||
failures: list[str] = []
|
||||
|
||||
for metric, _label, _precision in METRICS:
|
||||
if metric not in LOWER_IS_BETTER_METRICS:
|
||||
for policy in metric_policies:
|
||||
baseline = _baseline_metric(baseline_records, policy.key)
|
||||
curr = safe_float(current.get(policy.key))
|
||||
if baseline is None or curr is None:
|
||||
continue
|
||||
baseline = _baseline_metric(baseline_records, metric)
|
||||
curr = safe_float(current.get(metric))
|
||||
if baseline is None or curr is None or baseline <= 0:
|
||||
delta = regression_delta(policy, curr, baseline)
|
||||
if delta is None or not delta.regressed:
|
||||
continue
|
||||
regression = (curr - baseline) / baseline
|
||||
if regression > max_regression:
|
||||
failures.append(f"{current['model_id']} {metric} regressed by "
|
||||
f"{regression * 100:.1f}% "
|
||||
f"(current={curr:.3f}, baseline_median={baseline:.3f})")
|
||||
|
||||
baseline_tp = _baseline_metric(baseline_records, "throughput")
|
||||
curr_tp = safe_float(current.get("throughput"))
|
||||
if baseline_tp is not None and curr_tp is not None and baseline_tp > 0:
|
||||
regression = (baseline_tp - curr_tp) / baseline_tp
|
||||
if regression > max_regression:
|
||||
failures.append(f"{current['model_id']} throughput regressed by "
|
||||
f"{regression * 100:.1f}% "
|
||||
f"(current={curr_tp:.3f}, baseline_median={baseline_tp:.3f})")
|
||||
failures.append(
|
||||
f"{current['model_id']} {policy.key} regressed by "
|
||||
f"{delta.percent * 100:.1f}% and "
|
||||
f"{delta.absolute:.{policy.precision}f} "
|
||||
f"(current={curr:.{policy.precision}f}, "
|
||||
f"baseline_median={baseline:.{policy.precision}f}, "
|
||||
f"threshold={_metric_policy_summary(policy)})"
|
||||
)
|
||||
|
||||
return failures
|
||||
|
||||
|
||||
def _metric_delta_percent(
|
||||
metric: str,
|
||||
current: dict[str, Any],
|
||||
baseline_records: list[dict[str, Any]],
|
||||
) -> float | None:
|
||||
curr = safe_float(current.get(metric))
|
||||
baseline = _baseline_metric(baseline_records, metric)
|
||||
if curr is None or baseline is None or baseline <= 0:
|
||||
return None
|
||||
|
||||
if metric in LOWER_IS_BETTER_METRICS:
|
||||
return (curr - baseline) / baseline * 100.0
|
||||
if metric == "throughput":
|
||||
return (baseline - curr) / baseline * 100.0
|
||||
return None
|
||||
|
||||
|
||||
def _compact_value(value: float | None, precision: int = 3) -> str:
|
||||
if value is None:
|
||||
return "n/a"
|
||||
@@ -287,23 +274,42 @@ def _compact_value(value: float | None, precision: int = 3) -> str:
|
||||
def _build_summary_row(
|
||||
record: dict[str, Any],
|
||||
baseline_records: list[dict[str, Any]],
|
||||
metric_policies: tuple[MetricPolicy, ...],
|
||||
has_failed: bool,
|
||||
) -> dict[str, Any]:
|
||||
"""Format a single benchmark result as a row for the Markdown table."""
|
||||
|
||||
metric_values: dict[str, dict[str, float | None]] = {}
|
||||
metric_values: dict[str, dict[str, Any]] = {}
|
||||
regressions: list[float] = []
|
||||
for metric, _label, _precision in METRICS:
|
||||
curr = safe_float(record.get(metric))
|
||||
baseline = _baseline_metric(baseline_records, metric)
|
||||
regression = _metric_delta_percent(metric, record, baseline_records)
|
||||
metric_values[metric] = {
|
||||
failing_metrics: list[str] = []
|
||||
threshold_exceeded_metrics: list[str] = []
|
||||
for policy in metric_policies:
|
||||
curr = safe_float(record.get(policy.key))
|
||||
baseline = _baseline_metric(baseline_records, policy.key)
|
||||
delta = (
|
||||
regression_delta(policy, curr, baseline)
|
||||
if curr is not None and baseline is not None
|
||||
else None
|
||||
)
|
||||
regression = None if delta is None else delta.percent * 100.0
|
||||
absolute_delta = None if delta is None else delta.absolute
|
||||
metric_values[policy.key] = {
|
||||
"curr": curr,
|
||||
"base": baseline,
|
||||
"regression_pct": regression,
|
||||
"absolute_delta": absolute_delta,
|
||||
"threshold_percent": policy.threshold_percent * 100.0,
|
||||
"threshold_absolute": policy.threshold_absolute,
|
||||
"gated": policy.gated,
|
||||
"threshold_exceeded": False if delta is None else delta.threshold_exceeded,
|
||||
"regressed": False if delta is None else delta.regressed,
|
||||
}
|
||||
if regression is not None:
|
||||
regressions.append(regression)
|
||||
if delta is not None and delta.threshold_exceeded:
|
||||
threshold_exceeded_metrics.append(policy.key)
|
||||
if delta is not None and delta.regressed:
|
||||
failing_metrics.append(policy.key)
|
||||
|
||||
worst_regression_pct = max(regressions) if regressions else None
|
||||
|
||||
@@ -313,40 +319,50 @@ def _build_summary_row(
|
||||
"baseline_n": len(baseline_records),
|
||||
"metrics": metric_values,
|
||||
"worst_regression_pct": worst_regression_pct,
|
||||
"threshold_exceeded_metrics": threshold_exceeded_metrics,
|
||||
"failing_metrics": failing_metrics,
|
||||
"failed": has_failed,
|
||||
}
|
||||
|
||||
|
||||
def _build_markdown_summary(
|
||||
summary_rows: list[dict[str, Any]],
|
||||
max_regression: float,
|
||||
metric_policies: tuple[MetricPolicy, ...],
|
||||
) -> str:
|
||||
lines = [
|
||||
"## Performance Baseline Comparison",
|
||||
"",
|
||||
f"Threshold: regressions greater than {max_regression * 100:.1f}% fail",
|
||||
"Threshold: gated metrics fail only when both percent and absolute "
|
||||
"regression floors are exceeded.",
|
||||
"",
|
||||
("| Model | GPU | Baseline N | Latency (curr/base) | "
|
||||
"Throughput (curr/base) | Memory (curr/base) | "
|
||||
"Text Enc (curr/base) | DiT (curr/base) | "
|
||||
"VAE Decode (curr/base) | Worst Regression | Status |"),
|
||||
"|---|---|---:|---|---|---|---|---|---|---:|---|",
|
||||
"VAE Decode (curr/base) | Worst Regression | Exceeded Metrics | "
|
||||
"Failing Metrics | Status |"),
|
||||
"|---|---|---:|---|---|---|---|---|---|---:|---|---|---|",
|
||||
]
|
||||
|
||||
for row in summary_rows:
|
||||
metric_cells = []
|
||||
for metric, _label, precision in METRICS:
|
||||
values = row["metrics"][metric]
|
||||
metric_cells.append(f"{_compact_value(values['curr'], precision)} / "
|
||||
f"{_compact_value(values['base'], precision)}")
|
||||
for policy in metric_policies:
|
||||
values = row["metrics"][policy.key]
|
||||
metric_cells.append(f"{_compact_value(values['curr'], policy.precision)} / "
|
||||
f"{_compact_value(values['base'], policy.precision)}")
|
||||
|
||||
worst_reg = ("n/a" if row["worst_regression_pct"] is None else f"{row['worst_regression_pct']:.1f}%")
|
||||
exceeded_metrics = (
|
||||
", ".join(row["threshold_exceeded_metrics"])
|
||||
if row["threshold_exceeded_metrics"]
|
||||
else "none"
|
||||
)
|
||||
failing_metrics = ", ".join(row["failing_metrics"]) if row["failing_metrics"] else "none"
|
||||
status = "FAIL" if row["failed"] else "PASS"
|
||||
|
||||
lines.append(f"| {row['model_id']} | {row['gpu_type']} | "
|
||||
f"{row['baseline_n']} | "
|
||||
f"{' | '.join(metric_cells)} | "
|
||||
f"{worst_reg} | {status} |")
|
||||
f"{worst_reg} | {exceeded_metrics} | {failing_metrics} | {status} |")
|
||||
|
||||
return "\n".join(lines) + "\n"
|
||||
|
||||
@@ -400,6 +416,7 @@ def main() -> int:
|
||||
|
||||
for raw in current_results:
|
||||
record = _normalize_record(raw)
|
||||
metric_policies = resolve_metric_policies(record.get("regression_thresholds"))
|
||||
|
||||
baseline_records = load_records_for_model(
|
||||
TRACKING_ROOT,
|
||||
@@ -416,7 +433,7 @@ def main() -> int:
|
||||
failures: list[str] = []
|
||||
record["success"] = True
|
||||
else:
|
||||
failures = _check_regressions(record, baseline_records, MAX_REGRESSION)
|
||||
failures = _check_regressions(record, baseline_records, metric_policies)
|
||||
if static_threshold_failed:
|
||||
failures.append(f"{record['model_id']} fixed-threshold phase failed "
|
||||
f"(PERF_PYTEST_RC={os.environ.get('PERF_PYTEST_RC')})")
|
||||
@@ -434,10 +451,10 @@ def main() -> int:
|
||||
print("Tracking upload skipped for "
|
||||
f"{record['model_id']} ({record['run_source']}, success={record['success']})")
|
||||
|
||||
summary_rows.append(_build_summary_row(record, baseline_records, bool(failures)))
|
||||
summary_rows.append(_build_summary_row(record, baseline_records, metric_policies, bool(failures)))
|
||||
|
||||
commit_sha = os.environ.get("BUILDKITE_COMMIT", "unknown")[:7]
|
||||
markdown = _build_markdown_summary(summary_rows, MAX_REGRESSION)
|
||||
markdown = _build_markdown_summary(summary_rows, resolve_metric_policies(None))
|
||||
_emit_markdown_summary(markdown, commit_sha)
|
||||
|
||||
if all_failures:
|
||||
|
||||
@@ -1,12 +1,19 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
import os
|
||||
import sys
|
||||
from html import escape
|
||||
from datetime import datetime
|
||||
|
||||
import plotly.express as px
|
||||
import pandas as pd
|
||||
|
||||
from hf_store import sync_from_hf, load_as_dataframe
|
||||
try:
|
||||
from fastvideo.performance.hf_store import load_as_dataframe, sync_from_hf
|
||||
except ImportError:
|
||||
repo_root = os.path.abspath(os.path.join(os.path.dirname(__file__), "../../.."))
|
||||
if repo_root not in sys.path:
|
||||
sys.path.insert(0, repo_root)
|
||||
from fastvideo.performance.hf_store import load_as_dataframe, sync_from_hf
|
||||
|
||||
TRACKING_ROOT = os.environ.get(
|
||||
"PERFORMANCE_TRACKING_ROOT",
|
||||
|
||||
@@ -0,0 +1,116 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
import pytest
|
||||
|
||||
from fastvideo.tests.performance.test_inference_performance import (
|
||||
_benchmark_display_id,
|
||||
_config_identity_metadata,
|
||||
_is_v2_config,
|
||||
_validate_benchmark_config,
|
||||
)
|
||||
|
||||
|
||||
def _v2_config():
|
||||
return {
|
||||
"benchmark_id": "wan-t2v-1.3b-2gpu",
|
||||
"config_schema_version": 2,
|
||||
"workload_id": "wan-t2v-1.3b",
|
||||
"variant_id": "canonical",
|
||||
"benchmark_version": 1,
|
||||
}
|
||||
|
||||
|
||||
def test_v1_benchmark_config_without_schema_version_validates():
|
||||
cfg = {
|
||||
"benchmark_id": "legacy-benchmark",
|
||||
}
|
||||
|
||||
_validate_benchmark_config(cfg, "legacy.json")
|
||||
|
||||
assert _is_v2_config(cfg) is False
|
||||
assert _config_identity_metadata(cfg) == {}
|
||||
assert _benchmark_display_id(cfg) == "legacy-benchmark"
|
||||
|
||||
|
||||
def test_v2_benchmark_config_identity_validates_and_is_preserved():
|
||||
cfg = _v2_config()
|
||||
cfg["quality_metadata"] = {"some": "data"}
|
||||
|
||||
_validate_benchmark_config(cfg, "wan.json")
|
||||
|
||||
assert _is_v2_config(cfg) is True
|
||||
assert _config_identity_metadata(cfg) == {
|
||||
"config_schema_version": 2,
|
||||
"workload_id": "wan-t2v-1.3b",
|
||||
"variant_id": "canonical",
|
||||
"benchmark_version": 1,
|
||||
"quality_metadata": {"some": "data"},
|
||||
}
|
||||
|
||||
|
||||
def test_v2_benchmark_config_missing_identity_fields_fails_clearly():
|
||||
cfg = _v2_config()
|
||||
del cfg["variant_id"]
|
||||
del cfg["benchmark_version"]
|
||||
|
||||
expected = "wan.json: missing required v2 identity fields: variant_id, benchmark_version"
|
||||
with pytest.raises(ValueError, match=expected):
|
||||
_validate_benchmark_config(cfg, "wan.json")
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("field", "value"),
|
||||
[
|
||||
("workload_id", {}),
|
||||
("workload_id", ""),
|
||||
("workload_id", " "),
|
||||
("variant_id", []),
|
||||
("variant_id", ""),
|
||||
("variant_id", " "),
|
||||
],
|
||||
)
|
||||
def test_v2_benchmark_config_rejects_invalid_string_identity_values(field, value):
|
||||
cfg = _v2_config()
|
||||
cfg[field] = value
|
||||
|
||||
expected = f"wan.json: v2 identity field {field!r} must be a non-empty string"
|
||||
with pytest.raises(ValueError, match=expected):
|
||||
_validate_benchmark_config(cfg, "wan.json")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("value", [None, "1", 1.5, True])
|
||||
def test_v2_benchmark_config_rejects_invalid_benchmark_version_values(value):
|
||||
cfg = _v2_config()
|
||||
cfg["benchmark_version"] = value
|
||||
|
||||
expected = "wan.json: v2 identity field 'benchmark_version' must be an integer"
|
||||
with pytest.raises(ValueError, match=expected):
|
||||
_validate_benchmark_config(cfg, "wan.json")
|
||||
|
||||
|
||||
def test_partial_v2_identity_requires_schema_version():
|
||||
cfg = {
|
||||
"benchmark_id": "wan-t2v-1.3b-2gpu",
|
||||
"workload_id": "wan-t2v-1.3b",
|
||||
"variant_id": "canonical",
|
||||
"benchmark_version": 1,
|
||||
}
|
||||
|
||||
expected = "wan.json: v2 benchmark identity fields require config_schema_version=2"
|
||||
with pytest.raises(ValueError, match=expected):
|
||||
_validate_benchmark_config(cfg, "wan.json")
|
||||
|
||||
|
||||
def test_optional_v2_metadata_fields_must_be_objects():
|
||||
cfg = {
|
||||
"benchmark_id": "wan-t2v-1.3b-2gpu",
|
||||
"config_schema_version": 2,
|
||||
"workload_id": "wan-t2v-1.3b",
|
||||
"variant_id": "canonical",
|
||||
"benchmark_version": 1,
|
||||
"quality_metadata": ["not", "an", "object"],
|
||||
}
|
||||
|
||||
expected = "wan.json: optional v2 metadata field 'quality_metadata' must be an object"
|
||||
with pytest.raises(ValueError, match=expected):
|
||||
_validate_benchmark_config(cfg, "wan.json")
|
||||
@@ -1,6 +1,7 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
from fastvideo.tests.performance import compare_baseline
|
||||
from fastvideo.performance.metric_policy import resolve_metric_policies
|
||||
|
||||
|
||||
def _raw_result():
|
||||
@@ -72,9 +73,157 @@ def test_normalized_record_includes_source_metadata(monkeypatch):
|
||||
assert record["job_id"] == "job-1"
|
||||
|
||||
|
||||
def test_normalized_record_includes_effective_regression_thresholds(monkeypatch):
|
||||
monkeypatch.setenv("PERF_RUN_SOURCE", "pr")
|
||||
raw = _raw_result()
|
||||
raw["regression_thresholds"] = {
|
||||
"latency": {
|
||||
"threshold_percent": 0.09,
|
||||
"threshold_absolute": 0.75,
|
||||
"gated": True,
|
||||
},
|
||||
"throughput": {
|
||||
"gated": False,
|
||||
},
|
||||
}
|
||||
|
||||
record = compare_baseline.normalize_performance_result(raw)
|
||||
|
||||
assert record["regression_thresholds"]["latency"] == {
|
||||
"threshold_percent": 0.09,
|
||||
"threshold_absolute": 0.75,
|
||||
"gated": True,
|
||||
}
|
||||
assert record["regression_thresholds"]["throughput"]["gated"] is False
|
||||
|
||||
|
||||
def test_invalid_regression_threshold_container_uses_defaults():
|
||||
policies = resolve_metric_policies(["not", "a", "mapping"])
|
||||
|
||||
latency = next(policy for policy in policies if policy.key == "latency")
|
||||
assert latency.threshold_percent == 0.08
|
||||
assert latency.threshold_absolute == 0.5
|
||||
assert latency.gated is True
|
||||
|
||||
|
||||
def test_boolean_regression_threshold_values_are_ignored():
|
||||
policies = resolve_metric_policies({
|
||||
"latency": {
|
||||
"threshold_percent": True,
|
||||
"threshold_absolute": False,
|
||||
"gated": "false",
|
||||
}
|
||||
})
|
||||
|
||||
latency = next(policy for policy in policies if policy.key == "latency")
|
||||
assert latency.threshold_percent == 0.08
|
||||
assert latency.threshold_absolute == 0.5
|
||||
assert latency.gated is False
|
||||
|
||||
|
||||
def test_baseline_eligibility_only_for_successful_scheduled_main():
|
||||
assert compare_baseline._is_baseline_eligible("scheduled_main", True) is True
|
||||
assert compare_baseline._is_baseline_eligible("scheduled_main", False) is False
|
||||
assert compare_baseline._is_baseline_eligible("pr", True) is False
|
||||
assert compare_baseline._is_baseline_eligible("local", True) is False
|
||||
|
||||
|
||||
def test_latency_regression_requires_percent_and_absolute_floors():
|
||||
baseline = [{"latency": 10.0}]
|
||||
current = {"model_id": "wan", "latency": 10.6}
|
||||
|
||||
percent_only = resolve_metric_policies({
|
||||
"latency": {
|
||||
"threshold_percent": 0.05,
|
||||
"threshold_absolute": 0.75,
|
||||
}
|
||||
})
|
||||
absolute_only = resolve_metric_policies({
|
||||
"latency": {
|
||||
"threshold_percent": 0.10,
|
||||
"threshold_absolute": 0.5,
|
||||
}
|
||||
})
|
||||
both = resolve_metric_policies({
|
||||
"latency": {
|
||||
"threshold_percent": 0.05,
|
||||
"threshold_absolute": 0.5,
|
||||
}
|
||||
})
|
||||
|
||||
assert compare_baseline._check_regressions(current, baseline, percent_only) == []
|
||||
assert compare_baseline._check_regressions(current, baseline, absolute_only) == []
|
||||
|
||||
failures = compare_baseline._check_regressions(current, baseline, both)
|
||||
assert len(failures) == 1
|
||||
assert "latency regressed by 6.0% and 0.600" in failures[0]
|
||||
|
||||
|
||||
def test_throughput_regression_uses_higher_is_better_direction():
|
||||
baseline = [{"throughput": 10.0}]
|
||||
current = {"model_id": "wan", "throughput": 9.0}
|
||||
policies = resolve_metric_policies({
|
||||
"throughput": {
|
||||
"threshold_percent": 0.05,
|
||||
"threshold_absolute": 0.5,
|
||||
}
|
||||
})
|
||||
|
||||
failures = compare_baseline._check_regressions(current, baseline, policies)
|
||||
|
||||
assert len(failures) == 1
|
||||
assert "throughput regressed by 10.0% and 1.000" in failures[0]
|
||||
|
||||
|
||||
def test_memory_regression_uses_metric_specific_absolute_floor():
|
||||
baseline = [{"memory": 10000.0}]
|
||||
current = {"model_id": "wan", "memory": 10600.0}
|
||||
policies = resolve_metric_policies({
|
||||
"memory": {
|
||||
"threshold_percent": 0.05,
|
||||
"threshold_absolute": 256.0,
|
||||
}
|
||||
})
|
||||
|
||||
failures = compare_baseline._check_regressions(current, baseline, policies)
|
||||
|
||||
assert len(failures) == 1
|
||||
assert "memory regressed by 6.0% and 600.0" in failures[0]
|
||||
|
||||
|
||||
def test_component_metric_can_gate_independently():
|
||||
baseline = [{"dit_time_s": 8.0}]
|
||||
current = {"model_id": "wan", "dit_time_s": 8.6}
|
||||
policies = resolve_metric_policies({
|
||||
"dit_time_s": {
|
||||
"threshold_percent": 0.05,
|
||||
"threshold_absolute": 0.25,
|
||||
}
|
||||
})
|
||||
|
||||
failures = compare_baseline._check_regressions(current, baseline, policies)
|
||||
|
||||
assert len(failures) == 1
|
||||
assert "dit_time_s regressed by 7.5% and 0.600" in failures[0]
|
||||
|
||||
|
||||
def test_informational_metric_remains_visible_without_failing():
|
||||
baseline = [{"throughput": 10.0}]
|
||||
current = {"model_id": "wan", "gpu_type": "NVIDIA L40S", "throughput": 8.0}
|
||||
policies = resolve_metric_policies({
|
||||
"throughput": {
|
||||
"threshold_percent": 0.01,
|
||||
"threshold_absolute": 0.01,
|
||||
"gated": False,
|
||||
}
|
||||
})
|
||||
|
||||
row = compare_baseline._build_summary_row(current, baseline, policies, False)
|
||||
|
||||
assert compare_baseline._check_regressions(current, baseline, policies) == []
|
||||
assert row["metrics"]["throughput"]["regression_pct"] == 20.0
|
||||
assert row["metrics"]["throughput"]["gated"] is False
|
||||
assert row["metrics"]["throughput"]["threshold_exceeded"] is True
|
||||
assert row["metrics"]["throughput"]["regressed"] is False
|
||||
assert row["threshold_exceeded_metrics"] == ["throughput"]
|
||||
assert row["failing_metrics"] == []
|
||||
|
||||
@@ -68,6 +68,8 @@ def test_summary_endpoint_returns_latest_group_status():
|
||||
assert body["count"] == 1
|
||||
assert body["status_counts"] == {"pass": 1, "fail": 0}
|
||||
assert body["rows"][0]["metrics"]["latency"]["baseline"] == 10.0
|
||||
assert body["rows"][0]["metrics"]["latency"]["threshold_exceeded"] is True
|
||||
assert body["rows"][0]["threshold_exceeded_metrics"] == ["latency", "throughput"]
|
||||
assert body["rows"][0]["computed_regression_status"] == "fail"
|
||||
|
||||
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
from fastvideo.performance import hf_store
|
||||
from fastvideo.performance_dashboard.service import build_latest_summary, build_trends, filter_records
|
||||
from fastvideo.tests.performance import hf_store
|
||||
|
||||
|
||||
def _record(ts, commit, latency, throughput, success=True, **metadata):
|
||||
@@ -28,16 +28,23 @@ def test_build_latest_summary_uses_previous_successful_records_for_baseline():
|
||||
_record("2026-01-03T00:00:00+00:00", "c" * 40, 11.0, 9.0),
|
||||
]
|
||||
|
||||
rows = build_latest_summary(records, max_regression=0.05)
|
||||
rows = build_latest_summary(records)
|
||||
|
||||
assert len(rows) == 1
|
||||
row = rows[0]
|
||||
assert row["baseline_n"] == 1
|
||||
assert row["metrics"]["latency"]["baseline"] == 10.0
|
||||
assert row["metrics"]["latency"]["regression_pct"] == 10.0
|
||||
assert row["metrics"]["latency"]["absolute_delta"] == 1.0
|
||||
assert row["metrics"]["latency"]["threshold_percent"] == 8.0
|
||||
assert row["metrics"]["latency"]["threshold_absolute"] == 0.5
|
||||
assert row["metrics"]["latency"]["threshold_exceeded"] is True
|
||||
assert row["metrics"]["latency"]["regressed"] is True
|
||||
assert row["metrics"]["throughput"]["regression_pct"] == 10.0
|
||||
assert row["status"] == "pass"
|
||||
assert row["computed_regression_status"] == "fail"
|
||||
assert row["threshold_exceeded_metrics"] == ["latency", "throughput"]
|
||||
assert row["failing_metrics"] == ["latency", "throughput"]
|
||||
|
||||
|
||||
def test_build_latest_summary_status_uses_latest_record_success_field():
|
||||
@@ -73,7 +80,7 @@ def test_build_latest_summary_run_source_filter_keeps_canonical_baseline():
|
||||
),
|
||||
]
|
||||
|
||||
rows = build_latest_summary(records, max_regression=0.05, run_source="pr")
|
||||
rows = build_latest_summary(records, run_source="pr")
|
||||
|
||||
assert len(rows) == 1
|
||||
assert rows[0]["run_source"] == "pr"
|
||||
@@ -83,6 +90,60 @@ def test_build_latest_summary_run_source_filter_keeps_canonical_baseline():
|
||||
assert rows[0]["computed_regression_status"] == "fail"
|
||||
|
||||
|
||||
def test_build_latest_summary_requires_absolute_floor_for_computed_regression():
|
||||
records = [
|
||||
_record("2026-01-01T00:00:00+00:00", "a" * 40, 10.0, 10.0),
|
||||
_record(
|
||||
"2026-01-02T00:00:00+00:00",
|
||||
"b" * 40,
|
||||
10.6,
|
||||
10.0,
|
||||
regression_thresholds={
|
||||
"latency": {
|
||||
"threshold_percent": 0.05,
|
||||
"threshold_absolute": 0.75,
|
||||
"gated": True,
|
||||
}
|
||||
},
|
||||
),
|
||||
]
|
||||
|
||||
rows = build_latest_summary(records)
|
||||
|
||||
assert round(rows[0]["metrics"]["latency"]["regression_pct"], 1) == 6.0
|
||||
assert round(rows[0]["metrics"]["latency"]["absolute_delta"], 3) == 0.6
|
||||
assert rows[0]["metrics"]["latency"]["threshold_exceeded"] is False
|
||||
assert rows[0]["metrics"]["latency"]["regressed"] is False
|
||||
assert rows[0]["computed_regression_status"] == "pass"
|
||||
|
||||
|
||||
def test_build_latest_summary_separates_informational_threshold_crossing():
|
||||
records = [
|
||||
_record("2026-01-01T00:00:00+00:00", "a" * 40, 10.0, 10.0),
|
||||
_record(
|
||||
"2026-01-02T00:00:00+00:00",
|
||||
"b" * 40,
|
||||
10.6,
|
||||
10.0,
|
||||
regression_thresholds={
|
||||
"latency": {
|
||||
"threshold_percent": 0.05,
|
||||
"threshold_absolute": 0.5,
|
||||
"gated": False,
|
||||
}
|
||||
},
|
||||
),
|
||||
]
|
||||
|
||||
rows = build_latest_summary(records)
|
||||
|
||||
assert rows[0]["metrics"]["latency"]["threshold_exceeded"] is True
|
||||
assert rows[0]["metrics"]["latency"]["regressed"] is False
|
||||
assert rows[0]["threshold_exceeded_metrics"] == ["latency"]
|
||||
assert rows[0]["failing_metrics"] == []
|
||||
assert rows[0]["computed_regression_status"] == "pass"
|
||||
|
||||
|
||||
def test_filter_records_and_trends_preserve_metric_points():
|
||||
records = [
|
||||
_record("2026-01-01T00:00:00+00:00", "a" * 40, 10.0, 10.0),
|
||||
|
||||
@@ -28,6 +28,17 @@ STAGE_METRIC_MAP: dict[str, str] = {
|
||||
"DmdDenoisingStage": "dit_time_s",
|
||||
"DecodingStage": "vae_decode_time_s",
|
||||
}
|
||||
V2_CONFIG_SCHEMA_VERSION = 2
|
||||
V2_REQUIRED_IDENTITY_FIELDS = (
|
||||
"workload_id",
|
||||
"variant_id",
|
||||
"benchmark_version",
|
||||
)
|
||||
V2_OPTIONAL_METADATA_FIELDS = (
|
||||
"recipe",
|
||||
"metric_threshold_policy",
|
||||
"quality_metadata",
|
||||
)
|
||||
|
||||
# -- Config discovery -------------------------------------------------------
|
||||
|
||||
@@ -42,6 +53,71 @@ _BENCHMARKS_DIR = os.path.join(
|
||||
)
|
||||
|
||||
|
||||
def _has_v2_fields(cfg):
|
||||
v2_fields = V2_REQUIRED_IDENTITY_FIELDS + V2_OPTIONAL_METADATA_FIELDS
|
||||
return any(field in cfg for field in v2_fields)
|
||||
|
||||
|
||||
def _is_v2_config(cfg):
|
||||
return cfg.get("config_schema_version") == V2_CONFIG_SCHEMA_VERSION
|
||||
|
||||
|
||||
def _validate_non_empty_string(value, field, path):
|
||||
if not isinstance(value, str) or not value.strip():
|
||||
raise ValueError(f"{path}: v2 identity field {field!r} must be a non-empty string")
|
||||
|
||||
|
||||
def _validate_integer(value, field, path):
|
||||
if isinstance(value, bool) or not isinstance(value, int):
|
||||
raise ValueError(f"{path}: v2 identity field {field!r} must be an integer")
|
||||
|
||||
|
||||
def _validate_benchmark_config(cfg, path="<memory>"):
|
||||
missing_common = [field for field in ("benchmark_id",) if field not in cfg]
|
||||
if missing_common:
|
||||
raise ValueError(f"{path}: missing required benchmark config fields: {', '.join(missing_common)}")
|
||||
|
||||
schema_version = cfg.get("config_schema_version")
|
||||
if schema_version is None:
|
||||
if _has_v2_fields(cfg):
|
||||
raise ValueError(f"{path}: v2 benchmark identity fields require config_schema_version=2")
|
||||
return
|
||||
|
||||
if schema_version != V2_CONFIG_SCHEMA_VERSION:
|
||||
raise ValueError(f"{path}: unsupported benchmark config_schema_version={schema_version!r}")
|
||||
|
||||
missing_v2 = [field for field in V2_REQUIRED_IDENTITY_FIELDS if field not in cfg]
|
||||
if missing_v2:
|
||||
raise ValueError(f"{path}: missing required v2 identity fields: {', '.join(missing_v2)}")
|
||||
|
||||
_validate_non_empty_string(cfg["workload_id"], "workload_id", path)
|
||||
_validate_non_empty_string(cfg["variant_id"], "variant_id", path)
|
||||
_validate_integer(cfg["benchmark_version"], "benchmark_version", path)
|
||||
|
||||
for field in V2_OPTIONAL_METADATA_FIELDS:
|
||||
if field in cfg and not isinstance(cfg[field], Mapping):
|
||||
raise ValueError(f"{path}: optional v2 metadata field {field!r} must be an object")
|
||||
|
||||
|
||||
def _config_identity_metadata(cfg):
|
||||
if not _is_v2_config(cfg):
|
||||
return {}
|
||||
metadata = {
|
||||
"config_schema_version": cfg["config_schema_version"],
|
||||
"workload_id": cfg["workload_id"],
|
||||
"variant_id": cfg["variant_id"],
|
||||
"benchmark_version": cfg["benchmark_version"],
|
||||
}
|
||||
for field in V2_OPTIONAL_METADATA_FIELDS:
|
||||
if field in cfg:
|
||||
metadata[field] = cfg[field]
|
||||
return metadata
|
||||
|
||||
|
||||
def _benchmark_display_id(cfg):
|
||||
return cfg["benchmark_id"]
|
||||
|
||||
|
||||
def _discover_benchmarks():
|
||||
"""Glob benchmark JSON configs and return list of (id, config) tuples."""
|
||||
pattern = os.path.join(_BENCHMARKS_DIR, "*.json")
|
||||
@@ -49,6 +125,7 @@ def _discover_benchmarks():
|
||||
for path in sorted(glob.glob(pattern)):
|
||||
with open(path) as f:
|
||||
cfg = json.load(f)
|
||||
_validate_benchmark_config(cfg, path)
|
||||
configs.append(cfg)
|
||||
return configs
|
||||
|
||||
@@ -219,6 +296,7 @@ def _run_benchmark(cfg):
|
||||
|
||||
results = {
|
||||
"benchmark_id": cfg["benchmark_id"],
|
||||
**_config_identity_metadata(cfg),
|
||||
"model_short_name": model_info.get("model_short_name", ""),
|
||||
"device": device_name,
|
||||
"num_gpus": init_kwargs.get("num_gpus", 1),
|
||||
@@ -231,6 +309,7 @@ def _run_benchmark(cfg):
|
||||
"max_peak_memory_mb": round(max_peak_memory, 1),
|
||||
"individual_peak_memories_mb": [round(m, 1) for m in peak_memories],
|
||||
"thresholds": thresholds,
|
||||
"regression_thresholds": cfg.get("regression_thresholds", {}),
|
||||
"commit": os.environ.get("BUILDKITE_COMMIT", ""),
|
||||
"pr_number": os.environ.get("BUILDKITE_PULL_REQUEST", ""),
|
||||
"timestamp": datetime.now(timezone.utc).isoformat(),
|
||||
@@ -275,7 +354,7 @@ def _run_benchmark(cfg):
|
||||
@pytest.mark.parametrize(
|
||||
"cfg",
|
||||
_BENCHMARK_CONFIGS,
|
||||
ids=[c["benchmark_id"] for c in _BENCHMARK_CONFIGS],
|
||||
ids=[_benchmark_display_id(c) for c in _BENCHMARK_CONFIGS],
|
||||
)
|
||||
def test_inference_performance(cfg):
|
||||
"""Measure generation latency, peak GPU memory, and component-level timings
|
||||
|
||||
@@ -0,0 +1,40 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
from types import SimpleNamespace
|
||||
|
||||
from fastvideo.pipelines import ForwardBatch
|
||||
from fastvideo.pipelines.stages.input_validation import InputValidationStage
|
||||
|
||||
|
||||
def test_input_validation_preserves_explicit_dynamic_batch_seeds() -> None:
|
||||
batch = ForwardBatch(
|
||||
data_type="video",
|
||||
prompt=["one", "two"],
|
||||
seed=100,
|
||||
seeds=[17, 23],
|
||||
height=8,
|
||||
width=8,
|
||||
num_frames=1,
|
||||
num_inference_steps=1,
|
||||
)
|
||||
|
||||
InputValidationStage()._generate_seeds(batch, SimpleNamespace())
|
||||
|
||||
assert batch.seeds == [17, 23]
|
||||
assert [generator.initial_seed() for generator in batch.generator] == [17, 23]
|
||||
|
||||
|
||||
def test_input_validation_generates_one_seed_per_prompt() -> None:
|
||||
batch = ForwardBatch(
|
||||
data_type="video",
|
||||
prompt=["one", "two"],
|
||||
seed=100,
|
||||
height=8,
|
||||
width=8,
|
||||
num_frames=1,
|
||||
num_inference_steps=1,
|
||||
)
|
||||
|
||||
InputValidationStage()._generate_seeds(batch, SimpleNamespace())
|
||||
|
||||
assert batch.seeds == [100, 101]
|
||||
assert [generator.initial_seed() for generator in batch.generator] == [100, 101]
|
||||
@@ -13,7 +13,13 @@ class TensorDict(dict):
|
||||
return TensorDict({k: v.to(device) for k, v in self.items()})
|
||||
|
||||
class FakeTokenizer:
|
||||
def __init__(self):
|
||||
self.calls = []
|
||||
self.texts = []
|
||||
|
||||
def __call__(self, texts, **kwargs):
|
||||
self.calls.append(kwargs)
|
||||
self.texts.append(list(texts))
|
||||
B = len(texts)
|
||||
seq_len = int(kwargs.get("max_length", 4))
|
||||
return TensorDict({
|
||||
@@ -21,6 +27,26 @@ class FakeTokenizer:
|
||||
"attention_mask": torch.ones(B, seq_len, dtype=torch.long),
|
||||
})
|
||||
|
||||
|
||||
class FakeChatTokenizer:
|
||||
def __init__(self):
|
||||
self.last_messages = None
|
||||
self.last_kwargs = None
|
||||
|
||||
def apply_chat_template(self, messages, **kwargs):
|
||||
self.last_messages = messages
|
||||
self.last_kwargs = kwargs
|
||||
assert isinstance(messages[0], list)
|
||||
assert messages[0][0]["role"] == "system"
|
||||
assert messages[0][1]["role"] == "user"
|
||||
B = len(messages)
|
||||
seq_len = int(kwargs.get("max_length", 4))
|
||||
return TensorDict({
|
||||
"input_ids": torch.arange(B * seq_len).view(B, seq_len),
|
||||
"attention_mask": torch.ones(B, seq_len, dtype=torch.long),
|
||||
})
|
||||
|
||||
|
||||
class FakeTextEncoder(torch.nn.Module):
|
||||
def __init__(self, hidden_size=8):
|
||||
super().__init__()
|
||||
@@ -38,6 +64,14 @@ class FakeTextEncoder(torch.nn.Module):
|
||||
def id_preprocess(x: str) -> str:
|
||||
return x
|
||||
|
||||
|
||||
def chat_list_preprocess(x: str):
|
||||
return [
|
||||
{"role": "system", "content": "Describe the video."},
|
||||
{"role": "user", "content": x if x else " "},
|
||||
]
|
||||
|
||||
|
||||
def take_mean_postprocess(outputs: BaseEncoderOutput) -> torch.Tensor:
|
||||
# [B, T, H] -> [B, H]
|
||||
return outputs.last_hidden_state.mean(dim=1)
|
||||
@@ -131,6 +165,66 @@ def test_forward_integration_cfg_off_and_on():
|
||||
assert len(out2.prompt_attention_mask) == 2
|
||||
assert len(out2.negative_attention_mask) == 2
|
||||
|
||||
def test_encode_text_adds_padding_for_prompt_lists():
|
||||
fastvideo_args, hidden = make_args(num_encoders=1, text_len=4, hidden_size=8)
|
||||
stage = make_stage(num_encoders=1, hidden_size=hidden)
|
||||
|
||||
stage.encode_text(["short", "a longer prompt"], fastvideo_args, encoder_index=[0])
|
||||
|
||||
assert stage.tokenizers[0].calls[-1]["padding"] is True
|
||||
|
||||
|
||||
def test_forward_prompt_list_preserves_single_prompt_text_encoding_path():
|
||||
fastvideo_args, hidden = make_args(num_encoders=1, text_len=4, hidden_size=8)
|
||||
stage = make_stage(num_encoders=1, hidden_size=hidden)
|
||||
batch = ForwardBatch(
|
||||
data_type="video",
|
||||
prompt=["short", "a longer prompt"],
|
||||
negative_prompt="",
|
||||
do_classifier_free_guidance=False,
|
||||
prompt_embeds=[],
|
||||
negative_prompt_embeds=None,
|
||||
prompt_attention_mask=[],
|
||||
negative_attention_mask=None,
|
||||
)
|
||||
|
||||
out = stage.forward(batch, fastvideo_args)
|
||||
|
||||
assert stage.tokenizers[0].texts == [["short"], ["a longer prompt"]]
|
||||
assert out.prompt_embeds[0].shape == (2, hidden)
|
||||
assert out.prompt_attention_mask[0].shape == (2, 4)
|
||||
|
||||
|
||||
def test_encode_prompt_list_individually_pads_variable_length_embeds_and_audio():
|
||||
fastvideo_args, _hidden = make_args(num_encoders=1, text_len=4, hidden_size=8)
|
||||
stage = TextEncodingStage(text_encoders=[], tokenizers=[])
|
||||
lengths = {"short": 2, "a longer prompt": 4}
|
||||
|
||||
def fake_encode_text(text, *_args, **_kwargs):
|
||||
length = lengths[text]
|
||||
embeds = [torch.full((1, length, 3), fill_value=float(length))]
|
||||
masks = [torch.ones((1, length), dtype=torch.long)]
|
||||
stage._last_audio_embeds = [torch.full((1, length, 5), fill_value=float(length))]
|
||||
return embeds, masks
|
||||
|
||||
stage.encode_text = fake_encode_text
|
||||
|
||||
embeds, masks = stage._encode_prompt_list_individually(
|
||||
["short", "a longer prompt"],
|
||||
fastvideo_args,
|
||||
encoder_index=[0],
|
||||
return_attention_mask=True,
|
||||
)
|
||||
|
||||
assert embeds[0].shape == (2, 4, 3)
|
||||
assert masks[0].shape == (2, 4)
|
||||
assert stage._last_audio_embeds is not None
|
||||
assert stage._last_audio_embeds[0].shape == (2, 4, 5)
|
||||
assert torch.equal(embeds[0][0, :2], torch.full((2, 3), 2.0))
|
||||
assert torch.equal(embeds[0][0, 2:], torch.zeros((2, 3)))
|
||||
assert torch.equal(stage._last_audio_embeds[0][0, :2], torch.full((2, 5), 2.0))
|
||||
assert torch.equal(stage._last_audio_embeds[0][0, 2:], torch.zeros((2, 5)))
|
||||
|
||||
|
||||
def test_encode_text_hidden_state_flag_follows_encoder_config():
|
||||
fastvideo_args, hidden = make_args(num_encoders=1, text_len=4, hidden_size=8)
|
||||
@@ -156,3 +250,32 @@ def test_encode_text_does_not_force_hidden_states_for_ltx2_prefix():
|
||||
stage.encode_text("a", fastvideo_args, encoder_index=[0])
|
||||
|
||||
assert stage.text_encoders[0].last_output_hidden_states is False
|
||||
|
||||
|
||||
def test_chat_list_preprocess_output_is_not_stripped():
|
||||
fastvideo_args, hidden = make_args(num_encoders=1, text_len=5, hidden_size=8)
|
||||
encoder_config = fastvideo_args.pipeline_config.text_encoder_configs[0]
|
||||
encoder_config.is_chat_model = True
|
||||
encoder_config.treat_empty_as_dot = True
|
||||
fastvideo_args.pipeline_config.preprocess_text_funcs = (chat_list_preprocess, )
|
||||
|
||||
tokenizer = FakeChatTokenizer()
|
||||
stage = TextEncodingStage(
|
||||
text_encoders=[FakeTextEncoder(hidden_size=hidden)],
|
||||
tokenizers=[tokenizer],
|
||||
)
|
||||
|
||||
embeds, masks = stage.encode_text(
|
||||
"a robotic arm welding a metal structure",
|
||||
fastvideo_args,
|
||||
encoder_index=[0],
|
||||
return_attention_mask=True,
|
||||
)
|
||||
|
||||
assert embeds[0].shape == (1, hidden)
|
||||
assert masks[0].shape == (1, 5)
|
||||
assert tokenizer.last_messages == [[
|
||||
{"role": "system", "content": "Describe the video."},
|
||||
{"role": "user", "content": "a robotic arm welding a metal structure"},
|
||||
]]
|
||||
assert tokenizer.last_kwargs["return_tensors"] == "pt"
|
||||
|
||||
@@ -127,4 +127,12 @@ def test_wan_causal_dfsft_single_train_step(
|
||||
|
||||
# 5a-ii: device-keyed grad-norm regression on top of the same harness.
|
||||
# Skips when the current GPU has no seeded reference.
|
||||
check_grad_norm_regression("test_wan_causal_dfsft", model.transformer)
|
||||
# rtol above the harness default: the causal model compiles flex_attention
|
||||
# with max-autotune (required for Wan 1.3B's head config), and the
|
||||
# timing-based kernel selection is bimodal across L40S containers —
|
||||
# observed 3.2562 vs 3.5860 (10.13% apart) with identical code, straddling
|
||||
# the default 10%. 12% covers both winners; real wiring breakage (dead
|
||||
# grads, scale bugs) still lands far outside it.
|
||||
check_grad_norm_regression("test_wan_causal_dfsft",
|
||||
model.transformer,
|
||||
rtol=0.12)
|
||||
|
||||
@@ -1,225 +0,0 @@
|
||||
# FastVideo Runtime — Aggressive Implementation Plan
|
||||
|
||||
**Companion to** `design.md` (v19) and `design_summary.md` · **Stance:** this plan trades interface stability for
|
||||
speed. Where it deviates from design.md's conservative migration (§10), the deviation is flagged with **⚡**.
|
||||
design.md remains the architectural authority; this is the execution order.
|
||||
|
||||
---
|
||||
|
||||
## 1. Rules of engagement
|
||||
|
||||
**We break, freely and early:**
|
||||
|
||||
- The public Python API: `generate_video(**kwargs)` and `SamplingParam` are **deleted**, not deprecated.
|
||||
- Config schemas: `FastVideoArgs` (1,272 lines, 81 fields) stops being a public or threaded surface.
|
||||
- `fastvideo/api/compat.py` (651 lines): **deleted in M1** ⚡ (design.md §6.6 shrinks it monotonically to Phase 5 —
|
||||
that policy existed only to honor signatures we are now licensed to break).
|
||||
- CLI flags, YAML schemas, package layout, `fastvideo.api` exports, ComfyUI node params, every example.
|
||||
- In-repo dependents (`apps/dreamverse`, `comfyui/`, `examples/`, `scripts/`) get **fixed in the same PR train** —
|
||||
we own them; no deprecation period, no shims.
|
||||
|
||||
**We never break, at any speed:**
|
||||
|
||||
- **Numerics.** Bit-identical loop parity and SSIM gates are not "conservative" — they are the definition of
|
||||
correct. Aggression applies to interfaces, never to outputs.
|
||||
- **Model coverage** for the families that matter (tier list in §6 — the tail is a decision, not a casualty).
|
||||
- The frozen legacy `fastvideo/training/` stack (N2) and the bit-exact porting methodology (N3).
|
||||
- External users get **batched breakage**: all user-visible breaks land in at most two releases (R1 = request/config
|
||||
cut, R2 = engine default), each with a migration guide and a `fastvideo migrate` codemod — never a drip.
|
||||
|
||||
## 2. The sequencing argument (answering "fix omni request first, then separate the planes?")
|
||||
|
||||
**Yes to the first half. The second half should not be a project.** The three planes are not separated by moving
|
||||
code into plane-named directories — today's monolithic stages would just get reshuffled and then rewritten when
|
||||
loops invert. The planes are *born* from two cuts, and a third that is really a config change:
|
||||
|
||||
1. **The request-plane cut (M1)** — `OmniRequest` becomes the only currency crossing the boundary. Everything
|
||||
behind it is implementation. This is your "fix the omni input and request first," and it goes first because it
|
||||
is low-risk, it defines the vocabulary every later stage consumes, and it gets the user-facing pain over with
|
||||
while the codebase is still familiar.
|
||||
2. **Loop inversion (M2)** — this *is* the pipeline/execution plane separation. Once families expose
|
||||
`init/step/finalize` step bodies, something other than the family must own iteration; that owner is the
|
||||
executor, and the execution plane exists by construction. Before inversion there is nothing for an execution
|
||||
plane to schedule — "separating" it would be an empty directory.
|
||||
3. **The config cut (inside M1)** — the real coupling between planes today is `FastVideoArgs`: one 81-field object
|
||||
threaded through entrypoints, pipelines, stages, and executors, mixing deploy-time, model-time, and
|
||||
request-time concerns. Splitting it into `DeployConfig` / `ModelSpec` / `OmniRequest` (design.md §6.6's four
|
||||
layers) is the single highest-leverage "separation" action, and it's schema work, not architecture work.
|
||||
|
||||
So the order is: **M1 request+config cut → M2 loop inversion (planes now exist) → M3 engine on top.** Plane
|
||||
separation is the *outcome* of M1+M2, not a milestone.
|
||||
|
||||
## 3. Milestones
|
||||
|
||||
Timeline assumes 3–4 engineers on the runtime critical path. Overlap is deliberate; gates are not ⚡-able.
|
||||
|
||||
### M0 — Baselines, harness, enforcement (weeks 0–3, overlaps M1)
|
||||
|
||||
The license for everything aggressive afterward. Not skippable, not shrinkable. M0 does not block M1 (which
|
||||
changes no numerics) — the only hard rule is **no family's M2 migration starts before its baseline exists**.
|
||||
|
||||
- Merge the `feat/cosmos3-reasoning` chain (design.md sizes this alone at 2–3 engineer-months — it runs as its own
|
||||
track); seed SSIM references for the ~7 uncovered families.
|
||||
- ParityAligner v0: record/compare named taps on *current* pipelines (it must exist before anything changes).
|
||||
- **The enforcement package, on day one** (design.md §10 — the prior freeze was broken 19× for lack of exactly
|
||||
this): CI path gates (reject new `fastvideo/training/` files now; reject new `DenoisingStage` subclasses once the
|
||||
first M2 family lands), CODEOWNERS on the frozen and migrating paths, a named owner per milestone, and the
|
||||
inflow rule — new model families land on the new abstractions from the first Wan/Flux2 landing onward.
|
||||
- Announce the M1 freeze window for in-flight PRs touching `fastvideo/api/`, `fastvideo_args.py`, entrypoints.
|
||||
|
||||
*Gate: every tier-A family has a recorded SSIM + activation baseline; CI gates live.*
|
||||
|
||||
### M1 — The request-plane cut (weeks 1–4) → **breaking release R1**
|
||||
|
||||
The typed API is partway there: `VideoGenerator.generate(GenerationRequest)` is already the documented primary
|
||||
entrypoint (`generate_video` carries a deprecation warning), and `fastvideo/entrypoints/openai/` already serves
|
||||
`POST /v1/videos` and `POST /v1/images`. But the legacy path is still what's *used*: Dreamverse calls
|
||||
`generate_video(**kwargs)` (`apps/dreamverse/dreamverse/video_generation.py:508`), as do ComfyUI and most
|
||||
examples. M1 finishes the cut instead of bridging it:
|
||||
|
||||
- **`OmniRequest` / `OmniOutput` / `OmniEvent`**: evolve `api/schema.py`'s `GenerationRequest` in place — typed
|
||||
modality parts, `TaskType`, per-model `ModelOptions` registered blocks (formalizing the `api/matrixgame2.py`
|
||||
pattern), seeds/priority/streaming flags. `api/results.py`'s `Video*Event` types become `OmniEvent`.
|
||||
- **Config: four layers, one owner each** (§6.6): extract `DeployConfig` (placement, parallelism axes, memory/
|
||||
offload, compile, plugins) from the runtime third of `FastVideoArgs` + `EngineConfig`/`ParallelismConfig`;
|
||||
`ModelSpec` manifest v0 (manifest-first component resolution; today's name-detectors as fallback);
|
||||
`OmniRequest` absorbs every per-call field. CLI flags, OpenAI protocol models, and presets are **generated**
|
||||
from the schema.
|
||||
- **Delete** ⚡: `compat.py` (651), `sampling_param.py` (411), `generate_video()`, the `fastvideo.api` legacy
|
||||
exports, `FastVideoArgs` as a *public* type. Internally it survives as a boundary-constructed shim for as long
|
||||
as anything still receives it: migrated families drop it per-family in M2, but unmigrated tier-B stages
|
||||
(`LegacyPipelineNode`) and the frozen `training/` stack (whose `TrainingArgs` subclasses it) carry it until M6 —
|
||||
it dies as a type with the tail, not before.
|
||||
- **Fix in-train**: ComfyUI nodes (legacy-API callers), all `examples/` (~75 files, mostly mechanical),
|
||||
`scripts/`, docs. Ship `fastvideo migrate` (codemod: old kwargs/YAML → `OmniRequest`/`DeployConfig`).
|
||||
- Internals unchanged: `ForwardBatch` is built *from* `OmniRequest` at the boundary; the executor and stages are
|
||||
untouched in M1.
|
||||
|
||||
*Gate: all SSIM suites unchanged; Dreamverse + ComfyUI + examples green on the new surface; R1 notes + codemod
|
||||
published.*
|
||||
|
||||
### M2 — Loop inversion (weeks 4–10): the pipeline plane is born
|
||||
|
||||
- `DenoiseLoop` / `ARDecodeLoop` with `init/step/finalize`; runtime owns iteration; custom-step escape hatch from
|
||||
day one (the Cosmos3-port and self-forcing pattern is legitimate, §6.2.3).
|
||||
- **Family order** (each lands step body + policies, and **deletes its legacy stage code in the same PR** ⚡ —
|
||||
continuous deletion, no end-of-plan cliff): **Wan 2.1/2.2 + Flux2 first, jointly** — design.md's rationale
|
||||
stands: together they exercise CFG variants, expert routing, chunk-KV, and the image path, so the step-body
|
||||
contract freezes only after all four are exercised → Wan-causal (self-forcing student) → LTX-2 → HunyuanVideo →
|
||||
Stable Audio → remaining image families. Unmigrated families keep running via `LegacyPipelineNode`.
|
||||
- Policies: `CFGPolicy` (absorbs the 3 CFG copies), `AttnMetadataProvider`, `FlowShiftPolicy`, `PrecisionPolicy`.
|
||||
- Extension core lands with the loop (it's why the loop is being rebuilt): observer bus, ParityAligner promoted to
|
||||
per-request observer, Profiler/NaNWatch, and **cache-dit as the first interceptor** (retiring `enable_teacache`).
|
||||
- `forward_context.py` off the *migrated* inference path (194 references across ~68 files today: ~8 importer files
|
||||
in frozen `training/`, the rest spread across train/ models, tier-B inference stages, quantization, and tests);
|
||||
the module survives as a shim for frozen `training/` **and unmigrated tier-B stages** until M6 — what M2
|
||||
guarantees is that no migrated family and no new code touches it.
|
||||
- **`train/` migrates per-family, immediately behind inference**: DMD2 and the landed DiffusionNFT (#1450) adopt
|
||||
the shared step functions as each family's body lands — `rl/common/sampling.py`'s loop is deleted, #1396
|
||||
grad-norm refs extended to the migrated methods (RL included).
|
||||
|
||||
*Gate, per family: old-vs-new loop bit-identical (ParityAligner) + SSIM + a recorded loop-overhead / batch-of-1
|
||||
latency measurement (the baseline M3 gates against); for train/: seeded rollout latents identical, reward metrics
|
||||
- grad-norms neutral. No family is ever dual-maintained.*
|
||||
|
||||
### M3 — Execution plane: engine + scheduler (weeks 8–14, overlaps M2) → **breaking release R2**
|
||||
|
||||
- `AsyncEngine` (queue, admission, cancellation-as-common-path, failure isolation); offline `VideoGenerator` keeps
|
||||
its name, becomes a thin sync wrapper that can bypass the queue.
|
||||
- `StepScheduler` v0: multiplexes denoise steps across requests in a pool; budget currency = **predicted GPU-time**
|
||||
from a calibrated per-(model, phase, shape) cost table (the cost *model* matures later; the currency is right
|
||||
from day one). Carries the `ARDecodeLoop` contract; AR batching itself waits for its workload (N5).
|
||||
- CacheManager v0: per-request chunk-KV slabs behind `KVHandle`; CFG-parallel axis (2-branch in practice).
|
||||
- **Dynamo stock worker** (registration, health/drain, cost metrics), retiring the locked
|
||||
`dynamo/examples/diffusers/worker.py` pattern.
|
||||
- **Dreamverse hard-cut** (per design.md Phase 2; the aggressive delta is doing it in one PR): `gpu_pool.py`,
|
||||
queue, warmup, and stream relay deleted and replaced by engine-client calls; the duty-cycle concurrency study
|
||||
runs on the result.
|
||||
- Colocated weight-sync RPC + component-granular sleep/wake + `RolloutClient` (engine-client RL mode for #1450).
|
||||
|
||||
*Gate: serving load tests; batch-of-1 latency regression ≤ 2% vs the M2-recorded measurement; Dreamverse
|
||||
single-session parity; RL engine-client seeded final-latent parity vs in-process; deploys under stock Dynamo.*
|
||||
|
||||
### M4 — Graphs, parallelism, multi-session (weeks 14–20)
|
||||
|
||||
- `PipelineSpec` graph IR: per-family pipeline classes shrink to **spec + step body + policies**
|
||||
(`create_pipeline_stages()` retires); LTX-2 and Hunyuan15+SR land as real fan-out graphs.
|
||||
- Role pools + connectors (port `multimodal_gen`'s disagg state machine); declarative stacked-parallelism axes
|
||||
compiled to DeviceMesh; general cross-mesh `WeightSyncPlan`.
|
||||
- ComfyUI workflow→spec compiler MVP (tier-1 ~20-node vocabulary) + weight/adapter fleet cache.
|
||||
|
||||
*Gate (design.md Phase 3's, in full): ≥2 Dreamverse sessions/GPU on the recorded duty-cycle trace, p95 within SLO
|
||||
— this is also where the loop-inversion **falsifier** is evaluated (see §7); LTX-2 A/V full-fan-out end-to-end;
|
||||
disaggregated-vs-colocated throughput benchmark; CPU-only topology validation suite; ComfyUI tier-1 workflows
|
||||
compile and run with equivalence reports; spec-built pipelines SSIM-identical to M2 loop versions.*
|
||||
|
||||
### M5 — Omni/MoT native + RL hardening (weeks 20–30)
|
||||
|
||||
- Cosmos3 re-port onto specs: packed factored sequences, dual-pathway attention, reasoner paged KV, joint denoise,
|
||||
world-model `ChunkRollout`; `/v1/chat/completions`; AR continuous batching arrives **with** this workload (N5).
|
||||
- Consistency ladder enforced end-to-end: C1 default in CI, C2 bitwise mode for goldens, Behavior Record opt-in;
|
||||
first GRPO-class method lands on the engine-client rollout path (log-prob drift becomes the gated metric).
|
||||
|
||||
*Gate: Cosmos3 150-test parity suite on the new runtime; reasoner pool efficiency — tokens/s/GPU at target
|
||||
concurrent denoise throughput, with the ≥10×-vs-re-prefill sanity floor; C1 drift ≈ 0 on a Wan RL run with the
|
||||
drift dashboard live.*
|
||||
|
||||
### M6 — The tail and the precondition (week 30+)
|
||||
|
||||
Continuous deletion (M1/M2) shrinks the final phase but does not eliminate it: what remains by M5 is the tier-B
|
||||
tail on `LegacyPipelineNode` and the frozen `training/` stack — which is a *live consumer* of
|
||||
`ComposedPipelineBase` and `forward_context`, so its retirement is the precondition, exactly as design.md Phase 5
|
||||
states. M6 = execute the §6 tail decision (migrate or deprecate each tier-B family), retire `training/` per the
|
||||
checklist, then delete `ComposedPipelineBase`, the legacy `DenoisingStage`, `forward_context.py`,
|
||||
`FastVideoArgs`/`TrainingArgs`, and `RayDistributedExecutor` together. **4 loop copies → 1.**
|
||||
|
||||
## 4. Breakage manifest (user-visible)
|
||||
|
||||
| Release | What breaks | Replacement | Aid |
|
||||
|---|---|---|---|
|
||||
| **R1** (M1) | `generate_video(prompt, **kwargs)`, `SamplingParam`, `FastVideoArgs` as public type, `fastvideo.api` legacy exports, CLI flag names, YAML config schema, streaming event types (`Video*Event` → `OmniEvent`, `schema_version`'d from day one) | `VideoGenerator.generate(OmniRequest)`, `DeployConfig`, generated CLI/protocol, `OmniEvent` | `fastvideo migrate` codemod, migration guide, R0 pinned |
|
||||
| **R2** (M3) | Default execution path becomes the engine (offline bypass preserved); server lifecycle (queue/admission semantics, job states) | `AsyncEngine` | guide; `OmniEvent` schema unchanged from R1 |
|
||||
| after R2 | nothing user-visible — M4/M5 are additive | — | — |
|
||||
|
||||
## 5. Deviations from design.md §10, stated honestly
|
||||
|
||||
| design.md | this plan | why it's safe now |
|
||||
|---|---|---|
|
||||
| Phase 0 keeps `VideoGenerator`/CLI signatures; `compat.py` shrinks to Phase 5 | M1 breaks signatures, deletes `compat.py` ⚡ | the only argument for the shim was signature stability — explicitly revoked |
|
||||
| Legacy code deleted at Phase 5 | per-family deletion at parity, M2 onward ⚡ | parity gate is per-family anyway; carrying dead code to a final phase only invites the 19×-broken-freeze failure mode |
|
||||
| Phases strictly sequential | M2/M3 overlap ⚡ | the engine consumes step bodies, not finished families; the step-body contract freezes at the Wan+Flux2 landing |
|
||||
| Phases −1 through 4 sized at 36–54 engineer-months | ~21–28 engineer-months (3–4 eng × 30 wks) ⚡ | the delta is real deleted work — no compat maintenance, no adapter upkeep, no dual-stack carry — plus M2/M3 overlap; treat 30 weeks as the aggressive case and 36–40 as the planning case |
|
||||
| Unchanged | parity/SSIM gates (restored in full at every milestone), enforcement package (CI path gates, CODEOWNERS, inflow rule — now at M0), train/RL migration timing (design.md Phase 1 already migrates NFT), Dreamverse hard-cut (Phase 2 already prescribes it), N2/N3/N5, cost-model currency, Dynamo asks + fallbacks, schema versioning | aggression budget is spent on interfaces only |
|
||||
|
||||
## 6. Decisions needed before M0
|
||||
|
||||
1. **Tier the model zoo.** Tier A (migrated, coverage guaranteed): Wan 2.1/2.2, Wan-causal/self-forcing, LTX-2,
|
||||
Flux2, HunyuanVideo, Stable Audio, Cosmos3 (contingent on the M0 merge — it is not on `main` today), image
|
||||
families. Tier B (runs on `LegacyPipelineNode` until someone claims it, candidate for deprecation at M6):
|
||||
gen3c, matrixgame2/3, longcat, the rest. **Approve or edit the split** — it bounds M2.
|
||||
2. **Release framing.** R1 as `v0.3.0` (pre-1.0 semantics, loud notes) vs holding breaks for a `v1.0` story.
|
||||
Recommendation: `v0.3.0` now — waiting taxes every milestone.
|
||||
3. **Freeze windows.** M1 freezes `api/`/args/entrypoints PRs ~2 weeks; M2 freezes per-family stage PRs while that
|
||||
family migrates (days each). Needs maintainer sign-off.
|
||||
4. **Staffing.** Critical path is M2's per-family step bodies — parallelizable per family after the Wan+Flux2
|
||||
reference lands. 3–4 engineers ≈ 30 weeks to M5 in the aggressive case (design.md's own sizing implies 36–40
|
||||
weeks at the same staffing — see §5); 2 engineers ≈ stretch ~1.5×. The Cosmos3-chain merge (M0) is its own
|
||||
2–3 engineer-month track and should be staffed separately from the runtime critical path.
|
||||
|
||||
## 7. Risks specific to the aggressive posture
|
||||
|
||||
- **In-flight PR collisions** with layout/schema moves → freeze windows (above) + landing schema cuts at
|
||||
milestone *starts*, not ends.
|
||||
- **Community churn at R1** (ComfyUI users, script users) → codemod covers the mechanical 90%; the 10% that isn't
|
||||
mechanical (kwargs with changed semantics) is enumerated in the guide; previous version stays pinned and
|
||||
installable.
|
||||
- **Parity harness becomes the bottleneck** — every aggressive deletion is licensed by it. Mitigation: it is the
|
||||
*first* deliverable (M0), and per-family migration PRs are template-driven (record → port → compare → delete).
|
||||
- **Overlap risk (M2/M3)**: the engine team building against a moving step-body contract → the contract
|
||||
(`init/step/finalize` + `StepResult`) freezes at the *first* family (Wan), enforced by the same schema-version
|
||||
discipline as external surfaces.
|
||||
- **The known unknown**: loop inversion at scheduler granularity has no production precedent (design.md §1). The
|
||||
falsifier stands, on design.md §11.6's schedule: the M3 duty-cycle study *publishes the targets*; the falsifier
|
||||
is **evaluated at the M4 gate** — if step-level multiplexing can't beat request-level serialization on real
|
||||
Dreamverse traces, StepScheduler retreats to request-level dispatch and the loop contract keeps only its
|
||||
streaming/preemption seams, with no family code changing — step bodies and the M1/M2 cuts retain full value.
|
||||
+2
-3
@@ -223,9 +223,8 @@ follow_imports = "silent"
|
||||
skip = "./data,./wandb,apps/fastvideo_studio/package-lock.json,apps/performance_dashboard/frontend/package-lock.json,*/_vendored/*"
|
||||
# "tread" matches daVinci-MagiHuman's acronym "TReAD" (Token Routing and
|
||||
# Early Drop). codespell lowercases ignore-words entries, so the single
|
||||
# lowercase form silences all case variants. "mot" = Mixture-of-Transformers
|
||||
# (MoT); "clen" = a Content-Length local; "te" = a text-embeds local.
|
||||
ignore-words-list = "tread,passt,mot,clen,te"
|
||||
# lowercase form silences all case variants.
|
||||
ignore-words-list = "tread,passt"
|
||||
|
||||
[tool.ruff]
|
||||
# Allow lines to be as long as 120.
|
||||
|
||||
@@ -1,267 +0,0 @@
|
||||
# Adversarial Review of `design.md` (v12) — FastVideo Next-Generation Inference Runtime
|
||||
|
||||
**Date:** 2026-06-11
|
||||
**Method:** Multi-agent adversarial review. 9 fact-check agents verified 70 concrete claims against the repo, the local reference checkouts (`cosmos-framework/`, `dynamo/`, `vllm-omni/`, `~/sglang`, `~/vllm`, `~/miles`, `~/verl-omni`, `~/diffusers`, `~/torchtitan`, `~/xDiT`, `~/ComfyUI`, `~/sglang-omni`, `~/cosmos-rl`), and GitHub. 9 attack lenses (abstractions, scheduler/perf, memory/cache, training/RL, strategy, migration, internal consistency, omissions, external borrowings) plus a completeness critic raised 70 findings; every finding went to a refute-by-default verifier. 36 findings were refuted; this document contains only the 34 that survived (1 critical, 26 major — consolidated below where lenses converged — 7 minor), plus fact-check corrections.
|
||||
|
||||
---
|
||||
|
||||
## Verdict
|
||||
|
||||
The architecture survives its strongest attacks — loop inversion's expressibility, the typed-state hybrid, the N1/N5 scope discipline, the clean-room GPL posture, and the C2-for-batch-1-video argument all held under refutation attempts. What does not survive is:
|
||||
|
||||
1. **The migration plan**, which consumes its own substrate two phases before building it and rests on a "frozen legacy stack" premise this repo has already empirically falsified.
|
||||
2. **Two load-bearing factual errors** about reference systems (vLLM's BlockPool page sizes, diffusers' loop ownership) that each drove a recorded design decision.
|
||||
3. **A family of undesigned failure/memory/trust paths** that the multiplexing bet itself creates. One is critical.
|
||||
|
||||
---
|
||||
|
||||
## Critical
|
||||
|
||||
### C1. No failure-isolation or cancellation semantics for the multiplexed pool — the blast-radius problem the architecture itself creates
|
||||
**Where:** §6.3.1; absent from §12.
|
||||
|
||||
Today one request per pool means one request's CUDA error is its own problem. Step-multiplexing changes the failure class categorically: a mid-step OOM/illegal-access/NaN from one request poisons the CUDA context and desyncs in-flight NCCL collectives for *every* co-scheduled tenant on the pool, including resident Dreamverse session caches. The doc designs none of the machinery: no SPMD-consistent abort broadcast (the dual of its scheduling broadcast), no request-fatal vs pool-fatal classification, no pool re-init + cache-invalidation policy, no partial-artifact semantics for fan-out graphs. "OOM" and request cancellation appear nowhere in 1799 lines; "abort" appears once (RL stragglers).
|
||||
|
||||
Ordinary cancellation is also missing — and vibe directing makes abandoning in-flight generations the *common* path. Worse, Phase 2 retires Dreamverse's `gpu_pool.py`, which today has a working sentinel-fd worker-death watch (`gpu_pool.py:542-586`), into engine-client calls — a reliability regression for the flagship customer if the gate ships as written. vLLM v1, the doc's own scheduler template, needed first-class machinery for exactly this (`ENGINE_CORE_DEAD`, `EngineDeadError`, `abort_requests`). Risk 4 covers only scheduling-decision divergence; the long-job-resilience known-gap is single-job-framed.
|
||||
|
||||
The abort path shapes the StepScheduler loop, the worker RPC surface, and CacheManager handle lifetimes — it must be designed *with* Phase 2, and by the doc's own standard ("absence reads as a decision"), this absence is an oversight.
|
||||
|
||||
---
|
||||
|
||||
## Major — reference-system misreads that drove recorded decisions
|
||||
|
||||
### M1. The single-BlockPool CacheManager rests on a property vLLM explicitly does not have: per-group page sizes
|
||||
**Where:** §6.3.2 lines 555-559.
|
||||
|
||||
The sentence asserts two mutually exclusive properties. vLLM's one-pool/no-fragmentation guarantee exists *only because* physical bytes-per-block are uniform across all groups: `kv_cache_utils.py` asserts a single page size (`get_uniform_page_size`), and its docstring says verbatim that breaking this "is non-trivial due to memory fragmentation concerns." Groups differ only in tokens-per-block at equal byte size; the unification mechanism inflates the smaller group's `block_size`.
|
||||
|
||||
Apply that to FastVideo's groups: a text-KV page (~64 KB/layer) vs a latent-frame slab (9.6–32 MB/layer for 1.3B/14B causal Wan) is a 150–500× ratio — unification means a 500-token reasoner prompt strands a multi-MB slab per layer-group. The one vLLM path with multiple page sizes (DeepseekV4) statically partitions capacity at startup over a single global block-id free list, which is harmless when group demand is token-coupled (every token passes through all layer groups) but wasteful exactly when demand is workload-decoupled — FastVideo's regime, where text-KV and chunk-KV demand vary independently with request mix.
|
||||
|
||||
Since this misread is what reversed the two-pool sketch (recorded at line 280), the decision rests on a false premise: either chunk-KV stays uniformly fine-paged (losing the slab semantics the MoT "falls out naturally" story depends on), or the two-pool design returns and needs its own fragmentation/deadlock argument.
|
||||
|
||||
### M2. diffusers Modular is not loop inversion — the "strongest external validation" of the keystone doesn't validate it
|
||||
**Where:** §5 line 277, §6.2.3 lines 426-428.
|
||||
|
||||
`LoopSequentialPipelineBlocks.__call__` raises `NotImplementedError`; every concrete family hand-writes `for i, t in enumerate(timesteps)` inside its own blocking wrapper (`wan/denoise.py:434`, `stable_diffusion_xl/denoise.py:701` — SDXL ships four such wrappers, the subclass forest again). The iteration is block-owned, invisible to any runtime — no init/step/finalize, no external driver, none of the properties §6.2.2 says inversion exists for (scheduling, interleaving, preemption, streaming, fair sharing). In scheduling terms it is the current `DenoisingStage` with a refactored body — i.e., it validates the Guiders/policy pillar but as evidence for inversion it is *equally consistent with the alternative the design rejects* ("keep loops in stages, make bodies pluggable"). The class also carries an explicit experimental warning.
|
||||
|
||||
Consequence: no surveyed system — vLLM, sglang, multimodal_gen, diffusers — implements runtime-owned diffusion iteration at scheduler granularity. Loop inversion is the design's most novel element with zero production precedent, and risk 3 (which admits novelty only for the hybrid AR+denoise slice) should say so instead of borrowing validation the reference doesn't provide.
|
||||
|
||||
### M3. Cost-currency scheduling drops the memory half of vLLM's admission — and memory is never a scheduling resource anywhere in the design
|
||||
**Where:** §6.3.1 (lines 476-547), §6.3.2; two lenses converged here.
|
||||
|
||||
vLLM's token budget is not a prediction — it is an exact cap checked *in the same loop as memory admission* (`allocate_slots` per request, preempt on allocation failure; activation memory separately bounded by a profiled worst case). The design takes the accounting structure, swaps the currency for a *forecast* (predicted GPU-time), and drops the memory dimension entirely: latents, conditioning sets, CFG duplicates, and activation peaks live in `RequestState`, explicitly outside the CacheManager, and nothing bounds how many concurrent LoopStates a pool admits — for a workload the doc itself calls memory-bound (line 499). Two items that each fit alone can jointly OOM, and a GPU-seconds currency cannot see it; combined with C1, that OOM is a pool-wide event. "Preemption only at step boundaries" never defines what happens to a preempted request's multi-GB resident state (offload? drop-and-resume-from-LoopState? — different economics from KV recompute).
|
||||
|
||||
Related internal contradiction, verified: cost is "static and known at admission... a table lookup" (line 538), but the same cost model is cache-dit-aware (line 520) — DBCache skip decisions are runtime data-dependent residual comparisons, unknowable at admission.
|
||||
|
||||
**Fix:** the budget needs a memory axis (resident-state + peak-activation per schedulable item), admission needs a memory planner over RequestState, and preemption semantics must be specified. The Phase-2 "≥2 sessions per GPU" gate rests on unaccounted memory until then.
|
||||
|
||||
### M4. Punica cannot express ComfyUI LoRA semantics
|
||||
**Where:** §9.4 lines 1376-1380 (also §6.3.2 lines 575-579). *Verifier rated minor-to-major; grouped here with the borrowings cluster.*
|
||||
|
||||
vLLM's `LoRARequest` carries one `lora_int_id` and no strength field; scaling is baked into `lora_b` at registration; the Punica wrapper maps one adapter index per token. ComfyUI traffic — the workload §9.4 names — is N stacked LoRAs per request with continuous user-set `strength_model` *and* `strength_clip`, routinely tweaked per generation. Pushing that through Punica means registering each (ordered-set, strengths) tuple as a synthetic concatenated adapter: near-zero cache-hit rate across strength tweaks, registration churn in the stacked GPU weight slots, and concatenated ranks colliding with `max_lora_rank`. "Strictly better than hot-swap-only" is unsupported without a composition layer that doesn't exist anywhere, including in vLLM.
|
||||
|
||||
---
|
||||
|
||||
## Major — execution-plane gaps
|
||||
|
||||
### M5. MoT mode multiplexing has no parallelism answer
|
||||
**Where:** §6.3.1 lines 502-503 vs §6.3.4; Phase 4 gate.
|
||||
|
||||
The "mode multiplexer" claim assumes both loop types share one static pool layout (`parallel: [dp, cfg, sp, tp]`), but their optimal layouts are disjoint: denoise wants SP+CFG; AR decode is sequence-length-1 — SP has nothing to shard and CFG doesn't exist. On a `[cfg(2), sp(4)]` 8-GPU pool the reasoner either runs replicated (1/8 useful work, paged KV duplicated 8×) or needs TP — and TP-everywhere regresses the bread-and-butter denoise workload on the flagship pool. Per-phase re-layout of the same resident weights is not expressible in the §6.3.4 spec (one static stack per pool), and resharding machinery exists only for train↔rollout weight sync (§8.6). §6.3.1's own jumbo-step mitigation (split cost classes across pools) is structurally unavailable for MoT — AR steps and denoise steps are the same weights — so concurrent reasoner token latency is gated by indivisible 50–500 ms denoise steps.
|
||||
|
||||
A workable resolution exists (AR continuous batching data-parallel across the cfg×sp weight-replica axes onto TP subgroups, plus §6.3.3 per-pathway TP, plus routing pure-REASON traffic to differently-shaped pools), but the doc never states one, and the Phase-4 gate ("reasoner ≥10× faster than re-prefill") is measured against an O(n²) strawman baseline that certifies nothing about pool efficiency. Risk 3's "prototype early in Phase 4" defers a *design contradiction*, not an implementation unknown.
|
||||
|
||||
### M6. The engine's own multi-node story is unstated, and the Ray executor silently disappears
|
||||
**Where:** N1 line 143, §6.3.5 line 673, §6.0 line 304.
|
||||
|
||||
Whether one worker pool may span nodes is a load-bearing decision the doc never makes — Dynamo routes *between* workers; it does not own the NCCL mesh *inside* one. If pools are single-node by fiat, SP degree caps at ~8 GPUs, directly contradicting line 543's jumbo-step mitigation ("shrink jumbo step wall-time with SP"), capping MoT model scale — and `RayDistributedExecutor`, today's shipping multi-node path, is silently dropped: it appears in the §3.1 diagram and then never again in §6, §10, §11, or §12 (violating the plan's own "every phase deletes or freezes what it replaces" discipline). If pools may span nodes, the engine owns cross-node collective bring-up, NCCL-timeout fault domains, and a multi-node health/drain contract — none designed, and C1's recovery problem becomes a multi-node recovery problem. Either answer changes Phase 2/3 scope. "Node-group" appears once, undefined.
|
||||
|
||||
### M7. Policies carry per-request mutable state with no state-scoping contract — and the doc contradicts itself on when policies are resolved
|
||||
**Where:** §6.2.3 lines 412-417 vs §6.2.2 line 387 vs risk 2 line 1621; §6.4 lines 837-840.
|
||||
|
||||
The doc says policies are resolved at pipeline build (lines 412-413; risk 2: "resolved to bound methods at build time") *and* in `DenoiseLoop.init` (line 387) — a genuine contradiction on a load-bearing contract. It matters: AdaptiveGateCFG — a named CFGPolicy example and the Wan2.2 worked-example default — is per-request mutable state in shipped code (`denoising.py:338-343, 507-551`: `delta_cached`, `delta_cached_model_id`, gate counters). Build-time-resolved singletons mean request A's cached CFG delta gets applied to request B the moment Phase 2 interleaving lands — silent quality corruption no Phase-2 gate (load tests, latency budget) can catch. This is the *exact* failure mode §6.4 cites to justify interceptor state scoping ("silently corrupts under concurrent requests") — the contract was designed for the plugin tier and forgotten for the policy tier, which sits on a hotter path. Cheap fix (policy state into LoopState, same as plugins), but it must be in the spec.
|
||||
|
||||
### M8. The six-policy taxonomy does not factor the shipped step bodies — no step skeleton or cross-policy interaction contract is defined
|
||||
**Where:** §6.2.3 (policy table, line 424 claim); §6.2.2 lines 386-389; §6.4 lines 837-844.
|
||||
|
||||
The proposed step is three phases (forward → CFG combine → scheduler step); the shipped loops need ~six, with dependencies that cross policy boundaries. Verified examples:
|
||||
|
||||
- **Cosmos** conditioning-frame injection consumes the *sampler's* EDM coefficients, applies per-CFG-branch both pre-forward (input mix) and post-forward (x0 clamp), and the CFG combine runs in x0 space — ConditioningInjector × Sampler × CFGPolicy interleaved inside each branch, unownable by any one of them (`denoising.py:845-933`).
|
||||
- **TI2V** clamps latents *after* `scheduler.step` — a post-step constraint with no policy slot (`denoising.py:570-573`).
|
||||
- **Cosmos2.5** builds per-frame timestep vectors with a conditioned-frame override and re-clamps GT every step pre-forward.
|
||||
- **CausalDMD** renoises between steps choosing `add_noise` vs `add_noise_high` by expert boundary — Sampler × ExpertRouting (`causal_denoising.py:268-301`).
|
||||
- **AdaptiveGateCFG** must observe ExpertRouting's switch to invalidate its delta (today an inline `id(current_model)` check) — yet no channel for one policy to observe another is defined anywhere.
|
||||
- **LTX2** guidance is 1–4 runtime-decided passes whose branches alter the network via forward kwargs (`skip_cross_modal_attn`, `skip_video/audio_self_attn_blocks`) — colliding with BlockInterceptor's domain in a way the "two block-skippers conflict" pre-flight check cannot see, and breaking §6.4's per-CFG-branch state scoping, which assumes a fixed cond/uncond branch vocabulary (`ltx2_denoising.py:503-605, 620-631`).
|
||||
|
||||
None of the six policies covers prediction-space conversion, per-token timestep construction, post-step latent constraints, inter-step renoising, or chunk-boundary refresh. The fix is not abandoning policies — the Sampler registry is the natural home for some of this, and composition still strips the duplicated offload/attn-metadata/autocast/trajectory plumbing — but the design needs the fixed step skeleton with ordered, typed extension points and an explicit policy-interaction contract, worked through Cosmos2.5 and LTX2 *in the doc*. Until then, "a new model contributes policies + a graph spec; it does not edit shared loop code" (line 424) is asserted, not demonstrated.
|
||||
|
||||
### M9. OmniRequest cannot parameterize multi-loop graphs
|
||||
**Where:** §6.1 lines 318-334; §6.6 line 905; worked examples (c)(d) lines 943-950.
|
||||
|
||||
One flat `SamplingParams` + one flat `DiffusionParams` per request, while the design's own flagship examples are multi-loop graphs needing per-node knobs: LTX-2's refine loop has its own step count and guidance scale *today* as first-class fields (`fastvideo_args.py:204-205`, threaded through `compat.py` and `dynamo/examples/diffusers/worker.py:201-203`); a thinker and talker need different `max_tokens`/`temperature`/`stop`. No request→graph-node parameter binding is defined anywhere; the only escape hatch is line 905's per-model `ModelOptions` blocks — i.e., the `ltx2_*` field-leakage pattern the doc indicts at P3, with a type wrapper, regenerated into the OpenAI/CLI views that derive from the request schema (line 907). Needs a real decision — parameters keyed by graph-node id, or per-node override blocks validated against the PipelineSpec — made in Phase 0, because that schema ships first and external consumers build against it.
|
||||
|
||||
---
|
||||
|
||||
## Major — caches and weights
|
||||
|
||||
### M10. No feature-cache invalidation story under LoRA hot-swap — te-LoRAs make the embedding cache serve stale embeddings in the workflow cloud
|
||||
**Where:** §6.3.2 lines 570-574 vs §9.4 lines 1349-1380.
|
||||
|
||||
The only invalidation rule in the document is RL `update_weights` → `reset()`. But ComfyUI-grade LoRAs routinely patch the *text encoder* alongside the DiT (`comfy/lora.py` maintains `lora_te/lora_te1/lora_te2` key maps; `load_lora_for_models` takes a separate `strength_clip`), so a content-hash-keyed embedding cache returns embeddings computed under the wrong adapter state the moment two workflows share a prompt but differ in te-LoRA stacks — silent wrong output in the exact product (§9.4 "exact mode") whose trust claim is reproducibility. §11.8 even makes cross-request embedding reuse load-bearing as the radix-cache substitute. And once Punica-style batched multi-LoRA lands, requests with different adapter stacks coexist concurrently on one pool, so the cache must be key-*partitioned* by (encoder identity × adapter set × strengths), not flushed — a different design from the `EncoderCacheManager` reset() semantics being adopted, which come from a world where encoders are never patched per request. The key schema needs a weight-state epoch / adapter-set hash as a mandatory component, decided before Phase 3.
|
||||
|
||||
### M11. Checkpoint/LoRA patching mutates pool-shared weights — a pool-quiescing barrier the StepScheduler has no vocabulary for
|
||||
**Where:** §9.4 lines 1371-1380 vs §6.3.1 and §6.0 line 299.
|
||||
|
||||
Components are "one resident copy per worker pool"; patch/unpatch mutates that copy, which is global to every loop interleaved on the pool — yet step-interleaving is the engine's core Phase-2 value. Two interleaved loops requiring different patch states cannot coexist, so every cross-group transition is a drain barrier: finish in-flight steps, apply/undo `W += scale·BA` across 14–28 GB shard-consistently across TP/SP ranks (ComfyUI keeps weight backups for the undo — 2× weight memory or a CPU→GPU restore at PCIe seconds), re-admit. Under workflow-cloud traffic (long-tail checkpoints, per-request adapter stacks), transition frequency is the whole game — and the §6.3.1 cost model (lines 516-521) has no weight-state-transition term, no notion of weight state as schedulable state, and no quiesce-vs-queue policy, even though transition cost is exactly what A1 checkpoint-affinity routing must weigh. The §8.6 safe-point-swap pattern shows the doc knows the shape but never applies it here. §9.4 calls this "the one real new subsystem"; §12 carries no risk entry for it.
|
||||
|
||||
---
|
||||
|
||||
## Major — training/RL
|
||||
|
||||
### M12. "Step bodies are plain tensor programs, so autograd composes" is contradicted by the distillation code the substrate must absorb
|
||||
**Where:** §8.2 lines 1016-1019; §6.2.2; §6.3.2.
|
||||
|
||||
Self-forcing does not "drive `DenoiseLoop.step`": its rollout samples per-block exit indices broadcast across ranks, runs no-grad steps to the exit, runs exactly *one* grad-enabled forward, then a separate no-grad `store_kv=True` context-caching pass with context noise, gated by `start_gradient_frame`. None of this fits `init/step/finalize` + `StepResult(done, emit)` without grad-gating flags, per-step cache-write control, and per-block exit policies — training-only surface in substrate code, or the method keeps its own loop and the "3 copies → 1" dedup claim dies for the hardest case. The KV path needs grad/AC-aware semantics the engine pool lacks: today's causal model snapshots KV indices whenever `torch.is_grad_enabled()` so activation-checkpoint recompute doesn't double-advance the cache (`wan_causal.py:119-120,405-431`), and never recycles blocks mid-rollout — while §6.3.2 specs vLLM-style out-of-window block recycling, and §8.5's own profile taxonomy says "training forward … *no caches*," showing the grad+KV case was never designed. §8.3 explicitly stakes the architecture on ChunkKVPool serving self-forcing training.
|
||||
|
||||
(Note: the related forward-context-backward attack was refuted — the Phase-1 retirement of the global plus explicit metadata passing *helps* autograd composition. The surviving residue is the grad-window/cache-mode design above.)
|
||||
|
||||
### M13. Behavior Record cost is understated ~1.5 orders of magnitude for its own flagship case (MoE diffusion)
|
||||
**Where:** §8.5 lines 1156-1160; §5 miles row line 1088.
|
||||
|
||||
The miles ~60 MB/sample figure is per-token routing, one forward per generated token. Diffusion re-routes the *entire packed sequence at every denoise step, twice under CFG*: the record is steps × CFG × tokens × MoE-layers × top_k. For a Cosmos3-class request (Qwen3-VL-MoE config: 60 experts, top_k 4, ~24 sparse layers via `decoder_sparse_step=1`, ~50K packed tokens, 35-50 steps × 2 branches) that is ~1.3–1.9 GB/sample int32 — ~20–30 GB per 16-sample GRPO group, before latents. "Cheap because trajectory capture is already an OutputSpec feature" conflates plumbing cost with byte cost; at these sizes the Record forces a buffering/transport/storage design (GB-scale trajectories through connectors from disaggregated rollout fleets) that appears nowhere — not in §8.7's TrajectoryBuffer, not in §12, not in the known-gaps list. (The RNG-draws sub-claim was refuted: seeded generators in a single shared loop reproduce draws; uint8 expert IDs also cut 4×. The routing-record problem stands.)
|
||||
|
||||
### M14. The omni-RL pilot is a Phase-4 deliverable with no objective design
|
||||
**Where:** §8.7 lines 1236-1240; §10 Phase 4.
|
||||
|
||||
The section establishes *expressibility* (one trajectory, two segment types — true, and a real structural advantage over engine-per-stage stacks) and quietly upgrades it to a deliverable without posing the algorithm problem:
|
||||
|
||||
- **Scale mismatch:** token log-probs are O(1–10) nats over 10²–10³ tokens; per-step diffusion SDE log-probs are Gaussian densities over 10⁶–10⁷ latent dims — any joint clipped-ratio objective needs principled per-segment normalization that none of the cited recipes (FlowGRPO/DanceGRPO/NFT/AIPO/GSPO) provides; get it wrong and one modality silently dominates the shared trunk.
|
||||
- **Credit assignment:** the reasoner influences video reward only through *sampled discrete tokens* re-entering as conditioning — a non-differentiable boundary, so token segments get sparse trajectory-level REINFORCE signal while denoise segments get dense per-step ratios, both updating shared attention-trunk weights, with no interference analysis.
|
||||
- **Reasoning regression:** RL-updating the und pathway on video-reward-correlated signal risks degrading its reasoning; reference-model KL anchoring for hybrid episodes is never mentioned.
|
||||
|
||||
The entire treatment is the phrase "optimized with mixed objectives," and §12's 15 open questions contain nothing on it — for the capability marketed as "the capability nobody else has." Either it gets an algorithm sketch and an open-question entry with an owner, or the Phase-4 item should be demoted from "pilot" to "trajectory capture demonstrated."
|
||||
|
||||
---
|
||||
|
||||
## Major — the migration plan (the weakest section)
|
||||
|
||||
### M15. The "frozen legacy stack" premise is empirically false in this very repo
|
||||
**Where:** lines 5, 110, 1026; §11.4; risk 5.
|
||||
|
||||
The anti-third-stack defense is a declared freeze plus intent to delete — and this repo has already run that experiment and it failed within weeks. Verified from git: `fastvideo/train/` landed 2026-03-09 (#1159); since then **19 commits modified the "frozen" `fastvideo/training/`**, including a *brand-new* `cosmos2_5_training_pipeline.py` added to the legacy stack on 2026-05-11 (#1227) — **nine days after `training/AGENTS.md` explicitly forbade adding new models there**, and eleven days after the same model landed in `train/` (#1224). World-model training (#1179) and LongCat finetuning (#1244) also landed in the frozen stack in May; EMA bugfixes as recently as June 8-9; `AGENTS.md` still calls `training/` "authoritative for shipped models."
|
||||
|
||||
The doc invokes the training/-vs-train/ "lesson" but proposes nothing mechanically different from what was tried: no CI gate rejecting new files under legacy paths, no codeowner veto, no named owner per family, no calendar date for Phase 5. "Phase 5 is a scheduled deletion, not an aspiration" (risk 5) — but nothing in the document is scheduled. Under the same model-port pressure that broke the training/ freeze (measurably higher on the inference side), this freeze breaks the same way. Name the enforcement mechanism that did not exist last time, or the deprecation commitment is the prior failure restated with more confidence.
|
||||
|
||||
### M16. Phase dependency inversion: Phases 1–2 consume the substrate Phase 4 builds
|
||||
**Where:** §10 lines 1406-1446 vs §6.3.1 lines 487-489, §6.3.2; three lenses converged on this.
|
||||
|
||||
Phase 1 migrates causal Wan ("exercises chunk-KV"); Phase 2 ships "AR continuous batching" — which §6.3.1 *constitutively defines* as "(continuous batching; paged KV; chunked prefill)"; the CacheManager owning both lands in Phase 4, and risk 3 even defers the StepScheduler+KVPool prototype to "early in Phase 4," contradicting Phase 2. Compounding it: **no AR-pathway model exists on the new runtime before the Phase-4 Cosmos3 re-port** (Wan-causal is chunked denoise, not token AR; thinkers/talkers are Phase 4), so Phase 2's headline deliverable has neither a cache backing nor a workload — and none of Phase 2's gates (lines 1427-1430) tests AR batching.
|
||||
|
||||
The Phase-1 half is softenable: an interim per-request chunk-KV behind the unchanged `KVHandle` seam, with a Phase-4 allocator swap, is normal incremental staging — but the doc never states this, and its own "no third stack / every phase deletes what it replaces" principle cuts against unstated throwaway implementations. Fix structurally: pull a CacheManager v0 (chunk-KV slabs + minimal paged text-KV) into Phases 1–2, or move AR batching to Phase 4 and rewrite the Phase-2 gate to what it actually exercises.
|
||||
|
||||
### M17. Phase 4 re-ports a baseline that is not on main, and the plan schedules neither its merge nor its rebase
|
||||
**Where:** §10 Phase 0 line 1405, Phase 4 lines 1439-1446; §1 lines 42-49; Appendix.
|
||||
|
||||
`fastvideo/pipelines/basic/cosmos3/` on main contains only `__pycache__` — the design's forcing function exists solely as the unmerged 5-branch stacked chain (`feat/cosmos3-tier-a-port` → … → `feat/cosmos3-reasoning`). Phase 0's "Cosmos3 audio leaves `batch.extra`" cannot execute against main: it presupposes the chain is merged (a major-model review effort the plan never schedules) or means maintaining the migration on a side branch, continuously rebased across the most churn-heavy refactors in the repo's history (ForwardBatch→RequestState, loop inversion, executor→engine) — months of conflict-resolution work, unowned and unsized, on the artifact whose 150/150 bit-exactness is the design's proudest credential and whose parity suite the Phase-4 gate requires ("every phase ships green" cannot apply to a suite that is not in the tree). The plan sequences other in-flight work explicitly (`fastvideo/api/` in Phase 0, PR #1438 in Phase 1) but skips this. Needs an explicit merge milestone before Phase 0 touches the port.
|
||||
|
||||
### M18. G5's enforcement instrument has holes: ~6-7 shipped families have no SSIM test, and the CI-cost mitigation is incoherent for substrate PRs
|
||||
**Where:** G5 lines 128-129; Phase 0 gate line 1405; risk 6.
|
||||
|
||||
`fastvideo/tests/ssim/` covers ~14 of 20+ families. Cosmos(2/2.5), Hunyuan, Hunyuan15(+SR), HYWorld, MagiHuman, Waypoint, and MatrixGame-v1 have no SSIM test — "all SSIM suites unchanged" passes *vacuously* for roughly a third of shipped pipelines, exactly the ones sitting on the shared loop being refactored. And risk 6's "gated to touched families" mitigation is designed for model-local PRs; Phases 0–2 are by construction not model-local — the ForwardBatch adapter, loop inversion, and executor replacement sit under every family, so "touched families" = all of them on precisely the riskiest PRs. Either substrate PRs run the full GPU matrix (a cost the plan should budget — SSIM runs on Modal L40S today) or gating quietly degrades to sampling, which is how regressions slip through. Needs: a reference-seeding work item before Phase 1, or G5 restated as "zero regression for the SSIM-covered subset," plus a stated per-phase GPU-CI budget.
|
||||
|
||||
### M19. Phase 5's deletion milestone breaks the "frozen and untouched" legacy training/ stack
|
||||
**Where:** lines 5-6, 144, 1026-1027 vs Phase 5 line 1448.
|
||||
|
||||
The frozen stack is a live consumer of exactly the code Phase 5 deletes: `fastvideo/training/training_pipeline.py:39` imports `ComposedPipelineBase`/`ForwardBatch`/`LoRAPipeline`, holds `validation_pipeline: ComposedPipelineBase`, and its validation instantiates real legacy pipelines that run the legacy `DenoisingStage`; `distillation_pipeline.py:31` likewise. So Phase 5 cannot remove `ComposedPipelineBase` and `DenoisingStage` while leaving `training/` untouched — either the deletion milestone hollows to "delete except what legacy training/ needs" (the old path never dies — the very smell being fixed) or the scope statement is false and `training/` breaks on this plan's schedule. Relatedly, "loop inversion makes the step functions the single shared implementation" is arithmetically 3→2, not 3→1: the legacy inlined copies are out of scope forever. The doc needs an explicit answer: what happens to `fastvideo/training/` at Phase 5?
|
||||
|
||||
### M20. "Retire `fastvideo/forward_context.py` (Phase 1)" is infeasible as scheduled
|
||||
**Where:** §6.3.3 lines 618-621; Phase 1 lines 1412-1414; vs N2/N4; Appendix line 1791.
|
||||
|
||||
194 references across ~50 files. The global is read inside `fastvideo/attention/layer.py` — the shared Attention module on *every* family's hot path — and set in 27 places inside the frozen `training/` stack (8 module-level imports). Phase 1 migrates only Wan+Flux2; the other ~16 families run "unmodified" behind the legacy adapter (N4) and still set the global. So in Phase 1 the file cannot be deleted (touches the frozen stack, violating N2; breaks every unmigrated family), and `attention/layer.py` must serve both worlds simultaneously — a dual-sourcing branch in the hottest shared layer, undesigned. The honest description: Phase 1 *adds a second context mechanism beside the global*, and the global survives until Phase 5 at the earliest — where the deliverables list never mentions it. Appendix A states "retired Phase 1" as accomplished fact. Rewrite as "new-path-only StageContext; `forward_context` frozen for legacy consumers; deletion gated on Phase 5," and design the dual-mechanism cost.
|
||||
|
||||
### M21. §10 is a dependency ordering, not a plan — no timeline, no staffing, no sizing, and no policy for the ~1-2 new model ports per month that arrive during the migration
|
||||
**Where:** §10; N4 line 153; risk 1.
|
||||
|
||||
The scope — typed I/O, loop inversion + policies, extension system, async engine + StepScheduler + online-calibrated cost model, four-class CacheManager, PackedSeq/MoT layers, declarative parallelism compiler, workflow compiler, RL layer, Dynamo contract, config collapse — is plainly multi-engineer-years, with zero dates, headcount, per-phase sizing, or owners; "by Phase 2" decision deadlines (§11.1, risks 7/15) are unanchored because Phase 2 is not a date.
|
||||
|
||||
The sharper, unanswered problem is **inflow**: git shows ~1–2 new families landing per month (Flux2 Klein and Lucy Edit on 2026-06-09 alone; MatrixGame3 05-27; MagiHuman 05-12; Stable Audio 05-01; Gen3C 04-01…). Over multi-quarter Phases 0–4, another 10–15 models arrive, and the doc never says what they target: land them on legacy abstractions and the Phase-5 tail grows faster than phases retire it (negative net migration velocity); force them onto the new stack and every port blocks on machinery that doesn't exist until Phase 1/3/4. Either answer materially changes the plan; choosing neither means the terminal state recedes indefinitely. Minimum fix: per-phase engineer-month estimates, a named owner per phase, a calendar target for Phase 5, and an explicit "new ports target the new stack starting at Phase X" rule with its porting-velocity cost stated.
|
||||
|
||||
---
|
||||
|
||||
## Major — product/trust surfaces
|
||||
|
||||
### M22. Per-request plugin enablement is an unsandboxed third-party-code and noisy-neighbor surface; only workflow JSON is named untrusted
|
||||
**Where:** §6.4 lines 859-861 vs §12 input-hardening gap lines 1693-1695.
|
||||
|
||||
Entry-point plugins execute arbitrary code inside the serving engine, and the doc makes their selection part of the *request* (`diffusion.plugins=[{"name": "cache_dit", "Fn": 8, "Bn": 8}]`) in the same engine pitched as a multi-tenant cloud — and since the OpenAI protocol is *generated from the request schema* (lines 907-908), the field derives into the public API with no carve-out. Consequences forcing a design change: (a) **correctness** — a caller can attach a distribution-altering interceptor to a request the product has labeled "exact mode" (the §9.4 trust claim), or pass unvalidated kwargs into third-party code; (b) **isolation** — a `needs_eager` observer on one request drops compile/cudagraph capture for scopes shared with co-scheduled tenants (line 809), a noisy-neighbor vector with no cost attribution anywhere in the metrics design; (c) **supply chain** — entry-point resolution imports whatever package claims the name. The needed contract: enablement/allowlisting at DeployConfig scope only; requests merely parameterize pre-enabled plugins against per-plugin validated schemas; plugin overhead attributed per-request in the cost model. §12's input-hardening gap names only workflow JSON — a categorically different surface.
|
||||
|
||||
### M23. No versioning or stability contract for the serialized schemas shipped to external consumers mid-migration
|
||||
**Where:** §6.4 line 861; §6.6 lines 920-927; §10; open question 12.
|
||||
|
||||
By Phase 3 there are at least four externally consumed serialized surfaces: hub-published ModelSpec manifests (interchange with diffusers' `modular_model_index.json` — a format co-owned with an external party), compiled-workflow PipelineSpecs (content-hash-keyed in the weight-fleet cache — schema changes silently change hashes and invalidate fleet affinity), the OmniEvent streaming schema (Dreamverse's frontend; proposed as Dynamo ask A3's wire format), and per-model ModelOptions blocks. Phase 4 then lands PackedSeq, session-scoped inputs, and the Cosmos3 re-port — guaranteed churn after consumers exist. The migration plan gates *behavior* at every phase (SSIM, parity, load) and gates *interfaces* at none; the only versioning commitment in the document is hook-point names (open question 12 is scoped to hook points). Without per-surface decisions now — `schema_version` fields, frozen-vs-experimental tiers per phase, a deprecation window — Phase 4 either breaks published artifacts or gets paralyzed by accidental freezing. G5 protects only the Python `VideoGenerator` call.
|
||||
|
||||
---
|
||||
|
||||
## Minor (confirmed)
|
||||
|
||||
1. **ForwardBatch has 111 fields, not ~250** (AST-verified; stated twice, lines 33/188). P3 survives at 111, but the headline metric is inflated 2.3× in a doc that brands its pain points "evidence-backed" — it invites discounting of the numbers that *do* verify exactly (1381 lines and 35 probes both check out).
|
||||
2. **"Prediction is a table lookup" vs the design's own flagship features** (§6.3.1 vs §6.4): DBCache/FBCache/TaylorSeer decide per step from runtime residual similarity — a stochastic per-step cost multiplier unknowable at admission; VSA tile selection is content-dependent; and AR decode lengths are unbounded (the doc concedes vLLM "must guess decode lengths," then silently exempts its own AR group).
|
||||
3. **Worked example (g) is internally contradictory**: cache-dit + C1 + "identical trajectories" are pairwise incompatible under §8.5's own `distribution_altering` contract (§8.7 states the rule correctly: cache acceleration is C0). Matters because (g) is the template PR #1438 is told to target in Phase 1.
|
||||
4. **The Phase-2 Dreamverse gate is untestable as written**: at ~4.55 s GPU-saturating per 5 s clip (line 1263), "≥2 concurrent sessions per GPU at unchanged segment latency" is only passable under an unstated think-time/collision-rate assumption — the gate can be passed or failed at will by choosing the test's session behavior. More broadly, no quantitative multiplexing target (sessions/GPU under a stated load profile, GPU-utilization, cost/clip) exists anywhere, so there is no way to conclude after Phase 2 whether step-level scheduling earned its complexity over the §11.6-rejected simpler design.
|
||||
5. **The exec summary launders Dynamo contingencies into outcomes** (line 75: "each with a fallback — so Dynamo fronts both production serving and RL rollout fleets"): the body is honest (A1–A7 with fallbacks; §11.9; §12.15), but the asks are unfiled RFCs on an NVIDIA-governed roadmap; A5's own fallback "weakens fleet-scale async RL," and if A3 misses Phase 2, Dreamverse ships on the direct-WebSocket bypass and the production-hardened fallback becomes permanent — the exact "permanent workaround" dynamic §11.9 claims the direct relationship avoids. Ask-sequencing (§12.15) has no owner or decision dates.
|
||||
6. **diffusers as "convergent validation" cuts both ways** (see M2): its four-wrappers-per-family shape is the subclass forest again; the citation supports the rejected alternative as well as the chosen one.
|
||||
7. **Punica/ComfyUI LoRA semantics gap** — see M4.
|
||||
|
||||
---
|
||||
|
||||
## Fact-check corrections
|
||||
|
||||
70 concrete claims were checked; **none was fabricated**; 13 need correction. Everything else verified, including the claims most likely to be embellished: vLLM RFC #42770 (author/date/content/two-tier resolution), PR #42304 **merged** 2026-05-16 with `VLLM_USE_BREAKABLE_CUDAGRAPH`, vllm-omni RFC #4084, the Thinking Machines numbers (80/1000 unique outputs, divergence at token 103, 26s→42s, KL results), the Dynamo worker's `asyncio.Lock`, cache-dit, the cosmos-framework MoT details (PackedAttentionMoT, MoTDecoderLayer, ReasonerKVCache, MoE gen-MLP), miles/verl-omni/sglang-omni mechanics, sglang's cache-dit monkeypatch scars, and `enable_teacache` genuinely having no consumer.
|
||||
|
||||
| # | design.md says | Reality |
|
||||
|---|---|---|
|
||||
| 1 | "1381-line `DenoisingStage`" (lines 34, 201) | 1381 is the **file**; the class is ~670 lines (47–715) plus 6 subclasses in-file. The 35-probe count is exact for the file. |
|
||||
| 2 | "~250-field ForwardBatch" (33, 188) | **111 fields** (whole file incl. TrainingBatch/PreprocessBatch: ~153). |
|
||||
| 3 | "19 denoising-stage classes" (201) | **22** model/variant classes (+ base = 23); the list omits Magi-class and two other same-category stages predating the doc. |
|
||||
| 4 | "Cosmos2.5 clamping … hardcoded in the shared loop" (201) | Clamping lives in the `Cosmos25DenoisingStage` **subclass**; the Wan2.2 expert switch (`denoising.py:229-235, 352-376`) and TI2V inline VAE encode (`:239-268, 399-404, 570-572`) are in the shared loop as claimed. |
|
||||
| 5 | `SamplingParam` "~170 fields" (887) | **75**. The ~170 figure belongs to TrainingArgs (90 own + 81 inherited = 171). |
|
||||
| 6 | `FastVideoArgs` "~96 fields" (885) | **81** (TrainingArgs subclassing claim correct). |
|
||||
| 7 | "TP and SP (Ulysses/ring)" (192) | Main is **Ulysses-only** (`all_to_all_4D`); no ring-attention SP is wired into FastVideo. |
|
||||
| 8 | CFG "3 copies: `stages/conditioning.py` vs …" (993) | Right count, wrong citation: the inference-stack copy is in `denoising.py`, not `conditioning.py`. |
|
||||
| 9 | ComfyUI "~45 `comfy_extras` packs", "90+ blueprints" (1335-1339) | **117** packs (matching nodes.py's 117-entry registration list); **80** in-tree blueprints (the larger library ships via the registry). 64 core nodes, 39 API providers, GPL-3.0, FIFO-no-batching all verify. |
|
||||
| 10 | kv-router events "`{sequence_hash, block_hash, removed}`" (707, A1 733-739) | Paraphrase: actual shape is `KvCacheEventData::Stored{parent_hash, blocks[{block_hash, tokens_hash}]}` / `Removed` / `Cleared` (`protocols.rs:627-646`). Token-prefix-derived keying verifies. |
|
||||
| 11 | miles TIS clamp "to `[0.5, 2.0]`" (1086) | Configurable `[tis_clip_low, tis_clip]`, CLI defaults [0, 2.0]; the 0.5/2.0 pair comes from the MIS example config (`mis.yaml`). |
|
||||
| 12 | sglang-omni "`DllmScheduler` for a DiT talker" (269) | DllmScheduler serves the **LLaDA2-Uni thinker** (diffusion-LLM); the DiT talker is Ming-Omni's, on a different scheduler. |
|
||||
| 13 | `_iter_packed_batches` under `model/vfm/` (236); §11.3's claim that the port's "own status notes" list reasoning-KV/batching/streaming/prefix-reuse as "missing for production" | Lives at `cosmos_framework/inference/inference.py:66`. PORT_STATUS.md confirms 150/150 but contains no such missing-for-production list — that framing is the design doc's own and should not be attributed to the port's status notes. |
|
||||
|
||||
---
|
||||
|
||||
## Attacks that failed (the doc survives these)
|
||||
|
||||
The refute-by-default verifiers killed 36 findings, several of them attacks a hostile reviewer would lead with — worth knowing they don't land:
|
||||
|
||||
- **ChunkRollout/DenoiseLoop nesting is expressible** in the stated Stage/LoopStage/StepResult contracts ("one solver step / one token / one chunk" + composition).
|
||||
- **N1 vs engine-internal pools** is consistent on a careful read (N1 is about datacenter orchestration; §6.3.5 states the reconciliation).
|
||||
- **The trainer-scope line (N2 vs §8)** is drawn consistently — N2's own text enumerates exactly what §8 changes.
|
||||
- **G6 vs the ≤2% Phase-2 gate** is goal-vs-acceptance-gate, not contradiction (Phase 1 is gated bit-identical).
|
||||
- **The clean-room GPL posture holds**: sampler/scheduler math (DPM-Solver, Karras sigmas, flow-match shift) is published outside GPL sources.
|
||||
- **C2 for the video denoise path is fine**: batch-1 fixed shapes are trivially batch-invariant — the doc's own analysis at lines 1145-1147 is correct; the AR/image/sharding exposures are correctly identified there too.
|
||||
- **Self-forcing's cross-chunk gradients truncate by construction** (KV written under `no_grad` on detached context), so the engine KV pool is not blocked the way one might fear — the surviving residue is M12's grad-window cache mode.
|
||||
- **"Every phase deletes or freezes something" survives audit** at the phase-deliverable level (the failures are the specific items in M19/M20).
|
||||
- **The tier-1 ComfyUI vocabulary claim survives** blueprint-corpus measurement under the doc's actual claim (curated canonical workflows, not top-N node frequency).
|
||||
- **The sglang reconvergence deferral** is substantively defended in §11.1 with reasons valid under either outcome.
|
||||
- **WeightSyncPlan's "literal no-op"** is correctly scoped to colocated same-layout in the doc's own sentence; FSDP-vs-TP/SP is explicitly routed to in-place reshard.
|
||||
|
||||
---
|
||||
|
||||
## Ranked recommendations
|
||||
|
||||
1. **Design the abort/cancellation/OOM path with Phase 2** (C1) **and add memory as a budget axis with admission planning and preemption semantics** (M3). These two are the soundness conditions of the multiplexing bet; everything else in the execution plane sits on them.
|
||||
2. **Re-derive §6.3.2 from the real vLLM constraint** (M1). The two-pool→one-pool reversal was made on a false premise; either accept uniform page bytes (and redesign the slab story) or bring back two pools with an explicit fragmentation/deadlock argument.
|
||||
3. **Fix the migration plan's three structural defects**: CacheManager v0 into Phases 1–2 or AR batching out of Phase 2 (M16); a merge milestone for the cosmos3 chain before Phase 0 touches it (M17); a new-port inflow rule plus a freeze-enforcement mechanism that did not exist last time — CI path gate, codeowners, a date (M15, M21). Also reconcile Phase 5 with the frozen `training/` stack (M19) and restate the `forward_context` retirement honestly (M20).
|
||||
4. **Specify the step skeleton and the policy contracts** — ordered, typed extension points; a policy state-scoping rule (state in LoopState, like plugins); a policy-observation channel — and work the mapping through Cosmos2.5 and LTX2 in the doc (M7, M8). Decide per-node request parameter binding in Phase 0 (M9).
|
||||
5. **Give MoT a stated parallelism answer** (M5) and make the single-pool-spans-nodes decision explicit, including the fate of `RayDistributedExecutor` (M6).
|
||||
6. **Close the workflow-cloud trust/correctness holes before Phase 3**: adapter-aware feature-cache keys (M10), weight-state transitions as a scheduled, costed operation (M11), DeployConfig-scoped plugin allowlisting (M22), per-surface schema stability tiers (M23), and an honest assessment of Punica's fit (M4).
|
||||
7. **Right-size the RL claims**: design the grad+KV cache mode or scope self-forcing out of the shared loop (M12); budget the Behavior Record at real byte counts (M13); demote the omni-RL pilot or give it an objective sketch and an owner (M14); fix worked example (g).
|
||||
8. **Reclassify loop inversion as unprecedented at scheduler granularity** in risk 3 and drop the diffusers "validation" (M2). The bet may still be right — but it should be made with open eyes, and the parity-gate plan is then carrying more weight than the doc admits.
|
||||
9. **Correct the thirteen numbers above before circulating.** The doc's credibility rests on its "evidence-backed" brand; ~250-vs-111 is the kind of error that makes a reader re-check everything else — and most of everything else checks out.
|
||||
@@ -1,3 +0,0 @@
|
||||
__pycache__/
|
||||
*.pyc
|
||||
*.pyo
|
||||
-105
@@ -1,105 +0,0 @@
|
||||
# Handoff — GPU bring-up of the v2 torch backend
|
||||
|
||||
**For: an agent on a GPU box, branched from `will/mini-fastvideo`.**
|
||||
**Your job:** take the *written-not-run* `cuda` backend to *runs-and-generates*, then commit + push.
|
||||
|
||||
Everything below is committed on `will/mini-fastvideo` and CPU-tested (**204 tests pass**). The torch
|
||||
path was authored on a machine with **no GPU and no torch**, so it is grounded in the real
|
||||
`fastvideo` APIs and cross-checked against the source, but **never executed**. That's what you finish.
|
||||
|
||||
---
|
||||
|
||||
## 0. Orientation (read these first, in order)
|
||||
|
||||
1. **`v2/README.md`** — what the whole v2 mini is (the `(recipe, runtime)` runtime; "architecture is
|
||||
real, kernels are toys"). The "Honest scope" paragraph says exactly what's wired.
|
||||
2. **`v2/platform/backends/GPU_BRINGUP.md`** — *your checklist*: the ordered 10-step bring-up + the
|
||||
risk table (A–G), each tied to a `# BRINGUP` marker in the source. **This handoff is orientation +
|
||||
process; GPU_BRINGUP.md is the work.**
|
||||
3. This file — the meta-instructions (verify bar, commit/push, gotchas).
|
||||
|
||||
## 1. What's already done (commits on this branch)
|
||||
|
||||
```
|
||||
d6d0580a [fix] correct GPU adapters against real fastvideo API (cross-check findings)
|
||||
b8d78f40 [feat] real torch/CUDA backend (written-not-run) behind the cuda cells
|
||||
27791b51 [feat] static-buffer capture form for the cudagraph step body (Path A)
|
||||
ae6a170d [feat] piecewise CUDA-graph capture/replay at the step boundary (Path A)
|
||||
9308d87e [feat] route diffusion loops through the kernel table
|
||||
7490c590 [feat] multi-backend dispatch substrate (device/arch/kernel registries)
|
||||
```
|
||||
|
||||
The dispatch substrate (two tuple-keyed registries `COMPONENTS(kind,device,variant)` +
|
||||
`KERNELS(op,device,arch,variant)`, a detected `Platform`, numpy terminal + parity oracle), the
|
||||
universal kernel seam (diffusion loops go through `model.platform.kernels`), the
|
||||
piecewise cudagraph lifecycle, and the torch backend cells are all in place. On a GPU box,
|
||||
`Platform.detect()` returns a `cuda` platform and resolves the torch cells instead of the numpy toys —
|
||||
**the inference loops/policies/scheduler are unchanged**; only the resolved implementations differ.
|
||||
|
||||
## 2. The files you'll touch
|
||||
|
||||
| File | What it is |
|
||||
|---|---|
|
||||
| `v2/platform/backends/torch_adapters.py` | `TorchWanDiT` / `TorchWanVAE` / `TorchT5Encoder` — wrap the real `fastvideo.models.*` (named by each card's `load_id`) to the mini's duck-typed surface. Built via the real FastVideo loaders. |
|
||||
| `v2/platform/backends/torch_kernels.py` | torch `flow_match_step` / `flow_sde_step` (plain elementwise — there is **no** fused solver kernel in fastvideo-kernel; don't look for one). |
|
||||
| `v2/platform/backends/torch_cuda.py` | registers the `cuda` cells as lazy trampolines (torch imported only inside builder bodies). |
|
||||
| `v2/card/specs.py` | `ComponentSpec.checkpoint` — the per-component weights source (empty on toys; **you fill it in**). |
|
||||
|
||||
The surface the adapters must honor (what the loops call):
|
||||
`dit(latent, text_embed, sigma) -> velocity` · `vae.decode(latent)` / `vae.encode(video)` ·
|
||||
`text_encoder.encode(text)`. The CPU toys in `v2/models/backend.py` are the reference behavior.
|
||||
|
||||
## 3. Your task (the gating items — full detail in GPU_BRINGUP.md)
|
||||
|
||||
1. **Env:** install `torch` + the parent `fastvideo` package + weights. (`fastvideo` source lives at
|
||||
`/Users/willlin/src/FastVideo`.)
|
||||
2. **Risk A — the one blocking gap:** the builders call `_load_via_fastvideo(...)` → the real loaders
|
||||
need a **`FastVideoArgs`**, which `_fastvideo_args(spec)` builds minimally from `spec.checkpoint`.
|
||||
Confirm/extend its fields (model config, precision, parallelism). And stamp `ComponentSpec.checkpoint`
|
||||
onto the wan21 card — a tiny helper that maps a model root onto the three components is the cleanest
|
||||
way (the toy cards leave it `""`).
|
||||
3. **Work the risk list (A–G in GPU_BRINGUP.md).** The *interface* contracts were cross-checked as
|
||||
matching (DiT returns bare velocity; `timestep=sigma*1000`; `encode().mode()`; `.last_hidden_state`;
|
||||
no fused solver kernel) — confirm them numerically. The *construction* layer was fixed (real loaders,
|
||||
`set_forward_context`, latent normalization, UMT5-from-config). What's left is box-dependent:
|
||||
`FastVideoArgs` fields, `shift_factor` placement/sign, exact tokenizer kwargs, FSDP sharding.
|
||||
4. **Bring up in order:** build each component in isolation → one DiT step → one solver step → VAE
|
||||
decode → full t2v → SDE stochastic sampling → cudagraph capture (last).
|
||||
|
||||
## 4. The verification bar (how you know it's right)
|
||||
|
||||
- **CPU suite must stay green:** `python3 -m pytest v2/ -q` → still **204 passed**. The torch path is
|
||||
gated `available=False` off-GPU; importing the backends must never import torch. If you break either,
|
||||
you broke the substrate. (`v2/tests/test_torch_backend.py` pins these.)
|
||||
- **Parity oracle is the spec:** the substrate's whole point is that a real backend matches the numpy
|
||||
reference on the consistency ladder. On GPU, compare a full generation against a known-good fastvideo
|
||||
output — use the parent repo's SSIM regression harness (`fastvideo/tests/ssim/`). Target C4 (SSIM /
|
||||
artifact quality); component/trajectory parity (C0/C1) is bit-level vs the reference pipeline.
|
||||
- **Don't trust "it ran" — trust "it matched."** A wrong `timestep` scale or `shift_factor` produces
|
||||
plausible-but-wrong video, not a crash (risks B/D). Diff against a reference, don't eyeball.
|
||||
|
||||
## 5. Commit + push
|
||||
|
||||
- **You are on a GPU branch** (branched from `will/mini-fastvideo`). Commit your bring-up fixes there,
|
||||
focused by concern (e.g. one commit per confirmed risk), in the existing style (`[fix]`/`[feat] …`).
|
||||
- **NEVER add Claude as a co-author** (repo policy, `/Users/willlin/src/.claude/CLAUDE.md`).
|
||||
- **Do not rewrite or force-push** the six commits above — build on top.
|
||||
- When the CPU suite is green **and** a GPU generation matches the reference, **push your branch.**
|
||||
- If you launch inference with wandb logging enabled, log in with the token in the project
|
||||
`CLAUDE.md` (`/Users/willlin/src/.claude/CLAUDE.md`) — **do not paste it into any committed file.**
|
||||
|
||||
## 6. Gotchas (don't relearn these the hard way)
|
||||
|
||||
- **No fused solver kernel exists.** `fastvideo-kernel` ships only attention/norm/quant primitives;
|
||||
the cuda `flow_match_step`/`flow_sde_step` are plain torch by design. Don't hunt for a `.cu` solver.
|
||||
- **`from_pretrained` is not the loader.** `WanTransformer3DModel`/`AutoencoderKLWan` have none — the
|
||||
real path is the `*Loader().load(model_path, fastvideo_args)` classes in
|
||||
`fastvideo/models/loader/component_loader.py`. The loader resolves the class from the checkpoint
|
||||
config (this is what makes UMT5-vs-T5 correct without hardcoding).
|
||||
- **T5 needs `set_forward_context`.** A bare encoder forward reads stale/None global context.
|
||||
- **The loop surface stays numpy** for bring-up; adapters marshal numpy↔torch at the boundary. A
|
||||
torch-native surface (latent on-device through forward→combine→solver) is the **perf follow-up**
|
||||
(Risk G), not bring-up — don't rewrite `cfg.combine`/`precision.cast`/the samplers yet.
|
||||
- **cudagraph capture is last.** The wan21 loop declares `breakable_cudagraph`; the v2 capturer models
|
||||
the lifecycle with a numpy `StaticWorkspace`. Capturing a real `torch.cuda.CUDAGraph` is GPU-only
|
||||
work and the riskiest step — leave it until inference is verified.
|
||||
-148
@@ -1,148 +0,0 @@
|
||||
# FastVideo v2 - Inference Runtime Scope
|
||||
|
||||
**Status:** source of truth for `v2/`.
|
||||
|
||||
`v2` is the model-native inference runtime for FastVideo. It owns model cards,
|
||||
programs, loops, runtime execution, serving, cache/memory policy, backend dispatch,
|
||||
compile/cudagraph integration, and inference parity checks.
|
||||
|
||||
`v2` does **not** own training, finetuning, distillation, RL, optimizer steps, or
|
||||
checkpoint production. Training remains in the existing FastVideo stacks:
|
||||
|
||||
- `fastvideo/train/` - the new modular trainer.
|
||||
- `fastvideo/training/` - the legacy shipped training pipelines.
|
||||
|
||||
`v2` may record how a checkpoint was produced through `RecipeSpec` metadata
|
||||
(`method`, `parents`, `assumes_loop`, `assumes_precision`), because inference must
|
||||
know which runtime loop and precision policy a post-training checkpoint expects.
|
||||
That metadata is provenance, not a v2 training API.
|
||||
|
||||
## Design Center
|
||||
|
||||
FastVideo v2 is video-generation first. Wan/LTX-style diffusion video inference
|
||||
is the baseline path, and unified models such as BAGEL/Cosmos3 are first-class:
|
||||
one resident model may run AR, diffusion, VAE, and codec loops in one request.
|
||||
Audio and TTS are supported as additional modalities on the same stage/loop
|
||||
model, not as the reason to build a separate universal serving framework.
|
||||
|
||||
The core abstraction should stay small:
|
||||
|
||||
- `ModelCard` declares the resident components, loops, capabilities, precision,
|
||||
caches, and sampling defaults a checkpoint needs.
|
||||
- `Program` is an ordered list of typed nodes passing values through named
|
||||
slots. It is not a general DAG, Walk graph, or declarative control-flow IR.
|
||||
- `Loop` owns model semantics. The runtime drives the loop, handles admission,
|
||||
cancellation, streaming, cache access, and backend dispatch.
|
||||
|
||||
Do not add a public contract field until a runtime path consumes it. Future
|
||||
optimizations such as richer stage placement, multi-GPU transport, or paged KV
|
||||
should start behind a concrete Wan/BAGEL/Cosmos/Qwen use case and graduate only
|
||||
after they simplify at least two model recipes.
|
||||
|
||||
## Scope
|
||||
|
||||
In scope:
|
||||
|
||||
- Python inference entrypoint through `v2.VideoGenerator`.
|
||||
- Typed `ModelCard` declarations for components, loops, capabilities, parity,
|
||||
sampling defaults, precision, and checkpoint layout.
|
||||
- Driven inference loops such as diffusion denoise, AR decode, causal/world
|
||||
continuation, VAE/audio decode, and multi-stage programs.
|
||||
- Runtime execution through `Engine` and `AsyncEngine`.
|
||||
- Serving through OpenAI-compatible HTTP/SSE surfaces and deployment cards.
|
||||
- Backend dispatch through the CPU toy backend, accelerator stand-ins, and the
|
||||
real torch/CUDA backend.
|
||||
- Inference acceleration features such as FP8/NVFP4 loading, Sage/Flash/SDPA
|
||||
attention backend selection, `torch.compile`, cudagraph capture, cache policy,
|
||||
and component placement.
|
||||
- Inference parity and regression tests.
|
||||
|
||||
Out of scope:
|
||||
|
||||
- Training methods, optimizers, loss functions, RL rewards, rollout trainers,
|
||||
weight-sync training loops, and behavior records for policy updates.
|
||||
- Training examples under `v2_examples/`.
|
||||
- Any CLI/API that advertises v2 as a trainer.
|
||||
- A universal graph runtime, Walk/state-machine authoring layer, or parallelism
|
||||
vocabulary that is not consumed by the current inference runtime.
|
||||
|
||||
## Core Model
|
||||
|
||||
The atomic inference artifact is a `(recipe, runtime)` pair:
|
||||
|
||||
- `RecipeSpec` records what the weights assume: parent checkpoints, post-training
|
||||
method name, required loop, and required precision.
|
||||
- `ModelCard` declares the runtime surface: components, loops, capabilities,
|
||||
caches, precision, parallelism, sampling defaults, and checkpoint manifest.
|
||||
- `Program` composes component nodes and loop nodes into a user-facing task as
|
||||
an ordered named-slot stage list.
|
||||
- `ModelInstance` is the resident loaded card with shared components, caches,
|
||||
weight versions, and optional captured graphs.
|
||||
|
||||
This keeps post-training artifacts serveable without making `v2` responsible for
|
||||
creating them.
|
||||
|
||||
## Execution Model
|
||||
|
||||
Loops are model-owned state machines:
|
||||
|
||||
```python
|
||||
state = loop.init(req, model, ctx)
|
||||
while True:
|
||||
plan = loop.next(state)
|
||||
if isinstance(plan, Done):
|
||||
break
|
||||
result = ctx.execute(plan)
|
||||
state = loop.advance(state, result)
|
||||
return loop.finalize(state)
|
||||
```
|
||||
|
||||
The loop owns semantics. The runtime owns execution, admission, cancellation,
|
||||
streaming, cache access, graph capture, and backend dispatch.
|
||||
|
||||
Serving is pooled run-to-completion. `AsyncEngine` bounds concurrency by pool
|
||||
slots; each request runs its program to completion. The synchronous `Engine` is
|
||||
the offline path used by tests and `VideoGenerator`.
|
||||
|
||||
## Package Layout
|
||||
|
||||
```text
|
||||
v2/
|
||||
video_generator.py public inference facade
|
||||
registry.py model id -> card builder registry
|
||||
core/
|
||||
card/ ModelCard, specs, ModelInstance
|
||||
loop/ loop contracts, driver, sampler, policies
|
||||
program/ task programs and workflows
|
||||
request/ request params, tasks, outputs, sessions
|
||||
parity/ inference parity helpers
|
||||
parallel/ named parallel plans
|
||||
recipes/ model-specific cards, loops, and programs
|
||||
runtime/ Engine, AsyncEngine, cache, memory, cudagraph, transport
|
||||
serving/ HTTP/SSE server and deployment adapters
|
||||
platform/ backend/device/kernel dispatch
|
||||
_vendor/ vendored FastVideo model/loader/config pieces for inference
|
||||
tests/ v2 inference/runtime/serving/parity tests
|
||||
```
|
||||
|
||||
There is intentionally no `v2/training/` package.
|
||||
|
||||
## Current Inference Path
|
||||
|
||||
The torch backend builds real components from stamped checkpoint paths, keeps
|
||||
components in eval mode, and dispatches inference through the same cards and loops
|
||||
used by the CPU tests. Wan2.1 T2V inference is the primary real path today.
|
||||
Wan/FastWan inference supports:
|
||||
|
||||
- real Wan component loading through vendored component loaders,
|
||||
- FP8 post-load quantization for `FastVideo/FastWan-QAD-FP8-1.3B`,
|
||||
- attention backend selection, including SageAttention when installed,
|
||||
- `torch.compile` for inference DiT modules,
|
||||
- on-device latent residency for cards that set `device_io=True`.
|
||||
|
||||
## Boundary Rule
|
||||
|
||||
If a change adds training behavior, it belongs in `fastvideo/train/` or
|
||||
`fastvideo/training/`, not in `v2/`. If inference needs to consume the result of
|
||||
that training, add or update a v2 card, loop, registry entry, checkpoint loader,
|
||||
sampling defaults, and inference tests.
|
||||
@@ -1,89 +0,0 @@
|
||||
"""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.
|
||||
|
||||
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.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
from v2.core.enums import (
|
||||
Capability,
|
||||
ConsistencyLevel,
|
||||
ExecutionProfile,
|
||||
LoopKind,
|
||||
WorkUnitKind,
|
||||
)
|
||||
from v2.core.card import (
|
||||
CapabilityMatrix,
|
||||
ComponentSpec,
|
||||
LoopSpec,
|
||||
ModelCard,
|
||||
ModelInstance,
|
||||
ParitySpec,
|
||||
RecipeSpec,
|
||||
load_card,
|
||||
)
|
||||
from v2.core.program import ComponentNode, ModelLoopNode, Program, ProgramKind, when_opt, when_task
|
||||
from v2.core.request import (
|
||||
DiffusionParams,
|
||||
Output,
|
||||
Request,
|
||||
SamplingParams,
|
||||
Session,
|
||||
TaskType,
|
||||
make_request,
|
||||
)
|
||||
from v2.runtime import AsyncEngine, Engine
|
||||
|
||||
__version__ = "0.2.0"
|
||||
|
||||
__all__ = [
|
||||
"ModelCard",
|
||||
"ComponentSpec",
|
||||
"LoopSpec",
|
||||
"RecipeSpec",
|
||||
"ParitySpec",
|
||||
"CapabilityMatrix",
|
||||
"ModelInstance",
|
||||
"load_card",
|
||||
"Engine",
|
||||
"AsyncEngine",
|
||||
"Program",
|
||||
"ProgramKind",
|
||||
"ComponentNode",
|
||||
"ModelLoopNode",
|
||||
"when_task",
|
||||
"when_opt",
|
||||
"Request",
|
||||
"Session",
|
||||
"Output",
|
||||
"make_request",
|
||||
"TaskType",
|
||||
"SamplingParams",
|
||||
"DiffusionParams",
|
||||
"LoopKind",
|
||||
"WorkUnitKind",
|
||||
"ConsistencyLevel",
|
||||
"ExecutionProfile",
|
||||
"Capability",
|
||||
"VideoGenerator",
|
||||
"__version__",
|
||||
]
|
||||
|
||||
|
||||
def __getattr__(name: str):
|
||||
# Lazy: the GPU entrypoint imports torch / fastvideo, so resolve it only on access — plain
|
||||
# ``import v2`` (and the CPU-only mini) stay torch-free.
|
||||
if name == "VideoGenerator":
|
||||
from v2.video_generator import VideoGenerator
|
||||
return VideoGenerator
|
||||
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
|
||||
@@ -1 +0,0 @@
|
||||
"""Vendored fastvideo code (copied, standalone). Internal layout mirrors upstream for diffing; v2-native code must not edit these ad hoc."""
|
||||
@@ -1,28 +0,0 @@
|
||||
"""Slim vendored API config surface for the v2 VideoGenerator.
|
||||
|
||||
Only the inference-config dataclasses (schema) + result types are vendored. The fastvideo
|
||||
parser / presets / overrides modules are intentionally NOT vendored — they pull the fastvideo
|
||||
pipeline runtime, which v2 replaces. See v2/README.md (vendoring)."""
|
||||
from __future__ import annotations
|
||||
|
||||
from v2._vendor.api.results import GenerationResult
|
||||
from v2._vendor.api.schema import (
|
||||
CompileConfig,
|
||||
EngineConfig,
|
||||
GenerationRequest,
|
||||
GeneratorConfig,
|
||||
OffloadConfig,
|
||||
OutputConfig,
|
||||
SamplingConfig,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"CompileConfig",
|
||||
"EngineConfig",
|
||||
"GenerationRequest",
|
||||
"GeneratorConfig",
|
||||
"OffloadConfig",
|
||||
"OutputConfig",
|
||||
"SamplingConfig",
|
||||
"GenerationResult",
|
||||
]
|
||||
@@ -1,16 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
from __future__ import annotations
|
||||
|
||||
|
||||
class ConfigValidationError(ValueError):
|
||||
"""Validation error that keeps track of the nested config path."""
|
||||
|
||||
def __init__(self, path: str, message: str):
|
||||
self.path = path
|
||||
self.message = message
|
||||
super().__init__(str(self))
|
||||
|
||||
def __str__(self) -> str:
|
||||
if self.path:
|
||||
return f"{self.path}: {self.message}"
|
||||
return self.message
|
||||
@@ -1,15 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
from dataclasses import dataclass
|
||||
|
||||
from v2._vendor.api.sampling_param import SamplingParam
|
||||
|
||||
|
||||
@dataclass
|
||||
class MatrixGame2SamplingParam(SamplingParam):
|
||||
height: int = 352
|
||||
width: int = 640
|
||||
num_frames: int = 57
|
||||
fps: int = 25
|
||||
guidance_scale: float = 1.0
|
||||
num_inference_steps: int = 3
|
||||
negative_prompt: str | None = None
|
||||
@@ -1,233 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""Track which GenerationRequest fields the user explicitly provided.
|
||||
|
||||
When translating a GenerationRequest into a legacy SamplingParam we must
|
||||
distinguish user-provided values (which should override model defaults)
|
||||
from schema defaults (which should NOT override model defaults).
|
||||
|
||||
The mechanism: a single ``_fastvideo_explicit_paths`` set stored on the
|
||||
root ``GenerationRequest``. It holds dotted leaf paths (e.g.
|
||||
``"sampling.guidance_scale"``) the user has touched, either via raw
|
||||
config at bind time or via attribute assignment at runtime. A patched
|
||||
``__setattr__`` on the request dataclass types records assignments into
|
||||
this set.
|
||||
|
||||
The set holds leaf paths only. Nested dataclass or mapping assignments
|
||||
are flattened to their leaves at record time.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Callable, Mapping
|
||||
import dataclasses
|
||||
from typing import Any, cast
|
||||
|
||||
from v2._vendor.api.schema import (
|
||||
ContinuationState,
|
||||
GenerationPlan,
|
||||
GenerationRequest,
|
||||
InputConfig,
|
||||
OutputConfig,
|
||||
PlannedStage,
|
||||
RequestRuntimeConfig,
|
||||
RunConfig,
|
||||
SamplingConfig,
|
||||
ServeConfig,
|
||||
)
|
||||
|
||||
EXPLICIT_PATHS_ATTR = "_fastvideo_explicit_paths"
|
||||
|
||||
_TRACKING_ROOT_ATTR = "_fastvideo_request_tracking_root"
|
||||
_TRACKING_PATH_ATTR = "_fastvideo_request_tracking_path"
|
||||
_TRACKING_PATCHED_ATTR = "_fastvideo_request_tracking_patched"
|
||||
|
||||
_TRACKED_REQUEST_TYPES = (
|
||||
GenerationRequest,
|
||||
InputConfig,
|
||||
SamplingConfig,
|
||||
RequestRuntimeConfig,
|
||||
OutputConfig,
|
||||
ContinuationState,
|
||||
PlannedStage,
|
||||
GenerationPlan,
|
||||
)
|
||||
|
||||
|
||||
def bind_generation_request_raw(
|
||||
request: GenerationRequest,
|
||||
raw: Mapping[str, Any] | None,
|
||||
) -> GenerationRequest:
|
||||
"""Install explicit-path tracking on *request*.
|
||||
|
||||
*raw* is the parsed config dict (YAML/JSON/kwargs); every leaf key
|
||||
in it becomes an explicit path. Subsequent attribute assignments on
|
||||
*request* or its nested dataclasses are recorded automatically via a
|
||||
patched ``__setattr__``.
|
||||
"""
|
||||
_ensure_request_tracking()
|
||||
# Disable recording while we walk the tree to install roots.
|
||||
object.__setattr__(request, EXPLICIT_PATHS_ATTR, None)
|
||||
_set_tracking_roots(request, request, "")
|
||||
paths: set[str] = set()
|
||||
_record_value_paths(raw or {}, "", paths)
|
||||
object.__setattr__(request, EXPLICIT_PATHS_ATTR, paths)
|
||||
return request
|
||||
|
||||
|
||||
def bind_run_config_raw(
|
||||
config: RunConfig,
|
||||
raw: Mapping[str, Any],
|
||||
) -> RunConfig:
|
||||
request_raw = raw.get("request")
|
||||
if isinstance(request_raw, Mapping):
|
||||
bind_generation_request_raw(config.request, request_raw)
|
||||
else:
|
||||
bind_generation_request_raw(config.request, {})
|
||||
return config
|
||||
|
||||
|
||||
def bind_serve_config_raw(
|
||||
config: ServeConfig,
|
||||
raw: Mapping[str, Any],
|
||||
) -> ServeConfig:
|
||||
default_request_raw = raw.get("default_request")
|
||||
if isinstance(default_request_raw, Mapping):
|
||||
bind_generation_request_raw(config.default_request, default_request_raw)
|
||||
else:
|
||||
bind_generation_request_raw(config.default_request, {})
|
||||
return config
|
||||
|
||||
|
||||
def get_explicit_paths(request: GenerationRequest) -> frozenset[str]:
|
||||
"""Return a snapshot of the explicit paths set on *request*."""
|
||||
paths = getattr(request, EXPLICIT_PATHS_ATTR, None)
|
||||
if isinstance(paths, set | frozenset):
|
||||
return frozenset(paths)
|
||||
return frozenset()
|
||||
|
||||
|
||||
def reset_tracking_roots(request: GenerationRequest) -> None:
|
||||
"""Re-install tracking roots after a deepcopy or manual clone.
|
||||
|
||||
The paths set itself deepcopies correctly; we only need to repoint
|
||||
the tracking root on nested dataclasses at the new root.
|
||||
"""
|
||||
_ensure_request_tracking()
|
||||
_set_tracking_roots(request, request, "")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Path recording
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _record_value_paths(
|
||||
value: Any,
|
||||
prefix: str,
|
||||
out: set[str],
|
||||
) -> None:
|
||||
"""Add every leaf path under *value* to *out*.
|
||||
|
||||
A leaf is any terminal value (non-dataclass, non-mapping, or empty
|
||||
mapping/dataclass). ``prefix`` is the dotted path at which *value*
|
||||
sits. When called with an empty ``prefix`` (the root), leaves are
|
||||
recorded at their own key.
|
||||
"""
|
||||
if dataclasses.is_dataclass(value) and not isinstance(value, type):
|
||||
dc_fields = dataclasses.fields(value)
|
||||
if not dc_fields:
|
||||
if prefix:
|
||||
out.add(prefix)
|
||||
return
|
||||
for field in dc_fields:
|
||||
child = getattr(value, field.name)
|
||||
path = f"{prefix}.{field.name}" if prefix else field.name
|
||||
_record_value_paths(child, path, out)
|
||||
return
|
||||
if isinstance(value, Mapping):
|
||||
if not value:
|
||||
if prefix:
|
||||
out.add(prefix)
|
||||
return
|
||||
for key, child in value.items():
|
||||
path = f"{prefix}.{key}" if prefix else key
|
||||
_record_value_paths(child, path, out)
|
||||
return
|
||||
if prefix:
|
||||
out.add(prefix)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# __setattr__ patching
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _ensure_request_tracking() -> None:
|
||||
for config_type in _TRACKED_REQUEST_TYPES:
|
||||
_patch_tracking_setattr(config_type)
|
||||
|
||||
|
||||
def _patch_tracking_setattr(config_type: type[Any]) -> None:
|
||||
if getattr(config_type, _TRACKING_PATCHED_ATTR, False):
|
||||
return
|
||||
|
||||
original_setattr = cast(
|
||||
Callable[[Any, str, Any], None],
|
||||
config_type.__setattr__,
|
||||
)
|
||||
field_names = {field.name for field in dataclasses.fields(config_type)}
|
||||
|
||||
def _tracking_setattr(self: Any, name: str, value: Any) -> None:
|
||||
if name.startswith("_fastvideo_") or name not in field_names:
|
||||
original_setattr(self, name, value)
|
||||
return
|
||||
|
||||
original_setattr(self, name, value)
|
||||
|
||||
root = getattr(self, _TRACKING_ROOT_ATTR, None)
|
||||
if root is None:
|
||||
return
|
||||
paths = getattr(root, EXPLICIT_PATHS_ATTR, None)
|
||||
if not isinstance(paths, set):
|
||||
return
|
||||
|
||||
prefix = getattr(self, _TRACKING_PATH_ATTR, "")
|
||||
path = f"{prefix}.{name}" if prefix else name
|
||||
# Wholesale dataclass replacement: install roots on the new
|
||||
# instance so its future mutations are tracked too.
|
||||
if dataclasses.is_dataclass(value) and not isinstance(value, type):
|
||||
_set_tracking_roots(root, value, path)
|
||||
_record_value_paths(value, path, paths)
|
||||
|
||||
type.__setattr__(config_type, "__setattr__", _tracking_setattr)
|
||||
setattr(config_type, _TRACKING_PATCHED_ATTR, True)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tree walk to set tracking root/path on nested dataclasses
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _set_tracking_roots(
|
||||
root: GenerationRequest,
|
||||
obj: Any,
|
||||
prefix: str,
|
||||
) -> None:
|
||||
if not dataclasses.is_dataclass(obj) or isinstance(obj, type):
|
||||
return
|
||||
object.__setattr__(obj, _TRACKING_ROOT_ATTR, root)
|
||||
object.__setattr__(obj, _TRACKING_PATH_ATTR, prefix)
|
||||
for field in dataclasses.fields(obj):
|
||||
child = getattr(obj, field.name)
|
||||
child_path = f"{prefix}.{field.name}" if prefix else field.name
|
||||
if dataclasses.is_dataclass(child) and not isinstance(child, type):
|
||||
_set_tracking_roots(root, child, child_path)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"EXPLICIT_PATHS_ATTR",
|
||||
"bind_generation_request_raw",
|
||||
"bind_run_config_raw",
|
||||
"bind_serve_config_raw",
|
||||
"get_explicit_paths",
|
||||
"reset_tracking_roots",
|
||||
]
|
||||
@@ -1,173 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any
|
||||
from collections.abc import Mapping
|
||||
|
||||
from v2._vendor.api.schema import ContinuationState
|
||||
|
||||
|
||||
@dataclass
|
||||
class GenerationResult:
|
||||
prompt: str | None = None
|
||||
prompt_index: int | None = None
|
||||
samples: Any | None = None
|
||||
frames: Any | None = None
|
||||
audio: Any | None = None
|
||||
audio_sample_rate: int | None = None
|
||||
size: tuple[int, int, int] | None = None
|
||||
generation_time: float | None = None
|
||||
logging_info: Any | None = None
|
||||
trajectory: Any | None = None
|
||||
trajectory_timesteps: Any | None = None
|
||||
trajectory_decoded: Any | None = None
|
||||
video_path: str | None = None
|
||||
peak_memory_mb: float | None = None
|
||||
state: ContinuationState | None = None
|
||||
extra: dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
@classmethod
|
||||
def from_legacy_result(
|
||||
cls,
|
||||
result: Mapping[str, Any],
|
||||
) -> GenerationResult:
|
||||
prompt = result.get("prompt")
|
||||
if prompt is None:
|
||||
prompt = result.get("prompts")
|
||||
|
||||
extra = {
|
||||
key: value
|
||||
for key, value in result.items() if key not in {
|
||||
"prompt",
|
||||
"prompt_index",
|
||||
"prompts",
|
||||
"samples",
|
||||
"frames",
|
||||
"audio",
|
||||
"audio_sample_rate",
|
||||
"size",
|
||||
"generation_time",
|
||||
"logging_info",
|
||||
"trajectory",
|
||||
"trajectory_timesteps",
|
||||
"trajectory_decoded",
|
||||
"video_path",
|
||||
"peak_memory_mb",
|
||||
"state",
|
||||
}
|
||||
}
|
||||
|
||||
return cls(
|
||||
prompt=prompt,
|
||||
prompt_index=result.get("prompt_index"),
|
||||
samples=result.get("samples"),
|
||||
frames=result.get("frames"),
|
||||
audio=result.get("audio"),
|
||||
audio_sample_rate=result.get("audio_sample_rate"),
|
||||
size=result.get("size"),
|
||||
generation_time=result.get("generation_time"),
|
||||
logging_info=result.get("logging_info"),
|
||||
trajectory=result.get("trajectory"),
|
||||
trajectory_timesteps=result.get("trajectory_timesteps"),
|
||||
trajectory_decoded=result.get("trajectory_decoded"),
|
||||
video_path=result.get("video_path"),
|
||||
peak_memory_mb=result.get("peak_memory_mb"),
|
||||
state=result.get("state"),
|
||||
extra=extra,
|
||||
)
|
||||
|
||||
def to_legacy_dict(self) -> dict[str, Any]:
|
||||
result = {
|
||||
"prompts": self.prompt,
|
||||
"samples": self.samples,
|
||||
"frames": self.frames,
|
||||
"audio": self.audio,
|
||||
"audio_sample_rate": self.audio_sample_rate,
|
||||
"size": self.size,
|
||||
"generation_time": self.generation_time,
|
||||
"logging_info": self.logging_info,
|
||||
"trajectory": self.trajectory,
|
||||
"trajectory_timesteps": self.trajectory_timesteps,
|
||||
"trajectory_decoded": self.trajectory_decoded,
|
||||
"video_path": self.video_path,
|
||||
"peak_memory_mb": self.peak_memory_mb,
|
||||
}
|
||||
if self.prompt_index is not None:
|
||||
result["prompt_index"] = self.prompt_index
|
||||
result["prompt"] = self.prompt
|
||||
if self.state is not None:
|
||||
result["state"] = self.state
|
||||
result.update(self.extra)
|
||||
return result
|
||||
|
||||
|
||||
# Alias the canonical result type; matches the public docs.
|
||||
VideoResult = GenerationResult
|
||||
|
||||
|
||||
@dataclass
|
||||
class VideoProgressEvent:
|
||||
"""Per-step progress event emitted by :meth:`VideoGenerator.generate_async`.
|
||||
|
||||
Consumers treat these as best-effort telemetry; ``total_steps`` is
|
||||
the count the pipeline reported at the start of the run, not a
|
||||
rolling estimate.
|
||||
"""
|
||||
|
||||
step: int
|
||||
total_steps: int
|
||||
stage: str = "denoise"
|
||||
"""Logical stage name (``denoise`` | ``refine`` | ``decode`` | …)."""
|
||||
|
||||
|
||||
@dataclass
|
||||
class VideoPartialEvent:
|
||||
"""Chunk of decoded frames ready for streaming.
|
||||
|
||||
Emitted only on the streaming path; the aggregated code path never
|
||||
yields partials. ``frames`` is a numpy ``(N, H, W, 3)`` uint8
|
||||
ndarray; ``index`` is a monotonic chunk index starting at 0.
|
||||
"""
|
||||
|
||||
frames: Any
|
||||
index: int
|
||||
|
||||
|
||||
@dataclass
|
||||
class VideoFinalEvent:
|
||||
"""Terminal event carrying the generated video and metadata.
|
||||
|
||||
Exactly one ``VideoFinalEvent`` is emitted per request. When
|
||||
``request.output.return_state`` is True the event also carries the
|
||||
:class:`ContinuationState` the caller needs to resume.
|
||||
"""
|
||||
|
||||
video_bytes: bytes | None = None
|
||||
tensor: Any | None = None
|
||||
frames: Any | None = None
|
||||
metadata: dict[str, Any] = field(default_factory=dict)
|
||||
continuation_state: ContinuationState | None = None
|
||||
result: VideoResult | None = None
|
||||
"""The full :class:`VideoResult` for callers that want everything.
|
||||
|
||||
Streaming consumers typically only care about ``frames`` /
|
||||
``continuation_state``; keeping the full result here avoids a
|
||||
second code path."""
|
||||
|
||||
|
||||
VideoEvent = VideoProgressEvent | VideoPartialEvent | VideoFinalEvent
|
||||
"""Union of every event :meth:`VideoGenerator.generate_async` yields.
|
||||
|
||||
Consumers match by ``isinstance`` rather than ``type`` so subclasses
|
||||
(e.g. a future ``VideoAudioSegmentEvent``) slot in without breaking
|
||||
existing code."""
|
||||
|
||||
__all__ = [
|
||||
"GenerationResult",
|
||||
"VideoEvent",
|
||||
"VideoFinalEvent",
|
||||
"VideoPartialEvent",
|
||||
"VideoProgressEvent",
|
||||
"VideoResult",
|
||||
]
|
||||
@@ -1,411 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
from __future__ import annotations
|
||||
|
||||
import copy
|
||||
from dataclasses import dataclass, field, fields
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
from v2._vendor.logger import init_logger
|
||||
from v2._vendor.utils import StoreBoolean
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from v2._vendor.api.schema import ContinuationState
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class SamplingParam:
|
||||
"""
|
||||
Sampling parameters for video generation.
|
||||
"""
|
||||
# All fields below are copied from ForwardBatch
|
||||
data_type: str = "video"
|
||||
|
||||
# Image inputs
|
||||
image_path: str | None = None
|
||||
pil_image: Any | None = None
|
||||
|
||||
# Video inputs
|
||||
video_path: str | None = None
|
||||
|
||||
# Optional pre-generated diffusion latents. Used by parity/debug harnesses
|
||||
# and advanced callers that need deterministic latent reuse.
|
||||
latents: Any | None = None
|
||||
|
||||
# Action control inputs (Matrix-Game)
|
||||
mouse_cond: Any | None = None # Shape: (B, T, 2)
|
||||
keyboard_cond: Any | None = None # Shape: (B, T, K)
|
||||
grid_sizes: Any | None = None # Shape: (3,) [F,H,W]
|
||||
|
||||
# Camera control inputs (HYWorld)
|
||||
pose: str | None = None # Camera trajectory: pose string (e.g., 'w-31') or JSON file path
|
||||
prompt_attention_mask: list = field(default_factory=list)
|
||||
negative_attention_mask: list = field(default_factory=list)
|
||||
|
||||
# Camera/action control inputs (GameCraft)
|
||||
camera_states: Any | None = None # Plücker coordinates [B, T_video, 6, H, W]
|
||||
camera_trajectory: str | None = None
|
||||
action_list: list[str] | None = None
|
||||
action_speed_list: list[float] | None = None
|
||||
gt_latents: Any | None = None # Ground truth latents [B, 16, T, H, W]
|
||||
conditioning_mask: Any | None = None # Mask [B, 1, T, H, W]
|
||||
|
||||
# Camera control inputs (LingBotWorld)
|
||||
c2ws_plucker_emb: Any | None = None # Plucker embedding: [B, C, F_lat, H_lat, W_lat]
|
||||
|
||||
# Refine inputs (LongCat 480p->720p upscaling)
|
||||
# Path-based refine (load stage1 video from disk, e.g. MP4)
|
||||
refine_from: str | None = None # Path to stage1 video (480p output from distill)
|
||||
t_thresh: float = 0.5 # Threshold for timestep scheduling in refinement
|
||||
spatial_refine_only: bool = False # If True, only spatial (no temporal doubling)
|
||||
num_cond_frames: int = 0 # Number of conditioning frames
|
||||
# In-memory refine input (for two-stage pipeline where stage1 frames are already in memory)
|
||||
# This mirrors LongCat's demo where a list of frames (e.g. np.ndarray or PIL.Image)
|
||||
# is passed directly to the refinement pipeline instead of reloading from disk.
|
||||
stage1_video: Any | None = None
|
||||
|
||||
# Text inputs
|
||||
prompt: str | list[str] | None = None
|
||||
negative_prompt: str = "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards"
|
||||
max_sequence_length: int | None = None
|
||||
prompt_path: str | None = None
|
||||
output_path: str = "outputs/"
|
||||
output_video_name: str | None = None
|
||||
|
||||
# Batch info
|
||||
num_videos_per_prompt: int = 1
|
||||
seed: int = 1024
|
||||
|
||||
# Original dimensions (before VAE scaling)
|
||||
num_frames: int = 125
|
||||
height: int = 720
|
||||
width: int = 1280
|
||||
height_sr: int = 1072
|
||||
width_sr: int = 1920
|
||||
fps: int = 24
|
||||
|
||||
# Denoising parameters
|
||||
num_inference_steps: int = 50
|
||||
num_inference_steps_sr: int = 50
|
||||
guidance_scale: float = 1.0
|
||||
guidance_scale_2: float | None = None
|
||||
guidance_rescale: float = 0.0
|
||||
boundary_ratio: float | None = None
|
||||
sigmas: list[float] | None = None
|
||||
|
||||
# TeaCache parameters
|
||||
enable_teacache: bool = False
|
||||
|
||||
# GEN3C camera control
|
||||
trajectory_type: str | None = None
|
||||
movement_distance: float | None = None
|
||||
camera_rotation: str | None = None
|
||||
|
||||
# LTX-2 multi-modal CFG and STG.
|
||||
# Class-level defaults match the *distilled* LTX-2 schedule
|
||||
# (mirrors ``FastVideo-internal/.../LTX2DistilledSamplingParam``):
|
||||
# the distilled model expects neutral guidance scales — modality 1,
|
||||
# rescale 0, STG 0 — and explicit-CFG callers (full LTX-2) opt back
|
||||
# in by selecting the ``LTX2_BASE`` preset, which overrides these
|
||||
# to mod=3.0 / rescale=0.7 / stg=1.0 in its ``defaults`` dict.
|
||||
# cfg_scale defaults stay at 1.0 (CFG off) so
|
||||
# ``ForwardBatch.__post_init__`` doesn't force CFG on non-LTX-2
|
||||
# models that never override these fields.
|
||||
ltx2_cfg_scale_video: float = 1.0
|
||||
ltx2_cfg_scale_audio: float = 1.0
|
||||
ltx2_modality_scale_video: float = 1.0
|
||||
ltx2_modality_scale_audio: float = 1.0
|
||||
ltx2_rescale_scale: float = 0.0
|
||||
ltx2_stg_scale_video: float = 0.0
|
||||
ltx2_stg_scale_audio: float = 0.0
|
||||
ltx2_stg_blocks_video: list[int] = field(default_factory=lambda: [29])
|
||||
ltx2_stg_blocks_audio: list[int] = field(default_factory=lambda: [29])
|
||||
|
||||
# LTX-2 image / video / continuation conditioning. These flow from
|
||||
# generate_video(...) kwargs through ``sampling_param.update(kwargs)``
|
||||
# onto the ForwardBatch fields of the same name. ``ltx2_image_crf``
|
||||
# gates the conditioning-image H.264 re-encode; the streaming
|
||||
# session controller passes ``ltx2_image_crf=0.0`` because it
|
||||
# conditions on already-decoded VAE-quality frames.
|
||||
ltx2_images: list[tuple[str, int, float]] | None = None
|
||||
ltx2_image_crf: float = 33.0
|
||||
ltx2_conditioning_latent_stage1: Any | None = None
|
||||
ltx2_conditioning_latent_stage2: Any | None = None
|
||||
ltx2_video_conditions: list[tuple[list[str], int, float]] | None = None
|
||||
|
||||
# Stable Audio (T2A): clip start/end in seconds. Honored by
|
||||
# `StableAudioConditioningStage` + `StableAudioDecodingStage`. Other
|
||||
# families ignore them.
|
||||
audio_start_in_s: float | None = None
|
||||
audio_end_in_s: float | None = None
|
||||
|
||||
# Stable Audio audio-to-audio (variation):
|
||||
# `init_audio` -- a path or `[B, C, samples]` waveform at the model
|
||||
# sample rate; the pipeline encodes it via the VAE
|
||||
# and uses it as the starting latent.
|
||||
# `init_audio_strength` -- 0..1, higher = closer to the reference
|
||||
# (matches the convention of Stability's
|
||||
# commercial Stable Audio 2.0 UI). 1.0 ~=
|
||||
# VAE round-trip, 0.0 ~= plain T2A.
|
||||
# `init_noise_level` -- legacy raw `sigma_max` override (0.3..500,
|
||||
# higher = more freedom). Kept for callers
|
||||
# that already use it; prefer `init_audio_strength`.
|
||||
init_audio: Any = None
|
||||
init_audio_strength: float | None = None
|
||||
init_noise_level: float | None = None
|
||||
|
||||
# Stable Audio inpainting (RePaint-style): `inpaint_audio` is the
|
||||
# reference clip, `inpaint_mask` is a [samples] tensor in {0, 1} where
|
||||
# 1 means *keep the reference* and 0 means *regenerate*.
|
||||
inpaint_audio: Any = None
|
||||
inpaint_mask: Any = None
|
||||
|
||||
# Continuation state carried across streaming/multi-segment calls.
|
||||
continuation_state: ContinuationState | None = None
|
||||
# When True, the pipeline returns a ContinuationState on the result so
|
||||
# the caller can resume from the generated segment.
|
||||
return_continuation_state: bool = False
|
||||
|
||||
# Misc
|
||||
save_video: bool = True
|
||||
return_frames: bool = True
|
||||
return_trajectory_latents: bool = False # returns all latents for each timestep
|
||||
return_trajectory_decoded: bool = False # returns decoded latents for each timestep
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
self.data_type = "video" if self.num_frames > 1 else "image"
|
||||
|
||||
def check_sampling_param(self):
|
||||
if self.prompt_path and not self.prompt_path.endswith(".txt"):
|
||||
raise ValueError("prompt_path must be a txt file")
|
||||
|
||||
def update(self, source_dict: dict[str, Any]) -> None:
|
||||
valid_fields = {f.name for f in fields(self)}
|
||||
unknown = [key for key in source_dict if key not in valid_fields]
|
||||
if unknown:
|
||||
raise ValueError(f"{type(self).__name__}.update() received unknown field(s): "
|
||||
f"{sorted(unknown)}. All kwargs must correspond to declared "
|
||||
f"SamplingParam fields. If a kwarg is meant to flow into "
|
||||
f"ForwardBatch.extra (e.g. LTX2 audio conditioning), route it "
|
||||
f"via VideoGenerator._BATCH_EXTRA_PASSTHROUGH_KEYS instead.")
|
||||
for key, value in source_dict.items():
|
||||
setattr(self, key, value)
|
||||
|
||||
self.__post_init__()
|
||||
|
||||
@classmethod
|
||||
def from_pretrained(cls, model_path: str) -> SamplingParam:
|
||||
sampling_param = cls._from_preset(model_path)
|
||||
if sampling_param is not None:
|
||||
return sampling_param
|
||||
|
||||
logger.warning(
|
||||
"Couldn't find a preset for %s."
|
||||
" Using the default sampling param.",
|
||||
model_path,
|
||||
)
|
||||
return cls()
|
||||
|
||||
@classmethod
|
||||
def _from_preset(
|
||||
cls,
|
||||
model_path: str,
|
||||
) -> SamplingParam | None:
|
||||
"""Build a SamplingParam from preset defaults.
|
||||
|
||||
Returns ``None`` when no preset is configured for
|
||||
*model_path*, letting the caller fall back to the legacy
|
||||
subclass lookup.
|
||||
"""
|
||||
from v2.registry import get_preset_selection
|
||||
|
||||
try:
|
||||
preset_name, model_family = get_preset_selection(model_path)
|
||||
except (ValueError, RuntimeError):
|
||||
return None
|
||||
if preset_name is None or model_family is None:
|
||||
return None
|
||||
|
||||
from v2._vendor.api.presets import get_preset
|
||||
|
||||
preset = get_preset(preset_name, model_family)
|
||||
sp = cls()
|
||||
valid_fields = {f.name for f in fields(cls)}
|
||||
for key, value in preset.defaults.items():
|
||||
if key in valid_fields:
|
||||
setattr(sp, key, copy.deepcopy(value))
|
||||
sp.__post_init__()
|
||||
return sp
|
||||
|
||||
@staticmethod
|
||||
def add_cli_args(parser: Any) -> Any:
|
||||
"""Add CLI arguments for SamplingParam fields"""
|
||||
parser.add_argument(
|
||||
"--prompt",
|
||||
type=str,
|
||||
default=SamplingParam.prompt,
|
||||
help="Text prompt for video generation",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--negative-prompt",
|
||||
type=str,
|
||||
default=SamplingParam.negative_prompt,
|
||||
help="Negative text prompt for video generation",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--prompt-path",
|
||||
type=str,
|
||||
default=SamplingParam.prompt_path,
|
||||
help="Path to a text file containing the prompt",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output-path",
|
||||
type=str,
|
||||
default=SamplingParam.output_path,
|
||||
help="Path to save the generated video",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output-video-name",
|
||||
type=str,
|
||||
default=SamplingParam.output_video_name,
|
||||
help="Name of the output video",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--num-videos-per-prompt",
|
||||
type=int,
|
||||
default=SamplingParam.num_videos_per_prompt,
|
||||
help="Number of videos to generate per prompt",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--seed",
|
||||
type=int,
|
||||
default=SamplingParam.seed,
|
||||
help="Random seed for generation",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--num-frames",
|
||||
type=int,
|
||||
default=SamplingParam.num_frames,
|
||||
help="Number of frames to generate",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--height",
|
||||
type=int,
|
||||
default=SamplingParam.height,
|
||||
help="Height of generated video",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--width",
|
||||
type=int,
|
||||
default=SamplingParam.width,
|
||||
help="Width of generated video",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--fps",
|
||||
type=int,
|
||||
default=SamplingParam.fps,
|
||||
help="Frames per second for saved video",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--num-inference-steps",
|
||||
type=int,
|
||||
default=SamplingParam.num_inference_steps,
|
||||
help="Number of denoising steps",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--guidance-scale",
|
||||
type=float,
|
||||
default=SamplingParam.guidance_scale,
|
||||
help="Classifier-free guidance scale",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--guidance-rescale",
|
||||
type=float,
|
||||
default=SamplingParam.guidance_rescale,
|
||||
help="Guidance rescale factor",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--boundary-ratio",
|
||||
type=float,
|
||||
default=SamplingParam.boundary_ratio,
|
||||
help="Boundary timestep ratio",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--save-video",
|
||||
action="store_true",
|
||||
default=SamplingParam.save_video,
|
||||
help="Whether to save the video to disk",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--no-save-video",
|
||||
action="store_false",
|
||||
dest="save_video",
|
||||
help="Don't save the video to disk",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--return-frames",
|
||||
action="store_true",
|
||||
default=False,
|
||||
help="Whether to return the raw frames",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--image-path",
|
||||
type=str,
|
||||
default=SamplingParam.image_path,
|
||||
help="Path to input image for image-to-video generation",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--video-path",
|
||||
type=str,
|
||||
default=SamplingParam.video_path,
|
||||
help="Path to input video for video-to-video generation",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--refine-from",
|
||||
type=str,
|
||||
default=SamplingParam.refine_from,
|
||||
help="Path to stage1 video for refinement (LongCat 480p->720p)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--t-thresh",
|
||||
type=float,
|
||||
default=SamplingParam.t_thresh,
|
||||
help="Threshold for timestep scheduling in refinement (default: 0.5)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--spatial-refine-only",
|
||||
action=StoreBoolean,
|
||||
default=SamplingParam.spatial_refine_only,
|
||||
help="Only perform spatial super-resolution (no temporal doubling)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--num-cond-frames",
|
||||
type=int,
|
||||
default=SamplingParam.num_cond_frames,
|
||||
help="Number of conditioning frames for refinement",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--moba-config-path",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Path to a JSON file containing V-MoBA specific configurations.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--return-trajectory-latents",
|
||||
action="store_true",
|
||||
default=SamplingParam.return_trajectory_latents,
|
||||
help="Whether to return the trajectory",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--return-trajectory-decoded",
|
||||
action="store_true",
|
||||
default=SamplingParam.return_trajectory_decoded,
|
||||
help="Whether to return the decoded trajectory",
|
||||
)
|
||||
return parser
|
||||
|
||||
|
||||
@dataclass
|
||||
class CacheParams:
|
||||
cache_type: str = "none"
|
||||
@@ -1,307 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Literal
|
||||
|
||||
|
||||
@dataclass
|
||||
class ServerConfig:
|
||||
host: str = "0.0.0.0"
|
||||
port: int = 8000
|
||||
output_dir: str = "outputs/"
|
||||
|
||||
|
||||
@dataclass
|
||||
class ParallelismConfig:
|
||||
tp_size: int = -1
|
||||
sp_size: int = -1
|
||||
hsdp_replicate_dim: int = 1
|
||||
hsdp_shard_dim: int = -1
|
||||
dist_timeout: int | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class OffloadConfig:
|
||||
dit: bool = True
|
||||
dit_layerwise: bool = True
|
||||
text_encoder: bool = True
|
||||
image_encoder: bool = True
|
||||
vae: bool = True
|
||||
pin_cpu_memory: bool = True
|
||||
|
||||
|
||||
@dataclass
|
||||
class CompileConfig:
|
||||
"""Typed ``torch.compile`` configuration.
|
||||
|
||||
``backend``/``fullgraph``/``mode``/``dynamic`` are the four most
|
||||
common ``torch.compile`` knobs. ``extras`` holds any remaining
|
||||
``torch.compile`` kwargs (e.g. ``options``, ``disable``).
|
||||
|
||||
The ``enabled`` switch covers the DiT transformer path (including
|
||||
``transformer_2`` and the LTX-2 stage-2 ``transformer_refine``).
|
||||
Per-component flags below are independent overlays — set to ``True``
|
||||
to compile that component, ``None`` to leave it eager. Each
|
||||
``*_kwargs`` dict overrides the master ``backend``/``fullgraph``/
|
||||
``mode``/``dynamic``/``extras`` for that component when non-empty;
|
||||
leaving it empty inherits the master kwargs.
|
||||
"""
|
||||
|
||||
enabled: bool = False
|
||||
backend: str | None = None
|
||||
fullgraph: bool | None = None
|
||||
mode: str | None = None
|
||||
dynamic: bool | None = None
|
||||
extras: dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
text_encoder_enabled: bool | None = None
|
||||
vae_enabled: bool | None = None
|
||||
audio_vae_enabled: bool | None = None
|
||||
|
||||
dit_kwargs: dict[str, Any] = field(default_factory=dict)
|
||||
text_encoder_kwargs: dict[str, Any] = field(default_factory=dict)
|
||||
vae_kwargs: dict[str, Any] = field(default_factory=dict)
|
||||
audio_vae_kwargs: dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
|
||||
@dataclass
|
||||
class QuantizationConfig:
|
||||
text_encoder_quant: str | None = None
|
||||
transformer_quant: str | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class EngineConfig:
|
||||
num_gpus: int = 1
|
||||
execution_backend: Literal["mp", "ray"] = "mp"
|
||||
parallelism: ParallelismConfig = field(default_factory=ParallelismConfig)
|
||||
offload: OffloadConfig = field(default_factory=OffloadConfig)
|
||||
compile: CompileConfig = field(default_factory=CompileConfig)
|
||||
enable_stage_verification: bool = True
|
||||
use_fsdp_inference: bool = False
|
||||
disable_autocast: bool = False
|
||||
quantization: QuantizationConfig | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class ComponentConfig:
|
||||
config_root: str | None = None
|
||||
pipeline_config_path: str | None = None
|
||||
text_encoder_weights: str | None = None
|
||||
transformer_weights: str | None = None
|
||||
transformer_2_weights: str | None = None
|
||||
vae_weights: str | None = None
|
||||
upsampler_weights: str | None = None
|
||||
lora_path: str | None = None
|
||||
override_pipeline_cls_name: str | None = None
|
||||
override_transformer_cls_name: str | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class PipelineSelection:
|
||||
workload_type: Literal["t2v", "i2v", "t2i", "i2i"] | None = None
|
||||
preset: str | None = None
|
||||
preset_version: int | None = None
|
||||
components: ComponentConfig = field(default_factory=ComponentConfig)
|
||||
vae_tiling: bool | None = None
|
||||
"""Tile-based VAE decode. ``None`` keeps the model's default."""
|
||||
preset_overrides: dict[str, Any] = field(default_factory=dict)
|
||||
experimental: dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
|
||||
@dataclass
|
||||
class GeneratorConfig:
|
||||
model_path: str
|
||||
revision: str | None = None
|
||||
trust_remote_code: bool = False
|
||||
engine: EngineConfig = field(default_factory=EngineConfig)
|
||||
pipeline: PipelineSelection = field(default_factory=PipelineSelection)
|
||||
|
||||
|
||||
@dataclass
|
||||
class InputConfig:
|
||||
prompt_path: str | None = None
|
||||
image_path: str | list[str] | None = None
|
||||
video_path: str | list[str] | None = None
|
||||
pil_image: Any | None = None
|
||||
pose: str | None = None
|
||||
mouse_cond: Any | None = None
|
||||
keyboard_cond: Any | None = None
|
||||
grid_sizes: Any | None = None
|
||||
c2ws_plucker_emb: Any | None = None
|
||||
refine_from: str | None = None
|
||||
stage1_video: Any | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class SamplingConfig:
|
||||
num_videos_per_prompt: int = 1
|
||||
seed: int = 1024
|
||||
num_frames: int = 125
|
||||
height: int = 720
|
||||
width: int = 1280
|
||||
height_sr: int = 1072
|
||||
width_sr: int = 1920
|
||||
fps: int = 24
|
||||
num_inference_steps: int = 50
|
||||
num_inference_steps_sr: int = 50
|
||||
guidance_scale: float = 1.0
|
||||
guidance_scale_2: float | None = None
|
||||
guidance_rescale: float = 0.0
|
||||
true_cfg_scale: float | None = None
|
||||
boundary_ratio: float | None = None
|
||||
sigmas: list[float] | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class RequestRuntimeConfig:
|
||||
enable_teacache: bool = False
|
||||
return_trajectory_latents: bool = False
|
||||
return_trajectory_decoded: bool = False
|
||||
|
||||
|
||||
@dataclass
|
||||
class OutputConfig:
|
||||
output_path: str = "outputs/"
|
||||
output_video_name: str | None = None
|
||||
save_video: bool = True
|
||||
return_frames: bool = True
|
||||
return_state: bool = False
|
||||
|
||||
|
||||
@dataclass
|
||||
class ContinuationState:
|
||||
kind: str
|
||||
payload: dict[str, Any]
|
||||
|
||||
|
||||
@dataclass
|
||||
class PlannedStage:
|
||||
name: str
|
||||
kind: str
|
||||
source: str | None = None
|
||||
overrides: dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
|
||||
@dataclass
|
||||
class GenerationPlan:
|
||||
stages: list[PlannedStage]
|
||||
final_stage: str | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class GenerationRequest:
|
||||
prompt: str | list[str] | None = None
|
||||
negative_prompt: str | None = None
|
||||
inputs: InputConfig = field(default_factory=InputConfig)
|
||||
sampling: SamplingConfig = field(default_factory=SamplingConfig)
|
||||
runtime: RequestRuntimeConfig = field(default_factory=RequestRuntimeConfig)
|
||||
output: OutputConfig = field(default_factory=OutputConfig)
|
||||
stage_overrides: dict[str, Any] = field(default_factory=dict)
|
||||
state: ContinuationState | None = None
|
||||
plan: GenerationPlan | None = None
|
||||
extensions: dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
|
||||
@dataclass
|
||||
class RunConfig:
|
||||
generator: GeneratorConfig
|
||||
request: GenerationRequest
|
||||
|
||||
|
||||
@dataclass
|
||||
class WarmupConfig:
|
||||
enabled: bool = True
|
||||
prompt: str = ("A cinematic drone shot over coastal cliffs at sunrise, "
|
||||
"golden light, gentle ocean waves, ultra detailed")
|
||||
timeout_seconds: int = 2400
|
||||
|
||||
|
||||
@dataclass
|
||||
class GpuPoolConfig:
|
||||
num_workers: int | None = None
|
||||
enable_audio_reencode: bool = True
|
||||
conditioning_num_frames: int = 9
|
||||
conditioning_end_offset: int = 0
|
||||
|
||||
|
||||
@dataclass
|
||||
class PromptEnhancerConfig:
|
||||
enabled: bool = False
|
||||
provider: Literal["cerebras", "groq"] = "cerebras"
|
||||
model: str = "gpt-oss-120b"
|
||||
timeout_ms: int = 20000
|
||||
system_prompt_dir: str | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class PromptSafetyConfig:
|
||||
enabled: bool = False
|
||||
classifier_path: str | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class StreamingConfig:
|
||||
session_timeout_seconds: int = 300
|
||||
generation_segment_cap: int = 6
|
||||
stream_mode: Literal["av_fmp4", "legacy_jpeg"] = "av_fmp4"
|
||||
warmup: WarmupConfig = field(default_factory=WarmupConfig)
|
||||
pool: GpuPoolConfig = field(default_factory=GpuPoolConfig)
|
||||
prompt: PromptEnhancerConfig = field(default_factory=PromptEnhancerConfig)
|
||||
safety: PromptSafetyConfig = field(default_factory=PromptSafetyConfig)
|
||||
|
||||
|
||||
@dataclass
|
||||
class ServeConfig:
|
||||
"""Typed serve config loaded from ``fastvideo serve --config``.
|
||||
|
||||
``default_request`` is a full :class:`GenerationRequest` — the same type
|
||||
clients POST to ``/v1/videos``. At request time the server merges it into
|
||||
the incoming body as the operator-pinned baseline.
|
||||
|
||||
Important nuance: only fields the operator **explicitly wrote** in the
|
||||
serve YAML/JSON count as defaults. Although the in-memory object is
|
||||
fully populated (schema defaults fill every unset field), the merge
|
||||
walks ``_fastvideo_explicit_paths`` — populated during parse — so
|
||||
unset fields are *not* forced onto requests. Per-request precedence:
|
||||
|
||||
body (client-explicit) > default_request (operator-explicit)
|
||||
> hardcoded fallback (e.g. ``fps=24``)
|
||||
|
||||
See :func:`v2._vendor.api.compat.explicit_request_updates` for the
|
||||
projection and ``entrypoints/openai/video_api.py::_build_generation_kwargs``
|
||||
for the merge.
|
||||
"""
|
||||
generator: GeneratorConfig
|
||||
server: ServerConfig = field(default_factory=ServerConfig)
|
||||
default_request: GenerationRequest = field(default_factory=GenerationRequest)
|
||||
streaming: StreamingConfig | None = None
|
||||
|
||||
|
||||
__all__ = [
|
||||
"CompileConfig",
|
||||
"ComponentConfig",
|
||||
"ContinuationState",
|
||||
"EngineConfig",
|
||||
"GenerationPlan",
|
||||
"GenerationRequest",
|
||||
"GeneratorConfig",
|
||||
"GpuPoolConfig",
|
||||
"InputConfig",
|
||||
"OffloadConfig",
|
||||
"OutputConfig",
|
||||
"ParallelismConfig",
|
||||
"PipelineSelection",
|
||||
"PlannedStage",
|
||||
"PromptEnhancerConfig",
|
||||
"PromptSafetyConfig",
|
||||
"QuantizationConfig",
|
||||
"RequestRuntimeConfig",
|
||||
"RunConfig",
|
||||
"SamplingConfig",
|
||||
"ServeConfig",
|
||||
"ServerConfig",
|
||||
"StreamingConfig",
|
||||
"WarmupConfig",
|
||||
]
|
||||
@@ -1,58 +0,0 @@
|
||||
# `fastvideo/attention/` — Attention Backends
|
||||
|
||||
**Generated:** 2026-05-02
|
||||
|
||||
Backend registry + selector wrapping FlashAttn / SageAttn / SageAttn3 / SDPA / VSA / VMoBA / SLA / BSA.
|
||||
|
||||
## Layout
|
||||
|
||||
```
|
||||
attention/
|
||||
├── __init__.py # Exports DistributedAttention, LocalAttention, get_attn_backend
|
||||
├── layer.py # DistributedAttention, DistributedAttention_VSA, LocalAttention
|
||||
├── selector.py # get_attn_backend (cached) + env-var override
|
||||
├── backends/
|
||||
│ ├── abstract.py # AttentionBackend / AttentionMetadata / AttentionMetadataBuilder
|
||||
│ ├── flash_attn.py # FA2/FA3
|
||||
│ ├── sage_attn.py # SageAttention v1
|
||||
│ ├── sage_attn3.py # SageAttention v3
|
||||
│ ├── sdpa.py # torch SDPA fallback
|
||||
│ ├── video_sparse_attn.py # VSA (paper: Video Sparse Attention)
|
||||
│ ├── vmoba.py # Video-MoBA
|
||||
│ ├── sla.py # Sliding-window (STA)
|
||||
│ └── bsa_attn.py # Block-sparse
|
||||
└── utils/
|
||||
├── flash_attn_cute.py
|
||||
└── flash_attn_no_pad.py
|
||||
```
|
||||
|
||||
## Selection Order
|
||||
|
||||
`get_attn_backend()` resolves via:
|
||||
|
||||
1. Env-var override `FASTVIDEO_ATTENTION_BACKEND` (see `STR_BACKEND_ENV_VAR` in `fastvideo/utils.py`).
|
||||
2. Per-platform default from `fastvideo/platforms/`.
|
||||
3. Heuristic fallback to SDPA.
|
||||
|
||||
The result is `@lru_cache`d. Tests that need a specific backend must use the
|
||||
`global_force_attn_backend(...)` context manager from `selector.py`, never set
|
||||
the env var mid-process.
|
||||
|
||||
## Adding a Backend
|
||||
|
||||
1. Subclass `AttentionBackend` in `backends/<name>.py`.
|
||||
2. Implement `AttentionMetadata` + `AttentionMetadataBuilder` for the new path.
|
||||
3. Register the enum value in `fastvideo/platforms/interface.py` (`AttentionBackendEnum`).
|
||||
4. Wire string → class resolution in `selector.py`.
|
||||
5. Verify the new backend works with `DistributedAttention` (sequence parallel)
|
||||
and `LocalAttention` (single-rank). If it cannot support SP, document the
|
||||
gap in the backend file's module docstring.
|
||||
|
||||
## Anti-Patterns
|
||||
|
||||
- Calling `torch.nn.functional.scaled_dot_product_attention` directly inside a
|
||||
model's forward — go through `DistributedAttention` / `LocalAttention`.
|
||||
- Reading `os.environ[STR_BACKEND_ENV_VAR]` from arbitrary call sites. Use
|
||||
`get_env_variable_attn_backend()`.
|
||||
- Caching backend instances per-module. The selector cache is process-wide; do
|
||||
not duplicate it.
|
||||
@@ -1,16 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
from v2._vendor.attention.backends.abstract import (AttentionBackend, AttentionMetadata, AttentionMetadataBuilder)
|
||||
from v2._vendor.attention.layer import (DistributedAttention, DistributedAttention_VSA, LocalAttention)
|
||||
from v2._vendor.attention.selector import get_attn_backend
|
||||
|
||||
__all__ = [
|
||||
"DistributedAttention",
|
||||
"LocalAttention",
|
||||
"DistributedAttention_VSA",
|
||||
"AttentionBackend",
|
||||
"AttentionMetadata",
|
||||
"AttentionMetadataBuilder",
|
||||
# "AttentionState",
|
||||
"get_attn_backend",
|
||||
]
|
||||
@@ -1,177 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/attention/backends/abstract.py
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from dataclasses import dataclass, field, fields
|
||||
from typing import TYPE_CHECKING, Any, Generic, Protocol, TypeVar
|
||||
|
||||
if TYPE_CHECKING:
|
||||
pass
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
class AttentionBackend(ABC):
|
||||
"""Abstract class for attention backends."""
|
||||
# For some attention backends, we allocate an output tensor before
|
||||
# calling the custom op. When piecewise cudagraph is enabled, this
|
||||
# makes sure the output tensor is allocated inside the cudagraph.
|
||||
accept_output_buffer: bool = False
|
||||
|
||||
@staticmethod
|
||||
@abstractmethod
|
||||
def get_name() -> str:
|
||||
raise NotImplementedError
|
||||
|
||||
@staticmethod
|
||||
@abstractmethod
|
||||
def get_impl_cls() -> type["AttentionImpl"]:
|
||||
raise NotImplementedError
|
||||
|
||||
@staticmethod
|
||||
@abstractmethod
|
||||
def get_metadata_cls() -> type["AttentionMetadata"]:
|
||||
raise NotImplementedError
|
||||
|
||||
# @staticmethod
|
||||
# @abstractmethod
|
||||
# def get_state_cls() -> Type["AttentionState"]:
|
||||
# raise NotImplementedError
|
||||
|
||||
# @classmethod
|
||||
# def make_metadata(cls, *args, **kwargs) -> "AttentionMetadata":
|
||||
# return cls.get_metadata_cls()(*args, **kwargs)
|
||||
|
||||
@staticmethod
|
||||
@abstractmethod
|
||||
def get_builder_cls() -> type["AttentionMetadataBuilder"]:
|
||||
raise NotImplementedError
|
||||
|
||||
|
||||
@dataclass
|
||||
class AttentionMetadata:
|
||||
"""Attention metadata for prefill and decode batched together."""
|
||||
# Current step of diffusion process
|
||||
current_timestep: int
|
||||
VSA_sparsity: float = field(default=0.0, kw_only=True)
|
||||
|
||||
def __getattr__(self, name: str) -> Any:
|
||||
raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'")
|
||||
|
||||
def asdict_zerocopy(self, skip_fields: set[str] | None = None) -> dict[str, Any]:
|
||||
"""Similar to dataclasses.asdict, but avoids deepcopying."""
|
||||
if skip_fields is None:
|
||||
skip_fields = set()
|
||||
# Note that if we add dataclasses as fields, they will need
|
||||
# similar handling.
|
||||
return {field.name: getattr(self, field.name) for field in fields(self) if field.name not in skip_fields}
|
||||
|
||||
|
||||
T = TypeVar("T", bound=AttentionMetadata)
|
||||
|
||||
|
||||
class AttentionMetadataBuilder(ABC, Generic[T]):
|
||||
"""Abstract class for attention metadata builders."""
|
||||
|
||||
@abstractmethod
|
||||
def __init__(self) -> None:
|
||||
"""Create the builder, remember some configuration and parameters."""
|
||||
raise NotImplementedError
|
||||
|
||||
@abstractmethod
|
||||
def prepare(self) -> None:
|
||||
"""Prepare for one batch."""
|
||||
raise NotImplementedError
|
||||
|
||||
@abstractmethod
|
||||
def build(
|
||||
self,
|
||||
**kwargs: Any,
|
||||
) -> AttentionMetadata:
|
||||
"""Build attention metadata with on-device tensors."""
|
||||
raise NotImplementedError
|
||||
|
||||
|
||||
class AttentionLayer(Protocol):
|
||||
|
||||
_k_scale: torch.Tensor
|
||||
_v_scale: torch.Tensor
|
||||
_k_scale_float: float
|
||||
_v_scale_float: float
|
||||
|
||||
def forward(
|
||||
self,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
kv_cache: torch.Tensor,
|
||||
attn_metadata: AttentionMetadata,
|
||||
) -> torch.Tensor:
|
||||
...
|
||||
|
||||
|
||||
class AttentionImpl(ABC, Generic[T]):
|
||||
|
||||
@abstractmethod
|
||||
def __init__(
|
||||
self,
|
||||
num_heads: int,
|
||||
head_size: int,
|
||||
softmax_scale: float,
|
||||
causal: bool = False,
|
||||
num_kv_heads: int | None = None,
|
||||
prefix: str = "",
|
||||
**extra_impl_args,
|
||||
) -> None:
|
||||
raise NotImplementedError
|
||||
|
||||
def preprocess_qkv(self, qkv: torch.Tensor, attn_metadata: T) -> torch.Tensor:
|
||||
"""Preprocess QKV tensor before performing attention operation.
|
||||
|
||||
Default implementation returns the tensor unchanged.
|
||||
Subclasses can override this to implement custom preprocessing
|
||||
like reshaping, tiling, scaling, or other transformations.
|
||||
|
||||
Called AFTER all_to_all for distributed attention
|
||||
|
||||
Args:
|
||||
qkv: The query-key-value tensor
|
||||
attn_metadata: Metadata for the attention operation
|
||||
|
||||
Returns:
|
||||
Processed QKV tensor
|
||||
"""
|
||||
return qkv
|
||||
|
||||
def postprocess_output(
|
||||
self,
|
||||
output: torch.Tensor,
|
||||
attn_metadata: T,
|
||||
) -> torch.Tensor:
|
||||
"""Postprocess the output tensor after the attention operation.
|
||||
|
||||
Default implementation returns the tensor unchanged.
|
||||
Subclasses can override this to implement custom postprocessing
|
||||
like untiling, scaling, or other transformations.
|
||||
|
||||
Called BEFORE all_to_all for distributed attention
|
||||
|
||||
Args:
|
||||
output: The output tensor from the attention operation
|
||||
attn_metadata: Metadata for the attention operation
|
||||
|
||||
Returns:
|
||||
Postprocessed output tensor
|
||||
"""
|
||||
|
||||
return output
|
||||
|
||||
@abstractmethod
|
||||
def forward(
|
||||
self,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
attn_metadata: T,
|
||||
) -> torch.Tensor:
|
||||
raise NotImplementedError
|
||||
@@ -1,125 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
import importlib
|
||||
import sys
|
||||
from collections.abc import Callable
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
|
||||
from v2._vendor.attention.backends.abstract import (
|
||||
AttentionBackend,
|
||||
AttentionImpl,
|
||||
AttentionMetadata,
|
||||
AttentionMetadataBuilder,
|
||||
)
|
||||
from v2._vendor.logger import init_logger
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
_project_root = Path(__file__).resolve().parent.parent.parent.parent
|
||||
_kernel_root = _project_root / "fastvideo-kernel"
|
||||
_kernel_python_root = _kernel_root / "python"
|
||||
_attn_qat_infer: Callable[..., torch.Tensor] | None = None
|
||||
_attn_qat_infer_import_attempted = False
|
||||
|
||||
|
||||
def _ensure_kernel_paths() -> None:
|
||||
for path in (_project_root, _kernel_root, _kernel_python_root):
|
||||
path_str = str(path)
|
||||
if path_str not in sys.path:
|
||||
sys.path.insert(0, path_str)
|
||||
|
||||
|
||||
def _get_attn_qat_infer() -> Callable[..., torch.Tensor] | None:
|
||||
global _attn_qat_infer
|
||||
global _attn_qat_infer_import_attempted
|
||||
|
||||
if _attn_qat_infer_import_attempted:
|
||||
return _attn_qat_infer
|
||||
|
||||
_attn_qat_infer_import_attempted = True
|
||||
_ensure_kernel_paths()
|
||||
|
||||
try:
|
||||
# Prefer the in-repo kernel implementation during local development.
|
||||
_attn_qat_infer = importlib.import_module("attn_qat_infer").sageattn_blackwell
|
||||
except ImportError:
|
||||
_attn_qat_infer = None
|
||||
|
||||
return _attn_qat_infer
|
||||
|
||||
|
||||
def is_attn_qat_infer_available() -> bool:
|
||||
return _get_attn_qat_infer() is not None
|
||||
|
||||
|
||||
class AttnQatInferBackend(AttentionBackend):
|
||||
|
||||
accept_output_buffer: bool = True
|
||||
|
||||
@staticmethod
|
||||
def get_supported_head_sizes() -> list[int]:
|
||||
return [64, 128]
|
||||
|
||||
@staticmethod
|
||||
def get_name() -> str:
|
||||
return "ATTN_QAT_INFER"
|
||||
|
||||
@staticmethod
|
||||
def get_impl_cls() -> type["AttnQatInferImpl"]:
|
||||
return AttnQatInferImpl
|
||||
|
||||
@staticmethod
|
||||
def get_metadata_cls() -> type["AttentionMetadata"]:
|
||||
raise NotImplementedError
|
||||
|
||||
@staticmethod
|
||||
def get_builder_cls() -> type["AttentionMetadataBuilder[AttentionMetadata]"]:
|
||||
raise NotImplementedError
|
||||
|
||||
|
||||
class AttnQatInferImpl(AttentionImpl[AttentionMetadata]):
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
num_heads: int,
|
||||
head_size: int,
|
||||
causal: bool,
|
||||
softmax_scale: float,
|
||||
num_kv_heads: int | None = None,
|
||||
prefix: str = "",
|
||||
**extra_impl_args,
|
||||
) -> None:
|
||||
self.causal = causal
|
||||
self.softmax_scale = softmax_scale
|
||||
dropout_p = extra_impl_args.get("dropout_p", 0.0)
|
||||
if dropout_p > 0:
|
||||
raise NotImplementedError(f"attn_qat_infer does not support dropout (got dropout_p={dropout_p}). "
|
||||
"The QAT inference kernel applies no stochastic dropout.")
|
||||
|
||||
def forward(
|
||||
self,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
attn_metadata: AttentionMetadata,
|
||||
) -> torch.Tensor:
|
||||
attn_qat_infer = _get_attn_qat_infer()
|
||||
if attn_qat_infer is None:
|
||||
raise ImportError("attn_qat_infer is not available. Please ensure the "
|
||||
"attn_qat_infer kernel package is installed.")
|
||||
|
||||
query = query.transpose(1, 2).contiguous()
|
||||
key = key.transpose(1, 2).contiguous()
|
||||
value = value.transpose(1, 2).contiguous()
|
||||
|
||||
output = attn_qat_infer(
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
attn_mask=None,
|
||||
is_causal=self.causal,
|
||||
sm_scale=self.softmax_scale,
|
||||
)
|
||||
return output.transpose(1, 2).contiguous()
|
||||
@@ -1,154 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
import importlib
|
||||
import sys
|
||||
from collections.abc import Callable
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
|
||||
from v2._vendor.attention.backends.abstract import (
|
||||
AttentionBackend,
|
||||
AttentionImpl,
|
||||
AttentionMetadata,
|
||||
AttentionMetadataBuilder,
|
||||
)
|
||||
from v2._vendor.logger import init_logger
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
_project_root = Path(__file__).resolve().parent.parent.parent.parent
|
||||
_kernel_root = _project_root / "fastvideo-kernel"
|
||||
_kernel_python_root = _kernel_root / "python"
|
||||
_attn_qat_train_attention: Callable[..., torch.Tensor] | None = None
|
||||
_attn_qat_train_import_attempted = False
|
||||
|
||||
|
||||
def _ensure_kernel_paths() -> None:
|
||||
for path in (_project_root, _kernel_root, _kernel_python_root):
|
||||
path_str = str(path)
|
||||
if path_str not in sys.path:
|
||||
sys.path.insert(0, path_str)
|
||||
|
||||
|
||||
def _get_attn_qat_train_attention() -> Callable[..., torch.Tensor] | None:
|
||||
global _attn_qat_train_attention
|
||||
global _attn_qat_train_import_attempted
|
||||
|
||||
if _attn_qat_train_import_attempted:
|
||||
return _attn_qat_train_attention
|
||||
|
||||
_attn_qat_train_import_attempted = True
|
||||
_ensure_kernel_paths()
|
||||
|
||||
try:
|
||||
_attn_qat_train_attention = importlib.import_module("fastvideo_kernel.triton_kernels.attn_qat_train").attention
|
||||
except ImportError:
|
||||
_attn_qat_train_attention = None
|
||||
|
||||
return _attn_qat_train_attention
|
||||
|
||||
|
||||
def is_attn_qat_train_available() -> bool:
|
||||
return _get_attn_qat_train_attention() is not None
|
||||
|
||||
|
||||
def attn_qat_train(q_BLHD: torch.Tensor,
|
||||
k_BLHD: torch.Tensor,
|
||||
v_BLHD: torch.Tensor,
|
||||
is_causal: bool = False,
|
||||
sm_scale: float | None = None) -> torch.Tensor:
|
||||
attention = _get_attn_qat_train_attention()
|
||||
if attention is None:
|
||||
raise ImportError("fastvideo_kernel.triton_kernels.attn_qat_train is not available. "
|
||||
"Please ensure the FastVideo kernel package is installed.")
|
||||
|
||||
q_BHLD = q_BLHD.permute(0, 2, 1, 3).contiguous()
|
||||
k_BHLD = k_BLHD.permute(0, 2, 1, 3).contiguous()
|
||||
v_BHLD = v_BLHD.permute(0, 2, 1, 3).contiguous()
|
||||
|
||||
use_qat_qkv_backward = True
|
||||
smooth_k = False
|
||||
warp_specialize = True
|
||||
is_qat = True
|
||||
two_level_quant_p_sage3 = False
|
||||
fake_quant_p_bwd = True
|
||||
use_high_prec_o = True
|
||||
smooth_q = False
|
||||
if sm_scale is None:
|
||||
sm_scale = 1.0 / (q_BHLD.shape[-1]**0.5)
|
||||
use_global_sf_qkv = False
|
||||
use_global_sf_p = False
|
||||
|
||||
o_BHLD = attention(
|
||||
q_BHLD,
|
||||
k_BHLD,
|
||||
v_BHLD,
|
||||
is_causal,
|
||||
sm_scale,
|
||||
use_qat_qkv_backward,
|
||||
smooth_k,
|
||||
warp_specialize,
|
||||
is_qat,
|
||||
two_level_quant_p_sage3,
|
||||
fake_quant_p_bwd,
|
||||
use_high_prec_o,
|
||||
smooth_q,
|
||||
use_global_sf_p,
|
||||
use_global_sf_qkv,
|
||||
)
|
||||
return o_BHLD.permute(0, 2, 1, 3).contiguous()
|
||||
|
||||
|
||||
class AttnQatTrainBackend(AttentionBackend):
|
||||
|
||||
accept_output_buffer: bool = True
|
||||
|
||||
@staticmethod
|
||||
def get_supported_head_sizes() -> list[int]:
|
||||
return [64, 96, 128, 160, 192, 224, 256]
|
||||
|
||||
@staticmethod
|
||||
def get_name() -> str:
|
||||
return "ATTN_QAT_TRAIN"
|
||||
|
||||
@staticmethod
|
||||
def get_impl_cls() -> type["AttnQatTrainImpl"]:
|
||||
return AttnQatTrainImpl
|
||||
|
||||
@staticmethod
|
||||
def get_metadata_cls() -> type["AttentionMetadata"]:
|
||||
raise NotImplementedError
|
||||
|
||||
@staticmethod
|
||||
def get_builder_cls() -> type["AttentionMetadataBuilder[AttentionMetadata]"]:
|
||||
raise NotImplementedError
|
||||
|
||||
|
||||
class AttnQatTrainImpl(AttentionImpl[AttentionMetadata]):
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
num_heads: int,
|
||||
head_size: int,
|
||||
causal: bool,
|
||||
softmax_scale: float,
|
||||
num_kv_heads: int | None = None,
|
||||
prefix: str = "",
|
||||
**extra_impl_args,
|
||||
) -> None:
|
||||
self.causal = causal
|
||||
self.softmax_scale = softmax_scale
|
||||
dropout_p = extra_impl_args.get("dropout_p", 0.0)
|
||||
if dropout_p > 0:
|
||||
raise NotImplementedError(f"attn_qat_train does not support dropout (got dropout_p={dropout_p}). "
|
||||
"The QAT training kernel applies no stochastic dropout.")
|
||||
|
||||
def forward(
|
||||
self,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
attn_metadata: AttentionMetadata,
|
||||
) -> torch.Tensor:
|
||||
return attn_qat_train(query, key, value, is_causal=self.causal, sm_scale=self.softmax_scale)
|
||||
@@ -1,740 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""
|
||||
Bidirectional Sparse Attention (BSA) backend for FastVideo.
|
||||
|
||||
Pure-PyTorch reference implementation from:
|
||||
"Bidirectional Sparse Attention for Faster Video Diffusion Training"
|
||||
(arXiv:2509.01085)
|
||||
|
||||
BSA sparsifies both queries (pruning redundant tokens per block) and
|
||||
key-value pairs (keeping only relevant KV blocks per query block).
|
||||
|
||||
This is a training-free inference backend: it works with any model
|
||||
trained with full attention by applying BSA sparsity at inference time.
|
||||
"""
|
||||
|
||||
import functools
|
||||
import math
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
from v2._vendor.attention.backends.abstract import (
|
||||
AttentionBackend,
|
||||
AttentionImpl,
|
||||
AttentionMetadata,
|
||||
AttentionMetadataBuilder,
|
||||
)
|
||||
from v2._vendor.distributed import get_sp_group
|
||||
from v2._vendor.logger import init_logger
|
||||
|
||||
try:
|
||||
from v2._vendor.attention.utils.flash_attn_no_pad import (
|
||||
flash_attn_varlen_func_impl, )
|
||||
|
||||
FLASH_ATTN_AVAILABLE = True
|
||||
except ImportError:
|
||||
flash_attn_varlen_func_impl = None
|
||||
FLASH_ATTN_AVAILABLE = False
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
BSA_TILE_SIZE = (4, 4, 4)
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Cached index helpers (same pattern as VSA)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@functools.lru_cache(maxsize=10)
|
||||
def get_tile_partition_indices(
|
||||
dit_seq_shape: tuple[int, int, int],
|
||||
tile_size: tuple[int, int, int],
|
||||
device: torch.device,
|
||||
) -> torch.LongTensor:
|
||||
"""Map raster-order tokens to tile-contiguous order."""
|
||||
T, H, W = dit_seq_shape
|
||||
ts, hs, ws = tile_size
|
||||
indices = torch.arange(T * H * W, device=device, dtype=torch.long).reshape(T, H, W)
|
||||
ls = []
|
||||
for t in range(math.ceil(T / ts)):
|
||||
for h in range(math.ceil(H / hs)):
|
||||
for w in range(math.ceil(W / ws)):
|
||||
ls.append(indices[
|
||||
t * ts:min(t * ts + ts, T),
|
||||
h * hs:min(h * hs + hs, H),
|
||||
w * ws:min(w * ws + ws, W),
|
||||
].flatten())
|
||||
return torch.cat(ls, dim=0)
|
||||
|
||||
|
||||
@functools.lru_cache(maxsize=10)
|
||||
def get_reverse_tile_partition_indices(
|
||||
dit_seq_shape: tuple[int, int, int],
|
||||
tile_size: tuple[int, int, int],
|
||||
device: torch.device,
|
||||
) -> torch.LongTensor:
|
||||
"""Inverse mapping: tile-contiguous order back to raster order."""
|
||||
return torch.argsort(get_tile_partition_indices(dit_seq_shape, tile_size, device))
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# BSA core operations
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _prune_queries(
|
||||
q_blocks: torch.Tensor,
|
||||
keep_ratio: float,
|
||||
) -> tuple[torch.Tensor, torch.Tensor, int]:
|
||||
"""
|
||||
Prune redundant query tokens within each block.
|
||||
|
||||
Scores tokens by cosine similarity to the block center.
|
||||
Keeps the LEAST similar (most informative) tokens.
|
||||
|
||||
Args:
|
||||
q_blocks: [B, N_heads, N_blocks, block_size, D]
|
||||
keep_ratio: fraction of tokens to keep
|
||||
|
||||
Returns:
|
||||
sparse_q: [B, N_heads, N_blocks, keep_size, D]
|
||||
keep_indices: [B, N_heads, N_blocks, keep_size]
|
||||
keep_size: int
|
||||
"""
|
||||
B, H, N, S, D = q_blocks.shape
|
||||
keep_size = max(1, int(S * keep_ratio))
|
||||
|
||||
if keep_size >= S:
|
||||
idx = torch.arange(S, device=q_blocks.device)
|
||||
idx = idx.view(1, 1, 1, S).expand(B, H, N, S)
|
||||
return q_blocks, idx, S
|
||||
|
||||
center_idx = S // 2
|
||||
center = q_blocks[:, :, :, center_idx:center_idx + 1, :]
|
||||
|
||||
q_norm = F.normalize(q_blocks, dim=-1)
|
||||
c_norm = F.normalize(center, dim=-1)
|
||||
similarity = (q_norm * c_norm).sum(dim=-1) # [B, H, N, S]
|
||||
|
||||
# lowest similarity = most distinctive = keep
|
||||
_, indices = similarity.topk(keep_size, dim=-1, largest=False)
|
||||
indices, _ = indices.sort(dim=-1)
|
||||
|
||||
idx_expand = indices.unsqueeze(-1).expand(-1, -1, -1, -1, D)
|
||||
sparse_q = torch.gather(q_blocks, 3, idx_expand)
|
||||
|
||||
return sparse_q, indices, keep_size
|
||||
|
||||
|
||||
def _select_kv_blocks(
|
||||
sparse_q: torch.Tensor,
|
||||
k_blocks: torch.Tensor,
|
||||
cumulative_threshold: float,
|
||||
min_kv_blocks: int,
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Dynamically select KV blocks for each query block.
|
||||
|
||||
Mean-pools to block level, computes block attention scores,
|
||||
admits blocks in descending order until cumulative mass
|
||||
exceeds threshold.
|
||||
|
||||
Args:
|
||||
sparse_q: [B, H, N, Sq, D]
|
||||
k_blocks: [B, H, N, Sk, D]
|
||||
cumulative_threshold: e.g. 0.9
|
||||
min_kv_blocks: minimum blocks to keep
|
||||
|
||||
Returns:
|
||||
kv_mask: [B, H, N, N] boolean
|
||||
"""
|
||||
B, H, N, _, D = sparse_q.shape
|
||||
|
||||
q_repr = sparse_q.mean(dim=3)
|
||||
k_repr = k_blocks.mean(dim=3)
|
||||
|
||||
scores = torch.matmul(q_repr, k_repr.transpose(-1, -2)) / (D**0.5)
|
||||
block_attn = F.softmax(scores, dim=-1)
|
||||
|
||||
sorted_attn, sorted_idx = block_attn.sort(dim=-1, descending=True)
|
||||
cumsum = sorted_attn.cumsum(dim=-1)
|
||||
|
||||
keep_sorted = torch.ones_like(cumsum, dtype=torch.bool)
|
||||
keep_sorted[..., 1:] = cumsum[..., :-1] < cumulative_threshold
|
||||
|
||||
min_mask = torch.zeros_like(keep_sorted)
|
||||
min_mask[..., :min(min_kv_blocks, N)] = True
|
||||
keep_sorted = keep_sorted | min_mask
|
||||
|
||||
kv_mask = torch.zeros_like(block_attn, dtype=torch.bool)
|
||||
kv_mask.scatter_(-1, sorted_idx, keep_sorted)
|
||||
|
||||
return kv_mask
|
||||
|
||||
|
||||
def _compute_sparse_attention(
|
||||
sparse_q: torch.Tensor,
|
||||
k_blocks: torch.Tensor,
|
||||
v_blocks: torch.Tensor,
|
||||
kv_mask: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Compute attention for each query block against selected KV blocks.
|
||||
|
||||
Handles per-batch and per-head KV masks correctly.
|
||||
Uses flash_attn_varlen_func when available on GPU.
|
||||
Falls back to pure-PyTorch reference on CPU.
|
||||
|
||||
Args:
|
||||
sparse_q: [B, H, N, Sq, D]
|
||||
k_blocks: [B, H, N, Sk, D]
|
||||
v_blocks: [B, H, N, Sk, D]
|
||||
kv_mask: [B, H, N, N] boolean (per-batch, per-head)
|
||||
|
||||
Returns:
|
||||
output: [B, H, N, Sq, D]
|
||||
"""
|
||||
if FLASH_ATTN_AVAILABLE and sparse_q.is_cuda:
|
||||
return _compute_sparse_attention_flash(sparse_q, k_blocks, v_blocks, kv_mask)
|
||||
else:
|
||||
return _compute_sparse_attention_reference(sparse_q, k_blocks, v_blocks, kv_mask)
|
||||
|
||||
|
||||
def _compute_sparse_attention_reference(
|
||||
sparse_q: torch.Tensor,
|
||||
k_blocks: torch.Tensor,
|
||||
v_blocks: torch.Tensor,
|
||||
kv_mask: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
"""Pure-PyTorch fallback with per-batch, per-head mask support."""
|
||||
B, H, N, Sq, D = sparse_q.shape
|
||||
output = torch.zeros_like(sparse_q)
|
||||
|
||||
for b in range(B):
|
||||
for h in range(H):
|
||||
for qb in range(N):
|
||||
selected = kv_mask[b, h, qb] # [N] boolean
|
||||
sel_idx = selected.nonzero(as_tuple=True)[0]
|
||||
|
||||
if sel_idx.shape[0] == 0:
|
||||
continue
|
||||
|
||||
# [num_sel * Sk, D]
|
||||
sel_k = k_blocks[b, h, sel_idx].reshape(-1, D)
|
||||
sel_v = v_blocks[b, h, sel_idx].reshape(-1, D)
|
||||
|
||||
q = sparse_q[b, h, qb] # [Sq, D]
|
||||
scores = torch.matmul(q, sel_k.transpose(-1, -2)) / (D**0.5)
|
||||
weights = F.softmax(scores, dim=-1)
|
||||
output[b, h, qb] = torch.matmul(weights, sel_v)
|
||||
|
||||
return output
|
||||
|
||||
|
||||
def _compute_sparse_attention_flash(
|
||||
sparse_q: torch.Tensor,
|
||||
k_blocks: torch.Tensor,
|
||||
v_blocks: torch.Tensor,
|
||||
kv_mask: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
FlashAttention implementation with per-batch, per-head mask support.
|
||||
|
||||
Strategy: check if all heads share the same mask. If so, use a single
|
||||
FlashAttention call per batch (fast path). If not, process each head
|
||||
separately (correct path).
|
||||
|
||||
Args:
|
||||
sparse_q: [B, H, N, Sq, D]
|
||||
k_blocks: [B, H, N, Sk, D]
|
||||
v_blocks: [B, H, N, Sk, D]
|
||||
kv_mask: [B, H, N, N] boolean
|
||||
|
||||
Returns:
|
||||
output: [B, H, N, Sq, D]
|
||||
"""
|
||||
B, H, N, Sq, D = sparse_q.shape
|
||||
Sk = k_blocks.shape[3]
|
||||
device = sparse_q.device
|
||||
output = torch.zeros_like(sparse_q)
|
||||
|
||||
for b in range(B):
|
||||
# Check if all heads share the same mask for this batch element
|
||||
# Compare each head's mask to head 0's mask
|
||||
head0_mask = kv_mask[b, 0] # [N, N]
|
||||
all_heads_same = all(torch.equal(kv_mask[b, h], head0_mask) for h in range(1, H))
|
||||
|
||||
if all_heads_same:
|
||||
# Fast path: all heads share the same mask, single FA call
|
||||
_flash_attn_single_mask(
|
||||
sparse_q[b],
|
||||
k_blocks[b],
|
||||
v_blocks[b],
|
||||
head0_mask,
|
||||
output[b],
|
||||
H,
|
||||
N,
|
||||
Sq,
|
||||
Sk,
|
||||
D,
|
||||
device,
|
||||
)
|
||||
else:
|
||||
# Per-head path: process each head individually
|
||||
for h in range(H):
|
||||
head_mask = kv_mask[b, h] # [N, N]
|
||||
# Process single head: squeeze head dim, run FA, put back
|
||||
_flash_attn_single_head(
|
||||
sparse_q[b, h],
|
||||
k_blocks[b, h],
|
||||
v_blocks[b, h],
|
||||
head_mask,
|
||||
output,
|
||||
b,
|
||||
h,
|
||||
N,
|
||||
Sq,
|
||||
Sk,
|
||||
D,
|
||||
device,
|
||||
)
|
||||
|
||||
return output
|
||||
|
||||
|
||||
def _flash_attn_single_mask(
|
||||
sparse_q_b: torch.Tensor, # [H, N, Sq, D]
|
||||
k_blocks_b: torch.Tensor, # [H, N, Sk, D]
|
||||
v_blocks_b: torch.Tensor, # [H, N, Sk, D]
|
||||
mask: torch.Tensor, # [N, N] boolean
|
||||
output_b: torch.Tensor, # [H, N, Sq, D] (modified in-place)
|
||||
H: int,
|
||||
N: int,
|
||||
Sq: int,
|
||||
Sk: int,
|
||||
D: int,
|
||||
device: torch.device,
|
||||
) -> None:
|
||||
"""Run FlashAttention for all heads sharing the same KV mask."""
|
||||
q_list = []
|
||||
k_list = []
|
||||
v_list = []
|
||||
cu_seqlens_q = [0]
|
||||
cu_seqlens_k = [0]
|
||||
active_blocks = []
|
||||
|
||||
for qb in range(N):
|
||||
selected = mask[qb] # [N] boolean
|
||||
sel_idx = selected.nonzero(as_tuple=True)[0]
|
||||
|
||||
if sel_idx.shape[0] == 0:
|
||||
continue
|
||||
|
||||
active_blocks.append(qb)
|
||||
num_kv_tokens = sel_idx.shape[0] * Sk
|
||||
|
||||
# [H, Sq, D] -> [Sq, H, D]
|
||||
q_block = sparse_q_b[:, qb].permute(1, 0, 2)
|
||||
q_list.append(q_block)
|
||||
|
||||
# [H, num_sel, Sk, D] -> [num_kv_tokens, H, D]
|
||||
sel_k = k_blocks_b[:, sel_idx].permute(1, 2, 0, 3).reshape(num_kv_tokens, H, D)
|
||||
sel_v = v_blocks_b[:, sel_idx].permute(1, 2, 0, 3).reshape(num_kv_tokens, H, D)
|
||||
k_list.append(sel_k)
|
||||
v_list.append(sel_v)
|
||||
|
||||
cu_seqlens_q.append(cu_seqlens_q[-1] + Sq)
|
||||
cu_seqlens_k.append(cu_seqlens_k[-1] + num_kv_tokens)
|
||||
|
||||
if not q_list:
|
||||
return
|
||||
|
||||
flat_q = torch.cat(q_list, dim=0)
|
||||
flat_k = torch.cat(k_list, dim=0)
|
||||
flat_v = torch.cat(v_list, dim=0)
|
||||
|
||||
# Compute max_seqlen_k from the Python list before moving to GPU to
|
||||
# avoid a `.item()` round-trip that would force a host/device sync.
|
||||
max_seqlen_q = Sq
|
||||
max_seqlen_k = max(b - a for a, b in zip(cu_seqlens_k[:-1], cu_seqlens_k[1:], strict=False))
|
||||
|
||||
cu_seqlens_q_t = torch.tensor(cu_seqlens_q, dtype=torch.int32, device=device)
|
||||
cu_seqlens_k_t = torch.tensor(cu_seqlens_k, dtype=torch.int32, device=device)
|
||||
|
||||
orig_dtype = flat_q.dtype
|
||||
compute_dtype = orig_dtype
|
||||
if compute_dtype not in (torch.float16, torch.bfloat16):
|
||||
compute_dtype = torch.bfloat16
|
||||
flat_q = flat_q.to(compute_dtype)
|
||||
flat_k = flat_k.to(compute_dtype)
|
||||
flat_v = flat_v.to(compute_dtype)
|
||||
|
||||
flat_out = flash_attn_varlen_func_impl(
|
||||
flat_q,
|
||||
flat_k,
|
||||
flat_v,
|
||||
cu_seqlens_q_t,
|
||||
cu_seqlens_k_t,
|
||||
max_seqlen_q,
|
||||
max_seqlen_k,
|
||||
causal=False,
|
||||
)
|
||||
|
||||
if compute_dtype != orig_dtype:
|
||||
flat_out = flat_out.to(orig_dtype)
|
||||
|
||||
idx = 0
|
||||
for qb in active_blocks:
|
||||
block_out = flat_out[idx:idx + Sq] # [Sq, H, D]
|
||||
output_b[:, qb] = block_out.permute(1, 0, 2) # [H, Sq, D]
|
||||
idx += Sq
|
||||
|
||||
|
||||
def _flash_attn_single_head(
|
||||
sparse_q_bh: torch.Tensor, # [N, Sq, D]
|
||||
k_blocks_bh: torch.Tensor, # [N, Sk, D]
|
||||
v_blocks_bh: torch.Tensor, # [N, Sk, D]
|
||||
mask: torch.Tensor, # [N, N] boolean
|
||||
output: torch.Tensor, # [B, H, N, Sq, D] (modified in-place)
|
||||
b: int,
|
||||
h: int,
|
||||
N: int,
|
||||
Sq: int,
|
||||
Sk: int,
|
||||
D: int,
|
||||
device: torch.device,
|
||||
) -> None:
|
||||
"""Run FlashAttention for a single head with its own KV mask."""
|
||||
q_list = []
|
||||
k_list = []
|
||||
v_list = []
|
||||
cu_seqlens_q = [0]
|
||||
cu_seqlens_k = [0]
|
||||
active_blocks = []
|
||||
|
||||
for qb in range(N):
|
||||
selected = mask[qb]
|
||||
sel_idx = selected.nonzero(as_tuple=True)[0]
|
||||
|
||||
if sel_idx.shape[0] == 0:
|
||||
continue
|
||||
|
||||
active_blocks.append(qb)
|
||||
num_kv_tokens = sel_idx.shape[0] * Sk
|
||||
|
||||
# [Sq, D] -> [Sq, 1, D] (single head)
|
||||
q_block = sparse_q_bh[qb].unsqueeze(1)
|
||||
q_list.append(q_block)
|
||||
|
||||
# [num_sel, Sk, D] -> [num_kv_tokens, 1, D]
|
||||
sel_k = k_blocks_bh[sel_idx].reshape(num_kv_tokens, 1, D)
|
||||
sel_v = v_blocks_bh[sel_idx].reshape(num_kv_tokens, 1, D)
|
||||
k_list.append(sel_k)
|
||||
v_list.append(sel_v)
|
||||
|
||||
cu_seqlens_q.append(cu_seqlens_q[-1] + Sq)
|
||||
cu_seqlens_k.append(cu_seqlens_k[-1] + num_kv_tokens)
|
||||
|
||||
if not q_list:
|
||||
return
|
||||
|
||||
flat_q = torch.cat(q_list, dim=0)
|
||||
flat_k = torch.cat(k_list, dim=0)
|
||||
flat_v = torch.cat(v_list, dim=0)
|
||||
|
||||
# Compute max_seqlen_k from the Python list before moving to GPU to
|
||||
# avoid a `.item()` round-trip that would force a host/device sync.
|
||||
max_seqlen_q = Sq
|
||||
max_seqlen_k = max(b - a for a, b in zip(cu_seqlens_k[:-1], cu_seqlens_k[1:], strict=False))
|
||||
|
||||
cu_seqlens_q_t = torch.tensor(cu_seqlens_q, dtype=torch.int32, device=device)
|
||||
cu_seqlens_k_t = torch.tensor(cu_seqlens_k, dtype=torch.int32, device=device)
|
||||
|
||||
orig_dtype = flat_q.dtype
|
||||
compute_dtype = orig_dtype
|
||||
if compute_dtype not in (torch.float16, torch.bfloat16):
|
||||
compute_dtype = torch.bfloat16
|
||||
flat_q = flat_q.to(compute_dtype)
|
||||
flat_k = flat_k.to(compute_dtype)
|
||||
flat_v = flat_v.to(compute_dtype)
|
||||
|
||||
flat_out = flash_attn_varlen_func_impl(
|
||||
flat_q,
|
||||
flat_k,
|
||||
flat_v,
|
||||
cu_seqlens_q_t,
|
||||
cu_seqlens_k_t,
|
||||
max_seqlen_q,
|
||||
max_seqlen_k,
|
||||
causal=False,
|
||||
)
|
||||
|
||||
if compute_dtype != orig_dtype:
|
||||
flat_out = flat_out.to(orig_dtype)
|
||||
|
||||
idx = 0
|
||||
for qb in active_blocks:
|
||||
block_out = flat_out[idx:idx + Sq] # [Sq, 1, D]
|
||||
output[b, h, qb] = block_out.squeeze(1) # [Sq, D]
|
||||
idx += Sq
|
||||
|
||||
|
||||
def _reconstruct_pruned(
|
||||
sparse_output: torch.Tensor,
|
||||
keep_indices: torch.Tensor,
|
||||
block_size: int,
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Scatter sparse output back to full block size.
|
||||
Pruned positions get nearest kept token's output.
|
||||
|
||||
Handles per-batch, per-head indices correctly.
|
||||
|
||||
Args:
|
||||
sparse_output: [B, H, N, keep_size, D]
|
||||
keep_indices: [B, H, N, keep_size]
|
||||
block_size: original tokens per block
|
||||
|
||||
Returns:
|
||||
full_output: [B, H, N, block_size, D]
|
||||
"""
|
||||
B, H, N, keep_size, D = sparse_output.shape
|
||||
device = sparse_output.device
|
||||
|
||||
if keep_size >= block_size:
|
||||
return sparse_output
|
||||
|
||||
full_output = torch.zeros(B, H, N, block_size, D, device=device, dtype=sparse_output.dtype)
|
||||
|
||||
# Scatter kept tokens
|
||||
idx_expand = keep_indices.unsqueeze(-1).expand(-1, -1, -1, -1, D)
|
||||
full_output.scatter_(3, idx_expand, sparse_output)
|
||||
|
||||
# Fill pruned positions with nearest kept token (vectorized)
|
||||
all_pos = torch.arange(block_size, device=device)
|
||||
|
||||
for b in range(B):
|
||||
for h in range(H):
|
||||
for n in range(N):
|
||||
kept = keep_indices[b, h, n] # [keep_size]
|
||||
|
||||
# Distance from every position to every kept position
|
||||
dists = (all_pos.view(-1, 1) - kept.view(1, -1)).abs()
|
||||
nearest_local_idx = dists.argmin(dim=1) # [block_size]
|
||||
|
||||
# Identify pruned positions
|
||||
is_pruned = torch.ones(block_size, dtype=torch.bool, device=device)
|
||||
is_pruned[kept] = False
|
||||
pruned_indices = is_pruned.nonzero(as_tuple=True)[0]
|
||||
|
||||
if pruned_indices.numel() > 0:
|
||||
src_indices = nearest_local_idx[pruned_indices]
|
||||
full_output[b, h, n, pruned_indices] = sparse_output[b, h, n, src_indices]
|
||||
|
||||
return full_output
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# FastVideo backend classes
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class BSAAttentionBackend(AttentionBackend):
|
||||
|
||||
accept_output_buffer: bool = False
|
||||
|
||||
@staticmethod
|
||||
def get_supported_head_sizes() -> list[int]:
|
||||
return [64, 128]
|
||||
|
||||
@staticmethod
|
||||
def get_name() -> str:
|
||||
return "BSA_ATTN"
|
||||
|
||||
@staticmethod
|
||||
def get_impl_cls() -> type["BSAAttentionImpl"]:
|
||||
return BSAAttentionImpl
|
||||
|
||||
@staticmethod
|
||||
def get_metadata_cls() -> type["BSAAttentionMetadata"]:
|
||||
return BSAAttentionMetadata
|
||||
|
||||
@staticmethod
|
||||
def get_builder_cls() -> type["BSAAttentionMetadataBuilder"]:
|
||||
return BSAAttentionMetadataBuilder
|
||||
|
||||
|
||||
@dataclass
|
||||
class BSAAttentionMetadata(AttentionMetadata):
|
||||
current_timestep: int
|
||||
dit_seq_shape: tuple[int, int, int]
|
||||
total_seq_length: int
|
||||
num_blocks: int
|
||||
block_size: int
|
||||
tile_partition_indices: torch.LongTensor
|
||||
reverse_tile_partition_indices: torch.LongTensor
|
||||
# BSA-specific config
|
||||
query_keep_ratio: float
|
||||
kv_cumulative_threshold: float
|
||||
min_kv_blocks: int
|
||||
|
||||
|
||||
class BSAAttentionMetadataBuilder(AttentionMetadataBuilder):
|
||||
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
def prepare(self):
|
||||
pass
|
||||
|
||||
def build(
|
||||
self,
|
||||
current_timestep: int,
|
||||
raw_latent_shape: tuple[int, int, int],
|
||||
patch_size: tuple[int, int, int],
|
||||
device: torch.device,
|
||||
bsa_query_keep_ratio: float = 0.5,
|
||||
bsa_kv_cumulative_threshold: float = 0.9,
|
||||
bsa_min_kv_blocks: int = 4,
|
||||
**kwargs: dict[str, Any],
|
||||
) -> "BSAAttentionMetadata":
|
||||
# Ensure patching does not drop tokens silently.
|
||||
assert all(r % p == 0 for r, p in zip(raw_latent_shape, patch_size, strict=False)), (
|
||||
"raw_latent_shape must be divisible by patch_size for BSA", )
|
||||
|
||||
dit_seq_shape = (
|
||||
raw_latent_shape[0] // patch_size[0],
|
||||
raw_latent_shape[1] // patch_size[1],
|
||||
raw_latent_shape[2] // patch_size[2],
|
||||
)
|
||||
|
||||
total_seq_length = math.prod(dit_seq_shape)
|
||||
block_size = math.prod(BSA_TILE_SIZE)
|
||||
# Require exact tiling to avoid reshape failures later.
|
||||
assert all(d % t == 0 for d, t in zip(dit_seq_shape, BSA_TILE_SIZE, strict=False)), (
|
||||
"dit_seq_shape must be divisible by BSA_TILE_SIZE", )
|
||||
num_blocks = total_seq_length // block_size
|
||||
|
||||
tile_partition_indices = get_tile_partition_indices(dit_seq_shape, BSA_TILE_SIZE, device)
|
||||
reverse_tile_partition_indices = get_reverse_tile_partition_indices(dit_seq_shape, BSA_TILE_SIZE, device)
|
||||
|
||||
return BSAAttentionMetadata(
|
||||
current_timestep=current_timestep,
|
||||
dit_seq_shape=dit_seq_shape,
|
||||
total_seq_length=total_seq_length,
|
||||
num_blocks=num_blocks,
|
||||
block_size=block_size,
|
||||
tile_partition_indices=tile_partition_indices,
|
||||
reverse_tile_partition_indices=reverse_tile_partition_indices,
|
||||
query_keep_ratio=bsa_query_keep_ratio,
|
||||
kv_cumulative_threshold=bsa_kv_cumulative_threshold,
|
||||
min_kv_blocks=bsa_min_kv_blocks,
|
||||
)
|
||||
|
||||
|
||||
class BSAAttentionImpl(AttentionImpl):
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
num_heads: int,
|
||||
head_size: int,
|
||||
causal: bool,
|
||||
softmax_scale: float,
|
||||
num_kv_heads: int | None = None,
|
||||
prefix: str = "",
|
||||
**extra_impl_args,
|
||||
) -> None:
|
||||
self.prefix = prefix
|
||||
self.num_heads = num_heads
|
||||
self.head_size = head_size
|
||||
if num_kv_heads is not None and num_kv_heads != num_heads:
|
||||
raise ValueError("BSA backend does not support grouped-query attention")
|
||||
if causal:
|
||||
raise ValueError("BSA backend is bidirectional; causal=True is unsupported")
|
||||
if softmax_scale is not None:
|
||||
expected_scale = 1.0 / math.sqrt(self.head_size)
|
||||
if not math.isclose(softmax_scale, expected_scale, rel_tol=1e-4, abs_tol=1e-5):
|
||||
raise ValueError("softmax_scale must be default (1/sqrt(d)) for BSA")
|
||||
try:
|
||||
sp_group = get_sp_group()
|
||||
self.sp_size = sp_group.world_size
|
||||
except (AssertionError, RuntimeError):
|
||||
self.sp_size = 1
|
||||
|
||||
def preprocess_qkv(
|
||||
self,
|
||||
qkv: torch.Tensor,
|
||||
attn_metadata: BSAAttentionMetadata,
|
||||
) -> torch.Tensor:
|
||||
"""Reorder tokens from raster order to tile-contiguous order."""
|
||||
# qkv: [B, L, num_heads, D]
|
||||
return qkv[:, attn_metadata.tile_partition_indices]
|
||||
|
||||
def postprocess_output(
|
||||
self,
|
||||
output: torch.Tensor,
|
||||
attn_metadata: BSAAttentionMetadata,
|
||||
) -> torch.Tensor:
|
||||
"""Reorder tokens from tile-contiguous order back to raster order."""
|
||||
return output[:, attn_metadata.reverse_tile_partition_indices]
|
||||
|
||||
def forward(
|
||||
self,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
attn_metadata: BSAAttentionMetadata,
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
BSA attention forward pass.
|
||||
|
||||
Input tensors are already in tile-contiguous order from preprocess_qkv.
|
||||
|
||||
Args:
|
||||
query: [B, L, num_heads, D] (tile-ordered)
|
||||
key: [B, L, num_heads, D] (tile-ordered)
|
||||
value: [B, L, num_heads, D] (tile-ordered)
|
||||
attn_metadata: BSA metadata
|
||||
|
||||
Returns:
|
||||
output: [B, L, num_heads, D] (tile-ordered)
|
||||
"""
|
||||
B, L, H, D = query.shape
|
||||
block_size = attn_metadata.block_size
|
||||
num_blocks = attn_metadata.num_blocks
|
||||
assert num_blocks * block_size == L, "Sequence length must match tiling"
|
||||
|
||||
# Reshape to [B, H, L, D] for attention computation
|
||||
q = query.transpose(1, 2).contiguous() # [B, H, L, D]
|
||||
k = key.transpose(1, 2).contiguous()
|
||||
v = value.transpose(1, 2).contiguous()
|
||||
|
||||
# Reshape into blocks: [B, H, num_blocks, block_size, D]
|
||||
q_blocks = q.view(B, H, num_blocks, block_size, D)
|
||||
k_blocks = k.view(B, H, num_blocks, block_size, D)
|
||||
v_blocks = v.view(B, H, num_blocks, block_size, D)
|
||||
|
||||
# --- Query sparsification ---
|
||||
sparse_q, keep_indices, keep_size = _prune_queries(q_blocks, attn_metadata.query_keep_ratio)
|
||||
|
||||
# --- KV block selection ---
|
||||
kv_mask = _select_kv_blocks(
|
||||
sparse_q,
|
||||
k_blocks,
|
||||
attn_metadata.kv_cumulative_threshold,
|
||||
attn_metadata.min_kv_blocks,
|
||||
)
|
||||
|
||||
# --- Sparse attention ---
|
||||
sparse_output = _compute_sparse_attention(sparse_q, k_blocks, v_blocks, kv_mask)
|
||||
|
||||
# --- Reconstruct pruned positions ---
|
||||
full_output = _reconstruct_pruned(sparse_output, keep_indices, block_size)
|
||||
|
||||
# Reshape back: [B, H, num_blocks, block_size, D] -> [B, H, L, D] -> [B, L, H, D]
|
||||
hidden_states = full_output.view(B, H, L, D).transpose(1, 2)
|
||||
|
||||
return hidden_states
|
||||
@@ -1,341 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
import os
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from dataclasses import dataclass
|
||||
|
||||
try:
|
||||
from v2._vendor.attention.utils.flash_attn_cute import flash_attn_func
|
||||
|
||||
fa_version = "4"
|
||||
except ImportError:
|
||||
try:
|
||||
from flash_attn_interface import flash_attn_func as flash_attn_3_func
|
||||
|
||||
# flash_attn 3 no longer have a different API, see following commit:
|
||||
# https://github.com/Dao-AILab/flash-attention/commit/ed209409acedbb2379f870bbd03abce31a7a51b7
|
||||
flash_attn_func = flash_attn_3_func
|
||||
fa_version = "3"
|
||||
except ImportError:
|
||||
from flash_attn import flash_attn_func as flash_attn_2_func
|
||||
flash_attn_func = flash_attn_2_func
|
||||
fa_version = "2"
|
||||
|
||||
# torch.compile traceability: the FA4/cute path (fa_version=="4") is
|
||||
# already a registered torch.library custom op, so dynamo treats it as a
|
||||
# graph node. The external FA2/FA3 `flash_attn_func` is NOT — dynamo
|
||||
# breaks the graph at the call site (observed: wanvideo.py self-attn,
|
||||
# once per layer every step), which fragments the compiled region and
|
||||
# blocks CUDA-graph capture. Wrap the FA2/FA3 default call in a custom
|
||||
# op (mirrors the FP4 `flash_attn_cute` template) so it becomes an
|
||||
# opaque-but-traceable node. The kernel still runs eager inside the op
|
||||
# (correct — flash-attn must run eager); only dynamo's treatment of the
|
||||
# boundary changes, so numerics are unchanged (SSIM-gate to confirm).
|
||||
if fa_version in ("2", "3"):
|
||||
_fa_default = flash_attn_func
|
||||
|
||||
# Scope: this op covers exactly the q/k/v + softmax_scale + causal
|
||||
# call shape used by FlashAttentionImpl.forward's default branch
|
||||
# (see `flash_attn_func_compilable(...)` call site below). The
|
||||
# masked/no-pad and varlen / cross-attn paths use different
|
||||
# entry points (`flash_attn_no_pad`, `flash_attn_varlen_*`) which
|
||||
# are intentionally out of scope for this PR — wrapping them is a
|
||||
# natural follow-up. The wrapper's signature is the contract: any
|
||||
# extra kwarg (dropout_p, window_size, alibi_slopes, deterministic,
|
||||
# return_attn_probs, ...) raises TypeError at the call site, so
|
||||
# silent loss of kwargs is not a failure mode.
|
||||
@torch.library.custom_op(
|
||||
"fastvideo::_flash_attn_default_forward",
|
||||
mutates_args=(),
|
||||
device_types="cuda",
|
||||
)
|
||||
def _flash_attn_default_forward(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
v: torch.Tensor,
|
||||
softmax_scale: float | None,
|
||||
causal: bool,
|
||||
) -> torch.Tensor:
|
||||
return _fa_default(q, k, v, softmax_scale=softmax_scale, causal=causal)
|
||||
|
||||
@torch.library.register_fake("fastvideo::_flash_attn_default_forward")
|
||||
def _flash_attn_default_forward_fake(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
v: torch.Tensor,
|
||||
softmax_scale: float | None,
|
||||
causal: bool,
|
||||
) -> torch.Tensor:
|
||||
del softmax_scale, causal
|
||||
# FA2/FA3 default path: [batch, seqlen_q, nheads, head_dim_v],
|
||||
# same dtype/device as q (head dim taken from v).
|
||||
return q.new_empty(q.shape[0], q.shape[1], q.shape[2], v.shape[-1])
|
||||
|
||||
def flash_attn_func_compilable(q, k, v, softmax_scale=None, causal=False):
|
||||
# Autograd carve-out. The custom op above registers a forward + fake
|
||||
# kernel but NO backward (register_autograd), so it is opaque to
|
||||
# autograd. Inference runs under no_grad / inference_mode and routes
|
||||
# through the traceable custom op — that is the torch.compile win, and
|
||||
# the only path this PR claims. Training backprops through attention,
|
||||
# so route grad-enabled calls to the original FA2/FA3 `flash_attn_func`
|
||||
# (itself an autograd.Function, so backward is correct) at the cost of a
|
||||
# dynamo graph break on the training path — i.e. pre-PR behavior, no
|
||||
# regression. Full autograd parity for the custom op (mirroring the FP4
|
||||
# cute template) is a tracked follow-up.
|
||||
if torch.is_grad_enabled() and (q.requires_grad or k.requires_grad or v.requires_grad):
|
||||
return _fa_default(q, k, v, softmax_scale=softmax_scale, causal=causal)
|
||||
return torch.ops.v2._flash_attn_default_forward(q, k, v, softmax_scale, causal)
|
||||
elif fa_version == "4":
|
||||
# FA4 path: `flash_attn_func` is already a torch.library custom op
|
||||
# (registered in `v2._vendor.attention.utils.flash_attn_cute`), so a
|
||||
# passthrough is enough — no extra registration needed.
|
||||
def flash_attn_func_compilable(q, k, v, softmax_scale=None, causal=False):
|
||||
return flash_attn_func(q, k, v, softmax_scale=softmax_scale, causal=causal)
|
||||
else:
|
||||
# Defensive: the probe above only ever sets fa_version to "2", "3",
|
||||
# or "4"; an unexpected value means an import/probe regression and
|
||||
# we want a loud error at import, not a silent NameError later.
|
||||
raise RuntimeError(f"Unsupported FlashAttention version: {fa_version!r} — expected "
|
||||
f"'2', '3', or '4' from the import probe above.")
|
||||
|
||||
from v2._vendor.attention.backends.abstract import (
|
||||
AttentionBackend,
|
||||
AttentionImpl,
|
||||
AttentionMetadata,
|
||||
AttentionMetadataBuilder,
|
||||
)
|
||||
from v2._vendor.logger import init_logger
|
||||
|
||||
logger = init_logger(__name__)
|
||||
logger.info("Using FlashAttention-%s backend", fa_version)
|
||||
|
||||
# FP4 FA4 support: quantize Q/K to NVFP4 E2M1 for block-scaled MMA on Blackwell.
|
||||
# Requires: flash-attention-fp4, flashinfer, cutlass-dsl. Enable via nvfp4_fa4=True kwarg.
|
||||
# The FP4 path uses a dedicated custom_op wrapper (flash_attn_fp4_func) so that
|
||||
# torch.compile treats the CuTeDSL kernel as an opaque boundary.
|
||||
try:
|
||||
from v2._vendor.attention.utils.flash_attn_cute import flash_attn_fp4_func
|
||||
_FA4_FP4_AVAILABLE = True
|
||||
except ImportError:
|
||||
flash_attn_fp4_func = None
|
||||
_FA4_FP4_AVAILABLE = False
|
||||
|
||||
|
||||
def _nvfp4_quantize_for_fa4(tensor_4d: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Quantize a (batch, seqlen, nheads, headdim) BF16 tensor to FP4.
|
||||
|
||||
Returns:
|
||||
fp4_tensor: torch.float4_e2m1fn_x2, shape (batch, seqlen_padded, nheads, headdim//2)
|
||||
where seqlen_padded is seqlen rounded up to multiple of 128.
|
||||
Caller should slice [:, :orig_seqlen] before passing to FA4.
|
||||
sf_tensor: torch.uint8, shape (32, 4, rest_m, 4, rest_k, nheads, batch) with stride[3]=1
|
||||
"""
|
||||
from flashinfer.quantization import nvfp4_quantize, SfLayout
|
||||
|
||||
batch, seqlen, nheads, headdim = tensor_4d.shape
|
||||
sf_vec_size = 16
|
||||
|
||||
# Pad seqlen to multiple of 128 (required by nvfp4_quantize layout_128x4)
|
||||
tile_m = 128
|
||||
seqlen_padded = (seqlen + tile_m - 1) // tile_m * tile_m
|
||||
if seqlen_padded != seqlen:
|
||||
tensor_4d = F.pad(tensor_4d, (0, 0, 0, 0, 0, seqlen_padded - seqlen))
|
||||
|
||||
# Quantize with nheads squashed into K dimension so M=batch*seqlen (divisible by 128)
|
||||
# and K=nheads*headdim. This ensures 128-row SF tiles align with seqlen boundaries.
|
||||
t2d = tensor_4d.reshape(batch * seqlen_padded, nheads * headdim)
|
||||
one = torch.ones(1, device=t2d.device, dtype=torch.float32)
|
||||
fp4_data, sf_data = nvfp4_quantize(t2d, one, sfLayout=SfLayout.layout_128x4, do_shuffle=False)
|
||||
|
||||
# FP4 data: (batch*seqlen, nheads*headdim/2) → (batch, seqlen, nheads, headdim/2)
|
||||
fp4_tensor = (fp4_data.reshape(batch, seqlen_padded, nheads,
|
||||
headdim // 2).view(torch.int8).view(torch.float4_e2m1fn_x2))
|
||||
|
||||
# SF layout conversion: nvfp4_quantize layout_128x4 → FA4 MMA layout
|
||||
# layout_128x4 buffer: [mTile, kTile, 32, 4, 4]
|
||||
# FA4 expects: (32, 4, rest_m, 4, rest_k, nheads, batch) with stride[3]=1
|
||||
atom_m0, atom_m1, atom_k = 32, 4, 4
|
||||
rest_m = seqlen_padded // tile_m
|
||||
sf_k_per_head = headdim // sf_vec_size # 8 for headdim=128
|
||||
rest_k = sf_k_per_head // atom_k # 2
|
||||
|
||||
total_m_tiles = batch * rest_m
|
||||
total_k_tiles = (nheads * sf_k_per_head) // atom_k
|
||||
|
||||
sf_swizzled = sf_data.reshape(total_m_tiles, total_k_tiles, atom_m0, atom_m1, atom_k)
|
||||
sf_decomposed = sf_swizzled.reshape(batch, rest_m, nheads, rest_k, atom_m0, atom_m1, atom_k)
|
||||
sf_canonical = sf_decomposed.permute(0, 2, 1, 3, 4, 5, 6).contiguous()
|
||||
sf_mma = sf_canonical.permute(4, 5, 2, 6, 3, 1, 0)
|
||||
|
||||
return fp4_tensor, sf_mma
|
||||
|
||||
|
||||
class FlashAttentionBackend(AttentionBackend):
|
||||
accept_output_buffer: bool = True
|
||||
|
||||
@staticmethod
|
||||
def get_supported_head_sizes() -> list[int]:
|
||||
return [32, 64, 96, 128, 160, 192, 224, 256]
|
||||
|
||||
@staticmethod
|
||||
def get_name() -> str:
|
||||
return "FLASH_ATTN"
|
||||
|
||||
@staticmethod
|
||||
def get_impl_cls() -> type["FlashAttentionImpl"]:
|
||||
return FlashAttentionImpl
|
||||
|
||||
@staticmethod
|
||||
def get_metadata_cls() -> type["AttentionMetadata"]:
|
||||
raise NotImplementedError
|
||||
|
||||
@staticmethod
|
||||
def get_builder_cls() -> type["AttentionMetadataBuilder"]:
|
||||
raise NotImplementedError
|
||||
|
||||
|
||||
def _key_padding_mask_from_attn_mask(attn_mask: torch.Tensor, key_len: int) -> torch.Tensor:
|
||||
# Normalize attn_mask to [B, key_len] where True means valid token.
|
||||
if attn_mask.dim() == 4:
|
||||
attn_mask = attn_mask[:, 0, 0, :]
|
||||
elif attn_mask.dim() == 3:
|
||||
attn_mask = attn_mask[:, 0, :]
|
||||
elif attn_mask.dim() != 2:
|
||||
raise ValueError(f"Unsupported attn_mask shape for FLASH_ATTN: {attn_mask.shape}")
|
||||
|
||||
# SDPA additive mask convention: valid=0, masked=-inf/large negative.
|
||||
key_padding_mask = attn_mask if attn_mask.dtype == torch.bool else attn_mask >= 0
|
||||
|
||||
if key_padding_mask.shape[-1] != key_len:
|
||||
raise ValueError("Invalid key padding mask length for FLASH_ATTN: "
|
||||
f"expected {key_len}, got {key_padding_mask.shape[-1]}")
|
||||
return key_padding_mask
|
||||
|
||||
|
||||
@dataclass
|
||||
class FlashAttnMetadata(AttentionMetadata):
|
||||
current_timestep: int
|
||||
attn_mask: torch.Tensor | None = None
|
||||
|
||||
|
||||
class FlashAttnMetadataBuilder(AttentionMetadataBuilder):
|
||||
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
def prepare(self):
|
||||
pass
|
||||
|
||||
def build( # type: ignore
|
||||
self,
|
||||
current_timestep: int,
|
||||
attn_mask: torch.Tensor,
|
||||
) -> FlashAttnMetadata:
|
||||
return FlashAttnMetadata(current_timestep=current_timestep, attn_mask=attn_mask)
|
||||
|
||||
|
||||
class FlashAttentionImpl(AttentionImpl):
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
num_heads: int,
|
||||
head_size: int,
|
||||
causal: bool,
|
||||
softmax_scale: float,
|
||||
num_kv_heads: int | None = None,
|
||||
prefix: str = "",
|
||||
**extra_impl_args,
|
||||
) -> None:
|
||||
self.causal = causal
|
||||
self.softmax_scale = softmax_scale
|
||||
self.nvfp4_fa4 = extra_impl_args.get("nvfp4_fa4", False) or os.environ.get("FASTVIDEO_NVFP4_FA4", "0") == "1"
|
||||
if self.nvfp4_fa4:
|
||||
cap = torch.cuda.get_device_capability()
|
||||
assert cap in [(10, 0), (10, 3)], (f"NVFP4 FA4 requires Blackwell (sm100a/sm103a), got sm{cap[0]}{cap[1]}")
|
||||
assert _FA4_FP4_AVAILABLE, ("NVFP4 FA4 requires flash-attention-fp4 (flash_attn.cute). "
|
||||
"Install via instructions in docs/inference/optimizations.md")
|
||||
logger.info("NVFP4 FA4 enabled for FlashAttentionImpl (quant_qk only)")
|
||||
|
||||
def forward(
|
||||
self,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
attn_metadata: FlashAttnMetadata,
|
||||
):
|
||||
if (attn_metadata is not None and hasattr(attn_metadata, "attn_mask") and attn_metadata.attn_mask is not None):
|
||||
from v2._vendor.attention.utils.flash_attn_no_pad import (
|
||||
flash_attn_no_pad,
|
||||
flash_attn_varlen_qk_no_pad,
|
||||
)
|
||||
|
||||
attn_mask = attn_metadata.attn_mask
|
||||
|
||||
# flash_attn_no_pad packs q/k/v as one tensor and assumes equal q/k
|
||||
# sequence lengths. Cross-attention can violate this.
|
||||
if query.shape[1] != key.shape[1]:
|
||||
query_padding_mask = torch.ones(
|
||||
(query.shape[0], query.shape[1]),
|
||||
dtype=torch.bool,
|
||||
device=query.device,
|
||||
)
|
||||
key_padding_mask = _key_padding_mask_from_attn_mask(attn_mask, key.shape[1]).to(device=key.device)
|
||||
|
||||
return flash_attn_varlen_qk_no_pad(
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
query_padding_mask=query_padding_mask,
|
||||
key_padding_mask=key_padding_mask,
|
||||
causal=self.causal,
|
||||
dropout_p=0.0,
|
||||
softmax_scale=self.softmax_scale,
|
||||
)
|
||||
|
||||
qkv = torch.stack([query, key, value], dim=2)
|
||||
attn_mask_padded = F.pad(attn_mask, (qkv.shape[1] - attn_mask.shape[1], 0), value=True)
|
||||
output = flash_attn_no_pad(qkv, attn_mask_padded, causal=False, dropout_p=0, softmax_scale=None)
|
||||
elif self.nvfp4_fa4:
|
||||
output = self._forward_nvfp4(query, key, value)
|
||||
|
||||
else:
|
||||
# Route through the compilable wrapper so dynamo sees a
|
||||
# registered op (no graph break) for FA2/FA3; identical
|
||||
# kernel + numerics, op runs eager internally.
|
||||
output = flash_attn_func_compilable(
|
||||
query, # type: ignore[no-untyped-call]
|
||||
key,
|
||||
value,
|
||||
softmax_scale=self.softmax_scale,
|
||||
causal=self.causal,
|
||||
)
|
||||
return output
|
||||
|
||||
def _forward_nvfp4(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor) -> torch.Tensor:
|
||||
"""FP4 flash attention with quantized Q and K, BF16 V."""
|
||||
orig_seqlen_q = query.shape[1]
|
||||
orig_seqlen_k = key.shape[1]
|
||||
|
||||
# Quantize Q/K to FP4 (internally pads to multiple of 128 for SF layout)
|
||||
q_fp4, q_sf = _nvfp4_quantize_for_fa4(query)
|
||||
k_fp4, k_sf = _nvfp4_quantize_for_fa4(key)
|
||||
|
||||
# Pass original seqlen to FA4 — the kernel handles non-multiple-of-128
|
||||
# via boundary masking. FP4/SF data is padded to 128-multiple but FA4
|
||||
# only attends to orig_seqlen positions, avoiding softmax bias on padding.
|
||||
q_fp4 = q_fp4[:, :orig_seqlen_q]
|
||||
k_fp4 = k_fp4[:, :orig_seqlen_k]
|
||||
|
||||
output = flash_attn_fp4_func(
|
||||
q_fp4,
|
||||
k_fp4,
|
||||
value,
|
||||
q_sf,
|
||||
k_sf,
|
||||
softmax_scale=self.softmax_scale,
|
||||
causal=self.causal,
|
||||
)
|
||||
if isinstance(output, tuple):
|
||||
output = output[0]
|
||||
return output
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user