microsoft--unilm/ReSA/llm/kernel/flash_attention_with_kv_cache.py

import math
import torch

import triton
import triton.language as tl

def is_hip():
    return triton.runtime.driver.active.get_current_target().backend == "hip"


def num_splits_heuristic(total_mblocks, num_SMs, num_n_blocks, num_m_blocks, size_one_kv_head,
                         is_causal_or_local, max_splits):
    """
    Determines the optimal number of splits for maximizing GPU occupancy while balancing memory efficiency.

    Parameters:
    - total_mblocks (int): Total number of m_blocks.
    - num_SMs (int): Number of Streaming Multiprocessors (SMs) in the GPU.
    - num_n_blocks (int): Number of n_blocks.
    - num_m_blocks (int): Number of m_blocks.
    - size_one_kv_head (int): Size of one KV head in bytes.
    - is_causal_or_local (bool): Indicates whether the operation is causal or local.
    - max_splits (int): Maximum number of allowed splits.

    Returns:
    - int: The optimal number of splits.
    """
    # If we have enough m_blocks to almost fill the SMs, prefer 1 split unless memory constraints apply.
    if total_mblocks >= 0.8 * num_SMs:
        size_l2 = 50 * 1024 * 1024  # L2 cache size assumption (50MB)
        # Only split if each KV head is too large for L2 and there are enough m_blocks
        if size_one_kv_head > size_l2 and num_m_blocks >= num_SMs * 2 and not is_causal_or_local:
            return min((size_one_kv_head + size_l2 - 1) // size_l2, max_splits)
        else:
            return 1

    # If num_n_blocks is too small, we don't split
    if num_n_blocks <= 4:
        return 1

    # Limit max_splits to a reasonable range
    max_splits = min(max_splits, num_SMs, num_n_blocks)

    max_efficiency = 0.0
    efficiency = []

    # Compute efficiency for different splits
    for num_splits in range(1, max_splits + 1):
        n_waves = (total_mblocks * num_splits) / num_SMs
        eff = n_waves / math.ceil(n_waves)
        # Track max efficiency
        if eff > max_efficiency:
            max_efficiency = eff

        efficiency.append(eff)

    # Find the smallest number of splits that achieves at least 85% of max efficiency
    for num_splits in range(1, max_splits + 1):
        if efficiency[num_splits - 1] >= 0.95 * max_efficiency:
            return num_splits

    return 1

@triton.autotune(
    configs=[
        triton.Config({}, num_warps=num_warps)
        for num_warps in [1, 2, 4, 8, 16]
    ],
    key=['gqa_group_size', 'BLOCK_H', 'BLOCK_N', 'BLOCK_D', 'BLOCK_V'],
)
@triton.jit
def _fwd_kernel_with_kv_cache(
    Q, K, V, Out, L,
    sm_scale,
    cache_seqlens,
    stride_qz, stride_qt, stride_qh, stride_qd,
    stride_kz, stride_kt, stride_kh, stride_kd,
    stride_vz, stride_vt, stride_vh, stride_vd,
    stride_oz, stride_ot, stride_oh, stride_os, stride_od,
    stride_lz, stride_lt, stride_lh, stride_ls,
    num_splits: tl.constexpr,
    seqlen_q: tl.constexpr,
    num_m_blocks: tl.constexpr,
    gqa_group_size: tl.constexpr,
    BLOCK_H: tl.constexpr, 
    BLOCK_M: tl.constexpr,
    BLOCK_N: tl.constexpr,
    BLOCK_D: tl.constexpr,
    BLOCK_V: tl.constexpr,
):
    off_sm = tl.program_id(0).to(tl.int64)
    off_split, off_m = off_sm // num_m_blocks, off_sm % num_m_blocks
    off_h_for_kv = tl.program_id(1).to(tl.int64)
    off_z = tl.program_id(2).to(tl.int64)

    off_h_q = off_h_for_kv * gqa_group_size
    offs_h = tl.arange(0, BLOCK_H)
    offs_n = tl.arange(0, BLOCK_N)
    offs_d = tl.arange(0, BLOCK_D)
    offs_v = tl.arange(0, BLOCK_V)
    mask_h = offs_h < gqa_group_size
    seqlen_k = tl.load(cache_seqlens + off_z)

    Q += off_z * stride_qz + off_h_q * stride_qh
    K += off_z * stride_kz + off_h_for_kv * stride_kh
    V += off_z * stride_vz + off_h_for_kv * stride_vh
    L += off_z * stride_lz + off_h_q * stride_lh + off_split * stride_ls
    Out += off_z * stride_oz + off_h_q * stride_oh + off_split * stride_os
    
    num_kv_blocks = tl.cdiv(seqlen_k, BLOCK_N)
    blocks_per_split = num_kv_blocks // num_splits
    remaining_blocks = num_kv_blocks % num_splits
    loop_range = blocks_per_split + (1 if off_split < remaining_blocks else 0)
    start = blocks_per_split * off_split + min(off_split, remaining_blocks)

    offs_m = tl.arange(0, BLOCK_M) + off_m * BLOCK_M
    mask_q = offs_m < seqlen_q
    mask_qh = mask_q[:, None] & mask_h[None, :]
    q_idx = offs_m[:, None] + tl.zeros([BLOCK_H], dtype=tl.int32) + seqlen_k - seqlen_q
    q_idx = tl.reshape(q_idx, [BLOCK_M * BLOCK_H])

    q = tl.load(Q + offs_m[:, None, None] * stride_qt + offs_h[None, :, None] * stride_qh + offs_d[None, None, :] * stride_qd, mask=mask_qh[:, :, None]) ## padding to min 16
    q = tl.reshape(q, (BLOCK_M * BLOCK_H, BLOCK_D))

    m_i = tl.full([BLOCK_M * BLOCK_H], float("-inf"), dtype=tl.float32)
    l_i = tl.full([BLOCK_M * BLOCK_H], 1.0, dtype=tl.float32)
    acc = tl.zeros([BLOCK_M * BLOCK_H, BLOCK_V], dtype=tl.float32)

    k_ptrs = K + offs_n[None, :] * stride_kt + offs_d[:, None] * stride_kd
    v_ptrs = V + offs_n[:, None] * stride_vt + offs_v[None, :] * stride_vd

    for block_idx in range(start, start + loop_range):
        start_n = block_idx * BLOCK_N
        k = tl.load(k_ptrs + start_n * stride_kt, mask=offs_n[None, :] + start_n < seqlen_k, cache_modifier=".ca")
        qk = tl.dot(q, k)

        causal_mask = q_idx[:, None] >= start_n + offs_n[None, :]
        qk = tl.where(causal_mask, qk, -1.0e6)
        qk *= sm_scale

        m_ij = tl.maximum(m_i, tl.max(qk, 1))
        qk -= m_ij[:, None]
        p = tl.exp(qk)
        l_ij = tl.sum(p, 1)
        alpha = tl.exp(m_i - m_ij)
        l_i = l_i * alpha + l_ij
        acc = acc * alpha[:, None]
        
        v = tl.load(v_ptrs + start_n * stride_vt, mask=offs_n[:, None] + start_n < seqlen_k, cache_modifier=".ca")
        p = p.to(v.type.element_ty)

        acc += tl.dot(p, v)
        m_i = m_ij

    l_recip = 1 / l_i[:, None]
    acc = acc * l_recip
    m_i += tl.math.log(l_i)

    l_ptrs = L + offs_m[:, None] * stride_lt + offs_h * stride_lh
    m_i = tl.reshape(m_i, (BLOCK_M, BLOCK_H))
    tl.store(l_ptrs, m_i, mask=mask_qh)
    O_ptrs = Out + offs_m[:, None, None] * stride_ot + offs_h[None, :, None] * stride_oh + offs_v[None, None, :] * stride_od
    acc = tl.reshape(acc, (BLOCK_M, BLOCK_H, BLOCK_V))
    tl.store(O_ptrs, acc, mask=mask_qh[:, :, None])

@triton.autotune(
    configs=[
        triton.Config({}, num_warps=num_warps)
        for num_warps in [1, 2, 4, 8, 16]
    ],
    key=['BLOCK_V'],
)
@triton.jit
def combine(
    out_partial, out, L,
    stride_op_z, stride_op_t, stride_op_h, stride_op_s, stride_op_d,
    stride_o_z, stride_o_t, stride_o_h, stride_o_d,
    stride_l_z, stride_l_t, stride_l_h, stride_l_s,
    num_splits: tl.constexpr,
    num_splits_pow2: tl.constexpr,
    BLOCK_V: tl.constexpr,
):
    off_h = tl.program_id(0).to(tl.int64)
    off_t = tl.program_id(1).to(tl.int64)
    off_z = tl.program_id(2).to(tl.int64)

    split = tl.arange(0, num_splits_pow2)
    split_mask = split < num_splits

    L += off_z * stride_l_z + off_t * stride_l_t + off_h * stride_l_h
    out_partial += off_z * stride_op_z + off_t * stride_op_t + off_h * stride_op_h
    out += off_z * stride_o_z + off_t * stride_o_t + off_h * stride_o_h

    lse_local = tl.load(L + split * stride_l_s, mask=split_mask, other=float("-inf"))
    lse_max_local = tl.max(lse_local, axis=0)

    lse_logsum_local = tl.sum(tl.exp(lse_local - lse_max_local), axis=0)
    lse_logsum_local = tl.log(lse_logsum_local) + lse_max_local
    
    po_local = tl.load(out_partial + split[:, None] * stride_op_s + tl.arange(0, BLOCK_V) * stride_op_d, mask=split_mask[:, None])
    scale_local = tl.exp(lse_local - lse_logsum_local)
    accum_local = tl.sum(po_local * scale_local[:, None], axis=0)
    tl.store(out + tl.arange(0, BLOCK_V) * stride_o_d, accum_local)

def flash_attention_with_kv_cache(
    q, k, v,
    cache_seqlens,
    sm_scale=None,
):
    # split q to blocks
    batch, seqlen_q, n_heads, key_dim = q.shape
    _, _, n_kv_heads, head_dim = v.shape
    gqa_group_size = n_heads // n_kv_heads
    block_h = triton.next_power_of_2(gqa_group_size)
    block_m = max(256 // block_h, 1)
    block_n = 32

    # assert seqlen_q <= 32, "it seems the performance is not good when seqlen_q > 32"

    assert k.size(0) == v.size(0)
    assert q.size(3) == k.size(3)
    assert k.size(1) == v.size(1)
    assert key_dim in {64, 128, 256}
    assert head_dim in {64, 128, 256}

    props = torch.cuda.get_device_properties(torch.device("cuda:0"))
    num_sm = props.multi_processor_count
    num_m_blocks = triton.cdiv(seqlen_q, block_m)
    num_n_blocks = triton.cdiv(cache_seqlens.max(), block_n)
    size_one_kv_head = cache_seqlens.max() * block_n * (key_dim + head_dim) * 2
    total_mblocks = batch * n_kv_heads
    num_splits = num_splits_heuristic(
        total_mblocks, num_sm, num_n_blocks, num_m_blocks,
        size_one_kv_head, is_causal_or_local=True, max_splits=16
    )

    out_partial = torch.empty((batch, seqlen_q, n_heads, num_splits, head_dim), device=q.device, dtype=torch.float32)
    out = torch.empty((batch, seqlen_q, n_heads, head_dim), device=q.device, dtype=q.dtype)
    L = torch.empty((batch, seqlen_q, n_heads, num_splits), device=q.device, dtype=torch.float32)

    if is_hip():
        extra_kern_args = {"waves_per_eu": 1}
    else:
        extra_kern_args = {}

    with torch.cuda.device(q.device.index): 
        grid = lambda META: (num_splits * num_m_blocks, n_kv_heads, batch)
        _fwd_kernel_with_kv_cache[grid](
            q, k, v, out_partial, L,
            sm_scale if sm_scale is not None else key_dim ** -0.5,
            cache_seqlens.contiguous(),
            *q.stride(),
            *k.stride(),
            *v.stride(),
            *out_partial.stride(),
            *L.stride(),
            num_splits=num_splits,
            seqlen_q=seqlen_q,
            num_m_blocks=num_m_blocks,
            gqa_group_size=gqa_group_size,
            BLOCK_H = block_h,
            BLOCK_M = block_m,
            BLOCK_N = block_n,
            BLOCK_D = key_dim,
            BLOCK_V = head_dim,
            **extra_kern_args
        )
        grid = lambda META: (n_heads, seqlen_q, batch)
        combine[grid](
            out_partial, out, L,
            *out_partial.stride(),
            *out.stride(),
            *L.stride(),
            num_splits=num_splits,
            num_splits_pow2=triton.next_power_of_2(num_splits),
            BLOCK_V=head_dim,
            **extra_kern_args
        )
    return out

def ref_program_fa(query, key, value, cache_seqlens):
    # latency reference
    # from flash_attn_interface import flash_attn_with_kvcache, flash_attn_func # fa3
    from flash_attn import flash_attn_with_kvcache, flash_attn_func #fa2
    output = flash_attn_with_kvcache(query, key, value, cache_seqlens=cache_seqlens, causal=True)
    return output



def debug(name,expect, actual, atol=1e-3, rtol=1e-3):
    all_close = torch.allclose(expect, actual, atol=atol, rtol=rtol)
    print(name + "  all_close={}".format(all_close))
    if not all_close:
        # print(expect[3, 28])
        # print(actual[3, 28])
        diff = (expect - actual).abs()
        print("all_close={}, max={}, min={}, mean={}".format(all_close, diff.max().item(), diff.min().item(), diff.mean().item()))
        max_indices = torch.nonzero(diff == diff.max().item())
        first_index = tuple(max_indices[0].tolist())
        print(f"Index: {first_index}, expect: {expect[first_index]}, actual: {actual[first_index]}") 

if __name__ == "__main__":
    import argparse
    import time
    parser = argparse.ArgumentParser()
    parser.add_argument('--batch', type=int, default=8, help='batch size')
    parser.add_argument('--seqlen_q', type=int, default=128, help='sequence length')
    parser.add_argument('--heads', type=int, default=28, help='heads')
    parser.add_argument('--heads_kv', type=int, default=4, help='heads_kv')
    parser.add_argument('--max_cache_seqlen', type=int, default=65536, help='kvcache sequence length')
    parser.add_argument('--dim', type=int, default=128, help='dim')
    parser.add_argument('--dim_v', type=int, default=128, help='dim_v')
    parser.add_argument('--load_from_file', type=str, default=None, help='load from file')
    args = parser.parse_args()

    batch, seqlen_q, heads, heads_kv, max_cache_seqlen, dim, dim_v = args.batch, args.seqlen_q, args.heads, args.heads_kv, args.max_cache_seqlen, args.dim, args.dim_v

    dtype = torch.bfloat16

    Q = torch.randn((batch, seqlen_q, heads, dim), dtype=dtype, device='cuda')
    K = torch.randn((batch, max_cache_seqlen, heads_kv, dim), dtype=dtype, device='cuda')
    V = torch.randn((batch, max_cache_seqlen, heads_kv, dim_v), dtype=dtype, device='cuda')

    cache_seqlens = torch.randint(max_cache_seqlen - 32, max_cache_seqlen, (batch,), device='cuda', dtype=torch.int32)
    print("cache_seqlens: ", cache_seqlens)
    # parity reference
    ref = ref_program_fa(Q, K, V, cache_seqlens)
    # ref = ref_program_triton(Q, K, V, block_indices, cache_seqlens, max_cache_seqlen, max_num_blocks, block_size)
    # out = kernel(Q, K, V, block_indices, cache_seqlens, actual_num_blocks, glse, Output_partial)
    # out = sparse_gqa_decode_varlen_indice(Q, K, V, block_indices, cache_seqlens, max_cache_seqlen, block_size)
    out = flash_attention_with_kv_cache(Q, K, V, cache_seqlens)
    debug("output", ref, out, atol=1e-3, rtol=1e-3)

    ## latency reference
    for i in range(10):
        ref = ref_program_fa(Q, K, V, cache_seqlens)
    torch.cuda.synchronize()
    start = time.time()
    for i in range(100):
        ref = ref_program_fa(Q, K, V, cache_seqlens)
    torch.cuda.synchronize()
    print("dense time: ", (time.time() - start) / 100*1000)

    for i in range(10):
        out = flash_attention_with_kv_cache(Q, K, V, cache_seqlens)
    torch.cuda.synchronize()
    start = time.time()
    for i in range(100):
        out = flash_attention_with_kv_cache(Q, K, V, cache_seqlens)
    torch.cuda.synchronize()
    print("sparse time: ", (time.time() - start) / 100*1000)