Repository: pytorch
Issue: Illegal Memory Access when Using Trainable Biases in Flex Attention
Labels: module: crash,module: cuda,triaged,module: flex attention

### 🐛 Describe the bug

I'm back with another insane flex attention bug report.

Recently, I was playing around with the learnable biases in flex attention when I started hitting an illegal memory access in the backward pass.

After using `CUDA_LAUNCH_BLOCKING` I found it was happening in the following kernel during autotuning:

```python
File "/tmp/torchinductor_ubuntu/jy/cjy56z23prjyls6tnwn4ay4mmmb6vvergqxm4wmnv5l7zlfzk66e.py", line 3784, in call
    triton_poi_fused_zeros_1.run(buf0, 100663296, grid=grid(100663296), stream=stream0)
```

```python
import triton
import triton.language as tl
from triton.compiler.compiler import AttrsDescriptor

from torch._inductor.runtime import triton_helpers, triton_heuristics
from torch._inductor.runtime.triton_helpers import libdevice, math as tl_math
from torch._inductor.runtime.hints import AutotuneHint, ReductionHint, TileHint, DeviceProperties
triton_helpers.set_driver_to_gpu()

@triton_heuristics.pointwise(
    size_hints={'x': 134217728},
    filename=__file__,
    triton_meta={'signature': {'out_ptr0': '*fp32', 'xnumel': 'i32'}, 'device': DeviceProperties(type='cuda', index=0, multi_processor_count=114, cc=90, major=9, regs_per_multiprocessor=65536, max_threads_per_multi_processor=2048, warp_size=32), 'constants': {}, 'configs': [AttrsDescriptor(divisible_by_16=(0, 1), equal_to_1=())]},
    inductor_meta={'autotune_hints': set(), 'kernel_name': 'triton_poi_fused_zeros_1', 'mutated_arg_names': [], 'optimize_mem': True, 'no_x_dim': False, 'num_load': 0, 'num_reduction': 0, 'backend_hash': 'BC5F52D6E7923B9DC1733AF7005D933F35B86508E4142BC7F067F48E9C59404B', 'are_deterministic_algorithms_enabled': False, 'assert_indirect_indexing': True, 'autotune_local_cache': True, 'autotune_pointwise': True, 'autotune_remote_cache': None, 'force_disable_caches': False, 'dynamic_scale_rblock': True, 'max_autotune': False, 'max_autotune_pointwise': False, 'min_split_scan_rblock': 256, 'spill_threshold': 16, 'store_cubin': False},
    min_elem_per_thread=0
)
@triton.jit
def triton_poi_fused_zeros_1(out_ptr0, xnumel, XBLOCK : tl.constexpr):
    xnumel = 100663296
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = tl.full([XBLOCK], True, tl.int1)
    x0 = xindex
    tmp0 = 0.0
    tl.store(out_ptr0 + (x0), tmp0, None)
```

What might be also related is the following inductor warnings beforehand:

```
[rank0]:E0109 15:24:39.114000 535215 torch/_inductor/memory.py:646] [0/1] Failed to reorder for topological_sort_lpmf: SchedulerBuffer(scheduler=<torch._inductor.scheduler.Scheduler object at 0x750d4019caf0>, node=ComputedBuffer(name='buf0', layout=FixedLayout('cuda:0', torch.float32, size=[1, 1, 384, 256], stride=[98304, 98304, 1, 384]), data=Reduction(
[rank0]:E0109 15:24:39.114000 535215 torch/_inductor/memory.py:646] [0/1]   'cuda',
[rank0]:E0109 15:24:39.114000 535215 torch/_inductor/memory.py:646] [0/1]   torch.float32,
[rank0]:E0109 15:24:39.114000 535215 torch/_inductor/memory.py:646] [0/1]   def inner_fn(index, rindex):
[rank0]:E0109 15:24:39.114000 535215 torch/_inductor/memory.py:646] [0/1]       _, _, i2, i3 = index
[rank0]:E0109 15:24:39.114000 535215 torch/_inductor/memory.py:646] [0/1]       r0_0 = rindex
[rank0]:E0109 15:24:39.114000 535215 torch/_inductor/memory.py:646] [0/1]       tmp0 = ops.load(tangents_1, i2 + 384 * ModularIndexing(r0_0 + 256 * i3, 1, 4096) + 1572864 * ModularIndexing(r0_0 + 256 * i3, 4096, 16))
[rank0]:E0109 15:24:39.114000 535215 torch/_inductor/memory.py:646] [0/1]       tmp1 = ops.load(add_24, i2 + 384 * ModularIndexing(r0_0 + 256 * i3, 1, 4096) + 1572864 * ModularIndexing(r0_0 + 256 * i3, 4096, 16))
[rank0]:E0109 15:24:39.114000 535215 torch/_inductor/memory.py:646] [0/1]       tmp2 = ops.load(rsqrt_6, 4096 * ModularIndexing(r0_0 + 256 * i3, 4096, 16) + ModularIndexing(r0_0 + 256 * i3, 1, 4096))
[rank0]:E0109 15:24:39.114000 535215 torch/_inductor/memory.py:646] [0/1]       tmp3 = tmp1 * tmp2
[rank0]:E0109 15:24:39.114000 535215 torch/_inductor/memory.py:646] [0/1]       tmp4 = tmp0 * tmp3
[rank0]:E0109 15:24:39.114000 535215 torch/_inductor/memory.py:646] [0/1]       return tmp4
[rank0]:E0109 15:24:39.114000 535215 torch/_inductor/memory.py:646] [0/1]   ,
[rank0]:E0109 15:24:39.114000 535215 torch/_inductor/memory.py:646] [0/1]   ranges=[1, 1, 384, 256],
[rank0]:E0109 15:24:39.114000 535215 torch/_inductor/memory.py:646] [0/1]   reduction_ranges=[256],
[rank0]:E0109 15:24:39.114000 535215 torch/_inductor/memory.py:646] [0/1]   reduction_type=sum,
[rank0]:E0109 15:24:39.114000 535215 torch/_inductor/memory.py:646] [0/1]   origin_node=None,
[rank0]:E0109 15:24:39.114000 535215 torch/_inductor/memory.py:646] [0/1]   origins=OrderedSet([sum_21, mul_52, mul_49])
[rank0]:E0109 15:24:39.114000 535215 torch/_inductor/memory.py:646] [0/1] )), defining_op=SchedulerNode(name='op0'), users=[NodeUser(node=SchedulerNode(name='op1'), can_inplace=False, is_weak=False), NodeUser(node=SchedulerNode(name='op27'), can_inplace=False, is_weak=False)], mpi_buffer=MemoryPlanningInfoForBuffer(size_alloc=0, size_free=0, succ_nodes=OrderedSet([])))
```

I know this likely isn't enough to fix the issue and I'd love to provide as repo as soon as possible, I'm just having trouble tracing down what exactly is going wrong here.

### Versions

2.7.0.dev20250109+cu126

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