Torch.compile + flex_attention with custom score_mod fails (FlexibleLayout / NoValidChoicesError)

V1kstrand: <p>(I also made a pytorch github post <a href="https://github.com/pytorch/pytorch/issues/163300" rel="noopener nofollow ugc">here</a>)</p>
<h3><a name="p-475539-describe-the-bug-1" class="anchor" href="#p-475539-describe-the-bug-1"></a><img src="https://discuss.pytorch.org/images/emoji/apple/bug.png?v=14" title=":bug:" class="emoji" alt=":bug:" loading="lazy" width="20" height="20"> Describe the bug</h3>
<p>Compiling a module that wraps <code>torch.nn.attention.flex_attention</code> with a custom <code>score_mod</code> fails under <code>torch.compile</code> (Inductor).</p>
<ul>
<li>Eager execution works correctly.</li>
<li>With <code>torch.compile(dynamic=False)</code> I get:</li>
</ul>
<pre><code class="lang-auto">
BackendCompilerFailed: backend='inductor' raised:
LoweringException: AssertionError: convert FlexibleLayout to FixedLayout first

</code></pre>
<ul>
<li>With <code>torch.compile(dynamic=False, mode="max-autotune-no-cudagraphs")</code> I get:</li>
</ul>
<pre><code class="lang-auto">
BackendCompilerFailed: backend='inductor' raised:
LoweringException: NoValidChoicesError

</code></pre>
<ul>
<li>Tried with both <code>bfloat16</code> (autocast) and <code>float16</code> — same errors.</li>
</ul>
<h3><a name="p-475539-minimal-repro-2" class="anchor" href="#p-475539-minimal-repro-2"></a>Minimal repro</h3>
<pre data-code-wrap="python"><code class="lang-python">import torch
import torch.nn as nn
from torch.nn.attention.flex_attention import flex_attention
from torch.nn.attention import SDPBackend

class FlexAttentionCPB(nn.Module):
    def __init__(self, N: int, R: int, H: int = 6, hidden: int = 32):
        super().__init__()
        self.mlp = nn.Sequential(nn.Linear(2, hidden), nn.GELU(), nn.Linear(hidden, H, bias=False))
        self.gamma = nn.Parameter(torch.zeros(H))
        self.H = H
        self.init_tables(N, R)
        self.register_buffer("r_cutoff", torch.tensor(R, dtype=torch.long), persistent=False)

    def init_tables(self, N: int, R: int):
        # continuous position bias –  SwinV2
        P = N - R
        S = int(P**0.5)
        assert S * S == P
        rng = torch.arange(-(S - 1), S, dtype=torch.float32)
        dY, dX = torch.meshgrid(rng, rng, indexing="ij")
        rel = torch.stack([dY / max(S - 1, 1), dX / max(S - 1, 1)], dim=-1).reshape(-1, 2)
        rel_table = torch.sign(rel) * torch.log1p(rel.abs())
        self.register_buffer("rel_table", rel_table, persistent=False)

        yy, xx = torch.arange(S), torch.arange(S)
        Y, X = torch.meshgrid(yy, xx, indexing="ij")
        flat = torch.stack([Y, X], 0).flatten(1)
        d = flat[:, :, None] - flat[:, None, :]
        d = d.permute(1, 2, 0).contiguous()
        d[:, :, 0] += S - 1; d[:, :, 1] += S - 1
        d[:, :, 0] *= 2 * S - 1
        l_idx = d.sum(-1).to(torch.long)

        idx = torch.full((N, N), 0, dtype=torch.long)
        idx[R:, R:] = l_idx
        self.register_buffer("idx_table", idx, persistent=False)

    def _score_mod(self, mu: torch.Tensor):
        bt = self.mlp(self.rel_table)
        idx = self.idx_table
        mu_q, mu_k = mu.unbind(2)
        gam_sig = torch.sigmoid(self.gamma)

        def score_mod(score, b, h, q, kv):
            has_bias = (q &gt;= self.r_cutoff) &amp; (kv &gt;= self.r_cutoff)
            l2 = idx[q, kv]
            bias = bt[l2, h]
            w_gate = gam_sig[h] * (mu_q[b, h, q] + mu_k[b, h, kv])
            return score + has_bias.to(score.dtype) * w_gate * bias

        return score_mod

    def forward(self, q, k, v, mu):
        return flex_attention(q, k, v, score_mod=self._score_mod(mu))

def main():
    device = "cuda"
    B, N, R, d, H = 2, 18, 2, 32, 4
    mod = FlexAttentionCPB(N, R, H).to(device)
    mod = torch.compile(mod, dynamic=False)  # also tested "max-autotune-no-cudagraphs"

    q = torch.randn(B, H, N, d, device=device)
    k = torch.randn_like(q)
    v = torch.randn_like(q)
    mu = torch.randn(B, H, 2, N, device=device)

    with torch.nn.attention.sdpa_kernel(SDPBackend.FLASH_ATTENTION):
        with torch.amp.autocast("cuda", dtype=torch.bfloat16):
            out = mod(q, k, v, mu)
            out.norm().backward()
    print("done")

