🐛 Describe the bug

On torch nightly, we run into a new CI failure in torchtitan (see https://github.com/pytorch/torchtitan/issues/2722)

The crash looks like:

python: /source/llvm-project/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp:2690: int llvm::slpvectorizer::BoUpSLP::LookAheadHeuristics::getScoreAtLevelRec(Value *, Value *, Instruction *, Instruction *, int, ArrayRef<Value *>) const: Assertion `FromIdx <= ToIdx && "Bad index"' failed.
Aborted (core dumped)

We find a workaround is to disable llvm opt - further, we were able to minimize to just:

import os
# Ensure not set so we repro
os.environ.pop("DISABLE_LLVM_OPT", None)

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

_compiled_create_block_mask = torch.compile(create_block_mask)


def causal_mask(b, h, q_idx, kv_idx):
    return q_idx >= kv_idx


def main():
    print("Building attention mask...")
    block_mask = _compiled_create_block_mask(
        causal_mask, 8, None, 2048, 2048
    )
    print("Done — no crash.")


if __name__ == "__main__":
    main()

Additional Context

Changing batch size (8) or seq len (2048) smaller makes this repro disappear

Claude has the following analysis:

Workaround for Triton 3.7 LLVM SLP vectorizer crash on flex attention kernels. The Triton pin update (https://github.com/pytorch/pytorch/pull/174896) ships an LLVM build whose SLP vectorizer hits an assertion (FromIdx <= ToIdx, "Bad index") when compiling certain flex attention Triton kernels. The crash occurs in llvm.optimize_module() during make_llir, which runs the LLVM O3 pipeline including SLP vectorization. There is no Python-level API to disable only the SLP vectorizer pass (optimize_module ignores its flags parameter, and DISABLE_LLVM_OPT= only sets flags to true). Setting DISABLE_LLVM_OPT=1 skips the LLVM O3 pipeline entirely; the MLIR-level optimizations still run, so GPU kernel performance impact is minimal.

Versions

Collecting environment information... PyTorch version: 2.12.0a0+gite7bcaf2 Is debug build: False CUDA used to build PyTorch: 12.9 ROCM used to build PyTorch: N/A

OS: CentOS Stream 9 (x86_64) GCC version: (GCC) 11.5.0 20240719 (Red Hat 11.5.0-14) Clang version: 21.1.7 (CentOS 21.1.7-1.el9) CMake version: version 4.3.0 Libc version: glibc-2.34

Python version: 3.10.20 (main, Mar 11 2026, 17:46:40) [GCC 14.3.0] (64-bit runtime) Python platform: Linux-6.13.2-0_fbk12_0_g0b66b3635210-x86_64-with-glibc2.34 Is CUDA available: True CUDA runtime version: 12.9.86 CUDA_MODULE_LOADING set to: GPU models and configuration: GPU 0: NVIDIA H100 GPU 1: NVIDIA H100 GPU 2: NVIDIA H100 GPU 3: NVIDIA H100 GPU 4: NVIDIA H100 GPU 5: NVIDIA H100 GPU 6: NVIDIA H100 GPU 7: NVIDIA H100

Nvidia driver version: 580.82.07 cuDNN version: Could not collect 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: 52 bits physical, 57 bits virtual Byte Order: Little Endian CPU(s): 384 On-line CPU(s) list: 0-383 Vendor ID: AuthenticAMD Model name: AMD EPYC 9654 96-Core Processor CPU family: 25 Model: 17 Thread(s) per core: 2 Core(s) per socket: 96 Socket(s): 2 Stepping: 1 Frequency boost: enabled CPU(s) scaling MHz: 100% CPU max MHz: 2400.0000 CPU min MHz: 1500.0000 BogoMIPS: 4792.43 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 amd_lbr_v2 nopl xtopology nonstop_tsc cpuid extd_apicid aperfmperf rapl pni pclmulqdq monitor ssse3 fma cx16 pcid sse4_1 sse4_2 x2apic 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 perfmon_v2 ibrs ibpb stibp ibrs_enhanced vmmcall fsgsbase bmi1 avx2 smep bmi2 erms invpcid cqm rdt_a avx512f avx512dq rdseed adx smap avx512ifma clflushopt clwb avx512cd sha_ni avx512bw avx512vl xsaveopt xsavec xgetbv1 xsaves cqm_llc cqm_occup_llc cqm_mbm_total cqm_mbm_local avx512_bf16 clzero irperf xsaveerptr rdpru wbnoinvd amd_ppin cppc arat npt lbrv svm_lock nrip_save tsc_scale vmcb_clean flushbyasid decodeassists pausefilter pfthreshold avic v_vmsave_vmload vgif x2avic v_spec_ctrl vnmi avx512vbmi umip pku ospke avx512_vbmi2 gfni vaes vpclmulqdq avx512_vnni avx512_bitalg avx512_vpopcntdq rdpid overflow_recov succor smca fsrm flush_l1d debug_swap Virtualization: AMD-V L1d cache: 6 MiB (192 instances) L1i cache: 6 MiB (192 instances) L2 cache: 192 MiB (192 instances) L3 cache: 768 MiB (24 instances) NUMA node(s): 2 NUMA node0 CPU(s): 0-95,192-287 NUMA node1 CPU(s): 96-191,288-383 Vulnerability Gather data sampling: 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: Not affected Vulnerability Spec rstack overflow: Vulnerable Vulnerability Spec store bypass: Vulnerable Vulnerability Spectre v1: Vulnerable: __user pointer sanitization and usercopy barriers only; no swapgs barriers Vulnerability Spectre v2: Vulnerable; IBPB: disabled; STIBP: disabled; PBRSB-eIBRS: Not affected; BHI: Not affected Vulnerability Srbds: Not affected Vulnerability Tsx async abort: Not affected

Versions of relevant libraries: [pip3] intel-cmplr-lib-ur==2025.3.3 [pip3] intel-openmp==2025.3.3 [pip3] mkl-include==2025.3.1 [pip3] mkl-static==2025.3.1 [pip3] mypy_extensions==1.1.0 [pip3] numpy==2.2.6 [pip3] nvidia-cudnn-frontend==1.18.0 [pip3] onemkl-license==2025.3.1 [pip3] optree==0.19.0 [pip3] tbb==2022.3.1 [pip3] tbb-devel==2022.3.1 [pip3] tcmlib==1.4.1 [pip3] torch==2.12.0a0+gite7bcaf2 [pip3] torch_c_dlpack_ext==0.1.5 [pip3] torchdata==0.11.0 [pip3] torchmonarch==0.3.0 [pip3] torchstore==0.0.0.dev0 [pip3] torchtitan==0.2.2 [pip3] torchx-nightly==2026.3.26 [pip3] triton==3.6.0+git7e48d5df [pip3] umf==1.0.3 [conda] intel-cmplr-lib-ur 2025.3.3 pypi_0 pypi [conda] intel-openmp 2025.3.3 pypi_0 pypi [conda] mkl-include 2025.3.1 pypi_0 pypi [conda] mkl-static 2025.3.1 pypi_0 pypi [conda] numpy 2.2.6 pypi_0 pypi [conda] nvidia-cudnn-frontend 1.18.0 pypi_0 pypi [conda] onemkl-license 2025.3.1 pypi_0 pypi [conda] optree 0.19.0 pypi_0 pypi [conda] tbb 2022.3.1 pypi_0 pypi [conda] tbb-devel 2022.3.1 pypi_0 pypi [conda] tcmlib 1.4.1 pypi_0 pypi [conda] torch 2.12.0a0+gite7bcaf2 pypi_0 pypi [conda] torch-c-dlpack-ext 0.1.5 pypi_0 pypi [conda] torchdata 0.11.0 pypi_0 pypi [conda] torchmonarch 0.3.0 pypi_0 pypi [conda] torchstore 0.0.0.dev0 pypi_0 pypi [conda] torchtitan 0.2.2 pypi_0 pypi [conda] torchx-nightly 2026.3.26 pypi_0 pypi [conda] triton 3.6.0+git7e48d5df pypi_0 pypi [conda] umf 1.0.3 pypi_0 pypi

cc @chauhang @penguinwu @voznesenskym @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @ipiszy @kadeng @muchulee8 @amjames @aakhundov @coconutruben @jataylo @bertmaher @int3 @davidberard98 @nmacchioni @embg @peterbell10 @iupaikov-amd @Chillee @drisspg @yanboliang @BoyuanFeng @liangel-02 @howardzhang-cv