FlexAttention + custom score_mod fails under torch.compile (FlexibleLayout / NoValidChoicesError)

[v1kstrand]

Hi! I filed a compiler issue on PyTorch here.

I’m using torch.nn.attention.flex_attention with a small learned CPB score_mod. Eager works, but torch.compile (Inductor) fails with:

• AssertionError: convert FlexibleLayout to FixedLayout first
• NoValidChoicesError (with mode="max-autotune-no-cudagraphs")

Minimal repro: same as in the linked issue (inline below if preferred).

Environment (summary)

• torch 2.6.0+cu124, triton 3.2.0
• CUDA 12.4, cuDNN 90100
• GPU: A100 80GB
• Dtypes tried: bf16 (autocast) and fp16 → same compile errors
• Backend: SDPBackend.FLASH_ATTENTION (can try MATH/CUDNN if that helps)

Question

• Is this score_mod pattern/layout known to be unsupported for compile?
• Any Attention-Gym example or workaround (e.g., layout constraints, alternate backend, or a score_mod idiom) I should follow?

I’ll keep this thread updated and close it once the PyTorch issue is resolved. Thanks!