[bugfix]: scope regional attention compile opt-in

docs/inference/optimizations.md

@@ -345,10 +345,11 @@ training-side regional compile of
 loader wraps each `_compile_conditions` block in
 `torch.compile(fullgraph=True)` with inductor
 `options={"emulate_precision_casts": True}` right after the transformer
-loads. Attention backends that cannot be traced end-to-end (VSA,
-FLASH_ATTN on flash-attn 3, or the `FASTVIDEO_DISABLE_ATTENTION_COMPILE=1`
-escape hatch) degrade the transformer to eager with one warning instead of
-failing mid-denoise.
+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, FLASH_ATTN on flash-attn
+3, 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(
@@ -362,7 +363,12 @@ 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
+### 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 |
 |---|---|
@@ -378,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
@@ -396,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

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
@@ -18,21 +19,42 @@ from fastvideo.layers.rotary_embedding import _apply_rotary_emb
 def _attention_compile_disabled() -> bool:
     """Whether to keep attention ``forward`` out of the torch.compile graph.
 
-    Attention backends expose traceable custom-op boundaries, so compilation
-    is enabled by default. Set ``FASTVIDEO_DISABLE_ATTENTION_COMPILE=1`` for
-    an explicit eager escape hatch when debugging a backend.
+    Defaults to ``True`` (the historical behavior: attention runs eager via
+    ``torch.compiler.disable``). Set ``FASTVIDEO_DISABLE_ATTENTION_COMPILE=0``
+    to let attention instances constructed under that environment be traced.
     """
     val = os.environ.get("FASTVIDEO_DISABLE_ATTENTION_COMPILE")
     if val is None:
-        return False
+        return True
     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(

fastvideo/models/loader/fsdp_load.py

@@ -250,6 +250,8 @@ def maybe_load_fsdp_model(
             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
 
@@ -269,11 +271,11 @@ def _regional_compile_unsupported_reason(init_params: dict[str, Any]) -> str | N
     falls back to eager while compile-safe dense loads still compile.
     """
     try:
-        from fastvideo.attention.layer import _attention_compile_disabled
+        from fastvideo.attention.layer import _attention_compile_explicitly_disabled
     except Exception:  # pragma: no cover - attention stack not importable
         pass
     else:
-        if _attention_compile_disabled():
+        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
@@ -301,10 +303,22 @@ def _regional_compile_unsupported_reason(init_params: dict[str, Any]) -> str | N
         return ("attention backend resolved to FLASH_ATTN with flash-attn 3, "
                 "whose grad-enabled path graph-breaks (incompatible with "
                 "fullgraph regional compile); use FA2, FA4 (FASTVIDEO_FA4=1), "
-                "or TORCH_SDPA for compiled runs")
+                "or TORCH_SDPA for compiled runs; 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.
 
@@ -321,9 +335,9 @@ def _compile_model_regions(model: nn.Module, compile_kwargs: dict[str, Any]) ->
     if "mode" in compile_kwargs:
         # torch.compile forbids passing both `mode` and `options`, and
         # regional compile always injects options (emulate_precision_casts)
-        # for bf16 numerics parity. Fail here with an actionable message
-        # instead of letting torch raise a mode/options conflict about an
-        # `options` key the user never wrote.
+        # 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 "

fastvideo/tests/inference/test_inference_regional_compile.py

@@ -9,23 +9,31 @@ 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: torch.compile is monkeypatched, no CUDA needed.
+CPU-safe: the real capture check uses torch's eager recording backend; no CUDA
+or inductor kernel build is needed.
 """
 
-from types import SimpleNamespace
+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,
 )
 
@@ -57,6 +65,94 @@ def test_attention_compile_escape_hatch_degrades_to_eager(monkeypatch) -> None:
     assert "FASTVIDEO_DISABLE_ATTENTION_COMPILE" in reason
 
 
+@pytest.mark.parametrize(("fa_version", "allowed"), [("2", True), ("3", False), ("4", True)])
+def test_flash_attention_3_degrades_to_eager(fa_version, allowed, 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)
+
+    reason = _regional_compile_unsupported_reason(_init_params_for("FLASH_ATTN"))
+    assert (reason is None) is allowed
+    if not allowed:
+        assert "flash-attn 3" in reason
+        assert "eager" in reason
+
+
+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: