def METHOD_NAME(q, k, v, o, b_start_loc, b_seq_len, max_input_len, alibi=None):
    BLOCK = 128
    # shape constraints
    Lq, Lk, Lv = q.shape[-1], k.shape[-1], v.shape[-1]
    assert Lq == Lk, "context process only supports equal query, key, value length"
    assert Lk == Lv, "context process only supports equal query, key, value length"
    assert Lk in {16, 32, 64, 128}
    sm_scale = 1.0 / math.sqrt(Lk)
    batch, head = b_seq_len.shape[0], q.shape[1]
    grid = (batch, head, triton.cdiv(max_input_len, BLOCK))
    num_warps = 4 if Lk <= 64 else 8
    tmp = torch.empty((batch, head, max_input_len + 256), device=q.device, dtype=torch.float32)
    _context_flash_attention_kernel[grid](
        q, k, v, sm_scale, 
        b_start_loc, b_seq_len,
        tmp,
        alibi,
        o,
        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),
        o.stride(0), o.stride(1), o.stride(2),
        tmp.stride(0), tmp.stride(1), tmp.stride(2),
        # manually setting this blcok num, we can use tuning config to futher speed-up 
        BLOCK_M=BLOCK,
        BLOCK_DMODEL=Lk,
        BLOCK_N=BLOCK,
        num_warps=num_warps,
        num_stages=1,
    )
    return