pytorch - 💡(How to fix) Fix torch.compile(create_block_mask) has pathological cold compile time for FlexAttention causal block masks at some sequence lengths [1 participants]

🐛 Describe the bug

torch.compile(torch.nn.attention.flex_attention.create_block_mask) can take several minutes or appear to hang for a simple causal mask at some sequence lengths.

This reproduces with a minimal script using only create_block_mask and a trivial causal mask_mod.

Reproduction:

import argparse
import time

import torch
from torch.nn.attention.flex_attention import create_block_mask


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("seq_len", type=int)
    return parser.parse_args()


def causal_mask_mod(
    batch: torch.Tensor,
    head: torch.Tensor,
    query_index: torch.Tensor,
    kv_index: torch.Tensor,
) -> torch.Tensor:
    return query_index >= kv_index


args = parse_args()

print("starting compile")
compiled = torch.compile(create_block_mask)


print(f"calling compiled create_block_mask with Q_LEN=KV_LEN={args.seq_len}")
call_start = time.perf_counter()
block_mask = compiled(
    causal_mask_mod,
    B=1,
    H=1,
    Q_LEN=args.seq_len,
    KV_LEN=args.seq_len,
    device="cuda",
    BLOCK_SIZE=(128, 128),
)
call_elapsed = time.perf_counter() - call_start
print(f"compiled create_block_mask returned in {call_elapsed:.6f}s")

torch.cuda.synchronize()
print("done")

Running this on an A100 SXM4:

rm -rf /tmp/torchinductor_root && rm -rf /root/.triton/cache
python block_mask_debug.py 13567

we get:

starting compile
calling compiled create_block_mask with Q_LEN=KV_LEN=13567
compiled create_block_mask returned in 573.226887s
done

Same results for 13568, 13696. For 17407 it completes in 20s. Same results with fullgraph=True.

Versions

Collecting environment information... PyTorch version: 2.13.0.dev20260429+cu130 Is debug build: False CUDA used to build PyTorch: 13.0 ROCM used to build PyTorch: N/A

OS: Ubuntu 24.04.4 LTS (x86_64) GCC version: (Ubuntu 13.3.0-6ubuntu2~24.04.1) 13.3.0 Clang version: Could not collect CMake version: version 3.28.3 Libc version: glibc-2.39

Python version: 3.12.13 | packaged by conda-forge | (main, Mar 5 2026, 16:50:00) [GCC 14.3.0] (64-bit runtime) Python platform: Linux-5.15.0-161-generic-x86_64-with-glibc2.39 Is CUDA available: True CUDA runtime version: 13.0.88 CUDA_MODULE_LOADING set to: GPU models and configuration: GPU 0: NVIDIA A100-SXM4-40GB Nvidia driver version: 580.95.05 cuDNN version: Probably one of the following: /usr/lib/x86_64-linux-gnu/libcudnn.so.9.14.0 /usr/lib/x86_64-linux-gnu/libcudnn_adv.so.9.14.0 /usr/lib/x86_64-linux-gnu/libcudnn_cnn.so.9.14.0 /usr/lib/x86_64-linux-gnu/libcudnn_engines_precompiled.so.9.14.0 /usr/lib/x86_64-linux-gnu/libcudnn_engines_runtime_compiled.so.9.14.0 /usr/lib/x86_64-linux-gnu/libcudnn_graph.so.9.14.0 /usr/lib/x86_64-linux-gnu/libcudnn_heuristic.so.9.14.0 /usr/lib/x86_64-linux-gnu/libcudnn_ops.so.9.14.0 Is XPU available: False HIP runtime version: N/A MIOpen runtime version: N/A Is XNNPACK available: True Caching allocator config: N/A

