Function flash_attention_kernel

kernels/flash_attention.py:23–127

Flash attention with online softmax. One program per (batch, head, query-block).

(
    Q_ptr, K_ptr, V_ptr, O_ptr,
    stride_qz, stride_qh, stride_qm, stride_qk,
    stride_kz, stride_kh, stride_kn, stride_kk,
    stride_vz, stride_vh, stride_vn, stride_vk,
    stride_oz, stride_oh, stride_om, stride_ok,
    Z, H, M_size, N_size,
    D: tl.constexpr,
    sm_scale,
    IS_CAUSAL: tl.constexpr,
    BLOCK_M: tl.constexpr,
    BLOCK_N: tl.constexpr,
)

@triton.jit
def flash_attention_kernel(
    Q_ptr, K_ptr, V_ptr, O_ptr,
    stride_qz, stride_qh, stride_qm, stride_qk,
    stride_kz, stride_kh, stride_kn, stride_kk,
    stride_vz, stride_vh, stride_vn, stride_vk,
    stride_oz, stride_oh, stride_om, stride_ok,
    Z, H, M_size, N_size,
    D: tl.constexpr,
    sm_scale,
    IS_CAUSAL: tl.constexpr,
    BLOCK_M: tl.constexpr,
    BLOCK_N: tl.constexpr,
):
    """Flash attention with online softmax. One program per (batch, head, query-block)."""
    pid_z = tl.program_id(2)  # batch
    pid_h = tl.program_id(1)  # head
    pid_m = tl.program_id(0)  # query block

    # Offsets into the batch and head
    qkv_offset_z = pid_z * stride_qz
    qkv_offset_h = pid_h * stride_qh

    k_offset_z = pid_z * stride_kz
    k_offset_h = pid_h * stride_kh

    v_offset_z = pid_z * stride_vz
    v_offset_h = pid_h * stride_vh

    o_offset_z = pid_z * stride_oz
    o_offset_h = pid_h * stride_oh

    # Query block offsets
    offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
    offs_d = tl.arange(0, D)

    # Load Q block [BLOCK_M, D]
    q_ptrs = Q_ptr + qkv_offset_z + qkv_offset_h + offs_m[:, None] * stride_qm + offs_d[None, :] * stride_qk
    q_mask = offs_m[:, None] < M_size
    q = tl.load(q_ptrs, mask=q_mask, other=0.0)

    # Initialize running max and sum for online softmax
    m_i = tl.full((BLOCK_M,), float("-inf"), dtype=tl.float32)
    l_i = tl.zeros((BLOCK_M,), dtype=tl.float32)
    acc = tl.zeros((BLOCK_M, D), dtype=tl.float32)

    # Determine the range of KV blocks to iterate over
    if IS_CAUSAL:
        kv_end = tl.minimum(N_size, (pid_m + 1) * BLOCK_M)
    else:
        kv_end = N_size

    # Iterate over KV blocks
    for start_n in range(0, kv_end, BLOCK_N):
        offs_n = start_n + tl.arange(0, BLOCK_N)

        # Load K block [BLOCK_N, D]
        k_ptrs = K_ptr + k_offset_z + k_offset_h + offs_n[:, None] * stride_kn + offs_d[None, :] * stride_kk
        k_mask = offs_n[:, None] < N_size