Compare commits
22
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
5a06b460b2 | ||
|
|
a534ba20a0 | ||
|
|
e9bbaca07d | ||
|
|
9bfa585448 | ||
|
|
b2062556a9 | ||
|
|
c9c5585758 | ||
|
|
9212f4f218 | ||
|
|
6388db815b | ||
|
|
7a4285189f | ||
|
|
a837fe841a | ||
|
|
f9e3680f11 | ||
|
|
98f761ec45 | ||
|
|
c041318f2c | ||
|
|
604e0205a4 | ||
|
|
13213395b4 | ||
|
|
46afee5998 | ||
|
|
c488fa1211 | ||
|
|
d3cff517cd | ||
|
|
2f3d407406 | ||
|
|
bcffa4026e | ||
|
|
6d6a10be7a | ||
|
|
73dd105f3d |
@@ -0,0 +1,99 @@
|
||||
---
|
||||
name: ci-runner
|
||||
description: Work on FastVideo's Slurm-only, change-aware GPU CI lanes, static Buildkite graph, trusted ci-runner policy, lane scripts, and GB200 validation.
|
||||
---
|
||||
|
||||
# Slinky Slurm CI lanes
|
||||
|
||||
FastVideo's `ci-runner` Buildkite queue is the control plane for all active
|
||||
GPU CI. A private host-owned dispatcher leases GPUs from the Slinky Slurm tray
|
||||
and runs the immutable PR SHA inside an isolated Enroot container. Buildkite
|
||||
pipeline upload and Slurm submission occur on the login plane; every test
|
||||
payload executes on Slurm compute.
|
||||
|
||||
The files under `fastvideo/tests/modal/` and `.buildkite/scripts/pr_test.sh`
|
||||
are dormant rollback code. Never add an active Buildkite or slash-command
|
||||
route to them. `pr_test.sh` must continue to reject Buildkite invocations.
|
||||
|
||||
The private operator bundle is deliberately outside this repository because
|
||||
it contains site paths and credentials. See
|
||||
`docs/contributing/ci_architecture.md`; this skill covers the repository half
|
||||
and the coordination contract with that bundle.
|
||||
|
||||
## Invariants
|
||||
|
||||
- `.buildkite/pipeline.yml` contains exactly one static step for every active
|
||||
GPU lane. Each step pins a unique key and label, a 90-minute timeout, the
|
||||
trusted `/opt/fastvideo-ci-runner/run-ci` command (`run-unit` is the one
|
||||
compatibility wrapper), step-level internal `TEST_TYPE`, and
|
||||
`queue: "ci-runner"`.
|
||||
- Active CI contains no `pr_test.sh` command, Modal invocation, default queue,
|
||||
Buildkite plugin, `soft_fail`, or job-controlled artifact glob.
|
||||
- The six Fastcheck lanes use `:microscope:` labels. Full-Suite-only lanes use
|
||||
`:test_tube:` or `:bar_chart:` so direct reruns update the right aggregate.
|
||||
- SSIM and vanilla training request all four GPUs. Keep both in the
|
||||
`fastvideo/slinky/whole-tray` Buildkite concurrency group with a limit of one
|
||||
so the second job does not consume an agent or command timeout while waiting
|
||||
for the same tray.
|
||||
- `/test full` schedules all twenty lanes. `/merge`, `ready`, and new pushes to
|
||||
ready PRs use the trusted base-branch planner in
|
||||
`.github/scripts/plan_merge_ci.py`: automatic Fastcheck remains the universal
|
||||
six-lane baseline, and the merge build adds only path-relevant integration
|
||||
lanes. Unknown source/build paths fail closed to all fourteen additive lanes.
|
||||
The trusted uploader still normalizes and validates the complete static graph
|
||||
before Buildkite evaluates its plan conditions.
|
||||
- Focused merge builds may pass allowlisted golden-gate and SSIM test basenames.
|
||||
The private host validates the lane plan and basenames before staging them,
|
||||
and the in-container scripts validate them again. Direct `/test ssim`,
|
||||
explicit `/test full`, and the weekly main-branch schedule run the complete
|
||||
SSIM matrix.
|
||||
- The trusted uploader serves exactly three entry pipelines:
|
||||
`pr-fastcheck` for automatic PR builds, `ci` for slash-command/ready-label
|
||||
API builds, and `fastvideo-performance-lane` for the weekly schedule. Keep
|
||||
incoming GitHub webhook processing disabled on `ci` so it cannot duplicate
|
||||
`pr-fastcheck` on every PR update.
|
||||
- Test payloads live in `.buildkite/scripts/unit_test.sh` or executable
|
||||
`.buildkite/scripts/lanes/<lane>.sh`. Backend policy (GPU count, extras,
|
||||
secrets, kernel build, artifacts) stays in the agent-owned lane table.
|
||||
- Tests must preserve an inherited `MASTER_PORT`. Packed containers share the
|
||||
tray network namespace, so the private runner assigns a distinct port range
|
||||
per GPU lease and the SSIM scheduler assigns task offsets within its range.
|
||||
- The ARM64 runner image includes the pinned FA4 CuTe overlay validated on
|
||||
GB200. Keep SSIM at `FASTVIDEO_FA4=1` because its references were seeded with
|
||||
FA4; keep lanes with FA2 baselines at `FASTVIDEO_FA4=0`. A runner image change
|
||||
must revalidate both the FA4 import and an actual GB200 forward kernel.
|
||||
- `fastvideo/tests/ssim/ci_runner.py` is the active four-GPU SSIM scheduler.
|
||||
New SSIM files are discovered through `REQUIRED_GPUS` and
|
||||
`*_MODEL_TO_PARAMS`; do not wire them through the dormant Modal scheduler.
|
||||
- The host policy fail-closes unknown tuples. A repository-side lane change is
|
||||
inert until the operator updates the private lane table and uploader policy
|
||||
in the same rollout.
|
||||
|
||||
## Adding or changing a lane
|
||||
|
||||
1. Read the closest `AGENTS.md` and the domain-specific testing guide.
|
||||
2. Add or update the executable lane payload under `.buildkite/scripts/`.
|
||||
Keep it deterministic and free of host-specific paths or credential fetches.
|
||||
3. Add the static pipeline step and canonical `/test <name>` mapping. Keep the
|
||||
`<name>-ci` alias only when compatibility requires it.
|
||||
4. Add its source/test path ownership to `.github/scripts/plan_merge_ci.py`.
|
||||
Prefer the narrowest correctness-preserving lane set; leave unknown paths
|
||||
fail-closed. Extend `fastvideo/tests/contract/test_ci_test_collection.py`,
|
||||
`test_merge_ci_plan.py`, and focused CPU-only scheduler/policy tests.
|
||||
5. Coordinate the private lane row: GPU count (1-4), wall time, script, scope
|
||||
pairs, step key, command, HF cache/token, tracking mode, extras, attention
|
||||
backend policy, kernel policy, and artifact relay. Active training lanes
|
||||
keep W&B offline and do not stage a W&B credential.
|
||||
6. Update the trusted pipeline-uploader schema. A mismatch must reject the
|
||||
pipeline rather than silently skip a lane.
|
||||
7. Run `pre-commit run --files <changed paths>`, the planner's representative
|
||||
diff matrix, contract tests, private driver tests, and a real GB200 canary.
|
||||
Multi-GPU, hardware-reference, training, performance, and SSIM changes need
|
||||
their own target-hardware evidence.
|
||||
|
||||
## Rollback
|
||||
|
||||
Rollback the Slurm routing/configuration change or pause the `ci-runner` queue.
|
||||
Do not silently reactivate Modal. A manual Modal experiment requires the
|
||||
explicit local opt-in documented in `ci_architecture.md`; returning it to
|
||||
production CI needs a separate reviewed decision.
|
||||
@@ -1,6 +1,6 @@
|
||||
---
|
||||
name: reseed-ssim-references
|
||||
description: Re-seed HF reference videos for a single existing SSIM test on Modal L40S. Always backs up current refs locally first, regenerates on Modal, pauses for the user to eyeball before-vs-after quality, then overwrites the targeted `<model_id>` subtree on `FastVideo/ssim-reference-videos` with `--force`. Use when an intentional code change (model port fix, attention backend swap, kernel upgrade, hyperparameter change) has invalidated existing refs and they need to be regenerated. Pairs with `seed-ssim-references`, which is for first-time seeding only.
|
||||
description: Re-seed HF reference videos for a single existing SSIM test on Modal L40S. Always backs up current refs locally first, regenerates on Modal, pauses for the user to eyeball before-vs-after quality, then overwrites the targeted model subtree on `FastVideo/ssim-reference-videos` with `--force`. Use when an intentional code change (model port fix, attention backend swap, kernel upgrade, hyperparameter change) has invalidated existing refs and they need to be regenerated. Pairs with `seed-ssim-references`, which is for first-time seeding only.
|
||||
---
|
||||
|
||||
# Re-seed SSIM Reference Videos
|
||||
@@ -13,7 +13,7 @@ on HF — the old refs are overwritten — so the skill always:
|
||||
|
||||
1. Confirms intent with a one-liner the user has to type.
|
||||
2. Downloads the existing refs as a local, timestamped backup.
|
||||
3. Regenerates on Modal L40S (same code path that CI uses).
|
||||
3. Regenerates through the manual legacy Modal L40S maintenance path.
|
||||
4. Pauses for a side-by-side eyeball of backup vs new mp4s.
|
||||
5. Uploads with `--force`, scoped to the single `--model-id`.
|
||||
6. Reminds the user to keep the backup until the PR lands.
|
||||
@@ -51,8 +51,9 @@ harder to recover from than failing closed.
|
||||
|
||||
Hardcoded:
|
||||
|
||||
- Modal GPU: **L40S** (matches CI; re-seeding from another SKU produces refs
|
||||
that L40S CI cannot match).
|
||||
- Modal GPU: **L40S**. This is a manual reference-maintenance target, not the
|
||||
active Slurm CI compute path; changing the SKU also changes the historical
|
||||
`L40S_reference_videos` contract.
|
||||
- Quality tier: **`default`**. `full_quality` is a separate, deliberate
|
||||
operation.
|
||||
- HF repo: `FastVideo/ssim-reference-videos` (override via
|
||||
|
||||
@@ -35,7 +35,8 @@ The skill is run **manually**, once per new test. Before invoking it, the user
|
||||
has already sanity-tested the new test locally — it launches `VideoGenerator`
|
||||
and writes an artefact without crashing (the missing-reference assertion at
|
||||
the end is expected). The skill does not re-test locally; it goes straight
|
||||
to Modal L40S (which is what CI uses).
|
||||
to the manual legacy Modal L40S reference-maintenance target. Active CI runs
|
||||
on the Slinky Slurm cluster and only consumes the resulting references.
|
||||
|
||||
## When to use
|
||||
|
||||
@@ -61,7 +62,8 @@ Prompt the user for it if they didn't supply it.
|
||||
|
||||
Everything else is fixed:
|
||||
|
||||
- Modal runner GPU: **L40S** (hardcoded in `fastvideo/tests/modal/ssim_test.py`).
|
||||
- Modal maintenance GPU: **L40S** (hardcoded in
|
||||
`fastvideo/tests/modal/ssim_test.py`; this is not the active CI compute path).
|
||||
- Device folder: `L40S_reference_videos`.
|
||||
- Quality tier: `default` (the tier CI runs). The `full_quality` tier is not
|
||||
seeded by this skill.
|
||||
|
||||
+448
-528
@@ -1,7 +1,8 @@
|
||||
env:
|
||||
IMAGE_VERSION: "py3.12-latest"
|
||||
BUILDKITE_CLEAN_CHECKOUT: true
|
||||
# Buildkite only launches Modal; remote jobs initialize their own submodules.
|
||||
# Slurm workers clone the immutable commit and initialize submodules inside
|
||||
# their isolated container. The Buildkite login-plane checkout is a no-op.
|
||||
BUILDKITE_GIT_SUBMODULES: false
|
||||
|
||||
notify:
|
||||
@@ -10,539 +11,458 @@ notify:
|
||||
if: build.env("TEST_SCOPE") == "fastcheck" || build.env("TEST_SCOPE") == null
|
||||
- github_commit_status:
|
||||
context: "full-suite-passed"
|
||||
if: build.env("TEST_SCOPE") == "full"
|
||||
if: build.env("TEST_SCOPE") == "full" || build.env("TEST_SCOPE") == "merge"
|
||||
- github_commit_status:
|
||||
context: "direct-test-completed"
|
||||
if: build.env("TEST_SCOPE") == "direct"
|
||||
- github_commit_status:
|
||||
context: "scheduled-ssim-passed"
|
||||
if: build.env("TEST_SCOPE") == "scheduled"
|
||||
|
||||
# This is the complete active GPU CI surface. Every command is a trusted host
|
||||
# dispatcher, and every test payload executes inside the Slinky Slurm tray.
|
||||
# fastvideo/tests/modal remains available only for an explicit manual rollback;
|
||||
# no active pipeline or slash-command route invokes it.
|
||||
steps:
|
||||
# ============================================================
|
||||
# Direct test: triggered by /test <name> slash command.
|
||||
# Labels match fastcheck/full-suite counterparts so the GitHub
|
||||
# check status overwrites the original failed check.
|
||||
# Only ONE step executes per build (gated by TEST_TYPE).
|
||||
# ============================================================
|
||||
- label: ":microscope: Encoder Tests"
|
||||
key: "encoder"
|
||||
if: |
|
||||
build.env("TEST_SCOPE") == "full" ||
|
||||
(build.env("TEST_SCOPE") == "merge" &&
|
||||
build.env("MERGE_TEST_PLAN") =~ /,encoder,/) ||
|
||||
build.env("TEST_SCOPE") == "fastcheck" ||
|
||||
build.env("TEST_SCOPE") == null ||
|
||||
(build.env("TEST_SCOPE") == "direct" &&
|
||||
(build.env("TEST_TYPE") == "encoder" || build.env("TEST_TYPE") == "encoder_ci"))
|
||||
command: "/opt/fastvideo-ci-runner/run-ci"
|
||||
timeout_in_minutes: 90
|
||||
env:
|
||||
TEST_TYPE: "encoder_ci"
|
||||
retry:
|
||||
automatic:
|
||||
- exit_status: 128
|
||||
limit: 3
|
||||
- exit_status: -1
|
||||
limit: 2
|
||||
agents:
|
||||
queue: "ci-runner"
|
||||
|
||||
# --- Fastcheck-scope direct tests ---
|
||||
- label: ":microscope: Encoder Tests"
|
||||
if: build.env("TEST_SCOPE") == "direct" && build.env("TEST_TYPE") == "encoder"
|
||||
command: "timeout 90m .buildkite/scripts/pr_test.sh"
|
||||
retry:
|
||||
automatic:
|
||||
- exit_status: 128
|
||||
limit: 3
|
||||
- exit_status: -1
|
||||
limit: 2
|
||||
agents:
|
||||
queue: "default"
|
||||
- label: ":microscope: VAE Tests"
|
||||
if: build.env("TEST_SCOPE") == "direct" && build.env("TEST_TYPE") == "vae"
|
||||
command: "timeout 90m .buildkite/scripts/pr_test.sh"
|
||||
retry:
|
||||
automatic:
|
||||
- exit_status: 128
|
||||
limit: 3
|
||||
- exit_status: -1
|
||||
limit: 2
|
||||
agents:
|
||||
queue: "default"
|
||||
- label: ":microscope: Transformer Tests"
|
||||
if: build.env("TEST_SCOPE") == "direct" && build.env("TEST_TYPE") == "transformer"
|
||||
command: "timeout 90m .buildkite/scripts/pr_test.sh"
|
||||
retry:
|
||||
automatic:
|
||||
- exit_status: 128
|
||||
limit: 3
|
||||
- exit_status: -1
|
||||
limit: 2
|
||||
agents:
|
||||
queue: "default"
|
||||
- label: ":microscope: Kernel Tests"
|
||||
if: build.env("TEST_SCOPE") == "direct" && build.env("TEST_TYPE") == "kernel_tests"
|
||||
command: "timeout 90m .buildkite/scripts/pr_test.sh"
|
||||
retry:
|
||||
automatic:
|
||||
- exit_status: 128
|
||||
limit: 3
|
||||
- exit_status: -1
|
||||
limit: 2
|
||||
agents:
|
||||
queue: "default"
|
||||
- label: ":vertical_traffic_light: Golden-Gate Tests"
|
||||
if: build.env("TEST_SCOPE") == "direct" && build.env("TEST_TYPE") == "golden_gate"
|
||||
command: "timeout 90m .buildkite/scripts/pr_test.sh"
|
||||
retry:
|
||||
automatic:
|
||||
- exit_status: 128
|
||||
limit: 3
|
||||
- exit_status: -1
|
||||
limit: 2
|
||||
agents:
|
||||
queue: "default"
|
||||
- label: ":microscope: Unit Tests"
|
||||
if: build.env("TEST_SCOPE") == "direct" && build.env("TEST_TYPE") == "unit_test"
|
||||
command: "timeout 90m .buildkite/scripts/pr_test.sh"
|
||||
retry:
|
||||
automatic:
|
||||
- exit_status: 128
|
||||
limit: 3
|
||||
- exit_status: -1
|
||||
limit: 2
|
||||
agents:
|
||||
queue: "default"
|
||||
- label: ":microscope: DreamVerse App Tests"
|
||||
if: build.env("TEST_SCOPE") == "direct" && build.env("TEST_TYPE") == "dreamverse_app"
|
||||
command: "timeout 90m .buildkite/scripts/pr_test.sh"
|
||||
retry:
|
||||
automatic:
|
||||
- exit_status: 128
|
||||
limit: 3
|
||||
- exit_status: -1
|
||||
limit: 2
|
||||
agents:
|
||||
queue: "default"
|
||||
- label: ":microscope: VAE Tests"
|
||||
key: "vae"
|
||||
if: |
|
||||
build.env("TEST_SCOPE") == "full" ||
|
||||
(build.env("TEST_SCOPE") == "merge" &&
|
||||
build.env("MERGE_TEST_PLAN") =~ /,vae,/) ||
|
||||
build.env("TEST_SCOPE") == "fastcheck" ||
|
||||
build.env("TEST_SCOPE") == null ||
|
||||
(build.env("TEST_SCOPE") == "direct" &&
|
||||
(build.env("TEST_TYPE") == "vae" || build.env("TEST_TYPE") == "vae_ci"))
|
||||
command: "/opt/fastvideo-ci-runner/run-ci"
|
||||
timeout_in_minutes: 90
|
||||
env:
|
||||
TEST_TYPE: "vae_ci"
|
||||
retry:
|
||||
automatic:
|
||||
- exit_status: 128
|
||||
limit: 3
|
||||
- exit_status: -1
|
||||
limit: 2
|
||||
agents:
|
||||
queue: "ci-runner"
|
||||
|
||||
# --- Full-suite-scope direct tests ---
|
||||
- label: ":bar_chart: SSIM Tests"
|
||||
if: build.env("TEST_SCOPE") == "direct" && build.env("TEST_TYPE") == "ssim"
|
||||
command: "timeout 90m .buildkite/scripts/pr_test.sh"
|
||||
retry:
|
||||
automatic:
|
||||
- exit_status: 128
|
||||
limit: 3
|
||||
- exit_status: -1
|
||||
limit: 2
|
||||
- exit_status: 1
|
||||
limit: 2
|
||||
agents:
|
||||
queue: "default"
|
||||
- label: ":test_tube: LoRA Inference Tests"
|
||||
if: build.env("TEST_SCOPE") == "direct" && build.env("TEST_TYPE") == "inference_lora"
|
||||
command: "timeout 90m .buildkite/scripts/pr_test.sh"
|
||||
retry:
|
||||
automatic:
|
||||
- exit_status: 128
|
||||
limit: 3
|
||||
- exit_status: -1
|
||||
limit: 2
|
||||
agents:
|
||||
queue: "default"
|
||||
- label: ":test_tube: LoRA Extraction Tests"
|
||||
if: build.env("TEST_SCOPE") == "direct" && build.env("TEST_TYPE") == "lora_extraction"
|
||||
command: "timeout 90m .buildkite/scripts/pr_test.sh"
|
||||
retry:
|
||||
automatic:
|
||||
- exit_status: 128
|
||||
limit: 3
|
||||
- exit_status: -1
|
||||
limit: 2
|
||||
agents:
|
||||
queue: "default"
|
||||
- label: ":test_tube: Training Tests"
|
||||
if: build.env("TEST_SCOPE") == "direct" && build.env("TEST_TYPE") == "training"
|
||||
command: "timeout 90m .buildkite/scripts/pr_test.sh"
|
||||
retry:
|
||||
automatic:
|
||||
- exit_status: 128
|
||||
limit: 3
|
||||
- exit_status: -1
|
||||
limit: 2
|
||||
agents:
|
||||
queue: "default"
|
||||
- label: ":test_tube: Distillation DMD Tests"
|
||||
if: build.env("TEST_SCOPE") == "direct" && build.env("TEST_TYPE") == "distillation_dmd"
|
||||
command: "timeout 90m .buildkite/scripts/pr_test.sh"
|
||||
retry:
|
||||
automatic:
|
||||
- exit_status: 128
|
||||
limit: 3
|
||||
- exit_status: -1
|
||||
limit: 2
|
||||
agents:
|
||||
queue: "default"
|
||||
- label: ":test_tube: Self-Forcing Tests"
|
||||
if: build.env("TEST_SCOPE") == "direct" && build.env("TEST_TYPE") == "self_forcing"
|
||||
command: "timeout 90m .buildkite/scripts/pr_test.sh"
|
||||
retry:
|
||||
automatic:
|
||||
- exit_status: 128
|
||||
limit: 3
|
||||
- exit_status: -1
|
||||
limit: 2
|
||||
agents:
|
||||
queue: "default"
|
||||
- label: ":test_tube: LoRA Training Tests"
|
||||
if: build.env("TEST_SCOPE") == "direct" && build.env("TEST_TYPE") == "training_lora"
|
||||
command: "timeout 90m .buildkite/scripts/pr_test.sh"
|
||||
retry:
|
||||
automatic:
|
||||
- exit_status: 128
|
||||
limit: 3
|
||||
- exit_status: -1
|
||||
limit: 2
|
||||
- exit_status: 1
|
||||
limit: 2
|
||||
agents:
|
||||
queue: "default"
|
||||
- label: ":test_tube: Training Tests VSA"
|
||||
if: build.env("TEST_SCOPE") == "direct" && build.env("TEST_TYPE") == "training_vsa"
|
||||
command: "timeout 90m .buildkite/scripts/pr_test.sh"
|
||||
retry:
|
||||
automatic:
|
||||
- exit_status: 128
|
||||
limit: 3
|
||||
- exit_status: -1
|
||||
limit: 2
|
||||
- exit_status: 1
|
||||
limit: 2
|
||||
agents:
|
||||
queue: "default"
|
||||
- label: ":test_tube: Inference Tests VMoBA"
|
||||
if: build.env("TEST_SCOPE") == "direct" && build.env("TEST_TYPE") == "inference_vmoba"
|
||||
command: "timeout 90m .buildkite/scripts/pr_test.sh"
|
||||
retry:
|
||||
automatic:
|
||||
- exit_status: 128
|
||||
limit: 3
|
||||
- exit_status: -1
|
||||
limit: 2
|
||||
agents:
|
||||
queue: "default"
|
||||
- label: ":test_tube: Performance Tests"
|
||||
if: build.env("TEST_SCOPE") == "direct" && build.env("TEST_TYPE") == "performance"
|
||||
command: "timeout 90m .buildkite/scripts/pr_test.sh"
|
||||
retry:
|
||||
automatic:
|
||||
- exit_status: 128
|
||||
limit: 3
|
||||
- exit_status: -1
|
||||
limit: 2
|
||||
agents:
|
||||
queue: "default"
|
||||
- label: ":test_tube: API Server Tests"
|
||||
if: build.env("TEST_SCOPE") == "direct" && build.env("TEST_TYPE") == "api_server"
|
||||
command: "timeout 90m .buildkite/scripts/pr_test.sh"
|
||||
retry:
|
||||
automatic:
|
||||
- exit_status: 128
|
||||
limit: 3
|
||||
- exit_status: -1
|
||||
limit: 2
|
||||
agents:
|
||||
queue: "default"
|
||||
- label: ":test_tube: Train Framework Tests"
|
||||
if: build.env("TEST_SCOPE") == "direct" && build.env("TEST_TYPE") == "train_framework"
|
||||
command: "timeout 90m .buildkite/scripts/pr_test.sh"
|
||||
retry:
|
||||
automatic:
|
||||
- exit_status: 128
|
||||
limit: 3
|
||||
- exit_status: -1
|
||||
limit: 2
|
||||
agents:
|
||||
queue: "default"
|
||||
- label: ":test_tube: Eval Metrics Tests"
|
||||
if: build.env("TEST_SCOPE") == "direct" && build.env("TEST_TYPE") == "eval"
|
||||
command: "timeout 90m .buildkite/scripts/pr_test.sh"
|
||||
retry:
|
||||
automatic:
|
||||
- exit_status: 128
|
||||
limit: 3
|
||||
- exit_status: -1
|
||||
limit: 2
|
||||
agents:
|
||||
queue: "default"
|
||||
- label: ":microscope: Transformer Tests"
|
||||
key: "transformer"
|
||||
if: |
|
||||
build.env("TEST_SCOPE") == "full" ||
|
||||
(build.env("TEST_SCOPE") == "merge" &&
|
||||
build.env("MERGE_TEST_PLAN") =~ /,transformer,/) ||
|
||||
build.env("TEST_SCOPE") == "fastcheck" ||
|
||||
build.env("TEST_SCOPE") == null ||
|
||||
(build.env("TEST_SCOPE") == "direct" &&
|
||||
(build.env("TEST_TYPE") == "transformer" || build.env("TEST_TYPE") == "transformer_ci"))
|
||||
command: "/opt/fastvideo-ci-runner/run-ci"
|
||||
timeout_in_minutes: 90
|
||||
env:
|
||||
TEST_TYPE: "transformer_ci"
|
||||
retry:
|
||||
automatic:
|
||||
- exit_status: 128
|
||||
limit: 3
|
||||
- exit_status: -1
|
||||
limit: 2
|
||||
agents:
|
||||
queue: "ci-runner"
|
||||
|
||||
# ============================================================
|
||||
# Fastcheck: Runs on every PR (~10-15 min parallel)
|
||||
# Core component validation: encoders, VAEs, transformers,
|
||||
# CUDA kernels, and unit tests.
|
||||
# ============================================================
|
||||
- label: "Trigger Fastcheck"
|
||||
if: build.env("TEST_SCOPE") == "fastcheck" || build.env("TEST_SCOPE") == null
|
||||
retry:
|
||||
automatic:
|
||||
- exit_status: 128
|
||||
limit: 3
|
||||
- exit_status: -1
|
||||
limit: 2
|
||||
plugins:
|
||||
- monorepo-diff#v1.4.0:
|
||||
diff: 'git fetch origin "${BUILDKITE_PULL_REQUEST_BASE_BRANCH:-main}" && git diff --name-only "origin/${BUILDKITE_PULL_REQUEST_BASE_BRANCH:-main}...HEAD"'
|
||||
watch:
|
||||
- path:
|
||||
- "fastvideo/models/encoders/**"
|
||||
- "fastvideo/models/loader/**"
|
||||
- "fastvideo/tests/encoders/**"
|
||||
- "pyproject.toml"
|
||||
- "docker/Dockerfile"
|
||||
config:
|
||||
command: "timeout 20m .buildkite/scripts/pr_test.sh"
|
||||
label: ":microscope: Encoder Tests"
|
||||
env:
|
||||
- TEST_TYPE=encoder
|
||||
agents:
|
||||
queue: "default"
|
||||
- path:
|
||||
- "fastvideo/models/vaes/**"
|
||||
- "fastvideo/models/loader/**"
|
||||
- "fastvideo/tests/vaes/**"
|
||||
- "pyproject.toml"
|
||||
- "docker/Dockerfile"
|
||||
config:
|
||||
command: "timeout 20m .buildkite/scripts/pr_test.sh"
|
||||
label: ":microscope: VAE Tests"
|
||||
env:
|
||||
- TEST_TYPE=vae
|
||||
agents:
|
||||
queue: "default"
|
||||
- path:
|
||||
- "fastvideo/models/dits/**"
|
||||
- "fastvideo/models/loader/**"
|
||||
- "fastvideo/tests/transformers/**"
|
||||
- "fastvideo/layers/**"
|
||||
- "fastvideo/attention/**"
|
||||
- "pyproject.toml"
|
||||
- "docker/Dockerfile"
|
||||
config:
|
||||
command: "timeout 15m .buildkite/scripts/pr_test.sh"
|
||||
label: ":microscope: Transformer Tests"
|
||||
env:
|
||||
- TEST_TYPE=transformer
|
||||
agents:
|
||||
queue: "default"
|
||||
- path:
|
||||
- "fastvideo-kernel/**"
|
||||
- "pyproject.toml"
|
||||
- "docker/Dockerfile"
|
||||
config:
|
||||
command: "timeout 15m .buildkite/scripts/pr_test.sh"
|
||||
label: ":microscope: Kernel Tests"
|
||||
env:
|
||||
- TEST_TYPE=kernel_tests
|
||||
agents:
|
||||
queue: "default"
|
||||
- path:
|
||||
- "fastvideo/**"
|
||||
- ".buildkite/**"
|
||||
- ".github/**"
|
||||
- "pyproject.toml"
|
||||
- "docker/Dockerfile"
|
||||
config:
|
||||
command: "timeout 15m .buildkite/scripts/pr_test.sh"
|
||||
label: ":microscope: Unit Tests"
|
||||
env:
|
||||
- TEST_TYPE=unit_test
|
||||
agents:
|
||||
queue: "default"
|
||||
- path:
|
||||
- "apps/dreamverse/**"
|
||||
- "pyproject.toml"
|
||||
config:
|
||||
command: "timeout 30m .buildkite/scripts/pr_test.sh"
|
||||
label: ":microscope: DreamVerse App Tests"
|
||||
env:
|
||||
- TEST_TYPE=dreamverse_app
|
||||
agents:
|
||||
queue: "default"
|
||||
- label: ":microscope: Kernel Tests"
|
||||
key: "kernel-tests"
|
||||
if: |
|
||||
build.env("TEST_SCOPE") == "full" ||
|
||||
(build.env("TEST_SCOPE") == "merge" &&
|
||||
build.env("MERGE_TEST_PLAN") =~ /,kernel-tests,/) ||
|
||||
build.env("TEST_SCOPE") == "fastcheck" ||
|
||||
build.env("TEST_SCOPE") == null ||
|
||||
(build.env("TEST_SCOPE") == "direct" &&
|
||||
(build.env("TEST_TYPE") == "kernel_tests" || build.env("TEST_TYPE") == "kernel_tests_ci"))
|
||||
command: "/opt/fastvideo-ci-runner/run-ci"
|
||||
timeout_in_minutes: 90
|
||||
env:
|
||||
TEST_TYPE: "kernel_tests_ci"
|
||||
retry:
|
||||
automatic:
|
||||
- exit_status: 128
|
||||
limit: 3
|
||||
- exit_status: -1
|
||||
limit: 2
|
||||
agents:
|
||||
queue: "ci-runner"
|
||||
|
||||
# ============================================================
|
||||
# Full Suite: Runs when TEST_SCOPE=full
|
||||
# Triggered by adding the 'ready' label (via ci-trigger-full-suite.yml)
|
||||
# or on-demand via /test full slash command.
|
||||
# Includes integration tests, SSIM regression, training pipelines,
|
||||
# and performance benchmarks.
|
||||
# ============================================================
|
||||
- label: "Trigger Full Suite"
|
||||
if: build.env("TEST_SCOPE") == "full"
|
||||
retry:
|
||||
automatic:
|
||||
- exit_status: 128
|
||||
limit: 3
|
||||
- exit_status: -1
|
||||
limit: 2
|
||||
plugins:
|
||||
- monorepo-diff#v1.4.0:
|
||||
diff: 'git fetch origin "${BUILDKITE_PULL_REQUEST_BASE_BRANCH:-main}" && git diff --name-only "origin/${BUILDKITE_PULL_REQUEST_BASE_BRANCH:-main}...HEAD"'
|
||||
watch:
|
||||
- path:
|
||||
- "fastvideo/**/*.py"
|
||||
- "pyproject.toml"
|
||||
- "docker/Dockerfile"
|
||||
config:
|
||||
command: "timeout 90m .buildkite/scripts/pr_test.sh"
|
||||
label: ":bar_chart: SSIM Tests"
|
||||
env:
|
||||
- TEST_TYPE=ssim
|
||||
retry:
|
||||
automatic:
|
||||
- exit_status: 1
|
||||
limit: 2
|
||||
agents:
|
||||
queue: "default"
|
||||
- path:
|
||||
- "fastvideo/tests/lora/**"
|
||||
- "fastvideo/models/loader/**"
|
||||
- "fastvideo/tests/transformers/**"
|
||||
- "fastvideo/pipelines/**"
|
||||
- "fastvideo/layers/lora/**"
|
||||
- "pyproject.toml"
|
||||
- "docker/Dockerfile"
|
||||
config:
|
||||
command: "timeout 20m .buildkite/scripts/pr_test.sh"
|
||||
label: ":test_tube: LoRA Inference Tests"
|
||||
env:
|
||||
- TEST_TYPE=inference_lora
|
||||
agents:
|
||||
queue: "default"
|
||||
- path:
|
||||
- "scripts/lora_extraction/**"
|
||||
- "fastvideo/tests/lora_extraction/**"
|
||||
- "fastvideo/models/loader/**"
|
||||
- "fastvideo/training/training_utils.py"
|
||||
- "fastvideo/layers/lora/**"
|
||||
- "pyproject.toml"
|
||||
- "docker/Dockerfile"
|
||||
config:
|
||||
command: "timeout 90m .buildkite/scripts/pr_test.sh"
|
||||
label: ":test_tube: LoRA Extraction Tests"
|
||||
env:
|
||||
- TEST_TYPE=lora_extraction
|
||||
agents:
|
||||
queue: "default"
|
||||
- path:
|
||||
- "fastvideo/**"
|
||||
- "pyproject.toml"
|
||||
- "docker/Dockerfile"
|
||||
config:
|
||||
command: "timeout 25m .buildkite/scripts/pr_test.sh"
|
||||
label: ":test_tube: Training Tests"
|
||||
env:
|
||||
- TEST_TYPE=training
|
||||
agents:
|
||||
queue: "default"
|
||||
- path:
|
||||
- "fastvideo/training/*distillation_pipeline.py"
|
||||
- "pyproject.toml"
|
||||
- "docker/Dockerfile"
|
||||
config:
|
||||
command: "timeout 25m .buildkite/scripts/pr_test.sh"
|
||||
label: ":test_tube: Distillation DMD Tests"
|
||||
env:
|
||||
- TEST_TYPE=distillation_dmd
|
||||
agents:
|
||||
queue: "default"
|
||||
- path:
|
||||
- "fastvideo/training/*self_forcing_distillation_pipeline.py"
|
||||
- "fastvideo/tests/training/self-forcing/**"
|
||||
- "pyproject.toml"
|
||||
- "docker/Dockerfile"
|
||||
config:
|
||||
command: "timeout 30m .buildkite/scripts/pr_test.sh"
|
||||
label: ":test_tube: Self-Forcing Tests"
|
||||
env:
|
||||
- TEST_TYPE=self_forcing
|
||||
agents:
|
||||
queue: "default"
|
||||
- path:
|
||||
- "fastvideo/**"
|
||||
- "pyproject.toml"
|
||||
- "docker/Dockerfile"
|
||||
config:
|
||||
command: "timeout 25m .buildkite/scripts/pr_test.sh"
|
||||
label: ":test_tube: LoRA Training Tests"
|
||||
env:
|
||||
- TEST_TYPE=training_lora
|
||||
retry:
|
||||
automatic:
|
||||
- exit_status: 1
|
||||
limit: 2
|
||||
agents:
|
||||
queue: "default"
|
||||
- path:
|
||||
- "fastvideo/**"
|
||||
- "fastvideo-kernel/**"
|
||||
- "pyproject.toml"
|
||||
- "docker/Dockerfile"
|
||||
config:
|
||||
command: "timeout 15m .buildkite/scripts/pr_test.sh"
|
||||
label: ":test_tube: Training Tests VSA"
|
||||
env:
|
||||
- TEST_TYPE=training_vsa
|
||||
retry:
|
||||
automatic:
|
||||
- exit_status: 1
|
||||
limit: 2
|
||||
agents:
|
||||
queue: "default"
|
||||
- path:
|
||||
- "fastvideo-kernel/**"
|
||||
- "fastvideo/attention/backends/vmoba.py"
|
||||
- "pyproject.toml"
|
||||
- "docker/Dockerfile"
|
||||
config:
|
||||
command: "timeout 15m .buildkite/scripts/pr_test.sh"
|
||||
label: ":test_tube: Inference Tests VMoBA"
|
||||
env:
|
||||
- TEST_TYPE=inference_vmoba
|
||||
agents:
|
||||
queue: "default"
|
||||
- path:
|
||||
- "fastvideo/models/dits/**"
|
||||
- "fastvideo/pipelines/**"
|
||||
- "fastvideo/attention/**"
|
||||
- "fastvideo/layers/**"
|
||||
- "fastvideo/worker/**"
|
||||
- "fastvideo/entrypoints/**"
|
||||
- "fastvideo/performance/**"
|
||||
- "fastvideo/tests/performance/**"
|
||||
- ".buildkite/performance-benchmarks/**"
|
||||
- "pyproject.toml"
|
||||
- "docker/Dockerfile"
|
||||
config:
|
||||
command: "timeout 30m .buildkite/scripts/pr_test.sh"
|
||||
label: ":test_tube: Performance Tests"
|
||||
env:
|
||||
- TEST_TYPE=performance
|
||||
agents:
|
||||
queue: "default"
|
||||
- path:
|
||||
- "fastvideo/entrypoints/openai/**"
|
||||
- "fastvideo/entrypoints/cli/serve.py"
|
||||
- "fastvideo/tests/entrypoints/test_openai_api_integration.py"
|
||||
- "pyproject.toml"
|
||||
- "docker/Dockerfile"
|
||||
config:
|
||||
command: "timeout 30m .buildkite/scripts/pr_test.sh"
|
||||
label: ":test_tube: API Server Tests"
|
||||
env:
|
||||
- TEST_TYPE=api_server
|
||||
agents:
|
||||
queue: "default"
|
||||
- path:
|
||||
- "fastvideo/train/**"
|
||||
- "fastvideo/tests/train/models/**"
|
||||
- "fastvideo/tests/train/fixtures/**"
|
||||
- "fastvideo/models/dits/**"
|
||||
- "fastvideo/models/loader/**"
|
||||
- "pyproject.toml"
|
||||
- "docker/Dockerfile"
|
||||
config:
|
||||
command: "timeout 30m .buildkite/scripts/pr_test.sh"
|
||||
label: ":test_tube: Train Framework Tests"
|
||||
env:
|
||||
- TEST_TYPE=train_framework
|
||||
agents:
|
||||
queue: "default"
|
||||
- path:
|
||||
- "fastvideo/eval/**"
|
||||
- "fastvideo/tests/eval/**"
|
||||
- "pyproject.toml"
|
||||
- "docker/Dockerfile"
|
||||
config:
|
||||
command: "timeout 90m .buildkite/scripts/pr_test.sh"
|
||||
label: ":test_tube: Eval Metrics Tests"
|
||||
env:
|
||||
- TEST_TYPE=eval
|
||||
agents:
|
||||
queue: "default"
|
||||
- label: ":microscope: Unit Tests"
|
||||
key: "unit"
|
||||
if: |
|
||||
build.env("TEST_SCOPE") == "full" ||
|
||||
(build.env("TEST_SCOPE") == "merge" &&
|
||||
build.env("MERGE_TEST_PLAN") =~ /,unit,/) ||
|
||||
build.env("TEST_SCOPE") == "fastcheck" ||
|
||||
build.env("TEST_SCOPE") == null ||
|
||||
(build.env("TEST_SCOPE") == "direct" &&
|
||||
(build.env("TEST_TYPE") == "unit_test" || build.env("TEST_TYPE") == "unit_test_ci"))
|
||||
command: "/opt/fastvideo-ci-runner/run-unit"
|
||||
timeout_in_minutes: 90
|
||||
env:
|
||||
TEST_TYPE: "unit_test_ci"
|
||||
retry:
|
||||
automatic:
|
||||
- exit_status: 128
|
||||
limit: 3
|
||||
- exit_status: -1
|
||||
limit: 2
|
||||
agents:
|
||||
queue: "ci-runner"
|
||||
|
||||
- label: ":microscope: DreamVerse App Tests"
|
||||
key: "dreamverse"
|
||||
if: |
|
||||
build.env("TEST_SCOPE") == "full" ||
|
||||
(build.env("TEST_SCOPE") == "merge" &&
|
||||
build.env("MERGE_TEST_PLAN") =~ /,dreamverse,/) ||
|
||||
build.env("TEST_SCOPE") == "fastcheck" ||
|
||||
build.env("TEST_SCOPE") == null ||
|
||||
(build.env("TEST_SCOPE") == "direct" &&
|
||||
(build.env("TEST_TYPE") == "dreamverse_app" || build.env("TEST_TYPE") == "dreamverse_app_ci"))
|
||||
command: "/opt/fastvideo-ci-runner/run-ci"
|
||||
timeout_in_minutes: 90
|
||||
env:
|
||||
TEST_TYPE: "dreamverse_app_ci"
|
||||
retry:
|
||||
automatic:
|
||||
- exit_status: 128
|
||||
limit: 3
|
||||
- exit_status: -1
|
||||
limit: 2
|
||||
agents:
|
||||
queue: "ci-runner"
|
||||
|
||||
- label: ":test_tube: Golden-Gate Tests"
|
||||
key: "golden-gate"
|
||||
if: |
|
||||
build.env("TEST_SCOPE") == "full" ||
|
||||
(build.env("TEST_SCOPE") == "merge" &&
|
||||
build.env("MERGE_TEST_PLAN") =~ /,golden-gate,/) ||
|
||||
(build.env("TEST_SCOPE") == "direct" &&
|
||||
(build.env("TEST_TYPE") == "golden_gate" || build.env("TEST_TYPE") == "golden_gate_ci"))
|
||||
command: "/opt/fastvideo-ci-runner/run-ci"
|
||||
timeout_in_minutes: 90
|
||||
env:
|
||||
TEST_TYPE: "golden_gate_ci"
|
||||
retry:
|
||||
automatic:
|
||||
- exit_status: 128
|
||||
limit: 3
|
||||
- exit_status: -1
|
||||
limit: 2
|
||||
agents:
|
||||
queue: "ci-runner"
|
||||
|
||||
- label: ":bar_chart: SSIM Tests"
|
||||
key: "ssim"
|
||||
if: |
|
||||
build.env("TEST_SCOPE") == "full" ||
|
||||
build.env("TEST_SCOPE") == "scheduled" ||
|
||||
(build.env("TEST_SCOPE") == "merge" &&
|
||||
build.env("MERGE_TEST_PLAN") =~ /,ssim,/) ||
|
||||
(build.env("TEST_SCOPE") == "direct" &&
|
||||
(build.env("TEST_TYPE") == "ssim" || build.env("TEST_TYPE") == "ssim_ci"))
|
||||
command: "/opt/fastvideo-ci-runner/run-ci"
|
||||
timeout_in_minutes: 90
|
||||
concurrency: 1
|
||||
concurrency_group: "fastvideo/slinky/whole-tray"
|
||||
env:
|
||||
TEST_TYPE: "ssim_ci"
|
||||
retry:
|
||||
automatic:
|
||||
- exit_status: 128
|
||||
limit: 3
|
||||
- exit_status: -1
|
||||
limit: 2
|
||||
- exit_status: 1
|
||||
limit: 2
|
||||
agents:
|
||||
queue: "ci-runner"
|
||||
|
||||
- label: ":test_tube: LoRA Inference Tests"
|
||||
key: "lora-inference"
|
||||
if: |
|
||||
build.env("TEST_SCOPE") == "full" ||
|
||||
(build.env("TEST_SCOPE") == "merge" &&
|
||||
build.env("MERGE_TEST_PLAN") =~ /,lora-inference,/) ||
|
||||
(build.env("TEST_SCOPE") == "direct" &&
|
||||
(build.env("TEST_TYPE") == "inference_lora" || build.env("TEST_TYPE") == "inference_lora_ci"))
|
||||
command: "/opt/fastvideo-ci-runner/run-ci"
|
||||
timeout_in_minutes: 90
|
||||
env:
|
||||
TEST_TYPE: "inference_lora_ci"
|
||||
retry:
|
||||
automatic:
|
||||
- exit_status: 128
|
||||
limit: 3
|
||||
- exit_status: -1
|
||||
limit: 2
|
||||
agents:
|
||||
queue: "ci-runner"
|
||||
|
||||
- label: ":test_tube: LoRA Extraction Tests"
|
||||
key: "lora-extraction"
|
||||
if: |
|
||||
build.env("TEST_SCOPE") == "full" ||
|
||||
(build.env("TEST_SCOPE") == "merge" &&
|
||||
build.env("MERGE_TEST_PLAN") =~ /,lora-extraction,/) ||
|
||||
(build.env("TEST_SCOPE") == "direct" &&
|
||||
(build.env("TEST_TYPE") == "lora_extraction" || build.env("TEST_TYPE") == "lora_extraction_ci"))
|
||||
command: "/opt/fastvideo-ci-runner/run-ci"
|
||||
timeout_in_minutes: 90
|
||||
env:
|
||||
TEST_TYPE: "lora_extraction_ci"
|
||||
retry:
|
||||
automatic:
|
||||
- exit_status: 128
|
||||
limit: 3
|
||||
- exit_status: -1
|
||||
limit: 2
|
||||
agents:
|
||||
queue: "ci-runner"
|
||||
|
||||
- label: ":test_tube: Training Tests"
|
||||
key: "training"
|
||||
if: |
|
||||
build.env("TEST_SCOPE") == "full" ||
|
||||
(build.env("TEST_SCOPE") == "merge" &&
|
||||
build.env("MERGE_TEST_PLAN") =~ /,training,/) ||
|
||||
(build.env("TEST_SCOPE") == "direct" &&
|
||||
(build.env("TEST_TYPE") == "training" || build.env("TEST_TYPE") == "training_ci"))
|
||||
command: "/opt/fastvideo-ci-runner/run-ci"
|
||||
timeout_in_minutes: 90
|
||||
concurrency: 1
|
||||
concurrency_group: "fastvideo/slinky/whole-tray"
|
||||
env:
|
||||
TEST_TYPE: "training_ci"
|
||||
retry:
|
||||
automatic:
|
||||
- exit_status: 128
|
||||
limit: 3
|
||||
- exit_status: -1
|
||||
limit: 2
|
||||
agents:
|
||||
queue: "ci-runner"
|
||||
|
||||
- label: ":test_tube: Distillation DMD Tests"
|
||||
key: "distillation"
|
||||
if: |
|
||||
build.env("TEST_SCOPE") == "full" ||
|
||||
(build.env("TEST_SCOPE") == "merge" &&
|
||||
build.env("MERGE_TEST_PLAN") =~ /,distillation,/) ||
|
||||
(build.env("TEST_SCOPE") == "direct" &&
|
||||
(build.env("TEST_TYPE") == "distillation_dmd" || build.env("TEST_TYPE") == "distillation_dmd_ci"))
|
||||
command: "/opt/fastvideo-ci-runner/run-ci"
|
||||
timeout_in_minutes: 90
|
||||
env:
|
||||
TEST_TYPE: "distillation_dmd_ci"
|
||||
retry:
|
||||
automatic:
|
||||
- exit_status: 128
|
||||
limit: 3
|
||||
- exit_status: -1
|
||||
limit: 2
|
||||
agents:
|
||||
queue: "ci-runner"
|
||||
|
||||
- label: ":test_tube: Self-Forcing Tests"
|
||||
key: "self-forcing"
|
||||
if: |
|
||||
build.env("TEST_SCOPE") == "full" ||
|
||||
(build.env("TEST_SCOPE") == "merge" &&
|
||||
build.env("MERGE_TEST_PLAN") =~ /,self-forcing,/) ||
|
||||
(build.env("TEST_SCOPE") == "direct" &&
|
||||
(build.env("TEST_TYPE") == "self_forcing" || build.env("TEST_TYPE") == "self_forcing_ci"))
|
||||
command: "/opt/fastvideo-ci-runner/run-ci"
|
||||
timeout_in_minutes: 90
|
||||
env:
|
||||
TEST_TYPE: "self_forcing_ci"
|
||||
retry:
|
||||
automatic:
|
||||
- exit_status: 128
|
||||
limit: 3
|
||||
- exit_status: -1
|
||||
limit: 2
|
||||
agents:
|
||||
queue: "ci-runner"
|
||||
|
||||
- label: ":test_tube: LoRA Training Tests"
|
||||
key: "lora-training"
|
||||
if: |
|
||||
build.env("TEST_SCOPE") == "full" ||
|
||||
(build.env("TEST_SCOPE") == "merge" &&
|
||||
build.env("MERGE_TEST_PLAN") =~ /,lora-training,/) ||
|
||||
(build.env("TEST_SCOPE") == "direct" &&
|
||||
(build.env("TEST_TYPE") == "training_lora" || build.env("TEST_TYPE") == "training_lora_ci"))
|
||||
command: "/opt/fastvideo-ci-runner/run-ci"
|
||||
timeout_in_minutes: 90
|
||||
env:
|
||||
TEST_TYPE: "training_lora_ci"
|
||||
retry:
|
||||
automatic:
|
||||
- exit_status: 128
|
||||
limit: 3
|
||||
- exit_status: -1
|
||||
limit: 2
|
||||
- exit_status: 1
|
||||
limit: 2
|
||||
agents:
|
||||
queue: "ci-runner"
|
||||
|
||||
- label: ":test_tube: Training Tests VSA"
|
||||
key: "training-vsa"
|
||||
if: |
|
||||
build.env("TEST_SCOPE") == "full" ||
|
||||
(build.env("TEST_SCOPE") == "merge" &&
|
||||
build.env("MERGE_TEST_PLAN") =~ /,training-vsa,/) ||
|
||||
(build.env("TEST_SCOPE") == "direct" &&
|
||||
(build.env("TEST_TYPE") == "training_vsa" || build.env("TEST_TYPE") == "training_vsa_ci"))
|
||||
command: "/opt/fastvideo-ci-runner/run-ci"
|
||||
timeout_in_minutes: 90
|
||||
env:
|
||||
TEST_TYPE: "training_vsa_ci"
|
||||
retry:
|
||||
automatic:
|
||||
- exit_status: 128
|
||||
limit: 3
|
||||
- exit_status: -1
|
||||
limit: 2
|
||||
- exit_status: 1
|
||||
limit: 2
|
||||
agents:
|
||||
queue: "ci-runner"
|
||||
|
||||
- label: ":test_tube: Inference Tests VMoBA"
|
||||
key: "inference-vmoba"
|
||||
if: |
|
||||
build.env("TEST_SCOPE") == "full" ||
|
||||
(build.env("TEST_SCOPE") == "merge" &&
|
||||
build.env("MERGE_TEST_PLAN") =~ /,inference-vmoba,/) ||
|
||||
(build.env("TEST_SCOPE") == "direct" &&
|
||||
(build.env("TEST_TYPE") == "inference_vmoba" || build.env("TEST_TYPE") == "inference_vmoba_ci"))
|
||||
command: "/opt/fastvideo-ci-runner/run-ci"
|
||||
timeout_in_minutes: 90
|
||||
env:
|
||||
TEST_TYPE: "inference_vmoba_ci"
|
||||
retry:
|
||||
automatic:
|
||||
- exit_status: 128
|
||||
limit: 3
|
||||
- exit_status: -1
|
||||
limit: 2
|
||||
agents:
|
||||
queue: "ci-runner"
|
||||
|
||||
- label: ":test_tube: Performance Tests"
|
||||
key: "performance"
|
||||
if: |
|
||||
build.env("TEST_SCOPE") == "full" ||
|
||||
(build.env("TEST_SCOPE") == "merge" &&
|
||||
build.env("MERGE_TEST_PLAN") =~ /,performance,/) ||
|
||||
(build.env("TEST_SCOPE") == "direct" &&
|
||||
(build.env("TEST_TYPE") == "performance" || build.env("TEST_TYPE") == "performance_ci"))
|
||||
command: "/opt/fastvideo-ci-runner/run-ci"
|
||||
timeout_in_minutes: 90
|
||||
env:
|
||||
TEST_TYPE: "performance_ci"
|
||||
retry:
|
||||
automatic:
|
||||
- exit_status: 128
|
||||
limit: 3
|
||||
- exit_status: -1
|
||||
limit: 2
|
||||
agents:
|
||||
queue: "ci-runner"
|
||||
|
||||
- label: ":test_tube: API Server Tests"
|
||||
key: "api-server"
|
||||
if: |
|
||||
build.env("TEST_SCOPE") == "full" ||
|
||||
(build.env("TEST_SCOPE") == "merge" &&
|
||||
build.env("MERGE_TEST_PLAN") =~ /,api-server,/) ||
|
||||
(build.env("TEST_SCOPE") == "direct" &&
|
||||
(build.env("TEST_TYPE") == "api_server" || build.env("TEST_TYPE") == "api_server_ci"))
|
||||
command: "/opt/fastvideo-ci-runner/run-ci"
|
||||
timeout_in_minutes: 90
|
||||
env:
|
||||
TEST_TYPE: "api_server_ci"
|
||||
retry:
|
||||
automatic:
|
||||
- exit_status: 128
|
||||
limit: 3
|
||||
- exit_status: -1
|
||||
limit: 2
|
||||
agents:
|
||||
queue: "ci-runner"
|
||||
|
||||
- label: ":test_tube: Train Framework Tests"
|
||||
key: "train-framework"
|
||||
if: |
|
||||
build.env("TEST_SCOPE") == "full" ||
|
||||
(build.env("TEST_SCOPE") == "merge" &&
|
||||
build.env("MERGE_TEST_PLAN") =~ /,train-framework,/) ||
|
||||
(build.env("TEST_SCOPE") == "direct" &&
|
||||
(build.env("TEST_TYPE") == "train_framework" || build.env("TEST_TYPE") == "train_framework_ci"))
|
||||
command: "/opt/fastvideo-ci-runner/run-ci"
|
||||
timeout_in_minutes: 90
|
||||
env:
|
||||
TEST_TYPE: "train_framework_ci"
|
||||
retry:
|
||||
automatic:
|
||||
- exit_status: 128
|
||||
limit: 3
|
||||
- exit_status: -1
|
||||
limit: 2
|
||||
agents:
|
||||
queue: "ci-runner"
|
||||
|
||||
- label: ":test_tube: Eval Metrics Tests"
|
||||
key: "eval"
|
||||
if: |
|
||||
build.env("TEST_SCOPE") == "full" ||
|
||||
(build.env("TEST_SCOPE") == "merge" &&
|
||||
build.env("MERGE_TEST_PLAN") =~ /,eval,/) ||
|
||||
(build.env("TEST_SCOPE") == "direct" &&
|
||||
(build.env("TEST_TYPE") == "eval" || build.env("TEST_TYPE") == "eval_ci"))
|
||||
command: "/opt/fastvideo-ci-runner/run-ci"
|
||||
timeout_in_minutes: 90
|
||||
env:
|
||||
TEST_TYPE: "eval_ci"
|
||||
retry:
|
||||
automatic:
|
||||
- exit_status: 128
|
||||
limit: 3
|
||||
- exit_status: -1
|
||||
limit: 2
|
||||
agents:
|
||||
queue: "ci-runner"
|
||||
|
||||
Executable
+5
@@ -0,0 +1,5 @@
|
||||
#!/usr/bin/env bash
|
||||
# Canonical Slurm CI selection for the OpenAI-compatible API lane.
|
||||
set -euo pipefail
|
||||
|
||||
exec pytest ./fastvideo/tests/entrypoints/test_openai_api_integration.py -vs
|
||||
Executable
+5
@@ -0,0 +1,5 @@
|
||||
#!/usr/bin/env bash
|
||||
# Canonical Slurm CI selection for the distillation-DMD lane.
|
||||
set -euo pipefail
|
||||
|
||||
exec pytest ./fastvideo/tests/training/distill/test_distill_dmd.py -vs
|
||||
Executable
+87
@@ -0,0 +1,87 @@
|
||||
#!/usr/bin/env bash
|
||||
# DreamVerse needs a GPU for import-time device resolution, but it does not
|
||||
# build or exercise fastvideo-kernel. A checksummed Node archive is installed
|
||||
# in the disposable Slurm container because the shared CI image is
|
||||
# Python/CUDA focused.
|
||||
set -euo pipefail
|
||||
|
||||
node_version=v22.23.2
|
||||
case $(uname -m) in
|
||||
aarch64 | arm64)
|
||||
node_arch=arm64
|
||||
node_archive_sha256=013b59cfd2819703a6f4a14ab891fc46fc2a4e3f5bcd92de3fb4929b43e35b30
|
||||
;;
|
||||
x86_64 | amd64)
|
||||
node_arch=x64
|
||||
node_archive_sha256=b294a556e639d64338823920e5866c21c02741742d2e1529ee1a225c1ec9252a
|
||||
;;
|
||||
*)
|
||||
echo "Unsupported architecture for DreamVerse Node runtime: $(uname -m)" >&2
|
||||
exit 2
|
||||
;;
|
||||
esac
|
||||
node_archive="node-${node_version}-linux-${node_arch}.tar.gz"
|
||||
node_runtime_root=$(mktemp -d -t fastvideo-node.XXXXXX)
|
||||
node_archive_path="${node_runtime_root}/${node_archive}"
|
||||
node_install_dir="${node_runtime_root}/${node_archive%.tar.gz}"
|
||||
curl --proto '=https' --tlsv1.2 --retry 5 --retry-all-errors \
|
||||
--location --fail --silent --show-error \
|
||||
"https://nodejs.org/dist/${node_version}/${node_archive}" \
|
||||
--output "$node_archive_path"
|
||||
printf '%s %s\n' "$node_archive_sha256" "$node_archive_path" | sha256sum --check --status
|
||||
tar -xzf "$node_archive_path" -C "$node_runtime_root"
|
||||
export PATH="${node_install_dir}/bin:${PATH}"
|
||||
node --version
|
||||
npm --version
|
||||
|
||||
export PYTHONPATH="$(pwd)/apps/dreamverse${PYTHONPATH:+:$PYTHONPATH}"
|
||||
pytest apps/dreamverse/dreamverse/tests -q
|
||||
|
||||
cd apps/dreamverse/web
|
||||
npm ci
|
||||
npm run typecheck
|
||||
npm test
|
||||
machine_arch=$(uname -m)
|
||||
if [[ $machine_arch =~ ^(aarch64|arm64)$ ]]; then
|
||||
npx playwright install --with-deps chromium firefox
|
||||
else
|
||||
npx playwright install --with-deps chromium webkit firefox
|
||||
fi
|
||||
|
||||
master_port=${MASTER_PORT:-7959}
|
||||
BACKEND_PORT=${BACKEND_PORT:-$((master_port + 50))}
|
||||
python -m uvicorn dreamverse.mock_server:app --host 127.0.0.1 --port "$BACKEND_PORT" &
|
||||
mock_server_pid=$!
|
||||
cleanup() {
|
||||
kill "$mock_server_pid" 2>/dev/null || true
|
||||
wait "$mock_server_pid" 2>/dev/null || true
|
||||
}
|
||||
trap cleanup EXIT INT TERM
|
||||
|
||||
for _ in {1..30}; do
|
||||
curl -fsS "http://127.0.0.1:$BACKEND_PORT/healthz" && break
|
||||
sleep 1
|
||||
done
|
||||
curl -fsS "http://127.0.0.1:$BACKEND_PORT/healthz"
|
||||
|
||||
if [[ $machine_arch =~ ^(aarch64|arm64)$ ]]; then
|
||||
# Playwright WebKit traps before opening a page on Linux ARM64, and its
|
||||
# bundled Chromium lacks the H.264/AAC codecs used by the fMP4 assertions.
|
||||
# Firefox covers every flow, including streaming. Chromium and its mobile
|
||||
# profile still cover all codec-independent UI behavior on GB200.
|
||||
BACKEND_HOST=127.0.0.1 BACKEND_PORT="$BACKEND_PORT" CI=1 \
|
||||
npm run e2e -- --project=firefox
|
||||
BACKEND_HOST=127.0.0.1 BACKEND_PORT="$BACKEND_PORT" CI=1 \
|
||||
npm run e2e -- \
|
||||
--project=chromium \
|
||||
--project=mobile-chromium \
|
||||
--grep-invert='streams, plays, and surfaces a downloadable clip|starts a new project and switches back to the prior session|saved projects persist across a page reload'
|
||||
else
|
||||
BACKEND_HOST=127.0.0.1 BACKEND_PORT="$BACKEND_PORT" CI=1 \
|
||||
npm run e2e -- \
|
||||
--project=chromium \
|
||||
--project=webkit \
|
||||
--project=firefox \
|
||||
--project=mobile-safari \
|
||||
--project=mobile-chromium
|
||||
fi
|
||||
Executable
+5
@@ -0,0 +1,5 @@
|
||||
#!/usr/bin/env bash
|
||||
# Canonical Slurm CI selection for the encoder lane.
|
||||
set -euo pipefail
|
||||
|
||||
exec pytest ./fastvideo/tests/encoders -vs
|
||||
Executable
+5
@@ -0,0 +1,5 @@
|
||||
#!/usr/bin/env bash
|
||||
# Canonical Slurm CI selection for the evaluation lane.
|
||||
set -euo pipefail
|
||||
|
||||
exec pytest ./fastvideo/tests/eval -vs
|
||||
Executable
+35
@@ -0,0 +1,35 @@
|
||||
#!/usr/bin/env bash
|
||||
# Canonical Slurm CI selection for the golden-gate lane. Environment (HF_HOME
|
||||
# and authentication) is the runner's responsibility.
|
||||
set -euo pipefail
|
||||
|
||||
golden_root=./fastvideo/tests/golden_gate
|
||||
selected=${FASTVIDEO_GOLDEN_TEST_FILES-}
|
||||
if [ -z "$selected" ]; then
|
||||
if [ "${TEST_SCOPE:-}" = merge ]; then
|
||||
echo "Missing FASTVIDEO_GOLDEN_TEST_FILES for merge scope" >&2
|
||||
exit 2
|
||||
fi
|
||||
selected=all
|
||||
fi
|
||||
if [ "$selected" = all ]; then
|
||||
exec pytest "$golden_root" -vs
|
||||
fi
|
||||
|
||||
[[ $selected =~ ^test_[a-z0-9_]+\.py(,test_[a-z0-9_]+\.py)*$ ]] || {
|
||||
echo "Invalid FASTVIDEO_GOLDEN_TEST_FILES selection" >&2
|
||||
exit 2
|
||||
}
|
||||
|
||||
IFS=, read -r -a golden_files <<< "$selected"
|
||||
golden_paths=()
|
||||
for golden_file in "${golden_files[@]}"; do
|
||||
golden_path="$golden_root/$golden_file"
|
||||
[ -f "$golden_path" ] || {
|
||||
echo "Selected golden test does not exist: $golden_file" >&2
|
||||
exit 2
|
||||
}
|
||||
golden_paths+=("$golden_path")
|
||||
done
|
||||
|
||||
exec pytest "${golden_paths[@]}" -vs
|
||||
Executable
+5
@@ -0,0 +1,5 @@
|
||||
#!/usr/bin/env bash
|
||||
# Canonical Slurm CI selection for the LoRA-inference lane.
|
||||
set -euo pipefail
|
||||
|
||||
exec pytest ./fastvideo/tests/inference/lora/test_lora_inference_similarity.py -vs
|
||||
Executable
+5
@@ -0,0 +1,5 @@
|
||||
#!/usr/bin/env bash
|
||||
# Canonical Slurm CI selection for the VMoBA-inference lane.
|
||||
set -euo pipefail
|
||||
|
||||
exec python fastvideo/tests/inference/vmoba/test_vmoba_inference.py
|
||||
Executable
+5
@@ -0,0 +1,5 @@
|
||||
#!/usr/bin/env bash
|
||||
# Canonical Slurm CI selection for the custom-kernel lane.
|
||||
set -euo pipefail
|
||||
|
||||
exec pytest fastvideo-kernel/tests/ -vs
|
||||
Executable
+5
@@ -0,0 +1,5 @@
|
||||
#!/usr/bin/env bash
|
||||
# Canonical Slurm CI selection for the LoRA-extraction lane.
|
||||
set -euo pipefail
|
||||
|
||||
exec pytest ./fastvideo/tests/lora_extraction/test_lora_extraction.py -vs
|
||||
Executable
+52
@@ -0,0 +1,52 @@
|
||||
#!/usr/bin/env bash
|
||||
# Canonical Slurm performance lane. Reports are written outside the checkout
|
||||
# so the trusted host driver can upload them after untrusted code exits.
|
||||
set -uo pipefail
|
||||
|
||||
export PERFORMANCE_TRACKING_ROOT=/tmp/perf-tracking
|
||||
export PERF_REPORTS_DIR=/workspace/artifacts/performance
|
||||
mkdir -p "$PERF_REPORTS_DIR"
|
||||
|
||||
if [[ ${BUILDKITE_PULL_REQUEST:-false} =~ ^[1-9][0-9]*$ ]]; then
|
||||
export PERF_RUN_SOURCE=pr
|
||||
export PERF_UPLOAD_POLICY=pass
|
||||
elif [ "${BUILDKITE_BRANCH:-}" = main ] \
|
||||
&& { [ "${BUILDKITE_SOURCE:-}" = schedule ] || [ "${TEST_SCOPE:-}" = full ]; }; then
|
||||
export PERF_RUN_SOURCE=scheduled_main
|
||||
export PERF_UPLOAD_POLICY=always
|
||||
elif [ "${TEST_SCOPE:-}" = direct ]; then
|
||||
export PERF_RUN_SOURCE=unknown
|
||||
export PERF_UPLOAD_POLICY=pass
|
||||
else
|
||||
export PERF_RUN_SOURCE=unknown
|
||||
export PERF_UPLOAD_POLICY=never
|
||||
fi
|
||||
|
||||
nvidia-smi \
|
||||
--query-gpu=index,timestamp,clocks.sm,clocks.max.sm,power.draw,power.limit,temperature.gpu \
|
||||
--format=csv -l 10 > "$PERF_REPORTS_DIR/gpu_telemetry.csv" 2>/dev/null &
|
||||
telemetry_pid=$!
|
||||
cleanup() {
|
||||
kill "$telemetry_pid" 2>/dev/null || true
|
||||
wait "$telemetry_pid" 2>/dev/null || true
|
||||
}
|
||||
trap cleanup EXIT INT TERM
|
||||
|
||||
pytest ./fastvideo/tests/performance -vs
|
||||
pytest_rc=$?
|
||||
compare_rc=0
|
||||
if [ "$pytest_rc" -eq 0 ] || [ "$PERF_UPLOAD_POLICY" = always ]; then
|
||||
PERF_PYTEST_RC=$pytest_rc python ./fastvideo/tests/performance/compare_baseline.py
|
||||
compare_rc=$?
|
||||
fi
|
||||
python ./fastvideo/tests/performance/dashboard.py || true
|
||||
cp -f fastvideo/tests/performance/results/*.json "$PERF_REPORTS_DIR/" 2>/dev/null || true
|
||||
|
||||
echo "--- GPU telemetry (clocks.sm vs clocks.max.sm reveals capped hosts) ---"
|
||||
cat "$PERF_REPORTS_DIR/gpu_telemetry.csv" || true
|
||||
|
||||
final_rc=$pytest_rc
|
||||
if [ "$final_rc" -eq 0 ]; then
|
||||
final_rc=$compare_rc
|
||||
fi
|
||||
exit "$final_rc"
|
||||
Executable
+6
@@ -0,0 +1,6 @@
|
||||
#!/usr/bin/env bash
|
||||
# Canonical Slurm CI selection for the self-forcing lane.
|
||||
set -euo pipefail
|
||||
|
||||
export WANDB_MODE=offline
|
||||
exec pytest ./fastvideo/tests/training/self-forcing/test_self_forcing.py -vs
|
||||
Executable
+40
@@ -0,0 +1,40 @@
|
||||
#!/usr/bin/env bash
|
||||
# Canonical four-GPU SSIM lane for the Slinky Slurm worker.
|
||||
set -euo pipefail
|
||||
|
||||
args=()
|
||||
if [ "${FASTVIDEO_SSIM_BOOTSTRAP_MODE:-0}" = 1 ]; then
|
||||
args+=(--bootstrap-mode)
|
||||
fi
|
||||
selected=${FASTVIDEO_SSIM_TEST_FILES-}
|
||||
if [ -z "$selected" ]; then
|
||||
if [ "${TEST_SCOPE:-}" = merge ]; then
|
||||
echo "Missing FASTVIDEO_SSIM_TEST_FILES for merge scope" >&2
|
||||
exit 2
|
||||
fi
|
||||
selected=all
|
||||
fi
|
||||
if [ "$selected" != all ]; then
|
||||
[[ $selected =~ ^test_[a-z0-9_]+\.py(,test_[a-z0-9_]+\.py)*$ ]] || {
|
||||
echo "Invalid FASTVIDEO_SSIM_TEST_FILES selection" >&2
|
||||
exit 2
|
||||
}
|
||||
IFS=, read -r -a ssim_files <<< "$selected"
|
||||
for ssim_file in "${ssim_files[@]}"; do
|
||||
args+=(--test-file "$ssim_file")
|
||||
done
|
||||
fi
|
||||
|
||||
# MoGe's utils3d dependency builds glcontext from source on ARM64. The current
|
||||
# runner image predates the baked-in X11 headers below, so keep this guarded
|
||||
# bootstrap until every deployed image digest contains libx11-dev.
|
||||
if [ ! -f /usr/include/X11/Xlib.h ]; then
|
||||
apt-get -o Acquire::Retries=5 update
|
||||
apt-get -o Acquire::Retries=5 install -y --no-install-recommends libx11-dev
|
||||
rm -rf /var/lib/apt/lists/*
|
||||
fi
|
||||
|
||||
uv pip install git+https://github.com/microsoft/MoGe.git
|
||||
uv pip install k_diffusion einops_exts alias_free_torch torchsde
|
||||
|
||||
exec python fastvideo/tests/ssim/ci_runner.py "${args[@]}"
|
||||
Executable
+5
@@ -0,0 +1,5 @@
|
||||
#!/usr/bin/env bash
|
||||
# Canonical Slurm CI selection for the modular training-framework lane.
|
||||
set -euo pipefail
|
||||
|
||||
exec pytest ./fastvideo/tests/train/models ./fastvideo/tests/train/methods -vs
|
||||
Executable
+6
@@ -0,0 +1,6 @@
|
||||
#!/usr/bin/env bash
|
||||
# Canonical Slurm CI selection for the legacy vanilla-training lane.
|
||||
set -euo pipefail
|
||||
|
||||
export WANDB_MODE=offline
|
||||
exec pytest ./fastvideo/tests/training/Vanilla -srP
|
||||
Executable
+6
@@ -0,0 +1,6 @@
|
||||
#!/usr/bin/env bash
|
||||
# Canonical Slurm CI selection for the legacy LoRA-training lane.
|
||||
set -euo pipefail
|
||||
|
||||
export WANDB_MODE=offline
|
||||
exec pytest ./fastvideo/tests/training/lora/test_lora_training.py -srP
|
||||
Executable
+6
@@ -0,0 +1,6 @@
|
||||
#!/usr/bin/env bash
|
||||
# Canonical Slurm CI selection for the legacy VSA-training lane.
|
||||
set -euo pipefail
|
||||
|
||||
export WANDB_MODE=offline
|
||||
exec pytest ./fastvideo/tests/training/VSA -srP
|
||||
Executable
+5
@@ -0,0 +1,5 @@
|
||||
#!/usr/bin/env bash
|
||||
# Canonical Slurm CI selection for the transformer lane.
|
||||
set -euo pipefail
|
||||
|
||||
exec pytest ./fastvideo/tests/transformers -vs
|
||||
Executable
+5
@@ -0,0 +1,5 @@
|
||||
#!/usr/bin/env bash
|
||||
# Canonical Slurm CI selection for the VAE lane.
|
||||
set -euo pipefail
|
||||
|
||||
exec pytest ./fastvideo/tests/vaes -vs
|
||||
@@ -1,6 +1,19 @@
|
||||
#!/bin/bash
|
||||
set -uo pipefail
|
||||
|
||||
# DORMANT ROLLBACK ONLY. Active CI is Slurm-only and pipeline.yml never calls
|
||||
# this launcher. Refuse every Buildkite invocation even if a stale step or
|
||||
# operator typo reaches this file; local rollback experiments require an
|
||||
# explicit opt-in.
|
||||
if [ -n "${BUILDKITE:-}" ]; then
|
||||
echo "Legacy Modal CI is disabled; use the Slinky Slurm runner." >&2
|
||||
exit 2
|
||||
fi
|
||||
if [ "${FASTVIDEO_ENABLE_LEGACY_MODAL_CI:-0}" != 1 ]; then
|
||||
echo "Legacy Modal CI is dormant. Set FASTVIDEO_ENABLE_LEGACY_MODAL_CI=1 only for a manual rollback test." >&2
|
||||
exit 2
|
||||
fi
|
||||
|
||||
log() {
|
||||
echo "[$(date '+%Y-%m-%d %H:%M:%S')] $1"
|
||||
}
|
||||
|
||||
@@ -0,0 +1,25 @@
|
||||
#!/usr/bin/env bash
|
||||
set -euo pipefail
|
||||
|
||||
exec pytest \
|
||||
./fastvideo/tests/api/ \
|
||||
./fastvideo/tests/contract/ \
|
||||
./fastvideo/tests/dataset/ \
|
||||
./fastvideo/tests/workflow/ \
|
||||
./fastvideo/tests/entrypoints/ \
|
||||
./fastvideo/tests/loader/ \
|
||||
./fastvideo/tests/pipelines/ \
|
||||
./fastvideo/tests/platforms/ \
|
||||
./fastvideo/tests/train/ \
|
||||
./fastvideo/tests/stages/ \
|
||||
./fastvideo/tests/ops/ \
|
||||
./fastvideo/tests/worker/ \
|
||||
./fastvideo/tests/training/test_trackers.py \
|
||||
./fastvideo/tests/attention/test_sdpa_metadata_mask_contract.py \
|
||||
./fastvideo/tests/modal/test_kernel_build_cache.py \
|
||||
./fastvideo/tests/modal/test_pr_test.py \
|
||||
./fastvideo/tests/modal/test_ssim_test.py \
|
||||
--ignore=./fastvideo/tests/entrypoints/test_openai_api_integration.py \
|
||||
--ignore=./fastvideo/tests/train/models \
|
||||
--ignore=./fastvideo/tests/train/methods \
|
||||
-vs
|
||||
@@ -8,10 +8,10 @@ PR TITLE: Must start with a type tag, e.g.:
|
||||
MERGE WORKFLOW:
|
||||
1. Ensure pre-commit passes and you have at least 1 approval
|
||||
2. Comment /merge (or add the "ready" label) to enter the Merge Queue
|
||||
3. Full Test Suite runs automatically on a staging branch → auto-merge on success
|
||||
3. A path-aware merge gate runs only relevant integration tests → auto-merge on success
|
||||
|
||||
ON-DEMAND TESTING (write access required):
|
||||
/test full — Full Test Suite /test ssim — SSIM regression
|
||||
/test full — Explicit all-lane run /test ssim — Full SSIM regression
|
||||
/test training — Training pipeline /test encoder — Encoder tests
|
||||
/test transformer — Transformer tests /test vae — VAE tests
|
||||
/test kernel — CUDA kernel tests /test unit — Unit tests
|
||||
|
||||
@@ -1,14 +1,14 @@
|
||||
#!/usr/bin/env bash
|
||||
# Gate the expensive Buildkite full suite on the cheap GitHub checks.
|
||||
# Gate the path-aware Buildkite merge plan 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.
|
||||
# 'ready' label cannot burn path-selected 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)
|
||||
# no merge gate; 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
|
||||
@@ -29,7 +29,7 @@ set -euo pipefail
|
||||
: "${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
|
||||
# Workflow-level `name:` values that must be green before the merge gate
|
||||
# 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"]'
|
||||
@@ -56,7 +56,7 @@ recheck_ready_label() {
|
||||
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."
|
||||
"NOT triggering the Buildkite merge gate. Re-add the label to re-arm."
|
||||
exit 1
|
||||
fi
|
||||
else
|
||||
@@ -84,7 +84,7 @@ while true; do
|
||||
| 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'" \
|
||||
"NOT triggering the Buildkite merge gate. 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
|
||||
@@ -97,7 +97,7 @@ while true; do
|
||||
if [ "$pending" -eq 0 ]; then
|
||||
if [ -z "$missing" ]; then
|
||||
recheck_ready_label
|
||||
echo "All watched cheap checks are green — full suite may proceed."
|
||||
echo "All watched cheap checks are green — merge gate may proceed."
|
||||
exit 0
|
||||
fi
|
||||
case "$missing" in
|
||||
@@ -119,14 +119,14 @@ while true; do
|
||||
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."
|
||||
echo "::warning::FAILING OPEN: cannot query GitHub check status — triggering the merge gate 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."
|
||||
echo "::warning::FAILING OPEN: watched checks still pending after $(( MAX_WAIT_SECS / 60 )) min${missing:+ (never appeared: ${missing})} — triggering the merge gate anyway."
|
||||
exit 0
|
||||
fi
|
||||
sleep "$POLL_SECS"
|
||||
|
||||
@@ -0,0 +1,570 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Select the additive GPU integration lanes needed by a PR diff.
|
||||
|
||||
Fastcheck is the universal six-lane baseline and is intentionally not repeated
|
||||
here. This planner selects only the more expensive merge-gate lanes. Unknown
|
||||
source/build paths fail closed to the complete integration set, while explicit
|
||||
documentation and repository-metadata paths require no additional GPU work.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import fnmatch
|
||||
import re
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
from typing import TextIO
|
||||
|
||||
MERGE_LANES = (
|
||||
"golden-gate",
|
||||
"ssim",
|
||||
"lora-inference",
|
||||
"lora-extraction",
|
||||
"training",
|
||||
"distillation",
|
||||
"self-forcing",
|
||||
"lora-training",
|
||||
"training-vsa",
|
||||
"inference-vmoba",
|
||||
"performance",
|
||||
"api-server",
|
||||
"train-framework",
|
||||
"eval",
|
||||
)
|
||||
|
||||
LANE_SCRIPT_TO_KEY = {
|
||||
"api_server.sh": "api-server",
|
||||
"distillation_dmd.sh": "distillation",
|
||||
"eval.sh": "eval",
|
||||
"golden_gate.sh": "golden-gate",
|
||||
"inference_lora.sh": "lora-inference",
|
||||
"inference_vmoba.sh": "inference-vmoba",
|
||||
"lora_extraction.sh": "lora-extraction",
|
||||
"performance.sh": "performance",
|
||||
"self_forcing.sh": "self-forcing",
|
||||
"ssim.sh": "ssim",
|
||||
"train_framework.sh": "train-framework",
|
||||
"training.sh": "training",
|
||||
"training_lora.sh": "lora-training",
|
||||
"training_vsa.sh": "training-vsa",
|
||||
}
|
||||
|
||||
FASTCHECK_LANE_SCRIPTS = {
|
||||
"dreamverse.sh",
|
||||
"encoder.sh",
|
||||
"kernel_tests.sh",
|
||||
"transformer.sh",
|
||||
"vae.sh",
|
||||
}
|
||||
|
||||
LEGACY_TRAINING_LANES = (
|
||||
"training",
|
||||
"distillation",
|
||||
"self-forcing",
|
||||
"lora-training",
|
||||
"training-vsa",
|
||||
)
|
||||
|
||||
ALL_TRAINING_LANES = (*LEGACY_TRAINING_LANES, "train-framework")
|
||||
|
||||
SSIM_SMOKE_TESTS = (
|
||||
"test_flux_t2i_similarity.py",
|
||||
"test_wan_t2v_similarity.py",
|
||||
)
|
||||
|
||||
SAFE_PATTERNS = (
|
||||
"*.md",
|
||||
"*.rst",
|
||||
".agents/**",
|
||||
".claude/**",
|
||||
".codex/**",
|
||||
".github/ISSUE_TEMPLATE/**",
|
||||
".github/PULL_REQUEST_TEMPLATE.md",
|
||||
".github/dependabot.yml",
|
||||
".github/mergify.yml",
|
||||
".github/scripts/**",
|
||||
".github/workflows/**",
|
||||
".buildkite/scripts/pre_commit.sh",
|
||||
".git-blame-ignore-revs",
|
||||
".gitattributes",
|
||||
".gitignore",
|
||||
".pre-commit-config.yaml",
|
||||
"AGENTS.md",
|
||||
"CITATION.cff",
|
||||
"CODE_OF_CONDUCT.md",
|
||||
"CONTRIBUTING.md",
|
||||
"LICENSE",
|
||||
"NOTICE",
|
||||
"__init__.py",
|
||||
"collect_env.py",
|
||||
"SECURITY.md",
|
||||
"assets/**",
|
||||
"comfyui/**",
|
||||
"docs/**",
|
||||
"examples/**",
|
||||
"mkdocs.yml",
|
||||
"requirements-mkdocs.in",
|
||||
"requirements-mkdocs.txt",
|
||||
"scripts/**",
|
||||
"tests/__init__.py",
|
||||
"tests/local_tests/**",
|
||||
)
|
||||
|
||||
ALL_IMPACT_PATTERNS = (
|
||||
".buildkite/pipeline.yml",
|
||||
"docker/**",
|
||||
"pyproject.toml",
|
||||
"requirements*.txt",
|
||||
"setup.cfg",
|
||||
"setup.py",
|
||||
"uv.lock",
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class FamilyCoverage:
|
||||
pattern: re.Pattern[str]
|
||||
golden_tests: tuple[str, ...]
|
||||
ssim_tests: tuple[str, ...]
|
||||
|
||||
|
||||
FAMILY_COVERAGE = (
|
||||
FamilyCoverage(
|
||||
re.compile(r"(^|[/_.-])dreamx(_world)?([/_.-]|$)"),
|
||||
("test_dreamx.py", ),
|
||||
("test_dreamx_world_similarity.py", ),
|
||||
),
|
||||
FamilyCoverage(
|
||||
re.compile(r"(^|[/_.-])flux[_-]?2([/_.-]|$)"),
|
||||
("test_flux2_klein.py", ),
|
||||
("test_flux2_similarity.py", ),
|
||||
),
|
||||
FamilyCoverage(
|
||||
re.compile(r"(^|[/_.-])flux(?![_-]?2)([/_.-]|$)"),
|
||||
("test_flux.py", ),
|
||||
("test_flux_t2i_similarity.py", ),
|
||||
),
|
||||
FamilyCoverage(
|
||||
re.compile(r"(^|[/_.-])(hunyuan)?gamecraft([/_.-]|$)"),
|
||||
("test_gamecraft.py", ),
|
||||
("test_gamecraft_similarity.py", ),
|
||||
),
|
||||
FamilyCoverage(
|
||||
re.compile(r"(^|[/_.-])gen3c([/_.-]|$)"),
|
||||
("test_gen3c.py", ),
|
||||
("test_gen3c_similarity.py", ),
|
||||
),
|
||||
FamilyCoverage(
|
||||
re.compile(r"(^|[/_.-])glm[_-]?image([/_.-]|$)"),
|
||||
("test_glm_image.py", ),
|
||||
("test_glm_image_similarity.py", ),
|
||||
),
|
||||
FamilyCoverage(
|
||||
re.compile(r"(^|[/_.-])kandinsky[_-]?5([/_.-]|$)"),
|
||||
("test_kandinsky5.py", ),
|
||||
("test_kandinsky5_similarity.py", ),
|
||||
),
|
||||
FamilyCoverage(
|
||||
re.compile(r"(^|[/_.-])lingbot([a-z0-9_-]*)([/_.-]|$)"),
|
||||
("test_lingbot.py", ),
|
||||
("test_lingbot_similarity.py", ),
|
||||
),
|
||||
FamilyCoverage(
|
||||
re.compile(r"(^|[/_.-])longcat([/_.-]|$)"),
|
||||
("test_longcat.py", ),
|
||||
("test_longcat_similarity.py", ),
|
||||
),
|
||||
FamilyCoverage(
|
||||
re.compile(r"(^|[/_.-])ltx[_-]?2([/_.-]|$)"),
|
||||
("test_ltx2.py", ),
|
||||
("test_ltx2_similarity.py", ),
|
||||
),
|
||||
FamilyCoverage(
|
||||
re.compile(r"(^|[/_.-])matrixgame[_-]?2([/_.-]|$)"),
|
||||
("test_matrixgame.py", ),
|
||||
("test_matrixgame2_similarity.py", ),
|
||||
),
|
||||
FamilyCoverage(
|
||||
re.compile(r"(^|[/_.-])matrixgame[_-]?3([/_.-]|$)"),
|
||||
("test_matrixgame.py", ),
|
||||
("test_matrixgame3_similarity.py", ),
|
||||
),
|
||||
FamilyCoverage(
|
||||
re.compile(r"(^|[/_.-])minimax[_-]?h3([/_.-]|$)"),
|
||||
("test_minimax_h3_t2v.py", ),
|
||||
("test_minimax_h3_similarity.py", ),
|
||||
),
|
||||
FamilyCoverage(
|
||||
re.compile(r"(^|[/_.-])sd[_-]?3([._-]?5)?([/_.-]|$)"),
|
||||
("test_sd35.py", ),
|
||||
("test_sd35_similarity.py", ),
|
||||
),
|
||||
FamilyCoverage(
|
||||
re.compile(r"(^|[/_.-])stable[_-]?audio([/_.-]|$)"),
|
||||
("test_stable_audio.py", ),
|
||||
("test_stable_audio_similarity.py", ),
|
||||
),
|
||||
FamilyCoverage(
|
||||
re.compile(r"(^|[/_.-])turbo(diffusion)?([/_.-]|$)"),
|
||||
(),
|
||||
("test_turbodiffusion_similarity.py", ),
|
||||
),
|
||||
FamilyCoverage(
|
||||
re.compile(r"(^|[/_.-])wan(video)?([/_.-]|$)"),
|
||||
("test_wan_t2v.py", ),
|
||||
(
|
||||
"test_causal_similarity.py",
|
||||
"test_wan_i2v_similarity.py",
|
||||
"test_wan_t2v_similarity.py",
|
||||
),
|
||||
),
|
||||
FamilyCoverage(
|
||||
re.compile(r"(^|[/_.-])z[_-]?image([/_.-]|$)"),
|
||||
("test_zimage.py", ),
|
||||
("test_zimage_similarity.py", ),
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class MergePlan:
|
||||
lanes: set[str] = field(default_factory=set)
|
||||
golden_tests: set[str] = field(default_factory=set)
|
||||
ssim_tests: set[str] = field(default_factory=set)
|
||||
golden_all: bool = False
|
||||
ssim_all: bool = False
|
||||
reasons: list[str] = field(default_factory=list)
|
||||
|
||||
def add_lanes(self, *lanes: str, reason: str) -> None:
|
||||
unknown = set(lanes) - set(MERGE_LANES)
|
||||
if unknown:
|
||||
raise ValueError(f"Unknown merge lanes: {sorted(unknown)}")
|
||||
self.lanes.update(lanes)
|
||||
self.reasons.append(reason)
|
||||
|
||||
def add_golden(self, tests: tuple[str, ...], reason: str) -> None:
|
||||
self.add_lanes("golden-gate", reason=reason)
|
||||
self.golden_tests.update(tests)
|
||||
|
||||
def add_ssim(self, tests: tuple[str, ...], reason: str) -> None:
|
||||
self.add_lanes("ssim", reason=reason)
|
||||
self.ssim_tests.update(tests)
|
||||
|
||||
def require_all(self, reason: str) -> None:
|
||||
self.lanes.update(MERGE_LANES)
|
||||
self.golden_all = True
|
||||
self.ssim_all = True
|
||||
self.reasons.append(reason)
|
||||
|
||||
def ordered_lanes(self) -> tuple[str, ...]:
|
||||
return tuple(lane for lane in MERGE_LANES if lane in self.lanes)
|
||||
|
||||
def encoded_lanes(self) -> str:
|
||||
lanes = self.ordered_lanes()
|
||||
return "," + ",".join(lanes or ("none", )) + ","
|
||||
|
||||
def encoded_golden_tests(self) -> str:
|
||||
if "golden-gate" not in self.lanes:
|
||||
return "none"
|
||||
if self.golden_all or not self.golden_tests:
|
||||
return "all"
|
||||
return ",".join(sorted(self.golden_tests))
|
||||
|
||||
def encoded_ssim_tests(self) -> str:
|
||||
if "ssim" not in self.lanes:
|
||||
return "none"
|
||||
if self.ssim_all or not self.ssim_tests:
|
||||
return "all"
|
||||
return ",".join(sorted(self.ssim_tests))
|
||||
|
||||
|
||||
def _matches_any(path: str, patterns: tuple[str, ...]) -> bool:
|
||||
return any(fnmatch.fnmatchcase(path, pattern) for pattern in patterns)
|
||||
|
||||
|
||||
def _family_coverage(path: str) -> tuple[set[str], set[str]]:
|
||||
normalized = path.lower()
|
||||
golden: set[str] = set()
|
||||
ssim: set[str] = set()
|
||||
for family in FAMILY_COVERAGE:
|
||||
if family.pattern.search(normalized):
|
||||
golden.update(family.golden_tests)
|
||||
ssim.update(family.ssim_tests)
|
||||
return golden, ssim
|
||||
|
||||
|
||||
def _select_output_coverage(plan: MergePlan, path: str) -> None:
|
||||
golden, ssim = _family_coverage(path)
|
||||
if golden:
|
||||
plan.add_golden(tuple(sorted(golden)), reason=f"model-family golden coverage: {path}")
|
||||
else:
|
||||
plan.golden_all = True
|
||||
plan.add_lanes("golden-gate", reason=f"shared output golden coverage: {path}")
|
||||
if ssim:
|
||||
plan.add_ssim(tuple(sorted(ssim)), reason=f"model-family SSIM coverage: {path}")
|
||||
else:
|
||||
plan.add_ssim(SSIM_SMOKE_TESTS, reason=f"shared output SSIM smoke coverage: {path}")
|
||||
|
||||
|
||||
def classify_paths(paths: list[str]) -> MergePlan:
|
||||
plan = MergePlan()
|
||||
normalized_paths: list[str] = []
|
||||
for raw_path in paths:
|
||||
path = raw_path.strip()
|
||||
while path.startswith("./"):
|
||||
path = path[2:]
|
||||
if path:
|
||||
normalized_paths.append(path)
|
||||
normalized_paths = sorted(set(normalized_paths))
|
||||
if not normalized_paths:
|
||||
plan.require_all("changed-file list was empty; failing closed")
|
||||
return plan
|
||||
|
||||
for path in normalized_paths:
|
||||
if path == "__FASTVIDEO_CI_PLAN_ALL__":
|
||||
plan.require_all("changed-file API failed; failing closed")
|
||||
continue
|
||||
|
||||
if path in {"requirements-mkdocs.in", "requirements-mkdocs.txt"}:
|
||||
plan.reasons.append(f"documentation dependencies need no GPU integration: {path}")
|
||||
continue
|
||||
|
||||
if _matches_any(path, ALL_IMPACT_PATTERNS):
|
||||
plan.require_all(f"cross-cutting build/runtime surface: {path}")
|
||||
continue
|
||||
|
||||
lane_script_prefix = ".buildkite/scripts/lanes/"
|
||||
if path.startswith(lane_script_prefix):
|
||||
script_name = Path(path).name
|
||||
lane = LANE_SCRIPT_TO_KEY.get(script_name)
|
||||
if lane is None:
|
||||
if script_name in FASTCHECK_LANE_SCRIPTS:
|
||||
plan.reasons.append(f"covered by automatic Fastcheck lane: {path}")
|
||||
else:
|
||||
plan.require_all(f"unknown lane script: {path}")
|
||||
elif lane == "golden-gate":
|
||||
plan.golden_all = True
|
||||
plan.add_lanes(lane, reason=f"golden lane implementation: {path}")
|
||||
elif lane == "ssim":
|
||||
plan.ssim_all = True
|
||||
plan.add_lanes(lane, reason=f"SSIM lane implementation: {path}")
|
||||
else:
|
||||
plan.add_lanes(lane, reason=f"lane implementation: {path}")
|
||||
continue
|
||||
|
||||
if path.startswith("fastvideo/tests/golden_gate/"):
|
||||
name = Path(path).name
|
||||
if name.startswith("test_") and name.endswith(".py"):
|
||||
plan.add_golden((name, ), reason=f"changed golden test: {path}")
|
||||
elif name in {"AGENTS.md", "README.md"}:
|
||||
plan.reasons.append(f"golden documentation only: {path}")
|
||||
else:
|
||||
plan.golden_all = True
|
||||
plan.add_lanes("golden-gate", reason=f"shared golden harness/reference: {path}")
|
||||
continue
|
||||
|
||||
if path.startswith("fastvideo/tests/ssim/"):
|
||||
name = Path(path).name
|
||||
if name.startswith("test_") and name.endswith(".py"):
|
||||
plan.add_ssim((name, ), reason=f"changed SSIM test: {path}")
|
||||
elif path.endswith((".py", ".json", ".pt", ".png", ".mp4")):
|
||||
plan.ssim_all = True
|
||||
plan.add_lanes("ssim", reason=f"shared SSIM harness/reference: {path}")
|
||||
continue
|
||||
|
||||
if path.startswith("fastvideo/tests/performance/") or path.startswith(".buildkite/performance-benchmarks/"):
|
||||
plan.add_lanes("performance", reason=f"performance coverage: {path}")
|
||||
continue
|
||||
if path.startswith(("fastvideo/performance/", "fastvideo/performance_dashboard/",
|
||||
"apps/performance_dashboard/")):
|
||||
plan.add_lanes("performance", reason=f"performance implementation: {path}")
|
||||
continue
|
||||
if path.startswith("fastvideo/benchmarks/"):
|
||||
if "/mlx_" in path or Path(path).name.startswith("mlx_"):
|
||||
plan.reasons.append(f"covered by the path-filtered macOS MLX workflow: {path}")
|
||||
else:
|
||||
plan.add_lanes("performance", reason=f"benchmark implementation: {path}")
|
||||
continue
|
||||
if path.startswith("fastvideo/tests/eval/") or path.startswith("fastvideo/eval/"):
|
||||
plan.add_lanes("eval", reason=f"evaluation coverage: {path}")
|
||||
continue
|
||||
if path.startswith("fastvideo/third_party/eval/"):
|
||||
plan.add_lanes("eval", reason=f"vendored evaluation implementation: {path}")
|
||||
continue
|
||||
if path.startswith("fastvideo/tests/lora_extraction/") or path.startswith("scripts/lora_extraction/"):
|
||||
plan.add_lanes("lora-extraction", reason=f"LoRA extraction coverage: {path}")
|
||||
continue
|
||||
if path.startswith("fastvideo/tests/inference/lora/"):
|
||||
plan.add_lanes("lora-inference", reason=f"LoRA inference coverage: {path}")
|
||||
continue
|
||||
if path.startswith("fastvideo/tests/inference/vmoba/"):
|
||||
plan.add_lanes("inference-vmoba", reason=f"VMoBA inference coverage: {path}")
|
||||
continue
|
||||
if path.startswith(("fastvideo/dataset/", "fastvideo/workflow/", "fastvideo/pipelines/preprocess/",
|
||||
"fastvideo/pipelines/training/")):
|
||||
plan.add_lanes(*ALL_TRAINING_LANES, reason=f"shared data/training input surface: {path}")
|
||||
continue
|
||||
if path.startswith("fastvideo/tests/train/") or path.startswith("fastvideo/train/"):
|
||||
plan.add_lanes("train-framework", reason=f"modular training coverage: {path}")
|
||||
continue
|
||||
|
||||
if path.startswith("fastvideo/tests/training/"):
|
||||
lowered = path.lower()
|
||||
if "/vanilla/" in lowered:
|
||||
plan.add_lanes("training", reason=f"vanilla training coverage: {path}")
|
||||
elif "/distill/" in lowered:
|
||||
plan.add_lanes("distillation", reason=f"distillation coverage: {path}")
|
||||
elif "/self-forcing/" in lowered:
|
||||
plan.add_lanes("self-forcing", reason=f"self-forcing coverage: {path}")
|
||||
elif "/lora/" in lowered:
|
||||
plan.add_lanes("lora-training", reason=f"LoRA training coverage: {path}")
|
||||
elif "/vsa/" in lowered:
|
||||
plan.add_lanes("training-vsa", reason=f"VSA training coverage: {path}")
|
||||
else:
|
||||
plan.add_lanes(*LEGACY_TRAINING_LANES, reason=f"shared legacy training coverage: {path}")
|
||||
continue
|
||||
|
||||
if path.startswith("fastvideo/training/"):
|
||||
lowered = path.lower()
|
||||
if "self_forcing" in lowered:
|
||||
plan.add_lanes("self-forcing", reason=f"self-forcing implementation: {path}")
|
||||
elif "distill" in lowered:
|
||||
plan.add_lanes("distillation", reason=f"distillation implementation: {path}")
|
||||
elif "lora" in lowered:
|
||||
plan.add_lanes("lora-training", reason=f"LoRA training implementation: {path}")
|
||||
else:
|
||||
plan.add_lanes(*LEGACY_TRAINING_LANES, reason=f"shared legacy training implementation: {path}")
|
||||
continue
|
||||
|
||||
lowered = path.lower()
|
||||
if "vmoba" in lowered and path.startswith(("fastvideo/", ".buildkite/")):
|
||||
plan.add_lanes("inference-vmoba", reason=f"VMoBA implementation: {path}")
|
||||
plan.add_golden(("test_wan_t2v.py", ), reason=f"VMoBA end-to-end coverage: {path}")
|
||||
continue
|
||||
if "lora" in lowered and path.startswith("fastvideo/"):
|
||||
plan.add_lanes(
|
||||
"lora-inference",
|
||||
"lora-extraction",
|
||||
"lora-training",
|
||||
reason=f"shared LoRA implementation: {path}",
|
||||
)
|
||||
_select_output_coverage(plan, path)
|
||||
continue
|
||||
|
||||
if path.startswith("fastvideo/entrypoints/") or path.startswith("fastvideo/api/"):
|
||||
plan.add_lanes("api-server", reason=f"API/entrypoint integration: {path}")
|
||||
if "openai" not in lowered and "/cli/" not in lowered:
|
||||
_select_output_coverage(plan, path)
|
||||
continue
|
||||
if path.startswith("fastvideo/worker/"):
|
||||
plan.add_lanes("api-server", reason=f"worker/API integration: {path}")
|
||||
_select_output_coverage(plan, path)
|
||||
continue
|
||||
if path.startswith("fastvideo/distributed/"):
|
||||
plan.add_lanes(
|
||||
"training",
|
||||
"train-framework",
|
||||
reason=f"distributed runtime integration: {path}",
|
||||
)
|
||||
_select_output_coverage(plan, path)
|
||||
continue
|
||||
if path.startswith(("fastvideo/hooks/", "fastvideo/platforms/", "fastvideo/third_party/")):
|
||||
_select_output_coverage(plan, path)
|
||||
continue
|
||||
if path.startswith(("fastvideo/models/", "fastvideo/pipelines/", "fastvideo/configs/",
|
||||
"fastvideo/layers/", "fastvideo/attention/")):
|
||||
_select_output_coverage(plan, path)
|
||||
continue
|
||||
if path in {
|
||||
"fastvideo/fastvideo_args.py",
|
||||
"fastvideo/forward_context.py",
|
||||
"fastvideo/image_processor.py",
|
||||
"fastvideo/registry.py",
|
||||
"fastvideo/utils.py",
|
||||
}:
|
||||
_select_output_coverage(plan, path)
|
||||
continue
|
||||
if path.startswith("fastvideo/mlx_runtime/"):
|
||||
plan.reasons.append(f"covered by the path-filtered macOS MLX workflow: {path}")
|
||||
continue
|
||||
if path.startswith("fastvideo/logging_utils/") or path in {
|
||||
"fastvideo/__init__.py",
|
||||
"fastvideo/envs.py",
|
||||
"fastvideo/logger.py",
|
||||
"fastvideo/profiler.py",
|
||||
"fastvideo/version.py",
|
||||
}:
|
||||
plan.reasons.append(f"covered by automatic Fastcheck: {path}")
|
||||
continue
|
||||
if path.startswith(("fastvideo-kernel/", "csrc/")):
|
||||
plan.add_golden(("test_wan_t2v.py", ), reason=f"kernel integration smoke: {path}")
|
||||
plan.add_ssim(("test_wan_t2v_similarity.py", ), reason=f"kernel numerical smoke: {path}")
|
||||
continue
|
||||
|
||||
if path.startswith("apps/dreamverse/"):
|
||||
# DreamVerse is already one of the six automatic Fastcheck lanes.
|
||||
plan.reasons.append(f"covered by automatic DreamVerse Fastcheck: {path}")
|
||||
continue
|
||||
if path.startswith("fastvideo/tests/"):
|
||||
# The automatic unit/component Fastcheck lanes own the remaining
|
||||
# package tests. Domain-specific expensive test roots were handled
|
||||
# above.
|
||||
plan.reasons.append(f"covered by automatic Fastcheck: {path}")
|
||||
continue
|
||||
if path in {".buildkite/scripts/unit_test.sh", ".buildkite/scripts/pr_test.sh"}:
|
||||
plan.reasons.append(f"covered by automatic unit Fastcheck: {path}")
|
||||
continue
|
||||
if _matches_any(path, SAFE_PATTERNS):
|
||||
plan.reasons.append(f"no additional GPU integration needed: {path}")
|
||||
continue
|
||||
|
||||
plan.require_all(f"unclassified path; failing closed: {path}")
|
||||
|
||||
return plan
|
||||
|
||||
|
||||
def _write_github_output(output: TextIO, plan: MergePlan) -> None:
|
||||
output.write(f"merge_test_plan={plan.encoded_lanes()}\n")
|
||||
output.write(f"merge_golden_tests={plan.encoded_golden_tests()}\n")
|
||||
output.write(f"merge_ssim_tests={plan.encoded_ssim_tests()}\n")
|
||||
output.write(f"merge_plan_label={','.join(plan.ordered_lanes()) or 'none'}\n")
|
||||
|
||||
|
||||
def _write_summary(output: TextIO, plan: MergePlan) -> None:
|
||||
output.write("## Change-aware merge test plan\n\n")
|
||||
output.write("| Selection | Value |\n|---|---|\n")
|
||||
output.write(f"| Additional Slurm lanes | `{','.join(plan.ordered_lanes()) or 'none'}` |\n")
|
||||
output.write(f"| Golden tests | `{plan.encoded_golden_tests()}` |\n")
|
||||
output.write(f"| SSIM tests | `{plan.encoded_ssim_tests()}` |\n\n")
|
||||
output.write("Fastcheck remains the universal six-lane baseline.\n")
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("--paths-file", type=Path, required=True)
|
||||
parser.add_argument("--github-output", type=Path)
|
||||
parser.add_argument("--summary-file", type=Path)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main() -> int:
|
||||
args = parse_args()
|
||||
paths = args.paths_file.read_text(encoding="utf-8").splitlines()
|
||||
plan = classify_paths(paths)
|
||||
print(f"MERGE_TEST_PLAN={plan.encoded_lanes()}")
|
||||
print(f"MERGE_GOLDEN_TESTS={plan.encoded_golden_tests()}")
|
||||
print(f"MERGE_SSIM_TESTS={plan.encoded_ssim_tests()}")
|
||||
for reason in plan.reasons:
|
||||
print(f"- {reason}")
|
||||
if args.github_output:
|
||||
with args.github_output.open("a", encoding="utf-8") as output:
|
||||
_write_github_output(output, plan)
|
||||
if args.summary_file:
|
||||
with args.summary_file.open("a", encoding="utf-8") as output:
|
||||
_write_summary(output, plan)
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
@@ -53,7 +53,7 @@ PC_PENDING='{"name": "pre-commit", "id": 1, "status": "in_progress", "conclusion
|
||||
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}'
|
||||
OTHER='{"name": "Trigger Merge Gate", "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"}'
|
||||
|
||||
|
||||
@@ -190,6 +190,7 @@ jobs:
|
||||
if: ${{ !inputs.push_by_digest }}
|
||||
run: |
|
||||
echo "✅ Python ${{ inputs.python_version }} image successfully built and pushed to ${{ steps.image.outputs.name }}:${{ inputs.tag_suffix }}-sha-${GITHUB_SHA::7}"
|
||||
echo "Digest: ${{ steps.build-push.outputs.digest }}"
|
||||
echo "To run tests with this image, manually trigger the 'Run Tests' workflow."
|
||||
|
||||
- name: Digest success message
|
||||
|
||||
@@ -26,29 +26,48 @@ jobs:
|
||||
per_page: 100,
|
||||
});
|
||||
|
||||
const bkStatuses = data.statuses.filter(
|
||||
s => s.context.startsWith('buildkite/ci/')
|
||||
);
|
||||
|
||||
const FASTCHECK_PREFIX = 'buildkite/ci/microscope-';
|
||||
// Buildkite derives the GitHub context prefix from the label emoji.
|
||||
// Keep hard Full Suite lanes in test-tube/bar-chart namespaces and
|
||||
// Fastcheck lanes in microscope so targeted reruns cannot clear the
|
||||
// wrong aggregate status. Automatic PR jobs use pr-fastcheck while
|
||||
// slash-command and Full Suite jobs use ci; normalize the suffix
|
||||
// and keep the newest status for each logical lane.
|
||||
const FASTCHECK_PREFIXES = [
|
||||
'buildkite/pr-fastcheck/microscope-',
|
||||
'buildkite/ci/microscope-',
|
||||
];
|
||||
const FULL_SUITE_PREFIXES = [
|
||||
'buildkite/ci/test-tube-',
|
||||
'buildkite/ci/bar-chart-',
|
||||
];
|
||||
|
||||
const fastcheck = bkStatuses.filter(
|
||||
s => s.context.startsWith(FASTCHECK_PREFIX)
|
||||
);
|
||||
const fullSuite = bkStatuses.filter(
|
||||
s => FULL_SUITE_PREFIXES.some(p => s.context.startsWith(p))
|
||||
);
|
||||
function newestByLane(prefixes) {
|
||||
const statuses = new Map();
|
||||
for (const status of data.statuses) {
|
||||
const prefix = prefixes.find(p => status.context.startsWith(p));
|
||||
if (!prefix) continue;
|
||||
const lane = status.context.slice(prefix.length);
|
||||
const previous = statuses.get(lane);
|
||||
if (!previous || Date.parse(status.updated_at) > Date.parse(previous.updated_at)) {
|
||||
statuses.set(lane, status);
|
||||
}
|
||||
}
|
||||
return statuses;
|
||||
}
|
||||
|
||||
if (
|
||||
fastcheck.length > 0
|
||||
&& fastcheck.every(s => s.state === 'success')
|
||||
) {
|
||||
const fastcheck = newestByLane(FASTCHECK_PREFIXES);
|
||||
const fullSuiteOnly = newestByLane(FULL_SUITE_PREFIXES);
|
||||
const fastcheckPassed =
|
||||
fastcheck.size === 6
|
||||
&& [...fastcheck.values()].every(s => s.state === 'success');
|
||||
const fullSuitePassed =
|
||||
fastcheckPassed
|
||||
&& fullSuiteOnly.size === 14
|
||||
&& [...fullSuiteOnly.values()].every(s => s.state === 'success');
|
||||
|
||||
if (fastcheckPassed) {
|
||||
core.info(
|
||||
`All ${fastcheck.length} fastcheck tests passed — updating fastcheck-passed`
|
||||
`All ${fastcheck.size} fastcheck tests passed — updating fastcheck-passed`
|
||||
);
|
||||
await github.rest.repos.createCommitStatus({
|
||||
owner: context.repo.owner,
|
||||
@@ -56,17 +75,13 @@ jobs:
|
||||
sha,
|
||||
state: 'success',
|
||||
context: 'fastcheck-passed',
|
||||
description:
|
||||
`All ${fastcheck.length} fastcheck tests passed`,
|
||||
description: `All ${fastcheck.size} fastcheck tests passed`,
|
||||
});
|
||||
}
|
||||
|
||||
if (
|
||||
fullSuite.length > 0
|
||||
&& fullSuite.every(s => s.state === 'success')
|
||||
) {
|
||||
if (fullSuitePassed) {
|
||||
core.info(
|
||||
`All ${fullSuite.length} full suite tests passed — updating full-suite-passed`
|
||||
'All 20 full suite tests passed — updating full-suite-passed'
|
||||
);
|
||||
await github.rest.repos.createCommitStatus({
|
||||
owner: context.repo.owner,
|
||||
@@ -74,7 +89,6 @@ jobs:
|
||||
sha,
|
||||
state: 'success',
|
||||
context: 'full-suite-passed',
|
||||
description:
|
||||
`All ${fullSuite.length} full suite tests passed`,
|
||||
description: 'All 20 full suite tests passed',
|
||||
});
|
||||
}
|
||||
|
||||
@@ -78,6 +78,7 @@ jobs:
|
||||
fastvideo/tests/mlx/test_mlx_dit_parity.py \
|
||||
fastvideo/tests/mlx/test_mlx_compile_parity.py \
|
||||
fastvideo/tests/mlx/test_mlx_checkpoint.py \
|
||||
fastvideo/tests/mlx/test_mlx_checkpoint_compat.py \
|
||||
fastvideo/tests/mlx/test_mlx_fastwan_benchmark.py \
|
||||
fastvideo/tests/mlx/test_taehv_decode.py \
|
||||
fastvideo/tests/mlx/test_frame_upsample.py \
|
||||
@@ -134,6 +135,7 @@ jobs:
|
||||
fastvideo/tests/mlx/test_mlx_dit_parity.py \
|
||||
fastvideo/tests/mlx/test_mlx_compile_parity.py \
|
||||
fastvideo/tests/mlx/test_mlx_checkpoint.py \
|
||||
fastvideo/tests/mlx/test_mlx_checkpoint_compat.py \
|
||||
fastvideo/tests/mlx/test_mlx_fastwan_benchmark.py \
|
||||
fastvideo/tests/mlx/test_taehv_decode.py \
|
||||
fastvideo/tests/mlx/test_frame_upsample.py \
|
||||
|
||||
@@ -0,0 +1,44 @@
|
||||
name: Scheduled Full SSIM
|
||||
|
||||
on:
|
||||
schedule:
|
||||
- cron: "0 5 * * 0"
|
||||
workflow_dispatch:
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
|
||||
jobs:
|
||||
trigger:
|
||||
if: github.repository == 'hao-ai-lab/FastVideo'
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- name: Trigger weekly full SSIM on Slinky Slurm
|
||||
env:
|
||||
BUILDKITE_API_TOKEN: ${{ secrets.BUILDKITE_API_TOKEN }}
|
||||
SOURCE_SHA: ${{ github.sha }}
|
||||
SOURCE_BRANCH: ${{ github.event.repository.default_branch }}
|
||||
BK_ORG: ${{ vars.BUILDKITE_ORG_SLUG }}
|
||||
BK_PIPELINE: ${{ vars.BUILDKITE_PIPELINE_SLUG }}
|
||||
run: |
|
||||
set -euo pipefail
|
||||
curl -sS --fail-with-body -X POST \
|
||||
"https://api.buildkite.com/v2/organizations/${BK_ORG}/pipelines/${BK_PIPELINE}/builds" \
|
||||
-H "Authorization: Bearer $BUILDKITE_API_TOKEN" \
|
||||
-H "Content-Type: application/json" \
|
||||
--data-raw "$(jq -n \
|
||||
--arg commit "$SOURCE_SHA" \
|
||||
--arg branch "$SOURCE_BRANCH" \
|
||||
'{
|
||||
commit: $commit,
|
||||
branch: $branch,
|
||||
message: "Weekly full SSIM on Slinky Slurm",
|
||||
ignore_pipeline_branch_filters: true,
|
||||
env: {
|
||||
TEST_SCOPE: "scheduled",
|
||||
FULL_SUITE: "false",
|
||||
TEST_TYPE: "ssim",
|
||||
PR_NUMBER: "false",
|
||||
PR_TITLE: "Scheduled full SSIM"
|
||||
}
|
||||
}')"
|
||||
@@ -33,7 +33,6 @@ jobs:
|
||||
core.setOutput('has_write', String(hasWrite));
|
||||
|
||||
- name: Add ready label and react
|
||||
id: label
|
||||
if: steps.perm.outputs.has_write == 'true'
|
||||
uses: actions/github-script@60a0d83039c74a4aee543508d2ffcb1c3799cdea # v7.0.1
|
||||
with:
|
||||
@@ -48,47 +47,6 @@ jobs:
|
||||
comment_id: context.payload.comment.id,
|
||||
content: 'rocket',
|
||||
});
|
||||
const { data: pr } = await github.rest.pulls.get({ owner, repo, pull_number: prNumber });
|
||||
core.setOutput('pr_sha', pr.head.sha);
|
||||
core.setOutput('pr_branch', pr.head.ref);
|
||||
core.setOutput('pr_number', String(prNumber));
|
||||
core.setOutput('pr_title', pr.title);
|
||||
|
||||
- name: Trigger Full Suite
|
||||
if: steps.perm.outputs.has_write == 'true'
|
||||
env:
|
||||
BUILDKITE_API_TOKEN: ${{ secrets.BUILDKITE_API_TOKEN }}
|
||||
PR_SHA: ${{ steps.label.outputs.pr_sha }}
|
||||
PR_BRANCH: ${{ steps.label.outputs.pr_branch }}
|
||||
PR_NUMBER: ${{ steps.label.outputs.pr_number }}
|
||||
PR_TITLE: ${{ steps.label.outputs.pr_title }}
|
||||
BK_ORG: ${{ vars.BUILDKITE_ORG_SLUG }}
|
||||
BK_PIPELINE: ${{ vars.BUILDKITE_PIPELINE_SLUG }}
|
||||
run: |
|
||||
curl -sS --fail-with-body -X POST \
|
||||
"https://api.buildkite.com/v2/organizations/${BK_ORG}/pipelines/${BK_PIPELINE}/builds" \
|
||||
-H "Authorization: Bearer $BUILDKITE_API_TOKEN" \
|
||||
-H "Content-Type: application/json" \
|
||||
--data-raw "$(jq -n \
|
||||
--arg commit "$PR_SHA" \
|
||||
--arg branch "$PR_BRANCH" \
|
||||
--arg message "Full Suite for PR #${PR_NUMBER} (via /merge)" \
|
||||
--arg pr_title "$PR_TITLE" \
|
||||
--argjson pr_id "$PR_NUMBER" \
|
||||
'{
|
||||
commit: $commit,
|
||||
branch: $branch,
|
||||
message: $message,
|
||||
ignore_pipeline_branch_filters: true,
|
||||
pull_request_id: $pr_id,
|
||||
pull_request_base_branch: "main",
|
||||
env: {
|
||||
TEST_SCOPE: "full",
|
||||
FULL_SUITE: "true",
|
||||
PR_NUMBER: ($pr_id | tostring),
|
||||
PR_TITLE: $pr_title
|
||||
}
|
||||
}')"
|
||||
|
||||
parse-command:
|
||||
if: >-
|
||||
@@ -129,7 +87,7 @@ jobs:
|
||||
set -euo pipefail
|
||||
TEST_NAME=$(echo "$COMMENT" | grep -oP '(?<=/test\s)\S+' | head -1 || true)
|
||||
|
||||
VALID="encoder vae transformer kernel unit dreamverse ssim golden-gate training lora-inference lora-training lora-extraction distillation self-forcing vsa vmoba performance api train-framework eval full fastcheck pre-commit"
|
||||
VALID="encoder vae transformer kernel unit dreamverse ssim golden-gate training lora-inference lora-training lora-extraction distillation self-forcing vsa vmoba performance api train-framework eval unit-ci kernel-ci dreamverse-ci ssim-ci golden-gate-ci encoder-ci vae-ci transformer-ci lora-inference-ci lora-training-ci lora-extraction-ci training-ci distillation-ci self-forcing-ci vsa-ci vmoba-ci performance-ci api-ci train-framework-ci eval-ci full fastcheck pre-commit"
|
||||
if [ -z "$TEST_NAME" ] || ! echo "$VALID" | grep -qw "$TEST_NAME"; then
|
||||
echo "Unknown test: '$TEST_NAME'. Valid: $VALID"
|
||||
exit 1
|
||||
@@ -137,7 +95,17 @@ jobs:
|
||||
|
||||
declare -A MAP=(
|
||||
[encoder]=encoder [vae]=vae [transformer]=transformer
|
||||
[kernel]=kernel_tests [unit]=unit_test [dreamverse]=dreamverse_app
|
||||
[kernel]=kernel_tests [unit]=unit_test [unit-ci]=unit_test_ci
|
||||
[kernel-ci]=kernel_tests_ci [dreamverse-ci]=dreamverse_app_ci
|
||||
[ssim-ci]=ssim_ci [vmoba-ci]=inference_vmoba_ci
|
||||
[golden-gate-ci]=golden_gate_ci [training-ci]=training_ci
|
||||
[encoder-ci]=encoder_ci [vae-ci]=vae_ci [transformer-ci]=transformer_ci
|
||||
[lora-inference-ci]=inference_lora_ci [lora-training-ci]=training_lora_ci
|
||||
[lora-extraction-ci]=lora_extraction_ci [distillation-ci]=distillation_dmd_ci
|
||||
[self-forcing-ci]=self_forcing_ci [vsa-ci]=training_vsa_ci
|
||||
[performance-ci]=performance_ci [api-ci]=api_server_ci
|
||||
[train-framework-ci]=train_framework_ci [eval-ci]=eval_ci
|
||||
[dreamverse]=dreamverse_app
|
||||
[ssim]=ssim [golden-gate]=golden_gate [training]=training
|
||||
[lora-inference]=inference_lora [lora-training]=training_lora
|
||||
[lora-extraction]=lora_extraction
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
name: Trigger Full Suite
|
||||
name: Trigger Merge Gate
|
||||
|
||||
on:
|
||||
pull_request_target:
|
||||
@@ -10,7 +10,7 @@ permissions:
|
||||
actions: read
|
||||
|
||||
concurrency:
|
||||
group: full-suite-${{ github.event.pull_request.number }}
|
||||
group: merge-gate-${{ github.event.pull_request.number }}
|
||||
cancel-in-progress: false
|
||||
|
||||
jobs:
|
||||
@@ -34,29 +34,72 @@ jobs:
|
||||
});
|
||||
const hasReady = pr.labels.some(l => l.name === 'ready');
|
||||
core.setOutput('has_ready', String(hasReady));
|
||||
if (!hasReady) core.info('No ready label — skipping Full Suite trigger.');
|
||||
core.setOutput('changed_files', String(pr.changed_files));
|
||||
if (!hasReady) core.info('No ready label — skipping merge-gate trigger.');
|
||||
|
||||
- name: Cancel previous Buildkite builds
|
||||
if: steps.check.outputs.has_ready == 'true'
|
||||
env:
|
||||
BUILDKITE_API_TOKEN: ${{ secrets.BUILDKITE_API_TOKEN }}
|
||||
PR_BRANCH: ${{ github.event.pull_request.head.ref }}
|
||||
PR_NUMBER: ${{ github.event.pull_request.number }}
|
||||
run: |
|
||||
# Find running builds for this branch with TEST_SCOPE=full and cancel them
|
||||
builds=$(curl -sS -H "Authorization: Bearer $BUILDKITE_API_TOKEN" \
|
||||
"https://api.buildkite.com/v2/organizations/${{ vars.BUILDKITE_ORG_SLUG }}/pipelines/${{ vars.BUILDKITE_PIPELINE_SLUG }}/builds?branch=${PR_BRANCH}&state=running,scheduled" \
|
||||
| jq -r '.[] | select(try (.env.TEST_SCOPE == "full") catch false) | .number')
|
||||
# Match both branch and PR number: forks can reuse the same branch name.
|
||||
builds=$(curl -sS --get -H "Authorization: Bearer $BUILDKITE_API_TOKEN" \
|
||||
--data-urlencode "branch=$PR_BRANCH" \
|
||||
--data-urlencode "state=running,scheduled" \
|
||||
"https://api.buildkite.com/v2/organizations/${{ vars.BUILDKITE_ORG_SLUG }}/pipelines/${{ vars.BUILDKITE_PIPELINE_SLUG }}/builds" \
|
||||
| jq -r --arg pr_number "$PR_NUMBER" \
|
||||
'.[] | select((.env.TEST_SCOPE? == "merge") and (.env.PR_NUMBER? == $pr_number)) | .number')
|
||||
for build_num in $builds; do
|
||||
echo "Cancelling Buildkite build #$build_num"
|
||||
curl -sS -X PUT -H "Authorization: Bearer $BUILDKITE_API_TOKEN" \
|
||||
"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
|
||||
# Check out the immutable BASE SHA: pull_request_target must never run a
|
||||
# planner or gate script from the untrusted PR head.
|
||||
- name: Checkout trusted merge planner
|
||||
if: steps.check.outputs.has_ready == 'true'
|
||||
uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
|
||||
with:
|
||||
ref: ${{ github.event.pull_request.base.sha }}
|
||||
persist-credentials: false
|
||||
|
||||
- name: Collect changed paths
|
||||
if: steps.check.outputs.has_ready == 'true'
|
||||
env:
|
||||
GH_TOKEN: ${{ github.token }}
|
||||
PR_NUMBER: ${{ github.event.pull_request.number }}
|
||||
EXPECTED_CHANGED_FILES: ${{ steps.check.outputs.changed_files }}
|
||||
run: |
|
||||
set -euo pipefail
|
||||
changed_json="$RUNNER_TEMP/merge-changed-files.json"
|
||||
changed_paths="$RUNNER_TEMP/merge-changed-paths.txt"
|
||||
if gh api --paginate --slurp \
|
||||
"repos/${GITHUB_REPOSITORY}/pulls/${PR_NUMBER}/files?per_page=100" \
|
||||
> "$changed_json"; then
|
||||
observed=$(jq '[.[][] | .filename] | unique | length' "$changed_json")
|
||||
if [ "$observed" = "$EXPECTED_CHANGED_FILES" ]; then
|
||||
jq -r '.[][] | .filename, (.previous_filename // empty)' "$changed_json" \
|
||||
| sort -u > "$changed_paths"
|
||||
else
|
||||
echo "::warning::Changed-file API returned $observed of $EXPECTED_CHANGED_FILES paths; selecting all merge lanes."
|
||||
echo '__FASTVIDEO_CI_PLAN_ALL__' > "$changed_paths"
|
||||
fi
|
||||
else
|
||||
echo "::warning::Changed-file API failed; selecting all merge lanes."
|
||||
echo '__FASTVIDEO_CI_PLAN_ALL__' > "$changed_paths"
|
||||
fi
|
||||
|
||||
- name: Select minimal merge tests
|
||||
id: plan
|
||||
if: steps.check.outputs.has_ready == 'true'
|
||||
run: |
|
||||
python3 .github/scripts/plan_merge_ci.py \
|
||||
--paths-file "$RUNNER_TEMP/merge-changed-paths.txt" \
|
||||
--github-output "$GITHUB_OUTPUT" \
|
||||
--summary-file "$GITHUB_STEP_SUMMARY"
|
||||
|
||||
- name: Wait for pre-commit and docs build
|
||||
if: steps.check.outputs.has_ready == 'true'
|
||||
@@ -66,7 +109,7 @@ jobs:
|
||||
PR_NUMBER: ${{ github.event.pull_request.number }}
|
||||
run: bash .github/scripts/gate_full_suite.sh
|
||||
|
||||
- name: Trigger Buildkite Full Suite
|
||||
- name: Trigger Buildkite merge gate
|
||||
if: steps.check.outputs.has_ready == 'true'
|
||||
env:
|
||||
BUILDKITE_API_TOKEN: ${{ secrets.BUILDKITE_API_TOKEN }}
|
||||
@@ -76,6 +119,10 @@ jobs:
|
||||
PR_TITLE: ${{ github.event.pull_request.title }}
|
||||
BK_ORG: ${{ vars.BUILDKITE_ORG_SLUG }}
|
||||
BK_PIPELINE: ${{ vars.BUILDKITE_PIPELINE_SLUG }}
|
||||
MERGE_TEST_PLAN: ${{ steps.plan.outputs.merge_test_plan }}
|
||||
MERGE_GOLDEN_TESTS: ${{ steps.plan.outputs.merge_golden_tests }}
|
||||
MERGE_SSIM_TESTS: ${{ steps.plan.outputs.merge_ssim_tests }}
|
||||
MERGE_PLAN_LABEL: ${{ steps.plan.outputs.merge_plan_label }}
|
||||
run: |
|
||||
curl -sS --fail-with-body -X POST \
|
||||
"https://api.buildkite.com/v2/organizations/${BK_ORG}/pipelines/${BK_PIPELINE}/builds" \
|
||||
@@ -84,8 +131,11 @@ jobs:
|
||||
--data-raw "$(jq -n \
|
||||
--arg commit "$PR_SHA" \
|
||||
--arg branch "$PR_BRANCH" \
|
||||
--arg message "Full Suite for PR #${PR_NUMBER}" \
|
||||
--arg message "Merge gate [${MERGE_PLAN_LABEL}] for PR #${PR_NUMBER}" \
|
||||
--arg pr_title "$PR_TITLE" \
|
||||
--arg merge_test_plan "$MERGE_TEST_PLAN" \
|
||||
--arg merge_golden_tests "$MERGE_GOLDEN_TESTS" \
|
||||
--arg merge_ssim_tests "$MERGE_SSIM_TESTS" \
|
||||
--argjson pr_id "$PR_NUMBER" \
|
||||
'{
|
||||
commit: $commit,
|
||||
@@ -95,8 +145,11 @@ jobs:
|
||||
pull_request_id: $pr_id,
|
||||
pull_request_base_branch: "main",
|
||||
env: {
|
||||
TEST_SCOPE: "full",
|
||||
TEST_SCOPE: "merge",
|
||||
FULL_SUITE: "true",
|
||||
MERGE_TEST_PLAN: $merge_test_plan,
|
||||
MERGE_GOLDEN_TESTS: $merge_golden_tests,
|
||||
MERGE_SSIM_TESTS: $merge_ssim_tests,
|
||||
PR_NUMBER: ($pr_id | tostring),
|
||||
PR_TITLE: $pr_title
|
||||
}
|
||||
|
||||
@@ -38,17 +38,17 @@ jobs:
|
||||
|
||||
**How our CI works:**
|
||||
|
||||
PRs run a two-tier CI system:
|
||||
PRs run a three-tier CI system:
|
||||
1. **Pre-commit** — formatting (yapf), linting (ruff), type checking (mypy). Runs immediately on every PR.
|
||||
2. **Fastcheck** — core GPU tests (encoders, VAEs, transformers, kernels, unit tests). Runs automatically via Buildkite on relevant file changes (~10-15 min).
|
||||
3. **Full Suite** — integration tests, training pipelines, SSIM regression. Runs only when a reviewer adds the `ready` label.
|
||||
2. **Fastcheck** — six core GPU lanes run automatically via Buildkite (~10-15 min).
|
||||
3. **Merge gate** — a reviewer adds `ready`; changed paths select only the relevant integration, training, golden, or SSIM coverage.
|
||||
|
||||
**Before your PR is reviewed:**
|
||||
- [ ] `pre-commit run --all-files` passes locally
|
||||
- [ ] You've added or updated tests for your changes
|
||||
- [ ] The PR description explains what and why
|
||||
|
||||
If pre-commit fails, a bot comment will explain how to fix it. Fastcheck and Full Suite results appear in the Checks section below.
|
||||
If pre-commit fails, a bot comment will explain how to fix it. Fastcheck and merge-gate results appear in the Checks section below.
|
||||
|
||||
**Useful links:**
|
||||
- [Contributing Guide](https://hao-ai-lab.github.io/FastVideo/contributing/overview/)
|
||||
|
||||
@@ -13,6 +13,11 @@ on:
|
||||
required: false
|
||||
default: false
|
||||
type: boolean
|
||||
build_ci_runner_image:
|
||||
description: 'Build the ARM64 CUDA 13 CI runner image (sm_100)'
|
||||
required: false
|
||||
default: false
|
||||
type: boolean
|
||||
# Auto-rebuild the CUDA images when a repository-controlled image input
|
||||
# changes on main. This includes the trusted SM89 kernel artifact's source,
|
||||
# metadata/key helper, ABI dependency metadata, and build orchestration.
|
||||
@@ -198,6 +203,28 @@ jobs:
|
||||
docker buildx imagetools create "${TAG_ARGS[@]}" "${IMAGE_REFS[@]}"
|
||||
docker buildx imagetools inspect "${TAGS[0]}"
|
||||
|
||||
# The CI runner is ARM64 like DGX Spark, but targets sm_100 rather than sm_121.
|
||||
# Publish a single-architecture variant so the self-hosted CI runner can reuse
|
||||
# the exact prebuilt kernel instead of compiling it in every job.
|
||||
build-ci-runner-image:
|
||||
if: ${{ (github.event_name == 'push' && github.repository == 'hao-ai-lab/FastVideo') || github.event.inputs.build_ci_runner_image == 'true' }}
|
||||
uses: ./.github/workflows/_template-build-image.yml
|
||||
with:
|
||||
python_version: '3.12'
|
||||
dockerfile_path: docker/Dockerfile
|
||||
tag_suffix: py3.12-cuda13.0.0-sm100
|
||||
runner: ubuntu-24.04-arm
|
||||
architecture: arm64
|
||||
build_args: |
|
||||
PYTHON_VERSION=3.12
|
||||
CUDA_VERSION=13.0.0
|
||||
UV_TORCH_BACKEND=cu130
|
||||
TORCH_CUDA_ARCH_LIST=10.0
|
||||
CMAKE_BUILD_PARALLEL_LEVEL=1
|
||||
FLASH_ATTN_WHEEL_TAG=cu130torch2.12
|
||||
FLASH_ATTN_WHEEL_RELEASE_ARM64=https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.9.22
|
||||
secrets: inherit
|
||||
|
||||
# Dreamverse matrix: {backend, UI} x {12.6.3, 13.0.0}, Python 3.12. Torch backend
|
||||
# matches the base CUDA (cu126 / cu130). Keep these images amd64-only until the
|
||||
# required FA4 dependency stack is available and validated on arm64.
|
||||
|
||||
@@ -62,8 +62,9 @@ jobs:
|
||||
cuda-version: '13.0.0'
|
||||
torch-cuda-short: 'cu130'
|
||||
platform:
|
||||
# x86_64 builds the full cu126 + cu130 set (cu130 ships the consumer
|
||||
# Blackwell sm_120a FP4 kernels).
|
||||
# x86_64 builds the full cu126 + cu130 set. cu130 ships the
|
||||
# data-center Blackwell sm_100a VSA and consumer sm_120a FP4
|
||||
# kernels.
|
||||
- os: ubuntu-22.04
|
||||
arch: x86_64
|
||||
wheel-plat: manylinux_2_35_x86_64
|
||||
@@ -124,7 +125,7 @@ jobs:
|
||||
- name: Install dependencies (GCC, Clang, CUDA Paths, Git)
|
||||
run: |
|
||||
sudo apt update
|
||||
sudo apt install -y git patchelf gcc-11 g++-11 clang-11
|
||||
sudo apt install -y git gcc-11 g++-11 clang-11
|
||||
sudo update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-11 100 --slave /usr/bin/g++ g++ /usr/bin/g++-11
|
||||
|
||||
# Allow Git to Access Safe Directory
|
||||
@@ -168,7 +169,8 @@ jobs:
|
||||
# covers sm_120a; turbodiffusion covers sm_100a+sm_120a. The sm_100 FP4
|
||||
# forward is the FA4 CuTe DSL path in the fastvideo package (PR #1221),
|
||||
# JIT-compiled at runtime — not built into this wheel.
|
||||
# * x86_64 cu130 = Hopper TK + consumer Blackwell sm_120a FP4.
|
||||
# * x86_64 cu130 = Hopper TK + data-center Blackwell sm_100a VSA
|
||||
# + consumer Blackwell sm_120a FP4.
|
||||
# * x86_64 cu126 = Hopper TK only (older drivers; CUDA < 12.8 has no FP4).
|
||||
# The per-arch split in CMakeLists pins the FP4 targets to sm_120a and builds
|
||||
# the main extension for the full arch list. CMAKE_BUILD_PARALLEL_LEVEL caps
|
||||
@@ -178,7 +180,7 @@ jobs:
|
||||
export CMAKE_ARGS="${CMAKE_ARGS:-} -DFASTVIDEO_KERNEL_BUILD_TK=OFF -DFASTVIDEO_KERNEL_BUILD_ATTN_QAT_INFER=ON"
|
||||
export CMAKE_BUILD_PARALLEL_LEVEL=1
|
||||
elif [ "${{ matrix.torch-cuda.torch-cuda-short }}" = "cu130" ]; then
|
||||
export TORCH_CUDA_ARCH_LIST="9.0a;12.0a"
|
||||
export TORCH_CUDA_ARCH_LIST="9.0a;10.0a;12.0a"
|
||||
export CMAKE_ARGS="${CMAKE_ARGS:-} -DFASTVIDEO_KERNEL_BUILD_TK=ON -DFASTVIDEO_KERNEL_BUILD_ATTN_QAT_INFER=ON -DCMAKE_CUDA_ARCHITECTURES=90a"
|
||||
# A single FP4 TU (attn_qat_infer) can use ~8-12 GB on its own, so serialize.
|
||||
export CMAKE_BUILD_PARALLEL_LEVEL=1
|
||||
@@ -194,7 +196,11 @@ jobs:
|
||||
python -m build --wheel --outdir dist
|
||||
|
||||
# Fix the wheel to be manylinux compliant
|
||||
uv pip install --system auditwheel
|
||||
# Ubuntu 22.04 ships patchelf 0.14.3, while current auditwheel
|
||||
# requires at least 0.14.5. Use the stable PyPI binary on both
|
||||
# x86_64 and aarch64 release runners.
|
||||
uv pip install --system auditwheel patchelf==0.17.2.4
|
||||
patchelf --version
|
||||
# Point auditwheel at torch libs, but do not vendor them into the wheel.
|
||||
TORCH_LIB_DIR=$(python - <<'PY'
|
||||
import os
|
||||
@@ -211,7 +217,8 @@ jobs:
|
||||
--exclude libtorch.so \
|
||||
--exclude libc10.so \
|
||||
--exclude libc10_cuda.so \
|
||||
--exclude libtorch_python.so
|
||||
--exclude libtorch_python.so \
|
||||
--exclude libnccl.so.2
|
||||
# Move fixed wheels back to dist for upload consistency
|
||||
rm dist/*.whl
|
||||
mv fixed_dist/*.whl dist/
|
||||
|
||||
@@ -9,6 +9,7 @@
|
||||
**FastVideo is a unified post-training and real-time inference framework for accelerated video generation.**
|
||||
|
||||
## NEWS
|
||||
- `2026/08/23`: [FastH3 Preview v0.2](https://huggingface.co/FastVideo/FastVideo-Minimax-FastH3-Preview-v0.2) is a 4-step DMD2-distilled MiniMax-H3 checkpoint that generates synchronized video and audio. Run the verified [basic FastH3 example](examples/inference/basic/basic_fasth3.py), see the [inference guide](examples/inference/basic/README.md#fasth3-preview), or have a coding agent install FastVideo with the [agent setup prompt](#install-with-an-ai-coding-agent).
|
||||
- `2026/08/19`: FastVideo now supports MLX on Apple Silicon with [FastMetal-QAD](https://huggingface.co/collections/FastVideo/fastmetal), a family of 1.3B, 5B, and 14B models optimized for Mac—follow the [Apple Silicon guide](https://hao-ai-lab.github.io/FastVideo/getting_started/installation/mps/) and read the [Blog](https://haoailab.com/blogs/fastmetal/).
|
||||
- `2026/06/23`: Release FastWan-QAD: 5s of Video generated in 1.8s E2E. See the [FastWan-QAD models](https://huggingface.co/FastVideo/FastWan-QAD-FP8-1.3B), [Attn-QAT training guide](https://haoailab.com/FastVideo/training/attn_qat/), and [blog](https://haoailab.com/blogs/fastwan-qad/).
|
||||
- `2026/03/17`: Release demo: Into the Dreamverse: Vibe Directing in FastVideo, check out the [Blog](https://haoailab.com/blogs/dreamverse/).
|
||||
@@ -63,9 +64,10 @@ UV_TORCH_BACKEND=cu126 uv pip install fastvideo
|
||||
Use `UV_TORCH_BACKEND=cu130` on CUDA 13. Apple silicon users should follow the
|
||||
[MPS installation guide](https://hao-ai-lab.github.io/FastVideo/getting_started/installation/mps/).
|
||||
|
||||
> **On an Apple Silicon Mac?** FastVideo runs FastWan text-to-video natively
|
||||
> through an MLX runtime — a 5-second 480p clip generated locally, no cloud,
|
||||
> no discrete GPU. Install with `uv pip install -e '.[mlx]'` and follow the
|
||||
> **On an Apple Silicon Mac?** FastVideo runs FastMetal-QAD through an MLX
|
||||
> runtime. Install with `uv pip install -e '.[mlx]'`, download
|
||||
> [`FastVideo/FastMetal-1.3B-QAD`](https://huggingface.co/FastVideo/FastMetal-1.3B-QAD),
|
||||
> and follow the
|
||||
> [Apple Silicon guide](https://hao-ai-lab.github.io/FastVideo/getting_started/installation/mps/).
|
||||
|
||||
Please see our [docs](https://hao-ai-lab.github.io/FastVideo/getting_started/installation/) for more detailed installation instructions.
|
||||
|
||||
+26
-12
@@ -89,11 +89,29 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
|
||||
zsh \
|
||||
vim \
|
||||
curl \
|
||||
ffmpeg \
|
||||
libgl1 \
|
||||
libglib2.0-0 \
|
||||
libx11-dev \
|
||||
gcc-11 \
|
||||
g++-11 \
|
||||
clang-11 \
|
||||
cmake \
|
||||
pkg-config \
|
||||
build-essential \
|
||||
libssl-dev \
|
||||
&& rm -rf /var/lib/apt/lists/*
|
||||
|
||||
# Rust toolchain: some dependencies only ship sdists on aarch64 and need cargo
|
||||
# to build. The dormant legacy Modal image layers the identical apt set +
|
||||
# rustup on top of this image (fastvideo/tests/modal/pr_test.py); baking both
|
||||
# here keeps its manual rollback path reproducible without changing the Slurm
|
||||
# runner's package surface.
|
||||
RUN set -o pipefail && \
|
||||
curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y --default-toolchain stable --profile minimal && \
|
||||
/root/.cargo/bin/cargo --version && /root/.cargo/bin/rustc --version
|
||||
ENV PATH=/root/.cargo/bin:${PATH}
|
||||
|
||||
# Set up C++20 compilers for ThunderKittens
|
||||
RUN update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-11 100 --slave /usr/bin/g++ g++ /usr/bin/g++-11
|
||||
|
||||
@@ -138,6 +156,7 @@ RUN --mount=type=cache,target=/opt/uv/cache \
|
||||
source /opt/venv/bin/activate && \
|
||||
uv pip install --upgrade pip && \
|
||||
uv pip install --excludes docker/uv-excludes ".[dev]" && \
|
||||
python -c "import cv2; print('OpenCV', cv2.__version__)" && \
|
||||
PYTAG=cp$(echo "${PYTHON_VERSION}" | tr -d .) && \
|
||||
case "${TARGETARCH:-amd64}" in \
|
||||
amd64) \
|
||||
@@ -169,26 +188,21 @@ RUN --mount=type=cache,target=/opt/uv/cache \
|
||||
# flash_attn/__init__.py), so FA2/varlen/bert_padding stay from the install above;
|
||||
# rmtree clears the wheel's stale cute files first to avoid an install conflict.
|
||||
# Then verify both survive so a broken overlay fails the build instead of shipping
|
||||
# an FA2-less image. x86 only: the FA4 stack (quack-kernels etc.) is unvalidated on
|
||||
# arm64 / GB10 (sm_121), so there we skip the overlay; FA4 is opt-in
|
||||
# (FASTVIDEO_FA4=1) and errors if set without the overlay, so leave it unset on
|
||||
# arm64 and the image runs FA3/FA2 as usual.
|
||||
# an FA2-less image. The pinned stack is validated on ARM64 GB200 (sm_100) as well
|
||||
# as x86; FA4 remains opt-in through FASTVIDEO_FA4=1 so lanes with FA2 baselines
|
||||
# keep their existing numerics.
|
||||
RUN --mount=type=cache,target=/opt/uv/cache \
|
||||
source $HOME/.local/bin/env && \
|
||||
source /opt/venv/bin/activate && \
|
||||
if [ "${TARGETARCH}" = "arm64" ]; then \
|
||||
echo "Skipping FA4 cute overlay on arm64 (FA4 stack unvalidated there; do not set FASTVIDEO_FA4)"; \
|
||||
else \
|
||||
python -c "import glob, shutil; [shutil.rmtree(d, ignore_errors=True) for d in glob.glob('/opt/venv/lib/python*/site-packages/flash_attn/cute')]" && \
|
||||
uv pip install "flash-attn-4 @ git+https://github.com/Dao-AILab/flash-attention.git@${FA4_CUTE_REF}#subdirectory=flash_attn/cute" && \
|
||||
python -c "import flash_attn; assert hasattr(flash_attn, 'flash_attn_func'), 'FA2 was clobbered by the cute overlay'; import flash_attn.cute; print('FA2 + FA4 cute OK')"; \
|
||||
fi
|
||||
python -c "import glob, shutil; [shutil.rmtree(d, ignore_errors=True) for d in glob.glob('/opt/venv/lib/python*/site-packages/flash_attn/cute')]" && \
|
||||
uv pip install "flash-attn-4 @ git+https://github.com/Dao-AILab/flash-attention.git@${FA4_CUTE_REF}#subdirectory=flash_attn/cute" && \
|
||||
python -c "import flash_attn; assert hasattr(flash_attn, 'flash_attn_func'), 'FA2 was clobbered by the cute overlay'; import flash_attn.cute; print('FA2 + FA4 cute OK')"
|
||||
|
||||
COPY . .
|
||||
|
||||
# Build immutable FastVideo kernel wheels for the published image. The requested
|
||||
# architecture remains installed for normal image users; amd64 images also carry
|
||||
# an SM89 artifact so the predominant L40S Modal lanes can reuse it exactly.
|
||||
# an SM89 artifact for L40S users and the dormant legacy rollback path.
|
||||
ARG FASTVIDEO_KERNEL_PREBUILT_DIR=/opt/fastvideo-kernel-prebuilt
|
||||
RUN --mount=type=cache,target=/opt/uv/cache \
|
||||
source $HOME/.local/bin/env && \
|
||||
|
||||
@@ -6,7 +6,8 @@ lives in [Testing](testing.md).
|
||||
|
||||
## Overview
|
||||
|
||||
FastVideo splits validation across GitHub Actions, Buildkite, and Modal:
|
||||
FastVideo splits validation across GitHub Actions, Buildkite, Slinky Slurm,
|
||||
and Mergify:
|
||||
|
||||
```text
|
||||
PR opened or updated
|
||||
@@ -16,17 +17,24 @@ PR opened or updated
|
||||
| style, lint, type, spelling, Markdown, workflow syntax, filenames
|
||||
|
|
||||
|-- Tier 2: Fastcheck
|
||||
| Buildkite orchestrates Modal GPU jobs
|
||||
| path-filtered component and unit checks
|
||||
| Buildkite schedules six lanes on the Slinky Slurm cluster
|
||||
| encoder, VAE, transformer, kernel, unit, DreamVerse
|
||||
|
|
||||
|-- /merge, /test full, or ready label
|
||||
|-- /merge or ready label
|
||||
|
|
||||
`-- Tier 3: Full Suite
|
||||
Buildkite orchestrates Modal GPU jobs
|
||||
path-filtered integration, SSIM, training, eval, and performance checks
|
||||
`-- Tier 3: Change-aware merge gate
|
||||
trusted base-branch planner classifies the complete PR diff
|
||||
Buildkite adds only relevant integration, quality, training,
|
||||
API, or performance lanes on Slinky Slurm
|
||||
|
|
||||
pass -> Mergify squash-merges when all merge conditions pass
|
||||
fail -> fix, push, and re-run
|
||||
|
||||
|-- /test full
|
||||
`-- Explicit all-20-lane diagnostic run
|
||||
|
||||
`-- weekly schedule on main
|
||||
`-- Complete four-GPU SSIM matrix
|
||||
```
|
||||
|
||||
CI is not one monolithic job:
|
||||
@@ -34,10 +42,31 @@ CI is not one monolithic job:
|
||||
- GitHub Actions owns pre-commit, slash-command handling, aggregate status
|
||||
updates, docs deployment, image builds, package publishing, and community
|
||||
automations.
|
||||
- Buildkite owns the GPU test pipeline and path filtering.
|
||||
- Modal owns the actual GPU execution environment for test jobs.
|
||||
- Buildkite owns the GPU test graph, statuses, and trusted dispatch control
|
||||
plane. Its agent runs on the Slurm login plane; it does not execute test
|
||||
payloads.
|
||||
- Slinky Slurm is the only active CI compute backend. A host-owned dispatcher
|
||||
leases GPUs from a persistent four-GPU allocation and runs each lane in an
|
||||
isolated Enroot container at the immutable PR SHA.
|
||||
- Mergify owns merge protection, labeling, and the final squash merge.
|
||||
|
||||
The old files under `fastvideo/tests/modal/` are retained as dormant manual
|
||||
rollback code. `.buildkite/scripts/pr_test.sh` rejects Buildkite invocations,
|
||||
and no pipeline or slash-command route calls Modal.
|
||||
|
||||
Three Buildkite entry pipelines share the validated graph:
|
||||
|
||||
| Pipeline | Trigger | Scope |
|
||||
|---|---|---|
|
||||
| `pr-fastcheck` | Automatic pull-request webhook | Six Fastcheck lanes |
|
||||
| `ci` | `/merge`, `ready`, schedules, and `/test` API builds | Change-aware merge gates, scheduled SSIM, explicit Full Suite, Fastcheck reruns, or one direct lane |
|
||||
| `fastvideo-performance-lane` | Weekly scheduler | Direct performance lane |
|
||||
|
||||
Each entry pipeline starts with the same trusted `pipeline-upload` job on the
|
||||
`ci-runner` queue. The `ci` pipeline's incoming GitHub webhook is disabled;
|
||||
otherwise it would duplicate the automatic `pr-fastcheck` build. API and
|
||||
scheduled builds continue to work with webhook processing disabled.
|
||||
|
||||
## CI Tiers
|
||||
|
||||
### Tier 1: Pre-commit
|
||||
@@ -70,34 +99,37 @@ debugging a hook implementation.
|
||||
| Attribute | Value |
|
||||
|---|---|
|
||||
| Triggered by | Buildkite PR builds with `TEST_SCOPE=fastcheck` or unset |
|
||||
| Runner | Buildkite agent that launches Modal GPU jobs |
|
||||
| Compute | Slinky Slurm (`ci-runner` queue) |
|
||||
| Definition | `.buildkite/pipeline.yml` |
|
||||
| Entrypoint | `.buildkite/scripts/pr_test.sh` -> `fastvideo/tests/modal/pr_test.py` |
|
||||
| Entrypoint | Trusted host driver -> `.buildkite/scripts/unit_test.sh` or `.buildkite/scripts/lanes/*.sh` |
|
||||
|
||||
Fastcheck uses Buildkite's `monorepo-diff` plugin. Jobs whose watched paths did
|
||||
not change are skipped and do not block the aggregate `fastcheck-passed`
|
||||
status.
|
||||
Fastcheck always schedules these six lanes: encoder, VAE, transformer, custom
|
||||
kernels, unit tests, and DreamVerse. Static steps replace the former
|
||||
host-side path-filter plugin: the login plane never checks out or executes PR
|
||||
code.
|
||||
|
||||
| Buildkite label | `TEST_TYPE` | Main watched paths |
|
||||
|---|---|---|
|
||||
| Encoder Tests | `encoder` | `fastvideo/models/encoders/**`, `fastvideo/models/loader/**`, `fastvideo/tests/encoders/**`, `pyproject.toml`, `docker/Dockerfile` |
|
||||
| VAE Tests | `vae` | `fastvideo/models/vaes/**`, `fastvideo/models/loader/**`, `fastvideo/tests/vaes/**`, `pyproject.toml`, `docker/Dockerfile` |
|
||||
| Transformer Tests | `transformer` | `fastvideo/models/dits/**`, `fastvideo/models/loader/**`, `fastvideo/tests/transformers/**`, `fastvideo/layers/**`, `fastvideo/attention/**`, `pyproject.toml`, `docker/Dockerfile` |
|
||||
| Kernel Tests | `kernel_tests` | `fastvideo-kernel/**`, `pyproject.toml`, `docker/Dockerfile` |
|
||||
| Unit Tests | `unit_test` | `fastvideo/**`, `.buildkite/**`, `.github/**`, `pyproject.toml`, `docker/Dockerfile` |
|
||||
| DreamVerse App Tests | `dreamverse_app` | `apps/dreamverse/**`, `pyproject.toml` |
|
||||
|
||||
### Tier 3: Full Suite
|
||||
### Tier 3: Change-Aware Merge Gate
|
||||
|
||||
| Attribute | Value |
|
||||
|---|---|
|
||||
| Triggered by | `/merge`, adding `ready`, `/test full`, or a new push to a PR that already has `ready` |
|
||||
| Runner | Buildkite agent that launches Modal GPU jobs |
|
||||
| Triggered by | `/merge`, adding `ready`, or a new push to a PR that already has `ready` |
|
||||
| Compute | Slinky Slurm only (`ci-runner` queue) |
|
||||
| Definition | `.buildkite/pipeline.yml` |
|
||||
| Entrypoint | `.buildkite/scripts/pr_test.sh` -> `fastvideo/tests/modal/pr_test.py` |
|
||||
| Entrypoint | `/opt/fastvideo-ci-runner/run-ci` (`run-unit` is a compatibility wrapper) |
|
||||
|
||||
Full Suite is also path-filtered. It validates broader behavior before Mergify
|
||||
can merge a PR.
|
||||
Fastcheck is the universal six-lane baseline. The merge gate does not repeat
|
||||
those jobs: it classifies every changed path and adds only the relevant lanes
|
||||
from the fourteen-lane integration set below. Selected jobs are hard gates;
|
||||
there are no soft-fail hardware lanes. A documentation-only PR can therefore
|
||||
finish its merge build after the trusted uploader, while model-family changes
|
||||
typically add focused golden-gate and SSIM files and a training-only change
|
||||
adds only its owning training lane.
|
||||
|
||||
`.github/scripts/plan_merge_ci.py` is the canonical path policy. It runs from
|
||||
the immutable base SHA under `pull_request_target`; PR code is never executed
|
||||
on the GitHub runner. The changed-file list includes both sides of renames. An
|
||||
API failure, truncated response, empty list, unknown build input, or unknown
|
||||
source path fails closed to all fourteen integration lanes.
|
||||
|
||||
A `ready`-labeled PR does not hit Buildkite immediately:
|
||||
`ci-trigger-full-suite.yml` first runs `.github/scripts/gate_full_suite.sh`,
|
||||
@@ -106,25 +138,100 @@ 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` |
|
||||
| LoRA Inference Tests | `inference_lora` | LoRA tests, loader, transformer tests, pipelines, LoRA layers |
|
||||
| LoRA Extraction Tests | `lora_extraction` | LoRA extraction scripts/tests, loader, training utilities, LoRA layers |
|
||||
| Training Tests | `training` | `fastvideo/**`, `pyproject.toml`, `docker/Dockerfile` |
|
||||
| Distillation DMD Tests | `distillation_dmd` | `fastvideo/training/*distillation_pipeline.py` |
|
||||
| Self-Forcing Tests | `self_forcing` | self-forcing distillation pipeline and tests |
|
||||
| LoRA Training Tests | `training_lora` | `fastvideo/**`, `pyproject.toml`, `docker/Dockerfile` |
|
||||
| Training Tests VSA | `training_vsa` | `fastvideo/**`, `fastvideo-kernel/**`, `pyproject.toml`, `docker/Dockerfile` |
|
||||
| Inference Tests VMoBA | `inference_vmoba` | `fastvideo-kernel/**`, `fastvideo/attention/backends/vmoba.py` |
|
||||
| Performance Tests | `performance` | DiTs, pipelines, attention, layers, worker, entrypoints, performance tests/configs |
|
||||
| API Server Tests | `api_server` | OpenAI entrypoints, serve CLI, OpenAI API integration test |
|
||||
| Train Framework Tests | `train_framework` | `fastvideo/train/**`, train model/method tests, model loader, DiTs |
|
||||
| Eval Metrics Tests | `eval` | `fastvideo/eval/**`, `fastvideo/tests/eval/**`, `pyproject.toml`, `docker/Dockerfile` |
|
||||
The complete static graph remains available through `/test full`; path
|
||||
selection never deletes or dynamically invents a Buildkite step.
|
||||
|
||||
| Lane | Public `TEST_TYPE` | GPUs | Typical merge trigger |
|
||||
|---|---|---:|---|
|
||||
| Encoder | `encoder` | 1 | Universal Fastcheck |
|
||||
| VAE | `vae` | 1 | Universal Fastcheck |
|
||||
| Transformer | `transformer` | 1 | Universal Fastcheck |
|
||||
| Kernel | `kernel_tests` | 1 | Universal Fastcheck |
|
||||
| Unit | `unit_test` | 1 | Universal Fastcheck |
|
||||
| DreamVerse | `dreamverse_app` | 1 | Universal Fastcheck |
|
||||
| Golden gate | `golden_gate` | 1 | Model, pipeline, attention, layer, or output changes |
|
||||
| SSIM | `ssim` | 4 | Matching model/SSIM paths; focused files when possible |
|
||||
| LoRA inference | `inference_lora` | 1 | LoRA inference/shared LoRA paths |
|
||||
| LoRA extraction | `lora_extraction` | 1 | LoRA extraction/shared LoRA paths |
|
||||
| Vanilla training | `training` | 4 | Legacy vanilla/shared training paths |
|
||||
| DMD distillation | `distillation_dmd` | 2 | DMD/shared training paths |
|
||||
| Self-forcing | `self_forcing` | 2 | Self-forcing/shared training paths |
|
||||
| LoRA training | `training_lora` | 2 | LoRA/shared training paths |
|
||||
| VSA training | `training_vsa` | 2 | VSA/shared training paths |
|
||||
| VMoBA inference | `inference_vmoba` | 1 | VMoBA backend/config paths |
|
||||
| Performance | `performance` | 2 | Performance tests or benchmark policy |
|
||||
| API server | `api_server` | 1 | API, worker, or server entrypoint paths |
|
||||
| Modular train framework | `train_framework` | 1 | `fastvideo/train/` and its tests |
|
||||
| Eval metrics | `eval` | 1 | `fastvideo/eval/` and its tests |
|
||||
|
||||
Golden-gate and SSIM selections are basenames, not arbitrary pytest arguments.
|
||||
The private host checks the comma-separated allowlist before staging, and the
|
||||
container checks it again before invoking pytest. Shared quality-harness
|
||||
changes still run the complete owning lane. The full SSIM matrix also runs on
|
||||
`main` every Sunday at 05:00 UTC through `ci-scheduled-ssim.yml`; `/test ssim`
|
||||
and `/test full` remain available for deliberate complete runs.
|
||||
|
||||
The four Buildkite workers may accept multiple jobs concurrently. The
|
||||
agent-owned lease broker packs their requested GPU counts onto one persistent
|
||||
four-GPU Slurm allocation and waits when capacity is full. A four-GPU lane
|
||||
such as SSIM or vanilla training owns the whole tray; two two-GPU lanes or up
|
||||
to four one-GPU lanes can overlap without sharing devices. SSIM and vanilla
|
||||
training also share the `fastvideo/slinky/whole-tray` Buildkite concurrency
|
||||
group. That keeps the second whole-tray lane in Buildkite instead of consuming
|
||||
an agent and its command timeout while the first lane waits for four free GPUs.
|
||||
Because packed Enroot containers share the node network namespace, each GPU
|
||||
lease receives its own 100-port rendezvous range. Tests preserve the
|
||||
runner-assigned `MASTER_PORT`, and parallel SSIM tasks use distinct offsets
|
||||
inside that range.
|
||||
|
||||
See [Performance Benchmarks](performance_benchmarks.md) for the performance
|
||||
lane's thresholds, rolling baseline, artifacts, and reseeding process.
|
||||
|
||||
## `/merge` Request Flow
|
||||
|
||||
The Buildkite agent and Slurm have deliberately separate responsibilities:
|
||||
|
||||
```text
|
||||
/merge PR comment
|
||||
-> GitHub verifies write permission and refreshes the `ready` label
|
||||
-> base-branch `ci-trigger-full-suite` workflow fetches the PR file list,
|
||||
computes MERGE_TEST_PLAN plus focused golden/SSIM basenames, and gates on
|
||||
cheap checks
|
||||
-> sends the PR SHA to pipeline `ci` with TEST_SCOPE=merge and FULL_SUITE=true
|
||||
-> trusted `pipeline-upload` job on queue `ci-runner`
|
||||
fetch exact-SHA .buildkite/pipeline.yml
|
||||
normalize + validate the complete static 20-lane policy and conditions
|
||||
upload static Buildkite steps
|
||||
-> each accepted step reaches the host policy hook
|
||||
validate org/repo/SHA/ref/step key/command/scope/timeout
|
||||
skip checkout on the login plane
|
||||
stage a mode-0600 request on Lustre
|
||||
lease 1-4 GPUs and attach an `srun` step to the Slinky tray
|
||||
-> Enroot worker container
|
||||
clone and verify the exact PR SHA
|
||||
install the lane's project extras and cached kernel
|
||||
run the repository-owned lane script
|
||||
write numeric exit status and approved artifacts
|
||||
-> trusted host returns that status to Buildkite
|
||||
-> Buildkite publishes `full-suite-passed` only when every selected lane passes
|
||||
```
|
||||
|
||||
The pipeline uploader and dispatcher run on the Slurm login plane, but those
|
||||
are control-plane operations only. Python tests, model loading, CUDA kernels,
|
||||
Node/Playwright checks, inference, training, SSIM generation, and performance
|
||||
benchmarks all execute in Slurm allocations.
|
||||
|
||||
PR-controlled values never become host commands. The host policy accepts only
|
||||
the pinned pipeline uploader or a known lane tuple. It rejects plugins,
|
||||
artifact globs, shell injection variables, non-immutable commits, and unknown
|
||||
commands before checkout. Hugging Face credentials are added only for lanes
|
||||
that declare them, passed through a mode-0600 request file, and removed before
|
||||
the PR payload starts. Active training lanes keep W&B offline and do not stage
|
||||
a W&B credential. The ARM64 image includes the pinned FA4 CuTe overlay validated
|
||||
on GB200. SSIM opts into FA4 to preserve its reference-video numerics; lanes
|
||||
with FA2 baselines keep `FASTVIDEO_FA4=0`. Performance artifacts are relayed
|
||||
afterward by the trusted host from an allowlisted directory and extension set.
|
||||
|
||||
## Slash Commands
|
||||
|
||||
Slash commands are handled by `.github/workflows/ci-slash-commands.yml`.
|
||||
@@ -132,7 +239,7 @@ Repository write permission is required.
|
||||
|
||||
| Command | Effect |
|
||||
|---|---|
|
||||
| `/merge` | Adds `ready` and triggers Full Suite for the PR head branch. |
|
||||
| `/merge` | Adds `ready` and triggers the path-aware merge gate for the PR head. |
|
||||
| `/test full` | Runs the whole Full Suite with `TEST_SCOPE=full`. |
|
||||
| `/test fastcheck` | Runs the whole Fastcheck suite with `TEST_SCOPE=fastcheck`. |
|
||||
| `/test pre-commit` | Re-runs the pre-commit workflow on the PR merge ref. |
|
||||
@@ -149,6 +256,7 @@ Valid direct test names:
|
||||
| `/test unit` | `unit_test` |
|
||||
| `/test dreamverse` | `dreamverse_app` |
|
||||
| `/test ssim` | `ssim` |
|
||||
| `/test golden-gate` | `golden_gate` |
|
||||
| `/test training` | `training` |
|
||||
| `/test lora-inference` | `inference_lora` |
|
||||
| `/test lora-training` | `training_lora` |
|
||||
@@ -162,12 +270,18 @@ Valid direct test names:
|
||||
| `/test train-framework` | `train_framework` |
|
||||
| `/test eval` | `eval` |
|
||||
|
||||
The temporary `<name>-ci` spellings remain accepted as compatibility aliases;
|
||||
they select the same Slurm lane and do not identify a second backend.
|
||||
|
||||
When a direct test completes successfully, Buildkite posts
|
||||
`direct-test-completed`. `.github/workflows/ci-aggregate-status.yml` then reads
|
||||
the latest Buildkite statuses for the commit and updates `fastcheck-passed` or
|
||||
`full-suite-passed` if all jobs in that group are green.
|
||||
|
||||
Skipped path-filtered jobs have no status entry and do not block the aggregate.
|
||||
Buildkite label emojis define the status namespace used by that aggregation:
|
||||
`:microscope:` is reserved for the six Fastcheck lanes, while Full-Suite-only
|
||||
lanes use `:test_tube:` or `:bar_chart:`. Each active lane has exactly one
|
||||
label and therefore one status context.
|
||||
|
||||
## Merge Protection
|
||||
|
||||
@@ -177,7 +291,7 @@ Mergify enforces these conditions before it squash-merges to `main`:
|
||||
|---|---|
|
||||
| `check-success~=pre-commit` | Tier 1 passed. |
|
||||
| `check-success=fastcheck-passed` | All triggered Fastcheck jobs passed. |
|
||||
| `check-success=full-suite-passed` | All triggered Full Suite jobs passed. |
|
||||
| `check-success=full-suite-passed` | The selected merge gate or explicit Full Suite passed. |
|
||||
| `#approved-reviews-by>=1` | At least one approving review. |
|
||||
| Valid title regex | PR title starts with an accepted `[type]` tag. |
|
||||
| `label=ready` | The PR has entered the merge flow. |
|
||||
@@ -226,29 +340,40 @@ Process labels:
|
||||
|
||||
| Label | Who sets it | Meaning |
|
||||
|---|---|---|
|
||||
| `ready` | `/merge` or maintainer action | Triggers/keeps Full Suite active and enables auto-merge. |
|
||||
| `ready` | `/merge` or maintainer action | Triggers/keeps the change-aware merge gate active and enables auto-merge. |
|
||||
| `needs-rebase` | Mergify | PR has merge conflicts. |
|
||||
| `do-not-merge` | Maintainer | Blocks merge regardless of CI status. |
|
||||
|
||||
## Modal Test Entrypoints
|
||||
## Slurm Lane Entrypoints
|
||||
|
||||
All Buildkite test jobs go through `.buildkite/scripts/pr_test.sh`, which:
|
||||
Every active test selection lives in `.buildkite/scripts/unit_test.sh` or a
|
||||
focused `.buildkite/scripts/lanes/<lane>.sh`. The private, agent-owned lane
|
||||
table binds each internal `*_ci` type to that script, its GPU count, wall-clock
|
||||
limit, dependency extras, kernel-build policy, secrets, and artifacts. The
|
||||
internal suffix is an implementation detail; there is only one active backend.
|
||||
|
||||
1. Reads Buildkite secrets for Modal, Hugging Face, and W&B when needed.
|
||||
2. Selects a Modal function based on `TEST_TYPE`.
|
||||
3. Passes Buildkite metadata into the Modal container.
|
||||
4. Runs the selected test command from `fastvideo/tests/modal/pr_test.py` or
|
||||
`fastvideo/tests/modal/ssim_test.py`.
|
||||
5. Uploads performance artifacts for `TEST_TYPE=performance`.
|
||||
SSIM uses `fastvideo/tests/ssim/ci_runner.py` inside a single four-GPU lease.
|
||||
It discovers `REQUIRED_GPUS` and `*_MODEL_TO_PARAMS` with AST parsing, then
|
||||
packs independent pytest subprocesses across the visible GPUs with fail-fast
|
||||
termination. Performance writes reports to a host-mounted artifact directory;
|
||||
the host uploads only `.md`, `.html`, `.json`, and `.csv` files after the
|
||||
container exits.
|
||||
|
||||
The Modal launchers remain in the repository for manual rollback archaeology,
|
||||
but they are not CI entrypoints. `pr_test.sh` rejects Buildkite calls and needs
|
||||
`FASTVIDEO_ENABLE_LEGACY_MODAL_CI=1` even for a local manual invocation.
|
||||
|
||||
If you add a new CI test category:
|
||||
|
||||
1. Add the Modal function in `fastvideo/tests/modal/pr_test.py` or a focused
|
||||
companion module.
|
||||
2. Add the `TEST_TYPE` case in `.buildkite/scripts/pr_test.sh`.
|
||||
3. Add the Buildkite direct-test step and any Fastcheck/Full Suite path filters
|
||||
in `.buildkite/pipeline.yml`.
|
||||
4. Add or update the `/test` mapping in `.github/workflows/ci-slash-commands.yml`.
|
||||
1. Add an executable `.buildkite/scripts/lanes/<lane>.sh` containing the test
|
||||
payload only.
|
||||
2. Add the static Buildkite step in `.buildkite/pipeline.yml`, its source/test
|
||||
ownership in `.github/scripts/plan_merge_ci.py`, and the `/test` mapping in
|
||||
`.github/workflows/ci-slash-commands.yml`.
|
||||
3. Extend `fastvideo/tests/contract/test_ci_test_collection.py`,
|
||||
`test_merge_ci_plan.py`, and the trusted private runner's lane table plus
|
||||
uploader policy in the same rollout.
|
||||
4. Validate on the target GB200 hardware before making the lane a merge gate.
|
||||
5. Document the lane here and link any domain-specific authoring guide.
|
||||
|
||||
## CD And Release Workflows
|
||||
@@ -282,13 +407,18 @@ the explicit `py3.12-cuda13.0.0-latest` tag. This publication policy does not
|
||||
change the unparameterized `docker/Dockerfile` build defaults, which remain CUDA
|
||||
13 and `cu130`.
|
||||
|
||||
Published amd64 development images keep their configured Hopper kernel wheel
|
||||
installed and also carry an immutable SM89 wheel under
|
||||
`/opt/fastvideo-kernel-prebuilt`. Modal PR and SSIM jobs select the exact
|
||||
source, ABI, and GPU-architecture match from that directory, so L40S jobs reuse
|
||||
the trusted image artifact while kernel-changing PRs still build locally. Once
|
||||
a kernel or artifact-key change reaches `main`, the image workflow republishes
|
||||
the matching trusted artifact before later jobs consume the updated image tag.
|
||||
Published development images carry architecture-specific kernel wheels under
|
||||
`/opt/fastvideo-kernel-prebuilt`. The Slurm worker selects the exact source,
|
||||
ABI, and GPU-architecture match, so normal lanes reuse the trusted artifact
|
||||
while kernel-changing PRs still build locally. Once a kernel or artifact-key
|
||||
change reaches `main`, the image workflow republishes the matching artifact
|
||||
before later jobs consume the updated image pin.
|
||||
|
||||
The same workflow publishes a single-architecture ARM64, CUDA 13, SM100 image
|
||||
for the self-hosted CI runner under the
|
||||
`py3.12-cuda13.0.0-sm100-{latest,sha-*}` tags. It carries the matching prebuilt
|
||||
kernel wheel so runner jobs can validate and install the exact source and ABI
|
||||
match instead of recompiling it in every lane.
|
||||
|
||||
The optional Dreamverse matrix builds backend and UI images for CUDA 12.6 and
|
||||
CUDA 13 on `amd64`. Dreamverse remains `amd64`-only because its FA4 dependency
|
||||
@@ -320,15 +450,17 @@ The reusable implementation lives in
|
||||
| `.github/mergify.yml` | Merge protection, PR title validation, PR labels, conflict labels, auto-merge |
|
||||
| `.github/workflows/ci-precommit.yml` | Tier 1 pre-commit |
|
||||
| `.github/workflows/ci-slash-commands.yml` | `/merge` and `/test` handling |
|
||||
| `.github/workflows/ci-trigger-full-suite.yml` | Full Suite trigger for `ready` PRs and new pushes to ready PRs |
|
||||
| `.github/workflows/ci-aggregate-status.yml` | Aggregate Fastcheck/Full Suite commit statuses |
|
||||
| `.buildkite/pipeline.yml` | Buildkite test graph and path filters |
|
||||
| `.buildkite/scripts/pr_test.sh` | Buildkite-to-Modal test dispatcher |
|
||||
| `fastvideo/tests/modal/pr_test.py` | Modal functions for most GPU CI lanes |
|
||||
| `fastvideo/tests/modal/ssim_test.py` | Modal functions and partitioning for SSIM |
|
||||
| `.github/workflows/ci-trigger-full-suite.yml` | Change-aware merge-gate trigger for `ready` PRs and new pushes |
|
||||
| `.github/workflows/ci-scheduled-ssim.yml` | Weekly complete SSIM trigger on `main` |
|
||||
| `.github/scripts/plan_merge_ci.py` | Trusted changed-path to integration-lane planner |
|
||||
| `.github/workflows/ci-aggregate-status.yml` | Aggregate Fastcheck and explicit Full Suite direct-rerun statuses |
|
||||
| `.buildkite/pipeline.yml` | Static 20-lane Slurm Buildkite graph |
|
||||
| `.buildkite/scripts/unit_test.sh`, `.buildkite/scripts/lanes/*.sh` | Active Slurm lane payloads |
|
||||
| `fastvideo/tests/ssim/ci_runner.py` | Four-GPU Slurm SSIM scheduler |
|
||||
| `.buildkite/scripts/pr_test.sh`, `fastvideo/tests/modal/*.py` | Dormant manual Modal rollback path (disabled in Buildkite) |
|
||||
| `.buildkite/performance-benchmarks/tests/*.json` | Performance benchmark configs and thresholds |
|
||||
| `.github/workflows/infra-docs.yml` | Docs build and GitHub Pages deploy |
|
||||
| `.github/workflows/infra-build-image.yml` | Automatic CUDA matrix and manual Docker image builds |
|
||||
| `.github/workflows/infra-build-image.yml` | CUDA matrix, CI runner image, and manual Docker image builds |
|
||||
| `.github/workflows/publish-fastvideo.yml` | FastVideo PyPI publishing |
|
||||
| `.github/workflows/publish-kernel.yml` | FastVideo kernel PyPI publishing |
|
||||
| `.github/workflows/publish-comfyui.yml` | ComfyUI registry publishing |
|
||||
|
||||
@@ -23,8 +23,8 @@ It serves three audiences:
|
||||
pytest fastvideo/tests/performance/ -vs
|
||||
|
||||
# Optional: compare against the rolling HF baseline.
|
||||
# PERF_REPORTS_DIR defaults to /root/data/perf_reports for Modal/CI, so
|
||||
# override it when running outside the container.
|
||||
# PERF_REPORTS_DIR defaults to /root/data/perf_reports in a CI container, so
|
||||
# override it for a local run.
|
||||
PERF_REPORTS_DIR=/tmp/fastvideo_perf_reports \
|
||||
python fastvideo/tests/performance/compare_baseline.py
|
||||
|
||||
@@ -459,16 +459,18 @@ FlashInfer, Cutlass DSL, SageAttention, Triton, and xFormers when installed.
|
||||
| `FASTVIDEO_FA4` | `0` | `test_inference_performance.py` | FlashAttention-4 toggle included in `software_profile_id`. |
|
||||
| `FASTVIDEO_PERFORMANCE_PROFILE_VERSION` | unset | `test_inference_performance.py` | Optional explicit software cohort/profile version included in `software_profile_id`. |
|
||||
| `IMAGE_VERSION` | unset | `test_inference_performance.py` | CI container image/profile version included in `software_profile_id` when available. |
|
||||
| `FASTVIDEO_CONTAINER_IMAGE_REF` | unset | `pr_test.py`, `launch_l40s_job.py`, `test_inference_performance.py` | Resolved CI container image ref or digest recorded in `environment_metadata` for audit without changing `software_profile_id`. |
|
||||
| `FASTVIDEO_CONTAINER_IMAGE_REF` | unset | Slurm runner, `test_inference_performance.py` | Pinned CI container image digest recorded in `environment_metadata` and `software_profile_id`. |
|
||||
| `FASTVIDEO_STAGE_LOGGING` | set by the pytest test | `test_inference_performance.py` | Enables pipeline stage timing capture for component metrics during benchmark runs. |
|
||||
|
||||
## CI integration
|
||||
|
||||
The performance step can run on demand with `/test performance` and as part of
|
||||
the Full Suite (see [CI/CD Architecture](ci_architecture.md)). The Modal entry
|
||||
point is `fastvideo/tests/modal/pr_test.py:run_performance_tests` and the
|
||||
Buildkite artifact upload is in
|
||||
`.buildkite/scripts/pr_test.sh:upload_performance_artifacts`.
|
||||
The performance step can run on demand with `/test performance`, through a
|
||||
merge gate when performance tests or benchmark policy changed, and as part of
|
||||
an explicit `/test full` run (see [CI/CD Architecture](ci_architecture.md)).
|
||||
The weekly `fastvideo-performance-lane` schedule runs the same Slurm payload.
|
||||
The active entry point is `.buildkite/scripts/lanes/performance.sh`; the
|
||||
trusted host dispatcher relays its allowlisted reports to Buildkite after the
|
||||
isolated container exits.
|
||||
|
||||
Each performance build runs pytest first. PR and direct runs only continue to
|
||||
`compare_baseline.py` when that fixed-threshold phase passes; if pytest fails,
|
||||
|
||||
@@ -42,19 +42,20 @@ The important process labels are:
|
||||
|
||||
| Label | Meaning |
|
||||
|---|---|
|
||||
| `ready` | The PR is ready for Full Suite and auto-merge consideration. |
|
||||
| `ready` | The PR is ready for the change-aware merge gate and auto-merge consideration. |
|
||||
| `needs-rebase` | The PR has merge conflicts with `main`. |
|
||||
| `do-not-merge` | A maintainer has blocked merge. |
|
||||
|
||||
## CI Summary
|
||||
|
||||
FastVideo has three validation tiers:
|
||||
FastVideo has three routine validation tiers plus an explicit full diagnostic:
|
||||
|
||||
| Tier | Runs when | What it does |
|
||||
|---|---|---|
|
||||
| Pre-commit | Pull requests and `/test pre-commit` | Formatting, linting, typing, spelling, Markdown, workflow syntax, filename checks |
|
||||
| Fastcheck | PR Buildkite builds | Path-filtered component and unit checks on Modal GPU runners |
|
||||
| Full Suite | `/merge`, `ready`, `/test full`, or new pushes to ready PRs | Path-filtered integration, SSIM, training, eval, API, and performance checks |
|
||||
| Fastcheck | PR Buildkite builds | Six component, kernel, unit, and app lanes on Slinky Slurm |
|
||||
| Merge gate | `/merge`, `ready`, or new pushes to ready PRs | Only path-relevant integration lanes on Slinky Slurm; Fastcheck remains the baseline |
|
||||
| Full Suite | `/test full` | Explicit all-twenty-lane diagnostic run on Slinky Slurm |
|
||||
|
||||
See [CI/CD Architecture](ci_architecture.md#ci-tiers) for exact jobs, path
|
||||
filters, and workflow files.
|
||||
@@ -66,10 +67,11 @@ filters, and workflow files.
|
||||
3. Fix pre-commit failures locally with `pre-commit run --all-files`.
|
||||
4. Wait for at least one approving review.
|
||||
5. When the PR is approved and ready, comment `/merge`.
|
||||
6. `/merge` adds `ready` and triggers the Full Suite for the PR branch.
|
||||
6. `/merge` adds `ready`, waits for cheap checks, and triggers the minimal
|
||||
path-relevant integration lanes for the PR branch.
|
||||
7. If all required checks pass, Mergify squash-merges the PR to `main`.
|
||||
8. If Full Suite fails, fix the regression, push again, and re-run `/merge` or
|
||||
the failed test.
|
||||
8. If the merge gate fails, fix the regression, push again, and re-run
|
||||
`/merge`. Use a targeted `/test` command for diagnosis.
|
||||
|
||||
Only contributors with repository write permission can use slash commands. If
|
||||
you are an external contributor, ask a maintainer to run `/merge` or add
|
||||
@@ -128,7 +130,7 @@ git push --force-with-lease
|
||||
|
||||
Mergify removes `needs-rebase` after conflicts are resolved.
|
||||
|
||||
### Full Suite Fails
|
||||
### Merge Gate Or Full Suite Fails
|
||||
|
||||
The failing Buildkite step is the source of truth. Common causes are:
|
||||
|
||||
|
||||
@@ -138,47 +138,31 @@ pytest fastvideo/tests/ssim/ -vs
|
||||
|
||||
Use a machine whose GPU and backend match the reference folder you are testing.
|
||||
|
||||
## Modal Runs For SSIM
|
||||
## Slurm CI Runs For SSIM
|
||||
|
||||
For CI-like SSIM execution, use `fastvideo/tests/modal/ssim_test.py`:
|
||||
Comment `/test ssim` on a pull request to run the canonical four-GPU SSIM
|
||||
lane on the Slinky Slurm cluster. `fastvideo/tests/ssim/ci_runner.py`
|
||||
discovers the suite without importing test modules, packs independent pytest
|
||||
processes across the four assigned GPUs, and stops the lane on the first
|
||||
failure.
|
||||
|
||||
The change-aware `/merge` planner may run only the SSIM test basenames owned
|
||||
by the changed model family. Shared SSIM harness changes still select the
|
||||
complete lane. Independently, `main` runs the full SSIM matrix every Sunday at
|
||||
05:00 UTC so infrequently touched model families retain periodic coverage.
|
||||
|
||||
For a focused developer run, invoke pytest directly and optionally select one
|
||||
model from a parameterized test through `FASTVIDEO_SSIM_MODEL_ID`:
|
||||
|
||||
```bash
|
||||
python -m modal run fastvideo/tests/modal/ssim_test.py::run_ssim_tests
|
||||
pytest fastvideo/tests/ssim/test_wan_t2v_similarity.py -vs
|
||||
|
||||
FASTVIDEO_SSIM_MODEL_ID=Wan2.1-T2V-1.3B-Diffusers \
|
||||
pytest fastvideo/tests/ssim/test_wan_t2v_similarity.py -vs
|
||||
```
|
||||
|
||||
Target specific files or model ids:
|
||||
|
||||
```bash
|
||||
python -m modal run fastvideo/tests/modal/ssim_test.py::run_ssim_tests \
|
||||
--test-files test_wan_t2v_similarity.py \
|
||||
--model-ids Wan2.1-T2V-1.3B-Diffusers
|
||||
```
|
||||
|
||||
If `HF_API_KEY`, `HUGGINGFACE_HUB_TOKEN`, or `HF_TOKEN` is not set, the local
|
||||
entrypoint fails fast.
|
||||
|
||||
To export raw generated videos from Modal to the shared volume:
|
||||
|
||||
```bash
|
||||
python -m modal run fastvideo/tests/modal/ssim_test.py::run_ssim_tests \
|
||||
--sync-generated-to-volume
|
||||
```
|
||||
|
||||
The raw export path is quality-tiered:
|
||||
|
||||
- default params: `ssim_generated_videos/default/<subdir>/generated_videos`
|
||||
- full-quality params: `ssim_generated_videos/full_quality/<subdir>/generated_videos`
|
||||
|
||||
The printed `modal volume get` command downloads into
|
||||
`./generated_videos_modal/<quality-tier>`. Convert those outputs into local
|
||||
references with `copy-local`:
|
||||
|
||||
```bash
|
||||
python fastvideo/tests/ssim/reference_videos_cli.py copy-local \
|
||||
--quality-tier full_quality \
|
||||
--generated-dir ./generated_videos_modal/full_quality/L40S_reference_videos \
|
||||
--device-folder L40S_reference_videos
|
||||
```
|
||||
The files under `fastvideo/tests/modal/` are retained only as a disabled
|
||||
manual rollback implementation. No active CI trigger invokes them.
|
||||
|
||||
### SSIM Bootstrap Mode
|
||||
|
||||
@@ -206,28 +190,30 @@ python fastvideo/tests/ssim/reference_videos_cli.py promote-draft \
|
||||
|
||||
## CI Integration
|
||||
|
||||
FastVideo CI tests are orchestrated by Buildkite and run on Modal GPU
|
||||
instances. The main files are:
|
||||
FastVideo GPU CI is orchestrated by Buildkite and runs only on isolated Slinky
|
||||
Slurm workers. The main files are:
|
||||
|
||||
| File | Purpose |
|
||||
|---|---|
|
||||
| `.buildkite/pipeline.yml` | Buildkite test graph and path filters. |
|
||||
| `.buildkite/scripts/pr_test.sh` | Dispatches `TEST_TYPE` to a Modal function. |
|
||||
| `fastvideo/tests/modal/pr_test.py` | Modal functions for most test lanes. |
|
||||
| `fastvideo/tests/modal/ssim_test.py` | Modal functions and partitioning for SSIM. |
|
||||
| `.buildkite/pipeline.yml` | Static, validated 20-lane Slurm test graph. |
|
||||
| `.github/scripts/plan_merge_ci.py` | Trusted path-to-lane and focused quality-test policy for `/merge`. |
|
||||
| `.buildkite/scripts/unit_test.sh`, `.buildkite/scripts/lanes/*.sh` | Repository-owned test payloads executed inside Slurm containers. |
|
||||
| `fastvideo/tests/ssim/ci_runner.py` | Four-GPU SSIM task discovery and scheduling. |
|
||||
| `.buildkite/scripts/pr_test.sh`, `fastvideo/tests/modal/*.py` | Dormant manual rollback path; rejected in Buildkite. |
|
||||
|
||||
For exact tier membership, path filters, slash commands, and aggregate statuses,
|
||||
For exact tier membership, slash commands, runner isolation, and aggregate statuses,
|
||||
see [CI/CD Architecture](ci_architecture.md).
|
||||
|
||||
### Adding A New CI Test Category
|
||||
|
||||
If a new test does not fit an existing lane:
|
||||
|
||||
1. Add a Modal function in `fastvideo/tests/modal/pr_test.py` or a focused
|
||||
companion module.
|
||||
2. Add a matching `TEST_TYPE` case in `.buildkite/scripts/pr_test.sh`.
|
||||
3. Add Buildkite direct-test and path-filter entries in `.buildkite/pipeline.yml`.
|
||||
4. Add the `/test` mapping in `.github/workflows/ci-slash-commands.yml`.
|
||||
1. Put the test payload in an executable `.buildkite/scripts/lanes/<lane>.sh`.
|
||||
2. Add its static step to `.buildkite/pipeline.yml`, its changed-path ownership
|
||||
to `.github/scripts/plan_merge_ci.py`, and extend the CI contract tests.
|
||||
3. Add the `/test` mapping in `.github/workflows/ci-slash-commands.yml`.
|
||||
4. Coordinate the matching GPU, timeout, dependency, secret, and artifact
|
||||
policy in the private Slurm runner allowlist.
|
||||
5. Document the new category in [CI/CD Architecture](ci_architecture.md) and add
|
||||
authoring notes here if contributors need them.
|
||||
|
||||
|
||||
@@ -8,7 +8,7 @@ compute ~3.7×, so the wall-clock drops far more than 2×. RIFE (which estimates
|
||||
its own optical flow — no motion vectors needed) fills the dropped frames back
|
||||
in for ~1.4 s, and a light unsharp pass counters its softening.
|
||||
|
||||
Measured on the 1.3B INT8 QAD model (fox, 480×832×81, M4): generate 41 + RIFE→81
|
||||
Measured on the 1.3B INT8 QAD model (480×832×81, M4): generate 41 + RIFE→81
|
||||
runs in ~35 s of denoise vs ~90 s full, at reconstruction MS-SSIM **0.97**.
|
||||
Reproduce with `python -m fastvideo.benchmarks.eval_metalfx_rife --mode int8`.
|
||||
|
||||
@@ -26,10 +26,11 @@ uv pip install -e ".[mlx]" # RIFE ships vendored; this only needs MLX
|
||||
|
||||
```bash
|
||||
python examples/inference/basic/mlx_wan_prompt_to_video.py \
|
||||
--mlx-checkpoint <FastWan2.1-T2V-1.3B-INT8-QAD> \
|
||||
--prompt "A red fox trotting through a snowy pine forest at golden hour, cinematic" \
|
||||
--model-root ./FastMetal-1.3B-QAD \
|
||||
--mlx-checkpoint ./FastMetal-1.3B-QAD \
|
||||
--prompt "A bird's-eye view of a misty forest valley at dawn." \
|
||||
--num-frames 81 --fast \
|
||||
--output-path video_samples/fox_fast.mp4
|
||||
--output-path video_samples/forest_fast.mp4
|
||||
```
|
||||
|
||||
`--num-frames` stays the *target* length; fast mode generates the smallest
|
||||
@@ -57,9 +58,9 @@ of denoise. It composes with `--fast`; both together run the same clip in
|
||||
|
||||
```bash
|
||||
python examples/inference/basic/mlx_wan_prompt_to_video.py \
|
||||
--prompt "A red fox trotting through a snowy pine forest at golden hour, cinematic" \
|
||||
--prompt "A bird's-eye view of a misty forest valley at dawn." \
|
||||
--height 480 --width 832 --num-frames 81 --fast-spatial \
|
||||
--output-path video_samples/fox_fast_spatial.mp4
|
||||
--output-path video_samples/forest_fast_spatial.mp4
|
||||
```
|
||||
|
||||
| Flag | Default | Meaning |
|
||||
|
||||
@@ -77,6 +77,7 @@ surfaces:
|
||||
output_type: "Legacy output formatting surface pending GenerationResult cleanup."
|
||||
VSA_sparsity: "Model-specific inference optimization not yet represented in the typed public schema."
|
||||
VSA_tile_size: "VSA-H3 tile geometry request; model-specific optimization not yet represented in the typed public schema."
|
||||
inference_torch_compile: "Regional inference compile opt-in currently carried through PipelineSelection.experimental rather than CompileConfig."
|
||||
vae_parallel_decode: "MiniMax-H3 sequence-parallel VAE decode opt-in; model-specific optimization not yet represented in the typed public schema."
|
||||
vae_parallel_encode: "MiniMax-H3 sequence-parallel reference VAE encode opt-in; model-specific optimization not yet represented in the typed public schema."
|
||||
vae_parallel_decode_strategy: "Chunk-transport collective for vae_parallel_decode; model-specific optimization not yet represented in the typed public schema."
|
||||
@@ -123,12 +124,14 @@ surfaces:
|
||||
vae_precision: "Precision override pending dedicated typed component precision design."
|
||||
vae_decode_precision: "Decode-only precision override pending dedicated typed component precision design."
|
||||
image_encoder_precision: "Precision override pending dedicated typed component precision design."
|
||||
image_encoder_precisions: "Precision overrides pending dedicated typed component precision design."
|
||||
text_encoder_precisions: "Precision override pending dedicated typed component precision design."
|
||||
internal_only:
|
||||
dit_config: "Legacy internal component config object."
|
||||
upsampler_config: "Legacy internal component config object."
|
||||
vae_config: "Legacy internal component config object."
|
||||
image_encoder_config: "Legacy internal component config object."
|
||||
image_encoder_configs: "Legacy internal component config objects."
|
||||
text_encoder_configs: "Legacy internal component config object."
|
||||
preprocess_text_funcs: "Internal text preprocessing hooks."
|
||||
postprocess_text_funcs: "Internal text postprocessing hooks."
|
||||
@@ -365,6 +368,30 @@ surfaces:
|
||||
sources: [fastvideo.configs.pipelines.matrixgame2.MatrixGame2I2V480PConfig]
|
||||
num_frames_per_block:
|
||||
sources: [fastvideo.configs.pipelines.matrixgame2.MatrixGame2I2V480PConfig]
|
||||
duration_s:
|
||||
sources: [fastvideo.configs.pipelines.mmaudio.MMAudioV2AConfig]
|
||||
spectrogram_frame_rate:
|
||||
sources: [fastvideo.configs.pipelines.mmaudio.MMAudioV2AConfig]
|
||||
latent_downsample_rate:
|
||||
sources: [fastvideo.configs.pipelines.mmaudio.MMAudioV2AConfig]
|
||||
clip_frame_rate:
|
||||
sources: [fastvideo.configs.pipelines.mmaudio.MMAudioV2AConfig]
|
||||
sync_frame_rate:
|
||||
sources: [fastvideo.configs.pipelines.mmaudio.MMAudioV2AConfig]
|
||||
sync_segment_size:
|
||||
sources: [fastvideo.configs.pipelines.mmaudio.MMAudioV2AConfig]
|
||||
sync_segment_stride:
|
||||
sources: [fastvideo.configs.pipelines.mmaudio.MMAudioV2AConfig]
|
||||
sync_downsample_rate:
|
||||
sources: [fastvideo.configs.pipelines.mmaudio.MMAudioV2AConfig]
|
||||
clip_image_size:
|
||||
sources: [fastvideo.configs.pipelines.mmaudio.MMAudioV2AConfig]
|
||||
sync_image_size:
|
||||
sources: [fastvideo.configs.pipelines.mmaudio.MMAudioV2AConfig]
|
||||
clip_batch_size_multiplier:
|
||||
sources: [fastvideo.configs.pipelines.mmaudio.MMAudioV2AConfig]
|
||||
sync_batch_size_multiplier:
|
||||
sources: [fastvideo.configs.pipelines.mmaudio.MMAudioV2AConfig]
|
||||
audio_channels:
|
||||
sources:
|
||||
- fastvideo.configs.pipelines.stable_audio.StableAudioT2AConfig
|
||||
@@ -379,6 +406,7 @@ surfaces:
|
||||
- fastvideo.configs.pipelines.stable_audio.StableAudioOpenSmallConfig
|
||||
max_audio_duration_s:
|
||||
sources:
|
||||
- fastvideo.configs.pipelines.mmaudio.MMAudioV2AConfig
|
||||
- fastvideo.configs.pipelines.stable_audio.StableAudioT2AConfig
|
||||
- fastvideo.configs.pipelines.stable_audio.StableAudioOpenSmallConfig
|
||||
sample_size:
|
||||
@@ -387,6 +415,7 @@ surfaces:
|
||||
- fastvideo.configs.pipelines.stable_audio.StableAudioOpenSmallConfig
|
||||
sampling_rate:
|
||||
sources:
|
||||
- fastvideo.configs.pipelines.mmaudio.MMAudioV2AConfig
|
||||
- fastvideo.configs.pipelines.stable_audio.StableAudioT2AConfig
|
||||
- fastvideo.configs.pipelines.stable_audio.StableAudioOpenSmallConfig
|
||||
audio_txt_guidance_scale:
|
||||
|
||||
@@ -1,6 +1,10 @@
|
||||
# MPS (Apple Silicon)
|
||||
|
||||
Instructions to install FastVideo for Apple Silicon.
|
||||
Install FastVideo on Apple Silicon and run FastMetal-QAD.
|
||||
|
||||
Apple Silicon uses the MLX runtime and the FastMetal-QAD INT8 checkpoints.
|
||||
See the [FastMetal-QAD blog](https://haoailab.com/blogs/fastmetal/) and the
|
||||
[FastMetal collection](https://huggingface.co/collections/FastVideo/fastmetal).
|
||||
|
||||
## Requirements
|
||||
|
||||
@@ -49,7 +53,7 @@ brew install ffmpeg
|
||||
|
||||
### Installation
|
||||
|
||||
FastWan's native Apple Silicon runtime requires the `mlx` extra.
|
||||
FastMetal's native Apple Silicon runtime requires the `mlx` extra.
|
||||
|
||||
#### With uv (recommended)
|
||||
|
||||
@@ -87,6 +91,47 @@ Alternative with Conda environment:
|
||||
uv pip install -e ".[mlx]"
|
||||
```
|
||||
|
||||
## Run FastMetal-QAD
|
||||
|
||||
Each release is self-contained. Download one checkpoint and point both
|
||||
`--model-root` and `--mlx-checkpoint` at it (the example also auto-detects
|
||||
`mlx_dit.json` under `--model-root`).
|
||||
|
||||
| Checkpoint | Script | Mac tier |
|
||||
| --- | --- | --- |
|
||||
| [`FastVideo/FastMetal-1.3B-QAD`](https://huggingface.co/FastVideo/FastMetal-1.3B-QAD) | `mlx_wan_prompt_to_video.py` | 16 GB+ |
|
||||
| [`FastVideo/FastMetal-5B-QAD`](https://huggingface.co/FastVideo/FastMetal-5B-QAD) | `mlx_wan22_generate.py` | 16 GB+ |
|
||||
| [`FastVideo/FastMetal-14B-QAD`](https://huggingface.co/FastVideo/FastMetal-14B-QAD) | `mlx_wan_prompt_to_video.py` | 36 GB+ |
|
||||
|
||||
```bash
|
||||
hf download FastVideo/FastMetal-1.3B-QAD --local-dir ./FastMetal-1.3B-QAD
|
||||
|
||||
python examples/inference/basic/mlx_wan_prompt_to_video.py \
|
||||
--model-root ./FastMetal-1.3B-QAD \
|
||||
--mlx-checkpoint ./FastMetal-1.3B-QAD \
|
||||
--height 480 --width 832 --num-frames 81 \
|
||||
--prompt "A bird's-eye view of a misty forest valley at dawn."
|
||||
```
|
||||
|
||||
14B uses the same script. Point both flags at `./FastMetal-14B-QAD`. That repo also ships an EMA variant: keep `--model-root` at the repo root and set `--mlx-checkpoint ./FastMetal-14B-QAD/ema`.
|
||||
|
||||
Wan2.2 5B uses a different latent layout, so it has its own entrypoint:
|
||||
|
||||
```bash
|
||||
hf download FastVideo/FastMetal-5B-QAD --local-dir ./FastMetal-5B-QAD
|
||||
|
||||
python examples/inference/basic/mlx_wan22_generate.py \
|
||||
--mlx-checkpoint ./FastMetal-5B-QAD \
|
||||
--text-encoder-root ./FastMetal-5B-QAD \
|
||||
--vae-root ./FastMetal-5B-QAD/vae \
|
||||
--height 704 --width 1280 --num-frames 81 \
|
||||
--prompt "A cinematic portrait with soft neon lighting and smooth camera motion."
|
||||
```
|
||||
|
||||
CUDA FastWan-QAD (`FastVideo/FastWan-QAD-1.3B`, `FastVideo/FastWan-QAD-FP8-1.3B`) is a separate NVIDIA release. The MLX examples look for FastMetal packed weights (`mlx_dit.json`).
|
||||
|
||||
`basic_mps.py` is a generic PyTorch MPS demo. For local video on Mac, use the FastMetal commands above.
|
||||
|
||||
## Development Environment Setup
|
||||
|
||||
If you're planning to contribute to FastVideo please see the following page:
|
||||
@@ -94,9 +139,9 @@ If you're planning to contribute to FastVideo please see the following page:
|
||||
|
||||
## Hardware Requirements
|
||||
|
||||
### For Basic Inference
|
||||
|
||||
- Mac M1, M2, M3, or M4 (at least 32 GB RAM is preferable for high quality video generation)
|
||||
- **1.3B / 5B:** 16 GB unified memory and up (M1 and later)
|
||||
- **14B:** 36 GB unified memory and up
|
||||
- Fanless 13-inch MacBook Air can run 1.3B and 5B at the same resolutions
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
|
||||
@@ -156,8 +156,11 @@ is power-cycled. To avoid it:
|
||||
|
||||
- **Builds** (flash-attn, kernel): `nice -n 19`, `MAX_JOBS=2`, `nohup`. Never a
|
||||
bare foreground high-parallelism build.
|
||||
- Leave `*_cpu_offload` at the example defaults — "CPU" offload is the *same*
|
||||
unified RAM on the GB10, so the win is tiling + sane resolution, not offloading.
|
||||
- FastVideo automatically disables DiT layerwise/CPU offload and encoder/VAE CPU
|
||||
offload after each worker binds its GB10 device. Do not force those modes back
|
||||
on: "CPU" offload uses the same unified RAM. Multi-GPU FSDP sharding remains
|
||||
available because it partitions weights without parking them in a separate
|
||||
host pool.
|
||||
|
||||
## Gotchas specific to the GB10
|
||||
|
||||
|
||||
@@ -14,6 +14,13 @@ vae_cpu_offload: bool = True
|
||||
pin_cpu_memory: bool = True
|
||||
```
|
||||
|
||||
On unified-memory accelerators such as NVIDIA GB10 and Apple silicon, FastVideo
|
||||
detects the selected device inside each worker and disables all five host-offload
|
||||
modes before loading modules. Host and accelerator allocations share one physical
|
||||
pool there, so offload adds transfers and duplicate residency instead of freeing
|
||||
memory. CUDA FSDP sharding remains enabled when requested; MPS continues to
|
||||
disable FSDP. `pin_cpu_memory` is not an offload mode and is left unchanged.
|
||||
|
||||
## Behavior Explanation
|
||||
|
||||
!!! note
|
||||
|
||||
@@ -88,6 +88,7 @@ runtime on some GPU/shape combinations. To use FA4, install the pinned
|
||||
`flash-attn-4` build (see the `flash-attn-4` source in `pyproject.toml`) and set:
|
||||
|
||||
```bash
|
||||
UV_TORCH_BACKEND=cu130 uv pip install -e ".[fasth3]"
|
||||
export FASTVIDEO_FA4=1
|
||||
```
|
||||
|
||||
@@ -97,6 +98,21 @@ sm90+) and GQA attention (FA4's `pack_gqa` fails to JIT-compile below sm90).
|
||||
On sm90+ both run on FA4. If FA4 is unusable while `FASTVIDEO_FA4=1` is set,
|
||||
FastVideo fails loudly instead of silently falling back.
|
||||
|
||||
MiniMax-H3 can additionally use FA4's packed-varlen entry point for its long,
|
||||
single-sequence dense DiT self-attention:
|
||||
|
||||
```bash
|
||||
export FASTVIDEO_FA4=1
|
||||
export FASTVIDEO_MINIMAX_H3_FA4_PACKED_VARLEN=1
|
||||
```
|
||||
|
||||
This route is inference-only and remains disabled by default. Runtime guards
|
||||
keep masked attention, batch sizes above one, unequal query/key lengths,
|
||||
grad-enabled calls, and NVFP4 on their established paths. It does not apply to
|
||||
the Preview checkpoint's sparse VSA blocks. Packed-varlen changes floating-point
|
||||
reduction order relative to fixed-length FA4, so treat it as a speed/quality
|
||||
evaluation option rather than an exact-parity mode.
|
||||
|
||||
### FP4 Flash Attention 4 (Blackwell only)
|
||||
|
||||
**`FLASH_ATTN`** with **`--nvfp4_fa4`**
|
||||
@@ -337,7 +353,42 @@ Only DiT submodules that declare `_compile_conditions` are compiled
|
||||
(most shipped models). The text encoder and VAE are not compiled by this
|
||||
flag.
|
||||
|
||||
### What to expect
|
||||
### Regional fullgraph compile (experimental)
|
||||
|
||||
`inference_torch_compile` is a stricter, kwargs-free variant that ports the
|
||||
training-side regional compile of
|
||||
[#1718](https://github.com/hao-ai-lab/FastVideo/pull/1718) to inference: the
|
||||
loader wraps each `_compile_conditions` block in
|
||||
`torch.compile(fullgraph=True)` with inductor
|
||||
`options={"emulate_precision_casts": True}` right after the transformer
|
||||
loads. The ordinary compile path keeps its historical compiler-disabled
|
||||
attention boundary by default; regional compile opts in only the compatible
|
||||
attention instances owned by this transformer. MiniMax-H3 VSA is supported
|
||||
only by the inference-only sm_100a tile-64 route
|
||||
(`FASTVIDEO_VSA_SM100A=1` and `VSA_tile_size=64`); the loader probes that
|
||||
route before capture and keeps the transformer eager when the kernel or
|
||||
device is unsupported. Legacy VSA, MiniMax-H3 tile-256 VSA, and the explicit
|
||||
`FASTVIDEO_DISABLE_ATTENTION_COMPILE=1` escape hatch keep the transformer
|
||||
eager with one warning instead of failing mid-denoise.
|
||||
|
||||
```python
|
||||
generator = VideoGenerator.from_pretrained(
|
||||
"MiniMaxAI/MiniMax-H3",
|
||||
inference_torch_compile=True, # or FASTVIDEO_INFERENCE_TORCH_COMPILE=1
|
||||
)
|
||||
```
|
||||
|
||||
Do not combine it with `torch_compile_kwargs['mode']` (the loader injects
|
||||
inductor options, and torch.compile forbids mode+options); it is
|
||||
independent of `enable_torch_compile`, and when both are set the regional
|
||||
compile wins for the DiT.
|
||||
|
||||
### What to expect from generic compile
|
||||
|
||||
The Wan result below measures the existing generic
|
||||
`enable_torch_compile=True` path. It is useful evidence that compile can help,
|
||||
but it is **not** a benchmark or numerical gate for the stricter regional
|
||||
fullgraph path above.
|
||||
|
||||
| Config | Effect |
|
||||
|---|---|
|
||||
@@ -353,6 +404,24 @@ generations with the same input shapes. Always exclude the first
|
||||
(warmup) generation when measuring steady-state latency — measuring the
|
||||
warmup is the most common way to wrongly conclude "compile is slower".
|
||||
|
||||
**Regional MiniMax-H3 accuracy caveat (job 2660).** On one GB200 at
|
||||
768×1344×124, the native 50-point schedule ran exactly 49 transformer
|
||||
forwards. After one warmup, three fixed-prompt/fixed-seed repeats averaged
|
||||
**185.08s → 157.01s end to end** and **174.90s → 147.12s denoising**. Each
|
||||
leg was independently pixel-deterministic, but compiled output did **not**
|
||||
match eager: mean absolute pixel error **20.247/255**, PSNR **16.67 dB**,
|
||||
mean SSIM **0.7108**, and mean MS-SSIM **0.6370** across 124 frames. Treat
|
||||
regional MiniMax-H3 compile as an opt-in performance experiment, not an
|
||||
eager-parity-safe mode.
|
||||
|
||||
The same caveat applies to sparse MiniMax-H3 regional compile. Its mask
|
||||
compaction, sm_100a launch, trained compression gates, and inference-only H3
|
||||
fusions are fullgraph-compatible, but compilation can still change model
|
||||
numerics. The FastH3 `all` profile enables regional compile by default because
|
||||
it is the fastest measured route at 124, 243, and 345 frames. Use
|
||||
`--no-inference-torch-compile` when comparing against the eager sparse-DiT
|
||||
route.
|
||||
|
||||
**Numerics.** Inductor's lowering is designed to preserve eager
|
||||
semantics within floating-point tolerance, but per-model equivalence is
|
||||
not asserted by any standing SSIM regression here — the SSIM tests in
|
||||
@@ -371,9 +440,14 @@ MS-SSIM gate on *your* config, especially when combining
|
||||
hao-ai-lab/FastVideo#1365 — keep that fix to get a clean compiled
|
||||
region under the default offload path.
|
||||
- **`mode="reduce-overhead"` / CUDA graphs**: not yet supported
|
||||
end-to-end. The attention dispatch is an untraceable custom op and
|
||||
still breaks the graph, which CUDA-graph trees cannot span. Use the
|
||||
default inductor mode (shown above) until that is resolved.
|
||||
by regional compile because that path injects inductor `options`, and
|
||||
PyTorch rejects `mode` together with `options`. The generic compile path
|
||||
can accept `mode`, but CUDA-graph compatibility remains backend- and
|
||||
shape-dependent. FA2/FA3 inference and FA4 expose traceable custom-op
|
||||
boundaries. MiniMax-H3's sm_100a tile-64 inference route is the only VSA
|
||||
path in the regional support envelope; other VSA paths and the FA3
|
||||
grad-enabled path remain outside it. Use the default inductor mode shown
|
||||
above unless your exact configuration has its own gate.
|
||||
|
||||
Extra `torch.compile` options are passed through `torch_compile_kwargs`
|
||||
(a dict), accepted by `VideoGenerator.from_pretrained(...)` and by the
|
||||
|
||||
@@ -182,12 +182,14 @@ optimizations: absence means **untested**, not incompatible.
|
||||
|
||||
| Release path | Model | Mode | Validated hardware | Status |
|
||||
| --- | --- | --- | --- | --- |
|
||||
| MLX FastWan T2V | FastWan-QAD-INT8-1.3B `[release model ID pending]` | 480x832, 81 frames, 3-step DMD, INT8 DiT + TAEHV decode | Apple M4 Max, 36 GB unified-memory class, MLX 0.31.2 | Release candidate; requires release-owner visual sign-off |
|
||||
| MLX FastMetal T2V 1.3B | [`FastVideo/FastMetal-1.3B-QAD`](https://huggingface.co/FastVideo/FastMetal-1.3B-QAD) | 480x832, 81 frames, 3-step DMD, INT8 DiT + TAEHV decode | Apple M4 Max, 16 GB+ unified memory | Released |
|
||||
| MLX FastMetal TI2V 5B | [`FastVideo/FastMetal-5B-QAD`](https://huggingface.co/FastVideo/FastMetal-5B-QAD) | 480p / 720p, 81 frames, 3-step DMD, INT8 DiT + TAEHV decode | Apple M4 Max, 16 GB+ unified memory | Released |
|
||||
| MLX FastMetal T2V 14B | [`FastVideo/FastMetal-14B-QAD`](https://huggingface.co/FastVideo/FastMetal-14B-QAD) | 480p / 720p, 81 frames, 3-step DMD, INT8 DiT + TAEHV decode | Apple M4 Max, 36 GB+ unified memory | Released |
|
||||
|
||||
This is a text-to-video-only source-install release. It is validated on the
|
||||
hardware listed above; MLX allocator caps are not evidence of support for a
|
||||
physical 16 GB Mac. See [Apple Silicon FastWan](../getting_started/installation/mps.md)
|
||||
for the supported command and release gates.
|
||||
Apple Silicon uses FastMetal-QAD. CUDA FastWan-QAD (`FastVideo/FastWan-QAD-1.3B`,
|
||||
`FastVideo/FastWan-QAD-FP8-1.3B`) is the NVIDIA release. See the
|
||||
[Apple Silicon guide](../getting_started/installation/mps.md) and the
|
||||
[FastMetal-QAD blog](https://haoailab.com/blogs/fastmetal/).
|
||||
|
||||
**Note**: Wan2.2 TI2V 5B has some quality issues when performing I2V generation. We are working on fixing this issue.
|
||||
|
||||
@@ -216,9 +218,10 @@ Per the installation guides:
|
||||
[GPU install guide](../getting_started/installation/gpu.md).
|
||||
- **NVIDIA DGX Spark (GB10, aarch64)** — CUDA 13, from-source kernel build; see
|
||||
the [DGX Spark install guide](../getting_started/installation/spark.md).
|
||||
- **Apple silicon (MPS)** — macOS 14 or newer; see the
|
||||
[MPS install guide](../getting_started/installation/mps.md) and
|
||||
[`basic_mps.py`](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_mps.py).
|
||||
- **Apple silicon** — macOS 14 or newer; FastMetal-QAD via the MLX runtime. See the
|
||||
[Apple Silicon guide](../getting_started/installation/mps.md). The older
|
||||
[`basic_mps.py`](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_mps.py)
|
||||
demo is PyTorch MPS only.
|
||||
|
||||
Optimization-specific hardware constraints (e.g. STA requiring Hopper) are
|
||||
listed under [Special requirements](#special-requirements).
|
||||
|
||||
@@ -18,11 +18,27 @@ git clone https://github.com/hao-ai-lab/FastVideo.git && cd FastVideo
|
||||
python examples/inference/basic/basic.py
|
||||
```
|
||||
|
||||
For an example on Apple silicon:
|
||||
```
|
||||
python examples/inference/basic/basic_mps.py
|
||||
### Apple Silicon (FastMetal-QAD)
|
||||
|
||||
Use the MLX runtime with FastMetal-QAD. See the
|
||||
[Apple Silicon guide](https://hao-ai-lab.github.io/FastVideo/getting_started/installation/mps/).
|
||||
|
||||
```bash
|
||||
hf download FastVideo/FastMetal-1.3B-QAD --local-dir ./FastMetal-1.3B-QAD
|
||||
|
||||
python examples/inference/basic/mlx_wan_prompt_to_video.py \
|
||||
--model-root ./FastMetal-1.3B-QAD \
|
||||
--mlx-checkpoint ./FastMetal-1.3B-QAD \
|
||||
--prompt "A bird's-eye view of a misty forest valley at dawn."
|
||||
```
|
||||
|
||||
5B uses
|
||||
[`mlx_wan22_generate.py`](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/mlx_wan22_generate.py)
|
||||
with
|
||||
`FastVideo/FastMetal-5B-QAD`.
|
||||
|
||||
`examples/inference/basic/basic_mps.py` is the older PyTorch MPS demo.
|
||||
|
||||
For an example running DMD+VSA inference:
|
||||
```
|
||||
python examples/inference/basic/basic_dmd.py
|
||||
@@ -33,13 +49,95 @@ For the typed config/request path added during the inference API refactor:
|
||||
python examples/inference/basic/basic_dmd_new_api.py
|
||||
```
|
||||
|
||||
For the few-step (4-step, DMD2-distilled) MiniMax-H3 preview, generating synchronized video and audio, optionally with block-sparse VSA attention:
|
||||
```
|
||||
python examples/inference/basic/basic_fasth3.py --prompt "your prompt" [--vsa-sparsity 0.9]
|
||||
```
|
||||
The default checkpoint `FastVideo/FastVideo-Minimax-FastH3-Preview-v0.1` is private on the Hub while its license review completes; until it flips public, pass `--model-path` with a local snapshot of the release.
|
||||
### FastH3 Preview
|
||||
|
||||
On Blackwell (sm_100) GPUs with a `fastvideo-kernel` build that carries the sm_100a block-sparse extension, `--vsa-kernel sm100a` routes the tile-64 attention forwards through the CUDA kernel instead of Triton (it sets `FASTVIDEO_VSA_SM100A=1` before the pipeline boots); if the extension or the arch is missing, the run warns once and falls back to Triton.
|
||||
The verified [basic FastH3 example](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_fasth3.py)
|
||||
runs the few-step (4-forward, DMD2-distilled) MiniMax-H3 preview, generating
|
||||
synchronized video and audio with its trained block-sparse VSA attention:
|
||||
|
||||
```bash
|
||||
UV_TORCH_BACKEND=cu130 uv pip install -e ".[fasth3]"
|
||||
```
|
||||
|
||||
This installs the pinned FA4 CuTe package and FastVideo kernel release used by
|
||||
the measured GB200 profile. Then run:
|
||||
|
||||
```
|
||||
python examples/inference/basic/basic_fasth3.py --prompt "your prompt"
|
||||
```
|
||||
The default checkpoint, [FastH3 Preview v0.2](https://huggingface.co/FastVideo/FastVideo-Minimax-FastH3-Preview-v0.2), is public on the Hub under the MiniMax H3 Community License. Review its model card and license before use or redistribution.
|
||||
|
||||
The default `all` profile is the fastest measured four-GPU Preview recipe on GB200. It selects VSA sparsity 0.9 with 64-token tiles and the sm_100a sparse kernel, enables FA4 for eligible non-VSA paths, regionally compiles and replicates the sparse DiT, compiles and temporally parallelizes the video VAE with the `gather` strategy, and pins CPU-offloaded component memory. It also pins the benchmark protocol: five sigma-grid points (exactly four DiT forwards), one excluded seed-999 warmup, then three timed seed-1000 requests with distinct output paths.
|
||||
|
||||
The equivalent explicit command is:
|
||||
|
||||
```bash
|
||||
python examples/inference/basic/basic_fasth3.py \
|
||||
--prompt "your prompt" \
|
||||
--profile all \
|
||||
--num-gpus 4 \
|
||||
--steps 5 \
|
||||
--vsa-sparsity 0.9 \
|
||||
--vsa-tile-size 64 \
|
||||
--vsa-kernel sm100a \
|
||||
--compile-vae \
|
||||
--parallel-vae \
|
||||
--replicated-dit \
|
||||
--pin-cpu-memory \
|
||||
--fa4 \
|
||||
--no-torch-compile \
|
||||
--inference-torch-compile \
|
||||
--ulysses-a2a off \
|
||||
--warmup \
|
||||
--repeats 3 \
|
||||
--seed 1000 \
|
||||
--warmup-seed 999
|
||||
```
|
||||
|
||||
`all` enables the inference-only H3 fusions and regional compile. Both can change floating-point operation order, so this is a report-only performance profile rather than an exact-parity route. Use `--profile strict` to disable the H3 fusions while preserving regional compile, or `--profile strict --no-inference-torch-compile` for the eager strict route. Individual `--no-*` switches are available for portability and attribution; in particular, use `--vsa-kernel triton --no-fa4` if the Blackwell kernels are unavailable. The script preserves the warmup and each measured video under distinct paths, then prints per-request wall time plus a warmup-excluded median.
|
||||
|
||||
One script covers each validated duration; regional compile is the fastest
|
||||
measured DiT route for all three:
|
||||
|
||||
```bash
|
||||
# 5 s
|
||||
python examples/inference/basic/basic_fasth3.py \
|
||||
--prompt "your prompt" --output outputs/fasth3_5s
|
||||
# 10 s
|
||||
python examples/inference/basic/basic_fasth3.py \
|
||||
--prompt "your prompt" --num-frames 243 --output outputs/fasth3_10s
|
||||
# 15 s
|
||||
python examples/inference/basic/basic_fasth3.py \
|
||||
--prompt "your prompt" --num-frames 345 --output outputs/fasth3_15s
|
||||
```
|
||||
|
||||
Pass `--no-inference-torch-compile` to recover the eager sparse-DiT route.
|
||||
|
||||
### FastH3 Preview LoRAs
|
||||
|
||||
The LoRA release runs on top of `MiniMaxAI/MiniMax-H3` with the same default
|
||||
compile, fusion, FA4, VSA, and parallel-VAE profile as the full FastH3 example:
|
||||
|
||||
```bash
|
||||
bash examples/inference/basic/run_fasth3_lora_preview_vsa_datafree.sh \
|
||||
--prompt "your prompt"
|
||||
```
|
||||
|
||||
The four release launchers are:
|
||||
|
||||
- `run_fasth3_lora_preview_vsa_datafree.sh`
|
||||
- `run_fasth3_lora_preview_vsa_synthetic_step1300.sh`
|
||||
- `run_fasth3_lora_preview_vsa_synthetic_step1900.sh`
|
||||
- `run_fasth3_lora_preview_dense_datafree.sh`
|
||||
|
||||
Each downloads its exact private adapter file from
|
||||
`FastVideo/FastVideo-FastH3-4-step-Preview-v1-LoRA`; authenticate with `hf auth
|
||||
login` first. Pass `--lora-strength 0.5` to interpolate every adapter payload at
|
||||
half strength. Strength `1` applies the published rank-64 adapter at its trained
|
||||
scale and approximates the full student; `0` removes its weight deltas. VSA
|
||||
launchers still use sparse attention at strength `0` and require FastVideo's
|
||||
tile-64 VSA kernel; the dense launcher selects FA4. Each launcher writes to its
|
||||
own variant directory by default so comparison outputs do not collide.
|
||||
|
||||
## Basic Walkthrough
|
||||
|
||||
|
||||
@@ -1,27 +1,32 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""Few-step video+audio generation with the DMD2-distilled MiniMax H3 preview.
|
||||
|
||||
FastVideo/FastVideo-Minimax-FastH3-Preview-v0.1 is a 4-step distillation of
|
||||
MiniMaxAI/MiniMax-H3 (data-free DMD2): it walks a 4-step grid on the release
|
||||
sampler's shift-12 schedule instead of the base model's 50 steps, generating
|
||||
synchronized video and audio in one pipeline call.
|
||||
The default ``all`` profile reproduces the fastest measured FastH3 Preview
|
||||
recipe on four GB200 GPUs. It runs the checkpoint's native five-point sigma
|
||||
grid (exactly four DiT forwards), trained VSA policy, Blackwell sparse kernel,
|
||||
regional fullgraph DiT compile, compiled/parallel video VAE, and inference-only
|
||||
H3 fusions. One compile warmup is excluded before three measured requests.
|
||||
|
||||
The student was trained with block-sparse video attention (VSA, 64-token
|
||||
tiles) and its checkpoint carries the trained sparse-gate parameters
|
||||
(``attn.to_gate_compress``), so this script always runs the VSA-H3 attention
|
||||
backend. At the default ``--vsa-sparsity 0.0`` the attention math is exactly
|
||||
dense (every tile is selected); raise the sparsity for additional speedup.
|
||||
Both regional compile and the default fusions can change floating-point
|
||||
operation order, so ``all`` is a report-only performance profile.
|
||||
``--profile strict`` disables the H3 fusions but preserves regional compile;
|
||||
combine it with ``--no-inference-torch-compile`` for the eager strict route.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import importlib.util
|
||||
import os
|
||||
import statistics
|
||||
import time
|
||||
from collections.abc import Sequence
|
||||
from pathlib import Path
|
||||
|
||||
from fastvideo import VideoGenerator
|
||||
from fastvideo.api import (
|
||||
CompileConfig,
|
||||
ComponentConfig,
|
||||
EngineConfig,
|
||||
GenerationRequest,
|
||||
GeneratorConfig,
|
||||
@@ -32,149 +37,331 @@ from fastvideo.api import (
|
||||
SamplingConfig,
|
||||
)
|
||||
|
||||
DEFAULT_MODEL = "FastVideo/FastVideo-Minimax-FastH3-Preview-v0.2"
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("--model-path", default="FastVideo/FastVideo-Minimax-FastH3-Preview-v0.1")
|
||||
# The HF repo is private while the MiniMax H3 Community License review
|
||||
# completes; until it flips public, pass --model-path with a local
|
||||
# snapshot of the release instead (e.g. the team export at
|
||||
# /mnt/lustre/vlm-wlsaidhi/fastvideo/exports/FastVideo-Minimax-FastH3-Preview-v0.1).
|
||||
|
||||
def build_parser(description: str | None = None) -> argparse.ArgumentParser:
|
||||
"""Build the shared FastH3 preview CLI used by full and LoRA checkpoints."""
|
||||
parser = argparse.ArgumentParser(description=description or __doc__)
|
||||
parser.add_argument("--model-path", default=DEFAULT_MODEL)
|
||||
# The HF repo may require authentication while the MiniMax H3 Community
|
||||
# License review completes. A local snapshot can be passed here instead.
|
||||
parser.add_argument("--prompt", required=True)
|
||||
parser.add_argument("--output", default="outputs/fasth3")
|
||||
parser.add_argument("--profile",
|
||||
choices=("all", "strict"),
|
||||
default="all",
|
||||
help="all enables the fastest measured, non-parity H3 fusions; strict disables only them")
|
||||
parser.add_argument("--height", type=int, default=768)
|
||||
parser.add_argument("--width", type=int, default=1344)
|
||||
parser.add_argument("--num-frames", type=int, default=124)
|
||||
# num_inference_steps counts sigma-GRID POINTS, matching the base model's
|
||||
# convention ("50 steps" = a 50-point grid = 49 transformer forwards). The
|
||||
# student's distilled 4-step grid is 4 FORWARDS, i.e. a 5-point grid
|
||||
# (t = 1000, 750, 500, 250 -> 0 on the shift-12 schedule) — so the correct
|
||||
# default here is 5. Other grids are off-distribution.
|
||||
# num_inference_steps counts sigma-GRID POINTS. The distilled schedule is
|
||||
# t=1000,750,500,250 -> 0: five points and exactly four DiT forwards.
|
||||
parser.add_argument("--steps",
|
||||
type=int,
|
||||
default=5,
|
||||
help="num_inference_steps = sigma-grid points; N points run N-1 denoising "
|
||||
"forwards. 5 (default) is the distilled 4-forward grid")
|
||||
parser.add_argument("--seed", type=int, default=0)
|
||||
help="sigma-grid points; N points run N-1 DiT forwards (the trained default is 5)")
|
||||
parser.add_argument("--seed", type=int, default=1000, help="seed reused for every measured request")
|
||||
parser.add_argument("--warmup-seed", type=int, default=999)
|
||||
parser.add_argument("--repeats", type=int, default=3, help="number of measured requests after warmup")
|
||||
parser.add_argument("--warmup",
|
||||
action=argparse.BooleanOptionalAction,
|
||||
default=True,
|
||||
help="run one excluded request before timing")
|
||||
parser.add_argument("--num-gpus", type=int, default=4)
|
||||
parser.add_argument("--vsa-sparsity",
|
||||
type=float,
|
||||
default=0.0,
|
||||
help="Run-level VSA sparsity in [0, 1). 0.0 (default) selects every tile, which is "
|
||||
"exactly dense attention; the student was trained at 0.9")
|
||||
# 64 is the trained contract: the student was TRAINED with 64-token
|
||||
# (4,4,4) tiles, and its to_gate_compress gates were learned against
|
||||
# pooling at that granularity — keep 64 unless you are ablating.
|
||||
default=0.9,
|
||||
help="run-level VSA sparsity in [0, 1); 0.9 is the checkpoint's trained policy")
|
||||
parser.add_argument("--vsa-tile-size",
|
||||
type=int,
|
||||
choices=(64, 256),
|
||||
default=64,
|
||||
help="VSA-H3 tile size in tokens; 64 (default) is what the student was trained "
|
||||
"with and runs the native Triton block-sparse path, 256 is the FA4-CuTe-capable "
|
||||
"geometry for ablations")
|
||||
help="VSA-H3 tile size; 64 is the checkpoint's trained and measured geometry")
|
||||
parser.add_argument("--vsa-kernel",
|
||||
choices=("triton", "sm100a"),
|
||||
default="triton",
|
||||
help="Block-sparse kernel for the tile-64 attention forward: triton (default, "
|
||||
"portable fwd+bwd) or sm100a — the opt-in Blackwell CUDA forward "
|
||||
"(fastvideo_kernel.block_sparse_attn_sm100a). sm100a needs an sm_100 GPU and a "
|
||||
"fastvideo-kernel build that carries the extension; if a precondition fails at "
|
||||
"run time the attention layer logs one warning and falls back to Triton. Only "
|
||||
"meaningful with --vsa-tile-size 64")
|
||||
parser.add_argument("--torch-compile", action="store_true", help="torch.compile the DiT transformer path")
|
||||
default="sm100a",
|
||||
help="tile-64 sparse kernel; sm100a is the measured GB200 route and requires a compatible "
|
||||
"fastvideo-kernel build")
|
||||
parser.add_argument("--fa4",
|
||||
action=argparse.BooleanOptionalAction,
|
||||
default=True,
|
||||
help="use FA4 for eligible non-VSA attention paths")
|
||||
parser.add_argument("--h3-fusions",
|
||||
action=argparse.BooleanOptionalAction,
|
||||
default=None,
|
||||
help="override the profile's H3 fusion policy (changes model numerics when enabled)")
|
||||
parser.add_argument("--compile-vae",
|
||||
action=argparse.BooleanOptionalAction,
|
||||
default=True,
|
||||
help="compile the video VAE decoder independently of the DiT")
|
||||
parser.add_argument("--parallel-vae",
|
||||
action=argparse.BooleanOptionalAction,
|
||||
default=True,
|
||||
help="round-robin VAE temporal chunks across sequence-parallel ranks")
|
||||
parser.add_argument("--replicated-dit",
|
||||
action=argparse.BooleanOptionalAction,
|
||||
default=True,
|
||||
help="replicate DiT weights instead of FSDP-sharding them")
|
||||
parser.add_argument("--pin-cpu-memory",
|
||||
action=argparse.BooleanOptionalAction,
|
||||
default=True,
|
||||
help="pin CPU-offloaded text-encoder and VAE weights")
|
||||
parser.add_argument("--torch-compile",
|
||||
action=argparse.BooleanOptionalAction,
|
||||
default=False,
|
||||
help="compile the whole DiT path (off in the fastest FastH3 profile)")
|
||||
parser.add_argument("--inference-torch-compile",
|
||||
action=argparse.BooleanOptionalAction,
|
||||
default=True,
|
||||
help="regionally compile DiT blocks (enabled in the fastest FastH3 profile)")
|
||||
parser.add_argument("--ulysses-a2a",
|
||||
choices=("off", "auto"),
|
||||
default="off",
|
||||
help="sequence-parallel all-to-all route; off reproduces the fastest FastH3 profile, while "
|
||||
"auto opts into the fused NVLink kernel when the installed kernel package supports it")
|
||||
parser.add_argument("--compile-mode",
|
||||
default=None,
|
||||
help='torch.compile mode, e.g. "reduce-overhead" for CUDA graphs')
|
||||
parser.add_argument("--repeats",
|
||||
type=int,
|
||||
default=1,
|
||||
help="generate N times; with --torch-compile the first run pays "
|
||||
"compilation, so steady-state is the last repeat")
|
||||
return parser.parse_args()
|
||||
help='whole-DiT torch.compile mode, e.g. "reduce-overhead"; requires '
|
||||
"--no-inference-torch-compile")
|
||||
return parser
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
def validate_args(parser: argparse.ArgumentParser, args: argparse.Namespace) -> argparse.Namespace:
|
||||
if args.repeats < 1:
|
||||
parser.error("--repeats must be at least 1")
|
||||
if args.num_gpus < 1:
|
||||
parser.error("--num-gpus must be at least 1")
|
||||
if not 0.0 <= args.vsa_sparsity < 1.0:
|
||||
parser.error("--vsa-sparsity must be in [0, 1)")
|
||||
if args.compile_mode is not None and args.inference_torch_compile:
|
||||
parser.error("--compile-mode cannot be combined with regional compile; pass --no-inference-torch-compile")
|
||||
return args
|
||||
|
||||
|
||||
def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
|
||||
parser = build_parser()
|
||||
return validate_args(parser, parser.parse_args(argv))
|
||||
|
||||
|
||||
def _uses_vsa(args: argparse.Namespace) -> bool:
|
||||
"""Full FastH3 checkpoints use VSA; LoRA previews may select dense attention."""
|
||||
return bool(getattr(args, "vsa", True))
|
||||
|
||||
|
||||
def _h3_fusions_enabled(args: argparse.Namespace) -> bool:
|
||||
if args.h3_fusions is not None:
|
||||
return bool(args.h3_fusions)
|
||||
return args.profile == "all"
|
||||
|
||||
|
||||
def profile_environment(args: argparse.Namespace) -> dict[str, str | None]:
|
||||
"""Return the complete boot-time environment for this profile.
|
||||
|
||||
``None`` means the variable must be removed. Values are explicit even for
|
||||
disabled features so a shell's inherited experiment settings cannot
|
||||
silently change the advertised profile.
|
||||
"""
|
||||
use_vsa = _uses_vsa(args)
|
||||
return {
|
||||
"FASTVIDEO_ATTENTION_BACKEND": "VIDEO_SPARSE_ATTN_H3" if use_vsa else "FLASH_ATTN",
|
||||
"FASTVIDEO_VSA_SM100A": "1" if use_vsa and args.vsa_kernel == "sm100a" else "0",
|
||||
"FASTVIDEO_VSA_CUTEDSL": "0",
|
||||
# A non-empty output path enables the diagnostic probe.
|
||||
"FASTVIDEO_H3_VSA_PROBE": None,
|
||||
"FASTVIDEO_DISABLE_ATTENTION_COMPILE": "0",
|
||||
"FASTVIDEO_FA4": "1" if args.fa4 else "0",
|
||||
"FASTVIDEO_NVFP4_FA4": "0",
|
||||
"FASTVIDEO_MINIMAX_H3_FA4_PACKED_VARLEN": "0",
|
||||
"FASTVIDEO_MINIMAX_H3_FUSIONS": "all" if _h3_fusions_enabled(args) else "0",
|
||||
"FASTVIDEO_INFERENCE_TORCH_COMPILE": "1" if args.inference_torch_compile else "0",
|
||||
"FASTVIDEO_VAE_PARALLEL_DECODE": "1" if args.parallel_vae else "0",
|
||||
"FASTVIDEO_VAE_PARALLEL_ENCODE": "0",
|
||||
"FASTVIDEO_VAE_PARALLEL_DECODE_STRATEGY": "gather",
|
||||
"FASTVIDEO_ULYSSES_A2A": args.ulysses_a2a,
|
||||
"FASTVIDEO_STAGE_LOGGING": "1",
|
||||
}
|
||||
|
||||
|
||||
def configure_environment(args: argparse.Namespace) -> dict[str, str | None]:
|
||||
environment = profile_environment(args)
|
||||
for name, value in environment.items():
|
||||
if value is None:
|
||||
os.environ.pop(name, None)
|
||||
else:
|
||||
os.environ[name] = value
|
||||
return environment
|
||||
|
||||
|
||||
def _fa4_is_installed() -> bool:
|
||||
try:
|
||||
return importlib.util.find_spec("flash_attn.cute") is not None
|
||||
except (ImportError, ModuleNotFoundError):
|
||||
return False
|
||||
|
||||
|
||||
def _sm100a_kernel_is_installed() -> bool:
|
||||
try:
|
||||
from fastvideo_kernel import block_sparse_attn_sm100a
|
||||
except ImportError:
|
||||
return False
|
||||
return bool(getattr(block_sparse_attn_sm100a, "_HAS_VSA_SM100A", False))
|
||||
|
||||
|
||||
def validate_profile_dependencies(args: argparse.Namespace) -> None:
|
||||
"""Fail before model loading when the selected measured route is absent."""
|
||||
if args.fa4 and not _fa4_is_installed():
|
||||
raise RuntimeError(
|
||||
"FastH3's FA4 profile requires the pinned flash-attn-4 package. Install it with "
|
||||
"`UV_TORCH_BACKEND=cu130 uv pip install -e \".[fasth3]\"`, or pass --no-fa4.")
|
||||
if _uses_vsa(args) and args.vsa_kernel == "sm100a" and not _sm100a_kernel_is_installed():
|
||||
raise RuntimeError(
|
||||
"FastH3's sm100a profile requires fastvideo-kernel 0.3.4 built with the Blackwell VSA extension. "
|
||||
"Install this checkout with `UV_TORCH_BACKEND=cu130 uv pip install -e \".[fasth3]\"` (or run "
|
||||
"`cd fastvideo-kernel && ./build.sh`), or pass --vsa-kernel triton.")
|
||||
|
||||
|
||||
def build_generator_config(args: argparse.Namespace) -> GeneratorConfig:
|
||||
use_vsa = _uses_vsa(args)
|
||||
experimental: dict[str, object] = {
|
||||
"attention_backend": "VIDEO_SPARSE_ATTN_H3" if use_vsa else "FLASH_ATTN",
|
||||
"inference_torch_compile": args.inference_torch_compile,
|
||||
"vae_parallel_decode": args.parallel_vae,
|
||||
"vae_parallel_decode_strategy": "gather",
|
||||
}
|
||||
if use_vsa:
|
||||
experimental.update({
|
||||
"VSA_sparsity": args.vsa_sparsity,
|
||||
"VSA_tile_size": args.vsa_tile_size,
|
||||
})
|
||||
return GeneratorConfig(
|
||||
model_path=args.model_path,
|
||||
pipeline=PipelineSelection(
|
||||
components=ComponentConfig(
|
||||
lora_path=getattr(args, "lora_path", None),
|
||||
lora_strength=float(getattr(args, "lora_strength", 1.0)),
|
||||
),
|
||||
experimental=experimental,
|
||||
),
|
||||
engine=EngineConfig(
|
||||
num_gpus=args.num_gpus,
|
||||
use_fsdp_inference=args.num_gpus > 1 and not args.replicated_dit,
|
||||
parallelism=ParallelismConfig(tp_size=1, sp_size=args.num_gpus),
|
||||
offload=OffloadConfig(
|
||||
dit=False,
|
||||
dit_layerwise=False,
|
||||
text_encoder=True,
|
||||
vae=True,
|
||||
pin_cpu_memory=args.pin_cpu_memory,
|
||||
),
|
||||
compile=CompileConfig(
|
||||
enabled=args.torch_compile,
|
||||
mode=args.compile_mode,
|
||||
vae_enabled=args.compile_vae,
|
||||
),
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def build_request(args: argparse.Namespace, output_path: Path, seed: int) -> GenerationRequest:
|
||||
return GenerationRequest(
|
||||
prompt=args.prompt,
|
||||
negative_prompt="",
|
||||
sampling=SamplingConfig(
|
||||
height=args.height,
|
||||
width=args.width,
|
||||
num_frames=args.num_frames,
|
||||
fps=24,
|
||||
num_inference_steps=args.steps,
|
||||
# MiniMax-H3 is guidance-distilled; FastH3 inherits that contract.
|
||||
guidance_scale=1.0,
|
||||
batch_cfg=False,
|
||||
seed=seed,
|
||||
),
|
||||
output=OutputConfig(
|
||||
output_path=str(output_path),
|
||||
save_video=True,
|
||||
return_frames=False,
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def _actual_output_path(result: object, requested: Path) -> Path:
|
||||
video_path = getattr(result, "video_path", None)
|
||||
return Path(video_path) if video_path else requested
|
||||
|
||||
|
||||
def _denoise_seconds(result: object) -> float | None:
|
||||
stages = getattr(getattr(result, "logging_info", None), "stages", None)
|
||||
if not stages:
|
||||
return None
|
||||
for stage_name, metrics in stages.items():
|
||||
if "denois" not in stage_name.lower():
|
||||
continue
|
||||
execution_time = metrics.get("execution_time")
|
||||
return float(execution_time) if execution_time is not None else None
|
||||
return None
|
||||
|
||||
|
||||
def run(args: argparse.Namespace) -> list[float]:
|
||||
output_dir = Path(args.output)
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
environment = configure_environment(args)
|
||||
validate_profile_dependencies(args)
|
||||
|
||||
if args.vsa_kernel == "sm100a":
|
||||
# The attention backend reads FASTVIDEO_VSA_SM100A per forward; set it
|
||||
# before the pipeline boots so spawned GPU workers inherit it. The
|
||||
# kernel is forward-only and inference runs under no-grad, so every
|
||||
# denoising forward qualifies for the CUDA route.
|
||||
os.environ["FASTVIDEO_VSA_SM100A"] = "1"
|
||||
print(f"Profile: {args.profile} ({'non-parity fusions' if _h3_fusions_enabled(args) else 'fusions off'})")
|
||||
print(f"Output directory: {output_dir.resolve()}")
|
||||
print("Denoising contract: 5 sigma points = 4 DiT forwards" if args.steps == 5 else
|
||||
f"Denoising contract override: {args.steps} sigma points = {args.steps - 1} DiT forwards")
|
||||
print("Profile environment: " + " ".join(f"{key}={value if value is not None else '<unset>'}"
|
||||
for key, value in environment.items()))
|
||||
|
||||
# Boot-time run configuration, folded into FastVideoArgs (the same route
|
||||
# examples/inference/basic/basic_minimax_h3_t2v.py uses for sparsity):
|
||||
# - attention_backend: the checkpoint carries trained to_gate_compress
|
||||
# gates, which only exist under the VSA-H3 backend — a dense-backend
|
||||
# load would reject them as unexpected weights. Layers that do not
|
||||
# support VSA-H3 (e.g. the token refiner) fall back to flash attention.
|
||||
# - VSA_tile_size: forwarded even at sparsity 0.0 because the gate-compress
|
||||
# branch pools per tile, and the gates were trained at 64 tokens/tile.
|
||||
experimental: dict[str, object] = {
|
||||
"attention_backend": "VIDEO_SPARSE_ATTN_H3",
|
||||
"VSA_tile_size": args.vsa_tile_size,
|
||||
}
|
||||
if args.vsa_sparsity > 0.0:
|
||||
experimental["VSA_sparsity"] = args.vsa_sparsity
|
||||
|
||||
generator = VideoGenerator.from_config(
|
||||
GeneratorConfig(
|
||||
model_path=args.model_path,
|
||||
pipeline=PipelineSelection(experimental=experimental),
|
||||
engine=EngineConfig(
|
||||
num_gpus=args.num_gpus,
|
||||
use_fsdp_inference=args.num_gpus > 1,
|
||||
parallelism=ParallelismConfig(tp_size=1, sp_size=args.num_gpus),
|
||||
offload=OffloadConfig(
|
||||
dit=False,
|
||||
dit_layerwise=False,
|
||||
text_encoder=True,
|
||||
vae=True,
|
||||
pin_cpu_memory=False,
|
||||
),
|
||||
compile=CompileConfig(
|
||||
enabled=args.torch_compile,
|
||||
mode=args.compile_mode,
|
||||
),
|
||||
),
|
||||
))
|
||||
generator = VideoGenerator.from_config(build_generator_config(args))
|
||||
measured_wall_times: list[float] = []
|
||||
measured_denoise_times: list[float] = []
|
||||
try:
|
||||
request = GenerationRequest(
|
||||
prompt=args.prompt,
|
||||
negative_prompt="",
|
||||
sampling=SamplingConfig(
|
||||
height=args.height,
|
||||
width=args.width,
|
||||
num_frames=args.num_frames,
|
||||
fps=24,
|
||||
num_inference_steps=args.steps,
|
||||
# the base model is guidance-distilled; the student inherits it
|
||||
guidance_scale=1.0,
|
||||
batch_cfg=False,
|
||||
seed=args.seed,
|
||||
),
|
||||
output=OutputConfig(
|
||||
output_path=str(output_dir / "fasth3.mp4"),
|
||||
save_video=True,
|
||||
return_frames=False,
|
||||
),
|
||||
)
|
||||
result = generator.generate(request)
|
||||
print(f"Output written to: {result.video_path}")
|
||||
if result.generation_time is not None:
|
||||
# machine-readable: benchmark harnesses parse this line to separate
|
||||
# generation from model-load time (last occurrence = steady state)
|
||||
print(f"Generation time: {result.generation_time:.2f}s")
|
||||
for _ in range(args.repeats - 1):
|
||||
result = generator.generate(request)
|
||||
if result.generation_time is not None:
|
||||
print(f"Generation time: {result.generation_time:.2f}s")
|
||||
if args.warmup:
|
||||
warmup_path = output_dir / "_fasth3_warmup.mp4"
|
||||
print(f"[warmup] generating (excluded from timing summary): {warmup_path}")
|
||||
started = time.perf_counter()
|
||||
warmup_result = generator.generate(build_request(args, warmup_path, args.warmup_seed))
|
||||
warmup_wall = time.perf_counter() - started
|
||||
actual_warmup_path = _actual_output_path(warmup_result, warmup_path)
|
||||
print(f"[warmup] wall={warmup_wall:.3f}s (excluded)")
|
||||
print(f"Warmup output written to: {actual_warmup_path}")
|
||||
|
||||
for index in range(1, args.repeats + 1):
|
||||
requested_path = output_dir / f"fasth3_{args.profile}_run_{index:02d}.mp4"
|
||||
print(f"[measured {index}/{args.repeats}] generating: {requested_path}")
|
||||
started = time.perf_counter()
|
||||
result = generator.generate(build_request(args, requested_path, args.seed))
|
||||
wall = time.perf_counter() - started
|
||||
measured_wall_times.append(wall)
|
||||
actual_path = _actual_output_path(result, requested_path)
|
||||
print(f"Output written to: {actual_path}")
|
||||
print(f"E2E wall time: {wall:.3f}s")
|
||||
generation_time = getattr(result, "generation_time", None)
|
||||
if generation_time is not None:
|
||||
print(f"Generation time: {float(generation_time):.3f}s")
|
||||
denoise_time = _denoise_seconds(result)
|
||||
if denoise_time is not None:
|
||||
measured_denoise_times.append(denoise_time)
|
||||
print(f"Denoising time: {denoise_time:.3f}s")
|
||||
|
||||
median = statistics.median(measured_wall_times)
|
||||
print(f"Measured E2E wall times (n={len(measured_wall_times)}, warmup excluded): "
|
||||
f"{[round(value, 3) for value in measured_wall_times]}")
|
||||
print(f"Median E2E wall time: {median:.3f}s")
|
||||
if measured_denoise_times:
|
||||
print(f"Median denoising time: {statistics.median(measured_denoise_times):.3f}s")
|
||||
return measured_wall_times
|
||||
finally:
|
||||
generator.shutdown()
|
||||
|
||||
|
||||
def main() -> None:
|
||||
run(parse_args())
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
||||
@@ -0,0 +1,81 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""Run a FastH3 four-step Preview LoRA with the measured FastVideo defaults.
|
||||
|
||||
This is the LoRA counterpart of ``basic_fasth3.py``. Both routes share the
|
||||
same performance profile: four DiT forwards, regional fullgraph DiT compile,
|
||||
H3 fusions, compiled and sequence-parallel video VAE decode, replicated DiT,
|
||||
pinned CPU offload, FA4, and the sm100a tile-64 kernel for VSA adapters.
|
||||
|
||||
The FastH3 adapters include low-rank factors plus exact dense deltas. Some also
|
||||
provide the VSA compression gate that is absent from the base checkpoint. Pass
|
||||
the adapter at construction so all three payload types receive the same
|
||||
``--lora-strength``. The attention backend is inferred from that payload unless
|
||||
``--vsa`` or ``--no-vsa`` is specified explicitly.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import math
|
||||
from collections.abc import Sequence
|
||||
|
||||
try:
|
||||
from . import basic_fasth3
|
||||
except ImportError:
|
||||
# Direct script execution puts this directory, rather than ``examples``, on
|
||||
# sys.path. Keep both ``python file.py`` and module/importlib use working.
|
||||
import basic_fasth3 # type: ignore[no-redef]
|
||||
|
||||
BASE_MODEL = "MiniMaxAI/MiniMax-H3"
|
||||
|
||||
|
||||
def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
|
||||
parser = basic_fasth3.build_parser(description=__doc__)
|
||||
parser.set_defaults(model_path=BASE_MODEL, output="outputs/fasth3_lora_preview")
|
||||
parser.add_argument(
|
||||
"--lora-path",
|
||||
required=True,
|
||||
help="FastH3 adapter safetensors file or local adapter directory",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--lora-strength",
|
||||
type=float,
|
||||
default=1.0,
|
||||
help="adapter strength; 0 zeros its weights but keeps its backend, and 1 applies its published scale",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--vsa",
|
||||
action=argparse.BooleanOptionalAction,
|
||||
default=None,
|
||||
help="select VSA explicitly; by default it is inferred from the adapter's compression-gate payload",
|
||||
)
|
||||
args = basic_fasth3.validate_args(parser, parser.parse_args(argv))
|
||||
if not math.isfinite(args.lora_strength):
|
||||
parser.error("--lora-strength must be finite")
|
||||
return _resolve_attention_backend(parser, args)
|
||||
|
||||
|
||||
def _resolve_attention_backend(parser: argparse.ArgumentParser, args: argparse.Namespace) -> argparse.Namespace:
|
||||
# Header-only inspection keeps the payload on disk. A replacement compression
|
||||
# gate is an unambiguous VSA requirement; adapters without one default to dense.
|
||||
from fastvideo.models.loader.lora_patch import DenseLoRAPatch
|
||||
|
||||
patch = DenseLoRAPatch.from_adapter(args.lora_path, strength=args.lora_strength)
|
||||
needs_vsa = bool(patch and any("gate_compress" in name for name in patch.replacement_parameters))
|
||||
if args.vsa is None:
|
||||
args.vsa = needs_vsa
|
||||
elif needs_vsa and not args.vsa:
|
||||
parser.error(f"{args.lora_path} provides to_gate_compress and must run with VSA; drop --no-vsa")
|
||||
return args
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
print(f"FastH3 adapter: {args.lora_path}")
|
||||
print(f"LoRA strength: {args.lora_strength:g}")
|
||||
print(f"Attention: {'VSA-H3' if args.vsa else 'dense FA4'}")
|
||||
basic_fasth3.run(args)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -15,6 +15,7 @@ from fastvideo.api import (
|
||||
OffloadConfig,
|
||||
OutputConfig,
|
||||
ParallelismConfig,
|
||||
PipelineSelection,
|
||||
SamplingConfig,
|
||||
)
|
||||
|
||||
@@ -41,6 +42,12 @@ def parse_args() -> argparse.Namespace:
|
||||
parser.add_argument("--compile-mode",
|
||||
default=None,
|
||||
help='torch.compile mode, e.g. "reduce-overhead" for CUDA graphs')
|
||||
parser.add_argument("--inference-torch-compile",
|
||||
action="store_true",
|
||||
help="regional fullgraph torch.compile of each DiT block after load (the #1718 "
|
||||
"training-port semantics: no kwargs; fullgraph + emulate_precision_casts injected). "
|
||||
"First generation pays the inductor JIT (~1-2 min); use --repeats >= 2 and time "
|
||||
"the last repeat. FASTVIDEO_INFERENCE_TORCH_COMPILE=1 is equivalent")
|
||||
parser.add_argument("--repeats",
|
||||
type=int,
|
||||
default=1,
|
||||
@@ -54,9 +61,16 @@ def main() -> None:
|
||||
output_dir = Path(args.output)
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Boot-time run configuration folded into FastVideoArgs (the same
|
||||
# experimental-dict route basic_fasth3.py uses for the VSA knobs).
|
||||
experimental: dict[str, object] = {}
|
||||
if args.inference_torch_compile:
|
||||
experimental["inference_torch_compile"] = True
|
||||
|
||||
generator = VideoGenerator.from_config(
|
||||
GeneratorConfig(
|
||||
model_path=args.model_path,
|
||||
pipeline=PipelineSelection(experimental=experimental),
|
||||
engine=EngineConfig(
|
||||
num_gpus=args.num_gpus,
|
||||
use_fsdp_inference=args.num_gpus > 1,
|
||||
|
||||
@@ -1,14 +1,18 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""End-to-end Wan2.2-TI2V-5B generation on Apple Silicon (MLX DiT + MLX TAEHV).
|
||||
"""End-to-end FastMetal-5B-QAD generation on Apple Silicon (MLX DiT + MLX TAEHV).
|
||||
|
||||
This is the Wan2.2 TI2V entrypoint. Use FastVideo/FastMetal-5B-QAD:
|
||||
|
||||
hf download FastVideo/FastMetal-5B-QAD --local-dir ./FastMetal-5B-QAD
|
||||
python examples/inference/basic/mlx_wan22_generate.py \\
|
||||
--mlx-checkpoint ./FastMetal-5B-QAD \\
|
||||
--text-encoder-root ./FastMetal-5B-QAD \\
|
||||
--vae-root ./FastMetal-5B-QAD/vae
|
||||
|
||||
Pipeline: torch/MPS UMT5 encode (shared with 1.3B) → MLXWan22DiT 3-step DMD
|
||||
(warped schedule, flow_shift=5) → MLX TAEHV decode (taew2_2.pth). Fully MLX
|
||||
on the heavy DiT + decode path.
|
||||
|
||||
PYTHONPATH=$PWD python examples/inference/basic/mlx_wan22_generate.py \
|
||||
--prompt "A red fox trotting through a snowy pine forest at golden hour" \
|
||||
--output-path video_samples/demo_5b/fox_5b_mlx.mp4
|
||||
|
||||
Decoder backends: ``taehv`` (default, MLX, ~seconds), ``taehv-torch`` (parity),
|
||||
``wan-vae`` (full AutoencoderKLWan on MPS, slow).
|
||||
"""
|
||||
@@ -31,6 +35,11 @@ from fastvideo.mlx_runtime.prompt_cache import (
|
||||
save_prompt_cache,
|
||||
text_encoder_fingerprint,
|
||||
)
|
||||
from fastvideo.mlx_runtime.checkpoint_compat import (
|
||||
UnsupportedMLXCheckpointError,
|
||||
raise_if_unsupported_mlx_checkpoint,
|
||||
resolve_mlx_checkpoint,
|
||||
)
|
||||
from fastvideo.mlx_runtime.rife_interp import aligned_keyframe_count
|
||||
|
||||
FASTWAN21_MODEL_ID = "FastVideo/FastWan2.1-T2V-1.3B-Diffusers"
|
||||
@@ -165,7 +174,9 @@ def main() -> None:
|
||||
"--mlx-checkpoint",
|
||||
type=Path,
|
||||
default=None,
|
||||
help="Pre-quantized MLX DiT checkpoint directory. Rewrapped with Wan2.2 per-token conditioning.",
|
||||
help="Packed FastMetal-5B-QAD MLX DiT directory (mlx_dit.json + mlx_dit.safetensors). "
|
||||
"If omitted, a FastMetal directory passed as --text-encoder-root is used when it "
|
||||
"already contains those files.",
|
||||
)
|
||||
parser.add_argument("--vae-root", type=Path, default=None)
|
||||
parser.add_argument("--height", type=int, default=DEFAULT_HEIGHT)
|
||||
@@ -241,6 +252,18 @@ def main() -> None:
|
||||
# latent never leaves the grid it was denoised on and the mode is usable.
|
||||
if args.refine and args.fast_spatial:
|
||||
print("[wan22] --refine takes precedence over --fast-spatial")
|
||||
|
||||
args.mlx_checkpoint = resolve_mlx_checkpoint(args.mlx_checkpoint, args.text_encoder_root)
|
||||
if args.mlx_checkpoint is not None:
|
||||
if args.text_encoder_root is None and (args.mlx_checkpoint / "text_encoder").is_dir():
|
||||
args.text_encoder_root = args.mlx_checkpoint
|
||||
if args.vae_root is None and (args.mlx_checkpoint / "vae").is_dir():
|
||||
args.vae_root = args.mlx_checkpoint / "vae"
|
||||
try:
|
||||
raise_if_unsupported_mlx_checkpoint(args.mlx_checkpoint, args.dit_checkpoint)
|
||||
except UnsupportedMLXCheckpointError as exc:
|
||||
raise SystemExit(str(exc)) from exc
|
||||
|
||||
args.text_encoder_root, args.dit_checkpoint, args.dit_config, args.vae_root = _resolve_model_paths(
|
||||
text_encoder_root=args.text_encoder_root,
|
||||
dit_checkpoint=args.dit_checkpoint,
|
||||
|
||||
@@ -1,11 +1,25 @@
|
||||
"""Generate a FastWan text-to-video clip with the Apple Silicon MLX runtime.
|
||||
"""Generate a FastMetal text-to-video clip with the Apple Silicon MLX runtime.
|
||||
|
||||
This is the supported source-tree entrypoint for the FastWan-QAD-INT8-1.3B
|
||||
Apple release:
|
||||
This is the supported source-tree entrypoint for FastMetal-QAD (Wan2.1 1.3B
|
||||
and 14B). Use ``mlx_wan22_generate.py`` for FastMetal-5B-QAD.
|
||||
|
||||
Download FastMetal-QAD and point ``--model-root`` / ``--mlx-checkpoint`` at it:
|
||||
|
||||
hf download FastVideo/FastMetal-1.3B-QAD --local-dir ./FastMetal-1.3B-QAD
|
||||
python examples/inference/basic/mlx_wan_prompt_to_video.py \\
|
||||
--model-root ./FastMetal-1.3B-QAD --mlx-checkpoint ./FastMetal-1.3B-QAD
|
||||
|
||||
CUDA FastWan-QAD (``FastVideo/FastWan-QAD-1.3B``, ``FastVideo/FastWan-QAD-FP8-1.3B``)
|
||||
is a separate NVIDIA release.
|
||||
|
||||
FastMetal-QAD Hugging Face repos ship ``mlx_dit.json`` + ``mlx_dit.safetensors``,
|
||||
not a Diffusers ``transformer/`` tree. Do not copy ``transformer/config.json``
|
||||
from Wan2.1 or other checkpoints; point ``--mlx-checkpoint`` at the FastMetal
|
||||
directory and the example reads the DiT config from ``mlx_dit.json``.
|
||||
|
||||
- Hugging Face/torch encodes the prompt with UMT5 (bf16 by default: fp32
|
||||
exponent range without fp16 overflow risk, at fp16 memory cost).
|
||||
- MLX runs the FastWan DiT denoising loop (INT8 by default, compiled with
|
||||
- MLX runs the FastMetal DiT denoising loop (INT8 by default, compiled with
|
||||
``mx.compile`` unless ``--no-mlx-compile``).
|
||||
- TAEHV (default, fast/low-memory) or the full Wan VAE (``--decode-backend
|
||||
wan-vae``, higher fidelity, bf16) decodes the final latents.
|
||||
@@ -56,13 +70,18 @@ from fastvideo.mlx_runtime.prompt_cache import (
|
||||
save_prompt_cache,
|
||||
text_encoder_fingerprint,
|
||||
)
|
||||
from fastvideo.mlx_runtime.checkpoint_compat import (
|
||||
UnsupportedMLXCheckpointError,
|
||||
raise_if_unsupported_mlx_checkpoint,
|
||||
resolve_mlx_checkpoint,
|
||||
)
|
||||
from fastvideo.mlx_runtime.rife_interp import aligned_keyframe_count
|
||||
|
||||
if TYPE_CHECKING: # pragma: no cover - typing only
|
||||
from fastvideo.mlx_runtime.fast_spatial import FastSpatialPlan
|
||||
|
||||
|
||||
DEFAULT_MODEL_ID = "FastVideo/FastWan2.1-T2V-1.3B-Diffusers"
|
||||
DEFAULT_MODEL_ID = "FastVideo/FastMetal-1.3B-QAD"
|
||||
|
||||
# Legacy pinned-snapshot location, kept for callers that import it (the MLX
|
||||
# benchmark harness). New code should prefer resolve_model_root(None), which
|
||||
@@ -99,6 +118,10 @@ def resolve_model_root(
|
||||
"tokenizer/*",
|
||||
"text_encoder/*",
|
||||
"vae/*",
|
||||
"mlx_dit.json",
|
||||
"mlx_dit.safetensors",
|
||||
"ema/mlx_dit.json",
|
||||
"ema/mlx_dit.safetensors",
|
||||
"transformer/*" if include_transformer else "transformer/config.json",
|
||||
]
|
||||
return Path(snapshot_download(
|
||||
@@ -138,6 +161,7 @@ def encode_prompt(
|
||||
text_encoder = UMT5EncoderModel.from_pretrained(
|
||||
model_root / "text_encoder",
|
||||
torch_dtype=dtype,
|
||||
low_cpu_mem_usage=True,
|
||||
local_files_only=True,
|
||||
).to(device)
|
||||
text_encoder.eval()
|
||||
@@ -379,6 +403,7 @@ def decode_latents_to_video(
|
||||
vae = AutoencoderKLWan.from_pretrained(
|
||||
model_root / "vae",
|
||||
torch_dtype=dtype,
|
||||
low_cpu_mem_usage=True,
|
||||
local_files_only=True,
|
||||
).to(device)
|
||||
vae.eval()
|
||||
@@ -471,11 +496,16 @@ def _rife_interpolate_video(*, video_path: Path, target_frames: int, factor: int
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description="Prompt-to-video FastWan generation using MLX for the DiT")
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Prompt-to-video FastMetal-QAD generation using the Apple Silicon MLX runtime")
|
||||
parser.add_argument("--model-root", type=Path, default=None,
|
||||
help=f"Model directory. Defaults to the local HF cache for {DEFAULT_MODEL_ID} "
|
||||
help="FastMetal-QAD directory (tokenizer, UMT5, VAE, packed MLX DiT). "
|
||||
f"Defaults to the local HF cache for {DEFAULT_MODEL_ID} "
|
||||
"(downloading it if missing).")
|
||||
parser.add_argument("--prompt", default="A paper boat sails through a shallow stream in a mossy forest.")
|
||||
parser.add_argument(
|
||||
"--prompt",
|
||||
default="A bird's-eye view of a misty forest valley at dawn.",
|
||||
)
|
||||
parser.add_argument("--output-path", type=Path, default=Path("video_samples/mlx_fastwan_prompt_to_video.mp4"))
|
||||
parser.add_argument("--height", type=int, default=480)
|
||||
parser.add_argument("--width", type=int, default=832)
|
||||
@@ -639,9 +669,8 @@ def main() -> None:
|
||||
"keyed by (model, prompt, length, dtype), so repeat runs skip "
|
||||
"the text encoder entirely. Default: on.")
|
||||
parser.add_argument("--mlx-checkpoint", type=Path, default=None,
|
||||
help="Load the DiT from a pre-quantized MLX checkpoint directory "
|
||||
"(created with --save-mlx-checkpoint) instead of casting/quantizing "
|
||||
"the Diffusers weights on every run.")
|
||||
help="Packed FastMetal MLX DiT directory (mlx_dit.json + mlx_dit.safetensors). "
|
||||
"Defaults to --model-root when that directory already contains those files.")
|
||||
parser.add_argument("--save-mlx-checkpoint", type=Path, default=None,
|
||||
help="After loading the DiT, save it (cast + quantized) as an MLX "
|
||||
"checkpoint directory for fast reloads via --mlx-checkpoint.")
|
||||
@@ -690,6 +719,12 @@ def main() -> None:
|
||||
np.save(args.encode_prompt_only, prompt_embeds.cpu().numpy())
|
||||
return
|
||||
|
||||
mlx_checkpoint = resolve_mlx_checkpoint(args.mlx_checkpoint, model_root)
|
||||
try:
|
||||
raise_if_unsupported_mlx_checkpoint(mlx_checkpoint or model_root)
|
||||
except UnsupportedMLXCheckpointError as exc:
|
||||
raise SystemExit(str(exc)) from exc
|
||||
|
||||
import mlx.core as mx
|
||||
import torch
|
||||
from diffusers import UniPCMultistepScheduler
|
||||
@@ -719,17 +754,21 @@ def main() -> None:
|
||||
|
||||
config_path = model_root / "transformer/config.json"
|
||||
checkpoint_path = model_root / "transformer/diffusion_pytorch_model.safetensors"
|
||||
config = json.loads(config_path.read_text())
|
||||
# A pre-quantized MLX DiT can be paired with a lightweight asset root for
|
||||
# UMT5/TAEHV. In that case the model root is *not* the architecture
|
||||
# authority: use the checkpoint's embedded transformer config for the
|
||||
# sampler guard and latent geometry.
|
||||
dit_config = config
|
||||
if args.mlx_checkpoint is not None:
|
||||
mlx_config_path = Path(args.mlx_checkpoint) / "mlx_dit.json"
|
||||
if mlx_config_path.is_file():
|
||||
mlx_checkpoint_config = json.loads(mlx_config_path.read_text())
|
||||
dit_config = mlx_checkpoint_config.get("config", mlx_checkpoint_config)
|
||||
# Packed FastMetal checkpoints are the architecture authority. Do not
|
||||
# require transformer/config.json when mlx_dit.json is already present.
|
||||
if mlx_checkpoint is not None:
|
||||
mlx_checkpoint_config = json.loads((mlx_checkpoint / "mlx_dit.json").read_text())
|
||||
dit_config = mlx_checkpoint_config.get("config", mlx_checkpoint_config)
|
||||
config = dit_config
|
||||
else:
|
||||
if not config_path.is_file():
|
||||
raise SystemExit(
|
||||
f"No packed MLX DiT (mlx_dit.json) and no Diffusers transformer config at {config_path}. "
|
||||
"FastMetal-QAD checkpoints intentionally omit transformer/; download "
|
||||
"FastVideo/FastMetal-1.3B-QAD and pass --model-root / --mlx-checkpoint at that directory."
|
||||
)
|
||||
config = json.loads(config_path.read_text())
|
||||
dit_config = config
|
||||
if int(dit_config.get("in_channels", 0)) == 48 and int(dit_config.get("out_channels", 0)) == 48:
|
||||
raise SystemExit(
|
||||
"Wan2.2-TI2V-5B uses 48-channel, per-token timestep conditioning. "
|
||||
@@ -835,10 +874,10 @@ def main() -> None:
|
||||
load_start = time.perf_counter()
|
||||
mx.clear_cache()
|
||||
mx.reset_peak_memory()
|
||||
if args.mlx_checkpoint is not None:
|
||||
if mlx_checkpoint is not None:
|
||||
from fastvideo.mlx_runtime.checkpoint import load_mlx_dit_checkpoint
|
||||
|
||||
dit = load_mlx_dit_checkpoint(args.mlx_checkpoint, compile=args.mlx_compile)
|
||||
dit = load_mlx_dit_checkpoint(mlx_checkpoint, compile=args.mlx_compile)
|
||||
config = dit.config
|
||||
else:
|
||||
dit = mlx_dit_from_diffusers_safetensors(
|
||||
|
||||
@@ -0,0 +1,13 @@
|
||||
#!/usr/bin/env bash
|
||||
set -euo pipefail
|
||||
|
||||
repo="FastVideo/FastVideo-FastH3-4-step-Preview-v1-LoRA"
|
||||
adapter="dense-datafree/adapter_model.safetensors"
|
||||
adapter_path="$(hf download "$repo" "$adapter")"
|
||||
|
||||
python examples/inference/basic/basic_fasth3_lora_preview.py \
|
||||
--lora-path "$adapter_path" \
|
||||
--lora-strength "${FASTH3_LORA_STRENGTH:-1.0}" \
|
||||
--output "${FASTH3_LORA_OUTPUT:-outputs/fasth3_lora_preview/dense-datafree}" \
|
||||
"$@" \
|
||||
--no-vsa
|
||||
@@ -0,0 +1,13 @@
|
||||
#!/usr/bin/env bash
|
||||
set -euo pipefail
|
||||
|
||||
repo="FastVideo/FastVideo-FastH3-4-step-Preview-v1-LoRA"
|
||||
adapter="vsa-datafree/adapter_model.safetensors"
|
||||
adapter_path="$(hf download "$repo" "$adapter")"
|
||||
|
||||
python examples/inference/basic/basic_fasth3_lora_preview.py \
|
||||
--lora-path "$adapter_path" \
|
||||
--lora-strength "${FASTH3_LORA_STRENGTH:-1.0}" \
|
||||
--output "${FASTH3_LORA_OUTPUT:-outputs/fasth3_lora_preview/vsa-datafree}" \
|
||||
"$@" \
|
||||
--vsa
|
||||
@@ -0,0 +1,13 @@
|
||||
#!/usr/bin/env bash
|
||||
set -euo pipefail
|
||||
|
||||
repo="FastVideo/FastVideo-FastH3-4-step-Preview-v1-LoRA"
|
||||
adapter="vsa-synthetic-step1300/adapter_model.safetensors"
|
||||
adapter_path="$(hf download "$repo" "$adapter")"
|
||||
|
||||
python examples/inference/basic/basic_fasth3_lora_preview.py \
|
||||
--lora-path "$adapter_path" \
|
||||
--lora-strength "${FASTH3_LORA_STRENGTH:-1.0}" \
|
||||
--output "${FASTH3_LORA_OUTPUT:-outputs/fasth3_lora_preview/vsa-synthetic-step1300}" \
|
||||
"$@" \
|
||||
--vsa
|
||||
@@ -0,0 +1,13 @@
|
||||
#!/usr/bin/env bash
|
||||
set -euo pipefail
|
||||
|
||||
repo="FastVideo/FastVideo-FastH3-4-step-Preview-v1-LoRA"
|
||||
adapter="vsa-synthetic-step1900/adapter_model.safetensors"
|
||||
adapter_path="$(hf download "$repo" "$adapter")"
|
||||
|
||||
python examples/inference/basic/basic_fasth3_lora_preview.py \
|
||||
--lora-path "$adapter_path" \
|
||||
--lora-strength "${FASTH3_LORA_STRENGTH:-1.0}" \
|
||||
--output "${FASTH3_LORA_OUTPUT:-outputs/fasth3_lora_preview/vsa-synthetic-step1900}" \
|
||||
"$@" \
|
||||
--vsa
|
||||
@@ -0,0 +1,184 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""Review whether a FastH3 adapter reproduces the checkpoint it was extracted from.
|
||||
|
||||
The question this page answers is not "is the video good" but "does base + adapter land
|
||||
where the real checkpoint lands". So each row is one checkpoint, and the two players in
|
||||
it are the checkpoint itself and base MiniMax-H3 with that checkpoint's adapter merged
|
||||
in. They share a seed and a prompt, so anything you can see between them is the
|
||||
adapter's approximation error and nothing else.
|
||||
|
||||
The base model at four steps sits at the top as the floor. It is not distilled, so it
|
||||
should look clearly worse than everything below it -- if an adapter row looks like the
|
||||
floor instead of like its checkpoint, the adapter did not land.
|
||||
|
||||
python examples/inference/gradio/lora_review.py --runs /path/to/lora_review
|
||||
|
||||
Expects one directory per arm, each holding ``<index>_<case_id>.mp4``:
|
||||
|
||||
<runs>/base/ <runs>/v1-true/ <runs>/v1-lora-r64/ ...
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import subprocess
|
||||
from functools import lru_cache
|
||||
from pathlib import Path
|
||||
|
||||
import gradio as gr
|
||||
|
||||
FLOOR_ARM = "base"
|
||||
|
||||
|
||||
def discover_pairs(arms: list[str]) -> tuple[list[tuple[str, str, str]], list[str]]:
|
||||
"""Split the arms present into checkpoint/adapter pairs and everything else.
|
||||
|
||||
Pairs are found by name -- ``<x>-true`` next to ``<x>-lora-<rank>`` -- rather than
|
||||
listed, so adding an arm to the render directory is enough to get it on the page.
|
||||
Arms that pair with nothing (a third-party adapter with no checkpoint to compare
|
||||
against) still get shown, on their own row, instead of being silently dropped.
|
||||
"""
|
||||
pairs, used = [], set()
|
||||
for arm in sorted(arms):
|
||||
if not arm.endswith("-true"):
|
||||
continue
|
||||
stem = arm[:-len("-true")]
|
||||
partner = next((a for a in arms if a.startswith(f"{stem}-lora")), None)
|
||||
if partner is None:
|
||||
continue
|
||||
pairs.append((stem, arm, partner))
|
||||
used.update({arm, partner})
|
||||
standalone = [a for a in sorted(arms) if a not in used and a != FLOOR_ARM]
|
||||
return pairs, standalone
|
||||
|
||||
|
||||
def probe(path: Path) -> str:
|
||||
"""`WxH · Nf · Ds · MiB`, so a truncated or mis-sized render is visible as text."""
|
||||
try:
|
||||
out = subprocess.run(
|
||||
["ffprobe", "-v", "error", "-select_streams", "v:0", "-show_entries",
|
||||
"stream=width,height,nb_read_packets,duration", "-count_packets", "-of", "json",
|
||||
str(path)],
|
||||
capture_output=True, text=True, check=True).stdout
|
||||
stream = json.loads(out)["streams"][0]
|
||||
frames = stream.get("nb_read_packets", "?")
|
||||
duration = float(stream.get("duration", 0) or 0)
|
||||
return (f"{stream['width']}x{stream['height']} · {frames}f · {duration:.1f}s · "
|
||||
f"{path.stat().st_size / 2**20:.1f} MiB")
|
||||
except (subprocess.CalledProcessError, KeyError, IndexError, json.JSONDecodeError):
|
||||
return f"{path.stat().st_size / 2**20:.1f} MiB"
|
||||
|
||||
|
||||
class Runs:
|
||||
"""Which prompts rendered, and where each arm's clip for them lives."""
|
||||
|
||||
def __init__(self, runs_dir: Path, prompts_file: Path | None) -> None:
|
||||
self.root = runs_dir
|
||||
self.arms = sorted(d.name for d in runs_dir.iterdir() if d.is_dir())
|
||||
self.prompts: dict[str, str] = {}
|
||||
if prompts_file and prompts_file.exists():
|
||||
with prompts_file.open() as handle:
|
||||
for index, line in enumerate(handle):
|
||||
line = line.strip()
|
||||
if line:
|
||||
self.prompts[f"{index:03d}"] = json.loads(line).get("prompt", "")
|
||||
|
||||
self.clips: dict[str, dict[str, Path]] = {}
|
||||
for arm in self.arms:
|
||||
for mp4 in sorted((runs_dir / arm).glob("*.mp4")):
|
||||
self.clips.setdefault(mp4.stem.split("_")[0], {})[arm] = mp4
|
||||
if not self.clips:
|
||||
raise SystemExit(f"no clips under {runs_dir}")
|
||||
|
||||
def label(self, index: str) -> str:
|
||||
head = " ".join(self.prompts.get(index, "").split())[:90]
|
||||
return f"[{index}] {head}..." if head else f"[{index}]"
|
||||
|
||||
def by_label(self, label: str) -> str:
|
||||
return next(i for i in self.clips if self.label(i) == label)
|
||||
|
||||
|
||||
def build(runs: Runs, height: int) -> gr.Blocks:
|
||||
|
||||
@lru_cache(maxsize=512)
|
||||
def cached_probe(path: str) -> str:
|
||||
return probe(Path(path))
|
||||
|
||||
def player_update(index: str, arm: str, prefix: str):
|
||||
mp4 = runs.clips.get(index, {}).get(arm)
|
||||
if mp4 is None:
|
||||
return gr.update(value=None, label=f"{prefix} — not rendered")
|
||||
return gr.update(value=str(mp4), label=f"{prefix} — {cached_probe(str(mp4))}")
|
||||
|
||||
with gr.Blocks(title="FastH3 adapter review") as demo:
|
||||
gr.Markdown(
|
||||
"# FastH3 adapter review\n"
|
||||
"Each row is one checkpoint: **left** is the real checkpoint, **right** is base "
|
||||
"MiniMax-H3 with that checkpoint's rank-64 adapter merged in. Same prompt, same "
|
||||
"seed, same sampler. Differences between the two are the adapter's approximation "
|
||||
"error.\n\n"
|
||||
"The top player is undistilled base MiniMax-H3 at four steps — the floor. Every "
|
||||
"row below it should look clearly better than that; an adapter that landed looks "
|
||||
"like its own left-hand player, not like the floor.")
|
||||
|
||||
prompt_dd = gr.Dropdown(choices=[runs.label(i) for i in runs.clips],
|
||||
value=runs.label(next(iter(runs.clips))),
|
||||
label="Prompt")
|
||||
prompt_box = gr.Textbox(label="Prompt", lines=4, max_lines=6, interactive=False, show_copy_button=True)
|
||||
|
||||
with gr.Row():
|
||||
floor = gr.Video(label="base MiniMax-H3, 4 steps (floor)", height=height, loop=True,
|
||||
autoplay=False, interactive=False)
|
||||
|
||||
pairs, standalone = discover_pairs(runs.arms)
|
||||
players: list[tuple[gr.Video, str, str]] = []
|
||||
for row_label, true_arm, lora_arm in pairs:
|
||||
gr.Markdown(f"### {row_label}")
|
||||
with gr.Row():
|
||||
left = gr.Video(label=f"{row_label} — checkpoint", height=height, loop=True,
|
||||
autoplay=False, interactive=False)
|
||||
right = gr.Video(label=f"{row_label} — base + adapter", height=height, loop=True,
|
||||
autoplay=False, interactive=False)
|
||||
players.append((left, true_arm, "checkpoint"))
|
||||
players.append((right, lora_arm, "base + adapter"))
|
||||
|
||||
if standalone:
|
||||
gr.Markdown("### Other adapters (no matching checkpoint to compare against)")
|
||||
with gr.Row():
|
||||
for arm in standalone:
|
||||
players.append((gr.Video(label=arm, height=height, loop=True, autoplay=False,
|
||||
interactive=False), arm, arm))
|
||||
|
||||
def show(label: str):
|
||||
index = runs.by_label(label)
|
||||
return [
|
||||
player_update(index, FLOOR_ARM, "base, 4 steps"),
|
||||
gr.update(value=runs.prompts.get(index, "")),
|
||||
*[player_update(index, arm, prefix) for _, arm, prefix in players],
|
||||
]
|
||||
|
||||
outputs = [floor, prompt_box, *[p for p, _, _ in players]]
|
||||
gr.on(triggers=[prompt_dd.change, demo.load], fn=show, inputs=prompt_dd, outputs=outputs)
|
||||
|
||||
return demo
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
|
||||
parser.add_argument("--runs", default="/mnt/lustre/vlm-s4duan/arena_arms/lora_review")
|
||||
parser.add_argument("--prompts-file", default="/mnt/lustre/vlm-s4duan/FastVideo/prompts.jsonl")
|
||||
parser.add_argument("--host", default="0.0.0.0")
|
||||
parser.add_argument("--port", type=int, default=7865)
|
||||
parser.add_argument("--share", action="store_true")
|
||||
parser.add_argument("--video-height", type=int, default=420)
|
||||
args = parser.parse_args()
|
||||
|
||||
runs = Runs(Path(args.runs).resolve(), Path(args.prompts_file))
|
||||
print(f"arms: {runs.arms}")
|
||||
print(f"prompts with output: {sorted(runs.clips)}")
|
||||
build(runs, args.video_height).launch(server_name=args.host, server_port=args.port,
|
||||
share=args.share, allowed_paths=[str(runs.root)])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,234 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""Run a FastH3 distillation adapter on top of the base MiniMax-H3 checkpoint.
|
||||
|
||||
The FastH3 checkpoints are full fine-tunes of MiniMax-H3 distilled to four steps under
|
||||
video sparse attention. Published as adapters they are three things at once, and all
|
||||
three have to land for the result to match the checkpoint:
|
||||
|
||||
* low-rank factors for the attention, feed-forward, and AdaLN projections
|
||||
* exact ``.diff`` deltas for the norms and biases an SVD cannot usefully factor
|
||||
* ``.set_weight`` values for ``attn.to_gate_compress``, the VSA compression gate that
|
||||
does not exist in the base model at all
|
||||
|
||||
An adapter carrying that last one needs ``--vsa``: under any other attention backend the
|
||||
gate module is never constructed, so part of the distillation has nowhere to go. The
|
||||
requirement is read off the adapter rather than assumed, because community adapters
|
||||
built against the ComfyUI layout carry no gate and run fine either way -- run one of
|
||||
those with ``--no-vsa`` (see ``scripts/checkpoint_conversion/convert_minimax_h3_comfy_lora.py``
|
||||
for getting them into a layout this loads).
|
||||
|
||||
Because a parameter the base lacks has to be supplied while weights are still unsharded,
|
||||
the adapter is passed at construction rather than swapped in afterwards.
|
||||
|
||||
python examples/inference/lora/minimax_h3_lora_inference.py \\
|
||||
--lora-path /models/fasth3-loras-publish/FastH3-4-step-v1.1/rank-64 \\
|
||||
--prompts-file prompts.jsonl --output outputs/v1.1-rank64
|
||||
|
||||
Pass no ``--lora-path`` to render the unmodified base model as a control.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import math
|
||||
import os
|
||||
from collections.abc import Sequence
|
||||
from pathlib import Path
|
||||
|
||||
# MiniMax-H3 generates 5-15 s at 24 fps, on a latent grid that only admits frame counts
|
||||
# of the form 17n + 5. 124 is the 5-second point the FastH3 profile is measured at.
|
||||
FRAMES_PER_CHUNK = 17
|
||||
LATENTS_PER_CHUNK = 5
|
||||
FPS = 24
|
||||
MIN_DURATION, MAX_DURATION = 5.0, 15.0
|
||||
|
||||
|
||||
def align_num_frames(num_frames: int) -> int:
|
||||
"""Round up to the next 17n + 5 the latent grid accepts."""
|
||||
if num_frames <= LATENTS_PER_CHUNK:
|
||||
return LATENTS_PER_CHUNK
|
||||
chunks = -(-(num_frames - LATENTS_PER_CHUNK) // FRAMES_PER_CHUNK)
|
||||
return LATENTS_PER_CHUNK + chunks * FRAMES_PER_CHUNK
|
||||
|
||||
|
||||
def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
|
||||
parser.add_argument("--model-path", default="MiniMaxAI/MiniMax-H3", help="the BASE checkpoint the adapter targets")
|
||||
parser.add_argument("--lora-path", default=None, help="adapter file or directory; omit to render the base model")
|
||||
parser.add_argument("--lora-nickname", default="fasth3")
|
||||
parser.add_argument("--lora-strength", type=float, default=1.0)
|
||||
parser.add_argument("--prompt", default=None)
|
||||
parser.add_argument("--prompts-file", default=None, help="JSONL with a 'prompt' field per line")
|
||||
parser.add_argument("--limit", type=int, default=None, help="use only the first N prompts")
|
||||
parser.add_argument("--num-shards", type=int, default=1, help="split the prompt list across processes")
|
||||
parser.add_argument("--shard", type=int, default=0)
|
||||
parser.add_argument("--output", default="outputs/minimax_h3_lora")
|
||||
parser.add_argument("--skip-existing", action="store_true", help="leave already-rendered clips alone")
|
||||
parser.add_argument("--height", type=int, default=768)
|
||||
parser.add_argument("--width", type=int, default=1344)
|
||||
parser.add_argument("--num-frames", type=int, default=124)
|
||||
# Counts sigma-GRID POINTS: N points run N-1 DiT forwards. The distilled ladder is
|
||||
# t=1000,750,500,250 -> 0, which is five points and exactly four forwards.
|
||||
parser.add_argument("--steps", type=int, default=5)
|
||||
parser.add_argument("--seed", type=int, default=1000)
|
||||
parser.add_argument("--num-gpus", type=int, default=4)
|
||||
parser.add_argument("--vsa", action=argparse.BooleanOptionalAction, default=None,
|
||||
help="video sparse attention; inferred from the adapter when omitted")
|
||||
parser.add_argument("--vsa-sparsity", type=float, default=0.9)
|
||||
parser.add_argument("--vsa-tile-size", type=int, choices=(64, 256), default=64)
|
||||
parser.add_argument("--vsa-kernel", choices=("triton", "sm100a"), default="sm100a")
|
||||
parser.add_argument("--fa4", action=argparse.BooleanOptionalAction, default=True)
|
||||
args = parser.parse_args(argv)
|
||||
if not args.prompt and not args.prompts_file:
|
||||
parser.error("pass --prompt or --prompts-file")
|
||||
# Whether VSA is required is a property of the adapter, not of having one at all --
|
||||
# community adapters built against the ComfyUI layout carry no gate. Checked in
|
||||
# main(), once the path has been resolved.
|
||||
aligned = align_num_frames(args.num_frames)
|
||||
if not MIN_DURATION <= aligned / FPS <= MAX_DURATION:
|
||||
parser.error(f"MiniMax-H3 generates {MIN_DURATION}-{MAX_DURATION}s at {FPS} fps; "
|
||||
f"aligned num_frames={aligned} is {aligned / FPS:.1f}s")
|
||||
args.num_frames = aligned
|
||||
if not math.isfinite(args.lora_strength):
|
||||
parser.error("--lora-strength must be finite")
|
||||
return args
|
||||
|
||||
|
||||
def configure_environment(args: argparse.Namespace) -> None:
|
||||
"""Set the boot-time backend selection explicitly, including what is off.
|
||||
|
||||
An inherited FASTVIDEO_* from an earlier experiment would otherwise silently change
|
||||
which attention path the run actually took, which is the one thing this comparison
|
||||
cannot afford to be vague about.
|
||||
"""
|
||||
env: dict[str, str | None] = {
|
||||
"FASTVIDEO_ATTENTION_BACKEND": "VIDEO_SPARSE_ATTN_H3" if args.vsa else "FLASH_ATTN",
|
||||
"FASTVIDEO_VSA_SM100A": "1" if (args.vsa and args.vsa_kernel == "sm100a") else "0",
|
||||
"FASTVIDEO_VSA_CUTEDSL": "0",
|
||||
"FASTVIDEO_H3_VSA_PROBE": None,
|
||||
"FASTVIDEO_DISABLE_ATTENTION_COMPILE": "0",
|
||||
"FASTVIDEO_FA4": "1" if args.fa4 else "0",
|
||||
"FASTVIDEO_NVFP4_FA4": "0",
|
||||
"FASTVIDEO_MINIMAX_H3_FA4_PACKED_VARLEN": "0",
|
||||
"FASTVIDEO_MINIMAX_H3_FUSIONS": "all",
|
||||
"FASTVIDEO_INFERENCE_TORCH_COMPILE": "1",
|
||||
"FASTVIDEO_VAE_PARALLEL_DECODE": "1",
|
||||
"FASTVIDEO_VAE_PARALLEL_ENCODE": "0",
|
||||
"FASTVIDEO_VAE_PARALLEL_DECODE_STRATEGY": "gather",
|
||||
"FASTVIDEO_ULYSSES_A2A": "off",
|
||||
"FASTVIDEO_STAGE_LOGGING": "1",
|
||||
}
|
||||
for name, value in env.items():
|
||||
if value is None:
|
||||
os.environ.pop(name, None)
|
||||
else:
|
||||
os.environ[name] = value
|
||||
|
||||
|
||||
def load_prompts(args: argparse.Namespace) -> list[dict]:
|
||||
if args.prompt:
|
||||
records = [{"id": "000", "prompt": args.prompt}]
|
||||
else:
|
||||
records = []
|
||||
with open(args.prompts_file) as handle:
|
||||
for index, line in enumerate(handle):
|
||||
line = line.strip()
|
||||
if not line:
|
||||
continue
|
||||
item = json.loads(line)
|
||||
records.append({
|
||||
"id": str(item.get("id", item.get("sample_id", f"{index:03d}"))),
|
||||
"prompt": item["prompt"],
|
||||
})
|
||||
if args.limit is not None:
|
||||
records = records[:args.limit]
|
||||
return [r for i, r in enumerate(records) if i % args.num_shards == args.shard]
|
||||
|
||||
|
||||
def main(argv: Sequence[str] | None = None) -> None:
|
||||
args = parse_args(argv)
|
||||
# Backend selection is finalized from the adapter before model construction.
|
||||
from fastvideo.models.loader.lora_patch import DenseLoRAPatch
|
||||
from fastvideo import VideoGenerator
|
||||
from fastvideo.api import (CompileConfig, ComponentConfig, EngineConfig, GenerationRequest, GeneratorConfig,
|
||||
OffloadConfig, OutputConfig, ParallelismConfig, PipelineSelection, SamplingConfig)
|
||||
|
||||
# An adapter carrying to_gate_compress needs the VSA backend, because that is the
|
||||
# only configuration in which the module exists. One that does not carry it runs
|
||||
# fine either way, so the requirement is read off the file rather than assumed from
|
||||
# the presence of an adapter at all.
|
||||
patch = (DenseLoRAPatch.from_adapter(args.lora_path, strength=args.lora_strength)
|
||||
if args.lora_path else None)
|
||||
needs_vsa = bool(patch and any("gate_compress" in name for name in patch.replacement_parameters))
|
||||
if args.vsa is None:
|
||||
args.vsa = needs_vsa
|
||||
if needs_vsa and not args.vsa:
|
||||
raise SystemExit(f"{args.lora_path} carries to_gate_compress, which exists only under the VSA "
|
||||
"attention backend. Drop --no-vsa.")
|
||||
if args.vsa and args.lora_path and not needs_vsa:
|
||||
print(f"note: {args.lora_path} carries no VSA gate; running under VSA leaves the "
|
||||
"compression branch at its zero-initialized value.")
|
||||
configure_environment(args)
|
||||
|
||||
experimental: dict[str, object] = {
|
||||
"inference_torch_compile": True,
|
||||
"vae_parallel_decode": True,
|
||||
"vae_parallel_decode_strategy": "gather",
|
||||
}
|
||||
if args.vsa:
|
||||
experimental.update({
|
||||
"attention_backend": "VIDEO_SPARSE_ATTN_H3",
|
||||
"VSA_sparsity": args.vsa_sparsity,
|
||||
"VSA_tile_size": args.vsa_tile_size,
|
||||
})
|
||||
|
||||
config = GeneratorConfig(
|
||||
model_path=args.model_path,
|
||||
pipeline=PipelineSelection(
|
||||
components=ComponentConfig(lora_path=args.lora_path, lora_strength=args.lora_strength),
|
||||
experimental=experimental,
|
||||
),
|
||||
engine=EngineConfig(
|
||||
num_gpus=args.num_gpus,
|
||||
use_fsdp_inference=False,
|
||||
parallelism=ParallelismConfig(tp_size=1, sp_size=args.num_gpus),
|
||||
offload=OffloadConfig(dit=False, dit_layerwise=False, text_encoder=True, vae=True, pin_cpu_memory=True),
|
||||
compile=CompileConfig(enabled=False, vae_enabled=True),
|
||||
),
|
||||
)
|
||||
|
||||
records = load_prompts(args)
|
||||
out_dir = Path(args.output)
|
||||
out_dir.mkdir(parents=True, exist_ok=True)
|
||||
print(f"adapter: {args.lora_path or '(none, base model)'}")
|
||||
print(f"prompts: {len(records)} (shard {args.shard}/{args.num_shards})")
|
||||
|
||||
generator = VideoGenerator.from_config(config)
|
||||
for index, record in enumerate(records):
|
||||
stem = f"{index:03d}_{record['id']}"
|
||||
if args.skip_existing and (out_dir / f"{stem}.mp4").exists():
|
||||
print(f"[{index}] skip {stem}")
|
||||
continue
|
||||
generator.generate(
|
||||
GenerationRequest(
|
||||
prompt=record["prompt"],
|
||||
negative_prompt="",
|
||||
sampling=SamplingConfig(
|
||||
height=args.height,
|
||||
width=args.width,
|
||||
num_frames=args.num_frames,
|
||||
fps=FPS,
|
||||
num_inference_steps=args.steps,
|
||||
# MiniMax-H3 is guidance-distilled; FastH3 inherits that contract.
|
||||
guidance_scale=1.0,
|
||||
batch_cfg=False,
|
||||
seed=args.seed,
|
||||
),
|
||||
output=OutputConfig(output_path=str(out_dir / f"{stem}.mp4"), save_video=True, return_frames=False),
|
||||
))
|
||||
print(f"[{index}] wrote {stem}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -114,7 +114,63 @@ list(APPEND TORCH_INCLUDE_DIRS ${TORCH_INCLUDE_PATHS})
|
||||
# Find Torch package (still useful for libraries)
|
||||
find_package(Torch REQUIRED)
|
||||
|
||||
# Include directories
|
||||
# The Ulysses kernel needs NCCL's 2.29 device API. Keep it optional so ROCm,
|
||||
# older NCCL installs, and minimal CUDA builders still produce a usable wheel.
|
||||
set(FASTVIDEO_KERNEL_BUILD_ULYSSES_A2A "AUTO" CACHE STRING
|
||||
"Build the NCCL-device-API Ulysses all-to-all: AUTO/ON/OFF")
|
||||
set_property(CACHE FASTVIDEO_KERNEL_BUILD_ULYSSES_A2A PROPERTY STRINGS AUTO ON OFF)
|
||||
set(ENABLE_ULYSSES_A2A OFF)
|
||||
|
||||
if(NOT GPU_BACKEND STREQUAL "ROCM" AND NOT FASTVIDEO_KERNEL_BUILD_ULYSSES_A2A STREQUAL "OFF")
|
||||
# nvidia.nccl is a namespace package, so find_spec rather than __file__.
|
||||
execute_process(COMMAND "${Python_EXECUTABLE}" -c
|
||||
"import importlib.util as u;s=u.find_spec('nvidia.nccl');print(list(s.submodule_search_locations)[0] if s else '')"
|
||||
OUTPUT_VARIABLE NCCL_PIP_ROOT OUTPUT_STRIP_TRAILING_WHITESPACE ERROR_QUIET)
|
||||
find_path(NCCL_INCLUDE_DIR nccl_device/coop.h
|
||||
HINTS $ENV{NCCL_HOME}/include ${NCCL_PIP_ROOT}/include)
|
||||
find_library(NCCL_LIBRARY NAMES nccl
|
||||
HINTS $ENV{NCCL_HOME}/lib $ENV{NCCL_HOME}/lib64 ${NCCL_PIP_ROOT}/lib)
|
||||
# PyPI's nvidia-nccl-cu* wheels ship the SONAME but not the development
|
||||
# symlink (libnccl.so.2, no libnccl.so). Accept that exact versioned name;
|
||||
# target_link_libraries can link an absolute SONAME path directly.
|
||||
if(NOT NCCL_LIBRARY)
|
||||
find_file(NCCL_VERSIONED_LIBRARY NAMES libnccl.so.2
|
||||
HINTS $ENV{NCCL_HOME}/lib $ENV{NCCL_HOME}/lib64 ${NCCL_PIP_ROOT}/lib)
|
||||
if(NCCL_VERSIONED_LIBRARY)
|
||||
set(NCCL_LIBRARY "${NCCL_VERSIONED_LIBRARY}")
|
||||
endif()
|
||||
endif()
|
||||
|
||||
if(NCCL_INCLUDE_DIR AND NCCL_LIBRARY)
|
||||
include(CheckCXXSourceCompiles)
|
||||
set(_FASTVIDEO_REQUIRED_INCLUDES "${CMAKE_REQUIRED_INCLUDES}")
|
||||
# nccl.h includes cuda_runtime.h, so a host-compiler feature probe
|
||||
# needs the toolkit includes explicitly even though CUDA is enabled.
|
||||
set(CMAKE_REQUIRED_INCLUDES "${NCCL_INCLUDE_DIR};${CUDAToolkit_INCLUDE_DIRS}")
|
||||
unset(NCCL_HAS_REQUIRED_DEVICE_API CACHE)
|
||||
check_cxx_source_compiles("\
|
||||
#define NCCL_HOSTLIB_ONLY
|
||||
#include <cstddef>
|
||||
#include <nccl_device.h>
|
||||
int main() {
|
||||
ncclDevCommRequirements reqs = NCCL_DEV_COMM_REQUIREMENTS_INITIALIZER;
|
||||
ncclCommProperties props = NCCL_COMM_PROPERTIES_INITIALIZER;
|
||||
return (reqs.size == 0 || props.size == 0);
|
||||
}"
|
||||
NCCL_HAS_REQUIRED_DEVICE_API)
|
||||
set(CMAKE_REQUIRED_INCLUDES "${_FASTVIDEO_REQUIRED_INCLUDES}")
|
||||
endif()
|
||||
|
||||
if(NCCL_HAS_REQUIRED_DEVICE_API)
|
||||
set(ENABLE_ULYSSES_A2A ON)
|
||||
elseif(FASTVIDEO_KERNEL_BUILD_ULYSSES_A2A STREQUAL "ON")
|
||||
message(FATAL_ERROR
|
||||
"FASTVIDEO_KERNEL_BUILD_ULYSSES_A2A=ON requires NCCL device headers/library "
|
||||
"with NCCL_DEV_COMM_REQUIREMENTS_INITIALIZER (NCCL 2.29+). "
|
||||
"Resolved include='${NCCL_INCLUDE_DIR}', library='${NCCL_LIBRARY}'.")
|
||||
endif()
|
||||
endif()
|
||||
|
||||
include_directories(
|
||||
${CMAKE_SOURCE_DIR}/include
|
||||
${CMAKE_SOURCE_DIR}/include/cutlass/include
|
||||
@@ -124,6 +180,9 @@ include_directories(
|
||||
${CMAKE_SOURCE_DIR}/csrc/turbodiffusion
|
||||
${TORCH_INCLUDE_DIRS}
|
||||
)
|
||||
if(ENABLE_ULYSSES_A2A)
|
||||
include_directories(${NCCL_INCLUDE_DIR})
|
||||
endif()
|
||||
|
||||
# ---------------------------
|
||||
# ThunderKittens (TK) toggles
|
||||
@@ -157,6 +216,7 @@ endif()
|
||||
message(STATUS "TORCH_CUDA_ARCH_LIST (cmake/env): ${TORCH_CUDA_ARCH_LIST}")
|
||||
message(STATUS "FASTVIDEO_KERNEL_BUILD_TK: ${FASTVIDEO_KERNEL_BUILD_TK}")
|
||||
message(STATUS "FASTVIDEO_KERNEL_BUILD_ATTN_QAT_INFER: ${FASTVIDEO_KERNEL_BUILD_ATTN_QAT_INFER}")
|
||||
message(STATUS "FASTVIDEO_KERNEL_BUILD_ULYSSES_A2A: ${FASTVIDEO_KERNEL_BUILD_ULYSSES_A2A}")
|
||||
|
||||
set(ENABLE_TK_KERNELS OFF)
|
||||
if(FASTVIDEO_KERNEL_BUILD_TK STREQUAL "ON")
|
||||
@@ -307,6 +367,9 @@ if(BUILD_CXX_KERNELS)
|
||||
csrc/turbodiffusion/norm/layernorm.cu
|
||||
csrc/turbodiffusion/quant/quant.cu
|
||||
)
|
||||
if(ENABLE_ULYSSES_A2A)
|
||||
list(APPEND EXTENSION_SOURCES csrc/comm/ulysses_all_to_all.cu)
|
||||
endif()
|
||||
|
||||
# Conditionally add TK kernels
|
||||
if(ENABLE_TK_KERNELS)
|
||||
@@ -371,6 +434,9 @@ if(BUILD_CXX_KERNELS)
|
||||
if(ENABLE_TK_KERNELS)
|
||||
list(APPEND COMPILE_DEFS TK_COMPILE_ST_ATTN TK_COMPILE_BLOCK_SPARSE)
|
||||
endif()
|
||||
if(ENABLE_ULYSSES_A2A)
|
||||
list(APPEND COMPILE_DEFS FASTVIDEO_KERNEL_COMPILE_ULYSSES_A2A)
|
||||
endif()
|
||||
|
||||
|
||||
target_compile_definitions(fastvideo_kernel_ops PRIVATE ${COMPILE_DEFS})
|
||||
@@ -382,6 +448,9 @@ if(BUILD_CXX_KERNELS)
|
||||
# Link against Torch libraries to avoid undefined symbols at import time
|
||||
# (e.g., torch::autograd vtables) when loading the extension module.
|
||||
target_link_libraries(fastvideo_kernel_ops PRIVATE ${TORCH_LIBRARIES})
|
||||
if(ENABLE_ULYSSES_A2A)
|
||||
target_link_libraries(fastvideo_kernel_ops PRIVATE "${NCCL_LIBRARY}")
|
||||
endif()
|
||||
|
||||
# Also link against libtorch_python to satisfy Python-binding symbols
|
||||
# (e.g., torch::PyWarningHandler) required by torch/extension.h.
|
||||
@@ -485,6 +554,7 @@ message(STATUS "host / backend: ${CMAKE_SYSTEM_PROCESSOR} / ${GPU_BACKEND}
|
||||
message(STATUS "TORCH_CUDA_ARCH_LIST: ${TORCH_CUDA_ARCH_LIST}")
|
||||
message(STATUS "fastvideo_kernel_ops: ON (turbodiffusion int8-gemm/quant/rmsnorm/layernorm, all listed archs)")
|
||||
message(STATUS " + TK sta/block_sparse (sm_90a only): ${ENABLE_TK_KERNELS}")
|
||||
message(STATUS " + Ulysses NCCL-device all-to-all: ${ENABLE_ULYSSES_A2A}")
|
||||
message(STATUS "fp4attn/fp4quant (sm_120a only, CUDA >= 12.8): ${ENABLE_ATTN_QAT_INFER}")
|
||||
message(STATUS "Triton fallbacks ship in python/fastvideo_kernel/triton_kernels regardless.")
|
||||
message(STATUS "============================================================")
|
||||
|
||||
@@ -10,6 +10,8 @@ Compiled CUDA extensions (CMake, see the build summary printed at the end of eve
|
||||
|---|---|---|---|---|
|
||||
| `fastvideo_kernel._C.fastvideo_kernel_ops` | TurboDiffusion INT8 GEMM, quant, RMSNorm, LayerNorm | `csrc/turbodiffusion/` | every arch in `TORCH_CUDA_ARCH_LIST` | always built |
|
||||
| same extension, optional part | ThunderKittens sliding-tile attention (`sta_fwd`) and VSA block-sparse (`block_sparse_fwd/bwd`) | `csrc/attention/*_h100.cu` | Hopper `sm_90a` only | `FASTVIDEO_KERNEL_BUILD_TK` (AUTO = ON iff `9.0a` is in the arch list; always OFF on aarch64 hosts — TK headers don't compile there) |
|
||||
| same extension, optional part | MiniMax-H3 block-sparse VSA forward (64- and 128-token blocks) | `csrc/attention/block_sparse*_sm100a.cu` | Blackwell `sm_100a` only | ON iff `10.0a` is in `TORCH_CUDA_ARCH_LIST` |
|
||||
| same extension, optional part | fused NVLink Ulysses all-to-all | `csrc/comm/ulysses_all_to_all.cu` | CUDA | `FASTVIDEO_KERNEL_BUILD_ULYSSES_A2A` (AUTO = ON with NCCL 2.29+ device headers and library; always OFF on ROCm) |
|
||||
| `fp4attn_cuda`, `fp4quant_cuda` | FP4 attention + quantization ("attn_qat_infer", modified SageAttention3) | `attn_qat_infer/` | consumer Blackwell `sm_120a` only, CUDA ≥ 12.8 | `FASTVIDEO_KERNEL_BUILD_ATTN_QAT_INFER` (AUTO = ON iff `12.0a` is in the arch list) |
|
||||
|
||||
Runtime-JIT kernels (no build step, ship in every wheel/image):
|
||||
@@ -22,20 +24,23 @@ Runtime-JIT kernels (no build step, ship in every wheel/image):
|
||||
|
||||
## What gets built where, and when
|
||||
|
||||
| Surface | Trigger | Leg | `TORCH_CUDA_ARCH_LIST` | TK | FP4 |
|
||||
|---|---|---|---|---|---|
|
||||
| PyPI wheels (`.github/workflows/publish-kernel.yml`) | version bump in `fastvideo-kernel/pyproject.toml` on main, or manual dispatch | x86_64 cu126 | `9.0a` | ON | — (CUDA < 12.8) |
|
||||
| | | x86_64 cu130 | `9.0a;12.0a` | ON | ON |
|
||||
| | | aarch64 cu130 | `10.0a;12.0a` | — | ON |
|
||||
| Docker images `ghcr.io/hao-ai-lab/fastvideo/fastvideo-dev` (`.github/workflows/infra-build-image.yml`) | `docker/Dockerfile` changes on main, or manual dispatch | amd64 cuda12.6.3 + cuda13.0.0 | `9.0a` | ON | — |
|
||||
| | | arm64 cuda12.6.3 (GH200) | `9.0a` | — (aarch64) | — |
|
||||
| | | arm64 cuda13.0.0 (GB10 / DGX Spark) | `12.1` | — | — |
|
||||
| Local `./build.sh` | manual | probes the visible GPU via torch | detected | ON iff sm_90 (non-aarch64 host) | ON iff sm_120 |
|
||||
| Surface | Trigger | Leg | `TORCH_CUDA_ARCH_LIST` | TK | Ulysses | FP4 |
|
||||
|---|---|---|---|---|---|---|
|
||||
| PyPI wheels (`.github/workflows/publish-kernel.yml`) | version bump in `fastvideo-kernel/pyproject.toml` on main, or manual dispatch | x86_64 cu126 | `9.0a` | ON | AUTO | — (CUDA < 12.8) |
|
||||
| | | x86_64 cu130 | `9.0a;10.0a;12.0a` | ON | AUTO | ON |
|
||||
| | | aarch64 cu130 | `10.0a;12.0a` | — | AUTO | ON |
|
||||
| Docker images `ghcr.io/hao-ai-lab/fastvideo/fastvideo-dev` (`.github/workflows/infra-build-image.yml`) | `docker/Dockerfile` changes on main, or manual dispatch | amd64 cuda12.6.3 + cuda13.0.0 | `9.0a` | ON | AUTO | — |
|
||||
| | | arm64 cuda12.6.3 (GH200) | `9.0a` | — (aarch64) | AUTO | — |
|
||||
| | | arm64 cuda13.0.0 (GB10 / DGX Spark) | `12.1` | — | AUTO | — |
|
||||
| Local `./build.sh` | manual | probes the visible GPU via torch | detected | ON iff sm_90 (non-aarch64 host) | AUTO | ON iff sm_120 |
|
||||
|
||||
Notes:
|
||||
|
||||
- No Docker image ships the FP4 kernels; only the x86_64/aarch64 cu130 wheels do.
|
||||
- Both cu130 PyPI wheels ship the MiniMax-H3 sm_100a VSA forward.
|
||||
- On arm64 images (GH200 included) STA/VSA run on the Triton fallbacks, since TK never builds on aarch64.
|
||||
- Ulysses AUTO builds only when CMake finds a NCCL library and device-API headers with the 2.29 initializers. Use
|
||||
`-DFASTVIDEO_KERNEL_BUILD_ULYSSES_A2A=ON` to require it or `OFF` to test the portable build.
|
||||
- Kernel tests run on Buildkite GPU CI for PRs touching `fastvideo-kernel/**` (see `.buildkite/pipeline.yml`).
|
||||
|
||||
## Installation
|
||||
|
||||
@@ -0,0 +1,279 @@
|
||||
/*
|
||||
* Copyright (c) 2025 by FlashInfer team.
|
||||
*
|
||||
* Adapted from flashinfer-ai/flashinfer @ 8a94642d83cba0939035868fb6c309b4474a13d6
|
||||
* (PR #3820), csrc/ulysses_all_to_all.cu.
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
|
||||
// Torch host bindings for the fused Ulysses all-to-all. Kernel in
|
||||
// include/comm/ulysses_all_to_all.cuh.
|
||||
//
|
||||
// The per-group context is an ncclDevComm plus a registered symmetric window,
|
||||
// both created here from the caller's ncclComm_t.
|
||||
|
||||
#include <c10/cuda/CUDAGuard.h>
|
||||
#include <c10/cuda/CUDAStream.h>
|
||||
#include <torch/extension.h>
|
||||
|
||||
#include <algorithm>
|
||||
#include <cstdint>
|
||||
#include <memory>
|
||||
#include <vector>
|
||||
|
||||
#include <nccl.h>
|
||||
#include <nccl_device.h>
|
||||
|
||||
#include "comm/ulysses_all_to_all.cuh"
|
||||
|
||||
namespace fi = fastvideo::comm::ulysses;
|
||||
|
||||
namespace {
|
||||
|
||||
// A symmetric window every rank can store into, plus the ncclDevComm the
|
||||
// kernel opens barrier sessions on.
|
||||
struct UlyssesContext {
|
||||
ncclComm_t comm = nullptr;
|
||||
ncclWindow_t win = nullptr;
|
||||
void* buf = nullptr;
|
||||
size_t nbytes = 0;
|
||||
ncclDevComm devComm{};
|
||||
bool dev_comm_created = false;
|
||||
int device = -1;
|
||||
int rank = 0;
|
||||
int world = 0;
|
||||
};
|
||||
|
||||
#define NCCL_TRY(expr, what) \
|
||||
do { \
|
||||
ncclResult_t _r = (expr); \
|
||||
TORCH_CHECK(_r == ncclSuccess, what " failed: rc=", static_cast<int>(_r)); \
|
||||
} while (0)
|
||||
|
||||
} // namespace
|
||||
|
||||
// Allocate the local half of a context. This is intentionally separate from
|
||||
// registration so Python can vote after local allocation: if one rank is OOM,
|
||||
// no peer enters a collective window registration alone.
|
||||
int64_t allocate_ulysses_a2a(int64_t nbytes, int64_t rank, int64_t world_size,
|
||||
int64_t device_index) {
|
||||
TORCH_CHECK(world_size == 2 || world_size == 4 || world_size == 6 || world_size == 8,
|
||||
"ulysses a2a only supports world size in (2, 4, 6, 8), got ", world_size);
|
||||
TORCH_CHECK(rank >= 0 && rank < world_size, "invalid rank");
|
||||
TORCH_CHECK(nbytes > 0, "nbytes must be positive");
|
||||
TORCH_CHECK(device_index >= 0, "device index must be non-negative");
|
||||
|
||||
const at::cuda::CUDAGuard device_guard(
|
||||
c10::Device(c10::DeviceType::CUDA, static_cast<c10::DeviceIndex>(device_index)));
|
||||
auto ctx = std::make_unique<UlyssesContext>();
|
||||
ctx->nbytes = static_cast<size_t>(nbytes);
|
||||
ctx->device = static_cast<int>(device_index);
|
||||
ctx->rank = static_cast<int>(rank);
|
||||
ctx->world = static_cast<int>(world_size);
|
||||
|
||||
// The window must come from NCCL's allocator (4096B aligned per
|
||||
// NCCL_WIN_REQUIRED_ALIGNMENT), which is why this is not a torch tensor.
|
||||
NCCL_TRY(ncclMemAlloc(&ctx->buf, ctx->nbytes), "ncclMemAlloc");
|
||||
|
||||
return reinterpret_cast<int64_t>(ctx.release());
|
||||
}
|
||||
|
||||
// Register the user window. Collective: every rank must call together.
|
||||
void register_ulysses_a2a_window(int64_t handle, int64_t comm_ptr) {
|
||||
auto* ctx = reinterpret_cast<UlyssesContext*>(handle);
|
||||
TORCH_CHECK(ctx != nullptr, "handle must come from allocate_ulysses_a2a");
|
||||
TORCH_CHECK(ctx->buf != nullptr && ctx->win == nullptr && !ctx->dev_comm_created,
|
||||
"ulysses a2a context is not in the allocated state");
|
||||
|
||||
const at::cuda::CUDAGuard device_guard(
|
||||
c10::Device(c10::DeviceType::CUDA, static_cast<c10::DeviceIndex>(ctx->device)));
|
||||
ctx->comm = reinterpret_cast<ncclComm_t>(comm_ptr);
|
||||
NCCL_TRY(ncclCommWindowRegister(ctx->comm, ctx->buf, ctx->nbytes, &ctx->win,
|
||||
NCCL_WIN_COLL_SYMMETRIC),
|
||||
"ncclCommWindowRegister");
|
||||
}
|
||||
|
||||
// Create the device communicator only after Python has voted that every rank
|
||||
// registered its window. This operation is collective as well.
|
||||
void create_ulysses_a2a_dev_comm(int64_t handle) {
|
||||
auto* ctx = reinterpret_cast<UlyssesContext*>(handle);
|
||||
TORCH_CHECK(ctx != nullptr, "handle must come from allocate_ulysses_a2a");
|
||||
TORCH_CHECK(ctx->comm != nullptr && ctx->win != nullptr && !ctx->dev_comm_created,
|
||||
"ulysses a2a context is not in the window-registered state");
|
||||
|
||||
const at::cuda::CUDAGuard device_guard(
|
||||
c10::Device(c10::DeviceType::CUDA, static_cast<c10::DeviceIndex>(ctx->device)));
|
||||
|
||||
ncclDevCommRequirements reqs = NCCL_DEV_COMM_REQUIREMENTS_INITIALIZER;
|
||||
reqs.lsaBarrierCount = fi::kMaxBlocks;
|
||||
NCCL_TRY(ncclDevCommCreate(ctx->comm, &reqs, &ctx->devComm), "ncclDevCommCreate");
|
||||
ctx->dev_comm_created = true;
|
||||
}
|
||||
|
||||
static ncclResult_t first_error(ncclResult_t current, ncclResult_t next) {
|
||||
return current == ncclSuccess ? next : current;
|
||||
}
|
||||
|
||||
void dispose_ulysses_a2a(int64_t handle) {
|
||||
auto* ctx = reinterpret_cast<UlyssesContext*>(handle);
|
||||
if (ctx == nullptr) return;
|
||||
|
||||
const at::cuda::CUDAGuard device_guard(
|
||||
c10::Device(c10::DeviceType::CUDA, static_cast<c10::DeviceIndex>(ctx->device)));
|
||||
ncclResult_t result = ncclSuccess;
|
||||
if (ctx->comm != nullptr && ctx->dev_comm_created) {
|
||||
result = first_error(result, ncclDevCommDestroy(ctx->comm, &ctx->devComm));
|
||||
ctx->dev_comm_created = false;
|
||||
}
|
||||
if (ctx->comm != nullptr && ctx->win != nullptr) {
|
||||
result = first_error(result, ncclCommWindowDeregister(ctx->comm, ctx->win));
|
||||
ctx->win = nullptr;
|
||||
}
|
||||
if (ctx->buf != nullptr) {
|
||||
result = first_error(result, ncclMemFree(ctx->buf));
|
||||
ctx->buf = nullptr;
|
||||
}
|
||||
delete ctx;
|
||||
TORCH_CHECK(result == ncclSuccess, "ulysses a2a cleanup failed: rc=", static_cast<int>(result));
|
||||
}
|
||||
|
||||
// Whether the whole group is load-store accessible. NCCL determined this at
|
||||
// ncclCommInitRank.
|
||||
bool ulysses_lsa_covers_group(int64_t comm_ptr, int64_t world_size) {
|
||||
auto comm = reinterpret_cast<ncclComm_t>(comm_ptr);
|
||||
ncclCommProperties properties = NCCL_COMM_PROPERTIES_INITIALIZER;
|
||||
NCCL_TRY(ncclCommQueryProperties(comm, &properties), "ncclCommQueryProperties");
|
||||
ncclTeam_t lsa = ncclTeamLsa(comm);
|
||||
return properties.deviceApiSupport && lsa.nRanks == static_cast<int>(world_size);
|
||||
}
|
||||
|
||||
// Fused-transpose Ulysses all-to-all.
|
||||
// mode == 0: inp [B, S_local, H, D] -> out [B, S_global, H_local, D]
|
||||
// mode == 1: inp [B, S_global, H_local, D] -> out [B, S_local, H, D]
|
||||
// where H is the *global* head count and H_local = H / world_size.
|
||||
void ulysses_a2a(int64_t handle, torch::Tensor inp, torch::Tensor out, int64_t B, int64_t S_local,
|
||||
int64_t H, int64_t D, int64_t mode) {
|
||||
auto* ctx = reinterpret_cast<UlyssesContext*>(handle);
|
||||
TORCH_CHECK(ctx != nullptr, "handle must come from allocate_ulysses_a2a");
|
||||
|
||||
const at::cuda::CUDAGuard device_guard(inp.device());
|
||||
auto stream = at::cuda::getCurrentCUDAStream();
|
||||
|
||||
TORCH_CHECK(inp.is_cuda() && out.is_cuda(), "inp and out must be CUDA tensors");
|
||||
TORCH_CHECK(inp.is_contiguous() && out.is_contiguous(), "inp and out must be contiguous");
|
||||
TORCH_CHECK(inp.device() == out.device(), "inp and out must be on the same device");
|
||||
TORCH_CHECK(inp.scalar_type() == out.scalar_type(), "inp and out must share a dtype");
|
||||
TORCH_CHECK(inp.numel() == out.numel(), "inp and out must have equal element counts");
|
||||
TORCH_CHECK(inp.get_device() == ctx->device, "input is on CUDA device ", inp.get_device(),
|
||||
" but the Ulysses context belongs to device ", ctx->device);
|
||||
TORCH_CHECK(mode == 0 || mode == 1, "mode must be 0 or 1");
|
||||
TORCH_CHECK(inp.dim() == 4 && out.dim() == 4, "inp and out must be 4-D");
|
||||
|
||||
const int W = ctx->world;
|
||||
TORCH_CHECK(H % W == 0, "global head count must be divisible by world size");
|
||||
const int H_local = static_cast<int>(H / W);
|
||||
|
||||
const torch::Tensor& local_op = (mode == 0) ? inp : out; // [B, S_local, H, D]
|
||||
const torch::Tensor& global_op = (mode == 0) ? out : inp; // [B, S_global, H_local, D]
|
||||
TORCH_CHECK(local_op.size(0) == B && local_op.size(1) == S_local && local_op.size(2) == H &&
|
||||
local_op.size(3) == D,
|
||||
"the [B, S_local, H, D] operand of mode ", mode, " has shape (", local_op.size(0),
|
||||
", ", local_op.size(1), ", ", local_op.size(2), ", ", local_op.size(3),
|
||||
"), expected (", B, ", ", S_local, ", ", H, ", ", D, ")");
|
||||
TORCH_CHECK(global_op.size(0) == B && global_op.size(1) == W * S_local &&
|
||||
global_op.size(2) == H_local && global_op.size(3) == D,
|
||||
"the [B, S_global, H_local, D] operand of mode ", mode, " has shape (",
|
||||
global_op.size(0), ", ", global_op.size(1), ", ", global_op.size(2), ", ",
|
||||
global_op.size(3), "), expected (", B, ", ", W * S_local, ", ", H_local, ", ", D,
|
||||
")");
|
||||
|
||||
const size_t out_bytes = out.numel() * out.element_size();
|
||||
TORCH_CHECK(out_bytes <= ctx->nbytes, "operand of ", out_bytes,
|
||||
" bytes exceeds the window capacity ", ctx->nbytes);
|
||||
|
||||
const int64_t num_rows = B * static_cast<int64_t>(W) * S_local;
|
||||
const int blocks =
|
||||
static_cast<int>(std::max<int64_t>(1, std::min<int64_t>(fi::kMaxBlocks, num_rows)));
|
||||
const int threads = fi::kUlyssesThreads;
|
||||
|
||||
#define LAUNCH_ULYSSES_A2A(T, NG, MODE) \
|
||||
fi::ulysses_a2a_kernel<T, NG, MODE><<<blocks, threads, 0, stream>>>( \
|
||||
reinterpret_cast<const T*>(inp.data_ptr()), ctx->devComm, ctx->win, /*off=*/0, \
|
||||
ctx->rank, static_cast<int>(B), static_cast<int>(S_local), H_local, static_cast<int>(D))
|
||||
|
||||
#define DISPATCH_NGPUS(T, MODE) \
|
||||
switch (W) { \
|
||||
case 2: \
|
||||
LAUNCH_ULYSSES_A2A(T, 2, MODE); \
|
||||
break; \
|
||||
case 4: \
|
||||
LAUNCH_ULYSSES_A2A(T, 4, MODE); \
|
||||
break; \
|
||||
case 6: \
|
||||
LAUNCH_ULYSSES_A2A(T, 6, MODE); \
|
||||
break; \
|
||||
case 8: \
|
||||
LAUNCH_ULYSSES_A2A(T, 8, MODE); \
|
||||
break; \
|
||||
default: \
|
||||
TORCH_CHECK(false, "ulysses_a2a only supports world size in (2,4,6,8)"); \
|
||||
}
|
||||
|
||||
#define DISPATCH_DTYPE(MODE) \
|
||||
switch (out.scalar_type()) { \
|
||||
case at::ScalarType::Float: { \
|
||||
DISPATCH_NGPUS(float, MODE); \
|
||||
break; \
|
||||
} \
|
||||
case at::ScalarType::Half: { \
|
||||
DISPATCH_NGPUS(half, MODE); \
|
||||
break; \
|
||||
} \
|
||||
case at::ScalarType::BFloat16: { \
|
||||
DISPATCH_NGPUS(nv_bfloat16, MODE); \
|
||||
break; \
|
||||
} \
|
||||
default: \
|
||||
TORCH_CHECK(false, "ulysses_a2a only supports float32, float16 and bfloat16, got ", \
|
||||
out.scalar_type()); \
|
||||
}
|
||||
|
||||
if (mode == 0) {
|
||||
DISPATCH_DTYPE(0);
|
||||
} else {
|
||||
DISPATCH_DTYPE(1);
|
||||
}
|
||||
|
||||
#undef DISPATCH_DTYPE
|
||||
#undef DISPATCH_NGPUS
|
||||
#undef LAUNCH_ULYSSES_A2A
|
||||
|
||||
TORCH_CHECK(cudaGetLastError() == cudaSuccess, "ulysses_a2a kernel launch failed");
|
||||
// Copy this rank's completed result out of the window.
|
||||
auto status = cudaMemcpyAsync(out.data_ptr(), ctx->buf, out_bytes, cudaMemcpyDeviceToDevice,
|
||||
stream);
|
||||
TORCH_CHECK(status == cudaSuccess, "ulysses_a2a copy-out failed: ", cudaGetErrorString(status));
|
||||
}
|
||||
|
||||
void register_ulysses_a2a(pybind11::module_& m) {
|
||||
m.def("allocate_ulysses_a2a", &allocate_ulysses_a2a, "allocate a local ulysses a2a window");
|
||||
m.def("register_ulysses_a2a_window", ®ister_ulysses_a2a_window,
|
||||
"register the ulysses a2a window collectively");
|
||||
m.def("create_ulysses_a2a_dev_comm", &create_ulysses_a2a_dev_comm,
|
||||
"create the ulysses a2a device communicator collectively");
|
||||
m.def("dispose_ulysses_a2a", &dispose_ulysses_a2a, "release a ulysses a2a context");
|
||||
m.def("ulysses_lsa_covers_group", &ulysses_lsa_covers_group,
|
||||
"whether the whole group is load-store accessible");
|
||||
m.def("ulysses_a2a", &ulysses_a2a, "fused-transpose Ulysses all-to-all over NVLink");
|
||||
}
|
||||
@@ -22,6 +22,11 @@ extern std::vector<torch::Tensor> block_sparse_attention_backward(
|
||||
);
|
||||
#endif
|
||||
|
||||
#ifdef FASTVIDEO_KERNEL_COMPILE_ULYSSES_A2A
|
||||
// Ulysses sequence-parallel all-to-all (csrc/comm/)
|
||||
void register_ulysses_a2a(pybind11::module_ &);
|
||||
#endif
|
||||
|
||||
// TurboDiffusion kernels
|
||||
void register_quant(pybind11::module_ &);
|
||||
void register_rms_norm(pybind11::module_ &);
|
||||
@@ -61,6 +66,11 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
|
||||
m.def("block_sparse_bwd", torch::wrap_pybind_function(block_sparse_attention_backward), "block sparse attention backward (Hopper)");
|
||||
#endif
|
||||
|
||||
#ifdef FASTVIDEO_KERNEL_COMPILE_ULYSSES_A2A
|
||||
// Ulysses sequence-parallel all-to-all
|
||||
register_ulysses_a2a(m);
|
||||
#endif
|
||||
|
||||
// TurboDiffusion
|
||||
register_quant(m);
|
||||
register_rms_norm(m);
|
||||
|
||||
@@ -0,0 +1,165 @@
|
||||
/*
|
||||
* Copyright (c) 2025 by FlashInfer team.
|
||||
*
|
||||
* Adapted from flashinfer-ai/flashinfer @ 8a94642d83cba0939035868fb6c309b4474a13d6
|
||||
* (PR #3820), which in turn adapted ThunderKittens' NVLink all-to-all:
|
||||
* https://github.com/HazyResearch/ThunderKittens/blob/main/kernels/parallel/all_to_all/all_to_all.cu
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
|
||||
// Fused-transpose Ulysses all-to-all over NVLink. Peer addresses come from
|
||||
// ncclGetLsaPointer and synchronization from ncclLsaBarrierSession; the index
|
||||
// math and slab decomposition below are upstream's.
|
||||
//
|
||||
// head_dim == 2 layout, uniform sequence splits. With
|
||||
// W = ulysses world size
|
||||
// H_local = H / W
|
||||
// S_global = S_local * W
|
||||
//
|
||||
// mode == 0 (input a2a): [B, S_local, H, D] -> [B, S_global, H_local, D]
|
||||
// y_r[b, j*S_local + s, hl, d] = x_j[b, s, r*H_local + hl, d]
|
||||
// mode == 1 (output a2a): [B, S_global, H_local, D] -> [B, S_local, H, D]
|
||||
// out_j[b, s, r*H_local + hl, d] = u_r[b, j*S_local + s, hl, d]
|
||||
//
|
||||
// In both modes the unit of transfer is a contiguous (H_local * D) block, so
|
||||
// every cross-GPU store is fully coalesced.
|
||||
|
||||
#ifndef FASTVIDEO_COMM_ULYSSES_ALL_TO_ALL_CUH_
|
||||
#define FASTVIDEO_COMM_ULYSSES_ALL_TO_ALL_CUH_
|
||||
|
||||
#include <cstdint>
|
||||
|
||||
#include <nccl.h>
|
||||
#include <nccl_device.h>
|
||||
|
||||
namespace fastvideo {
|
||||
namespace comm {
|
||||
namespace ulysses {
|
||||
|
||||
constexpr int kUlyssesThreads = 512;
|
||||
// Deliberately modest: this is link-bandwidth bound, so a small grid leaves the
|
||||
// rest of the GPU free without costing throughput.
|
||||
constexpr int kMaxBlocks = 36;
|
||||
|
||||
// Shared movement body for the fused-transpose all-to-all (no barriers).
|
||||
//
|
||||
// Rows are ordered ((b * W + peer) * S_local + s), so consecutive rows share a
|
||||
// (batch, peer) and are contiguous on the gather side of the transpose. Each
|
||||
// block takes a contiguous slab of rows and flattens its threads over the 16B
|
||||
// units in it, so consecutive lanes address one peer buffer back to back and the
|
||||
// remote writes coalesce into large bursts rather than (H_local * D)-sized
|
||||
// scattered ones.
|
||||
template <typename T, int NGPUS, int MODE>
|
||||
__device__ __forceinline__ void ulysses_a2a_move(const T* __restrict__ local_in,
|
||||
void* const* peer_ptrs, int rank, int B,
|
||||
int S_local, int H_local, int D) {
|
||||
static_assert(MODE == 0 || MODE == 1, "MODE must be 0 or 1");
|
||||
const int W = NGPUS;
|
||||
const int64_t H = static_cast<int64_t>(H_local) * W;
|
||||
const int64_t S_global = static_cast<int64_t>(S_local) * W;
|
||||
const int64_t block_len = static_cast<int64_t>(H_local) * D; // elements/row
|
||||
const int64_t num_rows = static_cast<int64_t>(B) * W * S_local;
|
||||
|
||||
// 16B-vectorized fast path when every row is 16B aligned (the common case:
|
||||
// contiguous bf16/fp16/fp32 tensors with block_len * sizeof(T) % 16 == 0).
|
||||
using Vec = int4;
|
||||
constexpr int kVecBytes = sizeof(Vec);
|
||||
const int64_t row_bytes = block_len * static_cast<int64_t>(sizeof(T));
|
||||
const bool vec_ok =
|
||||
(row_bytes % kVecBytes) == 0 && (reinterpret_cast<uintptr_t>(local_in) % kVecBytes) == 0;
|
||||
|
||||
// Contiguous slab of rows for this block.
|
||||
const int64_t rows_per_block = (num_rows + gridDim.x - 1) / gridDim.x;
|
||||
const int64_t row_lo = static_cast<int64_t>(blockIdx.x) * rows_per_block;
|
||||
int64_t row_hi = row_lo + rows_per_block;
|
||||
if (row_hi > num_rows) row_hi = num_rows;
|
||||
if (row_lo >= row_hi) return;
|
||||
|
||||
const int tid = threadIdx.x;
|
||||
const int nthr = blockDim.x;
|
||||
|
||||
// Decode (b, peer, s) and compute src/dst element offsets for a given row.
|
||||
auto offsets = [&](int64_t row, int64_t& src_off, int64_t& dst_off) {
|
||||
const int64_t s = row % S_local;
|
||||
const int64_t tmp = row / S_local;
|
||||
const int64_t peer = tmp % W;
|
||||
const int64_t b = tmp / W;
|
||||
if constexpr (MODE == 0) {
|
||||
src_off = ((b * S_local + s) * H + peer * H_local) * D;
|
||||
dst_off = (b * S_global + static_cast<int64_t>(rank) * S_local + s) * block_len;
|
||||
} else {
|
||||
src_off = (b * S_global + peer * S_local + s) * block_len;
|
||||
dst_off = ((b * S_local + s) * H + static_cast<int64_t>(rank) * H_local) * D;
|
||||
}
|
||||
return peer;
|
||||
};
|
||||
|
||||
if (vec_ok) {
|
||||
const int64_t units_per_row = row_bytes / kVecBytes;
|
||||
const int64_t total_units = (row_hi - row_lo) * units_per_row;
|
||||
for (int64_t u = tid; u < total_units; u += nthr) {
|
||||
const int64_t local_row = u / units_per_row;
|
||||
const int64_t unit = u - local_row * units_per_row;
|
||||
const int64_t row = row_lo + local_row;
|
||||
int64_t src_off, dst_off;
|
||||
const int64_t peer = offsets(row, src_off, dst_off);
|
||||
const Vec* s4 = reinterpret_cast<const Vec*>(local_in + src_off);
|
||||
Vec* d4 = reinterpret_cast<Vec*>((T*)peer_ptrs[peer] + dst_off);
|
||||
d4[unit] = s4[unit];
|
||||
}
|
||||
} else {
|
||||
// Scalar fallback (unaligned / odd shapes).
|
||||
for (int64_t row = row_lo; row < row_hi; ++row) {
|
||||
int64_t src_off, dst_off;
|
||||
const int64_t peer = offsets(row, src_off, dst_off);
|
||||
const T* s_ptr = local_in + src_off;
|
||||
T* d_ptr = (T*)peer_ptrs[peer] + dst_off;
|
||||
for (int64_t i = tid; i < block_len; i += nthr) {
|
||||
d_ptr[i] = s_ptr[i];
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// The transfer mode is a compile-time template parameter so the address math
|
||||
// specializes and the coalesced slab decomposition is fully unrolled per mode.
|
||||
template <typename T, int NGPUS, int MODE>
|
||||
__global__ void __launch_bounds__(kUlyssesThreads, 1)
|
||||
ulysses_a2a_kernel(const T* __restrict__ local_in, ncclDevComm devComm, ncclWindow_t win,
|
||||
size_t win_offset, int rank, int B, int S_local, int H_local, int D) {
|
||||
// Resolved once; the movement loop would otherwise call this per 16B store.
|
||||
void* peer_ptrs[NGPUS];
|
||||
#pragma unroll
|
||||
for (int p = 0; p < NGPUS; ++p) {
|
||||
peer_ptrs[p] = ncclGetLsaPointer(win, win_offset, p);
|
||||
}
|
||||
|
||||
// Each CTA owns one barrier generation. Sharing index 0 across independently
|
||||
// scheduled CTAs races the generation counter and is unsupported by NCCL's
|
||||
// device API. The host reserves kMaxBlocks slots when it creates devComm.
|
||||
ncclLsaBarrierSession<ncclCoopCta> bar(
|
||||
ncclCoopCta(), devComm, ncclTeamTagLsa(), /*index=*/blockIdx.x);
|
||||
|
||||
// Every rank must have entered before anyone writes into peer buffers.
|
||||
bar.sync(ncclCoopCta(), cuda::memory_order_relaxed);
|
||||
ulysses_a2a_move<T, NGPUS, MODE>(local_in, peer_ptrs, rank, B, S_local, H_local, D);
|
||||
// Release-acquire: all peer writes visible before a rank reads its window.
|
||||
bar.sync(ncclCoopCta(), cuda::memory_order_acq_rel);
|
||||
}
|
||||
|
||||
} // namespace ulysses
|
||||
} // namespace comm
|
||||
} // namespace fastvideo
|
||||
|
||||
#endif // FASTVIDEO_COMM_ULYSSES_ALL_TO_ALL_CUH_
|
||||
@@ -9,7 +9,7 @@ build-backend = "scikit_build_core.build"
|
||||
|
||||
[project]
|
||||
name = "fastvideo-kernel"
|
||||
version = "0.3.2"
|
||||
version = "0.3.4"
|
||||
description = "Unified CUDA kernels for FastVideo"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.10"
|
||||
|
||||
@@ -49,6 +49,36 @@ def _force_tk() -> bool:
|
||||
return os.environ.get("FASTVIDEO_VSA_TK", "0") == "1"
|
||||
|
||||
|
||||
def _force_sm100a() -> bool:
|
||||
"""True iff the sm_100a (Blackwell) forward is explicitly opted into.
|
||||
|
||||
Opt-in only (same env the H3 backend honors): the sm_100a extension is
|
||||
forward-only, so this routing pairs it with the Triton backward -- its lse
|
||||
is already in Triton's M format. Honored only when
|
||||
``block_sparse_attn_sm100a.is_supported`` passes. Unsupported 64-token
|
||||
metadata falls through to the default selection; unsupported 128-token
|
||||
metadata raises because Triton has no compatible fallback.
|
||||
``FASTVIDEO_VSA_TRITON`` still wins.
|
||||
"""
|
||||
return os.environ.get("FASTVIDEO_VSA_SM100A", "0") == "1"
|
||||
|
||||
|
||||
def _sm100a_is_supported(q: torch.Tensor, variable_block_sizes: torch.Tensor) -> bool:
|
||||
try:
|
||||
from fastvideo_kernel import block_sparse_attn_sm100a as vsa_sm100a
|
||||
except Exception:
|
||||
return False
|
||||
return vsa_sm100a.is_supported(q, variable_block_sizes)
|
||||
|
||||
|
||||
def _infer_block_size(q: torch.Tensor, variable_block_sizes: torch.Tensor) -> int:
|
||||
num_blocks = variable_block_sizes.numel()
|
||||
seq_len = q.shape[2]
|
||||
if num_blocks == 0 or seq_len % num_blocks != 0:
|
||||
return 0
|
||||
return seq_len // num_blocks
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Index helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -339,6 +369,76 @@ def _backward_sm90(ctx, grad_o, grad_lse):
|
||||
|
||||
block_sparse_attn_sm90.register_autograd(_backward_sm90, setup_context=_setup_context_sm90)
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# SM100A backend custom op (index-native)
|
||||
#
|
||||
# Forward runs the sm_100a CUDA extension; backward reuses the Triton kernels.
|
||||
# The sm_100a forward emits lse in exactly Triton's M format (max*log2e +
|
||||
# log2(l)), so the pairing needs no conversion. The Triton backward is
|
||||
# hardcoded to 64-token blocks, hence the block-size assert below.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@torch.library.custom_op(
|
||||
"fastvideo_kernel::block_sparse_attn_sm100a",
|
||||
mutates_args=(),
|
||||
device_types="cuda",
|
||||
)
|
||||
def block_sparse_attn_sm100a_op(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
v: torch.Tensor,
|
||||
q2k_idx: torch.Tensor,
|
||||
q2k_num: torch.Tensor,
|
||||
variable_block_sizes: torch.Tensor,
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
from fastvideo_kernel.block_sparse_attn_sm100a import block_sparse_attn_sm100a
|
||||
|
||||
o, M = block_sparse_attn_sm100a(q, k, v, q2k_idx, q2k_num, variable_block_sizes,
|
||||
need_lse=True)
|
||||
return o, M
|
||||
|
||||
|
||||
@torch.library.register_fake("fastvideo_kernel::block_sparse_attn_sm100a")
|
||||
def _block_sparse_attn_sm100a_fake(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
v: torch.Tensor,
|
||||
q2k_idx: torch.Tensor,
|
||||
q2k_num: torch.Tensor,
|
||||
variable_block_sizes: torch.Tensor,
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
o = torch.empty_like(q)
|
||||
M = torch.empty(
|
||||
(q.shape[0], q.shape[1], q.shape[2]),
|
||||
device=q.device,
|
||||
dtype=torch.float32,
|
||||
)
|
||||
return o, M
|
||||
|
||||
|
||||
def _setup_context_sm100a(ctx, inputs, output):
|
||||
q, k, v, q2k_idx, q2k_num, variable_block_sizes = inputs
|
||||
o, M = output
|
||||
ctx.save_for_backward(q, k, v, o, M, q2k_idx, q2k_num, variable_block_sizes)
|
||||
|
||||
|
||||
def _backward_sm100a(ctx, grad_o, grad_M):
|
||||
q, k, v, o, M, q2k_idx, q2k_num, variable_block_sizes = ctx.saved_tensors
|
||||
block = q.shape[2] // variable_block_sizes.numel()
|
||||
if block != 64:
|
||||
raise RuntimeError(
|
||||
"block_sparse_attn_sm100a backward pairs the sm_100a forward with the "
|
||||
f"Triton backward, which is hardcoded to 64-token blocks; got {block}. "
|
||||
"Run 128-token-block metadata without grad, or use the Triton forward.")
|
||||
dq, dk, dv = block_sparse_attn_backward_triton(grad_o, q, k, v, o, M, q2k_idx,
|
||||
q2k_num, variable_block_sizes)
|
||||
return dq, dk, dv, None, None, None
|
||||
|
||||
|
||||
block_sparse_attn_sm100a_op.register_autograd(_backward_sm100a,
|
||||
setup_context=_setup_context_sm100a)
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Public API
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -364,11 +464,24 @@ def block_sparse_attn_from_indices(
|
||||
|
||||
# Backend resolution:
|
||||
# - FASTVIDEO_VSA_TRITON forces Triton everywhere.
|
||||
# - FASTVIDEO_VSA_SM100A opts into the sm_100a forward (Triton backward).
|
||||
# Unsupported 64-token metadata falls through; unsupported 128-token
|
||||
# metadata raises because Triton cannot consume it.
|
||||
# - FASTVIDEO_VSA_TK requests sm_90 TK; honored only when it's actually
|
||||
# available (else falls through to the default below).
|
||||
# - Otherwise: TK on sm_90 if available, else Triton.
|
||||
if _force_triton():
|
||||
use_sm90 = False
|
||||
elif _force_sm100a():
|
||||
if _sm100a_is_supported(q, variable_block_sizes):
|
||||
return block_sparse_attn_sm100a_op(q, k, v, q2k_idx, q2k_num,
|
||||
variable_block_sizes)
|
||||
if _infer_block_size(q, variable_block_sizes) == 128:
|
||||
raise NotImplementedError(
|
||||
"128-token block-sparse attention requires the sm_100a forward; "
|
||||
"the Triton fallback only supports 64-token blocks, and the "
|
||||
"sm_100a route is unavailable for this input.")
|
||||
use_sm90 = sm90_available
|
||||
elif _force_tk():
|
||||
use_sm90 = sm90_available
|
||||
else:
|
||||
|
||||
@@ -84,6 +84,107 @@ def is_supported(q: torch.Tensor, variable_block_sizes: torch.Tensor) -> bool:
|
||||
return True
|
||||
|
||||
|
||||
@torch.library.custom_op(
|
||||
"fastvideo_kernel::block_sparse_attn_sm100a_inference",
|
||||
mutates_args=(),
|
||||
device_types="cuda",
|
||||
)
|
||||
def _block_sparse_attn_sm100a_inference(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
v: torch.Tensor,
|
||||
q2k_idx: torch.Tensor,
|
||||
q2k_num: torch.Tensor,
|
||||
variable_block_sizes: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
"""Opaque no-LSE launch used by the inference-only sm_100a route.
|
||||
|
||||
The extension is exposed as a raw pybind function rather than a dispatcher
|
||||
op. Calling it directly makes Dynamo descend through a Python/C++ boundary
|
||||
that has no fake implementation, so ``torch.compile(fullgraph=True)``
|
||||
cannot capture a sparse H3 block. Keep that boundary inside this custom op;
|
||||
its inputs have already been normalized by the public wrapper below.
|
||||
"""
|
||||
fwd = _FWD_BY_BLOCK[_block_size(q, variable_block_sizes)]
|
||||
sm_scale = 1.0 / (q.shape[-1]**0.5)
|
||||
res = fwd(q, k, v, None, q2k_idx, q2k_num, variable_block_sizes, sm_scale, False)
|
||||
return res[0]
|
||||
|
||||
|
||||
@torch.library.register_fake("fastvideo_kernel::block_sparse_attn_sm100a_inference")
|
||||
def _block_sparse_attn_sm100a_inference_fake(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
v: torch.Tensor,
|
||||
q2k_idx: torch.Tensor,
|
||||
q2k_num: torch.Tensor,
|
||||
variable_block_sizes: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
# The C++ binding allocates its output with torch::empty_like(q).
|
||||
return torch.empty_like(q)
|
||||
|
||||
|
||||
@torch.library.custom_op(
|
||||
"fastvideo_kernel::block_sparse_attn_sm100a_from_mask_inference",
|
||||
mutates_args=(),
|
||||
device_types="cuda",
|
||||
)
|
||||
def _block_sparse_attn_sm100a_from_mask_inference(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
v: torch.Tensor,
|
||||
block_map: torch.Tensor,
|
||||
variable_block_sizes: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
"""Opaque mask compaction plus no-LSE sm_100a launch for inference.
|
||||
|
||||
H3 naturally produces a bool block map. Its Triton ``map_to_index`` call
|
||||
must live behind the same opaque boundary as the raw pybind launch;
|
||||
otherwise Dynamo sees that kernel before reaching the index-native custom
|
||||
op and full-graph capture still fails.
|
||||
"""
|
||||
from fastvideo_kernel.triton_kernels.index import map_to_index
|
||||
|
||||
q2k_idx, q2k_num = map_to_index(block_map)
|
||||
return _block_sparse_attn_sm100a_inference(
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
q2k_idx.to(torch.int32).contiguous(),
|
||||
q2k_num.to(torch.int32).contiguous(),
|
||||
variable_block_sizes,
|
||||
)
|
||||
|
||||
|
||||
@torch.library.register_fake("fastvideo_kernel::block_sparse_attn_sm100a_from_mask_inference")
|
||||
def _block_sparse_attn_sm100a_from_mask_inference_fake(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
v: torch.Tensor,
|
||||
block_map: torch.Tensor,
|
||||
variable_block_sizes: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
return torch.empty_like(q)
|
||||
|
||||
|
||||
def block_sparse_attn_sm100a_from_mask(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
v: torch.Tensor,
|
||||
block_map: torch.Tensor,
|
||||
variable_block_sizes: torch.Tensor,
|
||||
) -> Tuple[torch.Tensor, None]:
|
||||
"""Inference forward from a bool block map, with compaction kept opaque."""
|
||||
out = _block_sparse_attn_sm100a_from_mask_inference(
|
||||
q.contiguous(),
|
||||
k.contiguous(),
|
||||
v.contiguous(),
|
||||
block_map.to(torch.bool).contiguous(),
|
||||
variable_block_sizes.to(torch.int32).contiguous(),
|
||||
)
|
||||
return out, None
|
||||
|
||||
|
||||
def block_sparse_attn_sm100a(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
@@ -92,13 +193,24 @@ def block_sparse_attn_sm100a(
|
||||
q2k_num: torch.Tensor,
|
||||
variable_block_sizes: torch.Tensor,
|
||||
need_lse: bool = True,
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
) -> Tuple[torch.Tensor, torch.Tensor | None]:
|
||||
"""Forward pass. Returns ``(out, lse)``; ``out`` has q's layout."""
|
||||
fwd = _FWD_BY_BLOCK[_block_size(q, variable_block_sizes)]
|
||||
q = q.contiguous()
|
||||
k = k.contiguous()
|
||||
v = v.contiguous()
|
||||
idx = q2k_idx.to(torch.int32).contiguous()
|
||||
num = q2k_num.to(torch.int32).contiguous()
|
||||
vbs = variable_block_sizes.to(torch.int32).contiguous()
|
||||
|
||||
if not need_lse:
|
||||
# This is the production inference path. The custom op keeps the raw
|
||||
# pybind launch opaque to Dynamo while its fake kernel carries output
|
||||
# metadata through full-graph capture.
|
||||
return _block_sparse_attn_sm100a_inference(q, k, v, idx, num, vbs), None
|
||||
|
||||
# Preserve the established LSE-producing path for correctness tests and
|
||||
# any future forward/backward pairing; only inference needs the opaque op.
|
||||
fwd = _FWD_BY_BLOCK[_block_size(q, vbs)]
|
||||
sm_scale = 1.0 / (q.shape[-1]**0.5)
|
||||
res = fwd(q.contiguous(), k.contiguous(), v.contiguous(), None,
|
||||
idx, num, vbs, sm_scale, need_lse)
|
||||
return (res[0], res[1]) if need_lse else (res[0], None)
|
||||
res = fwd(q, k, v, None, idx, num, vbs, sm_scale, True)
|
||||
return res[0], res[1]
|
||||
|
||||
@@ -0,0 +1,86 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""Ulysses sequence-parallel all-to-all ops.
|
||||
|
||||
Thin wrappers over csrc/comm/ulysses_all_to_all.cu. The kernel stores directly
|
||||
into peers' memory through NCCL's device API, so the caller supplies an
|
||||
ncclComm_t for the group.
|
||||
"""
|
||||
|
||||
import torch
|
||||
|
||||
try:
|
||||
from fastvideo_kernel._C import fastvideo_kernel_ops as _ops
|
||||
except ImportError: # pragma: no cover - no compiled extension in this install
|
||||
_ops = None
|
||||
|
||||
_SUPPORTED_WORLD_SIZES = (2, 4, 6, 8)
|
||||
_REQUIRED_OPS = (
|
||||
"allocate_ulysses_a2a",
|
||||
"register_ulysses_a2a_window",
|
||||
"create_ulysses_a2a_dev_comm",
|
||||
"dispose_ulysses_a2a",
|
||||
"ulysses_lsa_covers_group",
|
||||
"ulysses_a2a",
|
||||
)
|
||||
|
||||
|
||||
def is_available() -> bool:
|
||||
"""Whether this wheel was built with the Ulysses all-to-all kernel."""
|
||||
return _ops is not None and all(hasattr(_ops, name) for name in _REQUIRED_OPS)
|
||||
|
||||
|
||||
def _require() -> None:
|
||||
if not is_available():
|
||||
raise RuntimeError(
|
||||
"the Ulysses all-to-all kernel is not present in this fastvideo-kernel build; "
|
||||
"rebuild with ./build.sh or install a wheel that includes csrc/comm/")
|
||||
|
||||
|
||||
def lsa_covers_group(comm_ptr: int, world_size: int) -> bool:
|
||||
"""Whether every rank in the group is load-store accessible to every other."""
|
||||
_require()
|
||||
return bool(_ops.ulysses_lsa_covers_group(int(comm_ptr), int(world_size)))
|
||||
|
||||
|
||||
def allocate(nbytes: int, rank: int, world_size: int, device_index: int) -> int:
|
||||
"""Allocate one rank's local symmetric window without a collective."""
|
||||
_require()
|
||||
if world_size not in _SUPPORTED_WORLD_SIZES:
|
||||
raise ValueError(f"ulysses a2a supports world sizes {_SUPPORTED_WORLD_SIZES}, "
|
||||
f"got {world_size}")
|
||||
return int(_ops.allocate_ulysses_a2a(int(nbytes), int(rank), int(world_size), int(device_index)))
|
||||
|
||||
|
||||
def register_window(handle: int, comm_ptr: int) -> None:
|
||||
"""Register an allocated window with the supplied communicator.
|
||||
|
||||
Every rank in ``comm_ptr`` must call this together.
|
||||
"""
|
||||
_require()
|
||||
_ops.register_ulysses_a2a_window(int(handle), int(comm_ptr))
|
||||
|
||||
|
||||
def create_dev_comm(handle: int) -> None:
|
||||
"""Create the device communicator for a registered window collectively."""
|
||||
_require()
|
||||
_ops.create_ulysses_a2a_dev_comm(int(handle))
|
||||
|
||||
|
||||
def dispose(handle: int) -> None:
|
||||
"""Release a handle from :func:`allocate`. It is dangling afterwards."""
|
||||
_require()
|
||||
_ops.dispose_ulysses_a2a(int(handle))
|
||||
|
||||
|
||||
def all_to_all(handle: int, inp: torch.Tensor, out: torch.Tensor, B: int, S_local: int, H: int,
|
||||
D: int, mode: int) -> None:
|
||||
"""Run one fused all-to-all on the current stream, writing into ``out``.
|
||||
|
||||
``mode == 0``: ``[B, S_local, H, D] -> [B, S_global, H_local, D]``
|
||||
``mode == 1``: ``[B, S_global, H_local, D] -> [B, S_local, H, D]``
|
||||
|
||||
``H`` is the global head count. Every rank must call with consistent
|
||||
geometry in the same order.
|
||||
"""
|
||||
_require()
|
||||
_ops.ulysses_a2a(int(handle), inp, out, int(B), int(S_local), int(H), int(D), int(mode))
|
||||
+11
-4
@@ -16,14 +16,21 @@ import triton.language as tl
|
||||
import math # small utility needed by the sparse wrapper
|
||||
# ──────────────────────────── SPARSE ADDITION END ─────────────────────────────
|
||||
|
||||
# We don't run auto-tuning every time to keep the tutorial fast. Keeping
|
||||
# the code below and commenting out the equivalent parameters is convenient for
|
||||
# re-tuning.
|
||||
# BLOCK_M / BLOCK_N are fixed at 64 because they are structural, not tunable:
|
||||
# the kernel indexes the top-k list per BLOCK_M q-tile and addresses keys as
|
||||
# kv_idx * BLOCK_N, so both must match the granularity q2k_index and
|
||||
# variable_block_sizes were built at.
|
||||
#
|
||||
# num_stages / num_warps ARE free, and the previous {3, 4, 7} was inherited from
|
||||
# the upstream tutorial rather than tuned here. It skips 5 and 6; on Blackwell
|
||||
# (sm_121) the optimum is num_stages=5, so the search could not reach it. Both
|
||||
# block paths independently select 5 once it is available. Autotune still picks
|
||||
# per architecture, so other GPUs re-tune rather than inheriting this choice.
|
||||
configs = [
|
||||
triton.Config({'BLOCK_M': BM, 'BLOCK_N': BN}, num_stages=s, num_warps=w) \
|
||||
for BM in [64]\
|
||||
for BN in [64]\
|
||||
for s in [3, 4, 7]\
|
||||
for s in [2, 3, 4, 5, 6, 7]\
|
||||
for w in [4, 8]\
|
||||
]
|
||||
|
||||
|
||||
@@ -1 +1 @@
|
||||
__version__ = "0.3.2"
|
||||
__version__ = "0.3.4"
|
||||
|
||||
@@ -165,6 +165,8 @@ def legacy_from_pretrained_to_config(
|
||||
pipeline["workload_type"] = value
|
||||
elif key == "lora_path":
|
||||
components["lora_path"] = value
|
||||
elif key == "lora_strength":
|
||||
components["lora_strength"] = value
|
||||
elif key == "override_pipeline_cls_name":
|
||||
components["override_pipeline_cls_name"] = value
|
||||
elif key == "override_transformer_cls_name":
|
||||
@@ -283,6 +285,7 @@ def generator_config_to_fastvideo_args(config: GeneratorConfig | Mapping[str, An
|
||||
kwargs["pipeline_config"] = components.pipeline_config_path
|
||||
if components.lora_path is not None:
|
||||
kwargs["lora_path"] = components.lora_path
|
||||
kwargs["lora_strength"] = components.lora_strength
|
||||
if components.override_pipeline_cls_name is not None:
|
||||
kwargs["override_pipeline_cls_name"] = components.override_pipeline_cls_name
|
||||
if components.override_transformer_cls_name is not None:
|
||||
|
||||
@@ -94,6 +94,7 @@ class ComponentConfig:
|
||||
vae_weights: str | None = None
|
||||
upsampler_weights: str | None = None
|
||||
lora_path: str | None = None
|
||||
lora_strength: float = 1.0
|
||||
override_pipeline_cls_name: str | None = None
|
||||
override_transformer_cls_name: str | None = None
|
||||
|
||||
|
||||
@@ -10,6 +10,14 @@ from fastvideo.attention.utils.flash_attn_default import (
|
||||
flash_attn_func_compilable,
|
||||
)
|
||||
|
||||
if fa_version == "4":
|
||||
# The FA4 varlen wrapper is already a compile-safe custom op. Keep the
|
||||
# import conditional so FA2/FA3 environments do not need flash_attn.cute.
|
||||
from fastvideo.attention.utils.flash_attn_cute import (
|
||||
flash_attn_varlen_func as flash_attn_varlen_func_compilable, )
|
||||
else:
|
||||
flash_attn_varlen_func_compilable = None
|
||||
|
||||
from fastvideo.attention.backends.abstract import (
|
||||
AttentionBackend,
|
||||
AttentionImpl,
|
||||
@@ -38,6 +46,21 @@ except ImportError:
|
||||
_FA4_FP4_AVAILABLE = False
|
||||
|
||||
_FA4_QUANT_OPS: tuple | None = None
|
||||
_FA4_PACKED_VARLEN_CONFIG_LOGGED = False
|
||||
|
||||
|
||||
def _log_fa4_packed_varlen_config() -> None:
|
||||
"""Emit one configuration receipt per worker process.
|
||||
|
||||
This intentionally does not claim that a particular forward used the
|
||||
kernel: masks, gradients, batch size, unequal Q/K lengths, and NVFP4 are
|
||||
runtime guards evaluated later in ``_forward_impl``.
|
||||
"""
|
||||
global _FA4_PACKED_VARLEN_CONFIG_LOGGED
|
||||
if not _FA4_PACKED_VARLEN_CONFIG_LOGGED:
|
||||
logger.info("MiniMax-H3 dense attention: FA4 packed-varlen route configured (runtime guards apply)",
|
||||
local_main_process_only=False)
|
||||
_FA4_PACKED_VARLEN_CONFIG_LOGGED = True
|
||||
|
||||
|
||||
def _import_fa4_quant_ops() -> tuple:
|
||||
@@ -228,7 +251,21 @@ class FlashAttentionImpl(AttentionImpl):
|
||||
) -> None:
|
||||
self.causal = causal
|
||||
self.softmax_scale = softmax_scale
|
||||
self.nvfp4_fa4 = extra_impl_args.get("nvfp4_fa4", False) or os.environ.get("FASTVIDEO_NVFP4_FA4", "0") == "1"
|
||||
# MiniMax-H3's dense DiT explicitly enables this faster FA4 entry
|
||||
# point. It remains off for every other model and for the H3 text
|
||||
# refiner; grad-enabled calls stay on the established fixed path.
|
||||
self.fa4_packed_varlen = bool(extra_impl_args.get("fa4_packed_varlen", False))
|
||||
if self.fa4_packed_varlen and fa_version == "4":
|
||||
_log_fa4_packed_varlen_config()
|
||||
# An explicit ``nvfp4_fa4`` impl arg wins over the process-wide
|
||||
# FASTVIDEO_NVFP4_FA4 env opt-in, so precision-sensitive layers (e.g.
|
||||
# the FP32-pinned H3 VAE attention) can force-disable FP4 Q/K
|
||||
# quantization while the DiT keeps it. When the arg is absent the env
|
||||
# keeps its previous semantics.
|
||||
nvfp4_fa4 = extra_impl_args.get("nvfp4_fa4")
|
||||
if nvfp4_fa4 is None:
|
||||
nvfp4_fa4 = os.environ.get("FASTVIDEO_NVFP4_FA4", "0") == "1"
|
||||
self.nvfp4_fa4 = bool(nvfp4_fa4)
|
||||
if self.nvfp4_fa4:
|
||||
cap = torch.cuda.get_device_capability()
|
||||
assert cap in [(10, 0), (10, 3)], (f"NVFP4 FA4 requires Blackwell (sm100a/sm103a), got sm{cap[0]}{cap[1]}")
|
||||
@@ -323,6 +360,29 @@ class FlashAttentionImpl(AttentionImpl):
|
||||
elif self.nvfp4_fa4:
|
||||
output = self._forward_nvfp4(query, key, value)
|
||||
|
||||
elif (self.fa4_packed_varlen and fa_version == "4" and not torch.is_grad_enabled() and query.shape[0] == 1
|
||||
and query.shape[1] == key.shape[1] == value.shape[1]):
|
||||
# FA4's packed-varlen entry point is materially faster for H3's
|
||||
# long, single-document self-attention. Flatten only the batch
|
||||
# dimension and describe that one sequence with CUDA int32
|
||||
# cumulative lengths; the existing custom-op wrapper keeps this
|
||||
# route traceable under torch.compile(fullgraph=True).
|
||||
assert flash_attn_varlen_func_compilable is not None
|
||||
sequence_length = query.shape[1]
|
||||
cu_seqlens = torch.arange(2, dtype=torch.int32, device=query.device) * sequence_length
|
||||
output = flash_attn_varlen_func_compilable(
|
||||
query.squeeze(0),
|
||||
key.squeeze(0),
|
||||
value.squeeze(0),
|
||||
cu_seqlens,
|
||||
cu_seqlens,
|
||||
sequence_length,
|
||||
sequence_length,
|
||||
dropout_p=0.0,
|
||||
softmax_scale=self.softmax_scale,
|
||||
causal=self.causal,
|
||||
).unsqueeze(0)
|
||||
|
||||
else:
|
||||
# Route through the compilable wrapper so dynamo sees a
|
||||
# registered op (no graph break) for FA2/FA3; identical
|
||||
|
||||
@@ -34,9 +34,13 @@ for the tile-64 FORWARD only: ``FASTVIDEO_VSA_SM100A=1`` sends no-grad
|
||||
forwards through the sm_100a CUDA block-sparse kernel
|
||||
(``fastvideo_kernel.block_sparse_attn_sm100a``, upstream PR #1719 plus
|
||||
our per-q-tile ``q2k_num`` fix) when the extension is built, the device
|
||||
is sm_100, and the geometry qualifies; grad-tracking forwards and every
|
||||
backward stay on Triton unchanged. If the env is set but a precondition
|
||||
fails, the route logs one warning and falls back.
|
||||
is sm_100, and the geometry qualifies. The CUDA kernel assigns adjacent
|
||||
pairs of query tiles to CTAs, so an odd logical tile count receives one
|
||||
internal, zero-valid partner tile for the no-grad call only. Score search,
|
||||
the trained mask, gate-compress, and the returned packed sequence remain on
|
||||
the original logical tiles. Grad-tracking forwards and every backward stay
|
||||
on Triton unchanged. If the env is set but a precondition fails, the route
|
||||
logs one warning and falls back.
|
||||
"""
|
||||
|
||||
import functools
|
||||
@@ -90,6 +94,65 @@ VSA_H3_TILE_SHAPES: dict[int, tuple[int, int, int]] = {
|
||||
}
|
||||
|
||||
|
||||
@torch.library.custom_op(
|
||||
"fastvideo::h3_vsa_sm100a_from_mask_compat",
|
||||
mutates_args=(),
|
||||
device_types="cuda",
|
||||
)
|
||||
def _h3_vsa_sm100a_from_mask_compat(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
v: torch.Tensor,
|
||||
block_map: torch.Tensor,
|
||||
variable_block_sizes: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
"""Compile-safe mask adapter for kernel wheels predating the mask API."""
|
||||
if _sm100a is None or map_to_index is None:
|
||||
raise RuntimeError("The sm100a compatibility route requires the raw kernel and map_to_index")
|
||||
q2k_idx, q2k_num = map_to_index(block_map)
|
||||
out, _ = _sm100a.block_sparse_attn_sm100a(
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
q2k_idx.to(torch.int32).contiguous(),
|
||||
q2k_num.to(torch.int32).contiguous(),
|
||||
variable_block_sizes.to(torch.int32).contiguous(),
|
||||
need_lse=False,
|
||||
)
|
||||
return out
|
||||
|
||||
|
||||
@torch.library.register_fake("fastvideo::h3_vsa_sm100a_from_mask_compat")
|
||||
def _h3_vsa_sm100a_from_mask_compat_fake(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
v: torch.Tensor,
|
||||
block_map: torch.Tensor,
|
||||
variable_block_sizes: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
del k, v, block_map, variable_block_sizes
|
||||
return torch.empty_like(q)
|
||||
|
||||
|
||||
def _sm100a_has_compile_safe_mask_route(sm100a_mod: Any) -> bool:
|
||||
return (callable(getattr(sm100a_mod, "block_sparse_attn_sm100a_from_mask", None))
|
||||
or (callable(getattr(sm100a_mod, "block_sparse_attn_sm100a", None)) and map_to_index is not None))
|
||||
|
||||
|
||||
def _sm100a_from_mask(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
v: torch.Tensor,
|
||||
block_map: torch.Tensor,
|
||||
variable_block_sizes: torch.Tensor,
|
||||
) -> tuple[torch.Tensor, None]:
|
||||
"""Use the native mask entry when installed, otherwise the local adapter."""
|
||||
native = getattr(_sm100a, "block_sparse_attn_sm100a_from_mask", None)
|
||||
if callable(native):
|
||||
return native(q, k, v, block_map, variable_block_sizes)
|
||||
return _h3_vsa_sm100a_from_mask_compat(q, k, v, block_map, variable_block_sizes), None
|
||||
|
||||
|
||||
def token_tile_and_valid(variable_block_sizes: torch.Tensor,
|
||||
tile_elems: int = _TILE_ELEMS) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Per padded-token tile id and pad-validity mask.
|
||||
@@ -225,6 +288,14 @@ class MiniMaxH3VSABackend(AttentionBackend):
|
||||
return MiniMaxH3VSAMetadataBuilder
|
||||
|
||||
|
||||
class _MiniMaxH3VSATileBufferHolder:
|
||||
"""Builder-owned no-grad tile scratch and its active geometry."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.buffer: torch.Tensor | None = None
|
||||
self.untile_geometry: torch.Tensor | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class MiniMaxH3VSAMetadata(AttentionMetadata):
|
||||
total_seq_length: int
|
||||
@@ -233,23 +304,25 @@ class MiniMaxH3VSAMetadata(AttentionMetadata):
|
||||
exempt: bool
|
||||
variable_block_sizes: torch.Tensor
|
||||
untile_combined_index: torch.Tensor
|
||||
# Device-side copy of ``dense_layers``. Regional fullgraph capture uses
|
||||
# this tensor with each implementation's tensor-valued layer index so the
|
||||
# shared block code does not specialize once per Python ``layer_idx``.
|
||||
dense_layers_tensor: torch.Tensor
|
||||
# tokens per tile (256 or 64); selects the tile geometry AND the kernel
|
||||
# route in forward() (256 -> VSA-256 CuTe/Triton, 64 -> native Triton)
|
||||
tile_elems: int = _TILE_ELEMS
|
||||
# layers forced dense regardless of sparsity (probe-guided opt-outs)
|
||||
dense_layers: tuple[int, ...] = ()
|
||||
# Single-slot holder for the padded tile buffer, owned by the BUILDER so
|
||||
# one buffer serves the whole denoising loop (pad slots stay zero and
|
||||
# every non-pad slot is fully overwritten per tile(), so cross-step reuse
|
||||
# is valid; saves a ~1.4 GB alloc+memset per step at 720p). VSA-H3 runs
|
||||
# eager today — revisit the reuse if it ever goes under cudagraphs.
|
||||
tile_buf_holder: list = None # type: ignore[assignment]
|
||||
# Builder-owned padded tile buffer. It records the geometry that last
|
||||
# populated the allocation so a same-shaped geometry change can clear
|
||||
# stale pad rows once while steady-state denoising reuses the buffer.
|
||||
tile_buf_holder: _MiniMaxH3VSATileBufferHolder | None = None
|
||||
|
||||
|
||||
class MiniMaxH3VSAMetadataBuilder(AttentionMetadataBuilder):
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._tile_buf_holder: list = [None]
|
||||
self._tile_buf_holder = _MiniMaxH3VSATileBufferHolder()
|
||||
|
||||
def prepare(self) -> None:
|
||||
pass
|
||||
@@ -278,6 +351,7 @@ class MiniMaxH3VSAMetadataBuilder(AttentionMetadataBuilder):
|
||||
(_tile_partition_indices, variable_block_sizes, untile_combined_index, num_prefix_tiles,
|
||||
num_video_tiles) = _h3_tile_geometry(prefix_segments, dit_seq_shape, device, tile_shape)
|
||||
|
||||
dense_layers = tuple(int(layer) for layer in dense_layers)
|
||||
return MiniMaxH3VSAMetadata(
|
||||
current_timestep=current_timestep,
|
||||
VSA_sparsity=VSA_sparsity,
|
||||
@@ -288,7 +362,8 @@ class MiniMaxH3VSAMetadataBuilder(AttentionMetadataBuilder):
|
||||
variable_block_sizes=variable_block_sizes,
|
||||
untile_combined_index=untile_combined_index,
|
||||
tile_elems=int(tile_size),
|
||||
dense_layers=tuple(int(layer) for layer in dense_layers),
|
||||
dense_layers=dense_layers,
|
||||
dense_layers_tensor=torch.tensor(dense_layers, device=device, dtype=torch.int64),
|
||||
tile_buf_holder=self._tile_buf_holder,
|
||||
)
|
||||
|
||||
@@ -366,26 +441,109 @@ class MiniMaxH3VSAImpl(AttentionImpl):
|
||||
) -> None:
|
||||
self.prefix = prefix
|
||||
self.layer_idx = layer_idx_from_prefix(prefix, default=-1)
|
||||
self.head_size = head_size
|
||||
# None means the regional-compile preparation hook has not run. The
|
||||
# eager path deliberately ignores this cache and preserves its
|
||||
# request-time env/probe/fallback behavior; only Dynamo capture reads
|
||||
# the prepared, static route.
|
||||
self._regional_compile_sm100a_enabled: bool | None = None
|
||||
self._regional_compile_layer_idx: torch.Tensor | None = None
|
||||
|
||||
def prepare_for_regional_compile(self, device: torch.device) -> str | None:
|
||||
"""Resolve the inference-only sm_100a route before fullgraph capture.
|
||||
|
||||
The ordinary eager route probes the environment, extension, device,
|
||||
and tensor contract at every call so it can warn and fall back. Those
|
||||
Python/device-capability checks are not safe inside a regional
|
||||
``fullgraph=True`` block. Probe one representative tile-64 input on
|
||||
the loaded model's device now, then let ``forward`` specialize on the
|
||||
resulting plain bool while Dynamo is compiling.
|
||||
"""
|
||||
requested = os.environ.get(VSA_SM100A_ENV, "0") == "1"
|
||||
enabled = False
|
||||
reason = None if requested else f"{VSA_SM100A_ENV}=1 is required for compile-safe VSA-H3 attention"
|
||||
if requested:
|
||||
if _sm100a is None:
|
||||
reason = "fastvideo_kernel.block_sparse_attn_sm100a is not installed"
|
||||
elif not _sm100a_has_compile_safe_mask_route(_sm100a):
|
||||
reason = ("neither a native block_sparse_attn_sm100a_from_mask entry nor the raw sm100a "
|
||||
"kernel plus map_to_index compatibility route is installed")
|
||||
else:
|
||||
# Two 64-token blocks exercise the exact sm_100a inference
|
||||
# specialization while keeping the one-time probe tiny. The
|
||||
# kernel predicate checks extension presence, CUDA capability,
|
||||
# dtype/layout, head size, block size, and even block count
|
||||
# without reading metadata tensor contents.
|
||||
probe_query = torch.empty((1, 1, 128, self.head_size), device=device, dtype=torch.bfloat16)
|
||||
probe_block_sizes = torch.full((2, ), 64, device=device, dtype=torch.int32)
|
||||
reason = _sm100a_unavailable_reason(
|
||||
_sm100a,
|
||||
probe_query,
|
||||
probe_block_sizes,
|
||||
grad_mode=False,
|
||||
)
|
||||
enabled = reason is None
|
||||
|
||||
self._regional_compile_sm100a_enabled = enabled
|
||||
# Keep this marker unset when preparation fails. Generic/training
|
||||
# torch.compile must retain the established Triton attention route.
|
||||
self._regional_compile_layer_idx = (torch.tensor(self.layer_idx, device=device, dtype=torch.int64)
|
||||
if enabled else None)
|
||||
if enabled:
|
||||
route = ("native fastvideo-kernel mask entry" if callable(
|
||||
getattr(_sm100a, "block_sparse_attn_sm100a_from_mask", None)) else
|
||||
"FastVideo compatibility mask adapter")
|
||||
logger.info_once(f"VSA-H3 regional compile mask route: {route}")
|
||||
if requested and reason is not None:
|
||||
logger.warning_once(f"VSA-H3 regional compile is unavailable and will stay eager: {reason}")
|
||||
return reason
|
||||
|
||||
def tile(self, x: torch.Tensor, attn_metadata: MiniMaxH3VSAMetadata) -> torch.Tensor:
|
||||
"""Scatter rows into the padded tile buffer (pad positions stay zero).
|
||||
|
||||
The returned tensor aliases the builder-owned buffer; callers must
|
||||
consume it before the next ``tile()`` (both call sites in
|
||||
``forward()`` read it immediately).
|
||||
``forward()`` read it immediately). Odd tile-64 no-grad sm100a
|
||||
requests carry one additional all-zero tile internally; metadata and
|
||||
all observable outputs retain the logical geometry.
|
||||
"""
|
||||
if x.shape[1] != attn_metadata.total_seq_length:
|
||||
raise ValueError(f"VSA-H3 metadata was built for sequence length {attn_metadata.total_seq_length}, "
|
||||
f"got {x.shape[1]}. A non-packed sequence (e.g. the token refiner) is "
|
||||
"routed to the VSA-H3 backend; exclude it from the supported backends.")
|
||||
n_tiles = attn_metadata.variable_block_sizes.numel()
|
||||
target_shape = (x.shape[0], n_tiles * attn_metadata.tile_elems, x.shape[-2], x.shape[-1])
|
||||
grad_mode = torch.is_grad_enabled() and x.requires_grad
|
||||
compiling = torch.compiler.is_compiling()
|
||||
regional_compiling = compiling and self._regional_compile_layer_idx is not None
|
||||
if regional_compiling:
|
||||
sm100a_requested = bool(self._regional_compile_sm100a_enabled)
|
||||
elif compiling:
|
||||
# Training/generic compile keeps the long-standing Triton route.
|
||||
sm100a_requested = False
|
||||
else:
|
||||
sm100a_requested = os.environ.get(VSA_SM100A_ENV, "0") == "1"
|
||||
needs_sm100a_pair = (attn_metadata.tile_elems == 64 and n_tiles % 2 != 0 and not grad_mode and sm100a_requested)
|
||||
kernel_tiles = n_tiles + int(needs_sm100a_pair)
|
||||
target_shape = (x.shape[0], kernel_tiles * attn_metadata.tile_elems, x.shape[-2], x.shape[-1])
|
||||
|
||||
# single scatter: untile_combined_index maps original row i to its
|
||||
# padded slot, so this is exactly the inverse of postprocess_output
|
||||
# ``untile_combined_index`` maps each packed row to a logical tile
|
||||
# slot. Different geometries can share one transport shape; clear a
|
||||
# reused allocation once when the mapping identity changes so no old
|
||||
# valid row can survive as padding.
|
||||
holder = attn_metadata.tile_buf_holder
|
||||
holder[0] = scatter_into_tile_buf(x, target_shape, attn_metadata.untile_combined_index, holder[0])
|
||||
return holder[0]
|
||||
if holder is None:
|
||||
raise RuntimeError("VSA-H3 metadata has no builder-owned tile buffer holder")
|
||||
buffer_matches = (holder.buffer is not None and holder.buffer.shape == target_shape
|
||||
and holder.buffer.dtype == x.dtype and holder.buffer.device == x.device)
|
||||
if buffer_matches and holder.untile_geometry is not attn_metadata.untile_combined_index:
|
||||
holder.buffer.zero_()
|
||||
holder.buffer = scatter_into_tile_buf(x, target_shape, attn_metadata.untile_combined_index, holder.buffer)
|
||||
holder.untile_geometry = attn_metadata.untile_combined_index
|
||||
if needs_sm100a_pair:
|
||||
# A prior even geometry can reuse this allocation and may have
|
||||
# written the last tile as logical data.
|
||||
holder.buffer[:, n_tiles * attn_metadata.tile_elems:].zero_()
|
||||
return holder.buffer
|
||||
|
||||
def preprocess_qkv(self, qkv: torch.Tensor, attn_metadata: MiniMaxH3VSAMetadata) -> torch.Tensor:
|
||||
return self.tile(qkv, attn_metadata)
|
||||
@@ -401,28 +559,71 @@ class MiniMaxH3VSAImpl(AttentionImpl):
|
||||
gate_compress: torch.Tensor | None,
|
||||
attn_metadata: MiniMaxH3VSAMetadata,
|
||||
) -> torch.Tensor:
|
||||
compiling = torch.compiler.is_compiling()
|
||||
regional_compiling = compiling and self._regional_compile_layer_idx is not None
|
||||
|
||||
tile_elems = attn_metadata.tile_elems
|
||||
if regional_compiling and tile_elems != 64:
|
||||
raise RuntimeError("VSA-H3 regional fullgraph compile requires 64-token tiles; disable "
|
||||
"inference_torch_compile for tile-256/CuTe runs.")
|
||||
if tile_elems == 64:
|
||||
if block_sparse_attn_64_bhsd is None:
|
||||
raise NotImplementedError("fastvideo_kernel.block_sparse_attn is not installed")
|
||||
elif block_sparse_attn_256_bshd is None:
|
||||
raise NotImplementedError("fastvideo_kernel.block_sparse_attn_256 is not installed")
|
||||
|
||||
# probe-guided per-layer opt-out: diffuse layers run dense (all-True
|
||||
# mask) while the rest keep the configured sparsity
|
||||
layer_sparsity = 0.0 if self.layer_idx in attn_metadata.dense_layers else attn_metadata.VSA_sparsity
|
||||
probe_dir = probe_enabled()
|
||||
# Probe recording performs filesystem writes and host synchronizations,
|
||||
# so the loader keeps probe-enabled runs eager. Avoid even reading that
|
||||
# environment switch while Dynamo captures a regional full graph.
|
||||
# The metadata always describes the trained logical geometry.
|
||||
# ``tile()`` may append exactly one transport-only partner for an odd
|
||||
# tile-64 sm100a call. Keep score selection and the gate branch on the
|
||||
# logical prefix, and reject every other shape before a kernel sees it.
|
||||
n_tiles = attn_metadata.variable_block_sizes.numel()
|
||||
logical_seq_len = n_tiles * tile_elems
|
||||
pair_pad_seq_len = logical_seq_len + tile_elems
|
||||
pair_pad_is_valid = tile_elems == 64 and n_tiles % 2 != 0
|
||||
allowed_seq_lengths = (logical_seq_len, pair_pad_seq_len) if pair_pad_is_valid else (logical_seq_len, )
|
||||
if query.shape[1] not in allowed_seq_lengths:
|
||||
expected = (f"the logical length {logical_seq_len} or one sm100a partner tile "
|
||||
f"({pair_pad_seq_len})" if pair_pad_is_valid else f"the logical length {logical_seq_len}")
|
||||
raise ValueError(f"VSA-H3 tiled query has length {query.shape[1]}, expected {expected}.")
|
||||
has_sm100a_pair = query.shape[1] == pair_pad_seq_len
|
||||
for name, tensor in (("key", key), ("value", value)):
|
||||
if tensor.shape[1] != query.shape[1]:
|
||||
raise ValueError(f"VSA-H3 tiled {name} length {tensor.shape[1]} does not match query "
|
||||
f"length {query.shape[1]}.")
|
||||
if gate_compress is not None and gate_compress.shape[1] != query.shape[1]:
|
||||
raise ValueError(f"VSA-H3 tiled gate length {gate_compress.shape[1]} does not match query "
|
||||
f"length {query.shape[1]}.")
|
||||
|
||||
logical_query = query[:, :logical_seq_len]
|
||||
logical_key = key[:, :logical_seq_len]
|
||||
logical_value = value[:, :logical_seq_len]
|
||||
logical_gate = gate_compress[:, :logical_seq_len] if gate_compress is not None else None
|
||||
|
||||
# Probe-guided per-layer opt-out: diffuse layers run dense (all-True
|
||||
# mask) while the rest keep the configured sparsity. During regional
|
||||
# capture, keep the layer decision tensor-valued so the 50 block
|
||||
# instances reuse one graph instead of specializing on layer_idx.
|
||||
force_dense = None
|
||||
if regional_compiling:
|
||||
assert self._regional_compile_layer_idx is not None
|
||||
force_dense = (attn_metadata.dense_layers_tensor == self._regional_compile_layer_idx).any()
|
||||
layer_sparsity = attn_metadata.VSA_sparsity
|
||||
else:
|
||||
layer_sparsity = 0.0 if self.layer_idx in attn_metadata.dense_layers else attn_metadata.VSA_sparsity
|
||||
probe_dir = None if compiling else probe_enabled()
|
||||
|
||||
scores = None
|
||||
if layer_sparsity > 0.0 or gate_compress is not None or probe_dir is not None:
|
||||
q_pooled = _pool_tiles(query, attn_metadata.variable_block_sizes, tile_elems)
|
||||
k_pooled = _pool_tiles(key, attn_metadata.variable_block_sizes, tile_elems)
|
||||
q_pooled = _pool_tiles(logical_query, attn_metadata.variable_block_sizes, tile_elems)
|
||||
k_pooled = _pool_tiles(logical_key, attn_metadata.variable_block_sizes, tile_elems)
|
||||
scores = torch.matmul(q_pooled, k_pooled.transpose(-2, -1)) / (query.shape[-1]**0.5)
|
||||
if probe_dir is not None:
|
||||
record_probe(probe_dir, self.layer_idx, query, key, scores, attn_metadata)
|
||||
record_probe(probe_dir, self.layer_idx, logical_query, logical_key, scores, attn_metadata)
|
||||
|
||||
if scores is None:
|
||||
n_tiles = attn_metadata.variable_block_sizes.numel()
|
||||
mask = torch.ones(query.shape[0], query.shape[2], n_tiles, n_tiles, dtype=torch.bool, device=query.device)
|
||||
else:
|
||||
mask = _build_block_mask(
|
||||
@@ -432,6 +633,10 @@ class MiniMaxH3VSAImpl(AttentionImpl):
|
||||
layer_sparsity,
|
||||
attn_metadata.exempt,
|
||||
)
|
||||
if force_dense is not None:
|
||||
# A scalar bool tensor broadcasts over the block map. This exactly
|
||||
# preserves the eager dense-layer contract without a Python branch.
|
||||
mask = mask | force_dense
|
||||
|
||||
if tile_elems == 64:
|
||||
# Native 64-token path: the block map is already at the kernels'
|
||||
@@ -442,40 +647,87 @@ class MiniMaxH3VSAImpl(AttentionImpl):
|
||||
k_bhsd = key.transpose(1, 2).contiguous()
|
||||
v_bhsd = value.transpose(1, 2).contiguous()
|
||||
|
||||
sm100a_mask = mask
|
||||
sm100a_variable_block_sizes = attn_metadata.variable_block_sizes
|
||||
if has_sm100a_pair:
|
||||
# The synthetic tile is neither a logical query nor key. Its
|
||||
# all-False row yields q2k_num=0, the all-False column keeps it
|
||||
# out of real rows, and vbs=0 masks all of its key slots.
|
||||
sm100a_mask = torch.nn.functional.pad(mask, (0, 1, 0, 1), value=False)
|
||||
sm100a_variable_block_sizes = torch.nn.functional.pad(
|
||||
attn_metadata.variable_block_sizes,
|
||||
(0, 1),
|
||||
value=0,
|
||||
)
|
||||
|
||||
# Opt-in sm_100a CUDA forward (upstream PR #1719 + per-q-tile
|
||||
# q2k_num fix). Forward-only: grad-tracking calls stay on Triton
|
||||
# so autograd keeps the Triton fwd+bwd pairing untouched. The
|
||||
# kernel does return an LSE in Triton's M format, so a future
|
||||
# fwd/bwd pairing is possible, but it is not built here.
|
||||
grad_mode = torch.is_grad_enabled() and (query.requires_grad or key.requires_grad or value.requires_grad)
|
||||
use_sm100a = False
|
||||
if os.environ.get(VSA_SM100A_ENV, "0") == "1":
|
||||
grad_mode = torch.is_grad_enabled() and (query.requires_grad or key.requires_grad
|
||||
or value.requires_grad)
|
||||
reason = _sm100a_unavailable_reason(_sm100a, q_bhsd, attn_metadata.variable_block_sizes, grad_mode)
|
||||
if regional_compiling:
|
||||
if self._regional_compile_sm100a_enabled is None:
|
||||
raise RuntimeError(
|
||||
"VSA-H3 sm_100a routing was not resolved before torch.compile; "
|
||||
"call prepare_for_regional_compile(device) on every MiniMaxH3VSAImpl after loading weights.")
|
||||
# The preparation probe established module/device/kernel
|
||||
# support. Keep only static tensor/geometry facts here; no
|
||||
# env access, device-capability query, or is_supported call may
|
||||
# enter the Dynamo graph.
|
||||
if not (self._regional_compile_sm100a_enabled and not grad_mode and q_bhsd.dtype == torch.bfloat16
|
||||
and q_bhsd.shape[-1] == 128 and sm100a_variable_block_sizes.numel() % 2 == 0):
|
||||
raise RuntimeError(
|
||||
"VSA-H3 regional fullgraph compile requires the prepared sm_100a BF16/head-128 route "
|
||||
"on a supported device; disable inference_torch_compile for this request.")
|
||||
use_sm100a = True
|
||||
elif not compiling and os.environ.get(VSA_SM100A_ENV, "0") == "1":
|
||||
reason = _sm100a_unavailable_reason(_sm100a, q_bhsd, sm100a_variable_block_sizes, grad_mode)
|
||||
if reason is None and map_to_index is None:
|
||||
reason = "fastvideo_kernel.triton_kernels.index (map_to_index) is not importable"
|
||||
if reason is None:
|
||||
use_sm100a = True
|
||||
elif not torch.compiler.is_compiling():
|
||||
else:
|
||||
logger.warning_once(f"{VSA_SM100A_ENV}=1 but falling back to the Triton-64 kernels: {reason}")
|
||||
|
||||
if use_sm100a:
|
||||
# The sm_100a entry is index-native; compact the bool map the
|
||||
# same way the Triton bool entry does internally. Per-row
|
||||
# counts are NON-uniform here (prefix query tiles are dense,
|
||||
# video tiles run prefix+top-k) -- legal for the fixed kernel,
|
||||
# silently wrong on the pre-fix upstream one.
|
||||
q2k_idx, q2k_num = map_to_index(mask)
|
||||
out_bhsd, _ = _sm100a.block_sparse_attn_sm100a(
|
||||
q_bhsd,
|
||||
k_bhsd,
|
||||
v_bhsd,
|
||||
q2k_idx,
|
||||
q2k_num,
|
||||
attn_metadata.variable_block_sizes.to(torch.int32),
|
||||
need_lse=False,
|
||||
)
|
||||
# Regional preparation emits the compile-route receipt before
|
||||
# capture. Logging from this branch would itself break a
|
||||
# ``fullgraph=True`` forward.
|
||||
if not compiling:
|
||||
logger.info_once("MiniMax-H3 VSA tile-64 forward: using the sm100a CUDA block-sparse kernel")
|
||||
if regional_compiling:
|
||||
# The compile-safe wrapper keeps both Triton mask
|
||||
# compaction and the raw pybind launch behind one
|
||||
# fake-backed custom-op boundary.
|
||||
out_bhsd, _ = _sm100a_from_mask(
|
||||
q_bhsd,
|
||||
k_bhsd,
|
||||
v_bhsd,
|
||||
sm100a_mask,
|
||||
sm100a_variable_block_sizes,
|
||||
)
|
||||
else:
|
||||
# Preserve the established eager/index-native route and
|
||||
# compatibility with older kernel wheels. Per-row counts
|
||||
# are non-uniform (prefix queries are dense; video queries
|
||||
# run prefix+top-k), which the fixed kernel supports.
|
||||
q2k_idx, q2k_num = map_to_index(sm100a_mask)
|
||||
out_bhsd, _ = _sm100a.block_sparse_attn_sm100a(
|
||||
q_bhsd,
|
||||
k_bhsd,
|
||||
v_bhsd,
|
||||
q2k_idx,
|
||||
q2k_num,
|
||||
sm100a_variable_block_sizes.to(torch.int32),
|
||||
need_lse=False,
|
||||
)
|
||||
else:
|
||||
if has_sm100a_pair:
|
||||
q_bhsd = q_bhsd[:, :, :logical_seq_len].contiguous()
|
||||
k_bhsd = k_bhsd[:, :, :logical_seq_len].contiguous()
|
||||
v_bhsd = v_bhsd[:, :, :logical_seq_len].contiguous()
|
||||
out_bhsd, _ = block_sparse_attn_64_bhsd(
|
||||
q_bhsd,
|
||||
k_bhsd,
|
||||
@@ -483,25 +735,32 @@ class MiniMaxH3VSAImpl(AttentionImpl):
|
||||
mask,
|
||||
attn_metadata.variable_block_sizes,
|
||||
)
|
||||
if has_sm100a_pair and use_sm100a:
|
||||
out_bhsd = out_bhsd[:, :, :logical_seq_len]
|
||||
out = out_bhsd.transpose(1, 2).contiguous()
|
||||
else:
|
||||
out, _ = block_sparse_attn_256_bshd(query, key, value, mask, attn_metadata.variable_block_sizes)
|
||||
out, _ = block_sparse_attn_256_bshd(
|
||||
logical_query,
|
||||
logical_key,
|
||||
logical_value,
|
||||
mask,
|
||||
attn_metadata.variable_block_sizes,
|
||||
)
|
||||
|
||||
if gate_compress is not None:
|
||||
if logical_gate is not None:
|
||||
# Wan-style compression branch: dense attention over pooled tiles,
|
||||
# broadcast to each tile's rows, scaled by the learned gate
|
||||
# (zero-initialized for H3 => branch contributes nothing until
|
||||
# finetuned; the model layer skips it entirely for all-zero gates).
|
||||
v_pooled = _pool_tiles(value, attn_metadata.variable_block_sizes, tile_elems)
|
||||
v_pooled = _pool_tiles(logical_value, attn_metadata.variable_block_sizes, tile_elems)
|
||||
out_c = torch.matmul(torch.softmax(scores, dim=-1), v_pooled) # [B, H, n_tiles, D]
|
||||
out_c = out_c.permute(0, 2, 1, 3).to(out.dtype) # [B, n_tiles, H, D]
|
||||
batch, seq_len, heads, dim = out.shape
|
||||
n_tiles = attn_metadata.variable_block_sizes.numel()
|
||||
# Out-of-place: on the CuTe backend ``out`` is the tensor FA4's
|
||||
# autograd node saved for its backward, so an in-place add here
|
||||
# bumps its version counter and backward dies with "one of the
|
||||
# variables needed for gradient computation has been modified".
|
||||
out_tiled = out.view(batch, n_tiles, tile_elems, heads, dim)
|
||||
gate_tiled = gate_compress.view(batch, n_tiles, tile_elems, heads, dim)
|
||||
gate_tiled = logical_gate.view(batch, n_tiles, tile_elems, heads, dim)
|
||||
out = (out_tiled + out_c.unsqueeze(2) * gate_tiled).view(batch, seq_len, heads, dim)
|
||||
return out
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
import os
|
||||
from functools import wraps
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
@@ -20,7 +21,7 @@ def _attention_compile_disabled() -> bool:
|
||||
|
||||
Defaults to ``True`` (the historical behavior: attention runs eager via
|
||||
``torch.compiler.disable``). Set ``FASTVIDEO_DISABLE_ATTENTION_COMPILE=0``
|
||||
to let attention be traced/compiled into the surrounding graph.
|
||||
to let attention instances constructed under that environment be traced.
|
||||
"""
|
||||
val = os.environ.get("FASTVIDEO_DISABLE_ATTENTION_COMPILE")
|
||||
if val is None:
|
||||
@@ -28,11 +29,32 @@ def _attention_compile_disabled() -> bool:
|
||||
return val.strip().lower() not in ("0", "false", "no", "off", "")
|
||||
|
||||
|
||||
def _attention_compile_explicitly_disabled() -> bool:
|
||||
"""Whether the environment explicitly requests the eager boundary.
|
||||
|
||||
Regional compile can override the historical default for one loaded
|
||||
transformer, but it must still honor an explicit debugging escape hatch.
|
||||
"""
|
||||
return "FASTVIDEO_DISABLE_ATTENTION_COMPILE" in os.environ and _attention_compile_disabled()
|
||||
|
||||
|
||||
def _maybe_compiler_disable(fn):
|
||||
"""Apply ``torch.compiler.disable`` unless disabled via env var."""
|
||||
if _attention_compile_disabled():
|
||||
return torch.compiler.disable(fn)
|
||||
return fn
|
||||
"""Defer the eager/traceable choice to each attention instance.
|
||||
|
||||
A class-definition-time choice makes a process-wide default the only
|
||||
option. The deferred wrapper keeps ordinary instances on the historical
|
||||
eager boundary while allowing the regional loader to opt in only the
|
||||
attention modules owned by the transformer it is compiling.
|
||||
"""
|
||||
disabled_fn = torch.compiler.disable(fn)
|
||||
|
||||
@wraps(fn)
|
||||
def _dispatch(self, *args, **kwargs):
|
||||
if self._compile_forward_enabled:
|
||||
return fn(self, *args, **kwargs)
|
||||
return disabled_fn(self, *args, **kwargs)
|
||||
|
||||
return _dispatch
|
||||
|
||||
|
||||
class DistributedAttention(nn.Module):
|
||||
@@ -77,6 +99,13 @@ class DistributedAttention(nn.Module):
|
||||
self.num_kv_heads = num_kv_heads
|
||||
self.backend = backend_name_to_enum(attn_backend.get_name())
|
||||
self.dtype = dtype
|
||||
# Preserve the historical compiler-disabled default. The regional
|
||||
# inference loader may enable this one instance after validating the
|
||||
# transformer's resolved backend; no process-global default changes.
|
||||
self._compile_forward_enabled = not _attention_compile_disabled()
|
||||
|
||||
def _set_compile_forward_enabled(self, enabled: bool) -> None:
|
||||
self._compile_forward_enabled = enabled
|
||||
|
||||
@_maybe_compiler_disable
|
||||
def forward(
|
||||
|
||||
@@ -5,6 +5,7 @@ import torch
|
||||
from torch.distributed import ProcessGroup
|
||||
|
||||
from fastvideo.distributed.device_communicators.base_device_communicator import (DeviceCommunicatorBase)
|
||||
from fastvideo.distributed.device_communicators.ulysses_a2a import maybe_create_helper
|
||||
|
||||
|
||||
class CudaCommunicator(DeviceCommunicatorBase):
|
||||
@@ -25,6 +26,10 @@ class CudaCommunicator(DeviceCommunicatorBase):
|
||||
device=self.device,
|
||||
)
|
||||
|
||||
# Capability is agreed once; the persistent window arms on first use.
|
||||
self.ulysses_a2a = maybe_create_helper(self.cpu_group, self.device_group, self.world_size, self.device,
|
||||
self.pynccl_comm)
|
||||
|
||||
def all_reduce(self, input_, op: torch.distributed.ReduceOp | None = None):
|
||||
pynccl_comm = self.pynccl_comm
|
||||
assert pynccl_comm is not None
|
||||
@@ -38,6 +43,14 @@ class CudaCommunicator(DeviceCommunicatorBase):
|
||||
torch.distributed.all_reduce(out, group=self.device_group, op=op)
|
||||
return out
|
||||
|
||||
def all_to_all_4D(self, input_: torch.Tensor, scatter_dim: int = 2, gather_dim: int = 1) -> torch.Tensor:
|
||||
"""All-to-all over the sequence parallel group, fused when available."""
|
||||
if self.ulysses_a2a is not None:
|
||||
output = self.ulysses_a2a.try_all_to_all_4D(input_, scatter_dim, gather_dim)
|
||||
if output is not None:
|
||||
return output
|
||||
return super().all_to_all_4D(input_, scatter_dim, gather_dim)
|
||||
|
||||
def send(self, tensor: torch.Tensor, dst: int | None = None) -> None:
|
||||
"""Sends a tensor to the destination rank in a non-blocking way"""
|
||||
"""NOTE: `dst` is the local rank of the destination rank."""
|
||||
@@ -65,5 +78,9 @@ class CudaCommunicator(DeviceCommunicatorBase):
|
||||
return tensor
|
||||
|
||||
def destroy(self) -> None:
|
||||
if self.ulysses_a2a is not None:
|
||||
# The helper still needs its NCCL communicator during teardown.
|
||||
self.ulysses_a2a.close()
|
||||
self.ulysses_a2a = None
|
||||
if self.pynccl_comm is not None:
|
||||
self.pynccl_comm = None
|
||||
|
||||
@@ -0,0 +1,413 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""Fused NVLink all-to-all for Ulysses sequence parallelism.
|
||||
|
||||
Drop-in replacement for DistributedAutograd.AllToAll4D when the group is a
|
||||
load-store accessible NVLink mesh: same layout, byte-identical results, fewer
|
||||
passes over local memory. Anything else falls back to the NCCL path.
|
||||
"""
|
||||
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
from torch.distributed import ProcessGroup
|
||||
|
||||
from fastvideo import envs
|
||||
from fastvideo.logger import init_logger
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
# The kernel is template-specialized on the world size, so only these dispatch.
|
||||
SUPPORTED_WORLD_SIZES = (2, 4, 6, 8)
|
||||
|
||||
_DTYPE_CODES = {
|
||||
torch.float16: 1,
|
||||
torch.bfloat16: 2,
|
||||
torch.float32: 3,
|
||||
}
|
||||
|
||||
# Bound persistent registered memory per rank. Larger operands use NCCL instead
|
||||
# of growing the window without limit.
|
||||
MAX_WINDOW_BYTES = 1024**3
|
||||
|
||||
# (scatter_dim, gather_dim) -> kernel mode.
|
||||
# 0: [B, S_local, H, D] -> [B, S_global, H_local, D]
|
||||
# 1: [B, S_global, H_local, D] -> [B, S_local, H, D]
|
||||
_MODE_FROM_DIMS = {(2, 1): 0, (1, 2): 1}
|
||||
|
||||
|
||||
def is_enabled() -> bool:
|
||||
"""Whether the fused path is opted in via FASTVIDEO_ULYSSES_A2A."""
|
||||
return envs.FASTVIDEO_ULYSSES_A2A == "auto"
|
||||
|
||||
|
||||
class _FusedUlyssesA2A(torch.autograd.Function):
|
||||
"""Differentiable fused all-to-all.
|
||||
|
||||
The two directions are exact inverses, and Ulysses redistributes activations
|
||||
rather than reducing them, so backward is the opposite mode with no scaling.
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def forward(ctx, helper: "UlyssesA2AHelper", x: torch.Tensor, mode: int) -> torch.Tensor: # type: ignore[override]
|
||||
ctx.helper = helper
|
||||
ctx.mode = mode
|
||||
return helper.run_armed(x, mode)
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, grad_output: torch.Tensor): # type: ignore[override]
|
||||
# Same numel and dtype as the forward output, so the window is already
|
||||
# sized for it; only contiguity needs restoring.
|
||||
grad_input = ctx.helper.run_armed(grad_output.contiguous(), 1 - ctx.mode)
|
||||
return None, grad_input, None
|
||||
|
||||
|
||||
class UlyssesA2AHelper:
|
||||
"""Owns the fused all-to-all context for one sequence-parallel group.
|
||||
|
||||
Group capability is agreed during construction; the NCCL window is
|
||||
registered on first use, once an operand size is known.
|
||||
"""
|
||||
|
||||
def __init__(self, cpu_group: ProcessGroup, device_group: ProcessGroup, world_size: int, device: torch.device,
|
||||
pynccl_comm):
|
||||
self.cpu_group = cpu_group
|
||||
self.device_group = device_group
|
||||
self.world_size = world_size
|
||||
self.device = device
|
||||
self.pynccl_comm = pynccl_comm
|
||||
|
||||
self._handle: int | None = None
|
||||
self._nbytes = 0
|
||||
self._disabled_reason: str | None = None
|
||||
|
||||
if world_size not in SUPPORTED_WORLD_SIZES:
|
||||
self._disabled_reason = (f"world size {world_size} is not one of "
|
||||
f"{SUPPORTED_WORLD_SIZES}")
|
||||
|
||||
# -- lifecycle -----------------------------------------------------------
|
||||
|
||||
def _disable(self, reason: str) -> None:
|
||||
if self._disabled_reason is None:
|
||||
self._disabled_reason = reason
|
||||
logger.info("Ulysses fused all-to-all disabled: %s", reason)
|
||||
|
||||
def _comm_ptr(self) -> int:
|
||||
comm = self.pynccl_comm.comm
|
||||
return int(getattr(comm, "value", comm))
|
||||
|
||||
def _can_attempt(self) -> tuple[bool, str]:
|
||||
"""Whether this rank could use the fused path, without allocating anything."""
|
||||
try:
|
||||
from fastvideo_kernel import comm_ops
|
||||
if not comm_ops.is_available():
|
||||
return False, "fastvideo-kernel was built without the Ulysses a2a kernel"
|
||||
if not comm_ops.lsa_covers_group(self._comm_ptr(), self.world_size):
|
||||
return False, "the group is not a load-store-accessible (NVLink) mesh"
|
||||
except Exception as e: # noqa: BLE001
|
||||
return False, f"backend unavailable ({type(e).__name__}: {e})"
|
||||
return True, ""
|
||||
|
||||
def _agree(self, ok: bool) -> bool:
|
||||
"""Reduce a local yes/no to a group-wide verdict: True only if all agree."""
|
||||
vote = torch.tensor([1 if ok else 0], dtype=torch.int32)
|
||||
dist.all_reduce(vote, op=dist.ReduceOp.MIN, group=self.cpu_group)
|
||||
return bool(vote.item())
|
||||
|
||||
def _allocate(self, nbytes: int) -> int:
|
||||
"""Allocate locally; split out so allocation-failure tests can inject."""
|
||||
from fastvideo_kernel import comm_ops
|
||||
|
||||
device_index = self.device.index
|
||||
if device_index is None:
|
||||
device_index = torch.cuda.current_device()
|
||||
return comm_ops.allocate(nbytes, self.pynccl_comm.rank, self.world_size, device_index)
|
||||
|
||||
def _register_window(self, handle: int) -> None:
|
||||
"""Register the user window collectively."""
|
||||
from fastvideo_kernel import comm_ops
|
||||
|
||||
comm_ops.register_window(handle, self._comm_ptr())
|
||||
|
||||
def _create_dev_comm(self, handle: int) -> None:
|
||||
"""Create the NCCL device communicator collectively."""
|
||||
from fastvideo_kernel import comm_ops
|
||||
|
||||
comm_ops.create_dev_comm(handle)
|
||||
|
||||
def _dispose(self, handle: int, *, synchronize: bool) -> None:
|
||||
from fastvideo_kernel import comm_ops
|
||||
|
||||
if synchronize:
|
||||
# Kernel launches and copy-out are asynchronous. Do not deregister a
|
||||
# window that a prior call on this device is still accessing.
|
||||
torch.cuda.synchronize(self.device)
|
||||
comm_ops.dispose(handle)
|
||||
|
||||
def _dispose_after_failure(self, handle: int | None) -> bool:
|
||||
"""Best-effort group cleanup after a setup phase failed.
|
||||
|
||||
Every rank votes after attempting cleanup, including a rank that never
|
||||
obtained a local allocation. This keeps the helper permanently disabled
|
||||
if teardown was not unanimous instead of re-entering with split state.
|
||||
"""
|
||||
cleanup_ok = True
|
||||
if handle is not None:
|
||||
try:
|
||||
self._dispose(handle, synchronize=False)
|
||||
except Exception: # noqa: BLE001 - converted to a group verdict below
|
||||
cleanup_ok = False
|
||||
logger.warning("Ulysses partial-context cleanup failed", exc_info=True)
|
||||
return self._agree(cleanup_ok)
|
||||
|
||||
def _call_signature(self, x: torch.Tensor, scatter_dim: int, gather_dim: int) -> tuple[tuple[int, ...], str]:
|
||||
"""Return a rank-comparable call contract and any local decline reason."""
|
||||
mode = _MODE_FROM_DIMS.get((scatter_dim, gather_dim))
|
||||
dtype_code = _DTYPE_CODES.get(x.dtype, 0)
|
||||
shape = tuple(int(dim) for dim in x.shape) if x.dim() == 4 else (0, 0, 0, 0)
|
||||
status = 1
|
||||
reason = ""
|
||||
|
||||
if self._disabled_reason is not None:
|
||||
status, reason = -1, self._disabled_reason
|
||||
elif not is_enabled():
|
||||
status, reason = 0, "FASTVIDEO_ULYSSES_A2A is not auto"
|
||||
elif x.is_cuda and torch.cuda.is_current_stream_capturing():
|
||||
status, reason = 0, "the current CUDA stream is being captured"
|
||||
elif mode is None:
|
||||
status, reason = 0, "unsupported scatter/gather dimensions"
|
||||
elif x.dim() != 4:
|
||||
status, reason = 0, "input is not 4-D"
|
||||
elif dtype_code == 0:
|
||||
status, reason = 0, f"unsupported dtype {x.dtype}"
|
||||
elif not x.is_cuda or x.device != self.device:
|
||||
status, reason = 0, f"input device {x.device} does not match {self.device}"
|
||||
elif not x.is_contiguous():
|
||||
status, reason = 0, "input is not contiguous"
|
||||
elif mode == 0 and shape[2] % self.world_size != 0:
|
||||
status, reason = 0, "head count is not divisible by the group"
|
||||
elif mode == 1 and shape[1] % self.world_size != 0:
|
||||
status, reason = 0, "sequence length is not divisible by the group"
|
||||
|
||||
nbytes = int(x.numel() * x.element_size())
|
||||
if status == 1 and nbytes == 0:
|
||||
status, reason = 0, "input is empty"
|
||||
elif status == 1 and nbytes > MAX_WINDOW_BYTES:
|
||||
status, reason = 0, f"operand exceeds the {MAX_WINDOW_BYTES}-byte window cap"
|
||||
|
||||
# status, armed, mode, dtype, B, S, H, D, bytes, capacity. Comparing the
|
||||
# whole vector prevents equal-size but differently-shaped ranks from
|
||||
# entering the fused kernel with incompatible address math. CUDA device
|
||||
# ordinals are deliberately absent: rank-local ordinals normally differ.
|
||||
signature = (status, int(self._handle is not None), -1 if mode is None else mode, dtype_code, *shape, nbytes,
|
||||
self._nbytes)
|
||||
return signature, reason
|
||||
|
||||
def _agree_call(self, signature: tuple[int, ...]) -> tuple[bool, bool, bool]:
|
||||
"""Agree on eligibility and the complete call signature across ranks.
|
||||
|
||||
Returns ``(use_fused, permanently_unavailable, lifecycle_consistent)``. This control
|
||||
collective is intentionally eager-only; compiled regions use the NCCL
|
||||
implementation before reaching here.
|
||||
"""
|
||||
# Host-side Gloo control keeps this agreement outside CUDA graph capture
|
||||
# and avoids inserting a second NCCL collective ahead of the data path.
|
||||
local = torch.tensor(signature, dtype=torch.int64)
|
||||
gathered = torch.empty(self.world_size * local.numel(), dtype=local.dtype)
|
||||
dist.all_gather_into_tensor(gathered, local, group=self.cpu_group)
|
||||
contracts = gathered.view(self.world_size, local.numel())
|
||||
identical = bool(torch.all(contracts == contracts[0]).item())
|
||||
statuses = contracts[:, 0]
|
||||
use_fused = identical and bool(torch.all(statuses == 1).item())
|
||||
permanently_unavailable = bool(torch.any(statuses < 0).item())
|
||||
lifecycle_consistent = (bool(torch.all(contracts[:, 1] == contracts[0, 1]).item())
|
||||
and bool(torch.all(contracts[:, -1] == contracts[0, -1]).item()))
|
||||
return use_fused, permanently_unavailable, lifecycle_consistent
|
||||
|
||||
def _build(self, nbytes: int) -> bool:
|
||||
"""Collectively register the window. Returns True if it is armed."""
|
||||
handle: int | None = None
|
||||
allocation_reason = ""
|
||||
try:
|
||||
handle = self._allocate(nbytes)
|
||||
except Exception as e: # noqa: BLE001 - converted to a group verdict below
|
||||
allocation_reason = f"window allocation failed ({type(e).__name__}: {e})"
|
||||
|
||||
# Allocation is local, so vote before any rank enters registration.
|
||||
if not self._agree(handle is not None):
|
||||
cleanup_ok = self._dispose_after_failure(handle)
|
||||
reason = allocation_reason or "a peer rank could not allocate the window"
|
||||
if not cleanup_ok:
|
||||
reason += "; partial-context cleanup failed on a peer"
|
||||
self._disable(reason)
|
||||
return False
|
||||
|
||||
assert handle is not None
|
||||
window_registered = False
|
||||
registration_reason = ""
|
||||
try:
|
||||
self._register_window(handle)
|
||||
window_registered = True
|
||||
except Exception as e: # noqa: BLE001 - converted to a group verdict below
|
||||
registration_reason = f"window registration failed ({type(e).__name__}: {e})"
|
||||
|
||||
if not self._agree(window_registered):
|
||||
cleanup_ok = self._dispose_after_failure(handle)
|
||||
reason = registration_reason or "a peer rank could not register the window"
|
||||
if not cleanup_ok:
|
||||
reason += "; partial-context cleanup failed on a peer"
|
||||
self._disable(reason)
|
||||
return False
|
||||
|
||||
dev_comm_created = False
|
||||
creation_reason = ""
|
||||
try:
|
||||
self._create_dev_comm(handle)
|
||||
dev_comm_created = True
|
||||
except Exception as e: # noqa: BLE001 - converted to a group verdict below
|
||||
creation_reason = f"device communicator creation failed ({type(e).__name__}: {e})"
|
||||
|
||||
if not self._agree(dev_comm_created):
|
||||
cleanup_ok = self._dispose_after_failure(handle)
|
||||
reason = creation_reason or "a peer rank could not create the device communicator"
|
||||
if not cleanup_ok:
|
||||
reason += "; partial-context cleanup failed on a peer"
|
||||
self._disable(reason)
|
||||
return False
|
||||
|
||||
self._handle = handle
|
||||
self._nbytes = nbytes
|
||||
logger.info("Ulysses fused all-to-all armed: world_size=%d window=%.0f MiB", self.world_size, nbytes / 2**20)
|
||||
return True
|
||||
|
||||
def close(self) -> bool:
|
||||
"""Collectively destroy the device communicator and its window.
|
||||
|
||||
Returns whether all ranks completed teardown. An armed/unarmed split
|
||||
cannot safely enter NCCL window deregistration, so that exceptional
|
||||
state is leaked until process exit and permanently disabled instead of
|
||||
risking a distributed deadlock.
|
||||
"""
|
||||
handle = self._handle
|
||||
all_armed = self._agree(handle is not None)
|
||||
all_unarmed = self._agree(handle is None)
|
||||
if all_unarmed:
|
||||
self._nbytes = 0
|
||||
return True
|
||||
if not all_armed:
|
||||
self._handle = None
|
||||
self._nbytes = 0
|
||||
self._disable("ranks disagreed on whether a fused window was armed during teardown")
|
||||
return False
|
||||
|
||||
assert handle is not None
|
||||
synchronize_ok = True
|
||||
try:
|
||||
torch.cuda.synchronize(self.device)
|
||||
except Exception: # noqa: BLE001 - converted to a group verdict below
|
||||
synchronize_ok = False
|
||||
logger.warning("Ulysses pre-teardown synchronization failed", exc_info=True)
|
||||
if not self._agree(synchronize_ok):
|
||||
self._disable("a peer rank could not synchronize before fused-window teardown")
|
||||
return False
|
||||
|
||||
dispose_ok = True
|
||||
try:
|
||||
self._dispose(handle, synchronize=False)
|
||||
except Exception: # noqa: BLE001 - teardown must not mask a real error
|
||||
dispose_ok = False
|
||||
logger.warning("Ulysses window deregistration failed", exc_info=True)
|
||||
|
||||
group_ok = self._agree(dispose_ok)
|
||||
# The native disposer consumes the handle even when a cleanup call
|
||||
# reports an error, so never retry a potentially dangling pointer.
|
||||
self._handle = None
|
||||
self._nbytes = 0
|
||||
if not group_ok:
|
||||
self._disable("fused-window teardown failed on a peer rank")
|
||||
return group_ok
|
||||
|
||||
# -- collective ----------------------------------------------------------
|
||||
|
||||
def run_armed(self, x: torch.Tensor, mode: int) -> torch.Tensor:
|
||||
"""Run one collective on an already-armed context."""
|
||||
assert self._handle is not None, "run_armed called on an unarmed helper"
|
||||
from fastvideo_kernel import comm_ops
|
||||
|
||||
w = self.world_size
|
||||
if mode == 0:
|
||||
B, S_local, H, D = x.shape
|
||||
out = torch.empty(B, S_local * w, H // w, D, dtype=x.dtype, device=x.device)
|
||||
else:
|
||||
B, S_global, H_local, D = x.shape
|
||||
S_local, H = S_global // w, H_local * w
|
||||
out = torch.empty(B, S_local, H, D, dtype=x.dtype, device=x.device)
|
||||
comm_ops.all_to_all(self._handle, x, out, B, S_local, H, D, mode)
|
||||
return out
|
||||
|
||||
def try_all_to_all_4D(self, x: torch.Tensor, scatter_dim: int, gather_dim: int) -> torch.Tensor | None:
|
||||
"""Fused collective, or None to let the caller use the NCCL path."""
|
||||
if self._disabled_reason is not None:
|
||||
return None
|
||||
|
||||
# Python lifecycle checks, votes, and pybind calls are not valid inside
|
||||
# a fullgraph region. The inherited NCCL path is compiler-visible, so
|
||||
# regional compile stays fullgraph by declining before any tensor read.
|
||||
if torch.compiler.is_compiling():
|
||||
return None
|
||||
|
||||
signature, reason = self._call_signature(x, scatter_dim, gather_dim)
|
||||
use_fused, permanently_unavailable, lifecycle_consistent = self._agree_call(signature)
|
||||
if not use_fused:
|
||||
if not lifecycle_consistent:
|
||||
self.close()
|
||||
self._disable("ranks disagreed on the fused-window lifecycle")
|
||||
if permanently_unavailable:
|
||||
self._disable(reason or "a peer rank cannot use the fused path")
|
||||
return None
|
||||
|
||||
mode = signature[2]
|
||||
nbytes = signature[-2]
|
||||
if self._handle is None:
|
||||
if not self._build(nbytes):
|
||||
return None
|
||||
elif nbytes > self._nbytes:
|
||||
logger.info("Ulysses window grow: %d -> %d bytes", self._nbytes, nbytes)
|
||||
if not self.close():
|
||||
return None
|
||||
if not self._build(nbytes):
|
||||
return None
|
||||
|
||||
return _FusedUlyssesA2A.apply(self, x, mode)
|
||||
|
||||
|
||||
def maybe_create_helper(cpu_group: ProcessGroup | None, device_group: ProcessGroup | None, world_size: int,
|
||||
device: torch.device | None, pynccl_comm) -> UlyssesA2AHelper | None:
|
||||
"""Collectively create a helper only when every rank can use it."""
|
||||
if (world_size <= 1 or cpu_group is None or device_group is None or device is None or device.type != "cuda"):
|
||||
return None
|
||||
if not dist.is_initialized():
|
||||
return None
|
||||
|
||||
helper = None
|
||||
reason = ""
|
||||
if not is_enabled():
|
||||
reason = "FASTVIDEO_ULYSSES_A2A is not auto"
|
||||
elif world_size not in SUPPORTED_WORLD_SIZES:
|
||||
reason = f"world size {world_size} is not one of {SUPPORTED_WORLD_SIZES}"
|
||||
elif pynccl_comm is None or pynccl_comm.disabled:
|
||||
reason = "the group has no usable PyNccl communicator"
|
||||
else:
|
||||
try:
|
||||
candidate = UlyssesA2AHelper(cpu_group, device_group, world_size, device, pynccl_comm)
|
||||
can_attempt, reason = candidate._can_attempt()
|
||||
if can_attempt:
|
||||
helper = candidate
|
||||
except Exception as e: # noqa: BLE001 - converted to a group verdict below
|
||||
reason = f"helper construction failed ({type(e).__name__}: {e})"
|
||||
|
||||
vote = torch.tensor([int(helper is not None)], dtype=torch.int32)
|
||||
dist.all_reduce(vote, op=dist.ReduceOp.MIN, group=cpu_group)
|
||||
if not bool(vote.item()):
|
||||
if dist.get_rank(cpu_group) == 0:
|
||||
logger.info("Ulysses fused all-to-all unavailable: %s", reason or "a peer rank declined")
|
||||
return None
|
||||
return helper
|
||||
@@ -635,14 +635,16 @@ class GroupCoordinator:
|
||||
return self.device_communicator.recv(size, dtype, src)
|
||||
|
||||
def destroy(self) -> None:
|
||||
# First: communicator teardown can be collective, so it needs the
|
||||
# process groups alive.
|
||||
if self.device_communicator is not None:
|
||||
self.device_communicator.destroy()
|
||||
if self.device_group is not None:
|
||||
torch.distributed.destroy_process_group(self.device_group)
|
||||
self.device_group = None
|
||||
if self.cpu_group is not None:
|
||||
torch.distributed.destroy_process_group(self.cpu_group)
|
||||
self.cpu_group = None
|
||||
if self.device_communicator is not None:
|
||||
self.device_communicator.destroy()
|
||||
if self.mq_broadcaster is not None:
|
||||
self.mq_broadcaster = None
|
||||
|
||||
|
||||
@@ -99,6 +99,8 @@ _FROM_PRETRAINED_CONVENIENCE_KWARGS = frozenset({
|
||||
"pin_cpu_memory",
|
||||
"enable_torch_compile",
|
||||
"torch_compile_kwargs",
|
||||
"lora_path",
|
||||
"lora_strength",
|
||||
"output_type",
|
||||
"nvfp4_fa4",
|
||||
})
|
||||
|
||||
@@ -21,10 +21,13 @@ if TYPE_CHECKING:
|
||||
FASTVIDEO_TRACE_FUNCTION: int = 0
|
||||
FASTVIDEO_ATTENTION_BACKEND: str | None = None
|
||||
FASTVIDEO_FA4: bool = False
|
||||
FASTVIDEO_MINIMAX_H3_FA4_PACKED_VARLEN: bool = False
|
||||
FASTVIDEO_INFERENCE_TORCH_COMPILE: bool = False
|
||||
FASTVIDEO_MINIMAX_H3_FUSIONS: str = ""
|
||||
FASTVIDEO_VAE_PARALLEL_DECODE: bool = False
|
||||
FASTVIDEO_VAE_PARALLEL_ENCODE: bool = False
|
||||
FASTVIDEO_VAE_PARALLEL_DECODE_STRATEGY: str | None = None
|
||||
FASTVIDEO_ULYSSES_A2A: str = "off"
|
||||
FASTVIDEO_WORKER_MULTIPROC_METHOD: str = "spawn"
|
||||
FASTVIDEO_TARGET_DEVICE: str = "cuda"
|
||||
MAX_JOBS: str | None = None
|
||||
@@ -222,6 +225,24 @@ environment_variables: dict[str, Callable[[], Any]] = {
|
||||
"FASTVIDEO_FA4":
|
||||
lambda: os.getenv("FASTVIDEO_FA4", "0") != "0",
|
||||
|
||||
# Use FA4's packed-varlen entry point for the long, single-document
|
||||
# MiniMax-H3 dense DiT self-attention path. This changes floating-point
|
||||
# reduction order relative to the fixed-length entry point, so it remains
|
||||
# an explicit inference-only speed/quality opt-in.
|
||||
"FASTVIDEO_MINIMAX_H3_FA4_PACKED_VARLEN":
|
||||
lambda: os.getenv("FASTVIDEO_MINIMAX_H3_FA4_PACKED_VARLEN", "0") != "0",
|
||||
|
||||
# If set (=1), enable regional (per-transformer-block) fullgraph
|
||||
# torch.compile for the DiT at inference — the inference-side counterpart
|
||||
# of the training regional-compile port of hao-ai-lab/FastVideo#1718.
|
||||
# Equivalent to FastVideoArgs.inference_torch_compile=True (e.g. via
|
||||
# PipelineSelection.experimental={"inference_torch_compile": True}). VSA
|
||||
# and other non-fullgraph-traceable attention backends degrade to eager
|
||||
# with one warning; see _regional_compile_unsupported_reason in
|
||||
# fastvideo/models/loader/fsdp_load.py.
|
||||
"FASTVIDEO_INFERENCE_TORCH_COMPILE":
|
||||
lambda: os.getenv("FASTVIDEO_INFERENCE_TORCH_COMPILE", "0") != "0",
|
||||
|
||||
# If set (=1), MiniMax-H3 VAE decode (and, with the ENCODE variant,
|
||||
# reference-video encode) round-robins its temporal chunks across the
|
||||
# sequence-parallel ranks instead of running serially on the output rank.
|
||||
@@ -242,6 +263,13 @@ environment_variables: dict[str, Callable[[], Any]] = {
|
||||
"FASTVIDEO_MINIMAX_H3_FUSIONS":
|
||||
lambda: os.getenv("FASTVIDEO_MINIMAX_H3_FUSIONS", ""),
|
||||
|
||||
# Sequence-parallel all-to-all backend.
|
||||
# - "off" (default): the NCCL path in DistributedAutograd.AllToAll4D
|
||||
# - "auto": fused NVLink kernel when the group is a load-store accessible
|
||||
# mesh of 2/4/6/8 ranks in eager execution, else the NCCL path.
|
||||
"FASTVIDEO_ULYSSES_A2A":
|
||||
lambda: os.getenv("FASTVIDEO_ULYSSES_A2A", "off").strip().lower(),
|
||||
|
||||
# Use dedicated multiprocess context for workers.
|
||||
"FASTVIDEO_WORKER_MULTIPROC_METHOD":
|
||||
lambda: os.getenv("FASTVIDEO_WORKER_MULTIPROC_METHOD", "spawn"),
|
||||
|
||||
+122
-14
@@ -4,6 +4,7 @@
|
||||
import argparse
|
||||
import dataclasses
|
||||
import json
|
||||
import math
|
||||
from contextlib import contextmanager
|
||||
from dataclasses import field
|
||||
from enum import Enum
|
||||
@@ -25,6 +26,17 @@ else:
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
# Offload flags that trade device memory for host memory. All of them are a loss
|
||||
# on a device where the two are the same physical pool. Keeping the policy
|
||||
# centralized lets every loader and stage share one worker-local decision.
|
||||
UNIFIED_MEMORY_OFFLOAD_FLAGS = (
|
||||
"dit_layerwise_offload",
|
||||
"dit_cpu_offload",
|
||||
"text_encoder_cpu_offload",
|
||||
"image_encoder_cpu_offload",
|
||||
"vae_cpu_offload",
|
||||
)
|
||||
|
||||
|
||||
class ExecutionMode(str, Enum):
|
||||
"""
|
||||
@@ -119,6 +131,7 @@ class FastVideoArgs:
|
||||
# (Wenxuan) prefer to keep it here instead of in pipeline config to not make it complicated.
|
||||
lora_path: str | None = None
|
||||
lora_nickname: str = "default" # for swapping adapters in the pipeline
|
||||
lora_strength: float = 1.0
|
||||
# can restrict layers to adapt, e.g. ["q_proj"]
|
||||
# Will adapt only q, k, v, o by default.
|
||||
lora_target_modules: list[str] | None = None
|
||||
@@ -177,6 +190,19 @@ class FastVideoArgs:
|
||||
torch_compile_kwargs_text_encoder: dict[str, Any] = field(default_factory=dict)
|
||||
torch_compile_kwargs_vae: dict[str, Any] = field(default_factory=dict)
|
||||
torch_compile_kwargs_audio_vae: dict[str, Any] = field(default_factory=dict)
|
||||
# Regional (per-transformer-block) fullgraph torch.compile of the DiT at
|
||||
# inference — the inference-side counterpart of the training regional
|
||||
# compile ported from hao-ai-lab/FastVideo#1718. Applied by the loader
|
||||
# right after the transformer loads, with fullgraph=True and inductor
|
||||
# options {emulate_precision_casts: True} injected (no user kwargs
|
||||
# needed). MiniMax-H3 VSA is supported only by its compile-safe sm_100a
|
||||
# tile-64 inference route; other VSA routes degrade the transformer to
|
||||
# eager with one warning. Dense FA2/FA3/FA4 inference uses compile-visible
|
||||
# custom-op boundaries. Opt-in via FASTVIDEO_INFERENCE_TORCH_COMPILE=1 (folded in
|
||||
# __post_init__) or PipelineSelection.experimental
|
||||
# {"inference_torch_compile": true}. Distinct from ``enable_torch_compile``,
|
||||
# which keeps the pipeline-level compile semantics.
|
||||
inference_torch_compile: bool = False
|
||||
|
||||
disable_autocast: bool = False
|
||||
|
||||
@@ -274,6 +300,8 @@ class FastVideoArgs:
|
||||
return not self.inference_mode
|
||||
|
||||
def __post_init__(self):
|
||||
if not math.isfinite(self.lora_strength):
|
||||
raise ValueError(f"lora_strength must be finite, got {self.lora_strength}")
|
||||
if self.moba_config_path:
|
||||
try:
|
||||
with open(self.moba_config_path) as f:
|
||||
@@ -285,6 +313,13 @@ class FastVideoArgs:
|
||||
self._apply_ltx2_vae_overrides()
|
||||
self._resolve_refine_args()
|
||||
self._apply_transformer_quant()
|
||||
if not self.inference_torch_compile:
|
||||
# Parse-once adapter (same pattern as attention_backend below): the
|
||||
# environment variable is an input read once here, so the loader
|
||||
# only ever consults the typed field.
|
||||
import fastvideo.envs as envs
|
||||
if envs.FASTVIDEO_INFERENCE_TORCH_COMPILE:
|
||||
self.inference_torch_compile = True
|
||||
if self.attention_backend is not None:
|
||||
# Fail fast on typos instead of silently auto-selecting later.
|
||||
from fastvideo.attention.selector import coerce_attn_backend
|
||||
@@ -571,6 +606,12 @@ class FastVideoArgs:
|
||||
default=FastVideoArgs.lora_nickname,
|
||||
help="Nickname to refer to the loaded LoRA adapter (useful for swapping).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--lora-strength",
|
||||
type=float,
|
||||
default=FastVideoArgs.lora_strength,
|
||||
help="Scale applied to every part of the inference LoRA adapter (default: 1.0).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--lora-target-modules",
|
||||
nargs="+",
|
||||
@@ -624,6 +665,15 @@ class FastVideoArgs:
|
||||
help=
|
||||
"JSON string of kwargs to pass to torch.compile. Example: '{\"backend\":\"inductor\",\"mode\":\"reduce-overhead\"}'",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--inference-torch-compile",
|
||||
action=StoreBoolean,
|
||||
default=FastVideoArgs.inference_torch_compile,
|
||||
help="Regional fullgraph torch.compile of each DiT transformer block at inference "
|
||||
"(port of the #1718 training-side regional compile). The loader injects fullgraph=True "
|
||||
"and inductor options {emulate_precision_casts: true}; non-traceable attention backends "
|
||||
"(VSA) degrade to eager with one warning. FASTVIDEO_INFERENCE_TORCH_COMPILE=1 is equivalent.",
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--dit-cpu-offload",
|
||||
@@ -821,20 +871,6 @@ class FastVideoArgs:
|
||||
|
||||
def check_fastvideo_args(self) -> None:
|
||||
"""Validate inference arguments for consistency"""
|
||||
from fastvideo.platforms import current_platform
|
||||
|
||||
if current_platform.is_mps():
|
||||
self.use_fsdp_inference = False
|
||||
self.dit_layerwise_offload = False
|
||||
|
||||
if self.dit_layerwise_offload:
|
||||
if self.use_fsdp_inference:
|
||||
logger.warning("dit_layerwise_offload is enabled, automatically disabling use_fsdp_inference.")
|
||||
self.use_fsdp_inference = False
|
||||
if self.dit_cpu_offload:
|
||||
logger.warning("dit_layerwise_offload is enabled, automatically disabling dit_cpu_offload.")
|
||||
self.dit_cpu_offload = False
|
||||
|
||||
# Validate mode and inference_mode consistency
|
||||
assert isinstance(self.mode, ExecutionMode), f"Mode must be an ExecutionMode enum, got {type(self.mode)}"
|
||||
assert self.mode in ExecutionMode.choices(), f"Invalid execution mode: {self.mode}"
|
||||
@@ -851,6 +887,14 @@ class FastVideoArgs:
|
||||
logger.warning("Mode is '%s' but inference_mode is False. Setting inference_mode to True.", self.mode)
|
||||
self.inference_mode = True
|
||||
|
||||
# Inference policy must wait until a worker owns and binds its device:
|
||||
# a unified-memory device disables layerwise offload before conflicts
|
||||
# are resolved, preserving an explicit FSDP request. Training does not
|
||||
# pass through the inference worker boundary, so retain its historical
|
||||
# constructor-time normalization.
|
||||
if not self.inference_mode:
|
||||
self._resolve_device_offload_conflicts()
|
||||
|
||||
if not self.inference_mode:
|
||||
assert self.hsdp_replicate_dim != -1, "hsdp_replicate_dim must be set for training"
|
||||
assert self.hsdp_shard_dim != -1, "hsdp_shard_dim must be set for training"
|
||||
@@ -885,6 +929,70 @@ class FastVideoArgs:
|
||||
self.pipeline_config.vae_config.load_encoder = True
|
||||
self.preprocess_config.check_preprocess_config()
|
||||
|
||||
def _resolve_device_offload_conflicts(self) -> None:
|
||||
"""Resolve offload modes after device-local policy has been applied."""
|
||||
from fastvideo.platforms import current_platform
|
||||
|
||||
if current_platform.is_mps():
|
||||
self.use_fsdp_inference = False
|
||||
self.dit_layerwise_offload = False
|
||||
|
||||
if self.dit_layerwise_offload:
|
||||
if self.use_fsdp_inference:
|
||||
logger.warning("dit_layerwise_offload is enabled, automatically disabling use_fsdp_inference.")
|
||||
self.use_fsdp_inference = False
|
||||
if self.dit_cpu_offload:
|
||||
logger.warning("dit_layerwise_offload is enabled, automatically disabling dit_cpu_offload.")
|
||||
self.dit_cpu_offload = False
|
||||
|
||||
def finalize_device_offload_policy(self, device_id: int = 0) -> bool:
|
||||
"""Apply device-local memory policy, then resolve incompatible modes."""
|
||||
has_unified_memory = self.disable_offload_on_unified_memory(device_id)
|
||||
self._resolve_device_offload_conflicts()
|
||||
return has_unified_memory
|
||||
|
||||
def disable_offload_on_unified_memory(self, device_id: int = 0, *, offload_flag: str | None = None) -> bool:
|
||||
"""Disable host offload after a worker has selected its device.
|
||||
|
||||
CUDA's unified-memory probe reads runtime device properties and may
|
||||
initialize a CUDA context. Callers must therefore use this only inside
|
||||
a device-owning process, after selecting and binding ``device_id``.
|
||||
Returning the classification lets direct component-loader callers
|
||||
apply the same policy to explicit per-call overrides. When
|
||||
``offload_flag`` is given, the return value says whether this policy
|
||||
covers that component role.
|
||||
"""
|
||||
from fastvideo.platforms import current_platform
|
||||
|
||||
cached_device_id = getattr(self, "_unified_memory_device_id", None)
|
||||
cached_result = getattr(self, "_unified_memory_result", None)
|
||||
if cached_device_id != device_id or cached_result is None:
|
||||
cached_result = current_platform.has_unified_memory(device_id)
|
||||
self._unified_memory_device_id = device_id
|
||||
self._unified_memory_result = cached_result
|
||||
|
||||
if not cached_result:
|
||||
return False
|
||||
|
||||
enabled_flags = [flag for flag in UNIFIED_MEMORY_OFFLOAD_FLAGS if getattr(self, flag)]
|
||||
if enabled_flags:
|
||||
try:
|
||||
device_name = current_platform.get_device_name(device_id)
|
||||
except Exception:
|
||||
# Device naming is diagnostic only. NVML can be unavailable on
|
||||
# an integrated GPU (for example Jetson), and its physical-
|
||||
# ordinal lookup cannot interpret CUDA_VISIBLE_DEVICES UUID/MIG
|
||||
# selectors. Neither case should undo an authoritative driver
|
||||
# classification.
|
||||
device_name = current_platform.device_name
|
||||
|
||||
for flag in enabled_flags:
|
||||
logger.info(
|
||||
"Disabling %s: %s has unified memory, so moving weights to the host duplicates "
|
||||
"them rather than freeing device memory.", flag, device_name)
|
||||
setattr(self, flag, False)
|
||||
return offload_flag is None or offload_flag in UNIFIED_MEMORY_OFFLOAD_FLAGS
|
||||
|
||||
|
||||
_current_fastvideo_args = None
|
||||
|
||||
|
||||
@@ -64,6 +64,29 @@ class BaseLayerWithLoRA(nn.Module):
|
||||
self.lora_A = None
|
||||
self.lora_B = None
|
||||
|
||||
@property
|
||||
def weight(self) -> torch.Tensor:
|
||||
"""The wrapped layer's weight.
|
||||
|
||||
Model code reads ``layer.weight`` for perfectly ordinary reasons -- MiniMax H3's
|
||||
AdaLN modulation casts its input with ``self.linear.weight.dtype``, and its
|
||||
``proj_in`` reads the tensor itself. Wrapping a layer should not change what
|
||||
reading it looks like from outside, and without this the wrap turns those into
|
||||
``AttributeError`` at the first forward. Forcing every model to be listed in
|
||||
``lora_target_modules`` around its own attribute access is the wrong fix: the
|
||||
list would have to be maintained against code it does not own, and getting it
|
||||
wrong fails at generation time rather than at load.
|
||||
|
||||
Resolved through ``base_layer`` on each access rather than cached, because
|
||||
merging replaces that module outright.
|
||||
"""
|
||||
return self.base_layer.weight
|
||||
|
||||
@property
|
||||
def bias(self) -> torch.Tensor | None:
|
||||
"""The wrapped layer's bias, or ``None`` when it has none."""
|
||||
return getattr(self.base_layer, "bias", None)
|
||||
|
||||
@torch.compile()
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
lora_A = self.lora_A
|
||||
@@ -80,6 +103,7 @@ class BaseLayerWithLoRA(nn.Module):
|
||||
delta = delta * (
|
||||
self.lora_alpha / self.lora_rank # type: ignore
|
||||
) # type: ignore
|
||||
delta = delta * self.lora_strength
|
||||
out, output_bias = self.base_layer(x)
|
||||
return out + delta, output_bias
|
||||
else:
|
||||
|
||||
+4
-6
@@ -65,8 +65,9 @@ DEFAULT_LOGGING_CONFIG = {
|
||||
|
||||
@lru_cache
|
||||
def _print_info_once(logger: Logger, msg: str) -> None:
|
||||
# Set the stacklevel to 2 to print the original caller's line info
|
||||
logger.info(msg, stacklevel=2)
|
||||
# The process-aware info wrapper owns stacklevel; passing it here would
|
||||
# supply the keyword twice when the wrapper delegates to Logger.log.
|
||||
logger.info(msg)
|
||||
|
||||
|
||||
@lru_cache
|
||||
@@ -114,10 +115,7 @@ def _info(logger: Logger,
|
||||
is_local_main_process = local_rank == 0
|
||||
|
||||
if (main_process_only and is_main_process) or (local_main_process_only and is_local_main_process):
|
||||
# Honor an explicit stacklevel (info_once routes through here with
|
||||
# stacklevel already set) instead of passing the keyword twice.
|
||||
stacklevel = kwargs.pop("stacklevel", 2)
|
||||
logger.log(logging.INFO, msg, *args, stacklevel=stacklevel, **kwargs)
|
||||
logger.log(logging.INFO, msg, *args, stacklevel=2, **kwargs)
|
||||
|
||||
global _warned_local_main_process, _warned_main_process
|
||||
|
||||
|
||||
@@ -24,6 +24,12 @@ from fastvideo.mlx_runtime.checkpoint import (
|
||||
load_mlx_dit_checkpoint,
|
||||
save_mlx_dit_checkpoint,
|
||||
)
|
||||
from fastvideo.mlx_runtime.checkpoint_compat import (
|
||||
UnsupportedMLXCheckpointError,
|
||||
discover_mlx_checkpoint,
|
||||
raise_if_unsupported_mlx_checkpoint,
|
||||
resolve_mlx_checkpoint,
|
||||
)
|
||||
from fastvideo.mlx_runtime.memory import (
|
||||
AppliedMemoryLimits,
|
||||
add_memory_limit_args,
|
||||
@@ -81,6 +87,7 @@ __all__ = [
|
||||
"MLXWanTransformerBlock",
|
||||
"RefinePlan",
|
||||
"TwoPassResult",
|
||||
"UnsupportedMLXCheckpointError",
|
||||
"UnsupportedMLXQuantizationError",
|
||||
"add_memory_limit_args",
|
||||
"apply_fast_spatial_upsample",
|
||||
@@ -92,8 +99,11 @@ __all__ = [
|
||||
"fastwan_shape",
|
||||
"fastwan_shape_from_config",
|
||||
"gib_to_bytes",
|
||||
"discover_mlx_checkpoint",
|
||||
"load_mlx_dit_checkpoint",
|
||||
"load_or_enhance_prompt",
|
||||
"raise_if_unsupported_mlx_checkpoint",
|
||||
"resolve_mlx_checkpoint",
|
||||
"mlx_dit_from_diffusers_safetensors",
|
||||
"mlx_block_weights_from_diffusers_safetensors",
|
||||
"mlx_block_weights_from_torch",
|
||||
|
||||
@@ -197,8 +197,8 @@ def load_mlx_dit_checkpoint(checkpoint_dir: str | Path, *, compile: bool = False
|
||||
manifest_path = checkpoint_dir / MANIFEST_FILENAME
|
||||
weights_path = checkpoint_dir / WEIGHTS_FILENAME
|
||||
if not manifest_path.exists() or not weights_path.exists():
|
||||
raise FileNotFoundError(f"Not an MLX DiT checkpoint directory: {checkpoint_dir} "
|
||||
f"(expected {MANIFEST_FILENAME} and {WEIGHTS_FILENAME}).")
|
||||
from fastvideo.mlx_runtime.checkpoint_compat import mlx_checkpoint_missing_hint
|
||||
raise FileNotFoundError(mlx_checkpoint_missing_hint(checkpoint_dir))
|
||||
|
||||
manifest = json.loads(manifest_path.read_text())
|
||||
version = manifest.get("format_version")
|
||||
|
||||
@@ -0,0 +1,181 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""Reject NVIDIA FastWan-QAD checkpoints on the Apple Silicon MLX path.
|
||||
|
||||
FastMetal-QAD is the Apple Silicon release: DMD2 students trained on the affine
|
||||
INT8 grid, shipped as packed ``mlx_dit.safetensors`` + ``mlx_dit.json``.
|
||||
``FastVideo/FastWan-QAD-1.3B`` and ``FastVideo/FastWan-QAD-FP8-1.3B`` are
|
||||
NVIDIA-only QAD checkpoints (NVFP4 / FP8). Loading them through the MLX
|
||||
Diffusers path silently requantizes the wrong weights and produces videos that
|
||||
ignore the prompt. Fail loudly instead.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import re
|
||||
from pathlib import Path
|
||||
|
||||
FASTMETAL_COLLECTION_URL = "https://huggingface.co/collections/FastVideo/fastmetal"
|
||||
FASTMETAL_BLOG_URL = "https://haoailab.com/blogs/fastmetal/"
|
||||
|
||||
FASTMETAL_MODEL_IDS = (
|
||||
"FastVideo/FastMetal-1.3B-QAD",
|
||||
"FastVideo/FastMetal-5B-QAD",
|
||||
"FastVideo/FastMetal-14B-QAD",
|
||||
)
|
||||
|
||||
MLX_DIT_MANIFEST = "mlx_dit.json"
|
||||
MLX_DIT_WEIGHTS = "mlx_dit.safetensors"
|
||||
CUDA_QAD_OVERLAY_DIR = "generator_inference_transformer"
|
||||
|
||||
# Directory / HF-cache names for the NVIDIA FastWan-QAD family. The INT8 Apple
|
||||
# snapshots used an older FastWan-QAD-INT8-* name; ``int8`` in the path is
|
||||
# excluded below so those packed MLX checkpoints still load.
|
||||
_NVIDIA_FASTWAN_QAD_RE = re.compile(
|
||||
r"(?:^|[/\\._-]|--)fastwan-qad-(?:fp8-)?1\.3b(?:-sa2)?(?:$|[/\\._-]|--)",
|
||||
re.IGNORECASE,
|
||||
)
|
||||
|
||||
_NVIDIA_QUANT_MARKERS = (
|
||||
"nvfp4",
|
||||
"nv_fp4",
|
||||
"sageattention3",
|
||||
"attn_qat_infer",
|
||||
)
|
||||
|
||||
|
||||
class UnsupportedMLXCheckpointError(ValueError):
|
||||
"""Raised when an NVIDIA FastWan-QAD (or similarly incompatible) tree is used on MLX."""
|
||||
|
||||
|
||||
def is_mlx_dit_checkpoint(path: str | Path) -> bool:
|
||||
"""Return True if ``path`` is a packed FastMetal / MLX DiT directory."""
|
||||
checkpoint_dir = Path(path)
|
||||
return (checkpoint_dir / MLX_DIT_MANIFEST).is_file() and (checkpoint_dir / MLX_DIT_WEIGHTS).is_file()
|
||||
|
||||
|
||||
def discover_mlx_checkpoint(*candidates: str | Path | None) -> Path | None:
|
||||
"""Return the first candidate that is a packed MLX DiT directory."""
|
||||
for candidate in candidates:
|
||||
if candidate is None:
|
||||
continue
|
||||
path = Path(candidate)
|
||||
if is_mlx_dit_checkpoint(path):
|
||||
return path
|
||||
return None
|
||||
|
||||
|
||||
def resolve_mlx_checkpoint(explicit: str | Path | None, *search_roots: str | Path | None) -> Path | None:
|
||||
"""Prefer an explicit ``--mlx-checkpoint``, otherwise scan search roots."""
|
||||
if explicit is not None:
|
||||
return Path(explicit)
|
||||
return discover_mlx_checkpoint(*search_roots)
|
||||
|
||||
|
||||
def mlx_checkpoint_missing_hint(checkpoint_dir: str | Path) -> str:
|
||||
"""Extra FileNotFoundError text when a directory is not a packed MLX DiT."""
|
||||
nvidia_reason = nvidia_fastwan_qad_reason(checkpoint_dir)
|
||||
prefix = (f"Not an MLX DiT checkpoint directory: {checkpoint_dir} "
|
||||
f"(expected {MLX_DIT_MANIFEST} and {MLX_DIT_WEIGHTS}).")
|
||||
if nvidia_reason is not None:
|
||||
return prefix + "\n\n" + _nvidia_fastwan_qad_message(Path(checkpoint_dir), nvidia_reason)
|
||||
return prefix + "\n\n" + _fastmetal_howto()
|
||||
|
||||
|
||||
def raise_if_unsupported_mlx_checkpoint(*paths: str | Path | None) -> None:
|
||||
"""Raise if any path is an NVIDIA FastWan-QAD tree that must not run on MLX.
|
||||
|
||||
Packed FastMetal / MLX DiT directories are always allowed, including the
|
||||
older FastWan-QAD-INT8 directory name, because those already contain
|
||||
``mlx_dit.json``.
|
||||
"""
|
||||
for path in paths:
|
||||
if path is None:
|
||||
continue
|
||||
checkpoint = Path(path)
|
||||
if is_mlx_dit_checkpoint(checkpoint):
|
||||
continue
|
||||
reason = nvidia_fastwan_qad_reason(checkpoint)
|
||||
if reason is None:
|
||||
continue
|
||||
raise UnsupportedMLXCheckpointError(_nvidia_fastwan_qad_message(checkpoint, reason))
|
||||
|
||||
|
||||
def nvidia_fastwan_qad_reason(path: str | Path) -> str | None:
|
||||
"""Return a short reason if ``path`` looks like NVIDIA FastWan-QAD, else None."""
|
||||
checkpoint = Path(path)
|
||||
haystack = _path_haystack(checkpoint)
|
||||
if "int8" in haystack and is_mlx_dit_checkpoint(checkpoint):
|
||||
return None
|
||||
if _NVIDIA_FASTWAN_QAD_RE.search(haystack) and "int8" not in haystack:
|
||||
return "NVIDIA FastWan-QAD checkpoint name (NVFP4/FP8, not FastMetal INT8)"
|
||||
if (checkpoint / CUDA_QAD_OVERLAY_DIR).is_dir() or (checkpoint.parent / CUDA_QAD_OVERLAY_DIR).is_dir():
|
||||
return f"CUDA QAD overlay directory ({CUDA_QAD_OVERLAY_DIR}/)"
|
||||
for config_path in _config_candidates(checkpoint):
|
||||
marker = _quant_marker_in_file(config_path)
|
||||
if marker is not None:
|
||||
return f"{config_path.name} contains NVIDIA quantization marker {marker!r}"
|
||||
return None
|
||||
|
||||
|
||||
def _path_haystack(path: Path) -> str:
|
||||
try:
|
||||
resolved = path.resolve()
|
||||
except OSError:
|
||||
resolved = path
|
||||
return str(resolved).replace("\\", "/").lower()
|
||||
|
||||
|
||||
def _config_candidates(path: Path) -> list[Path]:
|
||||
roots = [path.parent, path.parent.parent] if path.is_file() else [path, path.parent, path / "transformer"]
|
||||
seen: set[Path] = set()
|
||||
files: list[Path] = []
|
||||
for root in roots:
|
||||
for candidate in (root / "config.json", root / "model_index.json", root / "transformer" / "config.json"):
|
||||
if candidate in seen or not candidate.is_file():
|
||||
continue
|
||||
seen.add(candidate)
|
||||
files.append(candidate)
|
||||
return files
|
||||
|
||||
|
||||
def _quant_marker_in_file(config_path: Path) -> str | None:
|
||||
try:
|
||||
payload = config_path.read_text(encoding="utf-8")
|
||||
except OSError:
|
||||
return None
|
||||
lowered = payload.lower()
|
||||
for marker in _NVIDIA_QUANT_MARKERS:
|
||||
if marker in lowered:
|
||||
return marker
|
||||
try:
|
||||
parsed = json.loads(payload)
|
||||
except json.JSONDecodeError:
|
||||
return None
|
||||
quant = parsed.get("quantization_config") if isinstance(parsed, dict) else None
|
||||
if isinstance(quant, dict):
|
||||
blob = json.dumps(quant).lower()
|
||||
for marker in ("fp8", "float8", "nvfp4", "fp4"):
|
||||
if marker in blob:
|
||||
return marker
|
||||
return None
|
||||
|
||||
|
||||
def _fastmetal_howto() -> str:
|
||||
models = ", ".join(FASTMETAL_MODEL_IDS)
|
||||
return ("The Apple Silicon MLX runtime requires FastMetal-QAD INT8 checkpoints, "
|
||||
f"not CUDA FastWan-QAD weights.\n"
|
||||
f" Download: hf download FastVideo/FastMetal-1.3B-QAD --local-dir ./FastMetal-1.3B-QAD\n"
|
||||
" FastMetal repos ship mlx_dit.json (not transformer/config.json).\n"
|
||||
f" Run: python examples/inference/basic/mlx_wan_prompt_to_video.py "
|
||||
f"--model-root ./FastMetal-1.3B-QAD --mlx-checkpoint ./FastMetal-1.3B-QAD\n"
|
||||
f" 5B uses examples/inference/basic/mlx_wan22_generate.py with FastMetal-5B-QAD.\n"
|
||||
f" Models: {models}\n"
|
||||
f" Guide: {FASTMETAL_BLOG_URL}\n"
|
||||
f" Collection: {FASTMETAL_COLLECTION_URL}\n"
|
||||
"FastVideo/FastWan-QAD-1.3B and FastVideo/FastWan-QAD-FP8-1.3B are "
|
||||
"NVIDIA NVFP4/FP8 QAD checkpoints.")
|
||||
|
||||
|
||||
def _nvidia_fastwan_qad_message(path: Path, reason: str) -> str:
|
||||
return (f"Refusing to load {path} on the Apple Silicon MLX runtime ({reason}).\n\n" + _fastmetal_howto())
|
||||
@@ -885,6 +885,10 @@ def mlx_dit_from_diffusers_safetensors(
|
||||
import mlx.core as mx
|
||||
from safetensors import safe_open
|
||||
|
||||
from fastvideo.mlx_runtime.checkpoint_compat import raise_if_unsupported_mlx_checkpoint
|
||||
|
||||
raise_if_unsupported_mlx_checkpoint(checkpoint_path, config_path)
|
||||
|
||||
config = json.loads(Path(config_path).read_text())
|
||||
total_blocks = int(config["num_layers"])
|
||||
if num_blocks is None:
|
||||
|
||||
@@ -18,7 +18,10 @@ import torch.distributed as dist
|
||||
import fastvideo.envs as envs
|
||||
from fastvideo.attention import (DistributedAttention, LocalAttention)
|
||||
from fastvideo.configs.models.dits import WanVideoConfig
|
||||
from fastvideo.distributed.parallel_state import get_sp_world_size
|
||||
from fastvideo.distributed.communication_op import (
|
||||
sequence_model_parallel_all_gather)
|
||||
from fastvideo.distributed.parallel_state import (get_sp_parallel_rank,
|
||||
get_sp_world_size)
|
||||
from fastvideo.forward_context import get_forward_context
|
||||
from fastvideo.layers.layernorm import (FP32LayerNorm, LayerNormScaleShift, RMSNorm, ScaleResidual,
|
||||
ScaleResidualLayerNormScaleShift)
|
||||
@@ -98,6 +101,29 @@ class CausalWanSelfAttention(nn.Module):
|
||||
roped_query = _apply_rotary_emb(q, cos, sin, is_neox_style=False).type_as(v)
|
||||
roped_key = _apply_rotary_emb(k, cos, sin, is_neox_style=False).type_as(v)
|
||||
|
||||
# Ulysses sequence-parallel (SP): the causal pipeline replicates the full
|
||||
# (unsharded) token sequence on every rank, so RoPE, the block mask and
|
||||
# the KV cache are all full-length here. We parallelize the attention
|
||||
# across the head dimension (à la Ulysses): each rank computes attention
|
||||
# for its own slice of heads over the full sequence, then all-gathers the
|
||||
# per-head outputs to reconstruct all heads. This is numerically identical
|
||||
# to single-GPU attention (heads are independent) and avoids a sequence
|
||||
# all-to-all that would be incompatible with the causal KV-cache windowing
|
||||
# (expressed in global token coordinates). Applied uniformly to all three
|
||||
# attention paths below (flex, relativistic window, and the standard
|
||||
# KV-cache path).
|
||||
sp_world_size = get_sp_world_size()
|
||||
sp_rank = get_sp_parallel_rank()
|
||||
if sp_world_size > 1:
|
||||
assert self.num_heads % sp_world_size == 0, (
|
||||
f"num_heads ({self.num_heads}) must be divisible by sequence "
|
||||
f"parallel size ({sp_world_size})")
|
||||
heads_per_rank = self.num_heads // sp_world_size
|
||||
head_slice = slice(sp_rank * heads_per_rank,
|
||||
(sp_rank + 1) * heads_per_rank)
|
||||
else:
|
||||
head_slice = slice(None)
|
||||
|
||||
if kv_cache is None:
|
||||
# Padding for flex attention
|
||||
padded_length = math.ceil(q.shape[1] / 128) * 128 - q.shape[1]
|
||||
@@ -117,10 +143,13 @@ class CausalWanSelfAttention(nn.Module):
|
||||
[v, torch.zeros([v.shape[0], padded_length, v.shape[2], v.shape[3]], device=v.device, dtype=v.dtype)],
|
||||
dim=1)
|
||||
|
||||
x = flex_attention(query=padded_roped_query.transpose(2, 1),
|
||||
key=padded_roped_key.transpose(2, 1),
|
||||
value=padded_v.transpose(2, 1),
|
||||
x = flex_attention(query=padded_roped_query[:, :, head_slice].transpose(2, 1),
|
||||
key=padded_roped_key[:, :, head_slice].transpose(2, 1),
|
||||
value=padded_v[:, :, head_slice].transpose(2, 1),
|
||||
block_mask=block_mask)[:, :, :-padded_length].transpose(2, 1)
|
||||
if sp_world_size > 1:
|
||||
# Gather per-rank head outputs back to the full head dimension.
|
||||
x = sequence_model_parallel_all_gather(x.contiguous(), dim=2)
|
||||
else:
|
||||
current_end = current_start + q.shape[1]
|
||||
sink_tokens = self.sink_size * frame_seqlen
|
||||
@@ -178,7 +207,11 @@ class CausalWanSelfAttention(nn.Module):
|
||||
is_neox_style=False).type_as(v)
|
||||
key_window = _apply_rotary_emb(key_window, cos[:window_len], sin[:window_len],
|
||||
is_neox_style=False).type_as(v)
|
||||
x = self.attn(roped_query, key_window, value_window)
|
||||
x = self.attn(roped_query[:, :, head_slice], key_window[:, :, head_slice],
|
||||
value_window[:, :, head_slice])
|
||||
if sp_world_size > 1:
|
||||
# Gather per-rank head outputs back to the full head dimension.
|
||||
x = sequence_model_parallel_all_gather(x.contiguous(), dim=2)
|
||||
if isinstance(kv_cache["global_end_index"], torch.Tensor):
|
||||
kv_cache["global_end_index"].fill_(current_end)
|
||||
else:
|
||||
@@ -581,10 +614,16 @@ class CausalWanTransformer3DModel(BaseDiT):
|
||||
# fixed table over [0, max_attention_frames); attention slices it per step
|
||||
max_attention_frames = (GLOBAL_ATTN_COMPAT_MAX_LATENT_FRAMES
|
||||
if self.local_attn_size == -1 else self.local_attn_size)
|
||||
rope_num_frames = max_attention_frames * get_sp_world_size()
|
||||
# No multiplication by the sequence-parallel world size: the causal
|
||||
# pipeline replicates the full (unsharded) sequence on every rank, so
|
||||
# RoPE must cover the full local sequence length. Multiplying by the
|
||||
# SP world size produced a 2x-too-long RoPE and crashed
|
||||
# _apply_rotary_emb under SP. Attention is parallelized across heads
|
||||
# instead (see CausalWanSelfAttention.forward).
|
||||
rope_num_frames = max_attention_frames
|
||||
rope_start_frame = 0
|
||||
else:
|
||||
rope_num_frames = post_patch_num_frames * get_sp_world_size()
|
||||
rope_num_frames = post_patch_num_frames
|
||||
rope_start_frame = start_frame # 0 when kv_cache is None
|
||||
freqs_cos, freqs_sin = get_rotary_pos_embed((rope_num_frames, post_patch_height, post_patch_width),
|
||||
self.hidden_size,
|
||||
@@ -682,7 +721,12 @@ class CausalWanTransformer3DModel(BaseDiT):
|
||||
d = self.hidden_size // self.num_attention_heads
|
||||
rope_dim_list = [d - 4 * (d // 6), 2 * (d // 6), 2 * (d // 6)]
|
||||
freqs_cos, freqs_sin = get_rotary_pos_embed(
|
||||
(post_patch_num_frames * get_sp_world_size(), post_patch_height, post_patch_width),
|
||||
# No multiplication by the sequence-parallel world size: the causal
|
||||
# pipeline replicates the full (unsharded) sequence on every rank, so
|
||||
# RoPE must cover the full local sequence length. Attention is
|
||||
# parallelized across heads instead (see
|
||||
# CausalWanSelfAttention.forward).
|
||||
(post_patch_num_frames, post_patch_height, post_patch_width),
|
||||
self.hidden_size,
|
||||
self.num_attention_heads,
|
||||
rope_dim_list,
|
||||
|
||||
@@ -61,13 +61,13 @@ def _enabled_minimax_h3_fusions(value: str | None = None) -> frozenset[str]:
|
||||
|
||||
|
||||
def _can_run_minimax_h3_fusion(tensor: torch.Tensor) -> bool:
|
||||
"""Triton kernels are inference-only and stay outside Dynamo capture.
|
||||
"""Triton kernels are inference-only and trace as opaque custom ops.
|
||||
|
||||
The ``HAVE_TRITON`` check makes the eager fallback exact: on a CUDA build
|
||||
whose Triton failed to import, an enabled fusion falls back instead of
|
||||
hitting the strict wrappers' hard RuntimeError mid-forward.
|
||||
"""
|
||||
return (HAVE_TRITON and tensor.is_cuda and not torch.is_grad_enabled() and not torch.compiler.is_compiling())
|
||||
return HAVE_TRITON and tensor.is_cuda and not torch.is_grad_enabled()
|
||||
|
||||
|
||||
class MiniMaxH3RotaryPosEmbed(nn.Module):
|
||||
@@ -144,6 +144,7 @@ class MiniMaxH3Attention(nn.Module):
|
||||
quant_config: QuantizationConfig | None,
|
||||
prefix: str,
|
||||
fuse_qknorm_rope: bool = False,
|
||||
fa4_packed_varlen: bool = False,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.num_attention_heads = num_attention_heads
|
||||
@@ -196,6 +197,7 @@ class MiniMaxH3Attention(nn.Module):
|
||||
causal=False,
|
||||
supported_attention_backends=supported_attention_backends,
|
||||
prefix=prefix,
|
||||
fa4_packed_varlen=fa4_packed_varlen,
|
||||
)
|
||||
self.to_gate_compress: ReplicatedLinear | None = None
|
||||
# None = unchecked; the first forward tests the loaded weight once and
|
||||
@@ -221,12 +223,31 @@ class MiniMaxH3Attention(nn.Module):
|
||||
"""
|
||||
if torch.is_grad_enabled():
|
||||
return True
|
||||
if self._gate_compress_active is None:
|
||||
if torch.compiler.is_compiling():
|
||||
raise RuntimeError(
|
||||
"MiniMax H3 VSA compression gate was not resolved before torch.compile; "
|
||||
"call prepare_for_compile() after loading weights.")
|
||||
self._resolve_gate_compress_for_compile()
|
||||
assert self._gate_compress_active is not None
|
||||
return self._gate_compress_active
|
||||
|
||||
def _resolve_gate_compress_for_compile(self) -> None:
|
||||
"""Resolve the inference-only compression-gate branch eagerly.
|
||||
|
||||
Regional compilation wraps each transformer block with
|
||||
``fullgraph=True``. Resolving the loaded weight before capture keeps
|
||||
the GPU-to-host bool conversion and the cache mutation out of every
|
||||
compiled block graph. Grad-enabled forwards ignore this cached value
|
||||
in ``_gate_active`` so training can still learn from a zero gate.
|
||||
"""
|
||||
if self.to_gate_compress is None:
|
||||
return
|
||||
if self._gate_compress_active is None:
|
||||
weight = self.to_gate_compress.weight
|
||||
# bool() on a DTensor reduction resolves collectively, so every
|
||||
# rank caches the same answer.
|
||||
self._gate_compress_active = bool((weight != 0).any())
|
||||
return self._gate_compress_active
|
||||
|
||||
@staticmethod
|
||||
def _apply_rotary_emb(
|
||||
@@ -453,6 +474,7 @@ class MiniMaxH3TransformerBlock(nn.Module):
|
||||
fuse_modulate: bool = False,
|
||||
fuse_qknorm_rope: bool = False,
|
||||
fuse_swiglu: bool = False,
|
||||
fa4_packed_varlen: bool = False,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.norm1 = nn.RMSNorm(hidden_size, eps=norm_eps)
|
||||
@@ -465,6 +487,7 @@ class MiniMaxH3TransformerBlock(nn.Module):
|
||||
quant_config,
|
||||
prefix=f"{prefix}.attn",
|
||||
fuse_qknorm_rope=fuse_qknorm_rope,
|
||||
fa4_packed_varlen=fa4_packed_varlen,
|
||||
)
|
||||
self.norm2 = nn.RMSNorm(hidden_size, eps=norm_eps)
|
||||
self.ff = MiniMaxH3FeedForward(
|
||||
@@ -576,8 +599,8 @@ class MiniMaxH3Transformer3DModel(BaseDiT):
|
||||
if self.enabled_fusions:
|
||||
if HAVE_TRITON:
|
||||
logger.info(
|
||||
"MiniMax H3 inference fusions enabled: %s (CUDA inference-only; grad-enabled and "
|
||||
"torch.compile-captured forwards fall back to eager).",
|
||||
"MiniMax H3 inference fusions enabled: %s (CUDA inference-only; grad-enabled forwards "
|
||||
"fall back to eager; torch.compile captures opaque custom-op boundaries).",
|
||||
",".join(sorted(self.enabled_fusions)))
|
||||
else:
|
||||
logger.warning(
|
||||
@@ -683,6 +706,7 @@ class MiniMaxH3Transformer3DModel(BaseDiT):
|
||||
fuse_modulate="modulate" in self.enabled_fusions,
|
||||
fuse_qknorm_rope="qknorm_rope" in self.enabled_fusions,
|
||||
fuse_swiglu="swiglu" in self.enabled_fusions,
|
||||
fa4_packed_varlen=envs.FASTVIDEO_MINIMAX_H3_FA4_PACKED_VARLEN,
|
||||
) for index in range(arch.num_layers)
|
||||
])
|
||||
self.norm_out = MiniMaxH3AdaLayerNormOut(
|
||||
@@ -712,16 +736,61 @@ class MiniMaxH3Transformer3DModel(BaseDiT):
|
||||
def prepare_for_compile(self) -> None:
|
||||
"""Pipeline hook, called once right before torch.compile wraps the blocks.
|
||||
|
||||
Dynamo capture traces the eager branch of every fusion guard, so an
|
||||
enabled ``FASTVIDEO_MINIMAX_H3_FUSIONS`` set is silently inert inside
|
||||
compiled block forwards (H3 compiles per-block by default). Say so
|
||||
once instead of leaving the flag looking active.
|
||||
Resolve each loaded VSA compression gate eagerly. Generic and training
|
||||
compile retain their established attention dispatch; only the
|
||||
inference loader's separate ``prepare_for_regional_compile`` hook may
|
||||
preselect the inference-only sm_100a path.
|
||||
|
||||
The inference-only Triton fusions expose fake-backed custom operators,
|
||||
so Dynamo can keep them active as opaque nodes inside each fullgraph
|
||||
block instead of tracing into their launcher implementation.
|
||||
"""
|
||||
gate_states: list[bool] = []
|
||||
for block in self.transformer_blocks:
|
||||
attention = block.attn
|
||||
if attention.to_gate_compress is not None:
|
||||
attention._resolve_gate_compress_for_compile()
|
||||
assert attention._gate_compress_active is not None
|
||||
gate_states.append(attention._gate_compress_active)
|
||||
if gate_states:
|
||||
logger.info(
|
||||
"Resolved MiniMax H3 VSA compression gates before torch.compile: %d active, %d inactive",
|
||||
sum(gate_states),
|
||||
len(gate_states) - sum(gate_states),
|
||||
)
|
||||
if self.enabled_fusions:
|
||||
logger.warning(
|
||||
"torch.compile is enabled for MiniMax H3, so the requested inference fusions (%s) are "
|
||||
"inert inside compiled block forwards; the compiled eager path runs instead.",
|
||||
",".join(sorted(self.enabled_fusions)))
|
||||
logger.info(
|
||||
"MiniMax H3 inference fusions remain active under torch.compile through custom-op boundaries: %s",
|
||||
",".join(sorted(self.enabled_fusions)),
|
||||
)
|
||||
|
||||
def prepare_for_regional_compile(self) -> str | None:
|
||||
"""Resolve state used only by inference regional fullgraph compile."""
|
||||
self.prepare_for_compile()
|
||||
prepared_vsa_impls = 0
|
||||
unsupported_reasons: set[str] = set()
|
||||
for block in self.transformer_blocks:
|
||||
attention = block.attn
|
||||
prepare_vsa = getattr(attention.distributed_attention.attn_impl, "prepare_for_regional_compile", None)
|
||||
if not callable(prepare_vsa):
|
||||
continue
|
||||
# Post-load FP8 conversion may replace to_q.weight with packed
|
||||
# buffers. Either representation identifies the local device.
|
||||
query_state = next(attention.to_q.parameters(), None)
|
||||
if query_state is None:
|
||||
query_state = next(attention.to_q.buffers(), None)
|
||||
if query_state is None:
|
||||
raise RuntimeError("MiniMax H3 to_q has no materialized parameter or buffer for compile setup.")
|
||||
unsupported = prepare_vsa(query_state.device)
|
||||
if unsupported:
|
||||
unsupported_reasons.add(str(unsupported))
|
||||
prepared_vsa_impls += 1
|
||||
if prepared_vsa_impls:
|
||||
logger.info("Prepared %d MiniMax H3 VSA attention implementations for regional torch.compile",
|
||||
prepared_vsa_impls)
|
||||
if unsupported_reasons:
|
||||
return "; ".join(sorted(unsupported_reasons))
|
||||
return None
|
||||
|
||||
def materialize_non_persistent_buffers(
|
||||
self,
|
||||
|
||||
@@ -202,7 +202,7 @@ def _row_addressable(table: torch.Tensor) -> torch.Tensor:
|
||||
return table if table.stride(-1) == 1 else table.contiguous()
|
||||
|
||||
|
||||
def fused_rmsnorm_modulate(
|
||||
def _fused_rmsnorm_modulate_impl(
|
||||
x: torch.Tensor,
|
||||
weight: torch.Tensor,
|
||||
scale: torch.Tensor,
|
||||
@@ -210,13 +210,6 @@ def fused_rmsnorm_modulate(
|
||||
index: torch.Tensor,
|
||||
eps: float,
|
||||
) -> torch.Tensor:
|
||||
"""Run RMSNorm and row-indexed modulation in one strict Triton kernel.
|
||||
|
||||
``index`` values must lie in ``[0, table_rows)``. Unlike eager
|
||||
``index_select``, the kernel does not raise on out-of-range values (a
|
||||
device-side bounds check would synchronize); callers are safe by
|
||||
construction (``timestep_indices * 3 + token_tags``, SP pads with 0).
|
||||
"""
|
||||
_validate_contract(x, weight, (scale, shift), index, eps)
|
||||
_require_forward_only(x, weight, scale, shift)
|
||||
_require_triton_cuda(x)
|
||||
@@ -248,7 +241,7 @@ def fused_rmsnorm_modulate(
|
||||
return output.view_as(x)
|
||||
|
||||
|
||||
def fused_residual_gate_rmsnorm_modulate(
|
||||
def _fused_residual_gate_rmsnorm_modulate_impl(
|
||||
residual: torch.Tensor,
|
||||
branch: torch.Tensor,
|
||||
gate: torch.Tensor,
|
||||
@@ -258,11 +251,6 @@ def fused_residual_gate_rmsnorm_modulate(
|
||||
index: torch.Tensor,
|
||||
eps: float,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Fuse residual update, row-indexed gate, RMSNorm, and modulation.
|
||||
|
||||
``index`` values must lie in ``[0, table_rows)``; see
|
||||
:func:`fused_rmsnorm_modulate` for why the wrapper does not check them.
|
||||
"""
|
||||
_validate_residual_branch(residual, branch)
|
||||
_validate_contract(residual, weight, (gate, scale, shift), index, eps)
|
||||
_require_forward_only(residual, branch, gate, weight, scale, shift)
|
||||
@@ -300,3 +288,138 @@ def fused_residual_gate_rmsnorm_modulate(
|
||||
num_warps=_num_warps(block_size),
|
||||
)
|
||||
return hidden.view_as(residual), modulated.view_as(residual)
|
||||
|
||||
|
||||
# Dynamo must see the Triton launches as opaque nodes. In eager mode the
|
||||
# public wrappers below keep their strict, actionable input validation; while
|
||||
# compiling, these custom-op boundaries avoid tracing into Triton's launcher
|
||||
# and let the surrounding H3 block remain one full graph.
|
||||
@torch.library.custom_op(
|
||||
"fastvideo::_minimax_h3_rmsnorm_modulate",
|
||||
mutates_args=(),
|
||||
device_types="cuda",
|
||||
)
|
||||
def _fused_rmsnorm_modulate_op(
|
||||
x: torch.Tensor,
|
||||
weight: torch.Tensor,
|
||||
scale: torch.Tensor,
|
||||
shift: torch.Tensor,
|
||||
index: torch.Tensor,
|
||||
eps: float,
|
||||
) -> torch.Tensor:
|
||||
return _fused_rmsnorm_modulate_impl(x, weight, scale, shift, index, eps)
|
||||
|
||||
|
||||
@torch.library.register_fake("fastvideo::_minimax_h3_rmsnorm_modulate")
|
||||
def _fused_rmsnorm_modulate_fake(
|
||||
x: torch.Tensor,
|
||||
weight: torch.Tensor,
|
||||
scale: torch.Tensor,
|
||||
shift: torch.Tensor,
|
||||
index: torch.Tensor,
|
||||
eps: float,
|
||||
) -> torch.Tensor:
|
||||
del weight, scale, shift, index, eps
|
||||
return x.new_empty(x.shape)
|
||||
|
||||
|
||||
@torch.library.custom_op(
|
||||
"fastvideo::_minimax_h3_residual_gate_rmsnorm_modulate",
|
||||
mutates_args=(),
|
||||
device_types="cuda",
|
||||
)
|
||||
def _fused_residual_gate_rmsnorm_modulate_op(
|
||||
residual: torch.Tensor,
|
||||
branch: torch.Tensor,
|
||||
gate: torch.Tensor,
|
||||
weight: torch.Tensor,
|
||||
scale: torch.Tensor,
|
||||
shift: torch.Tensor,
|
||||
index: torch.Tensor,
|
||||
eps: float,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
return _fused_residual_gate_rmsnorm_modulate_impl(
|
||||
residual,
|
||||
branch,
|
||||
gate,
|
||||
weight,
|
||||
scale,
|
||||
shift,
|
||||
index,
|
||||
eps,
|
||||
)
|
||||
|
||||
|
||||
@torch.library.register_fake("fastvideo::_minimax_h3_residual_gate_rmsnorm_modulate")
|
||||
def _fused_residual_gate_rmsnorm_modulate_fake(
|
||||
residual: torch.Tensor,
|
||||
branch: torch.Tensor,
|
||||
gate: torch.Tensor,
|
||||
weight: torch.Tensor,
|
||||
scale: torch.Tensor,
|
||||
shift: torch.Tensor,
|
||||
index: torch.Tensor,
|
||||
eps: float,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
del branch, gate, weight, scale, shift, index, eps
|
||||
return residual.new_empty(residual.shape), residual.new_empty(residual.shape)
|
||||
|
||||
|
||||
def fused_rmsnorm_modulate(
|
||||
x: torch.Tensor,
|
||||
weight: torch.Tensor,
|
||||
scale: torch.Tensor,
|
||||
shift: torch.Tensor,
|
||||
index: torch.Tensor,
|
||||
eps: float,
|
||||
) -> torch.Tensor:
|
||||
"""Run RMSNorm and row-indexed modulation in one strict Triton kernel.
|
||||
|
||||
``index`` values must lie in ``[0, table_rows)``. Unlike eager
|
||||
``index_select``, the kernel does not raise on out-of-range values (a
|
||||
device-side bounds check would synchronize); callers are safe by
|
||||
construction (``timestep_indices * 3 + token_tags``, SP pads with 0).
|
||||
"""
|
||||
_require_forward_only(x, weight, scale, shift)
|
||||
if torch.compiler.is_compiling():
|
||||
return torch.ops.fastvideo._minimax_h3_rmsnorm_modulate(x, weight, scale, shift, index, eps)
|
||||
return _fused_rmsnorm_modulate_impl(x, weight, scale, shift, index, eps)
|
||||
|
||||
|
||||
def fused_residual_gate_rmsnorm_modulate(
|
||||
residual: torch.Tensor,
|
||||
branch: torch.Tensor,
|
||||
gate: torch.Tensor,
|
||||
weight: torch.Tensor,
|
||||
scale: torch.Tensor,
|
||||
shift: torch.Tensor,
|
||||
index: torch.Tensor,
|
||||
eps: float,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Fuse residual update, row-indexed gate, RMSNorm, and modulation.
|
||||
|
||||
``index`` values must lie in ``[0, table_rows)``; see
|
||||
:func:`fused_rmsnorm_modulate` for why the wrapper does not check them.
|
||||
"""
|
||||
_require_forward_only(residual, branch, gate, weight, scale, shift)
|
||||
if torch.compiler.is_compiling():
|
||||
return torch.ops.fastvideo._minimax_h3_residual_gate_rmsnorm_modulate(
|
||||
residual,
|
||||
branch,
|
||||
gate,
|
||||
weight,
|
||||
scale,
|
||||
shift,
|
||||
index,
|
||||
eps,
|
||||
)
|
||||
return _fused_residual_gate_rmsnorm_modulate_impl(
|
||||
residual,
|
||||
branch,
|
||||
gate,
|
||||
weight,
|
||||
scale,
|
||||
shift,
|
||||
index,
|
||||
eps,
|
||||
)
|
||||
|
||||
@@ -121,23 +121,13 @@ def _validate_inputs(
|
||||
return batch, seq_len, num_heads, head_dim, rotary_dim
|
||||
|
||||
|
||||
def fused_qknorm_rope(
|
||||
def _fused_qknorm_rope_impl(
|
||||
x: torch.Tensor,
|
||||
weight: torch.Tensor,
|
||||
cos: torch.Tensor,
|
||||
sin: torch.Tensor,
|
||||
eps: float,
|
||||
) -> torch.Tensor:
|
||||
"""Run per-head RMSNorm and partial RoPE in one Sol-Engine-style kernel.
|
||||
|
||||
RMSNorm reduction and RoPE arithmetic stay in FP32 registers until the
|
||||
final store. Triton's reduction order and the absence of eager's BF16
|
||||
intermediate materializations can produce small, expected rounding drift.
|
||||
|
||||
Row offsets are computed in int64, so inputs beyond 2**31 total elements
|
||||
(about 300k tokens per rank at H3's 56 heads x 128 head_dim) address
|
||||
correctly.
|
||||
"""
|
||||
batch, seq_len, num_heads, head_dim, rotary_dim = _validate_inputs(x, weight, cos, sin, eps)
|
||||
if not weight.is_contiguous():
|
||||
raise ValueError("weight must be contiguous")
|
||||
@@ -171,4 +161,56 @@ def fused_qknorm_rope(
|
||||
return flat_out.view(batch, seq_len, num_heads, head_dim)
|
||||
|
||||
|
||||
@torch.library.custom_op(
|
||||
"fastvideo::_minimax_h3_qknorm_rope",
|
||||
mutates_args=(),
|
||||
device_types="cuda",
|
||||
)
|
||||
def _fused_qknorm_rope_op(
|
||||
x: torch.Tensor,
|
||||
weight: torch.Tensor,
|
||||
cos: torch.Tensor,
|
||||
sin: torch.Tensor,
|
||||
eps: float,
|
||||
) -> torch.Tensor:
|
||||
return _fused_qknorm_rope_impl(x, weight, cos, sin, eps)
|
||||
|
||||
|
||||
@torch.library.register_fake("fastvideo::_minimax_h3_qknorm_rope")
|
||||
def _fused_qknorm_rope_fake(
|
||||
x: torch.Tensor,
|
||||
weight: torch.Tensor,
|
||||
cos: torch.Tensor,
|
||||
sin: torch.Tensor,
|
||||
eps: float,
|
||||
) -> torch.Tensor:
|
||||
del weight, cos, sin, eps
|
||||
return x.new_empty(x.shape)
|
||||
|
||||
|
||||
def fused_qknorm_rope(
|
||||
x: torch.Tensor,
|
||||
weight: torch.Tensor,
|
||||
cos: torch.Tensor,
|
||||
sin: torch.Tensor,
|
||||
eps: float,
|
||||
) -> torch.Tensor:
|
||||
"""Run per-head RMSNorm and partial RoPE in one Sol-Engine-style kernel.
|
||||
|
||||
RMSNorm reduction and RoPE arithmetic stay in FP32 registers until the
|
||||
final store. Triton's reduction order and the absence of eager's BF16
|
||||
intermediate materializations can produce small, expected rounding drift.
|
||||
|
||||
Row offsets are computed in int64, so inputs beyond 2**31 total elements
|
||||
(about 300k tokens per rank at H3's 56 heads x 128 head_dim) address
|
||||
correctly.
|
||||
"""
|
||||
if torch.is_grad_enabled() and any(
|
||||
isinstance(tensor, torch.Tensor) and tensor.requires_grad for tensor in (x, weight, cos, sin)):
|
||||
raise RuntimeError("fused_qknorm_rope is inference-only and does not implement autograd")
|
||||
if torch.compiler.is_compiling():
|
||||
return torch.ops.fastvideo._minimax_h3_qknorm_rope(x, weight, cos, sin, eps)
|
||||
return _fused_qknorm_rope_impl(x, weight, cos, sin, eps)
|
||||
|
||||
|
||||
__all__ = ["HAVE_TRITON", "fused_qknorm_rope"]
|
||||
|
||||
@@ -65,12 +65,7 @@ def _num_warps(block_size: int) -> int:
|
||||
return 4
|
||||
|
||||
|
||||
def minimax_h3_swiglu(x: torch.Tensor) -> torch.Tensor:
|
||||
"""Run the forward-only Triton fusion over an H3 ``(..., 2 * ffn_dim)`` input.
|
||||
|
||||
This is intentionally a strict kernel wrapper: callers own fallback policy and
|
||||
must only invoke it for a supported CUDA inference path.
|
||||
"""
|
||||
def _minimax_h3_swiglu_impl(x: torch.Tensor) -> torch.Tensor:
|
||||
ffn_dim = _validate_input(x)
|
||||
if not x.is_cuda:
|
||||
raise ValueError("MiniMax H3 fused SwiGLU requires a CUDA tensor")
|
||||
@@ -101,4 +96,31 @@ def minimax_h3_swiglu(x: torch.Tensor) -> torch.Tensor:
|
||||
return out.view(output_shape)
|
||||
|
||||
|
||||
@torch.library.custom_op(
|
||||
"fastvideo::_minimax_h3_swiglu",
|
||||
mutates_args=(),
|
||||
device_types="cuda",
|
||||
)
|
||||
def _minimax_h3_swiglu_op(x: torch.Tensor) -> torch.Tensor:
|
||||
return _minimax_h3_swiglu_impl(x)
|
||||
|
||||
|
||||
@torch.library.register_fake("fastvideo::_minimax_h3_swiglu")
|
||||
def _minimax_h3_swiglu_fake(x: torch.Tensor) -> torch.Tensor:
|
||||
return x.new_empty((*x.shape[:-1], x.shape[-1] // 2))
|
||||
|
||||
|
||||
def minimax_h3_swiglu(x: torch.Tensor) -> torch.Tensor:
|
||||
"""Run the forward-only Triton fusion over an H3 ``(..., 2 * ffn_dim)`` input.
|
||||
|
||||
This is intentionally a strict kernel wrapper: callers own fallback policy and
|
||||
must only invoke it for a supported CUDA inference path.
|
||||
"""
|
||||
if torch.is_grad_enabled() and x.requires_grad:
|
||||
raise RuntimeError("MiniMax H3 fused SwiGLU is forward-only and does not implement autograd")
|
||||
if torch.compiler.is_compiling():
|
||||
return torch.ops.fastvideo._minimax_h3_swiglu(x)
|
||||
return _minimax_h3_swiglu_impl(x)
|
||||
|
||||
|
||||
__all__ = ["HAVE_TRITON", "minimax_h3_swiglu"]
|
||||
|
||||
@@ -289,12 +289,27 @@ def _flashinfer_gemm_w8a8_block_fp8_linear_with_fallback(
|
||||
if backend == "cutlass":
|
||||
m, k = input_2d.shape
|
||||
n = weight.shape[0]
|
||||
expected_x_scale_shape = (k // block_k, m)
|
||||
expected_weight_scale_shape = (k // block_k, n // block_n)
|
||||
if x_scale.shape == (m, k // block_k):
|
||||
x_scale = x_scale.transpose(-1, -2).contiguous()
|
||||
if weight_scale.shape == (n // block_n, k // block_k):
|
||||
weight_scale = weight_scale.transpose(-1, -2).contiguous()
|
||||
# FlashInfer documents that ``m`` "should be padded to a multiple of 4
|
||||
# before calling this function", and the CUTLASS groupwise module
|
||||
# rejects unaligned m at dispatch (``cutlass gemm.can_implement
|
||||
# failed``). Prompt token counts are arbitrary, so pad the quantized
|
||||
# activation rows and their per-token scales with zeros, then slice
|
||||
# the padded rows off the GEMM output. The trtllm (sm100) route
|
||||
# tolerates any m and stays unpadded.
|
||||
padded_m = (m + 3) // 4 * 4
|
||||
if padded_m != m:
|
||||
padded_q_input = q_input.new_zeros((padded_m, q_input.shape[1]))
|
||||
padded_q_input[:m] = q_input
|
||||
q_input = padded_q_input
|
||||
padded_x_scale = x_scale.new_zeros((x_scale.shape[0], padded_m))
|
||||
padded_x_scale[:, :m] = x_scale
|
||||
x_scale = padded_x_scale
|
||||
expected_x_scale_shape = (k // block_k, padded_m)
|
||||
expected_weight_scale_shape = (k // block_k, n // block_n)
|
||||
if x_scale.shape != expected_x_scale_shape or weight_scale.shape != expected_weight_scale_shape:
|
||||
raise RuntimeError("FlashInfer CUTLASS block-FP8 scale layout mismatch: "
|
||||
f"x_scale={tuple(x_scale.shape)}, weight_scale={tuple(weight_scale.shape)}, "
|
||||
@@ -310,6 +325,8 @@ def _flashinfer_gemm_w8a8_block_fp8_linear_with_fallback(
|
||||
backend="cutlass",
|
||||
scale_major_mode="MN",
|
||||
)
|
||||
if padded_m != m:
|
||||
output = output[:m]
|
||||
else:
|
||||
expected_x_scale_shape = (input_2d.shape[0], input_2d.shape[1] // block_k)
|
||||
expected_weight_scale_shape = (weight.shape[0] // block_n, weight.shape[1] // block_k)
|
||||
|
||||
@@ -347,11 +347,24 @@ class TextEncoderLoader(ComponentLoader):
|
||||
dtype: str = "fp16",
|
||||
use_text_encoder_override: bool = False, # prevent subclasses from misusing
|
||||
cpu_offload: bool | None = None,
|
||||
offload_flag: str = "text_encoder_cpu_offload",
|
||||
):
|
||||
if cpu_offload is None:
|
||||
cpu_offload = fastvideo_args.text_encoder_cpu_offload
|
||||
use_cpu_offload = (cpu_offload and len(getattr(model_config, "_fsdp_shard_conditions", [])) > 0)
|
||||
runtime_device = get_local_torch_device()
|
||||
device_id = runtime_device.index if runtime_device.index is not None else 0
|
||||
requested_cpu_offload = getattr(fastvideo_args, offload_flag) if cpu_offload is None else cpu_offload
|
||||
disable_cpu_offload = fastvideo_args.disable_offload_on_unified_memory(device_id,
|
||||
offload_flag=offload_flag)
|
||||
|
||||
if requested_cpu_offload and disable_cpu_offload:
|
||||
# Direct loader callers can choose a CPU target before the worker
|
||||
# applies its device-local policy. Reset both the request and the
|
||||
# target so the model is never constructed on the host first.
|
||||
logger.info("Disabling %s on unified-memory device %d", offload_flag, device_id)
|
||||
cpu_offload = False
|
||||
target_device = runtime_device
|
||||
else:
|
||||
cpu_offload = requested_cpu_offload
|
||||
use_cpu_offload = (cpu_offload and len(getattr(model_config, "_fsdp_shard_conditions", [])) > 0)
|
||||
|
||||
from fastvideo.platforms import current_platform
|
||||
|
||||
@@ -555,6 +568,7 @@ class ImageEncoderLoader(TextEncoderLoader):
|
||||
fastvideo_args,
|
||||
encoder_precision,
|
||||
cpu_offload=fastvideo_args.image_encoder_cpu_offload,
|
||||
offload_flag="image_encoder_cpu_offload",
|
||||
)
|
||||
|
||||
|
||||
@@ -1126,6 +1140,13 @@ class TransformerLoader(ComponentLoader):
|
||||
training_mode=fastvideo_args.training_mode,
|
||||
enable_torch_compile=fastvideo_args.enable_torch_compile,
|
||||
torch_compile_kwargs=fastvideo_args.torch_compile_kwargs,
|
||||
inference_regional_compile=fastvideo_args.inference_torch_compile,
|
||||
inference_vsa_tile_size=fastvideo_args.VSA_tile_size,
|
||||
# Only the whole-parameter half of the adapter is applied here, while
|
||||
# tensors are still unsharded; LoRAPipeline merges the low-rank half
|
||||
# once the module tree exists.
|
||||
lora_path=getattr(fastvideo_args, "lora_path", None),
|
||||
lora_strength=getattr(fastvideo_args, "lora_strength", 1.0),
|
||||
)
|
||||
|
||||
total_params = sum(p.numel() for p in model.parameters())
|
||||
|
||||
@@ -7,6 +7,7 @@
|
||||
from __future__ import annotations
|
||||
import os
|
||||
import contextlib
|
||||
import re
|
||||
from collections.abc import Callable, Generator
|
||||
from itertools import chain
|
||||
from typing import Any
|
||||
@@ -19,6 +20,7 @@ from torch.distributed.fsdp import (CPUOffloadPolicy, FSDPModule, MixedPrecision
|
||||
from torch.nn.modules.module import _IncompatibleKeys
|
||||
|
||||
from fastvideo.logger import init_logger
|
||||
from fastvideo.models.loader.lora_patch import DenseLoRAPatch
|
||||
from fastvideo.models.loader.utils import (get_param_names_mapping, hf_to_custom_state_dict)
|
||||
from fastvideo.models.loader.weight_utils import safetensors_weights_iterator
|
||||
from fastvideo.utils import set_mixed_precision_policy, is_pin_memory_available
|
||||
@@ -26,6 +28,14 @@ from fastvideo.utils import set_mixed_precision_policy, is_pin_memory_available
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
def _summarize_param_names(names: set[str]) -> str:
|
||||
"""Collapse per-layer parameter names into one ``blocks.*.suffix xN`` entry each."""
|
||||
families: dict[str, int] = {}
|
||||
for name in names:
|
||||
families[re.sub(r"\.\d+\.", ".*.", name)] = families.get(re.sub(r"\.\d+\.", ".*.", name), 0) + 1
|
||||
return ", ".join(f"{family} x{count}" if count > 1 else family for family, count in sorted(families.items()))
|
||||
|
||||
|
||||
def _maybe_quantize_model(model: nn.Module) -> None:
|
||||
"""Quantize NVFP4- or FP8-tagged linear layers in-place after weights are loaded.
|
||||
|
||||
@@ -117,6 +127,23 @@ def set_default_dtype(dtype: torch.dtype) -> Generator[None, None, None]:
|
||||
torch.set_default_dtype(old_dtype)
|
||||
|
||||
|
||||
def _prepare_model_for_compile(model: nn.Module, *, regional: bool) -> str | None:
|
||||
"""Run a model compile hook, preferring its regional specialization."""
|
||||
prepare = None
|
||||
hook_name = "prepare_for_compile"
|
||||
if regional:
|
||||
prepare = getattr(model, "prepare_for_regional_compile", None)
|
||||
hook_name = "prepare_for_regional_compile"
|
||||
if not callable(prepare):
|
||||
prepare = getattr(model, "prepare_for_compile", None)
|
||||
hook_name = "prepare_for_compile"
|
||||
if not callable(prepare):
|
||||
return None
|
||||
logger.info("Running %s for %s", hook_name, type(model).__name__)
|
||||
unsupported = prepare()
|
||||
return unsupported if isinstance(unsupported, str) and unsupported else None
|
||||
|
||||
|
||||
# Supports optional torch.compile for FSDP-wrapped models during training
|
||||
def maybe_load_fsdp_model(
|
||||
model_cls: type[nn.Module],
|
||||
@@ -136,9 +163,20 @@ def maybe_load_fsdp_model(
|
||||
pin_cpu_memory: bool = True,
|
||||
enable_torch_compile: bool = False,
|
||||
torch_compile_kwargs: dict[str, Any] | None = None,
|
||||
inference_regional_compile: bool = False,
|
||||
inference_vsa_tile_size: int | None = None,
|
||||
lora_path: str | None = None,
|
||||
lora_strength: float = 1.0,
|
||||
) -> torch.nn.Module:
|
||||
"""
|
||||
Load the model with FSDP if is training, else load the model without FSDP.
|
||||
|
||||
``lora_path`` is consulted only for the part of an adapter that addresses whole
|
||||
parameters (``.diff`` / ``.set_weight``); see
|
||||
:mod:`fastvideo.models.loader.lora_patch`. The low-rank half is merged later by
|
||||
``LoRAPipeline``. Passing it here is what lets an adapter contribute a parameter the
|
||||
base checkpoint does not contain, which has to happen while the tensor is still
|
||||
unsharded.
|
||||
"""
|
||||
# NOTE(will): cast_forward_inputs=True shouldn't be needed as we are
|
||||
# manually casting the inputs to the model
|
||||
@@ -203,6 +241,23 @@ def maybe_load_fsdp_model(
|
||||
|
||||
weight_iterator = safetensors_weights_iterator(weight_dir_list, to_cpu=True)
|
||||
param_names_mapping_fn = get_param_names_mapping(model.param_names_mapping)
|
||||
dense_lora_patch = DenseLoRAPatch.from_adapter(
|
||||
lora_path,
|
||||
param_names_mapping_fn,
|
||||
strength=lora_strength,
|
||||
)
|
||||
if dense_lora_patch is not None:
|
||||
# H3's compression gate is created only by the VSA attention backend. Loading a
|
||||
# VSA student under dense attention would otherwise warn about 50 unmatched
|
||||
# replacements and continue with a silently incomplete model.
|
||||
model_parameter_names = {name for name, _ in model.named_parameters()}
|
||||
missing_vsa_gates = sorted(name for name in dense_lora_patch.replacement_parameters
|
||||
if "gate_compress" in name and name not in model_parameter_names)
|
||||
if missing_vsa_gates:
|
||||
raise ValueError(
|
||||
"This LoRA adapter provides MiniMax H3 VSA compression gates, but the selected attention backend "
|
||||
"did not construct them. Use attention_backend='VIDEO_SPARSE_ATTN_H3'. Missing parameters: "
|
||||
+ ", ".join(missing_vsa_gates[:3]) + (" ..." if len(missing_vsa_gates) > 3 else ""))
|
||||
load_model_from_full_model_state_dict(
|
||||
model,
|
||||
weight_iterator,
|
||||
@@ -211,6 +266,7 @@ def maybe_load_fsdp_model(
|
||||
strict=strict,
|
||||
cpu_offload=cpu_offload,
|
||||
param_names_mapping=param_names_mapping_fn,
|
||||
dense_lora_patch=dense_lora_patch,
|
||||
)
|
||||
if hasattr(model, "materialize_non_persistent_buffers"):
|
||||
model.materialize_non_persistent_buffers(device=device, dtype=default_dtype)
|
||||
@@ -231,13 +287,149 @@ def maybe_load_fsdp_model(
|
||||
|
||||
compile_in_loader = enable_torch_compile and training_mode
|
||||
if compile_in_loader:
|
||||
compile_kwargs = torch_compile_kwargs or {}
|
||||
logger.info("Enabling torch.compile for FSDP training module with kwargs=%s", compile_kwargs)
|
||||
model = torch.compile(model, **compile_kwargs)
|
||||
logger.info("torch.compile enabled for %s", type(model).__name__)
|
||||
unsupported = _prepare_model_for_compile(model, regional=False)
|
||||
if unsupported is not None:
|
||||
logger.warning("Training torch.compile requested but disabled: %s. Model stays eager.", unsupported)
|
||||
else:
|
||||
compile_kwargs = torch_compile_kwargs or {}
|
||||
logger.info("Enabling torch.compile for FSDP training module with kwargs=%s", compile_kwargs)
|
||||
model = torch.compile(model, **compile_kwargs)
|
||||
logger.info("torch.compile enabled for %s", type(model).__name__)
|
||||
elif inference_regional_compile and not training_mode:
|
||||
# Inference-side counterpart of the #1718 training regional compile:
|
||||
# per-block fullgraph compile right after the transformer loads, no
|
||||
# user kwargs needed (fullgraph + emulate_precision_casts injected).
|
||||
unsupported = _regional_compile_unsupported_reason(
|
||||
init_params,
|
||||
vsa_tile_size=inference_vsa_tile_size,
|
||||
)
|
||||
if unsupported is None:
|
||||
unsupported = _prepare_model_for_compile(model, regional=True)
|
||||
if unsupported is not None:
|
||||
logger.warning(
|
||||
"inference_torch_compile requested but disabled: %s. "
|
||||
"Inference continues in eager mode.", unsupported)
|
||||
else:
|
||||
attention_count = _enable_regional_attention_compile(model)
|
||||
logger.info("Enabled attention tracing for %d modules in %s", attention_count, type(model).__name__)
|
||||
_compile_model_regions(model, torch_compile_kwargs or {})
|
||||
return model
|
||||
|
||||
|
||||
def _regional_compile_unsupported_reason(
|
||||
init_params: dict[str, Any],
|
||||
*,
|
||||
vsa_tile_size: int | None = None,
|
||||
) -> str | None:
|
||||
"""Return why regional fullgraph compile cannot run, or None if it can.
|
||||
|
||||
Dense FA2, FA3, and FA4 inference all route through compile-visible
|
||||
custom-op boundaries. FA3's raw autograd.Function carve-out applies only
|
||||
to grad-enabled calls, outside this inference-only loader path.
|
||||
|
||||
The legacy VSA backend remains outside the fullgraph support envelope.
|
||||
MiniMax H3's VSA backend is supported only through the inference-only
|
||||
sm_100a tile-64 route; its regional hook resolves loaded compression
|
||||
gates and probes the kernel before block capture.
|
||||
"""
|
||||
try:
|
||||
from fastvideo.attention.layer import _attention_compile_explicitly_disabled
|
||||
except Exception: # pragma: no cover - attention stack not importable
|
||||
pass
|
||||
else:
|
||||
if _attention_compile_explicitly_disabled():
|
||||
# The escape hatch wraps attention forwards in
|
||||
# torch.compiler.disable, which is a hard dynamo error inside a
|
||||
# fullgraph region ("Skip inlining `torch.compiler.disable()`d
|
||||
# function"). Degrade to eager instead, matching the hatch's
|
||||
# debugging intent.
|
||||
return ("FASTVIDEO_DISABLE_ATTENTION_COMPILE=1 keeps attention "
|
||||
"forwards out of compiled graphs via torch.compiler."
|
||||
"disable, which fullgraph regional compile cannot trace; "
|
||||
"this model stays eager")
|
||||
config = init_params.get("config")
|
||||
resolved = getattr(config, "_resolved_attention_backend", None)
|
||||
resolved_name = getattr(resolved, "name", "")
|
||||
if resolved_name == "VIDEO_SPARSE_ATTN_H3":
|
||||
if os.environ.get("FASTVIDEO_H3_VSA_PROBE"):
|
||||
return ("FASTVIDEO_H3_VSA_PROBE records tensors and files from the VSA-H3 attention body, which "
|
||||
"regional fullgraph compile cannot capture; this model stays eager")
|
||||
if os.environ.get("FASTVIDEO_VSA_SM100A", "0") != "1":
|
||||
return ("VIDEO_SPARSE_ATTN_H3 regional compile requires the compile-safe sm_100a route "
|
||||
"(FASTVIDEO_VSA_SM100A=1); Triton/CuTe VSA stays eager")
|
||||
if vsa_tile_size != 64:
|
||||
return ("VIDEO_SPARSE_ATTN_H3 regional compile requires VSA_tile_size=64; "
|
||||
f"got {vsa_tile_size!r}, so tile-256/CuTe VSA stays eager")
|
||||
if resolved_name == "VIDEO_SPARSE_ATTN":
|
||||
return (f"attention backend resolved to {resolved_name}, whose Triton "
|
||||
"kernels, sequence-parallel collectives, and sync metadata "
|
||||
"guard graph-break (incompatible with fullgraph regional "
|
||||
"compile); this model stays eager")
|
||||
return None
|
||||
|
||||
|
||||
def _enable_regional_attention_compile(model: nn.Module) -> int:
|
||||
"""Opt in distributed-attention instances owned by ``model`` only."""
|
||||
from fastvideo.attention.layer import DistributedAttention
|
||||
|
||||
enabled_count = 0
|
||||
for submodule in model.modules():
|
||||
if isinstance(submodule, DistributedAttention):
|
||||
submodule._set_compile_forward_enabled(True)
|
||||
enabled_count += 1
|
||||
return enabled_count
|
||||
|
||||
|
||||
def _compile_model_regions(model: nn.Module, compile_kwargs: dict[str, Any]) -> int:
|
||||
"""Compile repeated mathematical regions of a loaded model.
|
||||
|
||||
Only the selected module ``forward`` is replaced. This keeps activation
|
||||
checkpoint wrappers structurally transparent while any module-level hooks
|
||||
(FSDP pre/post, layerwise offload) execute outside the compiled region.
|
||||
"""
|
||||
compile_conditions = getattr(model, "_compile_conditions", None)
|
||||
if not compile_conditions:
|
||||
raise ValueError(f"{type(model).__name__} does not declare _compile_conditions")
|
||||
|
||||
if compile_kwargs.get("fullgraph", True) is not True:
|
||||
raise ValueError("Regional compile requires fullgraph=True")
|
||||
if "mode" in compile_kwargs:
|
||||
# torch.compile forbids passing both `mode` and `options`, and
|
||||
# regional compile always injects options (emulate_precision_casts)
|
||||
# to match the training-side regional-compile configuration. Fail here
|
||||
# with an actionable message instead of letting torch raise a
|
||||
# mode/options conflict about an `options` key the user never wrote.
|
||||
raise ValueError("Regional compile sets inductor options "
|
||||
"(emulate_precision_casts) and cannot be combined "
|
||||
"with torch_compile_kwargs['mode']. Remove 'mode' or "
|
||||
"express its effect via torch_compile_kwargs['options'].")
|
||||
kwargs = {**compile_kwargs, "fullgraph": True}
|
||||
options = {"emulate_precision_casts": True}
|
||||
options.update(kwargs.get("options") or {})
|
||||
kwargs["options"] = options
|
||||
compiled_count = 0
|
||||
for name, submodule in list(model.named_modules()):
|
||||
if not name:
|
||||
continue
|
||||
if any(condition(name, submodule) for condition in compile_conditions):
|
||||
# Activation checkpoint wrappers are control-flow boundaries, not
|
||||
# mathematical regions. Keep their saved-tensor/recompute logic
|
||||
# eager and compile only the repeated block they own.
|
||||
compile_target = getattr(submodule, "_checkpoint_wrapped_module", submodule)
|
||||
compile_target.forward = torch.compile(compile_target.forward, **kwargs)
|
||||
compiled_count += 1
|
||||
|
||||
if compiled_count == 0:
|
||||
raise ValueError(f"No submodules in {type(model).__name__} matched _compile_conditions")
|
||||
logger.info(
|
||||
"Enabled regional torch.compile for %d submodules in %s with kwargs=%s",
|
||||
compiled_count,
|
||||
type(model).__name__,
|
||||
kwargs,
|
||||
)
|
||||
return compiled_count
|
||||
|
||||
|
||||
def shard_model(
|
||||
model,
|
||||
*,
|
||||
@@ -363,6 +555,7 @@ def load_model_from_full_model_state_dict(
|
||||
cpu_offload: bool = False,
|
||||
param_names_mapping: Callable[[str], tuple[str, Any, Any]] | None = None,
|
||||
training_mode: bool = True,
|
||||
dense_lora_patch: DenseLoRAPatch | None = None,
|
||||
) -> _IncompatibleKeys:
|
||||
"""
|
||||
Converting full state dict into a sharded state dict
|
||||
@@ -376,6 +569,9 @@ def load_model_from_full_model_state_dict(
|
||||
cpu_offload (bool): flag to check if FSDP offload is enabled
|
||||
param_names_mapping (Optional[Callable[[str], str]]): a function that maps full param name to sharded param name
|
||||
training_mode (bool): apply FSDP only for training
|
||||
dense_lora_patch (Optional[DenseLoRAPatch]): whole-parameter adapter payload,
|
||||
folded into each full tensor before it is sharded so that FSDP/TP placement
|
||||
is inherited from this function rather than reimplemented downstream
|
||||
Returns:
|
||||
``NamedTuple`` with ``missing_keys`` and ``unexpected_keys`` fields:
|
||||
* **missing_keys** is a list of str containing the missing keys
|
||||
@@ -390,7 +586,14 @@ def load_model_from_full_model_state_dict(
|
||||
sharded_sd = {}
|
||||
custom_param_sd, reverse_param_names_mapping = hf_to_custom_state_dict(full_sd_iterator,
|
||||
param_names_mapping) # type: ignore
|
||||
for target_param_name, full_tensor in custom_param_sd.items():
|
||||
# Drain rather than iterate. Production safetensors values may retain
|
||||
# memory-mapped shard storage, while mapped or merged parameters can own
|
||||
# ordinary allocations. Keeping the dict retains all of that source
|
||||
# storage until loading finishes; popping releases each reference as soon
|
||||
# as its conversion completes and lowers the host/unified-memory working
|
||||
# set.
|
||||
for target_param_name in list(custom_param_sd):
|
||||
full_tensor = custom_param_sd.pop(target_param_name)
|
||||
meta_sharded_param = meta_sd.get(target_param_name)
|
||||
if meta_sharded_param is None:
|
||||
# Some checkpoints include extra entries that are not part of the
|
||||
@@ -419,6 +622,10 @@ def load_model_from_full_model_state_dict(
|
||||
dtype_selector = getattr(model, "_get_parameter_dtype", None)
|
||||
if callable(dtype_selector):
|
||||
target_dtype = dtype_selector(target_param_name, param_dtype)
|
||||
if dense_lora_patch is not None:
|
||||
# Returns float32 when a delta was added, so the cast below is what lands
|
||||
# the parameter in its storage dtype.
|
||||
full_tensor = dense_lora_patch.apply_to(target_param_name, full_tensor)
|
||||
if not hasattr(meta_sharded_param, "device_mesh"):
|
||||
full_tensor = full_tensor.to(device=device, dtype=target_dtype)
|
||||
target_param = named_parameters.get(target_param_name)
|
||||
@@ -455,12 +662,28 @@ def load_model_from_full_model_state_dict(
|
||||
model.reverse_param_names_mapping = reverse_param_names_mapping
|
||||
unused_keys = set(meta_sd.keys()) - set(sharded_sd.keys())
|
||||
if unused_keys:
|
||||
logger.warning("Found unloaded parameters in meta state dict: %s", unused_keys)
|
||||
# Say which of these the adapter is about to fill in. Reporting all of them as
|
||||
# "unloaded" was accurate when zero-init was the only outcome; with an adapter
|
||||
# supplying real values it reads as a problem that is not one. Names are
|
||||
# summarized because a 50-layer model prints 50 near-identical lines otherwise.
|
||||
from_adapter = ({key for key in unused_keys if dense_lora_patch.provides(key)} if dense_lora_patch is not None
|
||||
else set())
|
||||
zero_init = unused_keys - from_adapter
|
||||
if from_adapter:
|
||||
logger.info("Parameters absent from the checkpoint and supplied by the LoRA adapter: %d (%s)",
|
||||
len(from_adapter), _summarize_param_names(from_adapter))
|
||||
if zero_init:
|
||||
logger.warning("Found unloaded parameters in meta state dict, zero-initializing: %d (%s)", len(zero_init),
|
||||
_summarize_param_names(zero_init))
|
||||
|
||||
# List of allowed parameter name patterns
|
||||
ALLOWED_NEW_PARAM_PATTERNS = ["gate_compress", "proj_l"] # Can be extended as needed
|
||||
for new_param_name in unused_keys:
|
||||
if not any(pattern in new_param_name for pattern in ALLOWED_NEW_PARAM_PATTERNS):
|
||||
# An adapter that ships the parameter outright both supplies the value and
|
||||
# authorizes it: the allowlist exists to catch a checkpoint silently missing a
|
||||
# weight, which is not the case when something deliberately provides one.
|
||||
adapter_value = (dense_lora_patch.replacement_for(new_param_name) if dense_lora_patch is not None else None)
|
||||
if adapter_value is None and not any(pattern in new_param_name for pattern in ALLOWED_NEW_PARAM_PATTERNS):
|
||||
logger.error("Unsupported new parameter: %s. Allowed patterns: %s", new_param_name,
|
||||
ALLOWED_NEW_PARAM_PATTERNS)
|
||||
raise ValueError(f"New parameter '{new_param_name}' is not supported. "
|
||||
@@ -470,7 +693,22 @@ def load_model_from_full_model_state_dict(
|
||||
dtype_selector = getattr(model, "_get_parameter_dtype", None)
|
||||
if callable(dtype_selector):
|
||||
target_dtype = dtype_selector(new_param_name, param_dtype)
|
||||
if not hasattr(meta_sharded_param, "device_mesh"):
|
||||
if adapter_value is not None:
|
||||
if tuple(adapter_value.shape) != tuple(meta_sharded_param.shape):
|
||||
raise ValueError(f"LoRA set_weight for {new_param_name} has shape {tuple(adapter_value.shape)}, "
|
||||
f"but the parameter is {tuple(meta_sharded_param.shape)}")
|
||||
full_tensor = adapter_value.to(device=device, dtype=target_dtype)
|
||||
if not hasattr(meta_sharded_param, "device_mesh"):
|
||||
sharded_tensor = full_tensor
|
||||
else:
|
||||
sharded_tensor = distribute_tensor(
|
||||
full_tensor,
|
||||
meta_sharded_param.device_mesh,
|
||||
meta_sharded_param.placements,
|
||||
)
|
||||
if cpu_offload:
|
||||
sharded_tensor = sharded_tensor.cpu()
|
||||
elif not hasattr(meta_sharded_param, "device_mesh"):
|
||||
# Initialize with zeros
|
||||
sharded_tensor = torch.zeros_like(meta_sharded_param, device=device, dtype=target_dtype)
|
||||
else:
|
||||
@@ -485,5 +723,8 @@ def load_model_from_full_model_state_dict(
|
||||
sharded_tensor = sharded_tensor.cpu()
|
||||
sharded_sd[new_param_name] = nn.Parameter(sharded_tensor)
|
||||
|
||||
if dense_lora_patch is not None:
|
||||
dense_lora_patch.report_unapplied()
|
||||
|
||||
# choose `assign=True` since we cannot call `copy_` on meta tensor
|
||||
return model.load_state_dict(sharded_sd, strict=strict, assign=True)
|
||||
|
||||
@@ -0,0 +1,271 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""The part of a LoRA adapter that is not a low-rank product.
|
||||
|
||||
A LoRA states a weight delta as ``B @ A``. That form needs two things to be true: the
|
||||
base checkpoint must already contain the parameter, and the delta must actually be low
|
||||
rank. Distilled video checkpoints break both often enough that dropping whatever does
|
||||
not fit silently loses real signal.
|
||||
|
||||
Two payload kinds cover the gap, named after the convention ComfyUI's loader already
|
||||
reads so one file works in both places:
|
||||
|
||||
``<module>.diff`` / ``<module>.diff_b``
|
||||
An exact additive delta for a parameter the base model has. Used where a rank-``r``
|
||||
factorization buys nothing or cannot be formed at all -- RMSNorm vectors, biases,
|
||||
and matrices whose smaller dimension is already at or below the rank that would be
|
||||
chosen. Factoring a length-``n`` vector into rank ``r`` costs ``r(1 + n) > n``.
|
||||
|
||||
``<module>.set_weight``
|
||||
A whole parameter the base model does not carry, so no delta is expressible. MiniMax
|
||||
H3's VSA ``to_gate_compress`` is the case that motivated this: it exists only under
|
||||
the sparse-attention backend, and :func:`load_model_from_full_model_state_dict`
|
||||
otherwise zero-initializes it, which is exactly the "gate contributes nothing"
|
||||
state a VSA-distilled student was trained away from.
|
||||
|
||||
Both kinds are applied to the *unsharded* tensor while the checkpoint is streaming in,
|
||||
so FSDP and tensor-parallel placement come from the surrounding loader instead of being
|
||||
reimplemented here. That ordering is not a preference: ``maybe_load_fsdp_model`` shards
|
||||
the model on the meta device before any weight is read, so there is no earlier window,
|
||||
and patching afterwards would mean gathering and redistributing every affected
|
||||
parameter one at a time.
|
||||
|
||||
Because the adapter has to be known when the transformer loads, this path applies the
|
||||
adapter a pipeline is constructed with. Swapping to a different adapter later still goes
|
||||
through :meth:`LoRAPipeline.set_lora_adapter`, which handles the low-rank half only.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
import os
|
||||
import re
|
||||
from collections.abc import Callable
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
from safetensors import safe_open
|
||||
|
||||
from fastvideo.logger import init_logger
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
# Suffix -> the parameter suffix it targets. ``.diff``/``.diff_b`` are additive,
|
||||
# ``.set_weight`` replaces. Ordered longest-first so ``.diff_b`` is tested before
|
||||
# ``.diff`` would match a truncated key.
|
||||
ADDITIVE_SUFFIXES: dict[str, str] = {".diff_b": ".bias", ".diff": ".weight"}
|
||||
REPLACEMENT_SUFFIXES: dict[str, str] = {".set_weight": ".weight"}
|
||||
|
||||
# Recognized elsewhere in an adapter and deliberately not our business: the low-rank
|
||||
# half, which ``LoRAPipeline`` merges through the wrapped-module path.
|
||||
_LOW_RANK_MARKERS = (".lora_A", ".lora_B", ".lora_up", ".lora_down", ".lora_alpha", ".lora_rank", ".alpha",
|
||||
".dora_scale")
|
||||
|
||||
|
||||
# One low-rank pair has many spellings. PEFT writes ``.lora_A.weight``, and interposes
|
||||
# the adapter's name when it is not the default (``.lora_A.default.weight``); kohya and
|
||||
# ComfyUI write ``.lora_down.weight`` with a bare ``.alpha``. Normalizing on the way in
|
||||
# means an adapter published in any of those layouts loads without a conversion pass.
|
||||
_ADAPTER_NAME_INFIX = re.compile(r"\.(lora_[AB])\.[^.]+\.weight$")
|
||||
_LOW_RANK_ALIASES = ((".lora_down", ".lora_A"), (".lora_up", ".lora_B"))
|
||||
|
||||
|
||||
def normalize_lora_key(name: str) -> str | None:
|
||||
"""Rewrite an adapter key to the ``<module>.lora_A|lora_B|lora_alpha`` spelling.
|
||||
|
||||
Returns ``None`` for keys the low-rank merge path deliberately does not handle --
|
||||
the dense payload above, which lands during checkpoint load, and bookkeeping like
|
||||
``.dora_scale``. Callers use that to tell "not mine" apart from "mine and
|
||||
unmatched", which is the difference between a quiet skip and a warning.
|
||||
"""
|
||||
if name.endswith(tuple(ADDITIVE_SUFFIXES) + tuple(REPLACEMENT_SUFFIXES)) or name.endswith(".dora_scale"):
|
||||
return None
|
||||
name = name.replace("diffusion_model.", "")
|
||||
name = _ADAPTER_NAME_INFIX.sub(r".\1.weight", name)
|
||||
for alias, canonical in _LOW_RANK_ALIASES:
|
||||
name = name.replace(alias + ".", canonical + ".")
|
||||
if name.endswith(alias):
|
||||
name = name[:-len(alias)] + canonical
|
||||
if name.endswith(".alpha"):
|
||||
name = name[:-len(".alpha")] + ".lora_alpha"
|
||||
return name
|
||||
|
||||
|
||||
def _adapter_files(lora_path: str) -> list[str]:
|
||||
"""Every safetensors shard belonging to an adapter given as a file or a directory."""
|
||||
if os.path.isfile(lora_path):
|
||||
return [lora_path]
|
||||
if os.path.isdir(lora_path):
|
||||
return sorted(
|
||||
os.path.join(lora_path, name) for name in os.listdir(lora_path) if name.endswith(".safetensors"))
|
||||
return []
|
||||
|
||||
|
||||
class DenseLoRAPatch:
|
||||
"""Adapter keys that address a whole parameter rather than a factor of one.
|
||||
|
||||
Construction only reads the safetensors headers, so the tensors themselves stay on
|
||||
disk until the loader asks for one. That keeps peak host memory at a single
|
||||
parameter even for adapters whose dense half is several GiB, which the VSA gates
|
||||
alone are.
|
||||
"""
|
||||
|
||||
def __init__(self, files: list[str], additive: dict[str, tuple[str, str]],
|
||||
replacement: dict[str, tuple[str, str]], strength: float = 1.0) -> None:
|
||||
if not math.isfinite(strength):
|
||||
raise ValueError(f"LoRA strength must be finite, got {strength}")
|
||||
self._files = files
|
||||
# target parameter name -> (file, adapter key)
|
||||
self._additive = additive
|
||||
self._replacement = replacement
|
||||
self._strength = float(strength)
|
||||
self._applied: set[str] = set()
|
||||
|
||||
@classmethod
|
||||
def from_adapter(
|
||||
cls,
|
||||
lora_path: str | None,
|
||||
param_names_mapping: Callable[[str], tuple[str, Any, Any]] | None = None,
|
||||
*,
|
||||
strength: float = 1.0,
|
||||
) -> DenseLoRAPatch | None:
|
||||
"""Build a patch from an adapter, or ``None`` when it carries no dense payload.
|
||||
|
||||
``param_names_mapping`` is the same callable the checkpoint loader uses, so
|
||||
adapter keys are resolved into the model's own parameter names by the identical
|
||||
rules -- an adapter written against the published checkpoint layout needs no
|
||||
separate conversion table.
|
||||
"""
|
||||
if not lora_path:
|
||||
return None
|
||||
# Deferred: fastvideo.utils pulls in enough of the package that importing it at
|
||||
# module scope would make this loader helper part of an import cycle.
|
||||
from fastvideo.utils import maybe_download_lora
|
||||
files = _adapter_files(maybe_download_lora(lora_path))
|
||||
if not files:
|
||||
logger.warning("LoRA path %s holds no safetensors file; no dense patch applied", lora_path)
|
||||
return None
|
||||
|
||||
additive: dict[str, tuple[str, str]] = {}
|
||||
replacement: dict[str, tuple[str, str]] = {}
|
||||
for path in files:
|
||||
with safe_open(path, framework="pt") as handle:
|
||||
for key in handle.keys():
|
||||
resolved = _resolve(key, param_names_mapping)
|
||||
if resolved is None:
|
||||
continue
|
||||
target, kind = resolved
|
||||
table = additive if kind == "add" else replacement
|
||||
if target in table:
|
||||
raise ValueError(f"LoRA adapter {lora_path} maps two keys onto parameter {target}; "
|
||||
f"the second is {key}")
|
||||
table[target] = (path, key)
|
||||
|
||||
if not additive and not replacement:
|
||||
return None
|
||||
logger.info(
|
||||
"LoRA adapter %s carries a dense payload: %d additive (.diff/.diff_b), %d replacement (.set_weight)",
|
||||
lora_path, len(additive), len(replacement))
|
||||
return cls(files, additive, replacement, strength)
|
||||
|
||||
def apply_to(self, param_name: str, tensor: torch.Tensor) -> torch.Tensor:
|
||||
"""Add this parameter's ``.diff``/``.diff_b`` delta, if the adapter has one.
|
||||
|
||||
The sum is taken in float32 whatever the operands are: both sides arrive in
|
||||
bfloat16, whose 8-bit significand would quantize away a delta three orders of
|
||||
magnitude below the base weight. The caller casts back to the target dtype.
|
||||
"""
|
||||
entry = self._additive.get(param_name)
|
||||
if entry is None:
|
||||
return tensor
|
||||
delta = self._read(entry)
|
||||
if delta.shape != tensor.shape:
|
||||
raise ValueError(f"LoRA diff for {param_name} has shape {tuple(delta.shape)}, "
|
||||
f"but the parameter is {tuple(tensor.shape)}")
|
||||
self._applied.add(param_name)
|
||||
return tensor.to(torch.float32) + delta.to(torch.float32) * self._strength
|
||||
|
||||
def provides(self, param_name: str) -> bool:
|
||||
"""Whether the adapter carries this parameter whole, without reading it."""
|
||||
return param_name in self._replacement
|
||||
|
||||
@property
|
||||
def replacement_parameters(self) -> frozenset[str]:
|
||||
"""Names of the parameters this adapter supplies outright.
|
||||
|
||||
Lets a caller decide what the adapter needs from the runtime -- an H3 adapter
|
||||
carrying ``to_gate_compress`` only works under the VSA backend -- without
|
||||
reading any tensor or guessing from the file name.
|
||||
"""
|
||||
return frozenset(self._replacement)
|
||||
|
||||
def replacement_for(self, param_name: str) -> torch.Tensor | None:
|
||||
"""The adapter's whole-tensor value for a parameter, or ``None``."""
|
||||
entry = self._replacement.get(param_name)
|
||||
if entry is None:
|
||||
return None
|
||||
self._applied.add(param_name)
|
||||
# Replacement payloads address parameters absent from the base checkpoint. The
|
||||
# loader initializes those parameters to zero, so scaling the supplied value is
|
||||
# the same interpolation contract as ``base + strength * delta``.
|
||||
return self._read(entry).to(torch.float32) * self._strength
|
||||
|
||||
def report_unapplied(self) -> None:
|
||||
"""Warn about dense keys that never reached a parameter.
|
||||
|
||||
An adapter key the loader silently ignores is the failure mode this whole path
|
||||
exists to fix, so it is said out loud rather than left to be inferred from a
|
||||
model that merely generates worse.
|
||||
"""
|
||||
pending = (set(self._additive) | set(self._replacement)) - self._applied
|
||||
if not pending:
|
||||
logger.info("LoRA dense payload fully applied: %d parameters", len(self._applied))
|
||||
return
|
||||
for target in sorted(pending):
|
||||
source = self._additive.get(target) or self._replacement[target]
|
||||
logger.warning("LoRA key not loaded: %s (targets %s, absent from the model)", source[1], target)
|
||||
logger.warning("LoRA dense payload: %d of %d parameters applied, %d keys unmatched", len(self._applied),
|
||||
len(self._applied) + len(pending), len(pending))
|
||||
|
||||
def _read(self, entry: tuple[str, str]) -> torch.Tensor:
|
||||
path, key = entry
|
||||
with safe_open(path, framework="pt") as handle:
|
||||
return handle.get_tensor(key)
|
||||
|
||||
|
||||
def _resolve(
|
||||
key: str,
|
||||
param_names_mapping: Callable[[str], tuple[str, Any, Any]] | None,
|
||||
) -> tuple[str, str] | None:
|
||||
"""Map an adapter key to ``(model parameter name, "add" | "set")``.
|
||||
|
||||
Returns ``None`` for anything that is not a dense payload key, which includes every
|
||||
low-rank factor -- those belong to ``LoRAPipeline``, not here.
|
||||
"""
|
||||
if any(marker in key for marker in _LOW_RANK_MARKERS):
|
||||
return None
|
||||
for suffix, param_suffix in ADDITIVE_SUFFIXES.items():
|
||||
if key.endswith(suffix):
|
||||
return _map_name(key[:-len(suffix)] + param_suffix, param_names_mapping, key), "add"
|
||||
for suffix, param_suffix in REPLACEMENT_SUFFIXES.items():
|
||||
if key.endswith(suffix):
|
||||
return _map_name(key[:-len(suffix)] + param_suffix, param_names_mapping, key), "set"
|
||||
return None
|
||||
|
||||
|
||||
def _map_name(
|
||||
param_name: str,
|
||||
param_names_mapping: Callable[[str], tuple[str, Any, Any]] | None,
|
||||
source_key: str,
|
||||
) -> str:
|
||||
"""Run the checkpoint loader's own renaming rules over a resolved parameter name."""
|
||||
param_name = param_name.replace("diffusion_model.", "")
|
||||
if param_names_mapping is None:
|
||||
return param_name
|
||||
mapped, merge_index, _ = param_names_mapping(param_name)
|
||||
if merge_index is not None:
|
||||
# A fused target (stacked QKV and friends) takes its value from several
|
||||
# checkpoint tensors. Splitting a whole-tensor payload across that fusion is
|
||||
# not something we can infer, and guessing would corrupt the weight quietly.
|
||||
raise NotImplementedError(f"LoRA dense key {source_key} resolves to fused parameter {mapped}; "
|
||||
"whole-tensor payloads for fused parameters are not supported")
|
||||
return mapped
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user