ColossalAI/colossalai/kernel/triton/context_attn_unpad.py

# Applying the FlashAttention V2 as described in:
# "FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning"
# by Tri Dao, 2023
# https://github.com/Dao-AILab/flash-attention
#
# Inspired and modified from Triton Tutorial - Fused Attention
# https://triton-lang.org/main/getting-started/tutorials/06-fused-attention.html

import torch
import triton
import triton.language as tl


# Triton 2.1.0
@triton.jit
def _fwd_context_paged_attention_kernel(
    Q,
    K,
    V,
    O,
    KCache,
    VCache,
    BLOCK_TABLES,  # [num_seqs, max_blocks_per_sequence]
    batch_size,
    stride_qt,
    stride_qh,
    stride_qd,
    stride_kt,
    stride_kh,
    stride_kd,
    stride_vt,
    stride_vh,
    stride_vd,
    stride_ot,
    stride_oh,
    stride_od,
    stride_cacheb,
    stride_cacheh,
    stride_cachebs,
    stride_cached,
    stride_bts,
    stride_btb,
    context_lengths,
    sm_scale,
    KV_GROUPS: tl.constexpr,
    BLOCK_SIZE: tl.constexpr,
    HEAD_DIM: tl.constexpr,
    BLOCK_M: tl.constexpr,
    BLOCK_N: tl.constexpr,
):
    cur_seq_idx = tl.program_id(0)
    if cur_seq_idx >= batch_size:
        return
    cur_head_idx = tl.program_id(1)
    block_start_m = tl.program_id(2)  # Br, max_input_len // Block_M
    cur_kv_head_idx = cur_head_idx // KV_GROUPS

    # NOTE It requires BLOCK_M, BLOCK_N, and BLOCK_SIZE to be the same
    tl.static_assert(BLOCK_M == BLOCK_N)
    tl.static_assert(BLOCK_N == BLOCK_SIZE)

    # get the current sequence length from provided context lengths tensor
    cur_seq_len = tl.load(context_lengths + cur_seq_idx)
    # NOTE when talking to fused QKV and a nopadding context attention,
    # we assume that the input Q/K/V is contiguous, and thus here `prev_seq_len_sum`
    # could be considered as the start index of the current sequence.
    # FIXME might want to explore better way to get the summation of prev seq lengths.
    # `tl.sum(tensor[:end])` is invalid as tensor slice is not supported in triton.
    prev_seq_len_sum = 0
    for i in range(0, cur_seq_idx):
        prev_seq_len_sum += tl.load(context_lengths + i)

    offset_q = prev_seq_len_sum * stride_qt + cur_head_idx * stride_qh
    offset_kv = prev_seq_len_sum * stride_kt + cur_kv_head_idx * stride_kh
    Q_block_ptr = tl.make_block_ptr(
        base=Q + offset_q,
        shape=(cur_seq_len, HEAD_DIM),
        strides=(stride_qt, stride_qd),
        offsets=(block_start_m * BLOCK_M, 0),
        block_shape=(BLOCK_M, HEAD_DIM),
        order=(1, 0),
    )
    K_block_ptr = tl.make_block_ptr(
        base=K + offset_kv,
        shape=(HEAD_DIM, cur_seq_len),
        strides=(stride_kd, stride_kt),
        offsets=(0, 0),
        block_shape=(HEAD_DIM, BLOCK_N),
        order=(0, 1),
    )
    V_block_ptr = tl.make_block_ptr(
        base=V + offset_kv,
        shape=(cur_seq_len, HEAD_DIM),
        strides=(stride_vt, stride_vd),
        offsets=(0, 0),
        block_shape=(BLOCK_N, HEAD_DIM),
        order=(1, 0),
    )
    O_block_ptr = tl.make_block_ptr(
        base=O + offset_q,
        shape=(cur_seq_len, HEAD_DIM),
        strides=(stride_ot, stride_od),
        offsets=(block_start_m * BLOCK_M, 0),
        block_shape=(BLOCK_M, HEAD_DIM),
        order=(1, 0),
    )

    # block table for the current sequence
    block_table_ptr = BLOCK_TABLES + cur_seq_idx * stride_bts
    # block indexes on block table (i.e. 0, 1, 2, ..., max_blocks_per_seq)
    # Consider `block_start_m` as the logical block idx in the current block table,
    # as we have BLOCK_M the same size as the block size.
    cur_block_table_idx = block_start_m
    cur_block_id = tl.load(block_table_ptr + cur_block_table_idx * stride_btb)
    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh

    offsets_m = block_start_m * BLOCK_M + tl.arange(0, BLOCK_M)
    offsets_n = tl.arange(0, BLOCK_N)
    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, HEAD_DIM], dtype=tl.float32)

    if block_start_m * BLOCK_M >= cur_seq_len:
        return

    Q_i = tl.load(Q_block_ptr, boundary_check=(1, 0))

    for block_start_n in range(0, (block_start_m + 1) * BLOCK_M, BLOCK_N):
        block_start_n = tl.multiple_of(block_start_n, BLOCK_N)

        k = tl.load(K_block_ptr, boundary_check=(0, 1))
        S_ij = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
        S_ij += tl.dot(Q_i, k)
        S_ij *= sm_scale
        S_ij += tl.where(offsets_m[:, None] >= (block_start_n + offsets_n[None, :]), 0, float("-inf"))

        m_ij = tl.max(S_ij, 1)  # rowmax(Sij)
        m_ij = tl.maximum(m_i, m_ij)  # m_ij
        S_ij -= m_ij[:, None]
        p_ij_hat = tl.exp(S_ij)
        scale = tl.exp(m_i - m_ij)
        l_ij = scale * l_i + tl.sum(p_ij_hat, 1)
        acc = acc * scale[:, None]

        v = tl.load(V_block_ptr, boundary_check=(1, 0))
        p_ij_hat = p_ij_hat.to(v.type.element_ty)

        acc += tl.dot(p_ij_hat, v)
        l_i = l_ij
        m_i = m_ij
        K_block_ptr = tl.advance(K_block_ptr, (0, BLOCK_N))
        V_block_ptr = tl.advance(V_block_ptr, (BLOCK_N, 0))

    acc = acc / l_i[:, None]
    tl.store(O_block_ptr, acc.to(O.type.element_ty), boundary_check=(1, 0))

    if cur_head_idx % KV_GROUPS == 0:
        # Copy k to corresponding cache block
        offsets_dmodel = tl.arange(0, HEAD_DIM)
        offsets_kt = block_start_m * BLOCK_M + tl.arange(0, BLOCK_M)
        offsets_k = K + offset_kv + offsets_dmodel[None, :] * stride_kd + offsets_kt[:, None] * stride_kt
        k = tl.load(offsets_k, mask=offsets_kt[:, None] < cur_seq_len, other=0.0)
        offsets_kcachebs = tl.arange(0, BLOCK_SIZE)
        offsets_kcache = (
            KCache
            + offset_kvcache
            + offsets_dmodel[None, :] * stride_cached
            + offsets_kcachebs[:, None] * stride_cachebs
        )
        tl.store(offsets_kcache, k, mask=offsets_kcachebs[:, None] < cur_seq_len - block_start_m * BLOCK_SIZE)
        # Copy v to corresponding cache block
        offsets_vd = offsets_dmodel
        offsets_vt = block_start_m * BLOCK_N + tl.arange(0, BLOCK_N)
        offsets_v = V + offset_kv + offsets_vt[None, :] * stride_vt + offsets_vd[:, None] * stride_vd
        v = tl.load(offsets_v, mask=offsets_vt[None, :] < cur_seq_len, other=0.0)
        offsets_vcachebs = offsets_kcachebs  # same block size range, just to notify here
        offsets_vcache = (
            VCache
            + offset_kvcache
            + offsets_vcachebs[None, :] * stride_cachebs
            + offsets_dmodel[:, None] * stride_cached
        )
        tl.store(offsets_vcache, v, mask=offsets_vcachebs[None, :] < cur_seq_len - block_start_m * BLOCK_SIZE)

    return


