Compare commits
35
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
b8abf6f36b | ||
|
|
b30903b649 | ||
|
|
7592b57ec0 | ||
|
|
e35fa5452a | ||
|
|
302f79b6f4 | ||
|
|
efc53ba93b | ||
|
|
5ed9c7b472 | ||
|
|
5e9c755e1e | ||
|
|
2ca941a942 | ||
|
|
457017e071 | ||
|
|
dcd2cba956 | ||
|
|
cdd0f7f48f | ||
|
|
6ab03e1745 | ||
|
|
06f354afa1 | ||
|
|
a4f54702c3 | ||
|
|
522b9f7e47 | ||
|
|
961383b077 | ||
|
|
49e4673f6e | ||
|
|
e6c84fe2fd | ||
|
|
733462b4da | ||
|
|
0f413c2e8f | ||
|
|
d8ece1f37f | ||
|
|
2633eb1525 | ||
|
|
8e2b9ced60 | ||
|
|
45d7b55456 | ||
|
|
400e57a9de | ||
|
|
5560eab1db | ||
|
|
ff75cb2c97 | ||
|
|
8d4627845b | ||
|
|
602d57c370 | ||
|
|
a6aab65d97 | ||
|
|
055586703d | ||
|
|
5d89f86675 | ||
|
|
19a51a1fe6 | ||
|
|
d3232cea5a |
@@ -481,6 +481,7 @@ steps:
|
||||
- "fastvideo/layers/**"
|
||||
- "fastvideo/worker/**"
|
||||
- "fastvideo/entrypoints/**"
|
||||
- "fastvideo/performance/**"
|
||||
- "fastvideo/tests/performance/**"
|
||||
- ".buildkite/performance-benchmarks/**"
|
||||
- "pyproject.toml"
|
||||
|
||||
Executable
+133
@@ -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
|
||||
Executable
+122
@@ -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"
|
||||
@@ -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:
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -159,6 +159,16 @@ def legacy_from_pretrained_to_config(
|
||||
preset_refine["guidance_scale"] = value
|
||||
elif key in {"enable_stage_verification", "use_fsdp_inference", "disable_autocast"}:
|
||||
engine[key] = value
|
||||
elif key == "batching_mode":
|
||||
engine.setdefault("batching", {})["mode"] = value
|
||||
elif key == "batching_max_size":
|
||||
engine.setdefault("batching", {})["max_size"] = value
|
||||
elif key == "batching_delay_ms":
|
||||
engine.setdefault("batching", {})["delay_ms"] = value
|
||||
elif key == "batching_config":
|
||||
engine.setdefault("batching", {})["config_path"] = value
|
||||
elif key == "enable_batching_metrics":
|
||||
engine.setdefault("batching", {})["enable_metrics"] = value
|
||||
elif key == "override_text_encoder_quant":
|
||||
quantization["text_encoder_quant"] = value
|
||||
elif key == "workload_type":
|
||||
@@ -244,6 +254,11 @@ def generator_config_to_fastvideo_args(config: GeneratorConfig | Mapping[str, An
|
||||
"enable_stage_verification": engine.enable_stage_verification,
|
||||
"use_fsdp_inference": engine.use_fsdp_inference,
|
||||
"disable_autocast": engine.disable_autocast,
|
||||
"batching_mode": engine.batching.mode,
|
||||
"batching_max_size": engine.batching.max_size,
|
||||
"batching_delay_ms": engine.batching.delay_ms,
|
||||
"batching_config": engine.batching.config_path,
|
||||
"enable_batching_metrics": engine.batching.enable_metrics,
|
||||
}
|
||||
if normalized.pipeline.workload_type is not None:
|
||||
kwargs["workload_type"] = normalized.pipeline.workload_type
|
||||
|
||||
@@ -71,6 +71,15 @@ class QuantizationConfig:
|
||||
transformer_quant: str | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class BatchingConfig:
|
||||
mode: Literal["disabled", "dynamic"] = "disabled"
|
||||
max_size: int = 1
|
||||
delay_ms: float = 0.0
|
||||
config_path: str | None = None
|
||||
enable_metrics: bool = False
|
||||
|
||||
|
||||
@dataclass
|
||||
class EngineConfig:
|
||||
num_gpus: int = 1
|
||||
@@ -82,6 +91,7 @@ class EngineConfig:
|
||||
use_fsdp_inference: bool = False
|
||||
disable_autocast: bool = False
|
||||
quantization: QuantizationConfig | None = None
|
||||
batching: BatchingConfig = field(default_factory=BatchingConfig)
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -280,6 +290,7 @@ class ServeConfig:
|
||||
|
||||
|
||||
__all__ = [
|
||||
"BatchingConfig",
|
||||
"CompileConfig",
|
||||
"ComponentConfig",
|
||||
"ContinuationState",
|
||||
|
||||
@@ -1,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.")
|
||||
@@ -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)
|
||||
|
||||
@@ -0,0 +1,26 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""Dynamic generation batching helpers."""
|
||||
|
||||
from fastvideo.batching.admission import (
|
||||
AdmissionLimit,
|
||||
BatchAdmissionController,
|
||||
BatchingRule,
|
||||
load_batching_config,
|
||||
)
|
||||
from fastvideo.batching.signature import (
|
||||
BatchCompatibility,
|
||||
can_dynamic_batch,
|
||||
dynamic_batch_signature,
|
||||
resolution_key,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"AdmissionLimit",
|
||||
"BatchAdmissionController",
|
||||
"BatchCompatibility",
|
||||
"BatchingRule",
|
||||
"can_dynamic_batch",
|
||||
"dynamic_batch_signature",
|
||||
"load_batching_config",
|
||||
"resolution_key",
|
||||
]
|
||||
@@ -0,0 +1,297 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
from dataclasses import dataclass
|
||||
from difflib import get_close_matches
|
||||
from typing import Any
|
||||
|
||||
from fastvideo.batching.signature import resolution_key
|
||||
from fastvideo.fastvideo_args import FastVideoArgs
|
||||
from fastvideo.logger import init_logger
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
_BYTES_PER_GB = 1024**3
|
||||
_BATCHING_RULE_KEYS = frozenset({
|
||||
"model",
|
||||
"model_contains",
|
||||
"resolution",
|
||||
"device_memory_gb_min",
|
||||
"device_memory_gb_max",
|
||||
"offload",
|
||||
"max_batch_size",
|
||||
"max_cost",
|
||||
"calibration",
|
||||
})
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class AdmissionLimit:
|
||||
max_batch_size: int
|
||||
max_cost: float | None = None
|
||||
cap_reason: str | None = None
|
||||
|
||||
def reject_reason(self, *, batch_size: int, batch_cost: float) -> str | None:
|
||||
if batch_size > self.max_batch_size:
|
||||
return self.cap_reason or f"config_cap:{self.max_batch_size}"
|
||||
if self.max_cost is not None and batch_cost > self.max_cost:
|
||||
return f"cost_budget:{batch_cost:.0f}>{self.max_cost:.0f}"
|
||||
return None
|
||||
|
||||
def stop_reason_for_next_cost(self, next_batch_cost: float) -> str | None:
|
||||
if self.max_cost is not None and next_batch_cost > self.max_cost:
|
||||
return f"cost_budget_next:{next_batch_cost:.0f}>{self.max_cost:.0f}"
|
||||
return None
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class BatchingRule:
|
||||
model: str | None = None
|
||||
model_contains: str | None = None
|
||||
resolution: str | None = None
|
||||
device_memory_gb_min: float | None = None
|
||||
device_memory_gb_max: float | None = None
|
||||
offload: bool | None = None
|
||||
max_batch_size: int = 1
|
||||
max_cost: float | None = None
|
||||
source: str = "user"
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: dict[str, Any], *, source: str) -> BatchingRule:
|
||||
if not isinstance(data, dict):
|
||||
raise ValueError(f"batching config rule from {source} must be an object, got {type(data).__name__}")
|
||||
_validate_rule_keys(data, source=source)
|
||||
if "max_batch_size" not in data:
|
||||
raise ValueError("batching config rule requires max_batch_size")
|
||||
|
||||
rule = cls(
|
||||
model=_optional_str(data.get("model")),
|
||||
model_contains=_optional_str(data.get("model_contains")),
|
||||
resolution=_optional_str(data.get("resolution")),
|
||||
device_memory_gb_min=_optional_float(data.get("device_memory_gb_min")),
|
||||
device_memory_gb_max=_optional_float(data.get("device_memory_gb_max")),
|
||||
offload=_optional_bool(data.get("offload")),
|
||||
max_batch_size=int(data["max_batch_size"]),
|
||||
max_cost=_optional_float(data.get("max_cost")),
|
||||
source=source,
|
||||
)
|
||||
rule.validate()
|
||||
return rule
|
||||
|
||||
def validate(self) -> None:
|
||||
if self.model is not None and self.model_contains is not None:
|
||||
raise ValueError("batching config rule cannot set both model and model_contains")
|
||||
if self.model is None and self.model_contains is None:
|
||||
raise ValueError("batching config rule requires model or model_contains")
|
||||
if self.max_batch_size < 1:
|
||||
raise ValueError("batching config rule max_batch_size must be >= 1")
|
||||
if self.max_cost is not None and self.max_cost <= 0.0:
|
||||
raise ValueError("batching config rule max_cost must be > 0")
|
||||
if (self.device_memory_gb_min is not None and self.device_memory_gb_max is not None
|
||||
and self.device_memory_gb_min > self.device_memory_gb_max):
|
||||
raise ValueError("batching config rule device_memory_gb_min must be <= device_memory_gb_max")
|
||||
|
||||
def matches(
|
||||
self,
|
||||
*,
|
||||
model_path: str,
|
||||
resolution: str | None,
|
||||
device_memory_gb: float | None,
|
||||
offload: bool,
|
||||
) -> bool:
|
||||
if self.model is not None and self.model != model_path:
|
||||
return False
|
||||
if self.model_contains is not None and self.model_contains not in model_path:
|
||||
return False
|
||||
if self.resolution not in (None, "*") and self.resolution != resolution:
|
||||
return False
|
||||
if self.offload is not None and self.offload != offload:
|
||||
return False
|
||||
if device_memory_gb is None:
|
||||
return True
|
||||
if self.device_memory_gb_min is not None and device_memory_gb < self.device_memory_gb_min:
|
||||
return False
|
||||
return not (self.device_memory_gb_max is not None and device_memory_gb > self.device_memory_gb_max)
|
||||
|
||||
|
||||
class BatchAdmissionController:
|
||||
|
||||
def __init__(self, fastvideo_args: FastVideoArgs, *, gpu_id: int = 0):
|
||||
self._mode = fastvideo_args.batching_mode
|
||||
self._user_max_batch_size = max(1, int(fastvideo_args.batching_max_size))
|
||||
self._model_path = fastvideo_args.model_path
|
||||
self._offload = bool(fastvideo_args.dit_cpu_offload or fastvideo_args.dit_layerwise_offload)
|
||||
self._device_memory_gb = self._get_device_memory_gb(gpu_id)
|
||||
self._rules = load_batching_config(fastvideo_args.batching_config)
|
||||
self._pipeline_config = fastvideo_args.pipeline_config
|
||||
|
||||
if self.enabled:
|
||||
logger.info(
|
||||
"Batch admission enabled: user_max=%d, device_memory=%.1fGiB, rules=%d",
|
||||
self._user_max_batch_size,
|
||||
self._device_memory_gb or 0.0,
|
||||
len(self._rules),
|
||||
)
|
||||
|
||||
@property
|
||||
def enabled(self) -> bool:
|
||||
return self._mode == "dynamic" and self._user_max_batch_size > 1
|
||||
|
||||
def reject_reason_for_candidate(self, current_requests: list[Any], candidate_request: Any) -> str | None:
|
||||
if not self.enabled:
|
||||
return None
|
||||
proposed = current_requests + [candidate_request]
|
||||
limit = self.limit_for(proposed[0])
|
||||
return limit.reject_reason(
|
||||
batch_size=len(proposed),
|
||||
batch_cost=self.estimate_batch_cost(proposed),
|
||||
)
|
||||
|
||||
def batch_is_full(self, requests: list[Any]) -> bool:
|
||||
if not self.enabled or not requests:
|
||||
return len(requests) >= self._user_max_batch_size
|
||||
|
||||
limit = self.limit_for(requests[0])
|
||||
if len(requests) >= limit.max_batch_size:
|
||||
return True
|
||||
|
||||
next_cost = self.estimate_batch_cost(requests + [requests[0]])
|
||||
return limit.max_cost is not None and next_cost > limit.max_cost
|
||||
|
||||
def limit_reason_for_batch(self, requests: list[Any]) -> str | None:
|
||||
if not self.enabled or not requests:
|
||||
return None
|
||||
|
||||
limit = self.limit_for(requests[0])
|
||||
if len(requests) >= limit.max_batch_size:
|
||||
return limit.cap_reason or f"config_cap:{limit.max_batch_size}"
|
||||
|
||||
next_cost = self.estimate_batch_cost(requests + [requests[0]])
|
||||
return limit.stop_reason_for_next_cost(next_cost)
|
||||
|
||||
def max_admissible_batch_size(self, request: Any) -> int:
|
||||
return self.limit_for(request).max_batch_size
|
||||
|
||||
def limit_for(self, request: Any) -> AdmissionLimit:
|
||||
rules = self._matching_rules(request)
|
||||
if not rules:
|
||||
return AdmissionLimit(max_batch_size=self._user_max_batch_size)
|
||||
|
||||
config_cap = min(rule.max_batch_size for rule in rules)
|
||||
max_batch_size = min(self._user_max_batch_size, config_cap)
|
||||
cap_reason = f"config_cap:{max_batch_size}" if max_batch_size < self._user_max_batch_size else None
|
||||
costs = [rule.max_cost for rule in rules if rule.max_cost is not None]
|
||||
return AdmissionLimit(
|
||||
max_batch_size=max(1, max_batch_size),
|
||||
max_cost=min(costs) if costs else None,
|
||||
cap_reason=cap_reason,
|
||||
)
|
||||
|
||||
def estimate_batch_cost(self, requests: list[Any]) -> float:
|
||||
return sum(float(self._pipeline_config.estimate_request_cost(request)) for request in requests)
|
||||
|
||||
def _matching_rules(self, request: Any) -> list[BatchingRule]:
|
||||
return [
|
||||
rule for rule in self._rules if rule.matches(
|
||||
model_path=self._model_path,
|
||||
resolution=resolution_key(request),
|
||||
device_memory_gb=self._device_memory_gb,
|
||||
offload=self._offload,
|
||||
)
|
||||
]
|
||||
|
||||
@staticmethod
|
||||
def _get_device_memory_gb(gpu_id: int) -> float | None:
|
||||
try:
|
||||
from fastvideo.platforms import current_platform
|
||||
|
||||
return current_platform.get_device_total_memory(gpu_id) / _BYTES_PER_GB
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
def load_batching_config(path: str | None) -> list[BatchingRule]:
|
||||
if path is None:
|
||||
return []
|
||||
|
||||
with open(path, encoding="utf-8") as f:
|
||||
payload = json.load(f)
|
||||
|
||||
source = os.path.abspath(path)
|
||||
entries = _config_entries(payload)
|
||||
rules = [BatchingRule.from_dict(entry, source=source) for entry in entries]
|
||||
if not rules:
|
||||
raise ValueError(f"batching config {source} does not contain any rules")
|
||||
return rules
|
||||
|
||||
|
||||
def _config_entries(payload: Any) -> list[dict[str, Any]]:
|
||||
if isinstance(payload, dict) and payload.get("schema_version") not in (None, 1):
|
||||
raise ValueError("batching config schema_version must be 1")
|
||||
if isinstance(payload, dict) and isinstance(payload.get("rules"), list):
|
||||
return payload["rules"]
|
||||
if isinstance(payload, list):
|
||||
return payload
|
||||
if isinstance(payload, dict):
|
||||
entries: list[dict[str, Any]] = []
|
||||
for key, value in payload.items():
|
||||
if key == "schema_version" or not isinstance(value, dict):
|
||||
continue
|
||||
model, _sep, resolution = key.partition("|")
|
||||
entry = dict(value)
|
||||
if model:
|
||||
entry.setdefault("model", model)
|
||||
if resolution:
|
||||
entry.setdefault("resolution", resolution)
|
||||
entries.append(entry)
|
||||
return entries
|
||||
raise ValueError("batching config must be a {'schema_version': 1, 'rules': [...]} object, "
|
||||
"a list of rules, or a mapping keyed by model|resolution")
|
||||
|
||||
|
||||
def _validate_rule_keys(data: dict[str, Any], *, source: str) -> None:
|
||||
unknown = sorted(set(data) - _BATCHING_RULE_KEYS)
|
||||
if not unknown:
|
||||
return
|
||||
|
||||
hints = []
|
||||
for key in unknown:
|
||||
matches = get_close_matches(key, _BATCHING_RULE_KEYS, n=1)
|
||||
if matches:
|
||||
hints.append(f"{key!r} (did you mean {matches[0]!r}?)")
|
||||
else:
|
||||
hints.append(repr(key))
|
||||
raise ValueError(f"batching config rule from {source} contains unknown key(s): {', '.join(hints)}")
|
||||
|
||||
|
||||
def _optional_str(value: Any) -> str | None:
|
||||
if value is None:
|
||||
return None
|
||||
return str(value)
|
||||
|
||||
|
||||
def _optional_float(value: Any) -> float | None:
|
||||
if value is None:
|
||||
return None
|
||||
return float(value)
|
||||
|
||||
|
||||
def _optional_bool(value: Any) -> bool | None:
|
||||
if value is None:
|
||||
return None
|
||||
if isinstance(value, bool):
|
||||
return value
|
||||
if isinstance(value, int | float):
|
||||
if value == 1.0:
|
||||
return True
|
||||
if value == 0.0:
|
||||
return False
|
||||
if isinstance(value, str):
|
||||
lowered = value.strip().lower()
|
||||
if lowered in ("1", "true", "yes", "y", "on"):
|
||||
return True
|
||||
if lowered in ("0", "false", "no", "n", "off"):
|
||||
return False
|
||||
raise ValueError(f"cannot parse boolean batching config value: {value!r}")
|
||||
@@ -0,0 +1,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
|
||||
@@ -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 = []
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -10,6 +10,7 @@ from fastapi.middleware.cors import CORSMiddleware
|
||||
|
||||
from fastvideo.api.presets import validate_preset_selection
|
||||
from fastvideo.api.schema import GenerationRequest
|
||||
from fastvideo.entrypoints.openai.batching import VideoBatchScheduler
|
||||
from fastvideo.entrypoints.openai.state import (
|
||||
DEFAULT_OUTPUT_DIR,
|
||||
clear_state,
|
||||
@@ -59,11 +60,29 @@ async def lifespan(app: FastAPI) -> AsyncIterator[None]:
|
||||
generator = VideoGenerator.from_fastvideo_args(args)
|
||||
logger.info("Model loaded successfully.")
|
||||
|
||||
set_state(generator, args, output_dir, default_request=default_request)
|
||||
video_batch_scheduler: VideoBatchScheduler | None = None
|
||||
if args.batching_mode == "dynamic" and args.batching_max_size > 1:
|
||||
video_batch_scheduler = VideoBatchScheduler(generator, args)
|
||||
await video_batch_scheduler.start()
|
||||
logger.info(
|
||||
"Started dynamic video batch scheduler: max_size=%d delay_ms=%.2f",
|
||||
args.batching_max_size,
|
||||
args.batching_delay_ms,
|
||||
)
|
||||
|
||||
set_state(
|
||||
generator,
|
||||
args,
|
||||
output_dir,
|
||||
default_request=default_request,
|
||||
video_batch_scheduler=video_batch_scheduler,
|
||||
)
|
||||
|
||||
yield # server is running
|
||||
|
||||
logger.info("Shutting down — releasing model resources ...")
|
||||
if video_batch_scheduler is not None:
|
||||
await video_batch_scheduler.stop()
|
||||
generator.shutdown()
|
||||
clear_state()
|
||||
logger.info("Shutdown complete.")
|
||||
|
||||
@@ -0,0 +1,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
|
||||
@@ -11,6 +11,7 @@ from typing import TYPE_CHECKING
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from fastvideo.api.schema import GenerationRequest
|
||||
from fastvideo.entrypoints.openai.batching import VideoBatchScheduler
|
||||
from fastvideo.entrypoints.video_generator import VideoGenerator
|
||||
from fastvideo.fastvideo_args import FastVideoArgs
|
||||
|
||||
@@ -20,6 +21,7 @@ _generator: VideoGenerator | None = None
|
||||
_fastvideo_args: FastVideoArgs | None = None
|
||||
_output_dir: str = DEFAULT_OUTPUT_DIR
|
||||
_default_request: GenerationRequest | None = None
|
||||
_video_batch_scheduler: VideoBatchScheduler | None = None
|
||||
|
||||
|
||||
def get_generator() -> VideoGenerator:
|
||||
@@ -44,23 +46,31 @@ def get_default_request() -> GenerationRequest | None:
|
||||
return _default_request
|
||||
|
||||
|
||||
def get_video_batch_scheduler() -> VideoBatchScheduler | None:
|
||||
"""Return the video batch scheduler when dynamic batching is enabled."""
|
||||
return _video_batch_scheduler
|
||||
|
||||
|
||||
def set_state(
|
||||
generator: VideoGenerator,
|
||||
fastvideo_args: FastVideoArgs,
|
||||
output_dir: str,
|
||||
default_request: GenerationRequest | None = None,
|
||||
video_batch_scheduler: VideoBatchScheduler | None = None,
|
||||
) -> None:
|
||||
"""Set all server state at once (called from lifespan)."""
|
||||
global _generator, _fastvideo_args, _output_dir, _default_request
|
||||
global _generator, _fastvideo_args, _output_dir, _default_request, _video_batch_scheduler
|
||||
_generator = generator
|
||||
_fastvideo_args = fastvideo_args
|
||||
_output_dir = output_dir
|
||||
_default_request = default_request
|
||||
_video_batch_scheduler = video_batch_scheduler
|
||||
|
||||
|
||||
def clear_state() -> None:
|
||||
"""Clear server state on shutdown."""
|
||||
global _generator, _fastvideo_args, _default_request
|
||||
global _generator, _fastvideo_args, _default_request, _video_batch_scheduler
|
||||
_generator = None
|
||||
_fastvideo_args = None
|
||||
_default_request = None
|
||||
_video_batch_scheduler = None
|
||||
|
||||
@@ -26,6 +26,7 @@ from fastvideo.entrypoints.openai.state import (
|
||||
get_generator,
|
||||
get_output_dir,
|
||||
get_server_args,
|
||||
get_video_batch_scheduler,
|
||||
)
|
||||
from fastvideo.entrypoints.openai.protocol import (
|
||||
VideoGenerationsRequest,
|
||||
@@ -151,15 +152,19 @@ async def _run_generation(request_id: str, kwargs: dict[str, Any]) -> None:
|
||||
is synchronous) and update the store on completion or failure.
|
||||
"""
|
||||
generator = get_generator()
|
||||
scheduler = get_video_batch_scheduler()
|
||||
loop = asyncio.get_running_loop()
|
||||
|
||||
try:
|
||||
start = time.perf_counter()
|
||||
|
||||
result = await loop.run_in_executor(
|
||||
None,
|
||||
lambda: generator.generate_video(**kwargs),
|
||||
)
|
||||
if scheduler is not None and scheduler.enabled:
|
||||
result = await scheduler.submit(request_id, kwargs)
|
||||
else:
|
||||
result = await loop.run_in_executor(
|
||||
None,
|
||||
lambda: generator.generate_video(**kwargs),
|
||||
)
|
||||
|
||||
elapsed = time.perf_counter() - start
|
||||
update: dict[str, Any] = {
|
||||
|
||||
@@ -18,6 +18,7 @@ import warnings
|
||||
from collections.abc import Mapping
|
||||
from contextlib import suppress
|
||||
from copy import deepcopy
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
import imageio
|
||||
@@ -26,6 +27,8 @@ import torch
|
||||
import torchvision
|
||||
from einops import rearrange
|
||||
|
||||
from fastvideo.batching.admission import BatchAdmissionController
|
||||
from fastvideo.batching.signature import can_dynamic_batch
|
||||
from fastvideo.api.compat import (
|
||||
expand_request_prompt_batch,
|
||||
generator_config_to_fastvideo_args,
|
||||
@@ -94,6 +97,11 @@ _FROM_PRETRAINED_CONVENIENCE_KWARGS = frozenset({
|
||||
"pin_cpu_memory",
|
||||
"enable_torch_compile",
|
||||
"torch_compile_kwargs",
|
||||
"batching_mode",
|
||||
"batching_max_size",
|
||||
"batching_delay_ms",
|
||||
"batching_config",
|
||||
"enable_batching_metrics",
|
||||
"output_type",
|
||||
"nvfp4_fa4",
|
||||
})
|
||||
@@ -112,6 +120,17 @@ def _infer_latent_batch_size(batch: ForwardBatch) -> int:
|
||||
return latent_batch_size
|
||||
|
||||
|
||||
@dataclass
|
||||
class _GenerationWorkItem:
|
||||
prompt: str
|
||||
sampling_param: SamplingParam
|
||||
fastvideo_args: FastVideoArgs
|
||||
batch: ForwardBatch
|
||||
output_path: str
|
||||
target_height: int
|
||||
target_width: int
|
||||
|
||||
|
||||
class VideoGenerator:
|
||||
"""
|
||||
A unified class for generating videos using diffusion models.
|
||||
@@ -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]:
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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...
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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."""
|
||||
|
||||
@@ -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
|
||||
@@ -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)
|
||||
|
||||
@@ -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"),
|
||||
})
|
||||
)
|
||||
|
||||
|
||||
@@ -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
@@ -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.
|
||||
|
||||
Reference in New Issue
Block a user