Repository: pytorch
Issue: FlexAttention `create_block_mask` contains a CUDA sync
Labels: triaged,oncall: pt2,module: higher order operators,module: pt2-dispatcher,module: flex attention

### 🐛 Describe the bug

I am trying to capture our model forward pass into a CUDA graph, but Flex Attention's `create_block_mask` contains a graph break.

I'm honestly not sure if this is a "bug" or a "feature request".

I have tested `create_block_mask` both with and without `_compile=True` and it happens in both cases.

Relevant stack trace:
```
   RuntimeError: called a synchronizing CUDA operation

   While executing %setitem : [num_users=0] = call_function[target=operator.setitem](args = (%dense_mask_2, (%row_indi
   ces, %valid_indices), 1), kwargs = {})
   Original traceback:
     File "/opt/conda/envs/main-env/lib/python3.11/site-packages/torch/nn/attention/flex_attention.py", line 893, in c
   reate_block_mask
       block_mask = _create_sparse_block_from_block_mask(
     File "/opt/conda/envs/main-env/lib/python3.11/site-packages/torch/nn/attention/flex_attention.py", line 765, in _
   create_sparse_block_from_block_mask
       return BlockMask.from_kv_blocks(
     File "/opt/conda/envs/main-env/lib/python3.11/site-packages/torch/nn/attention/flex_attention.py", line 353, in f
   rom_kv_blocks
       q_num_blocks, q_indices = _transpose_ordered(kv_num_blocks, kv_indices)
     File "/opt/conda/envs/main-env/lib/python3.11/site-packages/torch/nn/attention/flex_attention.py", line 187, in _
   transpose_ordered
       dense = _ordered_to_dense(num_blocks_in_row, col_indices)
     File "/opt/conda/envs/main-env/lib/python3.11/site-packages/torch/nn/attention/flex_attention.py", line 172, in _
   ordered_to_dense
       out = create_dense_batched(num_blocks_in_row, col_indices)
     File "/opt/conda/envs/main-env/lib/python3.11/site-packages/torch/_functorch/apis.py", line 203, in wrapped
       return vmap_impl(
     File "/opt/conda/envs/main-env/lib/python3.11/site-packages/torch/_functorch/vmap.py", line 331, in vmap_impl
       return _flat_vmap(
     File "/opt/conda/envs/main-env/lib/python3.11/site-packages/torch/_functorch/vmap.py", line 479, in _flat_vmap
       batched_outputs = func(*batched_inputs, **kwargs)
     File "/opt/conda/envs/main-env/lib/python3.11/site-packages/torch/_functorch/apis.py", line 203, in wrapped
       return vmap_impl(
     File "/opt/conda/envs/main-env/lib/python3.11/site-packages/torch/_functorch/vmap.py", line 331, in vmap_impl
       return _flat_vmap(
     File "/opt/conda/envs/main-env/lib/python3.11/site-packages/torch/_functorch/vmap.py", line 479, in _flat_vmap
       batched_outputs = func(*batched_inputs, **kwargs)
     File "/opt/conda/envs/main-env/lib/python3.11/site-packages/torch/nn/attention/flex_attention.py", line 165, in c
   reate_dense_one
       dense_mask[row_indices, valid_indices] = 1
```

### Versions

