From 19061188c396d851ef17bc34b526e2f2b4fc1479 Mon Sep 17 00:00:00 2001 From: Yuanheng Zhao <54058983+yuanheng-zhao@users.noreply.github.com> Date: Mon, 26 Feb 2024 16:17:47 +0800 Subject: [PATCH 1/7] [Infer/Fix] Fix Dependency in test - RMSNorm kernel (#5399) fix dependency in pytest --- tests/test_infer/test_ops/triton/test_rmsnorm_triton.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/tests/test_infer/test_ops/triton/test_rmsnorm_triton.py b/tests/test_infer/test_ops/triton/test_rmsnorm_triton.py index 5ce852164fa1..66e1745d85c5 100644 --- a/tests/test_infer/test_ops/triton/test_rmsnorm_triton.py +++ b/tests/test_infer/test_ops/triton/test_rmsnorm_triton.py @@ -3,13 +3,12 @@ import triton from packaging import version from transformers.models.llama.modeling_llama import LlamaRMSNorm -from vllm.model_executor.layers.layernorm import RMSNorm from colossalai.kernel.triton import rms_layernorm from colossalai.testing.utils import parameterize try: - pass + import triton # noqa HAS_TRITON = True except ImportError: @@ -85,6 +84,11 @@ def benchmark_rms_layernorm( SEQUENCE_TOTAL: int, HIDDEN_SIZE: int, ): + try: + from vllm.model_executor.layers.layernorm import RMSNorm + except ImportError: + raise ImportError("Please install vllm from https://github.com/vllm-project/vllm") + warmup = 10 rep = 1000 From a748841b848ac12a0b053c33bcaf787d5b4fb5c1 Mon Sep 17 00:00:00 2001 From: Yuanheng Date: Thu, 22 Feb 2024 21:10:18 +0800 Subject: [PATCH 2/7] resolve conflicts for revising flash-attn --- colossalai/kernel/triton/flash_decoding.py | 89 +++++++------- .../test_ops/triton/kernel_utils.py | 34 +++--- .../test_ops/triton/test_decoding_attn.py | 109 ++++++++++-------- 3 files changed, 123 insertions(+), 109 deletions(-) diff --git a/colossalai/kernel/triton/flash_decoding.py b/colossalai/kernel/triton/flash_decoding.py index d351b20dadfd..b12a6cef1a0d 100644 --- a/colossalai/kernel/triton/flash_decoding.py +++ b/colossalai/kernel/triton/flash_decoding.py @@ -9,13 +9,14 @@ # Triton 2.1.0 @triton.jit def _flash_decoding_fwd_kernel( - Q, # [batch_size, head_num, q_len(1), head_dim] + Q, # [batch_size * q_len, head_num, head_dim] KCache, # [num_blocks, num_kv_heads, block_size, head_dim] VCache, # [num_blocks, num_kv_heads, block_size, head_dim] block_tables, # [batch_size, max_blocks_per_sequence] - mid_o, # [batch_size, head_num, kv_split_num, head_dim] - mid_o_lse, # [batch_size, head_num, kv_split_num] + mid_o, # [batch_size * q_len, head_num, kv_split_num, head_dim] + mid_o_lse, # [batch_size * q_len, head_num, kv_split_num] kv_seq_len, # [batch_size] + q_len, batch_size, stride_qt, stride_qh, @@ -39,44 +40,37 @@ def _flash_decoding_fwd_kernel( BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr, ): - cur_seq_idx = tl.program_id(0) + cur_token_idx = tl.program_id(0) + cur_seq_idx = cur_token_idx // q_len if cur_seq_idx >= batch_size: return cur_head_idx = tl.program_id(1) block_start_kv = tl.program_id(2) # for splitting k/v - cur_kv_head_idx = cur_head_idx // KV_GROUPS - offsets_dmodel = tl.arange(0, HEAD_DIM) - # NOTE It requires BLOCK_KV and BLOCK_SIZE to be the same # TODO might want to replace with BLOCK_KV % BLOCK_SIZE == 0 (optimize BLOCK_KV as multiple of BLOCK_SIZE) # and then support calculating multiple kv cache blocks on an instance tl.static_assert(BLOCK_KV == BLOCK_SIZE) - - # get the current (kv) sequence length from provided context lengths tensor + # get the current (kv) sequence length cur_kv_seq_len = tl.load(kv_seq_len + cur_seq_idx) + if block_start_kv * BLOCK_KV >= cur_kv_seq_len: + return - offsets_q = cur_seq_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd + offsets_dmodel = tl.arange(0, HEAD_DIM) + offsets_q = cur_token_idx * stride_qt + cur_head_idx * stride_qh + offsets_dmodel * stride_qd q = tl.load(Q + offsets_q) - # block table for the current sequence block_table_ptr = block_tables + cur_seq_idx * stride_bts - - # actually current block table current block start idx # cur_bt_start_idx = block_start_kv * (BLOCK_KV // BLOCK_SIZE) - cur_bt_start_idx = block_start_kv - cur_block_id = tl.load(block_table_ptr + cur_bt_start_idx * stride_btb) - - if block_start_kv * BLOCK_KV >= cur_kv_seq_len: - return - + # cur_block_id = tl.load(block_table_ptr + cur_bt_start_idx * stride_btb) + cur_block_id = tl.load(block_table_ptr + block_start_kv * stride_btb) cur_occupied_size = tl.where( (block_start_kv + 1) * BLOCK_SIZE <= cur_kv_seq_len, BLOCK_SIZE, cur_kv_seq_len - block_start_kv * BLOCK_SIZE ) tl.device_assert(cur_occupied_size >= 0) + cur_kv_head_idx = cur_head_idx // KV_GROUPS offset_kvcache = cur_block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh - K_block_ptr = tl.make_block_ptr( base=KCache + offset_kvcache, shape=(cur_occupied_size, HEAD_DIM), @@ -115,14 +109,14 @@ def _flash_decoding_fwd_kernel( acc = acc / l offsets_mid_o = ( - cur_seq_idx * stride_mid_ot + cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + block_start_kv * stride_mid_ob + offsets_dmodel * stride_mid_od ) tl.store(mid_o + offsets_mid_o, acc) offsets_mid_o_lse = ( - cur_seq_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb + cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb ) # logsumexp L^(j) = m^(j) + log(l^(j)) tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l)) @@ -135,6 +129,7 @@ def _flash_decoding_fwd_reduce_kernel( mid_o_lse, # [batch_size, head_num, kv_split_num] O, # [batch_size, num_heads, head_dim] or [batch_size, 1, num_heads, head_dim] kv_seq_len, + q_len, batch_size, stride_mid_ot, stride_mid_oh, @@ -149,7 +144,8 @@ def _flash_decoding_fwd_reduce_kernel( BLOCK_KV: tl.constexpr, HEAD_DIM: tl.constexpr, ): - cur_seq_idx = tl.program_id(0) + cur_token_idx = tl.program_id(0) + cur_seq_idx = cur_token_idx // q_len if cur_seq_idx >= batch_size: return cur_head_idx = tl.program_id(1) @@ -164,8 +160,8 @@ def _flash_decoding_fwd_reduce_kernel( l = 0.0 # sum exp acc = tl.zeros([HEAD_DIM], dtype=tl.float32) - offsets_mid_o = cur_seq_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel - offset_mid_lse = cur_seq_idx * stride_o_lset + cur_head_idx * stride_o_lseh + offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel + offset_mid_lse = cur_token_idx * stride_o_lset + cur_head_idx * stride_o_lseh for block_i in range(0, kv_split_num, 1): mid_o_block = tl.load(mid_o + offsets_mid_o + block_i * stride_mid_ob) lse = tl.load(mid_o_lse + offset_mid_lse + block_i * stride_o_lseb) @@ -179,7 +175,7 @@ def _flash_decoding_fwd_reduce_kernel( m_i = m_ij acc = acc / l - offsets_O = cur_seq_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel + offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel tl.store(O + offsets_O, acc.to(O.type.element_ty)) return @@ -199,6 +195,7 @@ def flash_decoding_attention( mid_output_lse: torch.Tensor = None, sm_scale: int = None, kv_group_num: int = 1, + q_len: int = 1, ): """ Flash decoding implemented with a blocked KV Cache (PagedAttention) during decoding stage. @@ -218,13 +215,16 @@ def flash_decoding_attention( Log-sum-exp of intermediate output. `max_bsz` should be greater than or equal to `bsz`. block_size (int): Size of each block in the blocked key/value cache. num_kv_group (int, optional): Number of key/value groups. Defaults to 1. + q_length (int): Query length. Use for speculative decoding when `q_length` > 1 (i.e. the last n tokens). + Defaults to 1. Returns: - Output tensor with shape [bsz, num_heads * head_dim] + Output tensor with shape [bsz * q_len, num_heads * head_dim] """ q = q.squeeze() if q.dim() == 4 else q assert q.dim() == 3, f"Incompatible q dim: {q.dim()}" - bsz, num_heads, head_dim = q.shape + n_tokens, num_heads, head_dim = q.shape + bsz = n_tokens // q_len assert head_dim in {32, 64, 128, 256} assert kv_seq_len.shape[0] == block_tables.shape[0] == bsz, ( @@ -247,22 +247,24 @@ def flash_decoding_attention( max_seq_len_in_batch = kv_seq_len.max().item() if max_seq_len_in_batch is None else max_seq_len_in_batch # For compatibility (TODO revise modeling in future) kv_max_split_num = (max_seq_len_in_batch + BLOCK_KV - 1) // BLOCK_KV - mid_output = ( - torch.zeros(size=(bsz, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device) - if mid_output is None - else mid_output - ) - mid_output_lse = ( - torch.zeros(size=(bsz, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device) - if mid_output_lse is None - else mid_output_lse - ) + + if mid_output is None: + mid_output = torch.empty( + (bsz * q_len, num_heads, kv_max_split_num, head_dim), dtype=torch.float32, device=q.device + ) + if mid_output_lse is None: + mid_output_lse = torch.empty((bsz * q_len, num_heads, kv_max_split_num), dtype=torch.float32, device=q.device) + if output is None: + # A hack to prevent `view` operation in modeling + output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device) # 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(bsz), num_heads, triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), BLOCK_KV)) - output = torch.empty((bsz, num_heads * head_dim), dtype=q.dtype, device=q.device) if output is None else output - + grid = ( + triton.next_power_of_2(bsz * q_len), + num_heads, + triton.cdiv(triton.next_power_of_2(max_seq_len_in_batch), BLOCK_KV), + ) _flash_decoding_fwd_kernel[grid]( q, k_cache, @@ -271,6 +273,7 @@ def flash_decoding_attention( mid_output, mid_output_lse, kv_seq_len, + q_len, bsz, q.stride(0), q.stride(1), @@ -295,13 +298,13 @@ def flash_decoding_attention( HEAD_DIM=head_dim, ) - grid = (triton.next_power_of_2(bsz), num_heads) - + grid = (triton.next_power_of_2(bsz * q_len), num_heads) _flash_decoding_fwd_reduce_kernel[grid]( mid_output, mid_output_lse, output, kv_seq_len, + q_len, bsz, mid_output.stride(0), mid_output.stride(1), diff --git a/tests/test_infer/test_ops/triton/kernel_utils.py b/tests/test_infer/test_ops/triton/kernel_utils.py index 22167ded02be..f1ae45477386 100644 --- a/tests/test_infer/test_ops/triton/kernel_utils.py +++ b/tests/test_infer/test_ops/triton/kernel_utils.py @@ -19,12 +19,12 @@ def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: return hidden_states.reshape(bsz, num_key_value_heads * n_rep, seq_len, head_dim) -def prepare_padding_mask(kv_lengths: torch.Tensor, bsz: int, kv_seq_len: int, device="cuda"): - padding_mask = torch.zeros((bsz, 1, 1, kv_seq_len), dtype=torch.float32, device=device) +def prepare_padding_mask(kv_lengths: torch.Tensor, bsz: int, q_len: int, kv_len: int, device="cuda"): + padding_mask = torch.zeros((bsz, 1, q_len, kv_len), dtype=torch.float32, device=device) for i in range(bsz): cur_seq_len = kv_lengths[i].item() - assert cur_seq_len <= kv_seq_len - padding_mask[i, :, :, : kv_seq_len - cur_seq_len] = float("-inf") + assert cur_seq_len <= kv_len + padding_mask[i, :, :, : kv_len - cur_seq_len] = float("-inf") return padding_mask @@ -33,12 +33,12 @@ def prepare_padding_mask(kv_lengths: torch.Tensor, bsz: int, kv_seq_len: int, de # https://github.com/huggingface/transformers/blob/633215ba58fe5114d8c8d32e415a04600e010701/src/transformers/models/llama/modeling_llama.py#L350 def torch_attn_ref( q: torch.Tensor, # [bsz, num_heads, q_len, head_dim] - k: torch.Tensor, # [bsz, num_heads, kv_seq_len, head_dim] - v: torch.Tensor, # [bsz, num_heads, kv_seq_len, head_dim] - attention_mask: torch.Tensor, # [bsz, 1, seq_len, kv_seq_len] + k: torch.Tensor, # [bsz, num_heads, kv_len, head_dim] + v: torch.Tensor, # [bsz, num_heads, kv_len, head_dim] + attention_mask: torch.Tensor, # [bsz, 1, q_len, kv_len] bsz: int, - seq_len: int, - kv_seq_len: int, + q_len: int, + kv_len: int, num_heads: int, num_kv_heads: int, head_dim: int, @@ -54,22 +54,22 @@ def torch_attn_ref( qk = torch.matmul(q, k.transpose(2, 3)) attn_scores = qk / (head_dim**0.5) - assert attn_scores.shape == (bsz, num_heads, seq_len, kv_seq_len), "Invalid shape of attention scores" + + assert attn_scores.shape == (bsz, num_heads, q_len, kv_len), "Invalid shape of attention scores" # for left-side padding - if attention_mask.size() != (bsz, 1, seq_len, kv_seq_len): - raise ValueError( - f"Attention mask should be of size {(bsz, 1, seq_len, kv_seq_len)}, but is {attention_mask.size()}" - ) + if attention_mask.size() != (bsz, 1, q_len, kv_len): + raise ValueError(f"Attention mask should be of size {(bsz, 1, q_len, kv_len)}, but is {attention_mask.size()}") attn_scores = attn_scores + attention_mask attn_weights = F.softmax(attn_scores.to(dtype=torch.float32), dim=-1).to(dtype=q.dtype) out = torch.matmul(attn_weights, v) - if out.size() != (bsz, num_heads, seq_len, head_dim): + if out.size() != (bsz, num_heads, q_len, head_dim): raise ValueError( - f"`attn_output` should be of size {(bsz, num_heads, seq_len, head_dim)}, but is" f" {out.size()}" + f"`attn_output` should be of size {(bsz, num_heads, q_len, head_dim)}, but is" f" {out.size()}" ) out = out.transpose(1, 2).contiguous() - out = out.squeeze(1) + out = out.view(-1, out.size(-2), out.size(-1)) + # out [bsz * q_len, num_heads, head_dim] return out diff --git a/tests/test_infer/test_ops/triton/test_decoding_attn.py b/tests/test_infer/test_ops/triton/test_decoding_attn.py index 4b9b63f7da7b..75f9d4092357 100644 --- a/tests/test_infer/test_ops/triton/test_decoding_attn.py +++ b/tests/test_infer/test_ops/triton/test_decoding_attn.py @@ -21,7 +21,6 @@ TRITON_CUDA_SUPPORT = version.parse(torch.version.cuda) > version.parse("11.4") -Q_LEN = 1 HEAD_DIM = 128 @@ -64,6 +63,7 @@ def prepare_data( @pytest.mark.parametrize("num_attn_heads", [16]) @pytest.mark.parametrize("kv_group_num", [1, 2, 16]) @pytest.mark.parametrize("same_context_len", [True, False]) +@pytest.mark.parametrize("q_len", [1, 5]) def test_flash_decoding( bsz: int, block_size: int, @@ -71,6 +71,7 @@ def test_flash_decoding( num_attn_heads: int, kv_group_num: int, same_context_len: bool, + q_len: int, ): torch.manual_seed(123) torch.cuda.empty_cache() @@ -82,47 +83,57 @@ def test_flash_decoding( max_seq_len = block_size * max_num_blocks_per_seq dtype = torch.float16 device = get_current_device() + q, k_unpad, v_unpad, kv_lengths = prepare_data( + bsz, num_attn_heads, num_kv_heads, HEAD_DIM, same_context_len, q_len, max_seq_len, dtype, device + ) + # The maximum sequence length in the batch (if context lengths randomly generated) + max_kv_len_in_b = kv_lengths.max().item() - q, k_unpad, v_unpad, kv_seq_lengths = prepare_data( - bsz, num_attn_heads, num_kv_heads, HEAD_DIM, same_context_len, Q_LEN, max_seq_len, dtype, device + k_torch = convert_kv_unpad_to_padded(k_unpad, kv_lengths, bsz, max_kv_len_in_b) + v_torch = convert_kv_unpad_to_padded(v_unpad, kv_lengths, bsz, max_kv_len_in_b) + torch_padding_mask = prepare_padding_mask(kv_lengths, bsz, q_len, max_kv_len_in_b, q.device) + out_torch = torch_attn_ref( + q, k_torch, v_torch, torch_padding_mask, bsz, q_len, max_kv_len_in_b, num_attn_heads, num_kv_heads, HEAD_DIM ) + k_cache, v_cache, block_tables = generate_caches_and_block_tables_v2( - k_unpad, v_unpad, kv_seq_lengths, bsz, max_num_blocks_per_seq, block_size, dtype, device + k_unpad, v_unpad, kv_lengths, bsz, max_num_blocks_per_seq, block_size, dtype, device ) block_tables = block_tables.to(device=device) - # The maximum sequence length in the batch (if context lengths randomly generated) - max_seq_len_in_b = kv_seq_lengths.max().item() # The maximum block length splitted on kv should be the kv cache block size - kv_max_split_num = (max_seq_len_in_b + block_size - 1) // block_size - output = torch.empty((bsz, num_attn_heads, HEAD_DIM), dtype=q.dtype, device=q.device) + kv_max_split_num = (max_kv_len_in_b + block_size - 1) // block_size + output = torch.empty((bsz * q_len, num_attn_heads, HEAD_DIM), dtype=q.dtype, device=q.device) mid_output = torch.empty( - size=(bsz, num_attn_heads, kv_max_split_num, HEAD_DIM), dtype=torch.float32, device=q.device + size=(bsz * q_len, num_attn_heads, kv_max_split_num, HEAD_DIM), dtype=torch.float32, device=q.device + ) + mid_output_lse = torch.empty( + size=(bsz * q_len, num_attn_heads, kv_max_split_num), dtype=torch.float32, device=q.device ) - mid_output_lse = torch.empty(size=(bsz, num_attn_heads, kv_max_split_num), dtype=torch.float32, device=q.device) sm_scale = 1.0 / (HEAD_DIM**0.5) + # Here we use different methods to hide the q_len dimension, + # refer to attention forward function in modeling. + if q_len > 1: + q = q.transpose(1, 2).contiguous() # [bsz, q_len, num_heads, head_dim] + q = q.view(-1, q.size(-2), q.size(-1)) # [bsz * q_len, num_heads, head_dim] + else: + q = q.squeeze(2) + assert q.shape == (bsz * q_len, num_attn_heads, HEAD_DIM) + out_triton = flash_decoding_attention( - # Here we use q.squeeze(2) because we hide the q_len dimension (which is equivalent to 1), - # refer to attention forward in modeling. - q.squeeze(2), + q, k_cache, v_cache, - kv_seq_lengths, + kv_lengths, block_tables, block_size, - max_seq_len_in_b, + max_kv_len_in_b, output, mid_output, mid_output_lse, sm_scale=sm_scale, kv_group_num=kv_group_num, - ) # [bsz, 1, num_heads, head_dim] - - k_torch = convert_kv_unpad_to_padded(k_unpad, kv_seq_lengths, bsz, max_seq_len_in_b) - v_torch = convert_kv_unpad_to_padded(v_unpad, kv_seq_lengths, bsz, max_seq_len_in_b) - torch_padding_mask = prepare_padding_mask(kv_seq_lengths, bsz, max_seq_len_in_b, q.device) - out_torch = torch_attn_ref( - q, k_torch, v_torch, torch_padding_mask, bsz, 1, max_seq_len_in_b, num_attn_heads, num_kv_heads, HEAD_DIM - ) + q_len=q_len, + ) # [bsz * q_len, num_heads, head_dim] assert out_torch.shape == out_triton.shape assert torch.allclose(out_torch, out_triton, atol=1e-3, rtol=1e-4) @@ -144,20 +155,13 @@ def test_flash_decoding( styles=[("red", "-"), ("blue", "-")], ylabel="ms", plot_name=f"decoding-block_size-{BLOCK_SIZE}-batch{BATCH}", - args={"bsz": BATCH, "block_size": BLOCK_SIZE, "same_context_len": SAME_LEN, "kv_group_num": 1}, + args={"bsz": BATCH, "block_size": BLOCK_SIZE, "same_context_len": SAME_LEN, "kv_group_num": 1, "q_len": 1}, ) ] @triton.testing.perf_report(configs) -def bench_kernel( - bsz, - KV_LEN, - provider, - block_size: int, - kv_group_num: int, - same_context_len: bool, -): +def bench_kernel(bsz, KV_LEN, provider, block_size: int, kv_group_num: int, same_context_len: bool, q_len: int): num_attn_heads = 16 max_num_blocks_per_seq = triton.cdiv(KV_LEN, block_size) max_seq_len = block_size * max_num_blocks_per_seq @@ -167,19 +171,18 @@ def bench_kernel( block_size * max_num_blocks_per_seq dtype = torch.float16 device = get_current_device() - q, k_unpad, v_unpad, kv_lengths = prepare_data( - bsz, num_attn_heads, num_kv_heads, HEAD_DIM, same_context_len, Q_LEN, max_seq_len, dtype, device + bsz, num_attn_heads, num_kv_heads, HEAD_DIM, same_context_len, q_len, max_seq_len, dtype, device ) - max_seq_len_in_b = kv_lengths.max().item() # for random lengths + max_kv_len_in_b = kv_lengths.max().item() # for random lengths quantiles = [0.5, 0.2, 0.8] if provider == "torch": - k_torch = convert_kv_unpad_to_padded(k_unpad, kv_lengths, bsz, max_seq_len_in_b) - v_torch = convert_kv_unpad_to_padded(v_unpad, kv_lengths, bsz, max_seq_len_in_b) - torch_padding_mask = prepare_padding_mask(kv_lengths, bsz, max_seq_len_in_b, q.device) + k_torch = convert_kv_unpad_to_padded(k_unpad, kv_lengths, bsz, max_kv_len_in_b) + v_torch = convert_kv_unpad_to_padded(v_unpad, kv_lengths, bsz, max_kv_len_in_b) + torch_padding_mask = prepare_padding_mask(kv_lengths, bsz, q_len, max_kv_len_in_b, q.device) fn = lambda: torch_attn_ref( - q, k_torch, v_torch, torch_padding_mask, bsz, 1, max_seq_len_in_b, num_attn_heads, num_kv_heads, HEAD_DIM + q, k_torch, v_torch, torch_padding_mask, bsz, q_len, max_kv_len_in_b, num_attn_heads, num_kv_heads, HEAD_DIM ) ms, min_ms, max_ms = triton.testing.do_bench(fn, warmup=WARM_UPS, rep=REPS, quantiles=quantiles) if provider == "triton": @@ -188,34 +191,42 @@ def bench_kernel( ) block_tables = block_tables.to(device=device) # the maximum block length splitted on kv should be the kv cache block size - kv_max_split_num = (max_seq_len_in_b + block_size - 1) // block_size - output = torch.empty((bsz, num_attn_heads, HEAD_DIM), dtype=dtype, device=device) + kv_max_split_num = (max_kv_len_in_b + block_size - 1) // block_size + output = torch.empty((bsz * q_len, num_attn_heads, HEAD_DIM), dtype=dtype, device=device) mid_output = torch.empty( - size=(bsz, num_attn_heads, kv_max_split_num, HEAD_DIM), dtype=torch.float32, device=q.device + size=(bsz * q_len, num_attn_heads, kv_max_split_num, HEAD_DIM), dtype=torch.float32, device=q.device + ) + mid_output_lse = torch.empty( + size=(bsz * q_len, num_attn_heads, kv_max_split_num), dtype=torch.float32, device=q.device ) - mid_output_lse = torch.empty(size=(bsz, num_attn_heads, kv_max_split_num), dtype=torch.float32, device=q.device) sm_scale = 1.0 / (HEAD_DIM**0.5) + + if q_len > 1: + q = q.transpose(1, 2).contiguous() # [bsz, q_len, num_heads, head_dim] + q = q.view(-1, q.size(-2), q.size(-1)) # [bsz * q_len, num_heads, head_dim] + else: + q = q.squeeze(2) + assert q.shape == (bsz * q_len, num_attn_heads, HEAD_DIM) fn = lambda: flash_decoding_attention( - # Here we use q.squeeze(2) because we hide the q_len dimension (which is equivalent to 1), - # refer to attention forward in modeling. - q.squeeze(2), + q, k_cache, v_cache, kv_lengths, block_tables, block_size, - max_seq_len_in_b, + max_kv_len_in_b, output, mid_output, mid_output_lse, sm_scale=sm_scale, kv_group_num=kv_group_num, - ) # [bsz, 1, num_heads, head_dim] + q_len=q_len, + ) # [bsz * q_len, num_heads, head_dim] ms, min_ms, max_ms = triton.testing.do_bench(fn, warmup=WARM_UPS, rep=REPS, quantiles=quantiles) return ms, min_ms, max_ms if __name__ == "__main__": - test_flash_decoding(16, 32, 32, 16, 1, True) + test_flash_decoding(16, 32, 32, 16, 1, True, 1) # bench_kernel.run(save_path=".", print_data=True) From 0dfe87e49a1d1f3408bd61f8d7e7b555ce4526c2 Mon Sep 17 00:00:00 2001 From: Yuanheng Date: Fri, 23 Feb 2024 23:46:03 +0800 Subject: [PATCH 3/7] adapt kv cache copy kernel for spec-dec --- colossalai/kernel/triton/kvcache_copy.py | 109 +++++++++++++++++- .../test_ops/triton/test_kvcache_copy.py | 73 ++++++++---- 2 files changed, 156 insertions(+), 26 deletions(-) diff --git a/colossalai/kernel/triton/kvcache_copy.py b/colossalai/kernel/triton/kvcache_copy.py index 96ab922e3a9b..5374f3b62a1e 100644 --- a/colossalai/kernel/triton/kvcache_copy.py +++ b/colossalai/kernel/triton/kvcache_copy.py @@ -3,6 +3,50 @@ import triton.language as tl +# Triton 2.1.0 +@triton.jit +def _copy_to_kcache_seqlen_n_kernel( + KV, # K or V + KVCache, # KCache or VCache + BLOCK_TABLES, + context_lengths, + stride_kt, + stride_kh, + stride_kd, + stride_cacheb, + stride_cacheh, + stride_cachebs, + stride_cached, + stride_bts, + stride_btb, + block_size, + n, + HEAD_DIM: tl.constexpr, +): + cur_token_idx = tl.program_id(0) + cur_seq_idx = cur_token_idx // n + cur_token_shift = cur_token_idx - (n * (cur_seq_idx + 1)) + # cur_token_shift = cur_token_idx - n * cur_seq_idx + cur_kv_head_idx = tl.program_id(1) + + past_kv_seq_len = tl.load(context_lengths + cur_seq_idx) + cur_token_shift + last_bt_block_idx = past_kv_seq_len // block_size + block_table_ptr = BLOCK_TABLES + cur_seq_idx * stride_bts + block_id = tl.load(block_table_ptr + last_bt_block_idx * stride_btb) + offsets_in_last_block = past_kv_seq_len % block_size + offsets_dmodel = tl.arange(0, HEAD_DIM) + offsets_kv = cur_seq_idx * stride_kt + cur_kv_head_idx * stride_kh + offsets_dmodel * stride_kd + kv = tl.load(KV + offsets_kv) + offsets_kvcache = ( + block_id * stride_cacheb + + cur_kv_head_idx * stride_cacheh + + offsets_in_last_block * stride_cachebs + + offsets_dmodel * stride_cached + ) + tl.store(KVCache + offsets_kvcache, kv) + return + + # Triton 2.1.0 @triton.jit def _copy_to_kvcache_seqlen1_kernel( @@ -40,10 +84,11 @@ def _copy_to_kvcache_seqlen1_kernel( block_id = tl.load(block_table_ptr + last_bt_block_idx * stride_btb) offsets_in_last_block = past_kv_seq_len % block_size offsets_dmodel = tl.arange(0, HEAD_DIM) - offsets_kv = cur_seq_idx * stride_kt + cur_kv_head_idx * stride_kh + offsets_dmodel * stride_kd + offsets_k = cur_seq_idx * stride_kt + cur_kv_head_idx * stride_kh + offsets_dmodel * stride_kd + offsets_v = cur_seq_idx * stride_vt + cur_kv_head_idx * stride_vh + offsets_dmodel * stride_vd - k = tl.load(K + offsets_kv) - v = tl.load(V + offsets_kv) + k = tl.load(K + offsets_k) + v = tl.load(V + offsets_v) offsets_kcache = ( block_id * stride_cachekb @@ -63,6 +108,64 @@ def _copy_to_kvcache_seqlen1_kernel( return +def copy_k_to_blocked_cache( + k: torch.Tensor, k_cache: torch.Tensor, kv_lengths: torch.Tensor, block_tables: torch.Tensor, n: int = 1 +): + """ + Copy keys or values to the blocked key/value cache during decoding stage. + + Args: + k (torch.Tensor): [bsz, 1, num_kv_heads, head_dim]/[bsz, num_kv_heads, head_dim] - Keys or values during decoding with seq len 1. + [bsz * n, num_kv_heads, head_dim] - Keys or values with seq len n + k_cache (torch.Tensor): [num_blocks, num_kv_heads, block_size, head_dim] - Blocked key or value cache. + kv_lengths (torch.Tensor): [bsz] - Past key/value sequence lengths plus current sequence length for each sequence. + block_tables (torch.Tensor): [bsz, max_blocks_per_sequence] - Block tables for each sequence. + n (int): Number of tokens to copy. Default to 1. + """ + assert k.size(-1) == k_cache.size(-1), "Incompatible head dim" + assert k.dtype == k_cache.dtype, "Expected consistent dtype for tensor and cache." + + k = k.view(-1, k.size(-2), k.size(-1)) if k.dim() == 4 else k + assert k.dim() == 3, f"Invalid k dim {k.dim()}" + bsz, num_kv_heads, head_dim = k.shape + # NOTE when n > 1, the shape of k is [bsz * n, num_kv_heads, head_dim] + if n > 1: + assert bsz % n == 0, "Each sequence should have the same number of tokens to copy" + bsz = bsz // n + + assert kv_lengths.shape[0] == block_tables.shape[0] == bsz, ( + f"Got incompatible batch size (number of seqs):\n" + f" Past kv sequence lengths bsz {kv_lengths.shape[0]}; " + f" block tables bsz {block_tables.shape[0]}, input k batch size {bsz}" + ) + + # Modify if the shape of kv cahce is changed. + block_size = k_cache.size(-2) + + num_warps = 8 if head_dim > 128 else 4 + + grid = (bsz * n, num_kv_heads) + _copy_to_kcache_seqlen_n_kernel[grid]( + k, + k_cache, + block_tables, + kv_lengths, + k.stride(0), + k.stride(1), + k.stride(2), + k_cache.stride(0), + k_cache.stride(1), + k_cache.stride(2), + k_cache.stride(3), + block_tables.stride(0), + block_tables.stride(1), + block_size, + n=n, + HEAD_DIM=head_dim, + num_warps=num_warps, + ) + + def copy_kv_to_blocked_cache( k: torch.Tensor, v: torch.Tensor, diff --git a/tests/test_infer/test_ops/triton/test_kvcache_copy.py b/tests/test_infer/test_ops/triton/test_kvcache_copy.py index 53475270e867..a21de760b30a 100644 --- a/tests/test_infer/test_ops/triton/test_kvcache_copy.py +++ b/tests/test_infer/test_ops/triton/test_kvcache_copy.py @@ -3,7 +3,7 @@ from packaging import version from colossalai.inference.modeling.layers.attention import copy_to_cache -from colossalai.kernel.triton import copy_kv_to_blocked_cache +from colossalai.kernel.triton import copy_k_to_blocked_cache, copy_kv_to_blocked_cache from colossalai.utils import get_current_device from tests.test_infer.test_ops.triton.kernel_utils import generate_caches_and_block_tables_v2, mock_alloc_single_token @@ -28,15 +28,17 @@ def prepare_data( max_num_blocks_per_seq, same_context_len, max_seq_len, + n, device, dtype=torch.float16, ): + assert max_seq_len > n, "max_seq_len must be greater than or equal to n" # past_kv_seq_lengths in this test records the previous kv seq len # (not incorporating the current input whose seq len is 1) past_kv_seq_lengths = ( - torch.tensor([max_seq_len - 1 for _ in range(bsz)], dtype=torch.int32, device=device) + torch.tensor([max_seq_len - n for _ in range(bsz)], dtype=torch.int32, device=device) if same_context_len - else torch.randint(low=1, high=max_seq_len - 1, size=(bsz,), dtype=torch.int32, device=device) + else torch.randint(low=1, high=max_seq_len - n, size=(bsz,), dtype=torch.int32, device=device) ) num_tokens = torch.sum(past_kv_seq_lengths).item() @@ -49,8 +51,8 @@ def prepare_data( ) block_tables = block_tables.to(device=device) - new_k = torch.randn((bsz, 1, num_kv_heads, head_dim), dtype=dtype, device=device) - new_v = torch.randn((bsz, 1, num_kv_heads, head_dim), dtype=dtype, device=device) + new_k = torch.randn((bsz, n, num_kv_heads, head_dim), dtype=dtype, device=device) + new_v = torch.randn((bsz, n, num_kv_heads, head_dim), dtype=dtype, device=device) # mock allocating blocks for the new k/v and update block tables mock_alloc_single_token(block_tables, past_kv_seq_lengths, block_size) # kv seq len = past kv seq len + seq len (1 during decoding stage) @@ -65,12 +67,9 @@ def prepare_data( @pytest.mark.parametrize("max_num_blocks_per_seq", [8, 32]) @pytest.mark.parametrize("num_kv_heads", [16]) @pytest.mark.parametrize("same_context_len", [True, False]) +@pytest.mark.parametrize("n_tokens", [1, 5]) def test_copy_kv_to_caches( - bsz: int, - block_size: int, - max_num_blocks_per_seq: int, - num_kv_heads: int, - same_context_len: bool, + bsz: int, block_size: int, max_num_blocks_per_seq: int, num_kv_heads: int, same_context_len: bool, n_tokens: int ): torch.manual_seed(123) torch.cuda.empty_cache() @@ -89,28 +88,56 @@ def test_copy_kv_to_caches( max_num_blocks_per_seq, same_context_len, max_seq_len, + n_tokens, device=device, dtype=dtype, ) # k_cache_torch = k_cache.clone().detach() # copy_to_cache(new_k, k_cache_torch, lengths=kv_seq_lengths, block_tables=block_tables, type="decoding") - copy_kv_to_blocked_cache(new_k, new_v, k_cache, v_cache, kv_seq_lengths, block_tables) - - past_kv_seq_len = kv_seq_lengths - 1 + k_cache_copy = k_cache.detach().clone() + past_kv_seq_len = kv_seq_lengths - n_tokens target_block_ids = block_tables[range(0, block_tables.size(0)), past_kv_seq_len // block_size] offsets_in_block = past_kv_seq_len % block_size - k_target = k_cache[target_block_ids, :, offsets_in_block, :] - k_source = new_k.squeeze() - v_target = v_cache[target_block_ids, :, offsets_in_block, :] - v_source = new_v.squeeze() + print("offsets_in_block: ", offsets_in_block) + + # n_tokens > 1 + k_target = [] + for i in range(bsz): + block_id_start = target_block_ids[i].item() + offset = offsets_in_block[i].item() + tokens_left = n_tokens + while tokens_left > 0: + tokens_to_fill = min(block_size - offset, tokens_left) + k_target.append(k_cache[block_id_start, :, offset : offset + tokens_to_fill, :]) + block_id_start += 1 # + tokens_left -= tokens_to_fill + offset = 0 + k_target = torch.stack(k_target, dim=0) # [bsz, num_kv_heads, n, head_dim] + k_target = k_target.transpose(1, 2).contiguous() # [bsz, n, num_kv_heads, head_dim] + k_target = k_target.view(-1, k_target.size(-2), k_target.size(-1)) + + # Copy k (or v) to k (or v) cache + copy_k_to_blocked_cache(new_k, k_cache, kv_seq_lengths, block_tables, n=n_tokens) + # k_target = k_cache[target_block_ids, :, offsets_in_block, :] + k_source = new_k.view(-1, new_k.size(-2), new_k.size(-1)) assert k_target.shape == k_source.shape assert torch.equal(k_target, k_source) - assert v_target.shape == v_source.shape - assert torch.equal(v_target, v_source) - # target_torch = k_cache_copy[target_block_ids, :, offsets_in_block, :] - # assert target_torch.shape == source.shape - # assert torch.equal(target_torch, source) + + if n_tokens == 1: + # Copy k and v to k/v caches + k_cache = k_cache_copy + copy_kv_to_blocked_cache(new_k, new_v, k_cache, v_cache, kv_seq_lengths, block_tables) + k_target = k_cache_copy[target_block_ids, :, offsets_in_block, :] + v_target = v_cache[target_block_ids, :, offsets_in_block, :] + v_source = new_v.squeeze() + assert k_target.shape == k_source.shape + assert torch.equal(k_target, k_source) + assert v_target.shape == v_source.shape + assert torch.equal(v_target, v_source) + # target_torch = k_cache_copy[target_block_ids, :, offsets_in_block, :] + # assert target_torch.shape == source.shape + # assert torch.equal(target_torch, source) BATCH = 16 @@ -172,5 +199,5 @@ def benchmark_kvcache_copy( if __name__ == "__main__": - test_copy_kv_to_caches(4, 32, 8, 16, True) + test_copy_kv_to_caches(4, 32, 8, 16, True, 5) # benchmark_kvcache_copy.run(save_path=".", print_data=True) From 39b097e4e6c0900b9cde5875e2702c0fbe6b48cb Mon Sep 17 00:00:00 2001 From: Yuanheng Date: Sat, 24 Feb 2024 19:17:35 +0800 Subject: [PATCH 4/7] fix seqlen-n kvcache copy kernel/tests --- colossalai/kernel/triton/__init__.py | 3 +- colossalai/kernel/triton/kvcache_copy.py | 8 +-- .../test_ops/triton/test_kvcache_copy.py | 59 ++++++++----------- 3 files changed, 32 insertions(+), 38 deletions(-) diff --git a/colossalai/kernel/triton/__init__.py b/colossalai/kernel/triton/__init__.py index 8d41dff13619..35d61a5f7c7a 100644 --- a/colossalai/kernel/triton/__init__.py +++ b/colossalai/kernel/triton/__init__.py @@ -12,7 +12,7 @@ from .flash_decoding import flash_decoding_attention from .fused_rotary_embedding import fused_rotary_embedding from .gptq_triton import gptq_fused_linear_triton - from .kvcache_copy import copy_kv_to_blocked_cache + from .kvcache_copy import copy_k_to_blocked_cache, copy_kv_to_blocked_cache from .no_pad_rotary_embedding import decoding_fused_rotary_embedding, rotary_embedding from .rms_layernorm import rms_layernorm from .rotary_cache_copy import get_xine_cache @@ -21,6 +21,7 @@ __all__ = [ "context_attention_unpadded", "flash_decoding_attention", + "copy_k_to_blocked_cache", "copy_kv_to_blocked_cache", "softmax", "rms_layernorm", diff --git a/colossalai/kernel/triton/kvcache_copy.py b/colossalai/kernel/triton/kvcache_copy.py index 5374f3b62a1e..59e32baf4524 100644 --- a/colossalai/kernel/triton/kvcache_copy.py +++ b/colossalai/kernel/triton/kvcache_copy.py @@ -33,14 +33,14 @@ def _copy_to_kcache_seqlen_n_kernel( last_bt_block_idx = past_kv_seq_len // block_size block_table_ptr = BLOCK_TABLES + cur_seq_idx * stride_bts block_id = tl.load(block_table_ptr + last_bt_block_idx * stride_btb) - offsets_in_last_block = past_kv_seq_len % block_size + offset_last_block = past_kv_seq_len % block_size offsets_dmodel = tl.arange(0, HEAD_DIM) - offsets_kv = cur_seq_idx * stride_kt + cur_kv_head_idx * stride_kh + offsets_dmodel * stride_kd + offsets_kv = cur_token_idx * stride_kt + cur_kv_head_idx * stride_kh + offsets_dmodel * stride_kd kv = tl.load(KV + offsets_kv) offsets_kvcache = ( block_id * stride_cacheb + cur_kv_head_idx * stride_cacheh - + offsets_in_last_block * stride_cachebs + + offset_last_block * stride_cachebs + offsets_dmodel * stride_cached ) tl.store(KVCache + offsets_kvcache, kv) @@ -125,7 +125,7 @@ def copy_k_to_blocked_cache( assert k.size(-1) == k_cache.size(-1), "Incompatible head dim" assert k.dtype == k_cache.dtype, "Expected consistent dtype for tensor and cache." - k = k.view(-1, k.size(-2), k.size(-1)) if k.dim() == 4 else k + k = k.reshape(-1, k.size(-2), k.size(-1)) if k.dim() == 4 else k assert k.dim() == 3, f"Invalid k dim {k.dim()}" bsz, num_kv_heads, head_dim = k.shape # NOTE when n > 1, the shape of k is [bsz * n, num_kv_heads, head_dim] diff --git a/tests/test_infer/test_ops/triton/test_kvcache_copy.py b/tests/test_infer/test_ops/triton/test_kvcache_copy.py index a21de760b30a..24b05c93db3c 100644 --- a/tests/test_infer/test_ops/triton/test_kvcache_copy.py +++ b/tests/test_infer/test_ops/triton/test_kvcache_copy.py @@ -17,7 +17,7 @@ TRITON_CUDA_SUPPORT = version.parse(torch.version.cuda) > version.parse("11.4") -HEAD_DIM = 128 +HEAD_DIM = 32 def prepare_data( @@ -32,9 +32,8 @@ def prepare_data( device, dtype=torch.float16, ): - assert max_seq_len > n, "max_seq_len must be greater than or equal to n" - # past_kv_seq_lengths in this test records the previous kv seq len - # (not incorporating the current input whose seq len is 1) + assert max_seq_len > n, "max_seq_len must be greater than n" + past_kv_seq_lengths = ( torch.tensor([max_seq_len - n for _ in range(bsz)], dtype=torch.int32, device=device) if same_context_len @@ -54,11 +53,11 @@ def prepare_data( new_k = torch.randn((bsz, n, num_kv_heads, head_dim), dtype=dtype, device=device) new_v = torch.randn((bsz, n, num_kv_heads, head_dim), dtype=dtype, device=device) # mock allocating blocks for the new k/v and update block tables - mock_alloc_single_token(block_tables, past_kv_seq_lengths, block_size) - # kv seq len = past kv seq len + seq len (1 during decoding stage) - kv_seq_lengths = past_kv_seq_lengths + 1 + for _ in range(n): + mock_alloc_single_token(block_tables, past_kv_seq_lengths, block_size) + past_kv_seq_lengths += 1 - return new_k, new_v, k_cache, v_cache, kv_seq_lengths, block_tables + return new_k, new_v, k_cache, v_cache, past_kv_seq_lengths, block_tables @pytest.mark.skipif(not (HAS_TRITON and TRITON_CUDA_SUPPORT), reason="requires triton") @@ -92,37 +91,34 @@ def test_copy_kv_to_caches( device=device, dtype=dtype, ) - # k_cache_torch = k_cache.clone().detach() - # copy_to_cache(new_k, k_cache_torch, lengths=kv_seq_lengths, block_tables=block_tables, type="decoding") + k_source = new_k.view(-1, new_k.size(-2), new_k.size(-1)) + v_source = new_v.view(-1, new_v.size(-2), new_v.size(-1)) k_cache_copy = k_cache.detach().clone() - past_kv_seq_len = kv_seq_lengths - n_tokens - target_block_ids = block_tables[range(0, block_tables.size(0)), past_kv_seq_len // block_size] - offsets_in_block = past_kv_seq_len % block_size + past_kv_seq_lengths = kv_seq_lengths - n_tokens + target_block_ids = block_tables[range(0, block_tables.size(0)), past_kv_seq_lengths // block_size] + offsets_in_block = past_kv_seq_lengths % block_size - print("offsets_in_block: ", offsets_in_block) - - # n_tokens > 1 + # Copy k (or v) to k (or v) cache + copy_k_to_blocked_cache(new_k, k_cache, kv_seq_lengths, block_tables, n=n_tokens) + # Reshape target k from k cache to compare if matching with original tensor + # Mainly to handle cases of n_tokens > 1 k_target = [] for i in range(bsz): - block_id_start = target_block_ids[i].item() + block_table = block_tables[i] + curr_kv_len = past_kv_seq_lengths[i].item() offset = offsets_in_block[i].item() tokens_left = n_tokens while tokens_left > 0: tokens_to_fill = min(block_size - offset, tokens_left) - k_target.append(k_cache[block_id_start, :, offset : offset + tokens_to_fill, :]) - block_id_start += 1 # + curr_block_id = block_table[curr_kv_len // block_size] + k_target.append(k_cache[curr_block_id, :, offset : offset + tokens_to_fill, :]) + curr_kv_len += tokens_to_fill tokens_left -= tokens_to_fill offset = 0 - k_target = torch.stack(k_target, dim=0) # [bsz, num_kv_heads, n, head_dim] - k_target = k_target.transpose(1, 2).contiguous() # [bsz, n, num_kv_heads, head_dim] - k_target = k_target.view(-1, k_target.size(-2), k_target.size(-1)) + k_target = torch.concat(k_target, dim=1).transpose(0, 1).contiguous() # [bsz * n, num_kv_heads, head_dim] - # Copy k (or v) to k (or v) cache - copy_k_to_blocked_cache(new_k, k_cache, kv_seq_lengths, block_tables, n=n_tokens) - # k_target = k_cache[target_block_ids, :, offsets_in_block, :] - k_source = new_k.view(-1, new_k.size(-2), new_k.size(-1)) assert k_target.shape == k_source.shape - assert torch.equal(k_target, k_source) + assert torch.allclose(k_target, k_source) if n_tokens == 1: # Copy k and v to k/v caches @@ -130,14 +126,10 @@ def test_copy_kv_to_caches( copy_kv_to_blocked_cache(new_k, new_v, k_cache, v_cache, kv_seq_lengths, block_tables) k_target = k_cache_copy[target_block_ids, :, offsets_in_block, :] v_target = v_cache[target_block_ids, :, offsets_in_block, :] - v_source = new_v.squeeze() assert k_target.shape == k_source.shape assert torch.equal(k_target, k_source) assert v_target.shape == v_source.shape assert torch.equal(v_target, v_source) - # target_torch = k_cache_copy[target_block_ids, :, offsets_in_block, :] - # assert target_torch.shape == source.shape - # assert torch.equal(target_torch, source) BATCH = 16 @@ -183,6 +175,7 @@ def benchmark_kvcache_copy( max_seq_len // block_size, same_context_len, KV_SEQ_LEN, + n=1, device=device, dtype=dtype, ) @@ -199,5 +192,5 @@ def benchmark_kvcache_copy( if __name__ == "__main__": - test_copy_kv_to_caches(4, 32, 8, 16, True, 5) - # benchmark_kvcache_copy.run(save_path=".", print_data=True) + # test_copy_kv_to_caches(4, 16, 8, 16, True, 1) + benchmark_kvcache_copy.run(save_path=".", print_data=True) From 29042a8d06470c282fa630a881c9099295789458 Mon Sep 17 00:00:00 2001 From: Yuanheng Date: Sun, 25 Feb 2024 11:09:27 +0800 Subject: [PATCH 5/7] test kvcache copy - use torch.equal --- tests/test_infer/test_ops/triton/test_kvcache_copy.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_infer/test_ops/triton/test_kvcache_copy.py b/tests/test_infer/test_ops/triton/test_kvcache_copy.py index 24b05c93db3c..6044ad5cd8ac 100644 --- a/tests/test_infer/test_ops/triton/test_kvcache_copy.py +++ b/tests/test_infer/test_ops/triton/test_kvcache_copy.py @@ -118,7 +118,7 @@ def test_copy_kv_to_caches( k_target = torch.concat(k_target, dim=1).transpose(0, 1).contiguous() # [bsz * n, num_kv_heads, head_dim] assert k_target.shape == k_source.shape - assert torch.allclose(k_target, k_source) + assert torch.equal(k_target, k_source) if n_tokens == 1: # Copy k and v to k/v caches From f8a3afd9bea4c33dae1f4632f9d5dc914ada4dde Mon Sep 17 00:00:00 2001 From: Yuanheng Date: Tue, 27 Feb 2024 15:22:53 +0800 Subject: [PATCH 6/7] add assertions --- colossalai/kernel/triton/flash_decoding.py | 17 ++++++++++++++--- colossalai/kernel/triton/kvcache_copy.py | 4 ++-- 2 files changed, 16 insertions(+), 5 deletions(-) diff --git a/colossalai/kernel/triton/flash_decoding.py b/colossalai/kernel/triton/flash_decoding.py index b12a6cef1a0d..e1ccffe53b4a 100644 --- a/colossalai/kernel/triton/flash_decoding.py +++ b/colossalai/kernel/triton/flash_decoding.py @@ -201,7 +201,8 @@ def flash_decoding_attention( Flash decoding implemented with a blocked KV Cache (PagedAttention) during decoding stage. Args: - q (torch.Tensor): [bsz, num_heads, head_dim] + q (torch.Tensor): [bsz * q_len, num_heads, head_dim] + q_len > 1 only for verification process in speculative-decoding. 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] kv_seq_len (torch.Tensor): [batch_size] @@ -209,10 +210,12 @@ def flash_decoding_attention( block_tables (torch.Tensor): [batch_size, max_blocks_per_sequence] max_seq_len_in_batch (int): Maximum sequence length in the batch. output (torch.Tensor): [bsz, num_heads * head_dim] - mid_output (torch.Tensor): [ max_bsz , num_heads, kv_max_split_num, head_dim] + mid_output (torch.Tensor): [max_bsz * q_len, num_heads, kv_max_split_num, head_dim] Intermediate output tensor. `max_bsz` should be greater than or equal to `bsz`. - mid_output_lse (torch.Tensor): [ max_bsz , num_heads, kv_max_split_num] + q_len > 1 only for verification process in speculative-decoding. + mid_output_lse (torch.Tensor): [max_bsz * q_len, num_heads, kv_max_split_num] Log-sum-exp of intermediate output. `max_bsz` should be greater than or equal to `bsz`. + q_len > 1 only for verification process in speculative-decoding. block_size (int): Size of each block in the blocked key/value cache. num_kv_group (int, optional): Number of key/value groups. Defaults to 1. q_length (int): Query length. Use for speculative decoding when `q_length` > 1 (i.e. the last n tokens). @@ -224,6 +227,7 @@ def flash_decoding_attention( q = q.squeeze() if q.dim() == 4 else q assert q.dim() == 3, f"Incompatible q dim: {q.dim()}" n_tokens, num_heads, head_dim = q.shape + assert n_tokens % q_len == 0, "Invalid q_len" bsz = n_tokens // q_len assert head_dim in {32, 64, 128, 256} @@ -258,6 +262,13 @@ def flash_decoding_attention( # A hack to prevent `view` operation in modeling output = torch.empty((bsz * q_len, num_heads * head_dim), dtype=q.dtype, device=q.device) + assert ( + mid_output.size(2) == mid_output_lse.size(2) >= kv_max_split_num + ), "Incompatible kv split number of intermediate output tensors" + assert ( + mid_output.size(0) == mid_output_lse.size(0) >= output.size(0) == n_tokens + ), f"Incompatible first dimension of output tensors" + # 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 = ( diff --git a/colossalai/kernel/triton/kvcache_copy.py b/colossalai/kernel/triton/kvcache_copy.py index 59e32baf4524..871f1f6d8261 100644 --- a/colossalai/kernel/triton/kvcache_copy.py +++ b/colossalai/kernel/triton/kvcache_copy.py @@ -120,7 +120,7 @@ def copy_k_to_blocked_cache( k_cache (torch.Tensor): [num_blocks, num_kv_heads, block_size, head_dim] - Blocked key or value cache. kv_lengths (torch.Tensor): [bsz] - Past key/value sequence lengths plus current sequence length for each sequence. block_tables (torch.Tensor): [bsz, max_blocks_per_sequence] - Block tables for each sequence. - n (int): Number of tokens to copy. Default to 1. + n (int): Number of tokens to copy for each sequence. Default to 1. """ assert k.size(-1) == k_cache.size(-1), "Incompatible head dim" assert k.dtype == k_cache.dtype, "Expected consistent dtype for tensor and cache." @@ -130,7 +130,7 @@ def copy_k_to_blocked_cache( bsz, num_kv_heads, head_dim = k.shape # NOTE when n > 1, the shape of k is [bsz * n, num_kv_heads, head_dim] if n > 1: - assert bsz % n == 0, "Each sequence should have the same number of tokens to copy" + assert bsz % n == 0, "Each sequence should have the same number of tokens to be copied" bsz = bsz // n assert kv_lengths.shape[0] == block_tables.shape[0] == bsz, ( From 8f8e7f070f5bfddbcca3805d7ff325796f6880d6 Mon Sep 17 00:00:00 2001 From: Yuanheng Date: Tue, 27 Feb 2024 15:38:00 +0800 Subject: [PATCH 7/7] (trivial) comment out --- tests/test_infer/test_ops/triton/test_kvcache_copy.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_infer/test_ops/triton/test_kvcache_copy.py b/tests/test_infer/test_ops/triton/test_kvcache_copy.py index 6044ad5cd8ac..9394f6a43b1f 100644 --- a/tests/test_infer/test_ops/triton/test_kvcache_copy.py +++ b/tests/test_infer/test_ops/triton/test_kvcache_copy.py @@ -192,5 +192,5 @@ def benchmark_kvcache_copy( if __name__ == "__main__": - # test_copy_kv_to_caches(4, 16, 8, 16, True, 1) - benchmark_kvcache_copy.run(save_path=".", print_data=True) + test_copy_kv_to_caches(4, 32, 8, 16, True, 1) + # benchmark_kvcache_copy.run(save_path=".", print_data=True)