[perf] Add opt-in regional fullgraph compile for DiT inference (#1741)

This commit is contained in:
William Lin
2026-08-22 12:14:39 -07:00
committed by GitHub
parent 2f3d407406
commit d3cff517cd
10 changed files with 482 additions and 9 deletions
+8
View File
@@ -85,6 +85,12 @@ def parse_args() -> argparse.Namespace:
parser.add_argument("--compile-mode",
default=None,
help='torch.compile mode, e.g. "reduce-overhead" for CUDA graphs')
parser.add_argument("--inference-torch-compile",
action="store_true",
help="regional fullgraph torch.compile of each DiT block after load. 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,
@@ -119,6 +125,8 @@ def main() -> None:
}
if args.vsa_sparsity > 0.0:
experimental["VSA_sparsity"] = args.vsa_sparsity
if args.inference_torch_compile:
experimental["inference_torch_compile"] = True
generator = VideoGenerator.from_config(
GeneratorConfig(
@@ -15,6 +15,7 @@ from fastvideo.api import (
OffloadConfig,
OutputConfig,
ParallelismConfig,
PipelineSelection,
SamplingConfig,
)
@@ -41,6 +42,12 @@ def parse_args() -> argparse.Namespace:
parser.add_argument("--compile-mode",
default=None,
help='torch.compile mode, e.g. "reduce-overhead" for CUDA graphs')
parser.add_argument("--inference-torch-compile",
action="store_true",
help="regional fullgraph torch.compile of each DiT block after load (the #1718 "
"training-port semantics: no kwargs; fullgraph + emulate_precision_casts injected). "
"First generation pays the inductor JIT (~1-2 min); use --repeats >= 2 and time "
"the last repeat. FASTVIDEO_INFERENCE_TORCH_COMPILE=1 is equivalent")
parser.add_argument("--repeats",
type=int,
default=1,
@@ -54,9 +61,16 @@ def main() -> None:
output_dir = Path(args.output)
output_dir.mkdir(parents=True, exist_ok=True)
# Boot-time run configuration folded into FastVideoArgs (the same
# experimental-dict route basic_fasth3.py uses for the VSA knobs).
experimental: dict[str, object] = {}
if args.inference_torch_compile:
experimental["inference_torch_compile"] = True
generator = VideoGenerator.from_config(
GeneratorConfig(
model_path=args.model_path,
pipeline=PipelineSelection(experimental=experimental),
engine=EngineConfig(
num_gpus=args.num_gpus,
use_fsdp_inference=args.num_gpus > 1,