if __name__ == "__main__":
    main()
</code></pre>
<h3><a name="p-475539-ablation-3" class="anchor" href="#p-475539-ablation-3"></a>Ablation</h3>
<ul>
<li>Eager (no <code>torch.compile</code>) → <img src="https://discuss.pytorch.org/images/emoji/apple/white_check_mark.png?v=14" title=":white_check_mark:" class="emoji" alt=":white_check_mark:" loading="lazy" width="20" height="20"> works</li>
<li><code>torch.compile(dynamic=False)</code> → <img src="https://discuss.pytorch.org/images/emoji/apple/cross_mark.png?v=14" title=":cross_mark:" class="emoji" alt=":cross_mark:" loading="lazy" width="20" height="20"> FlexibleLayout→FixedLayout assertion</li>
<li><code>torch.compile(dynamic=False, mode="max-autotune-no-cudagraphs")</code> → <img src="https://discuss.pytorch.org/images/emoji/apple/cross_mark.png?v=14" title=":cross_mark:" class="emoji" alt=":cross_mark:" loading="lazy" width="20" height="20"> NoValidChoicesError</li>
<li>Backends: tested with <code>SDPBackend.FLASH_ATTENTION</code>.</li>
<li>Dtypes: bf16 autocast and fp16 both fail under compile.</li>
</ul>
<h3><a name="p-475539-environment-4" class="anchor" href="#p-475539-environment-4"></a>Environment</h3>
<pre><code class="lang-auto">python: 3.10.12 (GCC 11.4.0)
cuda: 12.4
cudnn: 90100
torch: 2.6.0+cu124
torchvision: 0.21.0+cu124
torchaudio: 2.6.0+cu124
triton: 3.2.0
GPU: NVIDIA A100-SXM4-80GB (sm 8.0)
</code></pre>
<p>(Full <code>python -m torch.utils.collect_env</code> output attached.)</p>
<h3><a name="p-475539-expected-behavior-5" class="anchor" href="#p-475539-expected-behavior-5"></a>Expected behavior</h3>
<p>FlexAttention with a custom score modifier should compile successfully, or at least provide a clear error if the pattern/layout is unsupported.</p>
<h3><a name="p-475539-actual-behavior-6" class="anchor" href="#p-475539-actual-behavior-6"></a>Actual behavior</h3>
<p>Inductor compilation fails with <code>LoweringException</code> under multiple modes.</p>
<h3><a name="p-475539-additional-context-7" class="anchor" href="#p-475539-additional-context-7"></a>Additional context</h3>
<ul>
<li>
<p>This looks specific to FlexAttention + custom <code>score_mod</code>.</p>
</li>
<li>
<p>Happy to retest on nightly if maintainers think this may already be addressed.</p>
</li>
<li>
<p>Attached logs:</p>
<ul>
<li><code>compile_default_log.txt</code></li>
<li><code>compile_autotune_no_cudagraphs_log.txt</code></li>
<li><code>torch_collect_env.txt</code></li>
<li><code>nvidia_smi.xml</code></li>
<li><code>pip_freeze.txt</code></li>
</ul>
</li>
</ul>
<p>(full list on github issues)</p>
<h3><a name="p-475539-error-logs-8" class="anchor" href="#p-475539-error-logs-8"></a>Error logs</h3>
<p>see attached –</p>
<h3><a name="p-475539-versions-9" class="anchor" href="#p-475539-versions-9"></a>Versions</h3>
<p>see attached –</p>
<p>[pip3] mypy_extensions==1.1.0<br>
[pip3] numpy==2.1.2<br>
[pip3] nvidia-cublas-cu12==12.4.5.8<br>
[pip3] nvidia-cuda-cupti-cu12==12.4.127<br>
[pip3] nvidia-cuda-nvrtc-cu12==12.4.127<br>
[pip3] nvidia-cuda-runtime-cu12==12.4.127<br>
[pip3] nvidia-cudnn-cu12==9.1.0.70<br>
[pip3] nvidia-cufft-cu12==11.2.1.3<br>
[pip3] nvidia-curand-cu12==10.3.5.147<br>
[pip3] nvidia-cusolver-cu12==11.6.1.9<br>
[pip3] nvidia-cusparse-cu12==12.3.1.170<br>
[pip3] nvidia-cusparselt-cu12==0.6.2<br>
[pip3] nvidia-nccl-cu12==2.21.5<br>
[pip3] nvidia-nvjitlink-cu12==12.4.127<br>
[pip3] nvidia-nvtx-cu12==12.4.127<br>
[pip3] torch==2.6.0+cu124<br>
[pip3] torchaudio==2.6.0+cu124<br>
[pip3] torchvision==0.21.0+cu124<br>
[pip3] triton==3.2.0</p>