Repository: pytorch
Issue: Custom attention recompilations
Labels: triaged,oncall: pt2,module: dynamic shapes,module: dynamo

### 🐛 Describe the bug

I have a compiler cache exhaustion with:
```python
  @torch.compile
    def forward(self, q, k, v):
        n, c, h, w = v.size()

        if self.use_linear:
            q = self.linear_Q(q)
            k = self.linear_K(k)
            v = self.linear_V(v)

        hidden_dim = self.hidden_dim

        if self.qk_mask is not None and (h, w) == self.last_size_2d:
            qk_mask = self.qk_mask
        else:
            memory_mask = torch.ones((1, 1, h, w), device=v.device).float()
            unfolded_k_mask = self.pad_and_unfold(memory_mask).view(
                1, 1, self.window_size * self.window_size, h * w)
            qk_mask = 1 - unfolded_k_mask
            self.qk_mask = qk_mask

        relative_emb = self.relative_emb_k(q)

        # Scale
        q = q / self.T

        q = q.view(-1, self.d_att, h, w)
        k = k.view(-1, self.d_att, h, w)
        v = v.view(-1, self.num_head, hidden_dim, h * w)

        relative_emb = relative_emb.view(n, self.num_head,
                                        self.window_size * self.window_size,
                                        h * w)

        if self.enable_corr:
            qk = self.correlation_sampler(q, k).view(
                n, self.num_head, self.window_size * self.window_size, h * w)
        else:
            unfolded_k = self.pad_and_unfold(k).view(
                n * self.num_head, hidden_dim,
                self.window_size * self.window_size, h, w)
            qk = (q.unsqueeze(2) * unfolded_k).sum(dim=1).view(
                n, self.num_head, self.window_size * self.window_size, h * w)
        if self.use_dis:
            qk = 2 * qk - self.pad_and_unfold(
                k.pow(2).sum(dim=1, keepdim=True)).view(
                    n, self.num_head, self.window_size * self.window_size,
                    h * w)

        qk = qk + relative_emb

        qk -= qk_mask * 1e+8 if qk.dtype == torch.float32 else qk_mask * 1e+4

        local_attn = torch.softmax(qk, dim=2)

        local_attn = self.dropout(local_attn)

        if self.use_relative_v:
            agg_bias = torch.einsum('bhwn,hcw->bhnc', local_attn,
                                    self.relative_emb_v)

        global_attn = self.local2global(local_attn, h, w)

        agg_value = (global_attn @ v.transpose(-2, -1))

        if self.use_relative_v:
            output = (agg_value + agg_bias).permute(2, 0, 1,
                                                    3).reshape(h * w, n, c)
        else:
            output = (agg_value).permute(2, 0, 1,
                                                    3).reshape(h * w, n, c)

        output = self.projection(output)

        self.last_size_2d = (h, w)
        return output, local_attn
```

### Error logs

```python
W0308 16:22:56.764000 135559255537472 torch/_dynamo/convert_frame.py:357] torch._dynamo hit config.cache_size_limit (8)
W0308 16:22:56.764000 135559255537472 torch/_dynamo/convert_frame.py:357]    function: 'forward' (/./networks/layers/attention.py:310)
W0308 16:22:56.764000 135559255537472 torch/_dynamo/convert_frame.py:357]    last reason: ___check_obj_id(L['self'].qk_mask, 7616288)                 
W0308 16:22:56.764000 135559255537472 torch/_dynamo/convert_frame.py:357] To log all recompilation reasons, use TORCH_LOGS="recompiles".
W0308 16:22:56.764000 135559255537472 torch/_dynamo/convert_frame.py:357] To diagnose recompilation issues, see https://pytorch.org/docs/master/compile/troubleshooting.html.
 W0308 16:22:56.767000 138421087676224 torch/_dynamo/convert_frame.py:357] torch._dynamo hit config.cache_size_limit (8)
 W0308 16:22:56.767000 138421087676224 torch/_dynamo/convert_frame.py:357]    function: 'forward' (/workspace/./networks/layers/attention.py:310)
 W0308 16:22:56.767000 138421087676224 torch/_dynamo/convert_frame.py:357]    last reason: ___check_obj_id(L['self'].qk_mask, 7616288)                 
W0308 16:22:56.767000 138421087676224 torch/_dynamo/convert_frame.py:357] To log all recompilation reasons, use TORCH_LOGS="recompiles".
W0308 16:22:56.767000 138421087676224 torch/_dynamo/convert_frame.py:357] To diagnose recompilation issues, see https://pytorch.org/docs/master/compile/troubleshooting.html.
```

### Minified repro

_No response_

### Versions

pytorch nighlty

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