# Triton 2.1.0
# TODO(yuanheng-zhao): This is a temporary dispatch to use the new layout for kcache
# merge `_fwd_context_paged_attention_kernel_v2` with `_fwd_context_paged_attention_kernel` later
# as the kcache layout has been supported in the whole triton flow.
@triton.jit
def _fwd_context_paged_attention_kernel_v2(
    Q,
    K,
    V,
    O,
    KCache,  # [num_blocks, num_kv_heads, head_dim // x, block_size, x]
    VCache,  # [num_blocks, num_kv_heads, block_size, head_dim]
    BLOCK_TABLES,  # [num_seqs, max_blocks_per_sequence]
    batch_size,
    stride_qt,
    stride_qh,
    stride_qd,
    stride_kt,
    stride_kh,
    stride_kd,
    stride_vt,
    stride_vh,
    stride_vd,
    stride_ot,
    stride_oh,
    stride_od,
    stride_cacheb,  # v cache stride(0) - num_blocks
    stride_cacheh,  # v cache stride(1) - num_kv_heads
    stride_cachebs,  # v cache stride(2) - block_size
    stride_cached,  # v cache stride(3) - head_dim
    stride_bts,
    stride_btb,
    context_lengths,
    sm_scale,
    KV_GROUPS: tl.constexpr,
    BLOCK_SIZE: tl.constexpr,
    HEAD_DIM: tl.constexpr,
    KCACHE_X: tl.constexpr,  # k stride on the second last dimension
    BLOCK_M: tl.constexpr,
    BLOCK_N: tl.constexpr,
):
    cur_seq_idx = tl.program_id(0)
    if cur_seq_idx >= batch_size:
        return
    cur_head_idx = tl.program_id(1)
    block_start_m = tl.program_id(2)  # Br, max_input_len // Block_M
    cur_kv_head_idx = cur_head_idx // KV_GROUPS

    # NOTE It requires BLOCK_M, BLOCK_N, and BLOCK_SIZE to be the same
    tl.static_assert(BLOCK_M == BLOCK_N)
    tl.static_assert(BLOCK_N == BLOCK_SIZE)

    # get the current sequence length from provided context lengths tensor
    cur_seq_len = tl.load(context_lengths + cur_seq_idx)
    # NOTE when talking to fused QKV and a nopadding context attention,
    # we assume that the input Q/K/V is contiguous, and thus here `prev_seq_len_sum`
    # could be considered as the start index of the current sequence.
    # FIXME might want to explore better way to get the summation of prev seq lengths.
    # `tl.sum(tensor[:end])` is invalid as tensor slice is not supported in triton.
    prev_seq_len_sum = 0
    for i in range(0, cur_seq_idx):
        prev_seq_len_sum += tl.load(context_lengths + i)

    offset_q = prev_seq_len_sum * stride_qt + cur_head_idx * stride_qh
    offset_kv = prev_seq_len_sum * stride_kt + cur_kv_head_idx * stride_kh
    Q_block_ptr = tl.make_block_ptr(
        base=Q + offset_q,
        shape=(cur_seq_len, HEAD_DIM),
        strides=(stride_qt, stride_qd),
        offsets=(block_start_m * BLOCK_M, 0),
        block_shape=(BLOCK_M, HEAD_DIM),
        order=(1, 0),
    )
    K_block_ptr = tl.make_block_ptr(
        base=K + offset_kv,
        shape=(HEAD_DIM, cur_seq_len),
        strides=(stride_kd, stride_kt),
        offsets=(0, 0),
        block_shape=(HEAD_DIM, BLOCK_N),
        order=(0, 1),
    )
    V_block_ptr = tl.make_block_ptr(
        base=V + offset_kv,
        shape=(cur_seq_len, HEAD_DIM),
        strides=(stride_vt, stride_vd),
        offsets=(0, 0),
        block_shape=(BLOCK_N, HEAD_DIM),
        order=(1, 0),
    )
    O_block_ptr = tl.make_block_ptr(
        base=O + offset_q,
        shape=(cur_seq_len, HEAD_DIM),
        strides=(stride_ot, stride_od),
        offsets=(block_start_m * BLOCK_M, 0),
        block_shape=(BLOCK_M, HEAD_DIM),
        order=(1, 0),
    )

    # block table for the current sequence
    block_table_ptr = BLOCK_TABLES + cur_seq_idx * stride_bts
    # block indexes on block table (i.e. 0, 1, 2, ..., max_blocks_per_seq)
    # Consider `block_start_m` as the logical block idx in the current block table,
    # as we have BLOCK_M the same size as the block size.
    cur_block_table_idx = block_start_m
    cur_block_id = tl.load(block_table_ptr + cur_block_table_idx * stride_btb)
    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh

    offsets_m = block_start_m * BLOCK_M + tl.arange(0, BLOCK_M)
    offsets_n = tl.arange(0, BLOCK_N)
    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, HEAD_DIM], dtype=tl.float32)

    if block_start_m * BLOCK_M >= cur_seq_len:
        return

    Q_i = tl.load(Q_block_ptr, boundary_check=(1, 0))

    for block_start_n in range(0, (block_start_m + 1) * BLOCK_M, BLOCK_N):
        block_start_n = tl.multiple_of(block_start_n, BLOCK_N)

        k = tl.load(K_block_ptr, boundary_check=(0, 1))
        S_ij = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
        S_ij += tl.dot(Q_i, k)
        S_ij *= sm_scale
        S_ij += tl.where(offsets_m[:, None] >= (block_start_n + offsets_n[None, :]), 0, float("-inf"))

        m_ij = tl.max(S_ij, 1)  # rowmax(Sij)
        m_ij = tl.maximum(m_i, m_ij)  # m_ij
        S_ij -= m_ij[:, None]
        p_ij_hat = tl.exp(S_ij)
        scale = tl.exp(m_i - m_ij)
        l_ij = scale * l_i + tl.sum(p_ij_hat, 1)
        acc = acc * scale[:, None]

        v = tl.load(V_block_ptr, boundary_check=(1, 0))
        p_ij_hat = p_ij_hat.to(v.type.element_ty)

        acc += tl.dot(p_ij_hat, v)
        l_i = l_ij
        m_i = m_ij
        K_block_ptr = tl.advance(K_block_ptr, (0, BLOCK_N))
        V_block_ptr = tl.advance(V_block_ptr, (BLOCK_N, 0))

    acc = acc / l_i[:, None]
    tl.store(O_block_ptr, acc.to(O.type.element_ty), boundary_check=(1, 0))

    if cur_head_idx % KV_GROUPS == 0:
        # Copy k to corresponding cache block
        block_range = tl.arange(0, BLOCK_SIZE)
        X_range = tl.arange(0, KCACHE_X)
        # unroll the loop aggressively
        for split_x in tl.static_range(HEAD_DIM // KCACHE_X):
            offsets_dmodel_x_partition = tl.arange(split_x * KCACHE_X, (split_x + 1) * KCACHE_X)
            offsets_k = K + offset_kv + offsets_dmodel_x_partition[None, :] * stride_kd + offsets_m[:, None] * stride_kt
            k = tl.load(offsets_k, mask=offsets_m[:, None] < cur_seq_len, other=0.0)
            # HACK: KCache must be contiguous in order to apply the following offsets calculation
            offsets_kcache = (
                KCache
                + offset_kvcache
                + split_x * BLOCK_SIZE * KCACHE_X
                + block_range[:, None] * KCACHE_X
                + X_range[None, :]
            )
            tl.store(offsets_kcache, k, mask=block_range[:, None] < cur_seq_len - block_start_m * BLOCK_SIZE)
        # Copy v to corresponding cache block
        offsets_vd = tl.arange(0, HEAD_DIM)  # offsets_dmodel
        offsets_vt = block_start_m * BLOCK_N + offsets_n
        offsets_v = V + offset_kv + offsets_vt[None, :] * stride_vt + offsets_vd[:, None] * stride_vd
        v = tl.load(offsets_v, mask=offsets_vt[None, :] < cur_seq_len, other=0.0)
        offsets_vcache = (
            VCache + offset_kvcache + block_range[None, :] * stride_cachebs + offsets_vd[:, None] * stride_cached
        )
        tl.store(offsets_vcache, v, mask=block_range[None, :] < cur_seq_len - block_start_m * BLOCK_SIZE)

    return


# Triton 2.1.0
@triton.jit
def _alibi_fwd_context_paged_attention_kernel(
    Q,
    K,
    V,
    O,
    KCache,
    VCache,
    BLOCK_TABLES,  # [num_seqs, max_blocks_per_sequence]
    batch_size,
    alibi_slopes,
    stride_qt,
    stride_qh,
    stride_qd,
    stride_kt,
    stride_kh,
    stride_kd,
    stride_vt,
    stride_vh,
    stride_vd,
    stride_ot,
    stride_oh,
    stride_od,
    stride_cacheb,
    stride_cacheh,
    stride_cachebs,
    stride_cached,
    stride_bts,
    stride_btb,
    context_lengths,
    sm_scale,
    KV_GROUPS: tl.constexpr,
    BLOCK_SIZE: tl.constexpr,
    HEAD_DIM: tl.constexpr,
    BLOCK_M: tl.constexpr,
    BLOCK_N: tl.constexpr,
):
    cur_seq_idx = tl.program_id(0)
    if cur_seq_idx >= batch_size:
        return
    cur_head_idx = tl.program_id(1)
    block_start_m = tl.program_id(2)  # Br, max_input_len // Block_M
    cur_kv_head_idx = cur_head_idx // KV_GROUPS

    global_block_start_offest = block_start_m * BLOCK_M

    # NOTE It requires BLOCK_M, BLOCK_N, and BLOCK_SIZE to be the same
    tl.static_assert(BLOCK_M == BLOCK_N)
    tl.static_assert(BLOCK_N == BLOCK_SIZE)

    # get the current sequence length from provided context lengths tensor
    cur_seq_len = tl.load(context_lengths + cur_seq_idx)
    # NOTE when talking to fused QKV and a nopadding context attention,
    # we assume that the input Q/K/V is contiguous, and thus here `prev_seq_len_sum`
    # could be considered as the start index of the current sequence.
    # FIXME might want to explore better way to get the summation of prev seq lengths.
    # `tl.sum(tensor[:end])` is invalid as tensor slice is not supported in triton.
    prev_seq_len_sum = 0
    for i in range(0, cur_seq_idx):
        prev_seq_len_sum += tl.load(context_lengths + i)

    offset_q = prev_seq_len_sum * stride_qt + cur_head_idx * stride_qh
    offset_kv = prev_seq_len_sum * stride_kt + cur_kv_head_idx * stride_kh
    Q_block_ptr = tl.make_block_ptr(
        base=Q + offset_q,
        shape=(cur_seq_len, HEAD_DIM),
        strides=(stride_qt, stride_qd),
        offsets=(global_block_start_offest, 0),
        block_shape=(BLOCK_M, HEAD_DIM),
        order=(1, 0),
    )
    K_block_ptr = tl.make_block_ptr(
        base=K + offset_kv,
        shape=(HEAD_DIM, cur_seq_len),
        strides=(stride_kd, stride_kt),
        offsets=(0, 0),
        block_shape=(HEAD_DIM, BLOCK_N),
        order=(0, 1),
    )
    V_block_ptr = tl.make_block_ptr(
        base=V + offset_kv,
        shape=(cur_seq_len, HEAD_DIM),
        strides=(stride_vt, stride_vd),
        offsets=(0, 0),
        block_shape=(BLOCK_N, HEAD_DIM),
        order=(1, 0),
    )
    O_block_ptr = tl.make_block_ptr(
        base=O + offset_q,
        shape=(cur_seq_len, HEAD_DIM),
        strides=(stride_ot, stride_od),
        offsets=(global_block_start_offest, 0),
        block_shape=(BLOCK_M, HEAD_DIM),
        order=(1, 0),
    )

    # block table for the current sequence
    block_table_ptr = BLOCK_TABLES + cur_seq_idx * stride_bts
    # block indexes on block table (i.e. 0, 1, 2, ..., max_blocks_per_seq)
    # Consider `block_start_m` as the logical block idx in the current block table,
    # as we have BLOCK_M the same size as the block size.
    cur_block_table_idx = block_start_m
    cur_block_id = tl.load(block_table_ptr + cur_block_table_idx * stride_btb)
    offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh

    offsets_m = global_block_start_offest + tl.arange(0, BLOCK_M)
    offsets_n = tl.arange(0, BLOCK_N)
    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, HEAD_DIM], dtype=tl.float32)

    # load alibi_slope
    alibi_slope = tl.load(alibi_slopes + cur_head_idx)
    m_alibi_offset = tl.arange(0, BLOCK_M)[:, None] + global_block_start_offest
    n_alibi_offset = tl.arange(0, BLOCK_N)[None, :]

    if global_block_start_offest >= cur_seq_len:
        return

    Q_i = tl.load(Q_block_ptr, boundary_check=(1, 0))

    for block_start_n in range(0, (block_start_m + 1) * BLOCK_M, BLOCK_N):
        block_start_n = tl.multiple_of(block_start_n, BLOCK_N)

        k = tl.load(K_block_ptr, boundary_check=(0, 1))
        S_ij = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
        S_ij += tl.dot(Q_i, k)
        S_ij *= sm_scale
        S_ij += tl.where(offsets_m[:, None] >= (block_start_n + offsets_n[None, :]), 0, float("-inf"))

        alibi = (n_alibi_offset + block_start_n - m_alibi_offset) * alibi_slope
        alibi = tl.where((alibi <= 0) & (m_alibi_offset < cur_seq_len), alibi, float("-inf"))
        S_ij += alibi

        m_ij = tl.max(S_ij, 1)  # rowmax(Sij)
        m_ij = tl.maximum(m_i, m_ij)  # m_ij
        S_ij -= m_ij[:, None]
        p_ij_hat = tl.exp(S_ij)
        scale = tl.exp(m_i - m_ij)
        l_ij = scale * l_i + tl.sum(p_ij_hat, 1)
        acc = acc * scale[:, None]

        v = tl.load(V_block_ptr, boundary_check=(1, 0))
        p_ij_hat = p_ij_hat.to(v.type.element_ty)

        acc += tl.dot(p_ij_hat, v)
        l_i = l_ij
        m_i = m_ij
        K_block_ptr = tl.advance(K_block_ptr, (0, BLOCK_N))
        V_block_ptr = tl.advance(V_block_ptr, (BLOCK_N, 0))

    acc = acc / l_i[:, None]
    tl.store(O_block_ptr, acc.to(O.type.element_ty), boundary_check=(1, 0))

    if cur_head_idx % KV_GROUPS == 0:
        # Copy k to corresponding cache block
        offsets_dmodel = tl.arange(0, HEAD_DIM)
        offsets_kt = global_block_start_offest + tl.arange(0, BLOCK_M)
        offsets_k = K + offset_kv + offsets_dmodel[None, :] * stride_kd + offsets_kt[:, None] * stride_kt
        k = tl.load(offsets_k, mask=offsets_kt[:, None] < cur_seq_len, other=0.0)
        offsets_kcachebs = tl.arange(0, BLOCK_SIZE)
        offsets_kcache = (
            KCache
            + offset_kvcache
            + offsets_dmodel[None, :] * stride_cached
            + offsets_kcachebs[:, None] * stride_cachebs
        )
        tl.store(offsets_kcache, k, mask=offsets_kcachebs[:, None] < cur_seq_len - block_start_m * BLOCK_SIZE)
        # Copy v to corresponding cache block
        offsets_vd = offsets_dmodel
        offsets_vt = block_start_m * BLOCK_N + tl.arange(0, BLOCK_N)
        offsets_v = V + offset_kv + offsets_vt[None, :] * stride_vt + offsets_vd[:, None] * stride_vd
        v = tl.load(offsets_v, mask=offsets_vt[None, :] < cur_seq_len, other=0.0)
        offsets_vcachebs = offsets_kcachebs  # same block size range, just to notify here
        offsets_vcache = (
            VCache
            + offset_kvcache
            + offsets_vcachebs[None, :] * stride_cachebs
            + offsets_dmodel[:, None] * stride_cached
        )
        tl.store(offsets_vcache, v, mask=offsets_vcachebs[None, :] < cur_seq_len - block_start_m * BLOCK_SIZE)

    return


