@triton.heuristics({'M_LT_N': lambda kwargs: kwargs['BLOCK_M'] < kwargs[
    'BLOCK_N']})
@triton.jit
def _fwd_kernel_batch_inference(Q, K, V, Out, sm_scale, q_batch_starts,
    q_batch_ends, k_batch_starts, k_batch_ends, q_batch_ids, q_start_sids,
    stride_qb, stride_qt, stride_qh, stride_qd, stride_kb, stride_kt,
    stride_kh, stride_kd, stride_vb, stride_vt, stride_vh, stride_vd,
    stride_ob, stride_ot, stride_oh, stride_od, layout_crow_ptr,
    layout_col_ptr, layout_crow_stride_h, layout_crow_stride_m,
    layout_col_stride_h, layout_col_stride_m, q_k_ratio, HAS_BATCH_DIM: tl.
    constexpr, D_HEAD: tl.constexpr, BLOCK_M: tl.constexpr, BLOCK_N: tl.
    constexpr, BLOCK_D: tl.constexpr, BLOCK_M_LOADING: tl.constexpr, EVEN_D:
    tl.constexpr, M_LT_N: tl.constexpr):
    """
    NOTATION:
    pid: position id
    sid: storage id
    sbid: storage block id
    pbid: position block id
    offs_m, offs_n: storage offsets of m-dim(q, row) and n-dim(k, col)

    TODO(linxihui):
    Optimize grouped-attn
    """
    off_zm = tl.program_id(0)
    off_h = tl.program_id(1)
    off_h_for_kv = off_h // q_k_ratio
    if HAS_BATCH_DIM:
        off_z = tl.program_id(2)
        Q += off_z * stride_qb
        K += off_z * stride_kb
        V += off_z * stride_vb
        Out += off_z * stride_ob
        start_m = off_zm
        q_start_sid = start_m * BLOCK_M
    else:
        off_z = tl.load(q_batch_ids + off_zm).to(tl.int32)
        q_start_sid = tl.load(q_start_sids + off_zm)
        start_m = q_start_sid // BLOCK_M
    offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M_LOADING)
    offs_n = tl.arange(0, BLOCK_N)
    offs_d = tl.arange(0, BLOCK_D)
    q_cu_start = tl.load(q_batch_starts + off_z).to(tl.int32)
    q_seqlen = tl.load(q_batch_ends + off_z).to(tl.int32) - q_cu_start
    k_cu_start = tl.load(k_batch_starts + off_z).to(tl.int32)
    k_seqlen = tl.load(k_batch_ends + off_z).to(tl.int32) - k_cu_start
    past_len = k_seqlen - q_seqlen
    Q += q_cu_start * stride_qt + off_h * stride_qh
    K += k_cu_start * stride_kt + off_h_for_kv * stride_kh
    V += k_cu_start * stride_vt + off_h_for_kv * stride_vh
    Out += q_cu_start * stride_ot + off_h * stride_oh
    q_pbid = (past_len + q_start_sid) // BLOCK_M
    if EVEN_D:
        q = tl.load(Q + offs_m[:, None] * stride_qt + offs_d[None, :] *
            stride_qd, mask=offs_m[:, None] < q_seqlen)
    else:
        q = tl.load(Q + offs_m[:, None] * stride_qt + offs_d[None, :] *
            stride_qd, mask=(offs_m[:, None] < q_seqlen) & (offs_d[None, :] <
            D_HEAD), other=0)
    sparse_crow_ptr = (layout_crow_ptr + off_h * layout_crow_stride_h + 
        q_pbid * layout_crow_stride_m)
    k_block_start = tl.load(sparse_crow_ptr).to(tl.int32)
    k_block_end = tl.load(sparse_crow_ptr + 1).to(tl.int32)
    m_i = tl.zeros([BLOCK_M_LOADING], dtype=tl.float32) - float('inf')
    l_i = tl.zeros([BLOCK_M_LOADING], dtype=tl.float32)
    acc = tl.zeros([BLOCK_M_LOADING, BLOCK_D], 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_d[None, :] * stride_vd
    sm_scale *= 1.44269504
    for k_block_col_idx in range(k_block_start, k_block_end - 1):
        acc, l_i, m_i = _fwd_kernel_inner(acc, l_i, m_i, q, Q,
            k_block_col_idx, layout_col_ptr, layout_col_stride_h,
            layout_col_stride_m, k_ptrs, v_ptrs, off_h, offs_m, offs_n,
            offs_d, stride_kt, stride_vt, sm_scale, k_seqlen, past_len, 
            False, BLOCK_M_LOADING, BLOCK_N, D_HEAD, EVEN_D, M_LT_N)
    acc, l_i, m_i = _fwd_kernel_inner(acc, l_i, m_i, q, Q, k_block_end - 1,
        layout_col_ptr, layout_col_stride_h, layout_col_stride_m, k_ptrs,
        v_ptrs, off_h, offs_m, offs_n, offs_d, stride_kt, stride_vt,
        sm_scale, k_seqlen, past_len, True, BLOCK_M_LOADING, BLOCK_N,
        D_HEAD, EVEN_D, M_LT_N)
    m_i += tl.math.log2(l_i)
    acc = acc / l_i[:, None]
    if EVEN_D:
        tl.store(Out + offs_m[:, None] * stride_ot + offs_d[None, :] *
            stride_od, acc, mask=offs_m[:, None] < q_seqlen)
    else:
        tl.store(Out + offs_m[:, None] * stride_ot + offs_d[None, :] *
            stride_od, acc, mask=(offs_m[:, None] < q_seqlen) & (offs_d[
            None, :] < D_HEAD))
