Compare commits

...
Author SHA1 Message Date
SolitaryThinker b8abf6f36b [bugfix]: classify supports_dynamic_batching in schema parity inventory 2026-07-11 18:35:45 -07:00
SolitaryThinker b30903b649 [refactor]: route dynamic batching through the typed API (#1453)
Typed generate(request) prompt lists previously always expanded to
sequential single-request calls, so dynamic batching only served the
deprecated legacy-kwargs surface. When batching is enabled
(batching_mode='dynamic', batching_max_size>1) a typed prompt list now
expands into the same _GenerationWorkItem path the legacy prompt-file
and server scheduler routes use: shared FastVideoArgs (merge identity),
reserved output paths, compatibility-gated merging, and typed
GenerationResult wrapping with prompt_index.

Default (batching disabled) typed behavior is byte-identical; the new
route is only taken under the existing opt-in flags. Adds one CPU
regression test asserting a typed prompt list merges into a single
forward. Batching-touched suite: 140 -> 141 passed, zero skips.
2026-07-11 18:35:45 -07:00
SolitaryThinker 7592b57ec0 [refactor]: dedup single/batch generation paths (#1453)
_generate_single_video was a ~250-line copy of the batch helpers
(_prepare_generation_work_item / _run_forward_batch /
_postprocess_generation_output); it now delegates to them, so save/mux
fixes land once. prompt_embeds passthrough moves into
_prepare_generation_work_item. The per-request pipeline-override block
(three copies) is factored into _resolve_request_fastvideo_args.

Only intentional non-output-affecting delta: the single path no longer
overlaps the pinned staging-buffer allocation with the forward thread
(the shared helper allocates after join, as the batch path already did).

No behavior change; batching-touched suite 140 passed before and after.
2026-07-11 18:35:45 -07:00
SolitaryThinker e35fa5452a [test]: CPU regression tests for dynamic batching scheduler and signature fixes (#1453)
Covers the review's adversarial probes: expired-deadline greedy drain
coalesces jobs queued behind an in-flight generation (delay>0 no longer
degrades to sequential), submit-after-stop raises, stop() fails
still-queued futures without hanging, a dispatch exception fans out to
every future in the group, a crashed _run task fails in-flight and
queued futures and rejects later submits, the pipeline capability gate
rejects non-opted-in pipelines (DMD/causal always vetoed), pre-set
latents are never merged, and distinct opaque objects are conservatively
incompatible. Verified the deadline and gate tests fail on the pre-fix
scheduler.
2026-07-11 18:35:45 -07:00
SolitaryThinker 302f79b6f4 [bugfix]: park the in-flight job back in _pending if the scheduler loop crashes (#1453)
If _run dies between popping a job off the queue and resolving its
future, that job belonged to neither _queue nor _pending, so the crash
drain could not fail it and its submitter hung forever. Re-park it in
_pending on the way out so _on_run_done fails it like everything else.
2026-07-11 18:35:45 -07:00
SolitaryThinker efc53ba93b [bugfix]: gate dynamic request batching behind a pipeline capability flag (#1453)
can_dynamic_batch only inspects request fields, but several shipped
pipelines are structurally batch-size-1 and can corrupt outputs when
requests are merged. Add PipelineConfig.supports_dynamic_batching
(default False) with dynamic_batching_supported() as the effective
gate: the Wan family opts in (the generic Wan denoising path is the
verified parity target), while DMD and causal denoising variants are
always vetoed because those stages do not consume per-item generators
and merged requests would share RNG. The scheduler refuses to start and
_generate_prepared_work_items refuses to merge with a clear error when
the loaded pipeline has not opted in.
2026-07-11 18:35:45 -07:00
SolitaryThinker 5ed9c7b472 [bugfix]: never merge requests carrying pre-set latents or unrecognized objects (#1453)
SamplingParam.latents was neither signature-compared nor listed as
unsupported, so two requests with different pre-set latents could merge
and corrupt each other. List it as unsupported for dynamic batching, and
harden the signature freezer: unrecognized values fell back to repr(),
which truncates large tensors and can collide for different payloads —
compare unrecognized objects by identity instead so they are
conservatively incompatible unless literally shared.
2026-07-11 18:35:45 -07:00
SolitaryThinker 5e9c755e1e [bugfix]: log when the batch compatibility probe fails instead of silently disabling batching (#1453)
_jobs_are_compatible swallowed every exception and returned False, so a
broken probe silently degraded dynamic batching to sequential with zero
operator signal. Log a deduplicated warning (still returning False).
2026-07-11 18:35:45 -07:00
SolitaryThinker 2ca941a942 [bugfix]: fail jobs still queued when the batch scheduler stops (#1453)
The _run exit path only failed futures parked in _pending; jobs still
sitting in _queue when stop() landed mid-dispatch hung forever. Drain
the queue too (skipping None sentinels) and fail those futures with the
same 'stopped before dispatch' RuntimeError.
2026-07-11 18:35:45 -07:00
SolitaryThinker 457017e071 [bugfix]: survive asyncio.TimeoutError on py3.10 and fail futures if the scheduler dies (#1453)
asyncio.wait_for raises asyncio.TimeoutError on Python 3.10 (aliased to
the builtin only from 3.11), so the first expired batching window killed
the _run task silently: every in-flight future and all later submits
hung while the server kept accepting requests. Catch both exception
names, and add a done-callback on the scheduler task that logs the
crash, marks the scheduler stopped, and fails every job still waiting
in the queue or pending deque with a RuntimeError instead of hanging.
2026-07-11 18:35:45 -07:00
SolitaryThinker dcd2cba956 [bugfix]: coalesce in-flight arrivals when the batching deadline has expired (#1453)
When batching_delay_ms > 0, the collect deadline anchors to the first
job's enqueue time. Any job that waited behind an in-flight generation
arrives with an already-expired deadline and was dispatched as a batch
of 1, degrading dynamic batching to sequential under sustained load.
Fall through to the same non-blocking greedy drain the delay=0 path
uses so queued compatible jobs still coalesce up to batching_max_size.
2026-07-11 18:35:45 -07:00
SolitaryThinker cdd0f7f48f [misc]: drop .agents/exploration journals added by #1453
AGENTS.md bans experiment journals and branch-state snapshots under
.agents/.
2026-07-11 18:35:45 -07:00
SolitaryThinker 6ab03e1745 [misc]: collect fastvideo/tests/batching in the Modal unit-test lane (#1453)
The new batching unit tests (admission, signature) were not collected by
any CI lane. stages/ is already collected on main, so only batching/ is
added.
2026-07-11 18:35:45 -07:00
SolitaryThinker 06f354afa1 [bugfix]: keep completed results when a prompt-file batch group fails (#1453)
The dynamic-batching prompt-file branch let a single failed group raise
out of generate_video, discarding every completed result (and tripping
the strict zip). Restore the legacy per-prompt tolerance: a failed group
now yields an error entry per prompt, keeping results aligned with the
input prompts while completed videos survive.
2026-07-11 18:35:45 -07:00
SolitaryThinker a4f54702c3 [bugfix]: fail fast when submitting to a stopped batch scheduler (#1453)
submit() after stop() enqueued jobs the drained run loop would never
resolve, hanging the caller forever on the future. Raise immediately
instead.
2026-07-11 18:35:45 -07:00
SolitaryThinker 522b9f7e47 [bugfix]: drain queued jobs greedily when batching delay is 0 (#1453)
With the default batching_delay_ms=0, _collect_batch called
wait_for(timeout=0) which timed out before draining the queue, so every
dispatched batch had size 1 and dynamic batching silently no-oped. When
delay is 0, drain already-queued jobs non-blocking (until empty,
max_size, or an incompatible job) instead of waiting.
2026-07-11 18:35:45 -07:00
SolitaryThinker 961383b077 [bugfix]: reserve batch output paths so duplicate prompts do not collide (#1453)
Batch mode resolves every output path via _prepare_output_path before any
file is written, so the os.path.exists dedupe never fired and duplicate
prompts (or prompts equal in their first 100 sanitized chars) overwrote
each other. Track paths reserved within the batch and suffix duplicates,
matching what the legacy sequential path produces for the same inputs.
2026-07-11 18:35:45 -07:00
Mac Lee 49e4673f6e [bugfix]: address batching review comments (#1453) 2026-07-11 18:35:45 -07:00
Mac Lee e6c84fe2fd [bugfix]: update batching schema test fixtures (#1453) 2026-07-11 18:35:45 -07:00
Mac Lee 733462b4da [docs]: record checklist closure (#1453) 2026-07-11 18:35:45 -07:00
Mac Lee 0f413c2e8f [misc]: satisfy all-files pre-commit formatting 2026-07-11 18:35:45 -07:00
Mac Lee d8ece1f37f [docs]: record SSIM validation follow-up 2026-07-11 18:35:45 -07:00
Mac Lee 2633eb1525 [docs]: record final batching validation result 2026-07-11 18:35:45 -07:00
Mac Lee 8e2b9ced60 [docs]: record multimodal batching validation report 2026-07-11 18:35:45 -07:00
Mac Lee 45d7b55456 [fix]: harden dynamic generation batching 2026-07-11 18:35:45 -07:00
Mac Lee 400e57a9de [misc]: record batching stage 5 state 2026-07-11 18:35:45 -07:00
Mac Lee 5560eab1db [feat]: add OpenAI video batching scheduler 2026-07-11 18:35:45 -07:00
Mac Lee ff75cb2c97 [misc]: record batching stage 2 state 2026-07-11 18:35:45 -07:00
Mac Lee 8d4627845b [feat]: add generator dynamic batching path 2026-07-11 18:35:45 -07:00
Mac Lee 602d57c370 [misc]: record batching stage 1 state 2026-07-11 18:35:45 -07:00
Mac Lee a6aab65d97 [feat]: add generation batching primitives 2026-07-11 18:35:45 -07:00
Raghav K 055586703d [perf]: register a real backward for FA2 default + masked/varlen custom ops (training-under-compile) (#1388) 2026-07-12 00:45:27 +00:00
Mac Lee 5d89f86675 [ci] Stop forcing FA4 in model-load lanes (#1561) 2026-07-11 14:10:36 -07:00
Satyam Srivastava 19a51a1fe6 [ci] Trigger performance benchmarks for performance code changes (#1583) 2026-07-10 20:21:33 -07:00
William Lin d3232cea5a [ci]: gate the full-suite trigger on pre-commit and docs build (#1572) 2026-07-11 02:56:49 +00:00
42 changed files with 3953 additions and 267 deletions
+1
View File
@@ -481,6 +481,7 @@ steps:
- "fastvideo/layers/**"
- "fastvideo/worker/**"
- "fastvideo/entrypoints/**"
- "fastvideo/performance/**"
- "fastvideo/tests/performance/**"
- ".buildkite/performance-benchmarks/**"
- "pyproject.toml"
+133
View File
@@ -0,0 +1,133 @@
#!/usr/bin/env bash
# Gate the expensive Buildkite full suite on the cheap GitHub checks.
#
# Polls the workflow runs for the PR head commit and only exits 0 once the
# watched cheap workflows (pre-commit, docs build) have succeeded, so the
# 'ready' label cannot burn ~20 GPU lanes on a head that a cheap check has
# already doomed.
#
# Semantics:
# - watched run completed with a bad conclusion -> exit 1 (fail CLOSED:
# no full suite; the next push re-arms via the 'synchronize' trigger)
# - watched run cancelled -> still pending: the docs
# workflow's repo-global 'pages' concurrency group cancels runs superseded
# by unrelated pushes, so 'cancelled' is not a verdict on this PR
# - watched runs pending -> poll until done
# - docs run absent -> not applicable after a
# short grace period ('Deploy Documentation' is path-filtered on PRs)
# - pre-commit run absent -> keep polling: pre-commit
# is never path-filtered, so its absence is always anomalous
# - 'ready' label removed while waiting -> exit 1 (fail CLOSED:
# un-labeling is a deliberate maintainer action)
# - GitHub API unreachable or timeout -> exit 0 (fail OPEN,
# loud warning: never brick CI on a GitHub outage)
#
# Required env: PR_SHA (PR head commit), PR_NUMBER, GITHUB_REPOSITORY, GH_TOKEN.
set -euo pipefail
: "${PR_SHA:?PR_SHA (PR head commit) is required}"
: "${PR_NUMBER:?PR_NUMBER (pull request number) is required}"
: "${GITHUB_REPOSITORY:?GITHUB_REPOSITORY is required}"
# Workflow-level `name:` values that must be green before the full suite
# may start. "Deploy Documentation" is path-filtered on PRs, so its run may
# legitimately never exist; pre-commit always runs, so it must appear.
WATCHED_NAMES='["pre-commit", "Deploy Documentation"]'
WATCHED_REGEX='^(pre-commit|Deploy Documentation)$'
POLL_SECS="${POLL_SECS:-20}"
GRACE_SECS="${GRACE_SECS:-60}"
MAX_WAIT_SECS="${MAX_WAIT_SECS:-1500}"
# Bound each API call so a hung connection hits the 3-strike fail-open path
# instead of pinning the loop until the job timeout (which would fail closed
# on exactly the GitHub-outage case this script is meant to survive).
if command -v timeout >/dev/null 2>&1; then
gh_api() { timeout 30 gh api "$@"; }
else
gh_api() { gh api "$@"; } # macOS dev boxes; CI always has coreutils timeout
fi
# The workflow checked the label before starting the gate, but the wait can
# last ~25 min: re-check once before any exit 0 and fail closed if 'ready'
# was removed in the meantime. An API error here proceeds (the label was
# present when the gate started; never brick CI on an outage).
recheck_ready_label() {
local pr_json
if pr_json=$(gh_api "repos/${GITHUB_REPOSITORY}/pulls/${PR_NUMBER}" 2>/dev/null); then
if ! jq -e '[.labels[]?.name] | index("ready")' <<<"$pr_json" >/dev/null 2>&1; then
echo "::error::PR #${PR_NUMBER} no longer has the 'ready' label —" \
"NOT triggering the Buildkite full suite. Re-add the label to re-arm."
exit 1
fi
else
echo "::warning::Could not re-check the 'ready' label on PR #${PR_NUMBER}; proceeding (it was present when the gate started)."
fi
}
start=$(date +%s)
api_fails=0
missing=""
while true; do
elapsed=$(( $(date +%s) - start ))
if runs_json=$(gh_api "repos/${GITHUB_REPOSITORY}/actions/runs?head_sha=${PR_SHA}&per_page=100" 2>/dev/null) \
&& state=$(jq --arg re "$WATCHED_REGEX" '
[.workflow_runs[]? | select(.name // "" | test($re))]
| group_by(.name) | map(max_by(.id))
| map({name, status, conclusion})' <<<"$runs_json" 2>/dev/null); then
api_fails=0
echo "t+${elapsed}s watched checks: $(jq -c . <<<"$state")"
failed=$(jq -r '[.[] | select(.status == "completed"
and (.conclusion | IN("success", "skipped", "neutral", "cancelled") | not))]
| map(.name) | join(", ")' <<<"$state")
if [ -n "$failed" ]; then
echo "::error::Cheap check(s) failed on ${PR_SHA}: ${failed}." \
"NOT triggering the Buildkite full suite. Push a fix (the 'ready'" \
"label re-arms on every push), or re-run the failed check and then" \
"re-run this workflow."
exit 1
fi
# 'cancelled' counts as pending: wait for a re-run to reach a real verdict
# (bounded by MAX_WAIT, then the fail-open below).
pending=$(jq '[.[] | select(.status != "completed" or .conclusion == "cancelled")] | length' <<<"$state")
missing=$(jq -r --argjson watched "$WATCHED_NAMES" '($watched - map(.name)) | join(", ")' <<<"$state")
if [ "$pending" -eq 0 ]; then
if [ -z "$missing" ]; then
recheck_ready_label
echo "All watched cheap checks are green — full suite may proceed."
exit 0
fi
case "$missing" in
*pre-commit*)
echo "pre-commit run not found for ${PR_SHA} yet; waiting (pre-commit is never path-filtered, so its absence is anomalous)."
;;
*)
if [ "$elapsed" -ge "$GRACE_SECS" ]; then
recheck_ready_label
echo "::warning::Watched run(s) never appeared for ${PR_SHA}: ${missing} (path-filtered, likely not applicable). Proceeding on the checks that did run."
exit 0
fi
echo "Waiting up to ${GRACE_SECS}s grace for path-filtered run(s) to appear: ${missing}."
;;
esac
fi
else
api_fails=$(( api_fails + 1 ))
echo "::warning::GitHub API error querying workflow runs for ${PR_SHA} (attempt ${api_fails}/3)."
if [ "$api_fails" -ge 3 ]; then
recheck_ready_label
echo "::warning::FAILING OPEN: cannot query GitHub check status — triggering the full suite WITHOUT the cheap-check gate."
exit 0
fi
fi
if [ "$elapsed" -ge "$MAX_WAIT_SECS" ]; then
recheck_ready_label
echo "::warning::FAILING OPEN: watched checks still pending after $(( MAX_WAIT_SECS / 60 )) min${missing:+ (never appeared: ${missing})} — triggering the full suite anyway."
exit 0
fi
sleep "$POLL_SECS"
done
+122
View File
@@ -0,0 +1,122 @@
#!/usr/bin/env bash
# Self-test for gate_full_suite.sh using a mocked `gh`. No network, runs on
# any dev box: bash .github/scripts/test_gate_full_suite.sh
set -u
here=$(cd "$(dirname "$0")" && pwd)
tmp=$(mktemp -d)
trap 'rm -rf "$tmp"' EXIT
# Mock gh. Asserts the exact endpoint (including head_sha) it is called
# with — an endpoint typo in the gate script fails the test rather than
# silently serving canned data. On the runs endpoint it serves
# $MOCK_DIR/response_<call#>.json, sticking on the highest existing file,
# and exits 1 if none exist (simulates a GitHub API outage). On the pulls
# endpoint it serves $MOCK_DIR/pr.json, defaulting to a 'ready'-labeled PR.
cat > "$tmp/gh" <<'EOF'
#!/usr/bin/env bash
if [ "${1:-}" != "api" ]; then
echo "unexpected gh invocation: $*" >> "$MOCK_DIR/endpoint_error"
exit 2
fi
case "${2:-}" in
"repos/o/r/actions/runs?head_sha=deadbeef&per_page=100")
n=$(( $(cat "$MOCK_DIR/count" 2>/dev/null || echo 0) + 1 ))
echo "$n" > "$MOCK_DIR/count"
while [ "$n" -gt 0 ]; do
if [ -f "$MOCK_DIR/response_$n.json" ]; then
cat "$MOCK_DIR/response_$n.json"
exit 0
fi
n=$(( n - 1 ))
done
echo "api outage" >&2
exit 1
;;
"repos/o/r/pulls/42")
if [ -f "$MOCK_DIR/pr.json" ]; then
cat "$MOCK_DIR/pr.json"
else
echo '{"labels": [{"name": "ready"}]}'
fi
;;
*)
echo "unexpected gh endpoint: $2" >> "$MOCK_DIR/endpoint_error"
exit 2
;;
esac
EOF
chmod +x "$tmp/gh"
PC_OK='{"name": "pre-commit", "id": 1, "status": "completed", "conclusion": "success"}'
PC_BAD='{"name": "pre-commit", "id": 1, "status": "completed", "conclusion": "failure"}'
PC_PENDING='{"name": "pre-commit", "id": 1, "status": "in_progress", "conclusion": null}'
DOCS_OK='{"name": "Deploy Documentation", "id": 2, "status": "completed", "conclusion": "success"}'
DOCS_BAD='{"name": "Deploy Documentation", "id": 2, "status": "completed", "conclusion": "failure"}'
DOCS_CANCELLED='{"name": "Deploy Documentation", "id": 2, "status": "completed", "conclusion": "cancelled"}'
OTHER='{"name": "Trigger Full Suite", "id": 3, "status": "in_progress", "conclusion": null}'
NULL_NAME='{"name": null, "id": 4, "status": "completed", "conclusion": "failure"}'
PC_OK_RERUN='{"name": "pre-commit", "id": 5, "status": "completed", "conclusion": "success"}'
fails=0
want_log="" # optional: expect() also greps out.log for this regex, then resets
pr_json="" # optional: served for the pulls (label re-check) endpoint, then resets
raw_body="" # optional: serve responses verbatim instead of wrapping in workflow_runs
expect() { # <name> <expected-exit> <response json>...
local name=$1 want=$2 dir i=1
shift 2
dir=$(mktemp -d "$tmp/test_XXXXXX")
for body in "$@"; do
if [ -n "$raw_body" ]; then
printf '%s' "$body" > "$dir/response_$i.json"
else
printf '{"workflow_runs": [%s]}' "$body" > "$dir/response_$i.json"
fi
i=$(( i + 1 ))
done
[ -n "$pr_json" ] && printf '%s' "$pr_json" > "$dir/pr.json"
( export PATH="$tmp:$PATH" MOCK_DIR="$dir" PR_SHA=deadbeef PR_NUMBER=42 \
GITHUB_REPOSITORY=o/r POLL_SECS=0 GRACE_SECS=1 MAX_WAIT_SECS=3
bash "$here/gate_full_suite.sh" > "$dir/out.log" 2>&1 )
local rc=$?
if [ "$rc" -ne "$want" ]; then
echo "FAIL: $name (exit $rc, want $want)"
cat "$dir/out.log"
fails=1
elif [ -f "$dir/endpoint_error" ]; then
echo "FAIL: $name (mock gh got an unexpected call)"
cat "$dir/endpoint_error"
fails=1
elif [ -n "$want_log" ] && ! grep -Eq "$want_log" "$dir/out.log"; then
echo "FAIL: $name (log does not match: $want_log)"
cat "$dir/out.log"
fails=1
else
echo "ok: $name"
fi
want_log="" pr_json="" raw_body=""
}
expect "both green -> proceed" 0 "$PC_OK, $DOCS_OK, $OTHER, $NULL_NAME"
expect "docs build failed -> blocked" 1 "$PC_OK, $DOCS_BAD"
expect "pre-commit failed -> blocked" 1 "$PC_BAD"
expect "pending then green -> proceed" 0 "$PC_PENDING" "$PC_OK, $DOCS_OK"
want_log="never appeared.*Deploy Documentation"
expect "docs run absent (path-filtered) -> proceed after grace" 0 "$PC_OK"
expect "API outage -> fail open" 0
want_log="FAILING OPEN"
expect "pending past MAX_WAIT -> fail open" 0 "$PC_PENDING"
want_log="FAILING OPEN"
expect "unrelated runs only -> no grace, fail open at MAX_WAIT" 0 "$OTHER"
expect "cancelled docs then green -> proceed" 0 \
"$PC_OK, $DOCS_CANCELLED" "$PC_OK, $DOCS_OK"
want_log="FAILING OPEN"
expect "cancelled docs forever -> fail open at MAX_WAIT" 0 "$PC_OK, $DOCS_CANCELLED"
want_log="FAILING OPEN"
expect "pre-commit absent -> no grace, fail open at MAX_WAIT" 0 "$DOCS_OK"
expect "duplicate run names -> latest wins" 0 "$PC_BAD, $PC_OK_RERUN, $DOCS_OK"
raw_body=1
expect "garbage response body -> fail open" 0 "this is not json"
pr_json='{"labels": [{"name": "other"}]}'
expect "ready label removed mid-gate -> blocked" 1 "$PC_OK, $DOCS_OK"
exit "$fails"
+3
View File
@@ -47,3 +47,6 @@ jobs:
- uses: pre-commit/action@v3.0.1
with:
extra_args: --all-files --hook-stage manual
# After pre-commit so a self-test failure cannot mask lint failures.
- name: Full-suite gate self-test
run: bash .github/scripts/test_gate_full_suite.sh
@@ -7,6 +7,7 @@ on:
permissions:
contents: read
pull-requests: read
actions: read
concurrency:
group: full-suite-${{ github.event.pull_request.number }}
@@ -18,6 +19,8 @@ jobs:
(github.event.action == 'labeled' && github.event.label.name == 'ready')
|| github.event.action == 'synchronize'
runs-on: ubuntu-latest
# Gate below may wait for cheap checks (up to MAX_WAIT_SECS = 25 min).
timeout-minutes: 35
steps:
- name: Check ready label
id: check
@@ -49,6 +52,20 @@ jobs:
"https://api.buildkite.com/v2/organizations/${{ vars.BUILDKITE_ORG_SLUG }}/pipelines/${{ vars.BUILDKITE_PIPELINE_SLUG }}/builds/${build_num}/cancel"
done
# Checks out the BASE branch (default for pull_request_target), so PR
# authors cannot tamper with the gate script.
- name: Checkout gate script
if: steps.check.outputs.has_ready == 'true'
uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
- name: Wait for pre-commit and docs build
if: steps.check.outputs.has_ready == 'true'
env:
GH_TOKEN: ${{ github.token }}
PR_SHA: ${{ github.event.pull_request.head.sha }}
PR_NUMBER: ${{ github.event.pull_request.number }}
run: bash .github/scripts/gate_full_suite.sh
- name: Trigger Buildkite Full Suite
if: steps.check.outputs.has_ready == 'true'
env:
+7
View File
@@ -99,6 +99,13 @@ status.
Full Suite is also path-filtered. It validates broader behavior before Mergify
can merge a PR.
A `ready`-labeled PR does not hit Buildkite immediately:
`ci-trigger-full-suite.yml` first runs `.github/scripts/gate_full_suite.sh`,
which waits for the cheap Tier-1 checks (pre-commit, docs build) on the PR
head. A red cheap check blocks the suite (fail closed; the next push re-arms
it), while a GitHub outage or a >25 min wait lets it run anyway (fail open).
`/test full` bypasses the gate.
| Buildkite label | `TEST_TYPE` | Main watched paths |
|---|---|---|
| SSIM Tests | `ssim` | `fastvideo/**/*.py`, `pyproject.toml`, `docker/Dockerfile` |
@@ -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
@@ -128,6 +133,7 @@ surfaces:
preprocess_text_funcs: "Internal text preprocessing hooks."
postprocess_text_funcs: "Internal text postprocessing hooks."
scheduler_step_in_fp32: "Runtime scheduler precision toggle; not part of the public typed inference API."
supports_dynamic_batching: "Per-pipeline-family capability flag for dynamic batching opt-in; not part of the public typed inference API."
pipeline_config_extensions:
preset_owned:
+15
View File
@@ -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
+11
View File
@@ -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",
+15 -118
View File
@@ -1,13 +1,15 @@
# SPDX-License-Identifier: Apache-2.0
import importlib.util
import os
import torch
import torch.nn.functional as F
from dataclasses import dataclass
from fastvideo import envs
from fastvideo.attention.utils.flash_attn_default import (
fa_version,
flash_attn_func_compilable,
)
from fastvideo.attention.backends.abstract import (
AttentionBackend,
AttentionImpl,
@@ -17,119 +19,6 @@ from fastvideo.attention.backends.abstract import (
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"
else:
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"
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
# 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.fastvideo._flash_attn_default_forward(q, k, v, softmax_scale, causal)
elif fa_version == "4":
# 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:
# 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.")
logger.info("Using FlashAttention-%s backend", fa_version)
# FP4 FA4 support: quantize Q/K to NVFP4 E2M1 for block-scaled MMA on Blackwell.
@@ -309,9 +198,17 @@ class FlashAttentionImpl(AttentionImpl):
attn_metadata: FlashAttnMetadata,
):
if (attn_metadata is not None and hasattr(attn_metadata, "attn_mask") and attn_metadata.attn_mask is not None):
# Route through the *_compilable wrappers so dynamo sees one
# traceable node for each masked entry point (the unpad/pad
# bookkeeping runs eager inside the custom op). On FA2 these
# wrappers go through ops with full register_autograd, so
# training also backprops through the op (no graph break on
# the training path); on FA3/FA4 they carve out to the
# autograd.Function for grad-enabled calls — see
# fastvideo/attention/utils/flash_attn_no_pad.py.
from fastvideo.attention.utils.flash_attn_no_pad import (
flash_attn_no_pad,
flash_attn_varlen_qk_no_pad,
flash_attn_no_pad_compilable as flash_attn_no_pad,
flash_attn_varlen_qk_no_pad_compilable as flash_attn_varlen_qk_no_pad,
)
attn_mask = attn_metadata.attn_mask
@@ -0,0 +1,251 @@
# SPDX-License-Identifier: Apache-2.0
"""torch.compile-traceable wrapper for the FA2/FA3/FA4 default attention path.
The FA4/cute path (`fa_version == "4"`) is already a registered
`torch.library.custom_op` in `fastvideo.attention.utils.flash_attn_cute`, 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-gated).
Autograd: FA2 has full ``register_autograd`` parity — the custom op's
backward calls flash_attn's ``_flash_attn_backward`` directly, so training
backprops *through* the op (no graph break on the training path either).
FA3 currently keeps the no-backward + carve-out pattern from PR #1373
because FA3's private backward signature wants validation on a real Hopper
box (gated on Kuan-Hao's Modal FA3 setup PR). Once that lands the FA3 path
can mirror FA2.
Lives in `attention/utils/` (sibling of `flash_attn_cute.py` and
`flash_attn_no_pad.py`) so it can be imported by any backend that wants the
traceable FA default call without pulling in backend dispatch logic. The
backend (`attention/backends/flash_attn.py`) just imports
`flash_attn_func_compilable` and `fa_version` from here.
"""
import importlib.util
import torch
from fastvideo import envs
from fastvideo.logger import init_logger
logger = init_logger(__name__)
# Pick the same backend the rest of FastVideo picked for `flash_attn_func`
# (FA4/cute → FA3 → FA2). Mirror the precedence used in
# `attention/utils/flash_attn_no_pad.py` so the two probes always agree.
#
# FA4 (flash_attn.cute) is explicit opt-in via FASTVIDEO_FA4=1: 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.
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"
else:
try:
from flash_attn_interface import flash_attn_func as flash_attn_3_func
# flash_attn 3 no longer has 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"
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
if fa_version == "2":
# 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 in
# `attention/backends/flash_attn.py`). The masked/no-pad and varlen /
# cross-attn paths use different entry points
# (`flash_attn_no_pad`, `flash_attn_varlen_*`) which live in
# `attention/utils/flash_attn_no_pad.py`. 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.
from flash_attn.flash_attn_interface import _flash_attn_backward as _fa2_backward
_fa_default = flash_attn_func
@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,
) -> tuple[torch.Tensor, torch.Tensor]:
# `return_attn_probs=True` asks FA2 to also return softmax_lse +
# S_dmask. We need softmax_lse to feed the backward; S_dmask is the
# dropout mask (always None here since dropout_p is fixed at 0).
out, softmax_lse, _ = _fa_default(q, k, v, softmax_scale=softmax_scale, causal=causal, return_attn_probs=True)
return out, softmax_lse
@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,
) -> tuple[torch.Tensor, torch.Tensor]:
del softmax_scale, causal
# FA2 default path: out = [batch, seqlen_q, nheads, head_dim_v],
# softmax_lse = [batch, nheads, seqlen_q], fp32 regardless of q dtype.
b, sq, hq = q.shape[0], q.shape[1], q.shape[2]
out = q.new_empty(b, sq, hq, v.shape[-1])
lse = q.new_empty(b, hq, sq, dtype=torch.float32)
return out, lse
def _flash_attn_default_setup_context(ctx, inputs, output):
q, k, v, softmax_scale, causal = inputs
out, lse = output
ctx.save_for_backward(q, k, v, out, lse)
# `lse` is an auxiliary output we save to feed FA2's backward; nobody
# should differentiate through it. Mark it non-differentiable so
# autograd errors loudly if a caller wires it into a loss, rather
# than silently producing zero/None grads through the `del grad_lse`
# in our backward.
ctx.mark_non_differentiable(lse)
# FA2's *forward* substitutes `1 / sqrt(head_dim)` for `softmax_scale=None`
# internally; FA2's *backward* (`_flash_attn_backward`) demands a concrete
# float in its C++ schema and rejects None at the binding boundary. Resolve
# the default here so the value saved on ctx (and passed to backward) is
# always a real float — matches what FA2's own autograd.Function does.
if softmax_scale is None:
softmax_scale = q.shape[-1]**-0.5
ctx.softmax_scale = softmax_scale
ctx.causal = causal
def _flash_attn_default_backward(ctx, grad_out, grad_lse):
# We only differentiate `out`; softmax_lse is saved-for-backward, not
# a real differentiable output. (Mirrors the FP4 cute template.)
del grad_lse
q, k, v, out, lse = ctx.saved_tensors
dq = torch.empty_like(q)
dk = torch.empty_like(k)
dv = torch.empty_like(v)
# FA2's `_flash_attn_backward` writes into dq/dk/dv in place. The
# extra kwargs (window_size_*, softcap, alibi_slopes, deterministic,
# rng_state) are pinned to the same defaults the forward wrapper
# uses — flash-attn==2.8.1 (the version FastVideo pins) requires
# all of them explicitly. `rng_state=None` is correct for our
# `dropout_p=0` configuration.
_fa2_backward(
grad_out,
q,
k,
v,
out,
lse,
dq,
dk,
dv,
dropout_p=0.0,
softmax_scale=ctx.softmax_scale,
causal=ctx.causal,
window_size_left=-1,
window_size_right=-1,
softcap=0.0,
alibi_slopes=None,
deterministic=False,
rng_state=None,
)
return dq, dk, dv, None, None
torch.library.register_autograd(
"fastvideo::_flash_attn_default_forward",
_flash_attn_default_backward,
setup_context=_flash_attn_default_setup_context,
)
def flash_attn_func_compilable(q, k, v, softmax_scale=None, causal=False):
# Backward is registered: autograd flows through the op (training
# path is also traceable; no carve-out needed). Public API matches
# `flash_attn_func` — returns just `out`; we drop the saved-for-
# backward `lse` here so callers see the original single-tensor
# contract.
out, _ = torch.ops.fastvideo._flash_attn_default_forward(q, k, v, softmax_scale, causal)
return out
elif fa_version == "3":
# FA3 path: same forward+fake custom op as the original PR #1373, with
# the autograd carve-out kept. The full backward (mirroring the FA2 leg
# above) wants a Hopper box for grad-check validation, which we don't
# have until Kuan-Hao's Modal FA3 setup PR lands. Until then this keeps
# inference traceable + training correct (via the original
# autograd.Function path + a pre-PR-style graph break on training).
_fa_default = flash_attn_func
@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
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.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.
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.")
+483 -5
View File
@@ -24,7 +24,7 @@ from flash_attn.bert_padding import pad_input, unpad_input
from fastvideo import envs
def _resolve_flash_attn_varlen_func() -> Any:
def _resolve_flash_attn_varlen_func() -> tuple[Any, str]:
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
@@ -39,20 +39,28 @@ def _resolve_flash_attn_varlen_func() -> Any:
"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
return flash_attn_varlen_func_cute, "4"
try:
from flash_attn_interface import (
flash_attn_varlen_func as flash_attn_varlen_func_interface, )
return flash_attn_varlen_func_interface
return flash_attn_varlen_func_interface, "3"
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, "2"
flash_attn_varlen_func_impl = _resolve_flash_attn_varlen_func()
flash_attn_varlen_func_impl, _FA_VARLEN_VERSION = _resolve_flash_attn_varlen_func()
# FA2-only: the private varlen backward we register against the custom ops
# below. FA3 / FA4 have different private signatures and validation paths
# (Hopper / Blackwell boxes) — those legs keep the autograd carve-out
# pattern from PR #1373 until their setup PRs land.
if _FA_VARLEN_VERSION == "2":
from flash_attn.flash_attn_interface import (
_flash_attn_varlen_backward as _fa2_varlen_backward, )
def flash_attn_no_pad(
@@ -191,3 +199,473 @@ def flash_attn_varlen_qk_no_pad(
h=nheads,
)
return output
# ---------------------------------------------------------------------------
# torch.compile traceability + register_autograd parity for the masked /
# varlen attention paths.
#
# Wraps the two entry points `FlashAttentionImpl.forward` calls
# (`flash_attn_no_pad`, `flash_attn_varlen_qk_no_pad`) as
# `torch.library.custom_op`s so dynamo sees one traceable node — the
# internal unpad / pad bookkeeping (data-dependent `nnz` shapes) runs
# eager inside the op, and the op's outputs are the statically-shaped
# padded tensors. This mirrors the FA2 default-path wrapper in
# `fastvideo/attention/backends/flash_attn.py`.
#
# Autograd: on FA2 we register a real backward (`register_autograd`)
# that calls FA2's `_flash_attn_varlen_backward` on the unpadded form
# — re-unpadding the saved padded tensors using the saved mask. The
# `softmax_lse` from the varlen forward is naturally unpadded
# (`[nheads, total_q]`); we pad it to `[batch, nheads, seqlen]` on
# the way out (statically shaped) and re-unpad in backward. So
# training backprops *through* the op (no graph break on the training
# path either).
#
# FA3 / FA4 keep the autograd carve-out pattern from PR #1373: the
# custom op has forward + fake only, and `*_compilable` falls back to
# the original autograd.Function for grad-enabled calls. Those legs
# are gated on Hopper-class / Blackwell-class boxes for backward
# validation and ship as separate follow-ups.
if _FA_VARLEN_VERSION == "2":
# ---------- masked self-attention: flash_attn_no_pad (FA2) ----------
@torch.library.custom_op(
"fastvideo::_flash_attn_no_pad_forward",
mutates_args=(),
device_types="cuda",
)
def _flash_attn_no_pad_forward(
qkv: torch.Tensor,
key_padding_mask: torch.Tensor,
causal: bool,
dropout_p: float,
softmax_scale: float | None,
deterministic: bool,
) -> tuple[torch.Tensor, torch.Tensor]:
b, s, _three, h, d = qkv.shape
x = rearrange(qkv, "b s three h d -> b s (three h d)")
x_unpad, indices, cu_seqlens, max_s, _ = unpad_input(x, key_padding_mask)
x_unpad = rearrange(x_unpad, "nnz (three h d) -> nnz three h d", three=3, h=h)
out_unpad, lse_unpad, _ = flash_attn_varlen_qkvpacked_func(x_unpad,
cu_seqlens,
max_s,
dropout_p,
softmax_scale=softmax_scale,
causal=causal,
deterministic=deterministic,
return_attn_probs=True)
# Pad out: [nnz, h, d] -> [b, s, h, d]
out_padded = rearrange(pad_input(rearrange(out_unpad, "nnz h d -> nnz (h d)"), indices, b, s),
"b s (h d) -> b s h d",
h=h)
# Pad lse: FA2 varlen returns [nheads, total_q]. Transpose to [total_q,
# nheads], pad to [b, s, nheads], permute to [b, nheads, s] — statically
# shaped so register_fake matches.
lse_padded = pad_input(lse_unpad.t().contiguous(), indices, b, s).permute(0, 2, 1).contiguous()
return out_padded, lse_padded
@torch.library.register_fake("fastvideo::_flash_attn_no_pad_forward")
def _flash_attn_no_pad_forward_fake(
qkv: torch.Tensor,
key_padding_mask: torch.Tensor,
causal: bool,
dropout_p: float,
softmax_scale: float | None,
deterministic: bool,
) -> tuple[torch.Tensor, torch.Tensor]:
del key_padding_mask, causal, dropout_p, softmax_scale, deterministic
b, s, _three, h, d = qkv.shape
out = qkv.new_empty(b, s, h, d)
lse = qkv.new_empty(b, h, s, dtype=torch.float32)
return out, lse
def _flash_attn_no_pad_setup_context(ctx, inputs, output):
qkv, key_padding_mask, causal, dropout_p, softmax_scale, deterministic = inputs
out, lse = output
ctx.save_for_backward(qkv, out, lse, key_padding_mask)
# Auxiliary output, not differentiable — see default-path note.
ctx.mark_non_differentiable(lse)
# FA2's varlen backward requires a concrete float for softmax_scale.
if softmax_scale is None:
softmax_scale = qkv.shape[-1]**-0.5 # head_dim from qkv's last dim
ctx.softmax_scale = softmax_scale
ctx.causal = causal
ctx.dropout_p = dropout_p
ctx.deterministic = deterministic
def _flash_attn_no_pad_backward(ctx, grad_out, grad_lse):
# lse is saved-for-backward, not differentiated.
del grad_lse
qkv, out_padded, lse_padded, key_padding_mask = ctx.saved_tensors
b, s, _three, h, d = qkv.shape
# One `unpad_input` call (on qkv) gives us indices + cu_seqlens + max_s;
# reuse those for out / dout / lse below via direct indexing instead
# of redundant `unpad_input` calls (each of which would re-run
# `nonzero` + `cumsum` + a `.max().item()` GPU→CPU sync).
x = rearrange(qkv, "b s three h d -> b s (three h d)")
x_unpad, indices, cu_seqlens, max_s, _ = unpad_input(x, key_padding_mask)
x_unpad = rearrange(x_unpad, "nnz (three h d) -> nnz three h d", three=3, h=h)
q_unpad, k_unpad, v_unpad = (t.contiguous() for t in x_unpad.unbind(dim=1))
# Direct-index variants reuse `indices` (computed above).
out_unpad = out_padded.flatten(0, 1)[indices].view(-1, h, d).contiguous()
dout_unpad = grad_out.flatten(0, 1)[indices].view(-1, h, d).contiguous()
# lse_padded [b, h, s] -> [b, s, h] -> [nnz, h] -> [h, nnz].
lse_unpad = lse_padded.permute(0, 2, 1).contiguous().flatten(0, 1)[indices].t().contiguous()
dq_unpad = torch.empty_like(q_unpad)
dk_unpad = torch.empty_like(k_unpad)
dv_unpad = torch.empty_like(v_unpad)
_fa2_varlen_backward(
dout_unpad,
q_unpad,
k_unpad,
v_unpad,
out_unpad,
lse_unpad,
dq_unpad,
dk_unpad,
dv_unpad,
cu_seqlens_q=cu_seqlens,
cu_seqlens_k=cu_seqlens,
max_seqlen_q=max_s,
max_seqlen_k=max_s,
dropout_p=ctx.dropout_p,
softmax_scale=ctx.softmax_scale,
causal=ctx.causal,
window_size_left=-1,
window_size_right=-1,
softcap=0.0,
alibi_slopes=None,
deterministic=ctx.deterministic,
rng_state=None,
)
# Re-pad each grad and stack into dqkv.
def _repad(dt_unpad: torch.Tensor) -> torch.Tensor:
padded = pad_input(rearrange(dt_unpad, "nnz h d -> nnz (h d)"), indices, b, s)
return rearrange(padded, "b s (h d) -> b s h d", h=h)
dqkv = torch.stack([_repad(dq_unpad), _repad(dk_unpad), _repad(dv_unpad)], dim=2)
# 6 inputs total: qkv, key_padding_mask, causal, dropout_p, softmax_scale, deterministic.
return dqkv, None, None, None, None, None
torch.library.register_autograd(
"fastvideo::_flash_attn_no_pad_forward",
_flash_attn_no_pad_backward,
setup_context=_flash_attn_no_pad_setup_context,
)
# ---------- cross-attention: flash_attn_varlen_qk_no_pad (FA2) ----------
@torch.library.custom_op(
"fastvideo::_flash_attn_varlen_qk_no_pad_forward",
mutates_args=(),
device_types="cuda",
)
def _flash_attn_varlen_qk_no_pad_forward(
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
query_padding_mask: torch.Tensor,
key_padding_mask: torch.Tensor,
causal: bool,
dropout_p: float,
softmax_scale: float | None,
deterministic: bool,
) -> tuple[torch.Tensor, torch.Tensor]:
b, sq, h, d = query.shape
q_unpad, q_indices, cu_seqlens_q, max_seqlen_q, _ = unpad_input(rearrange(query, "b s h d -> b s (h d)"),
query_padding_mask)
k_unpad, _, cu_seqlens_k, max_seqlen_k, _ = unpad_input(rearrange(key, "b s h d -> b s (h d)"),
key_padding_mask)
v_unpad, _, _, _, _ = unpad_input(rearrange(value, "b s h d -> b s (h d)"), key_padding_mask)
q_unpad = rearrange(q_unpad, "nnz (h d) -> nnz h d", h=h)
k_unpad = rearrange(k_unpad, "nnz (h d) -> nnz h d", h=h)
v_unpad = rearrange(v_unpad, "nnz (h d) -> nnz h d", h=h)
out_unpad, lse_unpad, _ = flash_attn_varlen_func_impl(q_unpad,
k_unpad,
v_unpad,
cu_seqlens_q,
cu_seqlens_k,
max_seqlen_q,
max_seqlen_k,
dropout_p=dropout_p,
softmax_scale=softmax_scale,
causal=causal,
deterministic=deterministic,
return_attn_probs=True)
# Pad out: [nnz_q, h, d] -> [b, sq, h, d]
out_padded = rearrange(pad_input(rearrange(out_unpad, "nnz h d -> nnz (h d)"), q_indices, b, sq),
"b s (h d) -> b s h d",
h=h)
# Pad lse: [h, nnz_q] -> [b, h, sq]
lse_padded = pad_input(lse_unpad.t().contiguous(), q_indices, b, sq).permute(0, 2, 1).contiguous()
return out_padded, lse_padded
@torch.library.register_fake("fastvideo::_flash_attn_varlen_qk_no_pad_forward")
def _flash_attn_varlen_qk_no_pad_forward_fake(
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
query_padding_mask: torch.Tensor,
key_padding_mask: torch.Tensor,
causal: bool,
dropout_p: float,
softmax_scale: float | None,
deterministic: bool,
) -> tuple[torch.Tensor, torch.Tensor]:
del key, query_padding_mask, key_padding_mask
del causal, dropout_p, softmax_scale, deterministic
b, sq, h, _ = query.shape
# `out`'s head_dim comes from value (d_v), matching the real forward's
# out_padded ([b, sq, h, d_v]); it can differ from query's d_q.
out = query.new_empty(b, sq, h, value.shape[-1])
lse = query.new_empty(b, h, sq, dtype=torch.float32)
return out, lse
def _flash_attn_varlen_qk_no_pad_setup_context(ctx, inputs, output):
(query, key, value, query_padding_mask, key_padding_mask, causal, dropout_p, softmax_scale,
deterministic) = inputs
out, lse = output
ctx.save_for_backward(query, key, value, out, lse, query_padding_mask, key_padding_mask)
# Auxiliary output, not differentiable — see default-path note.
ctx.mark_non_differentiable(lse)
if softmax_scale is None:
softmax_scale = query.shape[-1]**-0.5
ctx.softmax_scale = softmax_scale
ctx.causal = causal
ctx.dropout_p = dropout_p
ctx.deterministic = deterministic
def _flash_attn_varlen_qk_no_pad_backward(ctx, grad_out, grad_lse):
del grad_lse
(query, key, value, out_padded, lse_padded, query_padding_mask, key_padding_mask) = ctx.saved_tensors
b, sq, h, d = query.shape
sk = key.shape[1]
# One `unpad_input` call per distinct mask; reuse the returned
# indices via direct indexing for everything else that shares
# the same mask (v with k_mask; out/dout/lse with q_mask; the
# final repad of dk/dv also reuses k_indices). Avoids ~4
# redundant `unpad_input` calls + their GPU→CPU `.max().item()`
# syncs.
q_unpad, q_indices, cu_seqlens_q, max_seqlen_q, _ = unpad_input(rearrange(query, "b s h d -> b s (h d)"),
query_padding_mask)
k_unpad, k_indices, cu_seqlens_k, max_seqlen_k, _ = unpad_input(rearrange(key, "b s h d -> b s (h d)"),
key_padding_mask)
q_unpad = rearrange(q_unpad, "nnz (h d) -> nnz h d", h=h).contiguous()
k_unpad = rearrange(k_unpad, "nnz (h d) -> nnz h d", h=h).contiguous()
v_unpad = value.flatten(0, 1)[k_indices].view(-1, h, d).contiguous()
# out / dout / lse follow q's shape, so index with q_indices.
out_unpad = out_padded.flatten(0, 1)[q_indices].view(-1, h, d).contiguous()
dout_unpad = grad_out.flatten(0, 1)[q_indices].view(-1, h, d).contiguous()
lse_unpad = lse_padded.permute(0, 2, 1).contiguous().flatten(0, 1)[q_indices].t().contiguous()
dq_unpad = torch.empty_like(q_unpad)
dk_unpad = torch.empty_like(k_unpad)
dv_unpad = torch.empty_like(v_unpad)
_fa2_varlen_backward(
dout_unpad,
q_unpad,
k_unpad,
v_unpad,
out_unpad,
lse_unpad,
dq_unpad,
dk_unpad,
dv_unpad,
cu_seqlens_q=cu_seqlens_q,
cu_seqlens_k=cu_seqlens_k,
max_seqlen_q=max_seqlen_q,
max_seqlen_k=max_seqlen_k,
dropout_p=ctx.dropout_p,
softmax_scale=ctx.softmax_scale,
causal=ctx.causal,
window_size_left=-1,
window_size_right=-1,
softcap=0.0,
alibi_slopes=None,
deterministic=ctx.deterministic,
rng_state=None,
)
# k_indices is already available from the unpad_input above —
# no need to recompute it for the dk/dv repad.
def _repad(dt_unpad: torch.Tensor, indices: torch.Tensor, batch: int, seqlen: int) -> torch.Tensor:
padded = pad_input(rearrange(dt_unpad, "nnz h d -> nnz (h d)"), indices, batch, seqlen)
return rearrange(padded, "b s (h d) -> b s h d", h=h)
dq_padded = _repad(dq_unpad, q_indices, b, sq)
dk_padded = _repad(dk_unpad, k_indices, b, sk)
dv_padded = _repad(dv_unpad, k_indices, b, sk)
# 9 inputs total: query, key, value, q_mask, k_mask, causal, dropout_p,
# softmax_scale, deterministic.
return dq_padded, dk_padded, dv_padded, None, None, None, None, None, None
torch.library.register_autograd(
"fastvideo::_flash_attn_varlen_qk_no_pad_forward",
_flash_attn_varlen_qk_no_pad_backward,
setup_context=_flash_attn_varlen_qk_no_pad_setup_context,
)
# ---------- public dispatchers (FA2: autograd flows through the op) -----
def flash_attn_no_pad_compilable(qkv,
key_padding_mask,
causal=False,
dropout_p=0.0,
softmax_scale=None,
deterministic=False):
"""dynamo-traceable wrapper around ``flash_attn_no_pad`` (registered op,
full register_autograd on FA2 — both inference and training go through
the op, no graph break on either)."""
out, _ = torch.ops.fastvideo._flash_attn_no_pad_forward(qkv, key_padding_mask, causal, dropout_p, softmax_scale,
deterministic)
return out
def flash_attn_varlen_qk_no_pad_compilable(query,
key,
value,
query_padding_mask,
key_padding_mask,
causal=False,
dropout_p=0.0,
softmax_scale=None,
deterministic=False):
"""dynamo-traceable wrapper around ``flash_attn_varlen_qk_no_pad`` (registered
op, full register_autograd on FA2)."""
out, _ = torch.ops.fastvideo._flash_attn_varlen_qk_no_pad_forward(query, key, value, query_padding_mask,
key_padding_mask, causal, dropout_p,
softmax_scale, deterministic)
return out
else:
# ---------- FA3 / FA4: carve-out (forward+fake only, no real backward) ---
# Same pattern as the parked varlen-extension and the FA3 default leg in
# `fastvideo/attention/backends/flash_attn.py`. Real backward for these
# versions is a follow-up gated on Hopper / Blackwell box validation.
@torch.library.custom_op(
"fastvideo::_flash_attn_no_pad_forward",
mutates_args=(),
device_types="cuda",
)
def _flash_attn_no_pad_forward(
qkv: torch.Tensor,
key_padding_mask: torch.Tensor,
causal: bool,
dropout_p: float,
softmax_scale: float | None,
deterministic: bool,
) -> torch.Tensor:
return flash_attn_no_pad( # type: ignore[no-untyped-call]
qkv,
key_padding_mask,
causal=causal,
dropout_p=dropout_p,
softmax_scale=softmax_scale,
deterministic=deterministic)
@torch.library.register_fake("fastvideo::_flash_attn_no_pad_forward")
def _flash_attn_no_pad_forward_fake(
qkv: torch.Tensor,
key_padding_mask: torch.Tensor,
causal: bool,
dropout_p: float,
softmax_scale: float | None,
deterministic: bool,
) -> torch.Tensor:
del key_padding_mask, causal, dropout_p, softmax_scale, deterministic
b, s, _three, h, d = qkv.shape
return qkv.new_empty(b, s, h, d)
@torch.library.custom_op(
"fastvideo::_flash_attn_varlen_qk_no_pad_forward",
mutates_args=(),
device_types="cuda",
)
def _flash_attn_varlen_qk_no_pad_forward(
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
query_padding_mask: torch.Tensor,
key_padding_mask: torch.Tensor,
causal: bool,
dropout_p: float,
softmax_scale: float | None,
deterministic: bool,
) -> torch.Tensor:
return flash_attn_varlen_qk_no_pad( # type: ignore[no-untyped-call]
query,
key,
value,
query_padding_mask=query_padding_mask,
key_padding_mask=key_padding_mask,
causal=causal,
dropout_p=dropout_p,
softmax_scale=softmax_scale,
deterministic=deterministic)
@torch.library.register_fake("fastvideo::_flash_attn_varlen_qk_no_pad_forward")
def _flash_attn_varlen_qk_no_pad_forward_fake(
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
query_padding_mask: torch.Tensor,
key_padding_mask: torch.Tensor,
causal: bool,
dropout_p: float,
softmax_scale: float | None,
deterministic: bool,
) -> torch.Tensor:
del key, query_padding_mask, key_padding_mask
del causal, dropout_p, softmax_scale, deterministic
b, sq, h, _ = query.shape
# `out`'s head_dim comes from value (d_v), matching the real forward's
# output ([b, sq, h, d_v]); it can differ from query's d_q.
return query.new_empty(b, sq, h, value.shape[-1])
def flash_attn_no_pad_compilable(qkv,
key_padding_mask,
causal=False,
dropout_p=0.0,
softmax_scale=None,
deterministic=False):
if torch.is_grad_enabled() and qkv.requires_grad:
return flash_attn_no_pad(qkv,
key_padding_mask,
causal=causal,
dropout_p=dropout_p,
softmax_scale=softmax_scale,
deterministic=deterministic)
return torch.ops.fastvideo._flash_attn_no_pad_forward(qkv, key_padding_mask, causal, dropout_p, softmax_scale,
deterministic)
def flash_attn_varlen_qk_no_pad_compilable(query,
key,
value,
query_padding_mask,
key_padding_mask,
causal=False,
dropout_p=0.0,
softmax_scale=None,
deterministic=False):
if torch.is_grad_enabled() and (query.requires_grad or key.requires_grad or value.requires_grad):
return flash_attn_varlen_qk_no_pad(query,
key,
value,
query_padding_mask=query_padding_mask,
key_padding_mask=key_padding_mask,
causal=causal,
dropout_p=dropout_p,
softmax_scale=softmax_scale,
deterministic=deterministic)
return torch.ops.fastvideo._flash_attn_varlen_qk_no_pad_forward(query, key, value, query_padding_mask,
key_padding_mask, causal, dropout_p,
softmax_scale, deterministic)
+26
View File
@@ -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",
]
+297
View File
@@ -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}")
+183
View File
@@ -0,0 +1,183 @@
# 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",
"latents",
"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)
# Unrecognized types (tensors, PIL images, ...) cannot be reliably compared
# by value — repr() truncates large tensors, so two different payloads can
# collide. Treat them as equal only when both requests hold the exact same
# object; anything else is conservatively incompatible.
return (type(value).__qualname__, id(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
+37
View File
@@ -75,6 +75,12 @@ class PipelineConfig:
# DMD parameters
dmd_denoising_steps: list[int] | None = field(default=None)
# Whether this pipeline family is verified to merge multiple requests into
# one forward pass (dynamic batching). Conservative default: several
# pipelines are structurally batch-size-1, so each family must opt in
# explicitly after verifying output parity.
supports_dynamic_batching: bool = False
# Wan2.2 task modifiers
ti2v_task: bool = False
lucy_edit_task: bool = False
@@ -276,6 +282,37 @@ 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 dynamic_batching_supported(self) -> bool:
"""Whether the loaded pipeline may merge requests via dynamic batching.
Conservative gate: a pipeline family must opt in explicitly, and DMD /
causal denoising variants are always excluded because those denoising
stages do not consume per-item generators, so merged requests would
share RNG.
"""
return (self.supports_dynamic_batching and self.dmd_denoising_steps is None and not self.is_causal)
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 = []
+6
View File
@@ -59,6 +59,12 @@ class WanT2V480PConfig(PipelineConfig):
vae_decode_precision: str = "bf16"
text_encoder_precisions: tuple[str, ...] = field(default_factory=lambda: ("fp32", ))
# The generic Wan denoising path is the only one verified for dynamic
# request batching (parity runs on Wan2.1-T2V). DMD/causal Wan variants
# inherit this flag but are vetoed by
# PipelineConfig.dynamic_batching_supported().
supports_dynamic_batching: bool = True
# self-forcing params
warp_denoising_step: bool = True
+20 -1
View File
@@ -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.")
+227
View File
@@ -0,0 +1,227 @@
# 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
pipeline_config = self._fastvideo_args.pipeline_config
if self.enabled and not pipeline_config.dynamic_batching_supported():
raise RuntimeError(
f"Pipeline config {type(pipeline_config).__name__} does not support dynamic request batching; "
"run with --batching-mode disabled, or use a pipeline that opts in via "
"supports_dynamic_batching=True (DMD/causal denoising variants are always excluded)")
self._task = asyncio.create_task(self._run(), name="fastvideo-video-batch-scheduler")
self._task.add_done_callback(self._on_run_done)
def _on_run_done(self, task: asyncio.Task) -> None:
if task.cancelled():
exc: BaseException | None = asyncio.CancelledError()
else:
exc = task.exception()
if exc is None:
return
logger.error("Video batch scheduler task died; failing all pending requests", exc_info=exc)
self._stopped = True
self._drain_and_fail_waiting(RuntimeError(f"Video batch scheduler crashed and cannot dispatch: {exc!r}"))
def _drain_and_fail_waiting(self, error: BaseException) -> None:
"""Fail every job still waiting in ``_pending`` or ``_queue``."""
while self._pending:
job = self._pending.popleft()
if not job.future.done():
job.future.set_exception(error)
while True:
try:
queued = self._queue.get_nowait()
except asyncio.QueueEmpty:
return
if queued is not None and not queued.future.done():
queued.future.set_exception(error)
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
try:
batch = await self._collect_batch(job)
await self._dispatch(batch)
except BaseException:
# The job was already popped off the queue; park it back in
# _pending so the crash drain in _on_run_done fails its future
# instead of leaving it hanging forever.
self._pending.append(job)
raise
# Fail jobs stranded in _pending AND still sitting in _queue (e.g.
# submitted while the final dispatch was in flight when stop() landed).
self._drain_and_fail_waiting(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:
timeout = deadline - time.perf_counter()
if timeout > 0:
try:
candidate = await asyncio.wait_for(self._get_next_job(), timeout=timeout)
except (TimeoutError, asyncio.TimeoutError):
# asyncio.TimeoutError is only aliased to the builtin on
# Python 3.11+; catch both so 3.10 does not kill _run.
break
else:
# delay=0 means "don't wait", and an already-expired deadline
# means the first job waited behind an in-flight generation:
# in both cases greedily drain whatever is already queued so
# max_size still coalesces instead of dispatching batches of 1.
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 as exc:
# Deduplicated per message so a systematically broken probe (which
# silently disables batching) is visible without log spam.
logger.warning_once("Dynamic batch compatibility probe failed; treating requests as "
f"incompatible (batching disabled for these requests): {exc!r}")
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
+12 -2
View File
@@ -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
+9 -4
View File
@@ -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] = {
+411 -95
View File
@@ -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.
@@ -416,14 +435,7 @@ class VideoGenerator:
legacy_kwargs=kwargs,
)
fastvideo_args = self.fastvideo_args
pipeline_overrides = request_to_pipeline_overrides(request)
if pipeline_overrides:
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 = self._resolve_request_fastvideo_args(request)
resolved_sampling_param = request_to_sampling_param(
request,
@@ -439,6 +451,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,
@@ -446,6 +522,8 @@ class VideoGenerator:
if isinstance(request.prompt, list):
if request.inputs.prompt_path is not None:
raise ValueError("request.prompt list cannot be combined with request.inputs.prompt_path")
if (self._dynamic_batching_enabled(self.fastvideo_args) and self.fastvideo_args.prompt_txt is None):
return self._generate_batched_typed_requests(request)
results: list[GenerationResult] = []
for index, single_request in enumerate(expand_request_prompt_batch(request)):
prompt = single_request.prompt
@@ -461,18 +539,58 @@ class VideoGenerator:
return self._generate_single_request(request)
def _generate_batched_typed_requests(self, request: GenerationRequest) -> list[GenerationResult]:
"""Route a typed prompt-list request through the dynamic-batching work-item path.
Expanded requests share every field except prompt/image_path/video_path
(none of which are pipeline overrides), so one resolved FastVideoArgs is
shared across items — merge compatibility requires identity.
"""
single_requests = expand_request_prompt_batch(request)
fastvideo_args = self._resolve_request_fastvideo_args(request)
work_items: list[_GenerationWorkItem] = []
reserved_output_paths: set[str] = set()
for single_request in single_requests:
sampling_param = request_to_sampling_param(
single_request,
model_path=self.fastvideo_args.model_path,
)
output_path = self._prepare_output_path(sampling_param.output_path, single_request.prompt,
reserved_output_paths)
work_items.append(
self._prepare_generation_work_item(
prompt=single_request.prompt,
sampling_param=sampling_param,
fastvideo_args=fastvideo_args,
output_path=output_path,
))
results: list[GenerationResult] = []
for index, legacy_result in enumerate(self._generate_prepared_work_items(work_items)):
wrapped = GenerationResult.from_legacy_result(legacy_result)
wrapped.prompt_index = index
if wrapped.prompt is None:
wrapped.prompt = single_requests[index].prompt
results.append(wrapped)
return results
def _resolve_request_fastvideo_args(self, request: GenerationRequest) -> FastVideoArgs:
"""Return ``self.fastvideo_args``, deep-copied with the request's pipeline overrides applied (if any)."""
pipeline_overrides = request_to_pipeline_overrides(request)
if not pipeline_overrides:
return self.fastvideo_args
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))
return fastvideo_args
def _generate_single_request(
self,
request: GenerationRequest,
) -> GenerationResult | list[GenerationResult]:
fastvideo_args = self.fastvideo_args
pipeline_overrides = request_to_pipeline_overrides(request)
if pipeline_overrides:
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 = self._resolve_request_fastvideo_args(request)
sampling_param = request_to_sampling_param(
request,
@@ -535,6 +653,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 +729,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 +744,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,53 +791,57 @@ 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 _generate_single_video(
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,
sampling_param: SamplingParam | None = None,
fastvideo_args: FastVideoArgs | None = None,
prompt: str | list[str],
sampling_param: SamplingParam,
fastvideo_args: FastVideoArgs,
**kwargs,
) -> dict[str, Any]:
"""Internal method for single video generation"""
if fastvideo_args is None:
fastvideo_args = self.fastvideo_args
) -> _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)}")
# Validate inputs
if not isinstance(prompt, str):
raise TypeError(f"`prompt` must be a string, but got {type(prompt)}")
prompt = prompt.strip()
sampling_param = deepcopy(sampling_param)
output_path = kwargs["output_path"]
prompt_embeds = kwargs.get("prompt_embeds")
sampling_param.prompt = prompt
# Process negative prompt
sampling_param.prompt = prompt_value
if sampling_param.negative_prompt is not None:
sampling_param.negative_prompt = sampling_param.negative_prompt.strip()
# Validate dimensions
if (sampling_param.height <= 0 or sampling_param.width <= 0 or sampling_param.num_frames <= 0):
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}")
# Calculate sizes
target_height = align_to(sampling_param.height, 16)
target_width = align_to(sampling_param.width, 16)
# Calculate latent sizes
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]
# Log parameters
debug_str = f"""
height: {target_height}
width: {target_width}
@@ -712,7 +861,6 @@ class VideoGenerator:
""" # type: ignore[attr-defined]
logger.info(debug_str)
# Prepare batch
batch = ForwardBatch(
**shallow_asdict(sampling_param),
eta=0.0,
@@ -720,19 +868,30 @@ class VideoGenerator:
VSA_sparsity=fastvideo_args.VSA_sparsity,
)
# Allow precomputed prompt_embeds (e.g. from diffusers) to skip text encoding
prompt_embeds = kwargs.get("prompt_embeds")
if prompt_embeds is not None:
batch.prompt_embeds = (list(prompt_embeds) if isinstance(prompt_embeds, list | tuple) else [prompt_embeds])
extra_overrides = kwargs.pop("_extra_overrides", {})
extra_overrides = kwargs.get("_extra_overrides", {})
for _ek, _ev in extra_overrides.items():
batch.extra[_ek] = _ev
# Run inference
start_time = time.perf_counter()
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,
)
# Execute forward pass in a new thread for non-blocking tensor
# allocation. Capture thread exceptions so we can surface the true
# failure in the main thread instead of later hitting None outputs.
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": ""}
@@ -747,21 +906,6 @@ class VideoGenerator:
thread = threading.Thread(target=execute_forward_thread)
thread.start()
latent_batch_size = _infer_latent_batch_size(batch)
# When ``output_type == "latent"`` the forward output has latent
# shape (e.g. ``[B, C_latent, T_latent, H_latent, W_latent]``)
# rather than the pre-allocation's pixel shape. Skip the pinned
# ~50 MB buffer entirely; we always fall through to the
# ``samples = output_batch.output.cpu()`` branch below in that
# mode. ``skip_pixel_prealloc`` also gates the slow-path warning.
skip_pixel_prealloc = fastvideo_args.output_type == "latent"
if skip_pixel_prealloc:
samples = torch.empty(0, device='cpu')
else:
samples = torch.empty(
(latent_batch_size, 3, sampling_param.num_frames, sampling_param.height, sampling_param.width),
device='cpu',
pin_memory=fastvideo_args.pin_cpu_memory)
thread.join()
if thread_error["error"] is not None:
@@ -773,29 +917,51 @@ class VideoGenerator:
raise RuntimeError("Forward execution returned no output tensor. "
"This usually means the executor/pipeline failed earlier.")
if output_batch.output.shape == samples.shape:
samples.copy_(output_batch.output)
else:
if not skip_pixel_prealloc:
logger.warning("Output shape %s does not match expected shape %s; use slow path",
output_batch.output.shape, samples.shape)
samples = output_batch.output.cpu()
logging_info = output_batch.logging_info
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
# Three mutually-exclusive output modes determine whether (a) we
# build an RGB frame buffer and (b) what file we write to disk:
#
# 1. `output_type == "latent"` — VAE is bypassed in DecodingStage
# and `samples` holds raw latents (arbitrary channel count).
# The RGB grid / uint8 / mp4 / png pipeline below cannot
# consume those, so we skip it entirely and let callers work
# with the latent tensor directly via `result["samples"]`.
# 2. Audio-only workload — `samples` is a 1×3×1×8×8 placeholder
# no caller will use; skip the grid loop and save a `.wav`.
# 3. Pixel video / image — the historical happy path.
is_latent_output = fastvideo_args.output_type == "latent"
audio_only = bool(output_batch.extra.get("audio_only"))
@@ -821,9 +987,6 @@ class VideoGenerator:
audio_mux_time = 0.0
if save_to_disk:
if audio_only:
# Audio-only workload: write a standalone .wav rather than
# muxing the audio into a placeholder mp4 (which forces
# ffmpeg to round 8x8 placeholder frames up to 16x16).
output_path = self._rewrite_extension(output_path, ".wav")
save_start = time.perf_counter()
self._write_pcm_wav(
@@ -834,20 +997,16 @@ class VideoGenerator:
save_video_time = time.perf_counter() - save_start
logger.info("Saved audio to %s", output_path)
elif self._is_image_workload():
# Image workloads (t2i, i2i, …): save the first frame as PNG.
assert frames is not None # implied by save_to_disk and not audio_only
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 # implied by save_to_disk and not audio_only
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:
# Single-pass save path: encode video+audio once, avoiding
# second-pass remux overhead. AudioMuxStage remains 0.0 on
# success because audio is already present in the saved MP4.
save_start = time.perf_counter()
save_ok = self._save_video_with_audio_ffmpeg_pipe(
output_path=output_path,
@@ -895,17 +1054,14 @@ class VideoGenerator:
e2e_time = time.perf_counter() - start_time
logger.info("End-to-end latency: %.2f seconds", e2e_time)
result: dict[str, Any] = {
"prompts": prompt,
return {
"prompts": work_item.prompt,
"samples": samples if batch.return_frames else None,
"frames": frames if batch.return_frames else None,
# Audio is the primary output for audio workloads — return it
# whenever the pipeline produced one, regardless of
# `return_frames` (which gates the video-shaped buffers).
"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": (target_height, target_width, batch.num_frames),
"size": (work_item.target_height, work_item.target_width, batch.num_frames),
"generation_time": gen_time,
"e2e_latency": e2e_time,
"logging_info": logging_info,
@@ -916,8 +1072,168 @@ class VideoGenerator:
"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])]
pipeline_config = fastvideo_args.pipeline_config
if not pipeline_config.dynamic_batching_supported():
raise ValueError(
f"Dynamic batching is enabled but pipeline config {type(pipeline_config).__name__} does not "
"support merging requests; run with batching_mode='disabled', or use a pipeline that opts in "
"via supports_dynamic_batching=True (DMD/causal denoising variants are always excluded)")
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,
sampling_param: SamplingParam | None = None,
fastvideo_args: FastVideoArgs | None = None,
**kwargs,
) -> dict[str, Any]:
"""Internal method for single video generation"""
if fastvideo_args is None:
fastvideo_args = self.fastvideo_args
if not isinstance(prompt, str):
raise TypeError(f"`prompt` must be a string, but got {type(prompt)}")
work_item = self._prepare_generation_work_item(
prompt=prompt,
sampling_param=sampling_param,
fastvideo_args=fastvideo_args,
**kwargs,
)
return self._execute_single_work_item(work_item)
@staticmethod
def _wrap_legacy_result(
result: dict[str, Any] | list[dict[str, Any]], ) -> GenerationResult | list[GenerationResult]:
+46
View File
@@ -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
-1
View File
@@ -248,7 +248,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...
+110 -13
View File
@@ -68,12 +68,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
@@ -85,13 +93,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
@@ -104,6 +123,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()
@@ -253,6 +347,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)
+50 -1
View File
@@ -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:
@@ -153,6 +153,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
# -------------------------------------------------------------------
@@ -166,6 +208,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."""
+7
View File
@@ -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,
@@ -31,23 +31,29 @@ def fa_default_impls():
if not torch.cuda.is_available():
pytest.skip("CUDA is required for the FA2/FA3 default custom-op tests")
# The registration + dispatcher live in `attention/utils/`, alongside the
# FP4 cute template and the masked/varlen wrappers. The backend
# (`attention/backends/flash_attn.py`) just imports `flash_attn_func_compilable`
# from there, so test references the utils module directly.
try:
from fastvideo.attention.backends import flash_attn as fa_backend
from fastvideo.attention.utils import flash_attn_default as fa_module
except ImportError as exc:
pytest.skip(f"flash_attn backend not importable: {exc}")
pytest.skip(f"flash_attn_default not importable: {exc}")
if fa_backend.fa_version not in ("2", "3"):
if fa_module.fa_version not in ("2", "3"):
pytest.skip(
f"FA2/FA3 default custom op only exists for fa_version in (2, 3); "
f"got {fa_backend.fa_version!r}"
f"got {fa_module.fa_version!r}"
)
# compilable dispatcher, the original FA wrapper it falls back to, and
# the raw custom op for opcheck.
# compilable dispatcher, the original FA wrapper it falls back to, the
# raw custom op for opcheck, and the fa_version (FA2 has full register_
# autograd; FA3 keeps the carve-out so some tests gate on this).
return (
fa_backend.flash_attn_func_compilable,
fa_backend._fa_default,
fa_module.flash_attn_func_compilable,
fa_module._fa_default,
torch.ops.fastvideo._flash_attn_default_forward,
fa_module.fa_version,
)
@@ -71,7 +77,7 @@ def test_default_compilable_inference_matches_original(fa_default_impls, dtype,
"""No-grad path routes through the custom op and is numerically identical."""
if dtype == torch.bfloat16 and not torch.cuda.is_bf16_supported():
pytest.skip("bfloat16 is not supported on this GPU")
compilable, original, _ = fa_default_impls
compilable, original, _, _ = fa_default_impls
torch.manual_seed(0)
q, k, v = _qkv(dtype, requires_grad=False)
@@ -91,7 +97,7 @@ def test_default_compilable_training_backward_flows(fa_default_impls, dtype, cau
"""
if dtype == torch.bfloat16 and not torch.cuda.is_bf16_supported():
pytest.skip("bfloat16 is not supported on this GPU")
compilable, original, _ = fa_default_impls
compilable, original, _, _ = fa_default_impls
torch.manual_seed(0)
q_ref, k_ref, v_ref = _qkv(dtype, requires_grad=True)
@@ -115,7 +121,65 @@ def test_default_compilable_training_backward_flows(fa_default_impls, dtype, cau
@pytest.mark.parametrize("causal", [False, True])
def test_default_forward_opcheck(fa_default_impls, causal):
"""Schema / fake-kernel consistency for the custom op (forward only)."""
_, _, op = fa_default_impls
_, _, op, _ = fa_default_impls
torch.manual_seed(0)
q, k, v = _qkv(torch.float16, requires_grad=False)
torch.library.opcheck(op, (q, k, v, None, causal))
# --------------------------------------------------------------------------- #
# FA2-only: backward is registered on the custom op itself. Exercises the #
# `register_autograd` wiring directly (not via the dispatcher's carve-out). #
# Skipped on FA3 until Kuan-Hao's Modal FA3 setup PR lands and we mirror the #
# pattern there. #
# --------------------------------------------------------------------------- #
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
@pytest.mark.parametrize("causal", [False, True])
def test_default_op_backward_through_registered_autograd(fa_default_impls, dtype, causal):
"""FA2: gradients flow through ``torch.ops.fastvideo._flash_attn_default_forward``
itself (no dispatcher carve-out involved) and match the original
``flash_attn_func``'s gradients."""
if dtype == torch.bfloat16 and not torch.cuda.is_bf16_supported():
pytest.skip("bfloat16 is not supported on this GPU")
_, original, op, fa_version = fa_default_impls
if fa_version != "2":
pytest.skip(f"register_autograd is only wired for FA2 right now; got {fa_version!r}")
torch.manual_seed(0)
q_ref, k_ref, v_ref = _qkv(dtype, requires_grad=True)
q_test, k_test, v_test = _clone(q_ref, k_ref, v_ref)
# Reference grads via the original autograd.Function.
out_ref = original(q_ref, k_ref, v_ref, softmax_scale=None, causal=causal)
# Custom-op grads via the registered backward — unpack (out, lse), discard lse.
out_test, _ = op(q_test, k_test, v_test, None, causal)
torch.testing.assert_close(out_test, out_ref,
atol=0 if dtype == torch.float16 else 1e-3,
rtol=0 if dtype == torch.float16 else 1e-3)
dout = torch.randn_like(out_ref)
dq_ref, dk_ref, dv_ref = torch.autograd.grad(
(out_ref * dout).sum(), (q_ref, k_ref, v_ref))
dq_test, dk_test, dv_test = torch.autograd.grad(
(out_test * dout).sum(), (q_test, k_test, v_test))
atol = rtol = 6e-3 if dtype == torch.float16 else 2e-2
torch.testing.assert_close(dq_test, dq_ref, atol=atol, rtol=rtol)
torch.testing.assert_close(dk_test, dk_ref, atol=atol, rtol=rtol)
torch.testing.assert_close(dv_test, dv_ref, atol=atol, rtol=rtol)
@pytest.mark.parametrize("causal", [False, True])
def test_default_op_opcheck_with_grad_inputs(fa_default_impls, causal):
"""FA2: full ``opcheck`` including ``test_autograd_registration`` —
catches a missing/inconsistent backward at unit-test time, which was
exactly the gap that #1373's first revision shipped."""
_, _, op, fa_version = fa_default_impls
if fa_version != "2":
pytest.skip(f"autograd registration only wired for FA2; got {fa_version!r}")
torch.manual_seed(0)
q, k, v = _qkv(torch.float16, requires_grad=True)
torch.library.opcheck(op, (q, k, v, None, causal))
@@ -0,0 +1,239 @@
# SPDX-License-Identifier: Apache-2.0
"""Regression guard for the masked/varlen custom ops.
`flash_attn_no_pad_compilable` / `flash_attn_varlen_qk_no_pad_compilable` wrap
the whole masked-attention functions in `torch.library.custom_op`s so dynamo
sees one traceable node (the internal unpad/pad bookkeeping runs eager inside).
On FA2 these ops register a real backward — `softmax_lse` is padded back to a
statically-shaped `[batch, nheads, seqlen]` form on the way out and re-unpadded
in backward, which calls FA2's `_flash_attn_varlen_backward`. So training also
backprops through the op (no graph break on the training path either).
On FA3/FA4 these ops are forward+fake only and the `*_compilable` dispatchers
carve out to the autograd.Function for grad-enabled calls (PR #1373 pattern).
Tests gating on FA2 are skipped on FA3/FA4.
These tests pin:
- inference (no grad): output through the custom op is bit-identical to the
original function;
- training (requires_grad): gradients through the registered op match the
original autograd.Function;
- schema/fake-kernel consistency via torch.library.opcheck, both with and
without grad-requiring inputs (the latter exercises
`test_autograd_registration` — would have caught a missing/broken backward
at unit-test time).
GPU assumptions: requires CUDA and the FA2 `flash_attn` varlen package.
Skips on CPU and when flash_attn is unavailable.
"""
from __future__ import annotations
import pytest
import torch
@pytest.fixture(scope="module")
def no_pad_impls():
if not torch.cuda.is_available():
pytest.skip("CUDA is required for the masked/varlen custom-op tests")
try:
from fastvideo.attention.utils import flash_attn_no_pad as mod
except ImportError as exc:
pytest.skip(f"flash_attn_no_pad not importable (flash_attn missing?): {exc}")
return mod
def _dtype_skip(dtype):
if dtype == torch.bfloat16 and not torch.cuda.is_bf16_supported():
pytest.skip("bfloat16 is not supported on this GPU")
def _fa2_only(mod):
if mod._FA_VARLEN_VERSION != "2":
pytest.skip(
f"register_autograd is only wired for FA2; got "
f"_FA_VARLEN_VERSION={mod._FA_VARLEN_VERSION!r}"
)
def _padding_mask(batch, seqlen, valid_lens, device):
mask = torch.zeros(batch, seqlen, dtype=torch.bool, device=device)
for i, n in enumerate(valid_lens):
mask[i, :n] = True
return mask
# --------------------------------------------------------------------------- #
# flash_attn_no_pad (masked self-attention) #
# --------------------------------------------------------------------------- #
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
def test_no_pad_inference_matches_original(no_pad_impls, dtype):
"""No-grad path through the custom op is bit-identical to the original."""
_dtype_skip(dtype)
mod = no_pad_impls
torch.manual_seed(0)
device = torch.device("cuda")
b, s, h, d = 2, 64, 4, 64
qkv = torch.randn(b, s, 3, h, d, device=device, dtype=dtype)
mask = _padding_mask(b, s, [64, 48], device)
with torch.inference_mode():
out_ref = mod.flash_attn_no_pad(qkv, mask, causal=False, dropout_p=0.0)
out_test = mod.flash_attn_no_pad_compilable(qkv, mask, causal=False, dropout_p=0.0)
torch.testing.assert_close(out_test, out_ref, atol=0, rtol=0)
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
def test_no_pad_training_backward_through_registered_autograd(no_pad_impls, dtype):
"""FA2: grads flow through the registered op and match the original."""
_dtype_skip(dtype)
mod = no_pad_impls
_fa2_only(mod)
torch.manual_seed(0)
device = torch.device("cuda")
b, s, h, d = 2, 64, 4, 64
mask = _padding_mask(b, s, [64, 48], device)
qkv_ref = torch.randn(b, s, 3, h, d, device=device, dtype=dtype, requires_grad=True)
qkv_test = qkv_ref.detach().clone().requires_grad_(True)
out_ref = mod.flash_attn_no_pad(qkv_ref, mask, causal=False, dropout_p=0.0)
# Go through the compilable wrapper (which on FA2 unconditionally routes
# to the registered op — no carve-out).
out_test = mod.flash_attn_no_pad_compilable(qkv_test, mask, causal=False, dropout_p=0.0)
dout = torch.randn_like(out_ref)
(dqkv_ref,) = torch.autograd.grad((out_ref * dout).sum(), (qkv_ref,))
(dqkv_test,) = torch.autograd.grad((out_test * dout).sum(), (qkv_test,))
atol = rtol = 6e-3 if dtype == torch.float16 else 2e-2
torch.testing.assert_close(dqkv_test, dqkv_ref, atol=atol, rtol=rtol)
def test_no_pad_forward_opcheck(no_pad_impls):
"""Schema/fake-kernel consistency for the custom op (forward only)."""
torch.manual_seed(0)
device = torch.device("cuda")
b, s, h, d = 2, 64, 4, 64
qkv = torch.randn(b, s, 3, h, d, device=device, dtype=torch.float16)
mask = _padding_mask(b, s, [64, 48], device)
torch.library.opcheck(
torch.ops.fastvideo._flash_attn_no_pad_forward,
(qkv, mask, False, 0.0, None, False),
)
def test_no_pad_opcheck_with_grad_inputs(no_pad_impls):
"""FA2: full opcheck including ``test_autograd_registration`` — catches a
missing/inconsistent backward at unit-test time."""
mod = no_pad_impls
_fa2_only(mod)
torch.manual_seed(0)
device = torch.device("cuda")
b, s, h, d = 2, 64, 4, 64
qkv = torch.randn(b, s, 3, h, d, device=device, dtype=torch.float16, requires_grad=True)
mask = _padding_mask(b, s, [64, 48], device)
torch.library.opcheck(
torch.ops.fastvideo._flash_attn_no_pad_forward,
(qkv, mask, False, 0.0, None, False),
)
# --------------------------------------------------------------------------- #
# flash_attn_varlen_qk_no_pad (cross-attn / unequal q-k seqlen) #
# --------------------------------------------------------------------------- #
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
def test_varlen_qk_inference_matches_original(no_pad_impls, dtype):
_dtype_skip(dtype)
mod = no_pad_impls
# The varlen-qk custom op's real forward goes through FA's varlen func; off
# FA2 that path is the pre-existing FA3 `dropout_p` carve-out (out of scope
# here), so skip cleanly like the autograd tests until that's fixed.
_fa2_only(mod)
torch.manual_seed(1)
device = torch.device("cuda")
b, sq, sk, h, d = 2, 48, 64, 4, 64
q = torch.randn(b, sq, h, d, device=device, dtype=dtype)
k = torch.randn(b, sk, h, d, device=device, dtype=dtype)
v = torch.randn(b, sk, h, d, device=device, dtype=dtype)
qmask = _padding_mask(b, sq, [48, 40], device)
kmask = _padding_mask(b, sk, [64, 56], device)
with torch.inference_mode():
out_ref = mod.flash_attn_varlen_qk_no_pad(q, k, v, qmask, kmask, causal=False, dropout_p=0.0)
out_test = mod.flash_attn_varlen_qk_no_pad_compilable(q, k, v, qmask, kmask, causal=False, dropout_p=0.0)
torch.testing.assert_close(out_test, out_ref, atol=0, rtol=0)
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
def test_varlen_qk_training_backward_through_registered_autograd(no_pad_impls, dtype):
"""FA2: grads through q, k, v all flow via the registered op and match."""
_dtype_skip(dtype)
mod = no_pad_impls
_fa2_only(mod)
torch.manual_seed(1)
device = torch.device("cuda")
b, sq, sk, h, d = 2, 48, 64, 4, 64
qmask = _padding_mask(b, sq, [48, 40], device)
kmask = _padding_mask(b, sk, [64, 56], device)
q_ref = torch.randn(b, sq, h, d, device=device, dtype=dtype, requires_grad=True)
k_ref = torch.randn(b, sk, h, d, device=device, dtype=dtype, requires_grad=True)
v_ref = torch.randn(b, sk, h, d, device=device, dtype=dtype, requires_grad=True)
q_test = q_ref.detach().clone().requires_grad_(True)
k_test = k_ref.detach().clone().requires_grad_(True)
v_test = v_ref.detach().clone().requires_grad_(True)
out_ref = mod.flash_attn_varlen_qk_no_pad(q_ref, k_ref, v_ref, qmask, kmask, causal=False, dropout_p=0.0)
out_test = mod.flash_attn_varlen_qk_no_pad_compilable(q_test, k_test, v_test, qmask, kmask,
causal=False, dropout_p=0.0)
dout = torch.randn_like(out_ref)
dq_ref, dk_ref, dv_ref = torch.autograd.grad((out_ref * dout).sum(), (q_ref, k_ref, v_ref))
dq_test, dk_test, dv_test = torch.autograd.grad((out_test * dout).sum(), (q_test, k_test, v_test))
atol = rtol = 6e-3 if dtype == torch.float16 else 2e-2
torch.testing.assert_close(dq_test, dq_ref, atol=atol, rtol=rtol)
torch.testing.assert_close(dk_test, dk_ref, atol=atol, rtol=rtol)
torch.testing.assert_close(dv_test, dv_ref, atol=atol, rtol=rtol)
def test_varlen_qk_forward_opcheck(no_pad_impls):
mod = no_pad_impls
_fa2_only(mod)
torch.manual_seed(1)
device = torch.device("cuda")
b, sq, sk, h, d = 2, 48, 64, 4, 64
q = torch.randn(b, sq, h, d, device=device, dtype=torch.float16)
k = torch.randn(b, sk, h, d, device=device, dtype=torch.float16)
v = torch.randn(b, sk, h, d, device=device, dtype=torch.float16)
qmask = _padding_mask(b, sq, [48, 40], device)
kmask = _padding_mask(b, sk, [64, 56], device)
torch.library.opcheck(
torch.ops.fastvideo._flash_attn_varlen_qk_no_pad_forward,
(q, k, v, qmask, kmask, False, 0.0, None, False),
)
def test_varlen_qk_opcheck_with_grad_inputs(no_pad_impls):
"""FA2: full opcheck with requires_grad inputs (autograd-registration check)."""
mod = no_pad_impls
_fa2_only(mod)
torch.manual_seed(1)
device = torch.device("cuda")
b, sq, sk, h, d = 2, 48, 64, 4, 64
q = torch.randn(b, sq, h, d, device=device, dtype=torch.float16, requires_grad=True)
k = torch.randn(b, sk, h, d, device=device, dtype=torch.float16, requires_grad=True)
v = torch.randn(b, sk, h, d, device=device, dtype=torch.float16, requires_grad=True)
qmask = _padding_mask(b, sq, [48, 40], device)
kmask = _padding_mask(b, sk, [64, 56], device)
torch.library.opcheck(
torch.ops.fastvideo._flash_attn_varlen_qk_no_pad_forward,
(q, k, v, qmask, kmask, False, 0.0, None, False),
)
@@ -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
+226
View File
@@ -0,0 +1,226 @@
# SPDX-License-Identifier: Apache-2.0
"""CPU regression tests for the dynamic VideoBatchScheduler (no GPU needed)."""
from __future__ import annotations
import asyncio
import threading
from types import SimpleNamespace
import pytest
from fastvideo.configs.pipelines.base import PipelineConfig
from fastvideo.configs.pipelines.wan import (
FastWan2_1_T2V_480P_Config,
SelfForcingWanT2V480PConfig,
WanT2V480PConfig,
)
from fastvideo.entrypoints.openai.batching import VideoBatchScheduler
class _FakeBatchGenerator:
def __init__(self):
self.calls: list[list[dict]] = []
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]
class _BlockingFirstCallGenerator(_FakeBatchGenerator):
"""Blocks the first generate_video_batch call until released.
Lets a test deterministically hold one generation "in flight" while more
requests queue up behind it.
"""
def __init__(self):
super().__init__()
self.started = threading.Event()
self.release = threading.Event()
self._first_call = True
def generate_video_batch(self, request_kwargs):
if self._first_call:
self._first_call = False
self.started.set()
assert self.release.wait(timeout=10)
return super().generate_video_batch(request_kwargs)
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(supports_dynamic_batching=True),
)
defaults.update(overrides)
return SimpleNamespace(**defaults)
def _job_kwargs(tmp_path, index: int) -> dict:
return {
"prompt": f"p{index}",
"height": 256,
"width": 256,
"num_frames": 1,
"num_inference_steps": 2,
"seed": index,
"output_path": str(tmp_path / f"video_{index}.mp4"),
"save_video": False,
}
def test_expired_deadline_coalesces_jobs_queued_behind_inflight_generation(tmp_path):
"""Jobs whose delay deadline expired while a generation was in flight must
still coalesce (greedy drain), not dispatch as batches of 1."""
async def run():
generator = _BlockingFirstCallGenerator()
scheduler = VideoBatchScheduler(generator, _batch_scheduler_args(batching_max_size=4, batching_delay_ms=25.0))
await scheduler.start()
try:
loop = asyncio.get_running_loop()
first = asyncio.ensure_future(scheduler.submit("req-1", _job_kwargs(tmp_path, 1)))
await asyncio.wait_for(loop.run_in_executor(None, generator.started.wait, 5), timeout=10)
rest = [asyncio.ensure_future(scheduler.submit(f"req-{i}", _job_kwargs(tmp_path, i))) for i in (2, 3, 4)]
# Let the in-queue jobs outlive the 25 ms batching window while the
# first generation is still blocking the scheduler.
await asyncio.sleep(0.1)
generator.release.set()
await asyncio.wait_for(asyncio.gather(first, *rest), timeout=10)
finally:
await scheduler.stop()
return [[item["prompt"] for item in call] for call in generator.calls]
assert asyncio.run(run()) == [["p1"], ["p2", "p3", "p4"]]
def test_submit_after_stop_raises(tmp_path):
async def run():
scheduler = VideoBatchScheduler(_FakeBatchGenerator(), _batch_scheduler_args())
await scheduler.start()
await scheduler.stop()
with pytest.raises(RuntimeError, match="stopped"):
await scheduler.submit("req-1", _job_kwargs(tmp_path, 1))
asyncio.run(run())
def test_stop_fails_jobs_still_queued_without_hanging(tmp_path):
"""stop() during an in-flight generation must fail the still-queued jobs
instead of stranding their futures forever."""
async def run():
generator = _BlockingFirstCallGenerator()
scheduler = VideoBatchScheduler(generator, _batch_scheduler_args())
await scheduler.start()
loop = asyncio.get_running_loop()
first = asyncio.ensure_future(scheduler.submit("req-1", _job_kwargs(tmp_path, 1)))
await asyncio.wait_for(loop.run_in_executor(None, generator.started.wait, 5), timeout=10)
queued = [asyncio.ensure_future(scheduler.submit(f"req-{i}", _job_kwargs(tmp_path, i))) for i in (2, 3)]
await asyncio.sleep(0) # let the queued submits enqueue
stop_task = asyncio.ensure_future(scheduler.stop())
await asyncio.sleep(0) # let stop() mark the scheduler stopped
generator.release.set()
result = await asyncio.wait_for(first, timeout=10)
assert result["prompts"] == "p1"
for task in queued:
with pytest.raises(RuntimeError, match="stopped before dispatch"):
await asyncio.wait_for(task, timeout=10)
await asyncio.wait_for(stop_task, timeout=10)
asyncio.run(run())
def test_dispatch_exception_fails_all_group_futures(tmp_path):
class _ExplodingGenerator:
def generate_video_batch(self, request_kwargs):
raise ValueError("generation exploded")
async def run():
scheduler = VideoBatchScheduler(_ExplodingGenerator(), _batch_scheduler_args())
await scheduler.start()
try:
return await asyncio.wait_for(
asyncio.gather(
scheduler.submit("req-1", _job_kwargs(tmp_path, 1)),
scheduler.submit("req-2", _job_kwargs(tmp_path, 2)),
return_exceptions=True,
),
timeout=10,
)
finally:
await scheduler.stop()
results = asyncio.run(run())
assert len(results) == 2
for result in results:
assert isinstance(result, ValueError)
assert "generation exploded" in str(result)
def test_run_crash_fails_inflight_and_queued_futures(tmp_path):
"""If the _run task ever dies, every waiting future must fail instead of
hanging, and later submits must be rejected."""
async def run():
scheduler = VideoBatchScheduler(_FakeBatchGenerator(), _batch_scheduler_args())
async def boom(first):
raise RuntimeError("kaboom")
scheduler._collect_batch = boom # type: ignore[method-assign]
await scheduler.start()
results = await asyncio.wait_for(
asyncio.gather(
scheduler.submit("req-1", _job_kwargs(tmp_path, 1)),
scheduler.submit("req-2", _job_kwargs(tmp_path, 2)),
return_exceptions=True,
),
timeout=10,
)
with pytest.raises(RuntimeError, match="stopped"):
await scheduler.submit("req-3", _job_kwargs(tmp_path, 3))
return results
results = asyncio.run(run())
assert len(results) == 2
for result in results:
assert isinstance(result, RuntimeError)
assert "crashed" in str(result)
def test_scheduler_start_rejects_pipeline_without_capability(tmp_path):
async def run():
scheduler = VideoBatchScheduler(
_FakeBatchGenerator(),
_batch_scheduler_args(pipeline_config=PipelineConfig()),
)
with pytest.raises(RuntimeError, match="does not support dynamic request batching"):
await scheduler.start()
assert scheduler._task is None
asyncio.run(run())
def test_dynamic_batching_capability_gate_defaults():
# Default: every pipeline is excluded until its family opts in.
assert PipelineConfig().dynamic_batching_supported() is False
# The verified generic Wan denoising path opts in.
assert WanT2V480PConfig().dynamic_batching_supported() is True
# DMD and causal denoising variants are always vetoed (shared-RNG risk),
# even though they inherit the Wan opt-in flag.
assert FastWan2_1_T2V_480P_Config().dynamic_batching_supported() is False
assert SelfForcingWanT2V480PConfig().dynamic_batching_supported() is False
@@ -0,0 +1,83 @@
# 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_can_dynamic_batch_rejects_preset_latents() -> None:
first = _request()
second = _request(latents=object())
result = can_dynamic_batch(first, second)
assert result.can_batch is False
assert result.reason == "latents"
def test_signature_treats_distinct_opaque_objects_as_incompatible() -> None:
class Opaque:
pass
first = _request("first")
second = _request("second")
distinct = can_dynamic_batch(first, second, base_extra={"foo": Opaque()}, candidate_extra={"foo": Opaque()})
assert distinct.can_batch is False
shared = Opaque()
same_object = can_dynamic_batch(first, second, base_extra={"foo": shared}, candidate_extra={"foo": shared})
assert same_object.can_batch is True
def test_resolution_key_uses_generation_shape() -> None:
assert resolution_key(_request(height=720, width=1280, num_frames=81)) == "720x1280x81"
@@ -0,0 +1,58 @@
# SPDX-License-Identifier: Apache-2.0
"""Guard the Modal FA4 defaults that keep CI lanes on their intended backend.
Pure text/AST analysis: no fastvideo imports, no torch, no Modal client.
"""
from __future__ import annotations
import ast
from pathlib import Path
REPO_ROOT = Path(__file__).resolve().parents[3]
MODAL_ROOT = REPO_ROOT / "fastvideo" / "tests" / "modal"
PR_TEST = MODAL_ROOT / "pr_test.py"
LAUNCH_L40S_JOB = MODAL_ROOT / "launch_l40s_job.py"
SSIM_TEST = MODAL_ROOT / "ssim_test.py"
def _function_strings(path: Path, function_name: str) -> str:
source = path.read_text(encoding="utf-8")
tree = ast.parse(source)
for node in tree.body:
if isinstance(node, ast.FunctionDef) and node.name == function_name:
return "\n".join(
child.value
for child in ast.walk(node)
if isinstance(child, ast.Constant)
and isinstance(child.value, str)
)
raise AssertionError(f"{function_name} not found in {path}")
def test_generic_l40s_launcher_defaults_fa4_off():
source = LAUNCH_L40S_JOB.read_text(encoding="utf-8")
assert '"FASTVIDEO_FA4": os.environ.get("FASTVIDEO_FA4", "0")' in source
def test_ssim_launcher_keeps_fa4_enabled_by_default():
source = SSIM_TEST.read_text(encoding="utf-8")
assert '"FASTVIDEO_FA4": os.environ.get("FASTVIDEO_FA4", "1")' in source
def test_pr_model_load_and_training_lanes_disable_fa4():
lanes = {
"run_transformer_tests": "pytest ./fastvideo/tests/transformers -vs",
"run_training_tests": "pytest ./fastvideo/tests/training/Vanilla -srP",
"run_training_lora_tests": "pytest ./fastvideo/tests/training/lora/test_lora_training.py -srP",
"run_training_tests_VSA": "pytest ./fastvideo/tests/training/VSA -srP",
"run_distill_dmd_tests": "pytest ./fastvideo/tests/training/distill/test_distill_dmd.py -vs",
"run_self_forcing_tests": "pytest ./fastvideo/tests/training/self-forcing/test_self_forcing.py -vs",
"run_train_framework_tests": "pytest ./fastvideo/tests/train/models ./fastvideo/tests/train/methods -vs",
"seed_grad_norm_references": "pytest ./fastvideo/tests/train/methods -vs -rs",
}
for function_name, pytest_command in lanes.items():
function_strings = _function_strings(PR_TEST, function_name)
assert "FASTVIDEO_FA4=0" in function_strings
assert pytest_command in function_strings
@@ -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(supports_dynamic_batching=True),
)
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,18 +3,22 @@ from types import SimpleNamespace
import warnings
import pytest
import torch
from fastvideo.api import (
GenerationRequest,
GenerationResult,
GeneratorConfig,
InputConfig,
OutputConfig,
SamplingConfig,
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 +41,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(supports_dynamic_batching=True),
)
defaults.update(overrides)
return SimpleNamespace(**defaults)
def _patch_from_fastvideo_args(monkeypatch):
captured = {}
@@ -151,6 +174,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)
@@ -558,6 +692,46 @@ def test_generate_batched_request_fans_out_media_inputs(monkeypatch):
]
def test_generate_typed_prompt_list_routes_through_dynamic_batching(monkeypatch, tmp_path):
generator = _new_video_generator()
generator.fastvideo_args = _batching_fastvideo_args()
_patch_sampling_param_from_pretrained(monkeypatch)
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), 0.5, 10.0
monkeypatch.setattr(
"fastvideo.batching.admission.BatchAdmissionController._get_device_memory_gb",
staticmethod(fake_device_memory),
)
monkeypatch.setattr(generator, "_run_forward_batch", fake_run_forward)
results = generator.generate(
GenerationRequest(
prompt=["first prompt", "second prompt"],
sampling=SamplingConfig(num_frames=1, height=8, width=8),
output=OutputConfig(
output_path=str(tmp_path),
save_video=False,
return_frames=True,
),
)
)
assert len(calls) == 1
assert calls[0].prompt == ["first prompt", "second prompt"]
assert [result.prompt for result in results] == ["first prompt", "second prompt"]
assert [result.prompt_index for result in results] == [0, 1]
assert all(isinstance(result, GenerationResult) for result in results)
def test_generate_batched_request_rejects_mismatched_media_inputs(monkeypatch):
generator = _new_runtime_video_generator()
_patch_sampling_param_from_pretrained(monkeypatch)
+4 -3
View File
@@ -97,9 +97,10 @@ 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"),
# FA4 is opt-in (FASTVIDEO_FA4). Generic ad hoc jobs should follow the
# product default unless a caller opts in through the local env or
# --env-vars.
"FASTVIDEO_FA4": os.environ.get("FASTVIDEO_FA4", "0"),
})
)
+19 -11
View File
@@ -73,8 +73,10 @@ ci_env_secret = modal.Secret.from_dict({
**({
"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.
# FA4 is opt-in (FASTVIDEO_FA4). Keep the default enabled for
# inference/perf parity; model-load and training lanes that do not exercise
# FA4 explicitly set FASTVIDEO_FA4=0 in their command strings below.
# Caller override wins.
"FASTVIDEO_FA4": os.environ.get("FASTVIDEO_FA4", "1"),
})
@@ -184,7 +186,8 @@ def run_vae_tests():
volumes={"/root/data": model_vol})
def run_transformer_tests():
run_test(
"export HF_HOME='/root/data/.cache' && hf auth login --token $HF_API_KEY && pytest ./fastvideo/tests/transformers -vs"
"export HF_HOME='/root/data/.cache' && hf auth login --token $HF_API_KEY && "
"FASTVIDEO_FA4=0 pytest ./fastvideo/tests/transformers -vs"
)
@@ -197,7 +200,8 @@ def run_transformer_tests():
volumes={"/root/data": model_vol})
def run_training_tests():
run_test(
"export HF_HOME='/root/data/.cache' && wandb login $WANDB_API_KEY && pytest ./fastvideo/tests/training/Vanilla -srP"
"export HF_HOME='/root/data/.cache' && wandb login $WANDB_API_KEY && "
"FASTVIDEO_FA4=0 pytest ./fastvideo/tests/training/Vanilla -srP"
)
@@ -210,7 +214,8 @@ def run_training_tests():
volumes={"/root/data": model_vol})
def run_training_lora_tests():
run_test(
"export HF_HOME='/root/data/.cache' && wandb login $WANDB_API_KEY && pytest ./fastvideo/tests/training/lora/test_lora_training.py -srP"
"export HF_HOME='/root/data/.cache' && wandb login $WANDB_API_KEY && "
"FASTVIDEO_FA4=0 pytest ./fastvideo/tests/training/lora/test_lora_training.py -srP"
)
@@ -220,7 +225,7 @@ def run_training_lora_tests():
secrets=[wandb_secret, ci_env_secret])
def run_training_tests_VSA():
run_test(
"wandb login $WANDB_API_KEY && pytest ./fastvideo/tests/training/VSA -srP"
"wandb login $WANDB_API_KEY && FASTVIDEO_FA4=0 pytest ./fastvideo/tests/training/VSA -srP"
)
@@ -254,7 +259,7 @@ def run_inference_lora_tests():
@app.function(gpu="L40S:2", image=image, timeout=900, secrets=[ci_env_secret])
def run_distill_dmd_tests():
run_test(
"pytest ./fastvideo/tests/training/distill/test_distill_dmd.py -vs")
"FASTVIDEO_FA4=0 pytest ./fastvideo/tests/training/distill/test_distill_dmd.py -vs")
@app.function(gpu="L40S:2",
@@ -263,14 +268,15 @@ def run_distill_dmd_tests():
secrets=[wandb_secret, ci_env_secret])
def run_self_forcing_tests():
run_test(
"wandb login $WANDB_API_KEY && pytest ./fastvideo/tests/training/self-forcing/test_self_forcing.py -vs"
"wandb login $WANDB_API_KEY && "
"FASTVIDEO_FA4=0 pytest ./fastvideo/tests/training/self-forcing/test_self_forcing.py -vs"
)
@app.function(gpu="L40S:1", image=image, timeout=900, secrets=[ci_env_secret])
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/ ./fastvideo/tests/stages/ ./fastvideo/tests/ops/ ./fastvideo/tests/training/test_trackers.py ./fastvideo/tests/attention/test_sdpa_metadata_mask_contract.py --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/ ./fastvideo/tests/ops/ ./fastvideo/tests/training/test_trackers.py ./fastvideo/tests/attention/test_sdpa_metadata_mask_contract.py --ignore=./fastvideo/tests/entrypoints/test_openai_api_integration.py --ignore=./fastvideo/tests/train/models --ignore=./fastvideo/tests/train/methods -vs"
)
@@ -323,7 +329,8 @@ def run_dreamverse_app_tests():
volumes={"/root/data": model_vol})
def run_train_framework_tests():
run_test(
"export HF_HOME='/root/data/.cache' && hf auth login --token $HF_API_KEY && pytest ./fastvideo/tests/train/models ./fastvideo/tests/train/methods -vs"
"export HF_HOME='/root/data/.cache' && hf auth login --token $HF_API_KEY && "
"FASTVIDEO_FA4=0 pytest ./fastvideo/tests/train/models ./fastvideo/tests/train/methods -vs"
)
@@ -349,7 +356,8 @@ def seed_grad_norm_references():
the local command and the ``_DEVICE_MAPPINGS`` table.
"""
run_test(
"export HF_HOME='/root/data/.cache' && hf auth login --token $HF_API_KEY && FASTVIDEO_GRADNORM_UPDATE=1 pytest ./fastvideo/tests/train/methods -vs -rs"
"export HF_HOME='/root/data/.cache' && hf auth login --token $HF_API_KEY && "
"FASTVIDEO_FA4=0 FASTVIDEO_GRADNORM_UPDATE=1 pytest ./fastvideo/tests/train/methods -vs -rs"
)
@@ -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({
@@ -159,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)
+1 -1
View File
@@ -227,7 +227,7 @@ skip = "./data,./wandb,apps/fastvideo_studio/package-lock.json,apps/performance_
# "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.
ignore-words-list = "tread,passt"
ignore-words-list = "tread,passt,dout"
[tool.ruff]
# Allow lines to be as long as 120.