Repository: pytorch
Issue: `_adjust_num_blocks_and_indices` gives wrong adjusted block mask
Labels: triaged,oncall: pt2,module: higher order operators,module: pt2-dispatcher,module: flex attention

### 🐛 Describe the bug

Repro:
```python
import torch
from torch.nn.attention.flex_attention import (
    _adjust_num_blocks_and_indices,
    create_block_mask,
)
def mask_mod(b, h, q, kv):
    b1 = torch.where(kv >= 0, True, False) & torch.where(kv < 128, True, False)
    b2 = torch.where(kv >= 384, True, False) & torch.where(kv < 512, True, False)
    return b1 | b2
block_mask = create_block_mask(mask_mod, 1, 1, 128, 640)
adjusted_full_kv_num_blocks, adjusted_full_kv_indices = _adjust_num_blocks_and_indices(
    block_mask.full_kv_num_blocks,
    block_mask.full_kv_indices,
    new_num_rows=1,
    new_num_cols=3,
)
print(f"original full_kv_num_blocks: {block_mask.full_kv_num_blocks}")
print(f"original full_kv_indices: {block_mask.full_kv_indices}")
print(f"adjusted full_kv_num_blocks: {adjusted_full_kv_num_blocks}")
print(f"adjusted full_kv_indices: {adjusted_full_kv_indices}")

"""
original full_kv_num_blocks: tensor([[[2]]], device='cuda:0', dtype=torch.int32)
original full_kv_indices: tensor([[[[0, 3, 1, 2, 4]]]], device='cuda:0', dtype=torch.int32)
adjusted full_kv_num_blocks: tensor([[[2]]], device='cuda:0', dtype=torch.int32)
adjusted full_kv_indices: tensor([[[[0, 3, 1]]]], device='cuda:0', dtype=torch.int32)
Expected adjusted full_kv_num_blocks: tensor([[[1]]], device='cuda:0', dtype=torch.int32)
"""
```

### Versions

PyTorch version: 2.6.0a0+git853e0fa
Is debug build: False
CUDA used to build PyTorch: 12.0
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-2)
Clang version: Could not collect
CMake version: version 3.26.4
Libc version: glibc-2.34

Python version: 3.10.15 (main, Oct  3 2024, 07:27:34) [GCC 11.2.0] (64-bit runtime)
Python platform: Linux-5.19.0-0_fbk12_hardened_11583_g0bef9520ca2b-x86_64-with-glibc2.34
Is CUDA available: True
CUDA runtime version: 12.0.140
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: 
GPU 0: NVIDIA H100
GPU 1: NVIDIA H100
GPU 2: NVIDIA H100
GPU 3: NVIDIA H100

Nvidia driver version: 525.105.17
cuDNN version: Probably one of the following:
/usr/lib64/libcudnn.so.9.1.1
/usr/lib64/libcudnn_adv.so.9.1.1
/usr/lib64/libcudnn_cnn.so.9.1.1
/usr/lib64/libcudnn_engines_precompiled.so.9.1.1
/usr/lib64/libcudnn_engines_runtime_compiled.so.9.1.1
/usr/lib64/libcudnn_graph.so.9.1.1
/usr/lib64/libcudnn_heuristic.so.9.1.1
/usr/lib64/libcudnn_ops.so.9.1.1
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
Address sizes:                   52 bits physical, 57 bits virtual
Byte Order:                      Little Endian
CPU(s):                          184
On-line CPU(s) list:             0-183
Vendor ID:                       AuthenticAMD
Model name:                      AMD EPYC 9654 96-Core Processor
CPU family:                      25
Model:                           17
Thread(s) per core:              1
Core(s) per socket:              184
Socket(s):                       1
Stepping:                        1
BogoMIPS:                        4792.80
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 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 invpcid_single ssbd ibrs ibpb stibp vmmcall fsgsbase tsc_adjust bmi1 avx2 smep bmi2 erms invpcid avx512f avx512dq rdseed adx smap avx512ifma clflushopt clwb avx512cd sha_ni avx512bw avx512vl xsaveopt xsavec xgetbv1 xsaves avx512_bf16 clzero xsaveerptr wbnoinvd arat npt lbrv nrip_save tsc_scale vmcb_clean pausefilter pfthreshold v_vmsave_vmload vgif avx512vbmi umip pku ospke avx512_vbmi2 gfni vaes vpclmulqdq avx512_vnni avx512_bitalg avx512_vpopcntdq rdpid fsrm arch_capabilities
Virtualization:                  AMD-V
Hypervisor vendor:               KVM
Virtualization type:             full
L1d cache:                       11.5 MiB (184 instances)
L1i cache:                       11.5 MiB (184 instances)
L2 cache:                        92 MiB (184 instances)
L3 cache:                        2.9 GiB (184 instances)
NUMA node(s):                    1
NUMA node0 CPU(s):               0-183
Vulnerability Itlb multihit:     Not affected
Vulnerability L1tf:              Not affected
Vulnerability Mds:               Not affected
Vulnerability Meltdown:          Not affected
Vulnerability Mmio stale data:   Not affected
Vulnerability Retbleed:          Not affected
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
Vulnerability Srbds:             Not affected
Vulnerability Tsx async abort:   Not affected

Versions of relevant libraries:
[pip3] flake8==6.1.0
[pip3] flake8-bugbear==23.3.23
[pip3] flake8-comprehensions==3.15.0
[pip3] flake8-executable==2.1.3
[pip3] flake8-logging-format==0.9.0
[pip3] flake8-pyi==23.3.1
[pip3] flake8-simplify==0.19.3
[pip3] mypy==1.11.2
[pip3] mypy-extensions==1.0.0
[pip3] numpy==1.26.0
[pip3] optree==0.13.0
[pip3] pytorch-triton==3.1.0+cf34004b8a
[pip3] torch==2.6.0a0+git853e0fa
[conda] blas                      1.0                         mkl  
[conda] magma-cuda121             2.6.1                         1    pytorch
[conda] mkl                       2023.1.0         h213fc3f_46344  
[conda] mkl-include               2023.1.0         h06a4308_46344  
[conda] mkl-service               2.4.0           py310h5eee18b_1  
[conda] mkl_fft                   1.3.10          py310h5eee18b_0  
[conda] mkl_random                1.2.7           py310h1128e8f_0  
[conda] numpy                     1.26.0                   pypi_0    pypi
[conda] optree                    0.13.0                   pypi_0    pypi
[conda] pytorch-triton            3.1.0+cf34004b8a          pypi_0    pypi
[conda] torch                     2.6.0a0+git853e0fa           dev_0    <develop>
[conda] torchfix                  0.4.0                    pypi_0    pypi

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