```
   Collecting environment information...
   PyTorch version: 2.6.0.dev20241211+cu126
   Is debug build: False
   CUDA used to build PyTorch: 12.6
   ROCM used to build PyTorch: N/A

   OS: Ubuntu 20.04.6 LTS (x86_64)
   GCC version: (Ubuntu 9.4.0-1ubuntu1~20.04.2) 9.4.0
   Clang version: Could not collect
   CMake version: version 3.16.3
   Libc version: glibc-2.31

   Python version: 3.11.10 | packaged by conda-forge | (main, Oct 16 2024, 01:27:36) [GCC 13.3.0] (64-bit runtime)
   Python platform: Linux-6.8.0-49-generic-x86_64-with-glibc2.31
   Is CUDA available: True
   CUDA runtime version: 12.6.77
   CUDA_MODULE_LOADING set to: LAZY
   GPU models and configuration: GPU 0: NVIDIA A100-SXM4-40GB
   Nvidia driver version: 550.127.05
   cuDNN version: Could not collect
   HIP runtime version: N/A
   MIOpen runtime version: N/A
   Is XNNPACK available: True

   CPU:
   Architecture:                         x86_64
   CPU op-mode(s):                       32-bit, 64-bit
   Byte Order:                           Little Endian
   Address sizes:                        48 bits physical, 48 bits virtual
   CPU(s):                               30
   On-line CPU(s) list:                  0-29
   Thread(s) per core:                   1
   Core(s) per socket:                   1
   Socket(s):                            30
   NUMA node(s):                         1
   Vendor ID:                            AuthenticAMD
   CPU family:                           25
   Model:                                1
   Model name:                           AMD EPYC 7J13 64-Core Processor
   Stepping:                             1
   CPU MHz:                              2449.998
   BogoMIPS:                             4899.99
   Virtualization:                       AMD-V
   Hypervisor vendor:                    KVM
   Virtualization type:                  full
   L1d cache:                            1.9 MiB
   L1i cache:                            1.9 MiB
   L2 cache:                             15 MiB
   L3 cache:                             480 MiB
   NUMA node0 CPU(s):                    0-29
   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: Safe RET, no microcode
   Vulnerability Spec store bypass:      Mitigation; Speculative Store Bypass disabled via prctl
   Vulnerability Spectre v1:             Mitigation; usercopy/swapgs barriers and __user pointer sanitization
   Vulnerability Spectre v2:             Mitigation; Retpolines; IBPB conditional; IBRS_FW; STIBP disabled; RSB filling; PBRSB-eIBRS Not affected; BHI Not affected
   Vulnerability Srbds:                  Not affected
   Vulnerability Tsx async abort:        Not affected
   Flags:                                fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 syscall nx mmxext fxsr_opt pdpe1gb rdtscp lm rep_good nopl cpuid extd_apicid tsc_known_freq pni pclmulqdq
   ssse3 fma cx16 pcid sse4_1 sse4_2 x2apic movbe popcnt tsc_deadline_timer aes xsave avx f16c rdrand hypervisor lahf_lm cmp_legacy svm cr8_legacy abm sse4a misalignsse 3dnowprefetch osvw perfctr_core ssbd ibrs ibpb stibp vmmcall fsgsbase tsc_ad
   just bmi1 avx2 smep bmi2 erms invpcid rdseed adx smap clflushopt clwb sha_ni xsaveopt xsavec xgetbv1 xsaves clzero xsaveerptr wbnoinvd arat npt nrip_save umip pku ospke vaes vpclmulqdq rdpid fsrm arch_capabilities

   Versions of relevant libraries:
   [pip3] flake8==7.1.1
   [pip3] mypy-extensions==1.0.0
   [pip3] numpy==1.26.4
   [pip3] nvidia-cublas-cu12==12.6.4.1
   [pip3] nvidia-cuda-cupti-cu12==12.6.80
   [pip3] nvidia-cuda-nvrtc-cu12==12.6.77
   [pip3] nvidia-cuda-runtime-cu12==12.6.77
   [pip3] nvidia-cudnn-cu12==9.5.1.17
   [pip3] nvidia-cufft-cu12==11.3.0.4
   [pip3] nvidia-curand-cu12==10.3.7.77
   [pip3] nvidia-cusolver-cu12==11.7.1.2
   [pip3] nvidia-cusparse-cu12==12.5.4.2
   [pip3] nvidia-cusparselt-cu12==0.6.3
   [pip3] nvidia-dlprof-pytorch-nvtx==1.8.0
   [pip3] nvidia-nccl-cu12==2.21.5
   [pip3] nvidia-nvjitlink-cu12==12.6.85
   [pip3] nvidia-nvtx-cu12==12.6.77
   [pip3] onnx==1.17.0
   [pip3] pytorch-ignite==0.5.0.post2
   [pip3] pytorch-lightning==2.4.0
   [pip3] pytorch-triton==3.2.0+git35c6c7c6
   [pip3] torch==2.6.0.dev20241211+cu126
   [pip3] torch-stoi==0.2.3
   [pip3] torchaudio==2.5.0.dev20241211+cu126
   [pip3] torchcde==0.2.5
   [pip3] torchcfm==1.0.6
   [pip3] torchdiffeq==0.2.0
   [pip3] torchdyn==1.0.6
   [pip3] torchmetrics==1.5.2
   [pip3] torchsde==0.2.6
   [pip3] torchvision==0.20.0.dev20241211+cu126
   [pip3] triton==3.1.0
   [conda] blas                      1.0                         mkl    conda-forge
   [conda] cuda-cudart               12.1.105                      0    nvidia
   [conda] cuda-cupti                12.1.105                      0    nvidia
   [conda] cuda-libraries            12.1.0                        0    nvidia
   [conda] cuda-nvrtc                12.1.105                      0    nvidia
   [conda] cuda-nvtx                 12.1.105                      0    nvidia
   [conda] cuda-opencl               12.6.77                       0    nvidia
   [conda] cuda-runtime              12.1.0                        0    nvidia
   [conda] libblas                   3.9.0            16_linux64_mkl    conda-forge
   [conda] libcblas                  3.9.0            16_linux64_mkl    conda-forge
   [conda] libcublas                 12.1.0.26                     0    nvidia
   [conda] libcufft                  11.0.2.4                      0    nvidia
   [conda] libcurand                 10.3.7.77                     0    nvidia
   [conda] libcusolver               11.4.4.55                     0    nvidia
   [conda] libcusparse               12.0.2.55                     0    nvidia
   [conda] liblapack                 3.9.0            16_linux64_mkl    conda-forge
   [conda] libnvjitlink              12.1.105                      0    nvidia
   [conda] libopenvino-pytorch-frontend 2024.3.0             he02047a_0    conda-forge
   [conda] mkl                       2022.1.0           hc2b9512_224
   [conda] numpy                     1.26.4          py311h64a7726_0    conda-forge
   [conda] nvidia-cublas-cu12        12.6.4.1                 pypi_0    pypi
   [conda] nvidia-cuda-cupti-cu12    12.6.80                  pypi_0    pypi
   [conda] nvidia-cuda-nvrtc-cu12    12.6.77                  pypi_0    pypi
   [conda] nvidia-cuda-runtime-cu12  12.6.77                  pypi_0    pypi
   [conda] nvidia-cudnn-cu12         9.5.1.17                 pypi_0    pypi
   [conda] nvidia-cufft-cu12         11.3.0.4                 pypi_0    pypi
   [conda] nvidia-curand-cu12        10.3.7.77                pypi_0    pypi
   [conda] nvidia-cusolver-cu12      11.7.1.2                 pypi_0    pypi
   [conda] nvidia-cusparse-cu12      12.5.4.2                 pypi_0    pypi
   [conda] nvidia-cusparselt-cu12    0.6.3                    pypi_0    pypi
   [conda] nvidia-dlprof-pytorch-nvtx 1.8.0                    pypi_0    pypi
   [conda] nvidia-nccl-cu12          2.21.5                   pypi_0    pypi
   [conda] nvidia-nvjitlink-cu12     12.6.85                  pypi_0    pypi
   [conda] nvidia-nvtx-cu12          12.6.77                  pypi_0    pypi
   [conda] pytorch-cuda              12.1                 ha16c6d3_6    pytorch
   [conda] pytorch-ignite            0.5.0.post2              pypi_0    pypi
   [conda] pytorch-lightning         2.4.0              pyhd8ed1ab_0    conda-forge
   [conda] pytorch-mutex             1.0                        cuda    pytorch
   [conda] pytorch-triton            3.2.0+git35c6c7c6          pypi_0    pypi
   [conda] torch                     2.6.0.dev20241211+cu126          pypi_0    pypi
   [conda] torch-stoi                0.2.3                    pypi_0    pypi
   [conda] torchaudio                2.5.0.dev20241211+cu126          pypi_0    pypi
   [conda] torchcde                  0.2.5                    pypi_0    pypi
   [conda] torchcfm                  1.0.6                    pypi_0    pypi
   [conda] torchdiffeq               0.2.0                    pypi_0    pypi
   [conda] torchdyn                  1.0.6                    pypi_0    pypi
   [conda] torchmetrics              1.5.2              pyhe5570ce_0    conda-forge
   [conda] torchsde                  0.2.6                    pypi_0    pypi
   [conda] torchtriton               3.1.0                     py311    pytorch
   [conda] torchvision               0.20.0.dev20241211+cu126          pypi_0    pypi
```

cc @chauhang @penguinwu @zou3519 @ydwu4 @bdhirsh @yf225 @Chillee @drisspg @yanboliang @BoyuanFeng