CPU: Architecture: x86_64 CPU op-mode(s): 32-bit, 64-bit Address sizes: 48 bits physical, 48 bits virtual Byte Order: Little Endian CPU(s): 128 On-line CPU(s) list: 0-127 Vendor ID: AuthenticAMD Model name: AMD EPYC 7532 32-Core Processor CPU family: 23 Model: 49 Thread(s) per core: 2 Core(s) per socket: 32 Socket(s): 2 Stepping: 0 BogoMIPS: 4790.91 Flags: fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ht syscall nx mmxext fxsr_opt pdpe1gb rdtscp lm constant_tsc rep_good nopl nonstop_tsc cpuid extd_apicid aperfmperf rapl pni pclmulqdq monitor ssse3 fma cx16 sse4_1 sse4_2 movbe popcnt aes xsave avx f16c rdrand lahf_lm cmp_legacy svm extapic cr8_legacy abm sse4a misalignsse 3dnowprefetch osvw ibs skinit wdt tce topoext perfctr_core perfctr_nb bpext perfctr_llc mwaitx cpb cat_l3 cdp_l3 hw_pstate ssbd mba ibrs ibpb stibp vmmcall fsgsbase bmi1 avx2 smep bmi2 cqm rdt_a rdseed adx smap clflushopt clwb sha_ni xsaveopt xsavec xgetbv1 xsaves cqm_llc cqm_occup_llc cqm_mbm_total cqm_mbm_local clzero irperf xsaveerptr rdpru wbnoinvd amd_ppin arat npt lbrv svm_lock nrip_save tsc_scale vmcb_clean flushbyasid decodeassists pausefilter pfthreshold avic v_vmsave_vmload vgif v_spec_ctrl umip rdpid overflow_recov succor smca ibpb_exit_to_user Virtualization: AMD-V L1d cache: 2 MiB (64 instances) L1i cache: 2 MiB (64 instances) L2 cache: 32 MiB (64 instances) L3 cache: 512 MiB (32 instances) NUMA node(s): 8 NUMA node0 CPU(s): 0-7,64-71 NUMA node1 CPU(s): 8-15,72-79 NUMA node2 CPU(s): 16-23,80-87 NUMA node3 CPU(s): 24-31,88-95 NUMA node4 CPU(s): 32-39,96-103 NUMA node5 CPU(s): 40-47,104-111 NUMA node6 CPU(s): 48-55,112-119 NUMA node7 CPU(s): 56-63,120-127 Vulnerability Gather data sampling: Not affected Vulnerability Indirect target selection: Not affected Vulnerability Itlb multihit: Not affected Vulnerability L1tf: Not affected Vulnerability Mds: Not affected Vulnerability Meltdown: Not affected Vulnerability Mmio stale data: Not affected Vulnerability Reg file data sampling: Not affected Vulnerability Retbleed: Mitigation; untrained return thunk; SMT enabled with STIBP protection Vulnerability Spec rstack overflow: Mitigation; safe RET Vulnerability Spec store bypass: Mitigation; Speculative Store Bypass disabled via prctl and seccomp Vulnerability Spectre v1: Mitigation; usercopy/swapgs barriers and __user pointer sanitization Vulnerability Spectre v2: Mitigation; Retpolines; IBPB conditional; STIBP always-on; RSB filling; PBRSB-eIBRS Not affected; BHI Not affected Vulnerability Srbds: Not affected Vulnerability Tsa: Not affected Vulnerability Tsx async abort: Not affected Vulnerability Vmscape: Mitigation; IBPB before exit to userspace

Versions of relevant libraries: [pip3] numpy==2.4.3 [pip3] nvidia-cublas==13.1.0.3 [pip3] nvidia-cuda-cupti==13.0.85 [pip3] nvidia-cuda-nvrtc==13.0.88 [pip3] nvidia-cuda-runtime==13.0.96 [pip3] nvidia-cudnn-cu13==9.20.0.48 [pip3] nvidia-cufft==12.0.0.61 [pip3] nvidia-curand==10.4.0.35 [pip3] nvidia-cusolver==12.0.4.66 [pip3] nvidia-cusparse==12.6.3.3 [pip3] nvidia-cusparselt-cu13==0.8.1 [pip3] nvidia-nccl-cu13==2.29.7 [pip3] nvidia-nvjitlink==13.0.88 [pip3] nvidia-nvtx==13.0.85 [pip3] torch==2.13.0.dev20260429+cu130 [pip3] torchaudio==2.11.0+cu130 [pip3] torchcodec==0.11.0+cu130 [pip3] torchvision==0.27.0.dev20260430+cu130 [pip3] triton==3.7.0+git88b227e2 [conda] numpy 2.4.3 pypi_0 pypi [conda] nvidia-cublas 13.1.0.3 pypi_0 pypi [conda] nvidia-cuda-cupti 13.0.85 pypi_0 pypi [conda] nvidia-cuda-nvrtc 13.0.88 pypi_0 pypi [conda] nvidia-cuda-runtime 13.0.96 pypi_0 pypi [conda] nvidia-cudnn-cu13 9.20.0.48 pypi_0 pypi [conda] nvidia-cufft 12.0.0.61 pypi_0 pypi [conda] nvidia-curand 10.4.0.35 pypi_0 pypi [conda] nvidia-cusolver 12.0.4.66 pypi_0 pypi [conda] nvidia-cusparse 12.6.3.3 pypi_0 pypi [conda] nvidia-cusparselt-cu13 0.8.1 pypi_0 pypi [conda] nvidia-nccl-cu13 2.29.7 pypi_0 pypi [conda] nvidia-nvjitlink 13.0.88 pypi_0 pypi [conda] nvidia-nvtx 13.0.85 pypi_0 pypi [conda] torch 2.13.0.dev20260429+cu130 pypi_0 pypi [conda] torchaudio 2.11.0+cu130 pypi_0 pypi [conda] torchcodec 0.11.0+cu130 pypi_0 pypi [conda] torchvision 0.27.0.dev20260430+cu130 pypi_0 pypi [conda] triton 3.7.0+git88b227e2 pypi_0 pypi

cc @jerryzh168 @chauhang @penguinwu @oulgen @jamesjwu @aorenste @anijain2305 @laithsakka @masnesral @coconutruben @aditvenk @Chillee @drisspg @yanboliang @BoyuanFeng @liangel-02 @howardzhang-cv