remove rope fusion option, do automatically on torch compile

ghstack-source-id: c37cc48
Pull-Request: #4055

Author: howardzhang-cv

benchmarks/prototype/attention/eval_flux_model.py

@@ -43,12 +43,6 @@
         "fp8": True,
         "fp8_backend": AttentionBackend.FP8_FA3,
     },
-    "fa4": {"flash_impl": "FA4", "fp8": False},
-    "fa4_fp8": {
-        "flash_impl": "FA4",
-        "fp8": True,
-        "fp8_backend": AttentionBackend.FP8_FA4,
-    },
 }
 
 IMAGE_SIZE = (512, 512)  # (width, height) - resize for consistent LPIPS
@@ -68,7 +62,6 @@ def setup_backend(
     backend_name,
     compile_flag,
     orig_transformer,
-    fuse_rope_using_torch_compile=False,
 ):
     """Set up a backend and return the flash_impl name."""
     cfg = BACKENDS[backend_name]
@@ -79,16 +72,8 @@ def setup_backend(
         pipe.transformer = apply_low_precision_attention(
             pipe.transformer,
             backend=cfg["fp8_backend"],
-            fuse_rope_using_torch_compile=fuse_rope_using_torch_compile,
         )
-        if fuse_rope_using_torch_compile:
-            print(
-                f"Compiling transformer with torch.compile ({backend_name}, FP8 backend)..."
-            )
-            pipe.transformer = torch.compile(
-                pipe.transformer, backend=pipe.transformer.compile_backend
-            )
-        elif compile_flag:
+        if compile_flag:
             print(f"Compiling transformer with torch.compile ({backend_name})...")
             pipe.transformer = torch.compile(pipe.transformer)
         return cfg["flash_impl"]
@@ -161,7 +146,6 @@ def run_benchmark(
     debug_prompt: Optional[str] = None,
     warmup_iters: int = 2,
     compile: bool = False,
-    fuse_rope_using_torch_compile: bool = False,
 ):
     """Run the attention backend benchmark on FLUX.1-schnell."""
     compile_str = " + torch.compile" if compile else ""
@@ -218,7 +202,6 @@ def run_benchmark(
         baseline_backend,
         compile,
         orig_transformer,
-        fuse_rope_using_torch_compile=fuse_rope_using_torch_compile,
     )
 
     print(f"Warming up {baseline_backend} with {warmup_iters} iterations...")
@@ -282,7 +265,6 @@ def run_benchmark(
         test_backend,
         compile,
         orig_transformer,
-        fuse_rope_using_torch_compile=fuse_rope_using_torch_compile,
     )
 
     print(f"Warming up {test_backend} with {warmup_iters} iterations...")
@@ -422,11 +404,6 @@ def main():
         action="store_true",
         help="Wrap the model with torch.compile for both backends",
     )
-    parser.add_argument(
-        "--fuse_rope_using_torch_compile",
-        action="store_true",
-        help="Fuse RoPE into the FP8 kernel (compile path, off by default)",
-    )
     parser.add_argument(
         "--height",
         type=int,
@@ -452,7 +429,6 @@ def main():
         debug_prompt=args.debug_prompt,
         warmup_iters=args.warmup_iters,
         compile=args.compile,
-        fuse_rope_using_torch_compile=args.fuse_rope_using_torch_compile,
     )
 
 

benchmarks/prototype/attention/eval_llama3_model.py

@@ -57,17 +57,6 @@
         "fp8_backend": AttentionBackend.FP8_FA3,
         "label": "FA3 FP8",
     },
-    "fa4": {
-        "flash_impl": "FA4",
-        "fp8": False,
-        "label": "FA4 BF16",
-    },
-    "fa4_fp8": {
-        "flash_impl": "FA4",
-        "fp8": True,
-        "fp8_backend": AttentionBackend.FP8_FA4,
-        "label": "FA4 FP8",
-    },
 }
 
 RANDOM_SEED = 42
@@ -115,9 +104,7 @@ def mask_strip_backend(gm, example_inputs):
     return torch.compile(model, backend=mask_strip_backend)
 
 
-def setup_backend(
-    orig_model, backend_name, compile_flag, fuse_rope_using_torch_compile=False
-):
+def setup_backend(orig_model, backend_name, compile_flag):
     """Set up a backend and return (model, flash_impl)."""
     cfg = BACKENDS[backend_name]
 