def context_attention_unpadded(
    q: torch.Tensor,  # [num_tokens, num_heads, head_dim]
    k: torch.Tensor,  # [num_tokens, num_kv_heads, head_dim]
    v: torch.Tensor,  # [num_tokens, num_kv_heads, head_dim]
    k_cache: torch.Tensor,  # [num_blocks, num_kv_heads, block_size, head_dim]
    v_cache: torch.Tensor,  # [num_blocks, num_kv_heads, block_size, head_dim]
    context_lengths: torch.Tensor,  # [num_seqs]
    block_tables: torch.Tensor,  # [num_seqs, max_blocks_per_sequence],
    block_size: int,
    output: torch.Tensor = None,  # [num_tokens, num_heads, head_dim]
    alibi_slopes: torch.Tensor = None,  # [num_heads]
    max_seq_len: int = None,
    sm_scale: int = None,
    # NOTE(yuanheng-zhao): the following flag is used to determine whether to use the new layout for kcache
    # [num_blocks, num_kv_heads, head_dim // x, block_size, x] - must be contiguous
    use_new_kcache_layout: bool = False,
):
    Lq, Lk, Lv = q.shape[-1], k.shape[-1], v.shape[-1]
    assert Lq == Lk == Lv
    assert Lk in {32, 64, 128, 256}
    assert q.shape[0] == k.shape[0] == v.shape[0]
    k_cache_shape = k_cache.shape
    v_cache_shape = v_cache.shape
    if use_new_kcache_layout:
        assert (
            len(k_cache_shape) == 5
            and k_cache_shape[1] == v_cache_shape[1]
            and k_cache_shape[2] * k_cache_shape[4] == v_cache_shape[3]
        ), f"Invalid KCache shape {k_cache_shape} and VCache shape {v_cache_shape}"
    else:
        assert k_cache_shape == v_cache_shape, f"Invalid KCache shape {k_cache_shape} and VCache shape {v_cache_shape}"
    assert context_lengths.shape[0] == block_tables.shape[0]

    num_tokens, num_heads, head_dim = q.shape
    num_kv_heads = k.shape[-2]
    assert num_kv_heads > 0 and num_heads % num_kv_heads == 0
    num_kv_group = num_heads // num_kv_heads

    num_seqs, max_blocks_per_seq = block_tables.shape
    max_seq_len = context_lengths.max().item() if max_seq_len is None else max_seq_len
    sm_scale = 1.0 / (Lq**0.5) if sm_scale is None else sm_scale
    output = (
        torch.empty((num_tokens, num_heads * head_dim), dtype=q.dtype, device=q.device) if output is None else output
    )

    # NOTE For now, BLOCK_M and BLOCK_N are supposed to be equivalent with
    # the size of physical cache block (i.e. `block_size`)
    assert block_size in {16, 32, 64, 128}
    BLOCK_M = BLOCK_N = block_size

    # NOTE use `triton.next_power_of_2` here to utilize the cache mechanism of triton
    # To optimize, revise batching/scheduling to batch 2^n sequences in a batch (preferred)
    grid = (triton.next_power_of_2(num_seqs), num_heads, triton.cdiv(max_seq_len, BLOCK_M))

    if use_new_kcache_layout:
        # TODO(yuanheng-zhao): Since the alibi kernel is pretty similar to the original one,
        # the code (alibi kernel) will be refactored later to avoid code duplication, when
        # the whole triton flow with new k cache layout has been supported and tested.
        assert (
            alibi_slopes is None
        ), "Alibi Slopes will be supported with new kcache layout later when the whole triton flow is ready"
        x = k_cache_shape[4]  # Intuition: 16 // dtype_size

        _fwd_context_paged_attention_kernel_v2[grid](
            q,
            k,
            v,
            output,
            k_cache,
            v_cache,
            block_tables,
            num_seqs,
            q.stride(0),
            q.stride(1),
            q.stride(2),
            k.stride(0),
            k.stride(1),
            k.stride(2),
            v.stride(0),
            v.stride(1),
            v.stride(2),
            output.stride(0),
            head_dim,
            1,
            v_cache.stride(0),
            v_cache.stride(1),
            v_cache.stride(2),
            v_cache.stride(3),
            block_tables.stride(0),
            block_tables.stride(1),
            context_lengths,
            sm_scale,
            KV_GROUPS=num_kv_group,
            BLOCK_SIZE=block_size,
            HEAD_DIM=Lk,
            KCACHE_X=x,
            BLOCK_M=BLOCK_M,
            BLOCK_N=BLOCK_N,
        )
        return output

    if alibi_slopes is not None:
        _alibi_fwd_context_paged_attention_kernel[grid](
            q,
            k,
            v,
            output,
            k_cache,
            v_cache,
            block_tables,
            num_seqs,
            alibi_slopes,
            q.stride(0),
            q.stride(1),
            q.stride(2),
            k.stride(0),
            k.stride(1),
            k.stride(2),
            v.stride(0),
            v.stride(1),
            v.stride(2),
            output.stride(0),
            head_dim,
            1,
            k_cache.stride(0),
            k_cache.stride(1),
            k_cache.stride(2),
            k_cache.stride(3),
            block_tables.stride(0),
            block_tables.stride(1),
            context_lengths,
            sm_scale,
            num_kv_group,
            block_size,
            HEAD_DIM=Lk,
            BLOCK_M=BLOCK_M,
            BLOCK_N=BLOCK_N,
        )
    else:
        _fwd_context_paged_attention_kernel[grid](
            q,
            k,
            v,
            output,
            k_cache,
            v_cache,
            block_tables,
            num_seqs,
            q.stride(0),
            q.stride(1),
            q.stride(2),
            k.stride(0),
            k.stride(1),
            k.stride(2),
            v.stride(0),
            v.stride(1),
            v.stride(2),
            output.stride(0),
            head_dim,
            1,
            k_cache.stride(0),
            k_cache.stride(1),
            k_cache.stride(2),
            k_cache.stride(3),
            block_tables.stride(0),
            block_tables.stride(1),
            context_lengths,
            sm_scale,
            num_kv_group,
            block_size,
            HEAD_DIM=Lk,
            BLOCK_M=BLOCK_M,
            BLOCK_N=BLOCK_N,
        )

    return output