From d3cff517cd5b3d824ec7590e78b28624244afbcc Mon Sep 17 00:00:00 2001 From: William Lin Date: Sat, 22 Aug 2026 14:14:39 -0500 Subject: [PATCH] [perf] Add opt-in regional fullgraph compile for DiT inference (#1741) --- docs/inference/optimizations.md | 53 ++++- examples/inference/basic/basic_fasth3.py | 8 + .../inference/basic/basic_minimax_h3_t2v.py | 14 ++ fastvideo/attention/layer.py | 39 +++- fastvideo/envs.py | 12 + fastvideo/fastvideo_args.py | 29 +++ fastvideo/models/loader/component_loader.py | 1 + fastvideo/models/loader/fsdp_load.py | 118 ++++++++++ fastvideo/pipelines/composed_pipeline_base.py | 8 + .../test_inference_regional_compile.py | 209 ++++++++++++++++++ 10 files changed, 482 insertions(+), 9 deletions(-) create mode 100644 fastvideo/tests/inference/test_inference_regional_compile.py diff --git a/docs/inference/optimizations.md b/docs/inference/optimizations.md index 1ec8730e7..35dde2a60 100644 --- a/docs/inference/optimizations.md +++ b/docs/inference/optimizations.md @@ -337,7 +337,38 @@ Only DiT submodules that declare `_compile_conditions` are compiled (most shipped models). The text encoder and VAE are not compiled by this flag. -### What to expect +### Regional fullgraph compile (experimental) + +`inference_torch_compile` is a stricter, kwargs-free variant that ports the +training-side regional compile of +[#1718](https://github.com/hao-ai-lab/FastVideo/pull/1718) to inference: the +loader wraps each `_compile_conditions` block in +`torch.compile(fullgraph=True)` with inductor +`options={"emulate_precision_casts": True}` right after the transformer +loads. The ordinary compile path keeps its historical compiler-disabled +attention boundary by default; regional compile opts in only the compatible +attention instances owned by this transformer. VSA and the explicit +`FASTVIDEO_DISABLE_ATTENTION_COMPILE=1` escape hatch keep +the transformer eager with one warning instead of failing mid-denoise. + +```python +generator = VideoGenerator.from_pretrained( + "MiniMaxAI/MiniMax-H3", + inference_torch_compile=True, # or FASTVIDEO_INFERENCE_TORCH_COMPILE=1 +) +``` + +Do not combine it with `torch_compile_kwargs['mode']` (the loader injects +inductor options, and torch.compile forbids mode+options); it is +independent of `enable_torch_compile`, and when both are set the regional +compile wins for the DiT. + +### What to expect from generic compile + +The Wan result below measures the existing generic +`enable_torch_compile=True` path. It is useful evidence that compile can help, +but it is **not** a benchmark or numerical gate for the stricter regional +fullgraph path above. | Config | Effect | |---|---| @@ -353,6 +384,16 @@ generations with the same input shapes. Always exclude the first (warmup) generation when measuring steady-state latency — measuring the warmup is the most common way to wrongly conclude "compile is slower". +**Regional MiniMax-H3 accuracy caveat (job 2660).** On one GB200 at +768×1344×124, the native 50-point schedule ran exactly 49 transformer +forwards. After one warmup, three fixed-prompt/fixed-seed repeats averaged +**185.08s → 157.01s end to end** and **174.90s → 147.12s denoising**. Each +leg was independently pixel-deterministic, but compiled output did **not** +match eager: mean absolute pixel error **20.247/255**, PSNR **16.67 dB**, +mean SSIM **0.7108**, and mean MS-SSIM **0.6370** across 124 frames. Treat +regional MiniMax-H3 compile as an opt-in performance experiment, not an +eager-parity-safe mode. + **Numerics.** Inductor's lowering is designed to preserve eager semantics within floating-point tolerance, but per-model equivalence is not asserted by any standing SSIM regression here — the SSIM tests in @@ -371,9 +412,13 @@ MS-SSIM gate on *your* config, especially when combining hao-ai-lab/FastVideo#1365 — keep that fix to get a clean compiled region under the default offload path. - **`mode="reduce-overhead"` / CUDA graphs**: not yet supported - end-to-end. The attention dispatch is an untraceable custom op and - still breaks the graph, which CUDA-graph trees cannot span. Use the - default inductor mode (shown above) until that is resolved. + by regional compile because that path injects inductor `options`, and + PyTorch rejects `mode` together with `options`. The generic compile path + can accept `mode`, but CUDA-graph compatibility remains backend- and + shape-dependent. FA2/FA3 inference and FA4 expose traceable custom-op + boundaries; 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. 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/basic_fasth3.py b/examples/inference/basic/basic_fasth3.py index 178fef26c..5badf3a46 100644 --- a/examples/inference/basic/basic_fasth3.py +++ b/examples/inference/basic/basic_fasth3.py @@ -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( diff --git a/examples/inference/basic/basic_minimax_h3_t2v.py b/examples/inference/basic/basic_minimax_h3_t2v.py index 4af566a3d..3e4344ae9 100644 --- a/examples/inference/basic/basic_minimax_h3_t2v.py +++ b/examples/inference/basic/basic_minimax_h3_t2v.py @@ -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, diff --git a/fastvideo/attention/layer.py b/fastvideo/attention/layer.py index 773101564..ac53492b7 100644 --- a/fastvideo/attention/layer.py +++ b/fastvideo/attention/layer.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 import os +from functools import wraps import torch import torch.nn as nn @@ -20,7 +21,7 @@ def _attention_compile_disabled() -> bool: Defaults to ``True`` (the historical behavior: attention runs eager via ``torch.compiler.disable``). Set ``FASTVIDEO_DISABLE_ATTENTION_COMPILE=0`` - to let attention be traced/compiled into the surrounding graph. + to let attention instances constructed under that environment be traced. """ val = os.environ.get("FASTVIDEO_DISABLE_ATTENTION_COMPILE") if val is None: @@ -28,11 +29,32 @@ def _attention_compile_disabled() -> bool: return val.strip().lower() not in ("0", "false", "no", "off", "") +def _attention_compile_explicitly_disabled() -> bool: + """Whether the environment explicitly requests the eager boundary. + + Regional compile can override the historical default for one loaded + transformer, but it must still honor an explicit debugging escape hatch. + """ + return "FASTVIDEO_DISABLE_ATTENTION_COMPILE" in os.environ and _attention_compile_disabled() + + def _maybe_compiler_disable(fn): - """Apply ``torch.compiler.disable`` unless disabled via env var.""" - if _attention_compile_disabled(): - return torch.compiler.disable(fn) - return fn + """Defer the eager/traceable choice to each attention instance. + + A class-definition-time choice makes a process-wide default the only + option. The deferred wrapper keeps ordinary instances on the historical + eager boundary while allowing the regional loader to opt in only the + attention modules owned by the transformer it is compiling. + """ + disabled_fn = torch.compiler.disable(fn) + + @wraps(fn) + def _dispatch(self, *args, **kwargs): + if self._compile_forward_enabled: + return fn(self, *args, **kwargs) + return disabled_fn(self, *args, **kwargs) + + return _dispatch class DistributedAttention(nn.Module): @@ -77,6 +99,13 @@ class DistributedAttention(nn.Module): self.num_kv_heads = num_kv_heads self.backend = backend_name_to_enum(attn_backend.get_name()) self.dtype = dtype + # Preserve the historical compiler-disabled default. The regional + # inference loader may enable this one instance after validating the + # transformer's resolved backend; no process-global default changes. + self._compile_forward_enabled = not _attention_compile_disabled() + + def _set_compile_forward_enabled(self, enabled: bool) -> None: + self._compile_forward_enabled = enabled @_maybe_compiler_disable def forward( diff --git a/fastvideo/envs.py b/fastvideo/envs.py index e68b319ac..64c37a6b7 100644 --- a/fastvideo/envs.py +++ b/fastvideo/envs.py @@ -21,6 +21,7 @@ if TYPE_CHECKING: FASTVIDEO_TRACE_FUNCTION: int = 0 FASTVIDEO_ATTENTION_BACKEND: str | None = None FASTVIDEO_FA4: bool = False + FASTVIDEO_INFERENCE_TORCH_COMPILE: bool = False FASTVIDEO_MINIMAX_H3_FUSIONS: str = "" FASTVIDEO_WORKER_MULTIPROC_METHOD: str = "spawn" FASTVIDEO_TARGET_DEVICE: str = "cuda" @@ -219,6 +220,17 @@ environment_variables: dict[str, Callable[[], Any]] = { "FASTVIDEO_FA4": lambda: os.getenv("FASTVIDEO_FA4", "0") != "0", + # If set (=1), enable regional (per-transformer-block) fullgraph + # torch.compile for the DiT at inference — the inference-side counterpart + # of the training regional-compile port of hao-ai-lab/FastVideo#1718. + # Equivalent to FastVideoArgs.inference_torch_compile=True (e.g. via + # PipelineSelection.experimental={"inference_torch_compile": True}). VSA + # and other non-fullgraph-traceable attention backends degrade to eager + # with one warning; see _regional_compile_unsupported_reason in + # fastvideo/models/loader/fsdp_load.py. + "FASTVIDEO_INFERENCE_TORCH_COMPILE": + lambda: os.getenv("FASTVIDEO_INFERENCE_TORCH_COMPILE", "0") != "0", + # Opt-in MiniMax-H3 inference-only Triton fusions adapted from the # NVlabs/Sana Sol-Engine implementation. Accepts `all`, `1`, or a # comma-separated subset of `modulate,qknorm_rope,swiglu`. An empty value diff --git a/fastvideo/fastvideo_args.py b/fastvideo/fastvideo_args.py index 457458f6a..3b947659b 100644 --- a/fastvideo/fastvideo_args.py +++ b/fastvideo/fastvideo_args.py @@ -164,6 +164,19 @@ class FastVideoArgs: torch_compile_kwargs_text_encoder: dict[str, Any] = field(default_factory=dict) torch_compile_kwargs_vae: dict[str, Any] = field(default_factory=dict) torch_compile_kwargs_audio_vae: dict[str, Any] = field(default_factory=dict) + # Regional (per-transformer-block) fullgraph torch.compile of the DiT at + # inference — the inference-side counterpart of the training regional + # compile ported from hao-ai-lab/FastVideo#1718. Applied by the loader + # right after the transformer loads, with fullgraph=True and inductor + # options {emulate_precision_casts: True} injected (no user kwargs + # needed). 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 + # __post_init__) or PipelineSelection.experimental + # {"inference_torch_compile": true}. Distinct from ``enable_torch_compile``, + # which keeps the pipeline-level compile semantics. + inference_torch_compile: bool = False disable_autocast: bool = False @@ -272,6 +285,13 @@ class FastVideoArgs: self._apply_ltx2_vae_overrides() self._resolve_refine_args() self._apply_transformer_quant() + if not self.inference_torch_compile: + # Parse-once adapter (same pattern as attention_backend below): the + # environment variable is an input read once here, so the loader + # only ever consults the typed field. + import fastvideo.envs as envs + if envs.FASTVIDEO_INFERENCE_TORCH_COMPILE: + self.inference_torch_compile = True if self.attention_backend is not None: # Fail fast on typos instead of silently auto-selecting later. from fastvideo.attention.selector import coerce_attn_backend @@ -592,6 +612,15 @@ class FastVideoArgs: help= "JSON string of kwargs to pass to torch.compile. Example: '{\"backend\":\"inductor\",\"mode\":\"reduce-overhead\"}'", ) + parser.add_argument( + "--inference-torch-compile", + action=StoreBoolean, + default=FastVideoArgs.inference_torch_compile, + help="Regional fullgraph torch.compile of each DiT transformer block at inference " + "(port of the #1718 training-side regional compile). The loader injects fullgraph=True " + "and inductor options {emulate_precision_casts: true}; non-traceable attention backends " + "(VSA) degrade to eager with one warning. FASTVIDEO_INFERENCE_TORCH_COMPILE=1 is equivalent.", + ) parser.add_argument( "--dit-cpu-offload", diff --git a/fastvideo/models/loader/component_loader.py b/fastvideo/models/loader/component_loader.py index 3820376ef..08e74b8f8 100644 --- a/fastvideo/models/loader/component_loader.py +++ b/fastvideo/models/loader/component_loader.py @@ -1126,6 +1126,7 @@ class TransformerLoader(ComponentLoader): training_mode=fastvideo_args.training_mode, enable_torch_compile=fastvideo_args.enable_torch_compile, torch_compile_kwargs=fastvideo_args.torch_compile_kwargs, + inference_regional_compile=fastvideo_args.inference_torch_compile, ) 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 4dfca6efb..904201abc 100644 --- a/fastvideo/models/loader/fsdp_load.py +++ b/fastvideo/models/loader/fsdp_load.py @@ -136,6 +136,7 @@ def maybe_load_fsdp_model( pin_cpu_memory: bool = True, enable_torch_compile: bool = False, torch_compile_kwargs: dict[str, Any] | None = None, + inference_regional_compile: bool = False, ) -> torch.nn.Module: """ Load the model with FSDP if is training, else load the model without FSDP. @@ -235,9 +236,126 @@ def maybe_load_fsdp_model( 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) + 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: + """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. + """ + try: + from fastvideo.attention.layer import _attention_compile_explicitly_disabled + except Exception: # pragma: no cover - attention stack not importable + pass + else: + if _attention_compile_explicitly_disabled(): + # The escape hatch wraps attention forwards in + # torch.compiler.disable, which is a hard dynamo error inside a + # fullgraph region ("Skip inlining `torch.compiler.disable()`d + # function"). Degrade to eager instead, matching the hatch's + # debugging intent. + return ("FASTVIDEO_DISABLE_ATTENTION_COMPILE=1 keeps attention " + "forwards out of compiled graphs via torch.compiler." + "disable, which fullgraph regional compile cannot trace; " + "this model stays eager") + config = init_params.get("config") + resolved = getattr(config, "_resolved_attention_backend", None) + resolved_name = getattr(resolved, "name", "") + if resolved_name in ("VIDEO_SPARSE_ATTN", "VIDEO_SPARSE_ATTN_H3"): + return (f"attention backend resolved to {resolved_name}, whose Triton " + "kernels, sequence-parallel collectives, and sync metadata " + "guard graph-break (incompatible with fullgraph regional " + "compile); this model stays eager") + return None + + +def _enable_regional_attention_compile(model: nn.Module) -> int: + """Opt in distributed-attention instances owned by ``model`` only.""" + from fastvideo.attention.layer import DistributedAttention + + enabled_count = 0 + for submodule in model.modules(): + if isinstance(submodule, DistributedAttention): + submodule._set_compile_forward_enabled(True) + enabled_count += 1 + return enabled_count + + +def _compile_model_regions(model: nn.Module, compile_kwargs: dict[str, Any]) -> int: + """Compile repeated mathematical regions of a loaded model. + + Only the selected module ``forward`` is replaced. This keeps activation + checkpoint wrappers structurally transparent while any module-level hooks + (FSDP pre/post, layerwise offload) execute outside the compiled region. + """ + compile_conditions = getattr(model, "_compile_conditions", None) + if not compile_conditions: + raise ValueError(f"{type(model).__name__} does not declare _compile_conditions") + + if compile_kwargs.get("fullgraph", True) is not True: + raise ValueError("Regional compile requires fullgraph=True") + if "mode" in compile_kwargs: + # torch.compile forbids passing both `mode` and `options`, and + # regional compile always injects options (emulate_precision_casts) + # to match the training-side regional-compile configuration. Fail here + # with an actionable message instead of letting torch raise a + # mode/options conflict about an `options` key the user never wrote. + raise ValueError("Regional compile sets inductor options " + "(emulate_precision_casts) and cannot be combined " + "with torch_compile_kwargs['mode']. Remove 'mode' or " + "express its effect via torch_compile_kwargs['options'].") + kwargs = {**compile_kwargs, "fullgraph": True} + options = {"emulate_precision_casts": True} + options.update(kwargs.get("options") or {}) + kwargs["options"] = options + compiled_count = 0 + for name, submodule in list(model.named_modules()): + if not name: + continue + if any(condition(name, submodule) for condition in compile_conditions): + # Activation checkpoint wrappers are control-flow boundaries, not + # mathematical regions. Keep their saved-tensor/recompute logic + # eager and compile only the repeated block they own. + compile_target = getattr(submodule, "_checkpoint_wrapped_module", submodule) + compile_target.forward = torch.compile(compile_target.forward, **kwargs) + compiled_count += 1 + + if compiled_count == 0: + raise ValueError(f"No submodules in {type(model).__name__} matched _compile_conditions") + logger.info( + "Enabled regional torch.compile for %d submodules in %s with kwargs=%s", + compiled_count, + type(model).__name__, + kwargs, + ) + return compiled_count + + def shard_model( model, *, diff --git a/fastvideo/pipelines/composed_pipeline_base.py b/fastvideo/pipelines/composed_pipeline_base.py index 6e2cc2f82..bc2a0f9a1 100644 --- a/fastvideo/pipelines/composed_pipeline_base.py +++ b/fastvideo/pipelines/composed_pipeline_base.py @@ -203,6 +203,14 @@ class ComposedPipelineBase(ABC): vae_compile_kwargs = (self.fastvideo_args.torch_compile_kwargs_vae or global_compile_kwargs) audio_vae_compile_kwargs = (self.fastvideo_args.torch_compile_kwargs_audio_vae or global_compile_kwargs) + if compile_transformer and self.fastvideo_args.inference_torch_compile: + # The loader already applied the regional fullgraph + # compile to the DiT blocks (inference_torch_compile); + # wrapping the same forwards again here would stack + # compiled callables. + logger.info("inference_torch_compile already compiled the DiT regions in the " + "loader; skipping the pipeline-level DiT compile") + compile_transformer = False if compile_transformer: self._maybe_compile_pipeline_module( module_name="transformer", diff --git a/fastvideo/tests/inference/test_inference_regional_compile.py b/fastvideo/tests/inference/test_inference_regional_compile.py new file mode 100644 index 000000000..f827d74e1 --- /dev/null +++ b/fastvideo/tests/inference/test_inference_regional_compile.py @@ -0,0 +1,209 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Contract tests for the inference-side regional torch.compile port. + +The loader applies a per-transformer-block fullgraph compile after the +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. +- attention forward dispatch: ordinary instances retain the historical + compiler-disabled boundary; regional compile opts in only the selected + model's instances. +- ``_compile_model_regions``: exactly the ``_compile_conditions`` blocks are + compiled, fullgraph=True plus inductor ``emulate_precision_casts`` are + injected, and ``mode`` kwargs are rejected (torch.compile forbids + mode+options). + +CPU-safe: the real capture check uses torch's eager recording backend; no CUDA +or inductor kernel build is needed. +""" + +import sys +from types import ModuleType, SimpleNamespace + +import pytest +import torch +from torch import nn + +from fastvideo.attention import layer as attention_layer +from fastvideo.attention.layer import DistributedAttention +from fastvideo.models.loader import fsdp_load +from fastvideo.models.loader.fsdp_load import ( + _compile_model_regions, + _enable_regional_attention_compile, + _regional_compile_unsupported_reason, +) + + +def _init_params_for(backend_name: str | None) -> dict: + resolved = None if backend_name is None else SimpleNamespace(name=backend_name) + 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: + monkeypatch.delenv("FASTVIDEO_DISABLE_ATTENTION_COMPILE", raising=False) + reason = _regional_compile_unsupported_reason(_init_params_for(backend_name)) + assert reason is not None + assert backend_name in reason + assert "eager" in reason + + +@pytest.mark.parametrize("backend_name", [None, "TORCH_SDPA"]) +def test_dense_backends_allow_compile(backend_name, monkeypatch) -> None: + monkeypatch.delenv("FASTVIDEO_DISABLE_ATTENTION_COMPILE", raising=False) + assert _regional_compile_unsupported_reason(_init_params_for(backend_name)) is None + + +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")) + assert reason is not None + assert "FASTVIDEO_DISABLE_ATTENTION_COMPILE" in reason + + +@pytest.mark.parametrize("fa_version", ["2", "3", "4"]) +def test_dense_flash_attention_inference_allows_compile(fa_version, monkeypatch) -> None: + monkeypatch.delenv("FASTVIDEO_DISABLE_ATTENTION_COMPILE", raising=False) + fake_module = ModuleType("fastvideo.attention.utils.flash_attn_default") + fake_module.fa_version = fa_version + monkeypatch.setitem(sys.modules, fake_module.__name__, fake_module) + + assert _regional_compile_unsupported_reason(_init_params_for("FLASH_ATTN")) is None + + +def test_default_attention_dispatch_stays_compiler_disabled(monkeypatch) -> None: + monkeypatch.delenv("FASTVIDEO_DISABLE_ATTENTION_COMPILE", raising=False) + disabled_calls: list[str] = [] + + def _fake_disable(fn): + + def _disabled(self): + disabled_calls.append("disabled") + return fn(self) + + return _disabled + + monkeypatch.setattr(attention_layer.torch.compiler, "disable", _fake_disable) + + class _ToyAttention: + + def __init__(self) -> None: + self._compile_forward_enabled = not attention_layer._attention_compile_disabled() + + @attention_layer._maybe_compiler_disable + def forward(self) -> str: + return "forward" + + ordinary = _ToyAttention() + assert ordinary.forward() == "forward" + assert disabled_calls == ["disabled"] + + # Preserve the existing process-wide escape hatch for callers that opt in + # before constructing their attention modules. + monkeypatch.setenv("FASTVIDEO_DISABLE_ATTENTION_COMPILE", "0") + explicit = _ToyAttention() + assert explicit.forward() == "forward" + assert disabled_calls == ["disabled"] + + +class _BareDistributedAttention(DistributedAttention): + + def __init__(self) -> None: + nn.Module.__init__(self) + self._compile_forward_enabled = False + + +class _CompileProbe(nn.Module): + + def __init__(self, enabled: bool) -> None: + super().__init__() + self._compile_forward_enabled = enabled + + @attention_layer._maybe_compiler_disable + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x + 1 + + +def test_attention_instance_flag_controls_real_fullgraph_capture() -> None: + try: + ordinary = torch.compile(_CompileProbe(False), backend="eager", fullgraph=True) + with pytest.raises(torch._dynamo.exc.Unsupported, match="Skip calling"): + ordinary(torch.ones(2)) + + regional = torch.compile(_CompileProbe(True), backend="eager", fullgraph=True) + torch.testing.assert_close(regional(torch.ones(2)), torch.full((2,), 2.0)) + finally: + torch._dynamo.reset() + + +def test_regional_compile_enables_only_loaded_model_attention() -> None: + selected_model = nn.Sequential(_BareDistributedAttention(), _BareDistributedAttention()) + unrelated_model = nn.Sequential(_BareDistributedAttention()) + + assert _enable_regional_attention_compile(selected_model) == 2 + assert all(module._compile_forward_enabled for module in selected_model) + assert not unrelated_model[0]._compile_forward_enabled + + +class _Block(nn.Module): + + def __init__(self) -> None: + super().__init__() + self.linear = nn.Linear(4, 4) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.linear(x) + + +class _Toy(nn.Module): + _compile_conditions = [lambda name, module: name.startswith("blocks.") and name.count(".") == 1] + + def __init__(self) -> None: + super().__init__() + self.blocks = nn.ModuleList([_Block() for _ in range(3)]) + self.proj_out = nn.Linear(4, 4) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + for block in self.blocks: + x = block(x) + return self.proj_out(x) + + +def test_compile_model_regions_injects_fullgraph_and_precision_casts(monkeypatch) -> None: + captured: list[dict] = [] + + def _fake_compile(fn, **kwargs): + captured.append(kwargs) + return fn + + monkeypatch.setattr(fsdp_load.torch, "compile", _fake_compile) + model = _Toy() + count = _compile_model_regions(model, {}) + # The three repeated blocks compile; proj_out and the root stay eager. + assert count == 3 + assert len(captured) == 3 + for kwargs in captured: + assert kwargs["fullgraph"] is True + assert kwargs["options"] == {"emulate_precision_casts": True} + + +def test_compile_model_regions_rejects_mode_kwargs() -> None: + with pytest.raises(ValueError, match="mode"): + _compile_model_regions(_Toy(), {"mode": "reduce-overhead"}) + + +def test_compile_model_regions_requires_conditions_and_matches(monkeypatch) -> None: + monkeypatch.setattr(fsdp_load.torch, "compile", lambda fn, **kwargs: fn) + plain = nn.Linear(4, 4) + with pytest.raises(ValueError, match="_compile_conditions"): + _compile_model_regions(plain, {}) + + class _NoMatch(_Toy): + _compile_conditions = [lambda name, module: False] + + with pytest.raises(ValueError, match="matched"): + _compile_model_regions(_NoMatch(), {})