@@ -129,14 +116,8 @@ def setup_backend(
         model = apply_low_precision_attention(
             orig_model,
             backend=cfg["fp8_backend"],
-            fuse_rope_using_torch_compile=fuse_rope_using_torch_compile,
         )
-        if fuse_rope_using_torch_compile:
-            print(
-                f"  Compiling model with torch.compile ({backend_name}, FP8 backend)..."
-            )
-            model = torch.compile(model, backend=model.compile_backend)
-        elif compile_flag:
+        if compile_flag:
             print(f"  Compiling model with torch.compile ({backend_name})...")
             model = torch.compile(model)
         return model, cfg["flash_impl"]
@@ -221,7 +202,6 @@ def run_benchmark(
     num_warmup: int = 3,
     seq_lengths: list[int] | None = None,
     compile: bool = False,
-    fuse_rope_using_torch_compile: bool = False,
 ):
     if seq_lengths is None:
         seq_lengths = DEFAULT_SEQ_LENGTHS
@@ -266,7 +246,6 @@ def run_benchmark(
         orig_model,
         baseline_backend,
         compile,
-        fuse_rope_using_torch_compile=fuse_rope_using_torch_compile,
     )
     baseline_ppl = evaluate_perplexity(baseline_model, tokenizer, baseline_flash)
     print(f"  {baseline_label} perplexity: {baseline_ppl:.2f}")
@@ -277,7 +256,6 @@ def run_benchmark(
         orig_model,
         test_backend,
         compile,
-        fuse_rope_using_torch_compile=fuse_rope_using_torch_compile,
     )
     test_ppl = evaluate_perplexity(test_model, tokenizer, test_flash)
     print(f"  {test_label} perplexity: {test_ppl:.2f}")
@@ -299,7 +277,6 @@ def run_benchmark(
         orig_model,
         baseline_backend,
         compile,
-        fuse_rope_using_torch_compile=fuse_rope_using_torch_compile,
     )
     baseline_runtimes = {}
     for S in seq_lengths:
@@ -330,7 +307,6 @@ def run_benchmark(
         orig_model,
         test_backend,
         compile,
-        fuse_rope_using_torch_compile=fuse_rope_using_torch_compile,
     )
     test_runtimes = {}
     for S in baseline_runtimes:
@@ -459,11 +435,6 @@ def main():
         action="store_true",
         help="Wrap the model with torch.compile (applies to non-FP8 backends)",
     )
-    parser.add_argument(
-        "--fuse_rope_using_torch_compile",
-        action="store_true",
-        help="Fuse RoPE into the FP8 kernel (compile path, off by default)",
-    )
     args = parser.parse_args()
 
     run_benchmark(
@@ -474,7 +445,6 @@ def main():
         num_warmup=args.num_warmup,
         seq_lengths=args.seq_lengths,
         compile=args.compile,
-        fuse_rope_using_torch_compile=args.fuse_rope_using_torch_compile,
     )
 
 

test/prototype/attention/test_fp8_attention.py

@@ -190,7 +190,6 @@ def test_monkey_patch_model(self, dtype):
         fp8_model = apply_low_precision_attention(
             fp8_model,
             backend=AttentionBackend.FP8_FA3,
-            fuse_rope_using_torch_compile=False,
         )
 
         with torch.no_grad():
@@ -231,9 +230,8 @@ def test_rope_fusion_model(self, dtype):
         fp8_model = apply_low_precision_attention(
             fp8_model,
             backend=AttentionBackend.FP8_FA3,
-            fuse_rope_using_torch_compile=True,
         )
-        fp8_model = torch.compile(fp8_model, backend=fp8_model.compile_backend)
+        fp8_model = torch.compile(fp8_model)
 
         with torch.no_grad():
             out_fp8 = fp8_model(x, cos, sin)

torchao/prototype/attention/api.py

