Repository: pytorch
Issue: Triton segfault with compiled flexattention kernel
Labels: triaged,upstream triton

### 🐛 Describe the bug

```

# AOT ID: ['1_inference']
from ctypes import c_void_p, c_long
import torch
import math
import random
import os
import tempfile
from math import inf, nan
from torch._inductor.hooks import run_intermediate_hooks
from torch._inductor.utils import maybe_profile
from torch._inductor.codegen.memory_planning import _align as align

from torch import device, empty_strided
from torch._inductor.async_compile import AsyncCompile
from torch._inductor.select_algorithm import extern_kernels
from torch._inductor.codegen.multi_kernel import MultiKernelCall

aten = torch.ops.aten
inductor_ops = torch.ops.inductor
_quantized = torch.ops._quantized
assert_size_stride = torch._C._dynamo.guards.assert_size_stride
empty_strided_cpu = torch._C._dynamo.guards._empty_strided_cpu
empty_strided_cuda = torch._C._dynamo.guards._empty_strided_cuda
reinterpret_tensor = torch._C._dynamo.guards._reinterpret_tensor
alloc_from_pool = torch.ops.inductor._alloc_from_pool
async_compile = AsyncCompile()


# kernel path: /tmp/torchinductor_chilli/bu/cbuaxkh3e3wa6pbzh2hpnx6xnbiuxioy7o42vr3hnsyfzaiyxi53.py
# Source Nodes: [ones], Original ATen: [aten.ones]
# ones => full_default
triton_poi_fused_ones_0 = async_compile.triton('triton_poi_fused_ones_0', '''
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, instance_descriptor, DeviceProperties

@triton_heuristics.pointwise(
    size_hints=[1], 
    filename=__file__,
    triton_meta={'signature': {0: '*i32', 1: 'i32'}, 'device': DeviceProperties(type='cuda', index=0, cc=80, major=8, regs_per_multiprocessor=65536, max_threads_per_multi_processor=2048, multi_processor_count=108), 'constants': {1: 1}, 'configs': [AttrsDescriptor(divisible_by_16=(0,), equal_to_1=(1,))]},
    inductor_meta={'autotune_hints': set(), 'kernel_name': 'triton_poi_fused_ones_0', 'mutated_arg_names': [], 'no_x_dim': False, 'num_load': 0, 'num_reduction': 0, 'backend_hash': 'BA22DB8D050D7BC9D986BC460D4FA53EBD3AB04191AAD21917C2B34E43B3201E', 'are_deterministic_algorithms_enabled': False, 'assert_indirect_indexing': True, 'autotune_local_cache': True, 'autotune_pointwise': True, 'autotune_remote_cache': False, '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_ones_0(out_ptr0, xnumel, XBLOCK : tl.constexpr):
    xnumel = 1
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = tl.full([XBLOCK], True, tl.int1)
    tmp0 = tl.full([1], 1, tl.int32)
    tl.store(out_ptr0 + (tl.full([XBLOCK], 0, tl.int32)), tmp0, None)
''', device_str='cuda')

import triton
import triton.language as tl
from torch._inductor.runtime.triton_heuristics import grid, split_scan_grid, start_graph, end_graph
from torch._C import _cuda_getCurrentRawStream as get_raw_stream


# kernel path: /tmp/torchinductor_chilli/do/cdo2oen7nkfl4arrhos253rz7ugsie7epfyshn3qx2ewlcbvbvb6.py
# Source Nodes: [zeros], Original ATen: [aten.zeros]
# zeros => full_default_1
triton_poi_fused_zeros_1 = async_compile.triton('triton_poi_fused_zeros_1', '''
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, instance_descriptor, DeviceProperties

@triton_heuristics.pointwise(
    size_hints=[1], 
    filename=__file__,
    triton_meta={'signature': {0: '*i32', 1: 'i32'}, 'device': DeviceProperties(type='cuda', index=0, cc=80, major=8, regs_per_multiprocessor=65536, max_threads_per_multi_processor=2048, multi_processor_count=108), 'constants': {1: 1}, 'configs': [AttrsDescriptor(divisible_by_16=(0,), equal_to_1=(1,))]},
    inductor_meta={'autotune_hints': set(), 'kernel_name': 'triton_poi_fused_zeros_1', 'mutated_arg_names': [], 'no_x_dim': False, 'num_load': 0, 'num_reduction': 0, 'backend_hash': 'BA22DB8D050D7BC9D986BC460D4FA53EBD3AB04191AAD21917C2B34E43B3201E', 'are_deterministic_algorithms_enabled': False, 'assert_indirect_indexing': True, 'autotune_local_cache': True, 'autotune_pointwise': True, 'autotune_remote_cache': False, '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 = 1
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = tl.full([XBLOCK], True, tl.int1)
    tmp0 = tl.full([1], 0, tl.int32)
    tl.store(out_ptr0 + (tl.full([XBLOCK], 0, tl.int32)), tmp0, None)
''', device_str='cuda')


# kernel path: /tmp/torchinductor_chilli/lf/clflnjkxha7sm6vqrehrglivx4bhbsi4czpwhsfwmjucxrv7qgw5.py
# Source Nodes: [flex_attention, ones, ones_1, zeros, zeros_1], Original ATen: [aten.ones, aten.zeros]
# flex_attention => flex_attention
# ones => full_default
# ones_1 => full_default_2
# zeros => full_default_1
# zeros_1 => full_default_3
triton_tem_fused_ones_zeros_2 = async_compile.triton('triton_tem_fused_ones_zeros_2', '''
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, instance_descriptor, DeviceProperties

@triton_heuristics.template(
    num_stages=2,
    num_warps=8,
    triton_meta={'signature': {0: '*fp16', 1: '*fp16', 2: '*fp16', 3: '*fp32', 4: '*i32', 5: '*i32', 6: '*fp16'}, 'device': DeviceProperties(type='cuda', index=0, cc=80, major=8, regs_per_multiprocessor=65536, max_threads_per_multi_processor=2048, multi_processor_count=108), 'constants': {}, 'configs': [AttrsDescriptor(divisible_by_16=(0, 1, 2, 3, 4, 5, 6), equal_to_1=())]},
    inductor_meta={'kernel_name': 'triton_tem_fused_ones_zeros_2', 'backend_hash': 'BA22DB8D050D7BC9D986BC460D4FA53EBD3AB04191AAD21917C2B34E43B3201E', 'are_deterministic_algorithms_enabled': False, 'assert_indirect_indexing': True, 'autotune_local_cache': True, 'autotune_pointwise': True, 'autotune_remote_cache': False, '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},
)
@triton.jit
def triton_tem_fused_ones_zeros_2(arg_Q, arg_K, arg_V, arg_LSE, arg_SPARSE_KV_NUM_BLKS, arg_SPARSE_KV_IDX, out_ptr0):
    BLOCK_M : tl.constexpr = 128
    BLOCK_N : tl.constexpr = 128
    BLOCK_DMODEL : tl.constexpr = 128
    SCORE_MOD_IS_LINEAR : tl.constexpr = True
    ROWS_GUARANTEED_SAFE : tl.constexpr = True
    OUTPUT_LOGSUMEXP : tl.constexpr = True
    BLOCKSPARSE_Q : tl.constexpr = 512
    BLOCKSPARSE_KV : tl.constexpr = 512
    Q = arg_Q
    K = arg_K
    V = arg_V
    LSE = arg_LSE
    SPARSE_KV_NUM_BLKS = arg_SPARSE_KV_NUM_BLKS
    SPARSE_KV_IDX = arg_SPARSE_KV_IDX

    # Sub notation for this kernel:
    # Q: Query, K: Key, V: Value
    # M: Number of queries, N: Number of keys/values, D: Model dimension
    # z: Batch size, h: Number of heads, m: Number of queries per head, k: Number of keys per head
    # (Modifiable) Config options:
    # BLOCK_M
    # BLOCK_N
    # SCORE_MOD_IS_LINEAR: Is the score modifier linear? If so, we can lift the
    # change of base out of the loop
    # ROWS_GUARANTEED_SAFE: Is it guaranteed that at least one value in each row
    # is not masked out? If so, we can skip an extra safety check
    # OUTPUT_LOGSUMEXP: We only need to store the logsumexp if we require grad

    tl.static_assert(BLOCKSPARSE_Q >= BLOCK_M and BLOCKSPARSE_Q % BLOCK_M == 0)
    tl.static_assert(BLOCKSPARSE_KV >= BLOCK_N and BLOCKSPARSE_KV % BLOCK_N == 0)

    # Define Q Strides
    stride_qz = 1048576
    stride_qh = 128
    stride_qm = 2048
    stride_qk = 1
    # Define K Strides
    stride_kz = 1048576
    stride_kh = 128
    stride_kn = 2048
    stride_kk = 1
    # Define V Strides
    stride_vz = 1048576
    stride_vh = 128
    stride_vk = 2048
    stride_vn = 1

    Z = 16
    H = 16
    Q_LEN = 512
    KV_LEN = 512

    qk_scale = 1.0
    MATMUL_PRECISION = Q.dtype.element_ty

    start_m = tl.program_id(0)
    off_z = tl.program_id(1) // H
    off_h = tl.program_id(1) % H

    q_offset = off_z * stride_qz + off_h * stride_qh
    k_offset = off_z * stride_kz + off_h * stride_kh
    v_offset = off_z * stride_vz + off_h * stride_vh

    SPARSE_Z = 1
    SPARSE_H = 1

    sparse_idx_z = off_z % SPARSE_Z
    sparse_idx_h = off_h % SPARSE_H

    BLOCKSPARSE_Q_MULTIPLE: tl.constexpr = (BLOCKSPARSE_Q // BLOCK_M)
    BLOCKSPARSE_KV_MULTIPLE: tl.constexpr = (BLOCKSPARSE_KV // BLOCK_N)

    BLOCKSPARSE_Q_LEN: tl.constexpr = Q_LEN // BLOCKSPARSE_Q
    BLOCKSPARSE_KV_LEN: tl.constexpr = KV_LEN // BLOCKSPARSE_KV

    # SPARSE_KV_IDX and SPARSE_KV_NUM_BLKS are always contiguous.
    sparse_hz_offset = sparse_idx_z * SPARSE_H + sparse_idx_h
    sparse_kv_num_blks_offset = sparse_hz_offset * BLOCKSPARSE_Q_LEN + start_m // BLOCKSPARSE_Q_MULTIPLE
    sparse_kv_idx_offset = sparse_hz_offset * BLOCKSPARSE_Q_LEN * BLOCKSPARSE_KV_LEN + (start_m // BLOCKSPARSE_Q_MULTIPLE) * BLOCKSPARSE_KV_LEN  # noqa: B950
    kv_indices = SPARSE_KV_IDX + sparse_kv_idx_offset
    kv_start = tl.load(kv_indices) * BLOCKSPARSE_KV # first kv block we're loading
    sparse_kv_num_blocks = tl.load(SPARSE_KV_NUM_BLKS + sparse_kv_num_blks_offset)

    Q_block_ptr = tl.make_block_ptr(
        base=Q + q_offset,
        shape=(Q_LEN, BLOCK_DMODEL),
        strides=(stride_qm, stride_qk),
        offsets=(start_m * BLOCK_M, 0),
        block_shape=(BLOCK_M, BLOCK_DMODEL),
        order=(1, 0)
    )
    K_block_ptr = tl.make_block_ptr(
        base=K + k_offset,
        shape=(BLOCK_DMODEL, KV_LEN),
        strides=(stride_kk, stride_kn),
        offsets=(0, kv_start),
        block_shape=(BLOCK_DMODEL, BLOCK_N),
        order=(0, 1)
    )
    V_block_ptr = tl.make_block_ptr(
        base=V + v_offset,
        shape=(KV_LEN, BLOCK_DMODEL),
        strides=(stride_vk, stride_vn),
        offsets=(kv_start, 0),
        block_shape=(BLOCK_N, BLOCK_DMODEL),
        order=(1, 0)
    )
    # initialize offsets
    offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
    offs_n = kv_start + tl.arange(0, BLOCK_N)
    # initialize pointer to m and l
    m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf")
    l_i = tl.zeros([BLOCK_M], dtype=tl.float32)
    acc = tl.zeros([BLOCK_M, BLOCK_DMODEL], dtype=tl.float32)

    q = tl.load(Q_block_ptr)
    if SCORE_MOD_IS_LINEAR:
        qk_scale *= 1.44269504
    q = (q * qk_scale).to(MATMUL_PRECISION)

    # loop over k, v and update accumulator
    lo = 0
    hi = sparse_kv_num_blocks * BLOCKSPARSE_KV_MULTIPLE

    for start_n in range(0, hi):
        # -- load k --
        k = tl.load(K_block_ptr)
        # -- compute qk ---
        qk = tl.dot(q, k)
        # ~~~~~~~~~~~~~~~~~~~ Apply score modification  ~~~~~~~~~~~~~~~~~~~
        m = offs_m[:, None]
        n = offs_n[None, :]
        tmp0 = (m) >= (n)
        tmp1 = float("-inf")
        tmp2 = tl.where(tmp0, (qk), tmp1)
        post_mod_scores = tmp2

        # TODO: In the case that score_mod is linear, this can be LICMed
        if not SCORE_MOD_IS_LINEAR:
            post_mod_scores *= 1.44269504
        # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

        # -- compute scaling constant ---
        m_ij = tl.maximum(m_i, tl.max(post_mod_scores, 1))

        alpha = tl.math.exp2(m_i - m_ij)
        p = tl.math.exp2(post_mod_scores - m_ij[:, None])
        if not ROWS_GUARANTEED_SAFE:
            masked_out_rows = (m_ij == float("-inf"))
            alpha = tl.where(masked_out_rows, 0, alpha)
            p = tl.where(masked_out_rows[:, None], 0, p)

        # NB: l_i update is pulled up here since it's a bit faster
        # NB: For headdim=256, it's faster to move it back down to after m_i =
        # m_ij
        l_i = l_i * alpha + tl.sum(p, 1)
        # # -- scale and update acc --
        acc = acc * alpha[:, None]
        v = tl.load(V_block_ptr)
        acc = tl.dot(p.to(MATMUL_PRECISION), v, acc)

        # -- update m_i
        m_i = m_ij

        # update pointers
        indices_idx = start_n // BLOCKSPARSE_KV_MULTIPLE

        cur_block = kv_start
        next_block = tl.load(kv_indices + indices_idx + 1)
        needs_jump = (start_n + 1) % BLOCKSPARSE_KV_MULTIPLE == 0
        jump_to_block = (next_block - cur_block ) * BLOCKSPARSE_KV - (BLOCKSPARSE_KV_MULTIPLE - 1) * BLOCK_N
        kv_start = needs_jump * next_block + (1 - needs_jump) * cur_block

        offset = jump_to_block * needs_jump + (1 - needs_jump) * BLOCK_N
        # offset = BLOCK_N

        V_block_ptr = tl.advance(V_block_ptr, (offset, 0))
        K_block_ptr = tl.advance(K_block_ptr, (0, offset))

        offs_n = offs_n + offset

    # Store output and logsumexp
    acc = acc / l_i[:, None]
    idx_z = tl.program_id(1) // H
    idx_h = tl.program_id(1) % H
    idx_m = offs_m[:, None]
    idx_d = tl.arange(0, BLOCK_DMODEL)[None, :]

    mask = idx_m < Q_LEN
    # TODO generalize and add proper mask support
    xindex = idx_d + (128*idx_m) + (65536*idx_h) + (1048576*idx_z)
    tl.store(out_ptr0 + (tl.broadcast_to(idx_d + (128*idx_h) + (2048*idx_m) + (1048576*idx_z), acc.shape)), acc, mask)

    # TODO dont want to write this if we dont require grad
    if OUTPUT_LOGSUMEXP:
        off_hz = tl.program_id(1)
        l_ptrs = LSE + off_hz * Q_LEN + offs_m
        lse = m_i + tl.math.log2(l_i)
        tl.store(l_ptrs, lse)
''', device_str='cuda')
import torch._inductor.kernel.flex_attention
meta0 = {'BLOCK_M': 128, 'BLOCK_N': 128, 'BLOCK_DMODEL': 128, 'SCORE_MOD_IS_LINEAR': True, 'ROWS_GUARANTEED_SAFE': True, 'OUTPUT_LOGSUMEXP': True, 'BLOCKSPARSE_Q': 512, 'BLOCKSPARSE_KV': 512}


async_compile.wait(globals())
del async_compile

def call(args):
    arg0_1, arg1_1, arg2_1 = args
    args.clear()
    assert_size_stride(arg0_1, (16, 16, 512, 128), (1048576, 128, 2048, 1))
    assert_size_stride(arg1_1, (16, 16, 512, 128), (1048576, 128, 2048, 1))
    assert_size_stride(arg2_1, (16, 16, 512, 128), (1048576, 128, 2048, 1))
    with torch.cuda._DeviceGuard(0):
        torch.cuda.set_device(0)
        buf0 = empty_strided_cuda((1, 1, 1), (1, 1, 1), torch.int32)
        # Source Nodes: [ones], Original ATen: [aten.ones]
        stream0 = get_raw_stream(0)
        triton_poi_fused_ones_0.run(buf0, 1, grid=grid(1), stream=stream0)
        buf1 = empty_strided_cuda((1, 1, 1, 1), (1, 1, 1, 1), torch.int32)
        # Source Nodes: [zeros], Original ATen: [aten.zeros]
        triton_poi_fused_zeros_1.run(buf1, 1, grid=grid(1), stream=stream0)
        buf4 = empty_strided_cuda((16, 16, 512), (8192, 512, 1), torch.float32)
        buf5 = empty_strided_cuda((16, 16, 512, 128), (1048576, 128, 2048, 1), torch.float16)
        # Source Nodes: [flex_attention, ones, ones_1, zeros, zeros_1], Original ATen: [aten.ones, aten.zeros]
        print("Running!")
        triton_tem_fused_ones_zeros_2.run(arg0_1, arg1_1, arg2_1, buf4, buf0, buf1, buf5, grid=torch._inductor.kernel.flex_attention.flex_attention_grid(16, 16, 512, 128, meta0), stream=stream0)
        del arg0_1
        del arg1_1
        del arg2_1
        del buf0
        del buf1
    return (buf5, )


def benchmark_compiled_module(times=10, repeat=10):
    from torch._dynamo.testing import rand_strided
    from torch._inductor.utils import print_performance
    arg0_1 = rand_strided((16, 16, 512, 128), (1048576, 128, 2048, 1), device='cuda:0', dtype=torch.float16)
    arg1_1 = rand_strided((16, 16, 512, 128), (1048576, 128, 2048, 1), device='cuda:0', dtype=torch.float16)
    arg2_1 = rand_strided((16, 16, 512, 128), (1048576, 128, 2048, 1), device='cuda:0', dtype=torch.float16)
    fn = lambda: call([arg0_1, arg1_1, arg2_1])
    return print_performance(fn, times=times, repeat=repeat)


if __name__ == "__main__":
    from torch._inductor.wrapper_benchmark import compiled_module_main
    compiled_module_main('None', benchmark_compiled_module)

```

Stacktrace from `cuda-gdb`
<img width="1233" alt="image" src="https://github.com/pytorch/pytorch/assets/6355099/89612b38-d1a0-410e-9f55-5942c3f160da">

I get a segfault when running this.

cc: @htyu @bertmaher 

### Versions

N/A

cc @bertmaher