Repository: pytorch
Issue: Support AC with graph break
Labels: triaged,enhancement,oncall: pt2,module: dynamo,module: higher order operators,module: pt2-dispatcher,activation-checkpointing

### 🐛 Describe the bug

As mentioned in this [blog](https://dev-discuss.pytorch.org/t/higher-order-operators-2023-10/1565), HigherOrderOperator does not support graph break inside the input/output function, because Dynamo cannot determine whether the operations inside the function are "safe". However, in most cases, graph break within the function is safe and acceptable. Is there a way to make Dynamo support these cases without falling back to eager mode? For example, this scenario frequently occurs in LLM models:

```
import torch
import torch.nn as nn
import flash_attn


class Attention(nn.Module):
    def __init__(self):
        super().__init__()
        self.q_proj = nn.Linear(4096, 4096)
        self.k_proj = nn.Linear(4096, 4096)
        self.v_proj = nn.Linear(4096, 4096)
        self.o_proj = nn.Linear(4096, 4096)
    
    def forward(self, x):
        def attn_forward(x):
            q = self.q_proj(x)
            k = self.k_proj(x)
            v = self.v_proj(x)
            q = q.view(x.shape[0], x.shape[1], 64, 64).transpose(1, 2)
            k = k.view(x.shape[0], x.shape[1], 64, 64).transpose(1, 2)
            v = v.view(x.shape[0], x.shape[1], 64, 64).transpose(1, 2)
            return flash_attn.flash_attn_func(q, k, v)
        return torch.utils.checkpoint.checkpoint(attn_forward, x, use_reentrant=False)

model = Attention()
model.cuda()

def custom_backand(gm, example_inputs):
    print(gm)
    return gm

opt_model = torch.compile(model, backend=custom_backand)

x = torch.randn(1, 32, 4096).cuda()

with torch.autocast("cuda", dtype=torch.bfloat16):
    out = opt_model(x)
```

In LLM models, flash attention is often used along with activation checkpointing, which can cause the entire attention or decoder layer execution to fall back to eager mode.

Is there a way to make Dynamo support graph breaking within HigherOrderOperator, such as providing some APIs to mark these functions as safe?

I know that we can use PyTorch's SDPA (Scaled Dot-Product Attention) op as a replacement, but our users typically hard-code the use of flash attention in their code. I hope there is a way to achieve this without requiring significant changes to the user's code.



### Error logs

_No response_

### Minified repro

_No response_

### Versions

PyTorch version: 2.5.0a0+git24dee99
Is debug build: False
CUDA used to build PyTorch: 12.1
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.6 LTS (x86_64)
GCC version: (Ubuntu 11.4.0-2ubuntu1~20.04) 11.4.0
Clang version: Could not collect
CMake version: version 3.28.3
Libc version: glibc-2.31

Python version: 3.10.13 (main, Aug 25 2023, 13:20:03) [GCC 9.4.0] (64-bit runtime)

Versions of relevant libraries:
[pip3] numpy==1.26.4
[pip3] torch==2.5.0a0+git417a076
[pip3] torchvision==0.18.0

cc @ezyang @chauhang @penguinwu @voznesenskym @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @chenyang78 @kadeng @amjames @zou3519 @ydwu4 @bdhirsh @yf225