@@ -60,20 +60,19 @@ def _check_backend_available(backend: AttentionBackend) -> None:
 def apply_low_precision_attention(
     model: nn.Module,
     backend: Optional[AttentionBackend] = None,
-    fuse_rope_using_torch_compile: bool = False,
 ) -> nn.Module:
     """Apply low-precision attention to a model.
 
     Must be called before ``torch.compile``. KV caching should be
     disabled before calling (e.g., ``config.use_cache = False`` for
     HuggingFace models).
 
-    When ``fuse_rope_using_torch_compile=True``, the returned wrapper
-    exposes a ``compile_backend`` attribute. You must compile with it to get
-    the RoPE fusion::
+    This replaces ``F.scaled_dot_product_attention`` with an FP8 SDPA
+    for eager execution and sets a global pre-grad pass so that
+    ``torch.compile`` will automatically fuse RoPE where detected::
 
-        model = apply_low_precision_attention(model, fuse_rope_using_torch_compile=True)
-        model = torch.compile(model, backend=model.compile_backend)
+        model = apply_low_precision_attention(model)
+        model = torch.compile(model)   # RoPE fusion happens automatically
     """
     if isinstance(model, _LowPrecisionAttentionWrapper):
         raise RuntimeError(
@@ -90,6 +89,6 @@ def apply_low_precision_attention(
         _check_backend_available(backend)
 
     if backend == AttentionBackend.FP8_FA3:
-        return setup_fp8_backend(model, "FA3", fuse_rope_using_torch_compile)
+        return setup_fp8_backend(model, "FA3")
 
     raise ValueError(f"Unknown backend: {backend}")

torchao/prototype/attention/fp8_fa3/attention.py

@@ -23,6 +23,9 @@
     _fp8_rope_sdpa,
     _fp8_sdpa,
 )
+from torchao.prototype.attention.shared_utils.custom_ops import (
+    register_fp8_attention_ops,
+)
 
 fp8_fa3_sdpa = partial(_fp8_sdpa, backend_name="FA3")
 fp8_fa3_sdpa.__doc__ = _fp8_sdpa.__doc__
@@ -33,3 +36,9 @@
 fp8_fa3_rope_sdpa.__doc__ = _fp8_rope_sdpa.__doc__
 fp8_fa3_rope_sdpa.__name__ = "fp8_fa3_rope_sdpa"
 fp8_fa3_rope_sdpa.__qualname__ = "fp8_fa3_rope_sdpa"
+
+_ops = register_fp8_attention_ops(
+    backend_name="fa3",
+    rope_sdpa_fn=fp8_fa3_rope_sdpa,
+    sdpa_fn=fp8_fa3_sdpa,
+)

torchao/prototype/attention/fp8_fa3/fusion_pass.py

@@ -12,19 +12,9 @@
 tells the compiler to treat them as opaque nodes with known shapes/dtypes.
 """
 
-from functools import partial
 from typing import Callable, NamedTuple
 
 import torch
-import torch._inductor.config as inductor_config
-import torch.nn as nn
-
-from torchao.prototype.attention.shared_utils.fusion_utils import (
-    detect_causal_mask,
-)
-from torchao.prototype.attention.shared_utils.fusion_utils import (
-    rope_sdpa_fusion_pass as _shared_fusion_pass,
-)
 
 
 class RegisteredOps(NamedTuple):
@@ -120,42 +110,3 @@ def _sdpa_fake(
     fp8_sdpa_op = getattr(getattr(torch.ops, "torchao"), f"fp8_{backend}_sdpa").default
 
     return RegisteredOps(rope_sdpa_op=rope_sdpa_op, fp8_sdpa_op=fp8_sdpa_op)
-
-
-def make_backend_fn(
-    ops: RegisteredOps,
-    backend_name: str,
-    flash_impl_name: str,
-    max_head_dim: int = 256,
-) -> Callable:
-    """Return a ``make_fp8_backend(model, fuse_rope_using_torch_compile)`` function for a backend."""
-
-    def make_fp8_backend(
-        model: nn.Module,
-        fuse_rope_using_torch_compile: bool,
-    ) -> Callable:
-        from torch._inductor.compile_fx import compile_fx
-
-        strip_causal_mask = detect_causal_mask(model, flash_impl_name=flash_impl_name)
-
-        pass_fn = partial(
-            _shared_fusion_pass,
-            rope_sdpa_op=ops.rope_sdpa_op,
-            fp8_sdpa_op=ops.fp8_sdpa_op,
-            max_head_dim=max_head_dim,
-            backend_name=backend_name,
-            fuse_rope=fuse_rope_using_torch_compile,
-            strip_causal_mask=strip_causal_mask,
-        )
-
-        def fp8_attention_backend(gm, example_inputs):
-            old_pass = inductor_config.pre_grad_custom_pass
-            inductor_config.pre_grad_custom_pass = pass_fn
-            try:
-                return compile_fx(gm, example_inputs)
-            finally:
-                inductor_config.pre_grad_custom_pass = old_pass
-
-        return fp8_attention_backend
-
-    return make_fp8_backend