From f9e3680f1103413826c427379967fbbc88fbf5fe Mon Sep 17 00:00:00 2001 From: William Lin Date: Sun, 23 Aug 2026 04:01:03 -0500 Subject: [PATCH] [perf] Align FastH3 optimized inference profile (#1748) --- .github/workflows/publish-kernel.yml | 18 +- docs/inference/optimizations.md | 26 +- examples/inference/basic/README.md | 48 +- examples/inference/basic/basic_fasth3.py | 412 ++++++++++++------ fastvideo-kernel/README.md | 4 +- fastvideo-kernel/pyproject.toml | 2 +- .../block_sparse_attn_sm100a.py | 122 +++++- .../python/fastvideo_kernel/version.py | 2 +- .../backends/video_sparse_attn_h3.py | 233 ++++++++-- fastvideo/fastvideo_args.py | 8 +- fastvideo/models/dits/minimax_h3.py | 90 +++- .../dits/minimax_h3_fusions/modulation.py | 151 ++++++- .../dits/minimax_h3_fusions/qknorm_rope.py | 64 ++- .../models/dits/minimax_h3_fusions/swiglu.py | 34 +- fastvideo/models/loader/component_loader.py | 1 + fastvideo/models/loader/fsdp_load.py | 69 ++- .../attention/test_vsa_h3_sm100a_route.py | 237 ++++++++++ .../inference/test_basic_fasth3_profile.py | 238 ++++++++++ .../test_inference_regional_compile.py | 196 ++++++++- .../test_minimax_h3_fusion_compile.py | 159 +++++++ pyproject.toml | 25 +- tests/test_block_sparse_sm100a.py | 69 +++ tests/test_fasth3_packaging.py | 38 ++ 23 files changed, 1981 insertions(+), 265 deletions(-) create mode 100644 fastvideo/tests/inference/test_basic_fasth3_profile.py create mode 100644 fastvideo/tests/transformers/test_minimax_h3_fusion_compile.py create mode 100644 tests/test_fasth3_packaging.py diff --git a/.github/workflows/publish-kernel.yml b/.github/workflows/publish-kernel.yml index cf0282ea8..e73956c70 100644 --- a/.github/workflows/publish-kernel.yml +++ b/.github/workflows/publish-kernel.yml @@ -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 diff --git a/docs/inference/optimizations.md b/docs/inference/optimizations.md index c774dab53..753fef283 100644 --- a/docs/inference/optimizations.md +++ b/docs/inference/optimizations.md @@ -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 ``` @@ -362,9 +363,13 @@ loader wraps each `_compile_conditions` block in `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. VSA and the explicit -`FASTVIDEO_DISABLE_ATTENTION_COMPILE=1` escape hatch keep -the transformer eager with one warning instead of failing mid-denoise. +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( @@ -409,6 +414,14 @@ 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 @@ -431,9 +444,10 @@ MS-SSIM gate on *your* config, especially when combining 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; VSA and the FA3 grad-enabled path remain outside the regional - support envelope. Use the default inductor mode shown above unless your - exact configuration has its own gate. + 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 diff --git a/examples/inference/basic/README.md b/examples/inference/basic/README.md index 2964b709c..e25022fa3 100644 --- a/examples/inference/basic/README.md +++ b/examples/inference/basic/README.md @@ -33,20 +33,28 @@ 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: +For 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]" ``` -python examples/inference/basic/basic_fasth3.py --prompt "your prompt" [--vsa-sparsity 0.9] + +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 `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. -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 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. -For the fastest measured four-GPU Preview profile on GB200, keep the sparse DiT eager, compile and parallelize only the video VAE, replicate the DiT weights, and enable the inference-only H3 fusions explicitly: +The equivalent explicit command is: ```bash -FASTVIDEO_MINIMAX_H3_FUSIONS=all \ python examples/inference/basic/basic_fasth3.py \ --prompt "your prompt" \ + --profile all \ --num-gpus 4 \ --steps 5 \ --vsa-sparsity 0.9 \ @@ -54,10 +62,36 @@ python examples/inference/basic/basic_fasth3.py \ --vsa-kernel sm100a \ --compile-vae \ --parallel-vae \ - --replicated-dit + --replicated-dit \ + --pin-cpu-memory \ + --fa4 \ + --no-torch-compile \ + --inference-torch-compile \ + --ulysses-a2a off \ + --warmup \ + --repeats 3 \ + --seed 1000 \ + --warmup-seed 999 ``` -The fusion profile changes floating-point operation order and is a speed/quality evaluation route, not an exact-parity route. Omit `FASTVIDEO_MINIMAX_H3_FUSIONS` for strict sparse-DiT numerics; VAE chunk parallelism itself is bit-exact with the eager decoder. +`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. ## Basic Walkthrough diff --git a/examples/inference/basic/basic_fasth3.py b/examples/inference/basic/basic_fasth3.py index 9c7198ea1..7f0f4b6f4 100644 --- a/examples/inference/basic/basic_fasth3.py +++ b/examples/inference/basic/basic_fasth3.py @@ -1,22 +1,26 @@ # 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 @@ -32,170 +36,306 @@ from fastvideo.api import ( SamplingConfig, ) +DEFAULT_MODEL = "FastVideo/FastVideo-Minimax-FastH3-Preview-v0.1" -def parse_args() -> argparse.Namespace: + +def parse_args(argv: Sequence[str] | None = None) -> 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). + 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="store_true", - help="torch.compile the video VAE decoder independently of the DiT") + action=argparse.BooleanOptionalAction, + default=True, + help="compile the video VAE decoder independently of the DiT") parser.add_argument("--parallel-vae", - action="store_true", - help="round-robin video VAE temporal chunks across sequence-parallel ranks") + action=argparse.BooleanOptionalAction, + default=True, + help="round-robin VAE temporal chunks across sequence-parallel ranks") parser.add_argument("--replicated-dit", - action="store_true", - help="replicate DiT weights on every GPU instead of FSDP-sharding them; faster when the " - "checkpoint fits on each GPU") + 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("--inference-torch-compile", - action="store_true", - help="regional fullgraph torch.compile of each DiT block after load. NOTE: this " - "script always runs the VSA-H3 backend, which is not fullgraph-traceable — the " - "loader logs one warning and keeps the transformer eager. The flag is exposed " - "here to exercise exactly that guard") - 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") + args = parser.parse_args(argv) + 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 main() -> None: - args = parse_args() - output_dir = Path(args.output) - output_dir.mkdir(parents=True, exist_ok=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" - 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" - # 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. +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. + """ + return { + "FASTVIDEO_ATTENTION_BACKEND": "VIDEO_SPARSE_ATTN_H3", + "FASTVIDEO_VSA_SM100A": "1" if 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 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: experimental: dict[str, object] = { "attention_backend": "VIDEO_SPARSE_ATTN_H3", + "VSA_sparsity": args.vsa_sparsity, "VSA_tile_size": args.vsa_tile_size, + "inference_torch_compile": args.inference_torch_compile, + "vae_parallel_decode": args.parallel_vae, + "vae_parallel_decode_strategy": "gather", } - if args.vsa_sparsity > 0.0: - experimental["VSA_sparsity"] = args.vsa_sparsity - if args.inference_torch_compile: - experimental["inference_torch_compile"] = True - if args.parallel_vae: - experimental["vae_parallel_decode"] = True + return GeneratorConfig( + model_path=args.model_path, + pipeline=PipelineSelection(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, + ), + ), + ) - 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 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=False, - ), - 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) + + 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 ''}" + for key, value in environment.items())) + + 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() diff --git a/fastvideo-kernel/README.md b/fastvideo-kernel/README.md index d1603a92a..2e14b7749 100644 --- a/fastvideo-kernel/README.md +++ b/fastvideo-kernel/README.md @@ -10,6 +10,7 @@ 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) | @@ -26,7 +27,7 @@ Runtime-JIT kernels (no build step, ship in every wheel/image): | 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;12.0a` | ON | AUTO | ON | +| | | 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 | — | @@ -36,6 +37,7 @@ Runtime-JIT kernels (no build step, ship in every wheel/image): 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. diff --git a/fastvideo-kernel/pyproject.toml b/fastvideo-kernel/pyproject.toml index 2b9dbfb54..b9a7a73fb 100644 --- a/fastvideo-kernel/pyproject.toml +++ b/fastvideo-kernel/pyproject.toml @@ -9,7 +9,7 @@ build-backend = "scikit_build_core.build" [project] name = "fastvideo-kernel" -version = "0.3.3" +version = "0.3.4" description = "Unified CUDA kernels for FastVideo" readme = "README.md" requires-python = ">=3.10" diff --git a/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_sm100a.py b/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_sm100a.py index 1ea406efc..a254b6586 100644 --- a/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_sm100a.py +++ b/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_sm100a.py @@ -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] diff --git a/fastvideo-kernel/python/fastvideo_kernel/version.py b/fastvideo-kernel/python/fastvideo_kernel/version.py index e19434e2e..334b89956 100644 --- a/fastvideo-kernel/python/fastvideo_kernel/version.py +++ b/fastvideo-kernel/python/fastvideo_kernel/version.py @@ -1 +1 @@ -__version__ = "0.3.3" +__version__ = "0.3.4" diff --git a/fastvideo/attention/backends/video_sparse_attn_h3.py b/fastvideo/attention/backends/video_sparse_attn_h3.py index a775e8722..38f0daa04 100644 --- a/fastvideo/attention/backends/video_sparse_attn_h3.py +++ b/fastvideo/attention/backends/video_sparse_attn_h3.py @@ -94,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. @@ -245,6 +304,10 @@ 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 @@ -288,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, @@ -298,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, ) @@ -376,6 +441,62 @@ 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). @@ -392,8 +513,16 @@ class MiniMaxH3VSAImpl(AttentionImpl): "routed to the VSA-H3 backend; exclude it from the supported backends.") n_tiles = attn_metadata.variable_block_sizes.numel() grad_mode = torch.is_grad_enabled() and x.requires_grad - needs_sm100a_pair = (attn_metadata.tile_elems == 64 and n_tiles % 2 != 0 and not grad_mode - and os.environ.get(VSA_SM100A_ENV, "0") == "1") + 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]) @@ -430,13 +559,22 @@ 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 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 @@ -464,10 +602,18 @@ class MiniMaxH3VSAImpl(AttentionImpl): 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 - layer_sparsity = 0.0 if self.layer_idx in attn_metadata.dense_layers else attn_metadata.VSA_sparsity - probe_dir = probe_enabled() + # 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: @@ -487,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' @@ -515,35 +665,64 @@ class MiniMaxH3VSAImpl(AttentionImpl): # 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) + 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: - logger.info_once("MiniMax-H3 VSA tile-64 forward: using the sm100a CUDA block-sparse kernel") - # 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(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, - ) + # 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() diff --git a/fastvideo/fastvideo_args.py b/fastvideo/fastvideo_args.py index fc6ba9358..5bb9e0fce 100644 --- a/fastvideo/fastvideo_args.py +++ b/fastvideo/fastvideo_args.py @@ -182,10 +182,10 @@ class FastVideoArgs: # 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). Attention backends that cannot be fullgraph-traced (VSA) - # 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 + # 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. diff --git a/fastvideo/models/dits/minimax_h3.py b/fastvideo/models/dits/minimax_h3.py index 12fa30145..ec168cec1 100644 --- a/fastvideo/models/dits/minimax_h3.py +++ b/fastvideo/models/dits/minimax_h3.py @@ -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): @@ -223,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( @@ -580,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( @@ -717,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, diff --git a/fastvideo/models/dits/minimax_h3_fusions/modulation.py b/fastvideo/models/dits/minimax_h3_fusions/modulation.py index 5cc3213b3..961250edb 100644 --- a/fastvideo/models/dits/minimax_h3_fusions/modulation.py +++ b/fastvideo/models/dits/minimax_h3_fusions/modulation.py @@ -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, + ) diff --git a/fastvideo/models/dits/minimax_h3_fusions/qknorm_rope.py b/fastvideo/models/dits/minimax_h3_fusions/qknorm_rope.py index d8aad9bef..1575cbf57 100644 --- a/fastvideo/models/dits/minimax_h3_fusions/qknorm_rope.py +++ b/fastvideo/models/dits/minimax_h3_fusions/qknorm_rope.py @@ -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"] diff --git a/fastvideo/models/dits/minimax_h3_fusions/swiglu.py b/fastvideo/models/dits/minimax_h3_fusions/swiglu.py index b7884639e..9c703ce16 100644 --- a/fastvideo/models/dits/minimax_h3_fusions/swiglu.py +++ b/fastvideo/models/dits/minimax_h3_fusions/swiglu.py @@ -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"] diff --git a/fastvideo/models/loader/component_loader.py b/fastvideo/models/loader/component_loader.py index 08e74b8f8..6d6aadc79 100644 --- a/fastvideo/models/loader/component_loader.py +++ b/fastvideo/models/loader/component_loader.py @@ -1127,6 +1127,7 @@ class TransformerLoader(ComponentLoader): 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, ) total_params = sum(p.numel() for p in model.parameters()) diff --git a/fastvideo/models/loader/fsdp_load.py b/fastvideo/models/loader/fsdp_load.py index 904201abc..186d37f19 100644 --- a/fastvideo/models/loader/fsdp_load.py +++ b/fastvideo/models/loader/fsdp_load.py @@ -117,6 +117,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], @@ -137,6 +154,7 @@ def maybe_load_fsdp_model( 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, ) -> torch.nn.Module: """ Load the model with FSDP if is training, else load the model without FSDP. @@ -232,41 +250,50 @@ 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) + 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: - prepare_for_compile = getattr(model, "prepare_for_compile", None) - if callable(prepare_for_compile): - logger.info("Running prepare_for_compile for %s", type(model).__name__) - prepare_for_compile() 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]) -> str | None: +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. - The VSA backends (Triton block-sparse kernels behind sequence-parallel - all-to-alls plus a host-synced metadata guard) are likewise not - fullgraph-traceable; a VSA-backed transformer (e.g. the FastH3 student) - falls back to eager while compile-safe dense loads still compile. - 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 @@ -286,7 +313,17 @@ def _regional_compile_unsupported_reason(init_params: dict[str, Any]) -> str | N config = init_params.get("config") resolved = getattr(config, "_resolved_attention_backend", None) resolved_name = getattr(resolved, "name", "") - if resolved_name in ("VIDEO_SPARSE_ATTN", "VIDEO_SPARSE_ATTN_H3"): + 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 " diff --git a/fastvideo/tests/attention/test_vsa_h3_sm100a_route.py b/fastvideo/tests/attention/test_vsa_h3_sm100a_route.py index a2baaed84..904c3ca3e 100644 --- a/fastvideo/tests/attention/test_vsa_h3_sm100a_route.py +++ b/fastvideo/tests/attention/test_vsa_h3_sm100a_route.py @@ -28,6 +28,7 @@ def _build_meta( device: torch.device = torch.device("cpu"), prefix_segments: tuple[int, ...] = _SPEC["prefix_segments"], sparsity: float = 0.0, + dense_layers: tuple[int, ...] = (), builder: MiniMaxH3VSAMetadataBuilder | None = None, ): builder = builder or MiniMaxH3VSAMetadataBuilder() @@ -38,6 +39,7 @@ def _build_meta( VSA_sparsity=sparsity, prefix_segments=prefix_segments, device=device, + dense_layers=dense_layers, tile_size=64, ) @@ -56,8 +58,11 @@ class _FakeSm100a: def __init__(self, supported=True): self.supported = supported self.calls = [] + self.support_calls = [] + self.mask_calls = [] def is_supported(self, q, variable_block_sizes): + self.support_calls.append((q, variable_block_sizes)) return self.supported def block_sparse_attn_sm100a(self, q, k, v, q2k_idx, q2k_num, variable_block_sizes, need_lse=True): @@ -65,6 +70,10 @@ class _FakeSm100a: need_lse=need_lse)) return q.clone(), None + def block_sparse_attn_sm100a_from_mask(self, q, k, v, block_map, variable_block_sizes): + self.mask_calls.append(dict(q=q, block_map=block_map, vbs=variable_block_sizes)) + return q.clone(), None + def _fake_map_to_index(block_map): """Pure-torch stand-in for the Triton map_to_index (same contract).""" @@ -117,6 +126,188 @@ def test_reason_covers_every_precondition(): assert _sm100a_unavailable_reason(ok, q, vbs, grad_mode=False) is None +def test_prepare_for_regional_compile_resolves_supported_route(monkeypatch): + fake_sm = _FakeSm100a(supported=True) + monkeypatch.setattr(vsa_h3, "_sm100a", fake_sm) + monkeypatch.setenv(VSA_SM100A_ENV, "1") + impl = MiniMaxH3VSAImpl(num_heads=_HEADS, head_size=_DIM, causal=False, softmax_scale=_DIM**-0.5) + + unsupported = impl.prepare_for_regional_compile(torch.device("cpu")) + + assert unsupported is None + assert impl._regional_compile_sm100a_enabled is True + assert len(fake_sm.support_calls) == 1 + probe_q, probe_vbs = fake_sm.support_calls[0] + assert probe_q.shape == (1, 1, 128, _DIM) + assert probe_q.dtype == torch.bfloat16 + assert probe_vbs.dtype == torch.int32 + assert probe_vbs.tolist() == [64, 64] + assert impl._regional_compile_layer_idx is not None + assert impl._regional_compile_layer_idx.item() == -1 + + +def test_prepare_for_regional_compile_env_off_skips_probe(monkeypatch): + fake_sm = _FakeSm100a(supported=True) + monkeypatch.setattr(vsa_h3, "_sm100a", fake_sm) + monkeypatch.delenv(VSA_SM100A_ENV, raising=False) + impl = MiniMaxH3VSAImpl(num_heads=_HEADS, head_size=_DIM, causal=False, softmax_scale=_DIM**-0.5) + + unsupported = impl.prepare_for_regional_compile(torch.device("cpu")) + + assert unsupported is not None + assert VSA_SM100A_ENV in unsupported + assert impl._regional_compile_sm100a_enabled is False + assert impl._regional_compile_layer_idx is None + assert fake_sm.support_calls == [] + + +def test_prepare_for_regional_compile_requires_mask_entry(monkeypatch): + + class _IndexOnlySm100a: + + def is_supported(self, q, variable_block_sizes): + raise AssertionError("missing mask entry must be rejected before the support probe") + + monkeypatch.setattr(vsa_h3, "_sm100a", _IndexOnlySm100a()) + monkeypatch.setenv(VSA_SM100A_ENV, "1") + warnings = [] + monkeypatch.setattr(vsa_h3.logger, "warning_once", warnings.append) + impl = MiniMaxH3VSAImpl(num_heads=_HEADS, head_size=_DIM, causal=False, softmax_scale=_DIM**-0.5) + + unsupported = impl.prepare_for_regional_compile(torch.device("cpu")) + + assert unsupported is not None + assert impl._regional_compile_sm100a_enabled is False + assert impl._regional_compile_layer_idx is None + assert len(warnings) == 1 + assert "compatibility route" in warnings[0] + + +def test_mask_dispatch_uses_local_compatibility_route_for_older_kernel_wheel(monkeypatch): + class _RawOnlySm100a: + + def block_sparse_attn_sm100a(self, *args, **kwargs): + raise AssertionError("raw entry must remain behind the compatibility custom op") + + fake_sm = _RawOnlySm100a() + monkeypatch.setattr(vsa_h3, "_sm100a", fake_sm) + monkeypatch.setattr(vsa_h3, "map_to_index", lambda mask: (mask, mask)) + calls = [] + + def compat(q, k, v, block_map, variable_block_sizes): + calls.append((block_map, variable_block_sizes)) + return q + 1 + + monkeypatch.setattr(vsa_h3, "_h3_vsa_sm100a_from_mask_compat", compat) + q = torch.zeros(1, 1, 128, _DIM) + mask = torch.ones(1, 1, 2, 2, dtype=torch.bool) + vbs = torch.full((2, ), 64, dtype=torch.int32) + + out, lse = vsa_h3._sm100a_from_mask(q, q, q, mask, vbs) + + assert vsa_h3._sm100a_has_compile_safe_mask_route(fake_sm) + torch.testing.assert_close(out, q + 1) + assert lse is None and calls == [(mask, vbs)] + + +def test_prepared_route_fullgraph_avoids_eager_dispatch(monkeypatch): + """Capture must not revisit env/capability checks or raw map_to_index.""" + fake_sm = _FakeSm100a(supported=True) + monkeypatch.setattr(vsa_h3, "_sm100a", fake_sm) + monkeypatch.setenv(VSA_SM100A_ENV, "1") + monkeypatch.setattr(vsa_h3, "probe_enabled", lambda: None) + impl = MiniMaxH3VSAImpl(num_heads=_HEADS, head_size=_DIM, causal=False, softmax_scale=_DIM**-0.5) + impl.prepare_for_regional_compile(torch.device("cpu")) + + # Four total tiles (two prefix + two video) satisfy the real kernel's + # adjacent-query-block contract. + meta = MiniMaxH3VSAMetadataBuilder().build( + current_timestep=0, + raw_latent_shape=_SPEC["raw_latent_shape"], + patch_size=_SPEC["patch_size"], + VSA_sparsity=0.0, + prefix_segments=(64, 64), + device=torch.device("cpu"), + tile_size=64, + ) + q, k, v = _tiled_qkv(meta) + + def fail_eager_dispatch(*args, **kwargs): + raise AssertionError("compiled route re-entered eager dispatch") + + def compile_safe_from_mask(q, k, v, block_map, variable_block_sizes): + assert block_map.dtype == torch.bool + return q.clone(), None + + monkeypatch.setattr(fake_sm, "is_supported", fail_eager_dispatch) + monkeypatch.setattr(fake_sm, "block_sparse_attn_sm100a", fail_eager_dispatch) + monkeypatch.setattr(fake_sm, "block_sparse_attn_sm100a_from_mask", compile_safe_from_mask) + monkeypatch.setattr(vsa_h3, "map_to_index", fail_eager_dispatch) + monkeypatch.setattr(vsa_h3, "_sm100a_unavailable_reason", fail_eager_dispatch) + monkeypatch.setattr(vsa_h3, "block_sparse_attn_64_bhsd", fail_eager_dispatch) + + def run(q, k, v): + return impl.forward(q, k, v, None, meta) + + try: + compiled = torch.compile(run, backend="eager", fullgraph=True) + with torch.inference_mode(): + actual = compiled(q, k, v) + torch.testing.assert_close(actual, q, atol=0, rtol=0) + finally: + torch._dynamo.reset() + + +def test_prepared_route_reuses_graph_across_layer_indices_and_preserves_dense_overrides(monkeypatch): + """Fifty H3 blocks must not specialize the shared graph on layer_idx.""" + fake_sm = _FakeSm100a(supported=True) + monkeypatch.setattr(vsa_h3, "_sm100a", fake_sm) + monkeypatch.setenv(VSA_SM100A_ENV, "1") + monkeypatch.setattr(vsa_h3, "probe_enabled", lambda: None) + meta = _build_meta(sparsity=0.5, dense_layers=(0, 17), prefix_segments=(64, 64)) + assert meta.dense_layers_tensor.tolist() == [0, 17] + q, k, v = _tiled_qkv(meta) + + def fail_eager_dispatch(*args, **kwargs): + raise AssertionError("compiled route re-entered eager dispatch") + + def compile_safe_from_mask(q, k, v, block_map, variable_block_sizes): + del k, v, variable_block_sizes + # Encode the dense-layer decision in the result so this is a semantic + # test as well as a recompile-cache stress test. + return q + block_map.all().to(q.dtype), None + + implementations = [] + for layer_idx in range(20): + impl = MiniMaxH3VSAImpl( + num_heads=_HEADS, + head_size=_DIM, + causal=False, + softmax_scale=_DIM**-0.5, + prefix=f"transformer_blocks.{layer_idx}.attn", + ) + impl.prepare_for_regional_compile(torch.device("cpu")) + implementations.append(impl) + + monkeypatch.setattr(fake_sm, "is_supported", fail_eager_dispatch) + monkeypatch.setattr(fake_sm, "block_sparse_attn_sm100a", fail_eager_dispatch) + monkeypatch.setattr(fake_sm, "block_sparse_attn_sm100a_from_mask", compile_safe_from_mask) + monkeypatch.setattr(vsa_h3, "map_to_index", fail_eager_dispatch) + monkeypatch.setattr(vsa_h3, "_sm100a_unavailable_reason", fail_eager_dispatch) + monkeypatch.setattr(vsa_h3, "block_sparse_attn_64_bhsd", fail_eager_dispatch) + + torch._dynamo.reset() + try: + compiled = [torch.compile(impl.forward, backend="eager", fullgraph=True) for impl in implementations] + with torch.inference_mode(): + for layer_idx, run in enumerate(compiled): + actual = run(q, k, v, None, meta) + expected_delta = 1.0 if layer_idx in meta.dense_layers else 0.0 + torch.testing.assert_close(actual, q + expected_delta, atol=0, rtol=0) + finally: + torch._dynamo.reset() + + def test_default_off_routes_triton(routed, monkeypatch): fake_sm, fake_triton, run, _ = routed monkeypatch.delenv(VSA_SM100A_ENV, raising=False) @@ -125,6 +316,17 @@ def test_default_off_routes_triton(routed, monkeypatch): assert fake_sm.calls == [] +def test_unprepared_training_compile_keeps_existing_triton_route(routed, monkeypatch): + fake_sm, fake_triton, run, _ = routed + monkeypatch.setenv(VSA_SM100A_ENV, "1") + monkeypatch.setattr(torch.compiler, "is_compiling", lambda: True) + + run(requires_grad=True) + + assert fake_triton.calls == 1 + assert fake_sm.calls == [] + + def test_env_on_routes_sm100a_with_index_metadata(routed, monkeypatch): fake_sm, fake_triton, run, meta = routed monkeypatch.setenv(VSA_SM100A_ENV, "1") @@ -382,3 +584,38 @@ def test_real_sm100a_odd_route_matches_triton_oracle(monkeypatch): assert sm100a_tiled.shape[1] == triton_tiled.shape[1] + 64 assert sm100a_output.shape == triton_output.shape == (1, meta.total_seq_length, _HEADS, _DIM) torch.testing.assert_close(sm100a_output.float(), triton_output.float(), atol=0.04, rtol=0.02) + + +def test_real_sm100a_odd_preprocess_forward_postprocess_fullgraph(monkeypatch): + """Capture #1745's odd transport buffer mutation with real Inductor.""" + if not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 0): + pytest.skip("requires a GB200 (sm_100a) compute node") + if vsa_h3._sm100a is None or not vsa_h3._sm100a._HAS_VSA_SM100A: + pytest.skip("requires a fastvideo_kernel build containing sm_100a VSA") + + device = torch.device("cuda") + monkeypatch.setenv(VSA_SM100A_ENV, "1") + meta = _build_meta(device=device, sparsity=0.5) + assert meta.variable_block_sizes.numel() % 2 == 1 + impl = MiniMaxH3VSAImpl(num_heads=_HEADS, head_size=_DIM, causal=False, softmax_scale=_DIM**-0.5) + assert impl.prepare_for_regional_compile(device) is None + torch.manual_seed(17) + raw_qkv = torch.randn(3, meta.total_seq_length, _HEADS, _DIM, device=device, dtype=torch.bfloat16) + + def run(raw): + tiled = impl.preprocess_qkv(raw, meta) + query, key, value = tiled.chunk(3, dim=0) + output = impl.forward(query, key, value, None, meta) + return impl.postprocess_output(output, meta) + + torch._dynamo.reset() + try: + with torch.inference_mode(): + compiled = torch.compile(run, fullgraph=True) + output = compiled(raw_qkv) + torch.cuda.synchronize() + finally: + torch._dynamo.reset() + + assert output.shape == (1, meta.total_seq_length, _HEADS, _DIM) + assert torch.isfinite(output).all() diff --git a/fastvideo/tests/inference/test_basic_fasth3_profile.py b/fastvideo/tests/inference/test_basic_fasth3_profile.py new file mode 100644 index 000000000..e4f8bc250 --- /dev/null +++ b/fastvideo/tests/inference/test_basic_fasth3_profile.py @@ -0,0 +1,238 @@ +# SPDX-License-Identifier: Apache-2.0 +from __future__ import annotations + +import importlib.util +from pathlib import Path +from types import SimpleNamespace + +import pytest + +from fastvideo.models.schedulers.scheduling_minimax_h3 import MiniMaxH3Scheduler + + +REPO_ROOT = Path(__file__).resolve().parents[3] +EXAMPLE_PATH = REPO_ROOT / "examples" / "inference" / "basic" / "basic_fasth3.py" + + +def _load_example(): + spec = importlib.util.spec_from_file_location("basic_fasth3_contract", EXAMPLE_PATH) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +fasth3 = _load_example() + + +def _args(*overrides: str): + return fasth3.parse_args(["--prompt", "a test prompt", *overrides]) + + +def test_default_all_profile_matches_fastest_contract(tmp_path): + args = _args() + config = fasth3.build_generator_config(args) + request = fasth3.build_request(args, tmp_path / "result.mp4", args.seed) + environment = fasth3.profile_environment(args) + + assert args.profile == "all" + assert args.num_gpus == 4 + assert args.vsa_sparsity == 0.9 + assert args.vsa_tile_size == 64 + assert args.vsa_kernel == "sm100a" + assert args.seed == 1000 + assert args.warmup_seed == 999 + assert args.warmup is True + assert args.repeats == 3 + assert args.inference_torch_compile is True + assert args.ulysses_a2a == "off" + + assert config.engine.num_gpus == 4 + assert config.engine.parallelism.tp_size == 1 + assert config.engine.parallelism.sp_size == 4 + assert config.engine.use_fsdp_inference is False + assert config.engine.offload.dit is False + assert config.engine.offload.text_encoder is True + assert config.engine.offload.vae is True + assert config.engine.offload.pin_cpu_memory is True + assert config.engine.compile.enabled is False + assert config.engine.compile.vae_enabled is True + assert config.pipeline.experimental == { + "attention_backend": "VIDEO_SPARSE_ATTN_H3", + "VSA_sparsity": 0.9, + "VSA_tile_size": 64, + "inference_torch_compile": True, + "vae_parallel_decode": True, + "vae_parallel_decode_strategy": "gather", + } + + assert environment["FASTVIDEO_VSA_SM100A"] == "1" + assert environment["FASTVIDEO_VSA_CUTEDSL"] == "0" + assert environment["FASTVIDEO_DISABLE_ATTENTION_COMPILE"] == "0" + assert environment["FASTVIDEO_FA4"] == "1" + assert environment["FASTVIDEO_NVFP4_FA4"] == "0" + assert environment["FASTVIDEO_MINIMAX_H3_FA4_PACKED_VARLEN"] == "0" + assert environment["FASTVIDEO_MINIMAX_H3_FUSIONS"] == "all" + assert environment["FASTVIDEO_INFERENCE_TORCH_COMPILE"] == "1" + assert environment["FASTVIDEO_VAE_PARALLEL_DECODE"] == "1" + assert environment["FASTVIDEO_VAE_PARALLEL_DECODE_STRATEGY"] == "gather" + assert environment["FASTVIDEO_ULYSSES_A2A"] == "off" + assert environment["FASTVIDEO_H3_VSA_PROBE"] is None + + assert request.sampling.num_inference_steps == 5 + assert request.sampling.seed == 1000 + assert request.sampling.guidance_scale == 1.0 + assert request.sampling.batch_cfg is False + assert request.output.output_path == str(tmp_path / "result.mp4") + + +@pytest.mark.parametrize("num_frames", (124, 243, 345)) +def test_fast_profile_supports_measured_durations_without_separate_scripts(tmp_path, num_frames): + args = _args("--num-frames", str(num_frames)) + config = fasth3.build_generator_config(args) + request = fasth3.build_request(args, tmp_path / f"result_{num_frames}.mp4", args.seed) + + assert args.inference_torch_compile is True + assert args.ulysses_a2a == "off" + assert config.pipeline.experimental["inference_torch_compile"] is True + assert request.sampling.num_frames == num_frames + + +def test_compile_mode_requires_regional_compile_opt_out(capsys): + with pytest.raises(SystemExit): + _args("--compile-mode", "reduce-overhead") + assert "--compile-mode cannot be combined with regional compile" in capsys.readouterr().err + + args = _args("--compile-mode", "reduce-overhead", "--no-inference-torch-compile") + assert args.compile_mode == "reduce-overhead" + assert args.inference_torch_compile is False + + +def test_default_sigma_grid_runs_exactly_four_dit_forwards(): + scheduler = MiniMaxH3Scheduler(shift=12.0) + scheduler.set_timesteps(5) + + assert scheduler.sigmas is not None and len(scheduler.sigmas) == 5 + assert scheduler.timesteps is not None and len(scheduler.timesteps) == 4 + + +def test_strict_profile_changes_only_non_parity_fusions(): + all_args = _args("--profile", "all") + strict_args = _args("--profile", "strict") + all_environment = fasth3.profile_environment(all_args) + strict_environment = fasth3.profile_environment(strict_args) + + changed = {key for key in all_environment if all_environment[key] != strict_environment[key]} + assert changed == {"FASTVIDEO_MINIMAX_H3_FUSIONS"} + assert all_environment["FASTVIDEO_MINIMAX_H3_FUSIONS"] == "all" + assert strict_environment["FASTVIDEO_MINIMAX_H3_FUSIONS"] == "0" + assert fasth3.build_generator_config(all_args) == fasth3.build_generator_config(strict_args) + + +def test_opt_outs_override_inherited_environment(monkeypatch): + args = _args( + "--vsa-kernel", + "triton", + "--no-fa4", + "--no-h3-fusions", + "--no-compile-vae", + "--no-parallel-vae", + "--no-replicated-dit", + "--no-pin-cpu-memory", + "--no-inference-torch-compile", + ) + expected = fasth3.profile_environment(args) + for name in expected: + monkeypatch.setenv(name, "inherited-experiment-value") + + fasth3.configure_environment(args) + config = fasth3.build_generator_config(args) + + assert "FASTVIDEO_H3_VSA_PROBE" not in fasth3.os.environ + assert fasth3.os.environ["FASTVIDEO_VSA_SM100A"] == "0" + assert fasth3.os.environ["FASTVIDEO_VSA_CUTEDSL"] == "0" + assert fasth3.os.environ["FASTVIDEO_DISABLE_ATTENTION_COMPILE"] == "0" + assert fasth3.os.environ["FASTVIDEO_FA4"] == "0" + assert fasth3.os.environ["FASTVIDEO_NVFP4_FA4"] == "0" + assert fasth3.os.environ["FASTVIDEO_MINIMAX_H3_FUSIONS"] == "0" + assert fasth3.os.environ["FASTVIDEO_MINIMAX_H3_FA4_PACKED_VARLEN"] == "0" + assert fasth3.os.environ["FASTVIDEO_INFERENCE_TORCH_COMPILE"] == "0" + assert fasth3.os.environ["FASTVIDEO_VAE_PARALLEL_DECODE"] == "0" + assert fasth3.os.environ["FASTVIDEO_ULYSSES_A2A"] == "off" + assert config.engine.compile.vae_enabled is False + assert config.pipeline.experimental["vae_parallel_decode"] is False + assert config.engine.use_fsdp_inference is True + assert config.engine.offload.pin_cpu_memory is False + + +def test_selected_fast_profile_requires_its_optional_routes(monkeypatch): + monkeypatch.setattr(fasth3, "_fa4_is_installed", lambda: False) + monkeypatch.setattr(fasth3, "_sm100a_kernel_is_installed", lambda: False) + + with pytest.raises(RuntimeError, match="flash-attn-4"): + fasth3.validate_profile_dependencies(_args()) + with pytest.raises(RuntimeError, match="fastvideo-kernel 0.3.4"): + fasth3.validate_profile_dependencies(_args("--no-fa4")) + + fasth3.validate_profile_dependencies(_args("--no-fa4", "--vsa-kernel", "triton")) + + monkeypatch.setattr(fasth3, "_fa4_is_installed", lambda: True) + monkeypatch.setattr(fasth3, "_sm100a_kernel_is_installed", lambda: True) + fasth3.validate_profile_dependencies(_args()) + + +def test_run_excludes_warmup_and_uses_distinct_measured_outputs(monkeypatch, tmp_path, capsys): + calls = [] + + class FakeGenerator: + shutdown_called = False + + def generate(self, request): + calls.append(request) + Path(request.output.output_path).touch() + return SimpleNamespace( + video_path=request.output.output_path, + generation_time=1.25, + logging_info=SimpleNamespace(stages={"denoising": {"execution_time": 2.5}}), + ) + + def shutdown(self): + self.shutdown_called = True + + fake_generator = FakeGenerator() + + class FakeVideoGenerator: + + @classmethod + def from_config(cls, config): + assert config == fasth3.build_generator_config(args) + return fake_generator + + for name in fasth3.profile_environment(_args()): + monkeypatch.setenv(name, "inherited-experiment-value") + monkeypatch.setattr(fasth3, "validate_profile_dependencies", lambda args: None) + monkeypatch.setattr(fasth3, "VideoGenerator", FakeVideoGenerator) + clock = iter((10.0, 15.0, 20.0, 26.0, 30.0, 37.0, 40.0, 48.0)) + monkeypatch.setattr(fasth3.time, "perf_counter", lambda: next(clock)) + args = _args("--output", str(tmp_path)) + + measured = fasth3.run(args) + + assert measured == [6.0, 7.0, 8.0] + assert [request.sampling.seed for request in calls] == [999, 1000, 1000, 1000] + assert [Path(request.output.output_path).name for request in calls] == [ + "_fasth3_warmup.mp4", + "fasth3_all_run_01.mp4", + "fasth3_all_run_02.mp4", + "fasth3_all_run_03.mp4", + ] + assert (tmp_path / "_fasth3_warmup.mp4").is_file() + assert fake_generator.shutdown_called is True + output = capsys.readouterr().out + assert "[warmup] wall=5.000s (excluded)" in output + assert f"Warmup output written to: {tmp_path / '_fasth3_warmup.mp4'}" in output + assert output.count("Output written to:") == 3 + assert output.count("Denoising time: 2.500s") == 3 + assert "Measured E2E wall times (n=3, warmup excluded): [6.0, 7.0, 8.0]" in output + assert "Median E2E wall time: 7.000s" in output + assert "Median denoising time: 2.500s" in output diff --git a/fastvideo/tests/inference/test_inference_regional_compile.py b/fastvideo/tests/inference/test_inference_regional_compile.py index f827d74e1..ffb5d0b4b 100644 --- a/fastvideo/tests/inference/test_inference_regional_compile.py +++ b/fastvideo/tests/inference/test_inference_regional_compile.py @@ -6,9 +6,9 @@ transformer loads (``FastVideoArgs.inference_torch_compile``, env ``FASTVIDEO_INFERENCE_TORCH_COMPILE=1``). These tests pin the two pieces that must not drift from the #1718 training-port semantics: -- ``_regional_compile_unsupported_reason``: VSA backends (and the attention - eager escape hatch) degrade to eager with a reason instead of hard-failing - fullgraph capture at the first denoising forward. +- ``_regional_compile_unsupported_reason``: legacy VSA (and the attention + eager escape hatch) degrades to eager, while prepared MiniMax-H3 VSA is + admitted to regional fullgraph capture. - attention forward dispatch: ordinary instances retain the historical compiler-disabled boundary; regional compile opts in only the selected model's instances. @@ -30,10 +30,12 @@ from torch import nn from fastvideo.attention import layer as attention_layer from fastvideo.attention.layer import DistributedAttention +from fastvideo.models.dits.minimax_h3 import MiniMaxH3Attention, MiniMaxH3Transformer3DModel from fastvideo.models.loader import fsdp_load from fastvideo.models.loader.fsdp_load import ( _compile_model_regions, _enable_regional_attention_compile, + _prepare_model_for_compile, _regional_compile_unsupported_reason, ) @@ -43,9 +45,9 @@ def _init_params_for(backend_name: str | None) -> dict: return {"config": SimpleNamespace(_resolved_attention_backend=resolved)} -@pytest.mark.parametrize("backend_name", ["VIDEO_SPARSE_ATTN", "VIDEO_SPARSE_ATTN_H3"]) -def test_vsa_backends_degrade_to_eager(backend_name, monkeypatch) -> None: +def test_legacy_vsa_backend_degrades_to_eager(monkeypatch) -> None: monkeypatch.delenv("FASTVIDEO_DISABLE_ATTENTION_COMPILE", raising=False) + backend_name = "VIDEO_SPARSE_ATTN" reason = _regional_compile_unsupported_reason(_init_params_for(backend_name)) assert reason is not None assert backend_name in reason @@ -53,11 +55,193 @@ def test_vsa_backends_degrade_to_eager(backend_name, monkeypatch) -> None: @pytest.mark.parametrize("backend_name", [None, "TORCH_SDPA"]) -def test_dense_backends_allow_compile(backend_name, monkeypatch) -> None: +def test_supported_backends_allow_compile(backend_name, monkeypatch) -> None: monkeypatch.delenv("FASTVIDEO_DISABLE_ATTENTION_COMPILE", raising=False) + monkeypatch.delenv("FASTVIDEO_H3_VSA_PROBE", raising=False) assert _regional_compile_unsupported_reason(_init_params_for(backend_name)) is None +def test_h3_vsa_sm100a_tile64_allows_compile(monkeypatch) -> None: + monkeypatch.delenv("FASTVIDEO_DISABLE_ATTENTION_COMPILE", raising=False) + monkeypatch.delenv("FASTVIDEO_H3_VSA_PROBE", raising=False) + monkeypatch.setenv("FASTVIDEO_VSA_SM100A", "1") + + reason = _regional_compile_unsupported_reason( + _init_params_for("VIDEO_SPARSE_ATTN_H3"), + vsa_tile_size=64, + ) + + assert reason is None + + +@pytest.mark.parametrize(("sm100a", "tile_size"), [(False, 64), (True, 256), (True, None)]) +def test_h3_vsa_unsupported_compile_route_degrades_to_eager(sm100a, tile_size, monkeypatch) -> None: + monkeypatch.delenv("FASTVIDEO_DISABLE_ATTENTION_COMPILE", raising=False) + monkeypatch.delenv("FASTVIDEO_H3_VSA_PROBE", raising=False) + if sm100a: + monkeypatch.setenv("FASTVIDEO_VSA_SM100A", "1") + else: + monkeypatch.delenv("FASTVIDEO_VSA_SM100A", raising=False) + + reason = _regional_compile_unsupported_reason( + _init_params_for("VIDEO_SPARSE_ATTN_H3"), + vsa_tile_size=tile_size, + ) + + assert reason is not None + assert "eager" in reason + + +def test_h3_vsa_probe_degrades_regional_compile_to_eager(monkeypatch) -> None: + monkeypatch.delenv("FASTVIDEO_DISABLE_ATTENTION_COMPILE", raising=False) + monkeypatch.setenv("FASTVIDEO_H3_VSA_PROBE", "/tmp/h3-vsa-probe") + + reason = _regional_compile_unsupported_reason(_init_params_for("VIDEO_SPARSE_ATTN_H3")) + + assert reason is not None + assert "FASTVIDEO_H3_VSA_PROBE" in reason + assert "eager" in reason + + +class _RegionalPrepareProbe: + + def __init__(self, unsupported: str | None = None) -> None: + self.devices: list[torch.device] = [] + self.unsupported = unsupported + + def prepare_for_regional_compile(self, device: torch.device) -> str | None: + self.devices.append(device) + return self.unsupported + + +class _GateBlock(nn.Module): + + def __init__(self, gate_weight: float) -> None: + super().__init__() + attention = MiniMaxH3Attention.__new__(MiniMaxH3Attention) + nn.Module.__init__(attention) + attention.to_q = nn.Linear(1, 1, bias=False) + attention.to_gate_compress = nn.Linear(1, 1, bias=False) + with torch.no_grad(): + attention.to_gate_compress.weight.fill_(gate_weight) + attention._gate_compress_active = None + attention.distributed_attention = SimpleNamespace(attn_impl=_RegionalPrepareProbe()) + self.attn = attention + + +class _GateProbe(nn.Module): + + def __init__(self, attention: MiniMaxH3Attention) -> None: + super().__init__() + self.attention = attention + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if self.attention._gate_active(): + return x + 1 + return x - 1 + + +def _gate_model(*weights: float) -> MiniMaxH3Transformer3DModel: + model = MiniMaxH3Transformer3DModel.__new__(MiniMaxH3Transformer3DModel) + nn.Module.__init__(model) + model.transformer_blocks = nn.ModuleList([_GateBlock(weight) for weight in weights]) + model.enabled_fusions = frozenset() + return model + + +def test_minimax_h3_prepare_for_compile_resolves_loaded_vsa_gates() -> None: + model = _gate_model(0.0, 2.0) + + model.prepare_for_compile() + + assert [block.attn._gate_compress_active for block in model.transformer_blocks] == [False, True] + for block in model.transformer_blocks: + impl = block.attn.distributed_attention.attn_impl + assert impl.devices == [] + + +def test_training_compile_prepare_does_not_probe_inference_kernel() -> None: + model = _gate_model(2.0) + + reason = _prepare_model_for_compile(model, regional=False) + + assert reason is None + attention = model.transformer_blocks[0].attn + assert attention._gate_compress_active is True + assert attention.distributed_attention.attn_impl.devices == [] + + +def test_regional_compile_prepare_prefers_specialized_hook() -> None: + model = _gate_model(2.0) + expected_device = next(model.parameters()).device + + reason = _prepare_model_for_compile(model, regional=True) + + assert reason is None + impl = model.transformer_blocks[0].attn.distributed_attention.attn_impl + assert impl.devices == [expected_device] + + +def test_minimax_h3_prepare_for_regional_compile_does_not_require_quantized_q_weight() -> None: + model = _gate_model(2.0) + attention = model.transformer_blocks[0].attn + removed_weight = attention.to_q._parameters.pop("weight") + attention.to_q.register_buffer("_fp8_weight", removed_weight.detach(), persistent=False) + expected_device = next(model.parameters()).device + + reason = model.prepare_for_regional_compile() + + assert reason is None + impl = attention.distributed_attention.attn_impl + assert impl.devices == [expected_device] + + +def test_minimax_h3_prepare_for_regional_compile_propagates_backend_rejection() -> None: + model = _gate_model(2.0) + impl = model.transformer_blocks[0].attn.distributed_attention.attn_impl + impl.unsupported = "sm100a probe failed" + + reason = model.prepare_for_regional_compile() + + assert reason == "sm100a probe failed" + + +def test_unprepared_vsa_gate_cannot_mutate_cache_during_compile(monkeypatch) -> None: + attention = _gate_model(2.0).transformer_blocks[0].attn + monkeypatch.setattr(torch.compiler, "is_compiling", lambda: True) + + with torch.no_grad(), pytest.raises(RuntimeError, match="not resolved before torch.compile"): + attention._gate_active() + assert attention._gate_compress_active is None + + +@pytest.mark.parametrize(("gate_weight", "expected"), [(0.0, -1.0), (2.0, 1.0)]) +def test_prepared_vsa_gate_is_static_under_fullgraph_compile(gate_weight, expected) -> None: + model = _gate_model(gate_weight) + model.prepare_for_compile() + attention = model.transformer_blocks[0].attn + resolved = attention._gate_compress_active + + try: + compiled = torch.compile(_GateProbe(attention), backend="eager", fullgraph=True) + with torch.no_grad(): + actual = compiled(torch.zeros(1)) + torch.testing.assert_close(actual, torch.full((1, ), expected)) + assert attention._gate_compress_active is resolved + finally: + torch._dynamo.reset() + + +def test_prepared_zero_vsa_gate_still_runs_in_grad_enabled_training() -> None: + model = _gate_model(0.0) + model.prepare_for_compile() + attention = model.transformer_blocks[0].attn + + with torch.enable_grad(): + assert attention._gate_active() is True + assert attention._gate_compress_active is False + + def test_attention_compile_escape_hatch_degrades_to_eager(monkeypatch) -> None: monkeypatch.setenv("FASTVIDEO_DISABLE_ATTENTION_COMPILE", "1") reason = _regional_compile_unsupported_reason(_init_params_for("TORCH_SDPA")) diff --git a/fastvideo/tests/transformers/test_minimax_h3_fusion_compile.py b/fastvideo/tests/transformers/test_minimax_h3_fusion_compile.py new file mode 100644 index 000000000..d6e8c1679 --- /dev/null +++ b/fastvideo/tests/transformers/test_minimax_h3_fusion_compile.py @@ -0,0 +1,159 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Fullgraph trace contract for the MiniMax-H3 inference fusions.""" + +from __future__ import annotations + +import pytest +import torch +from torch._subclasses.fake_tensor import FakeTensorMode + +from fastvideo.models.dits.minimax_h3_fusions import ( + fused_qknorm_rope, + fused_residual_gate_rmsnorm_modulate, + fused_rmsnorm_modulate, + minimax_h3_swiglu, +) + + +def _all_fusions( + x: torch.Tensor, + branch: torch.Tensor, + weight: torch.Tensor, + scale: torch.Tensor, + shift: torch.Tensor, + gate: torch.Tensor, + index: torch.Tensor, + q: torch.Tensor, + q_weight: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + packed_mlp: torch.Tensor, +) -> tuple[torch.Tensor, ...]: + norm1 = fused_rmsnorm_modulate(x, weight, scale, shift, index, 1e-5) + hidden, norm2 = fused_residual_gate_rmsnorm_modulate( + x, + branch, + gate, + weight, + scale, + shift, + index, + 1e-5, + ) + q_out = fused_qknorm_rope(q, q_weight, cos, sin, 1e-5) + mlp_out = minimax_h3_swiglu(packed_mlp) + return norm1, hidden, norm2, q_out, mlp_out + + +def test_minimax_h3_fusions_are_dynamic_fullgraph_traceable() -> None: + """Dynamo sees four opaque ops without requiring a physical CUDA GPU.""" + graphs: list[torch.fx.GraphModule] = [] + + def capture_graph(graph: torch.fx.GraphModule, example_inputs: list[torch.Tensor]): + del example_inputs + graphs.append(graph) + return graph.forward + + compiled = torch.compile(_all_fusions, backend=capture_graph, fullgraph=True, dynamic=True) + + with torch.no_grad(), FakeTensorMode(): + x = torch.empty((2, 3, 8), device="cuda", dtype=torch.bfloat16) + weight = torch.empty((8,), device="cuda", dtype=torch.bfloat16) + table = torch.empty((6, 8), device="cuda", dtype=torch.bfloat16) + index = torch.empty((3,), device="cuda", dtype=torch.int64) + q = torch.empty((2, 3, 4, 8), device="cuda", dtype=torch.bfloat16) + rotary = torch.empty((3, 8), device="cuda", dtype=torch.bfloat16) + packed_mlp = torch.empty((2, 3, 16), device="cuda", dtype=torch.bfloat16) + + outputs = compiled( + x, + x, + weight, + table, + table, + table, + index, + q, + weight, + rotary, + rotary, + packed_mlp, + ) + + assert tuple(output.shape for output in outputs) == ( + (2, 3, 8), + (2, 3, 8), + (2, 3, 8), + (2, 3, 4, 8), + (2, 3, 8), + ) + assert all(output.is_contiguous() for output in outputs) + assert len(graphs) == 1 + call_targets = {node.target for node in graphs[0].graph.nodes if node.op == "call_function"} + assert { + torch.ops.fastvideo._minimax_h3_rmsnorm_modulate, + torch.ops.fastvideo._minimax_h3_residual_gate_rmsnorm_modulate, + torch.ops.fastvideo._minimax_h3_qknorm_rope, + torch.ops.fastvideo._minimax_h3_swiglu, + } <= call_targets + + +@pytest.mark.gpu +def test_minimax_h3_fusion_custom_ops_match_real_fullgraph_cuda() -> None: + """Cross-check fake metadata and execute the real CUDA dispatcher path.""" + if not torch.cuda.is_available() or not torch.cuda.is_bf16_supported(): + pytest.skip("BF16 CUDA is required") + pytest.importorskip("triton") + + device = torch.device("cuda") + generator = torch.Generator(device=device).manual_seed(2026) + x = torch.randn((2, 3, 128), generator=generator, device=device, dtype=torch.bfloat16) + branch = torch.randn_like(x) + weight = torch.randn((128,), generator=generator, device=device, dtype=torch.bfloat16) + scale = torch.randn((6, 128), generator=generator, device=device, dtype=torch.bfloat16) + shift = torch.randn((6, 128), generator=generator, device=device, dtype=torch.bfloat16) + gate = torch.randn((6, 128), generator=generator, device=device, dtype=torch.bfloat16) + index = torch.tensor((0, 3, 5), device=device, dtype=torch.int64) + q = torch.randn((2, 3, 4, 128), generator=generator, device=device, dtype=torch.bfloat16) + q_weight = torch.randn((128,), generator=generator, device=device, dtype=torch.bfloat16).contiguous() + rotary = torch.randn((3, 96), generator=generator, device=device, dtype=torch.bfloat16).contiguous() + packed_mlp = torch.randn((2, 3, 256), generator=generator, device=device, dtype=torch.bfloat16) + + opchecks = ( + ( + torch.ops.fastvideo._minimax_h3_rmsnorm_modulate, + (x, weight, scale, shift, index, 1e-5), + ), + ( + torch.ops.fastvideo._minimax_h3_residual_gate_rmsnorm_modulate, + (x, branch, gate, weight, scale, shift, index, 1e-5), + ), + ( + torch.ops.fastvideo._minimax_h3_qknorm_rope, + (q, q_weight, rotary, rotary.clone(), 1e-5), + ), + (torch.ops.fastvideo._minimax_h3_swiglu, (packed_mlp,)), + ) + for op, args in opchecks: + torch.library.opcheck(op, args) + + compiled = torch.compile(_all_fusions, fullgraph=True, dynamic=True) + inputs = ( + x, + branch, + weight, + scale, + shift, + gate, + index, + q, + q_weight, + rotary, + rotary.clone(), + packed_mlp, + ) + with torch.no_grad(): + expected = _all_fusions(*inputs) + actual = compiled(*inputs) + for actual_output, expected_output in zip(actual, expected, strict=True): + torch.testing.assert_close(actual_output, expected_output, atol=0.0, rtol=0.0) diff --git a/pyproject.toml b/pyproject.toml index 95eeb9b14..b22486f04 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -81,7 +81,11 @@ dependencies = [ # (MPS) and other non-Linux platforms don't try to build it from source; # every fastvideo_kernel import elsewhere is already guarded to degrade # gracefully when the package is absent. - "fastvideo-kernel==0.3.2; sys_platform == 'linux'", + # Keep ordinary installs compatible with the currently published wheel + # while the in-repo FastH3 extra selects this checkout's kernel source. + # Fresh registry installs select 0.3.4 once the version-bump publisher + # completes after merge. + "fastvideo-kernel>=0.3.2,<0.3.5; sys_platform == 'linux'", "wheel", # Training Dependencies @@ -108,12 +112,15 @@ members = ["apps/dreamverse"] [tool.uv.sources] # PyPI's aarch64 fastvideo-kernel wheels do not cover the NVIDIA DGX Spark / -# GB10 (sm_121), so on ARM build from the in-tree source instead; x86_64 keeps -# the published PyPI wheel. This table is uv-only: plain pip ignores it. Pass -# --no-sources to bypass it (e.g. when the kernel is already installed, as in -# the Spark install guide). +# GB10 (sm_121), so on ARM build from the in-tree source instead. The FastH3 +# extra also selects the in-tree source on x86_64 so a PR checkout can build +# its matching sm_100a route before that version reaches PyPI. Other x86_64 +# installs keep the published wheel. This table is uv-only: plain pip ignores +# it. Pass --no-sources to bypass it (e.g. when the kernel is already +# installed, as in the Spark install guide). fastvideo-kernel = [ { path = "fastvideo-kernel", marker = "platform_machine == 'aarch64'" }, + { path = "fastvideo-kernel", marker = "platform_machine == 'x86_64'", extra = "fasth3" }, ] # torch / torchvision / torchaudio are intentionally NOT pinned to a PyTorch # index here. Pick the supported build explicitly at install time: @@ -174,6 +181,14 @@ eval-full = ["fastvideo[eval]", "fastvideo[eval-audio]", "qwen-omni-utils"] dev = [ "fastvideo[lint]", "fastvideo[test]", ] +# Dependencies required by the measured FastH3 Preview profile. The pinned +# flash-attn-4 source above keeps the CuTe DSL stack reproducible for uv +# installs without putting a direct URL in published package metadata. +fasth3 = [ + "flash-attn-4", + "fastvideo-kernel>=0.3.2,<0.3.5; sys_platform == 'linux'", +] + prompt-safety = [ "fasttext", ] diff --git a/tests/test_block_sparse_sm100a.py b/tests/test_block_sparse_sm100a.py index 6f90f4017..479a59aff 100644 --- a/tests/test_block_sparse_sm100a.py +++ b/tests/test_block_sparse_sm100a.py @@ -49,6 +49,15 @@ def make_case(block, num_blocks=8, topk=4, heads=4, batch=1, ragged=False, seed= return q, k, v, idx, num, vbs +def make_block_map(idx, num, batch, heads, num_blocks): + """Expand the test's compact rows to the bool map H3 produces.""" + block_map = torch.zeros((batch, heads, num_blocks, num_blocks), dtype=torch.bool, device=idx.device) + rows = block_map.view(-1, num_blocks) + for row in range(rows.shape[0]): + rows[row, idx[row, :int(num[row])].long()] = True + return block_map + + def reference(q, k, v, idx, num, vbs, block): """Dense attention restricted to the selected blocks, with padded keys masked.""" if not vsa.BHSD: @@ -123,6 +132,66 @@ def test_lse_is_not_vacuous(block): assert (got_lse.float() - (ref_lse + 1.0)).abs().max().item() > 0.5 +@pytest.mark.parametrize("block", BLOCK_SIZES) +def test_inference_no_lse_matches_eager_under_fullgraph_compile(block): + """The production no-LSE route is one opaque, fake-backed graph node.""" + q, k, v, idx, num, vbs = make_case(block, ragged=True) + + def inference_path(q, k, v, idx, num, vbs): + out, lse = vsa.block_sparse_attn_sm100a(q, k, v, idx, num, vbs, need_lse=False) + assert lse is None + return out + + with torch.inference_mode(): + expected = inference_path(q, k, v, idx, num, vbs) + compiled = torch.compile(inference_path, backend="eager", fullgraph=True) + actual = compiled(q, k, v, idx, num, vbs) + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + + +@pytest.mark.parametrize("block", BLOCK_SIZES) +def test_inference_custom_op_fake_matches_real(block): + """Pin the custom op schema and fake output metadata used by Dynamo.""" + q, k, v, idx, num, vbs = make_case(block, ragged=True) + torch.library.opcheck( + torch.ops.fastvideo_kernel.block_sparse_attn_sm100a_inference.default, + (q, k, v, idx, num, vbs), + test_utils=("test_schema", "test_faketensor", "test_aot_dispatch_dynamic"), + ) + + +@pytest.mark.parametrize("block", BLOCK_SIZES) +def test_mask_inference_matches_index_route_under_fullgraph_compile(block): + """H3's bool mask compaction and sm_100a launch stay in one opaque op.""" + batch, heads, num_blocks = 1, 4, 8 + q, k, v, idx, num, vbs = make_case(block, batch=batch, heads=heads, num_blocks=num_blocks, ragged=True) + block_map = make_block_map(idx, num, batch, heads, num_blocks) + + def inference_path(q, k, v, block_map, vbs): + out, lse = vsa.block_sparse_attn_sm100a_from_mask(q, k, v, block_map, vbs) + assert lse is None + return out + + with torch.inference_mode(): + expected, _ = vsa.block_sparse_attn_sm100a(q, k, v, idx, num, vbs, need_lse=False) + compiled = torch.compile(inference_path, backend="eager", fullgraph=True) + actual = compiled(q, k, v, block_map, vbs) + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + + +@pytest.mark.parametrize("block", BLOCK_SIZES) +def test_mask_inference_custom_op_fake_matches_real(block): + """Pin schema/fake metadata for H3's mask-taking opaque boundary.""" + batch, heads, num_blocks = 1, 4, 8 + q, k, v, idx, num, vbs = make_case(block, batch=batch, heads=heads, num_blocks=num_blocks, ragged=True) + block_map = make_block_map(idx, num, batch, heads, num_blocks) + torch.library.opcheck( + torch.ops.fastvideo_kernel.block_sparse_attn_sm100a_from_mask_inference.default, + (q, k, v, block_map, vbs), + test_utils=("test_schema", "test_faketensor", "test_aot_dispatch_dynamic"), + ) + + def test_unsupported_is_rejected(): q, _, _, _, _, vbs = make_case(64) assert not vsa.is_supported(q.float(), vbs) # wrong dtype diff --git a/tests/test_fasth3_packaging.py b/tests/test_fasth3_packaging.py new file mode 100644 index 000000000..6162ffbc8 --- /dev/null +++ b/tests/test_fasth3_packaging.py @@ -0,0 +1,38 @@ +# SPDX-License-Identifier: Apache-2.0 +from __future__ import annotations + +import tomllib +from pathlib import Path + +from packaging.requirements import Requirement +from packaging.version import Version + + +REPO_ROOT = Path(__file__).resolve().parents[1] + + +def test_fasth3_extra_and_root_kernel_pin_match_source_release(): + root_project = tomllib.loads((REPO_ROOT / "pyproject.toml").read_text(encoding="utf-8")) + kernel_project = tomllib.loads((REPO_ROOT / "fastvideo-kernel" / "pyproject.toml").read_text(encoding="utf-8")) + kernel_version = kernel_project["project"]["version"] + + dependencies = root_project["project"]["dependencies"] + kernel_requirement = next(Requirement(value) for value in dependencies if value.startswith("fastvideo-kernel")) + assert Version(kernel_version) in kernel_requirement.specifier + fasth3_extra = root_project["project"]["optional-dependencies"]["fasth3"] + assert "flash-attn-4" in fasth3_extra + assert any(value.startswith("fastvideo-kernel") for value in fasth3_extra) + kernel_sources = root_project["tool"]["uv"]["sources"]["fastvideo-kernel"] + assert { + "path": "fastvideo-kernel", + "marker": "platform_machine == 'x86_64'", + "extra": "fasth3", + } in kernel_sources + + +def test_kernel_release_matrix_can_publish_sm100a_wheels(): + workflow = (REPO_ROOT / ".github" / "workflows" / "publish-kernel.yml").read_text(encoding="utf-8") + + assert 'TORCH_CUDA_ARCH_LIST="9.0a;10.0a;12.0a"' in workflow + assert 'TORCH_CUDA_ARCH_LIST="10.0a;12.0a"' in workflow + assert "patchelf==0.17.2.4" in workflow