[Inference/Kernel] refactor kvcache manager and rotary_embedding and kvcache_memcpy oper… (#5663)

* refactor kvcache manager and rotary_embedding and kvcache_memcpy operator

* refactor decode_kv_cache_memcpy

* enable alibi in pagedattention

examples/inference/benchmark_ops/benchmark_flash_decoding_attention.py

@@ -20,7 +20,7 @@ inference_ops = InferenceOpsLoader().load()
 configs = [
     triton.testing.Benchmark(
         x_names=["MAX_NUM_BLOCKS_PER_SEQ"],
-        x_vals=[2**i for i in range(3, 8)],
+        x_vals=[2**i for i in range(2, 8)],
         line_arg="provider",
         line_vals=[
             "vllm_paged_decoding_attention",
@@ -113,6 +113,8 @@ def benchmark_flash_decoding_attention(
     kv_max_split_num = (max_seq_len_across_batch + BLOCK_SIZE - 1) // BLOCK_SIZE
     output = torch.empty((BATCH_SIZE, NUM_ATTN_HEADS, HEAD_SIZE), dtype=dtype, device=device)
     sm_scale = 1.0 / (HEAD_SIZE**0.5)
+    alibi_slopes = None
+    kv_scale = 1.0
 
     mid_output = torch.empty(
         size=(BATCH_SIZE, NUM_ATTN_HEADS, kv_max_split_num, HEAD_SIZE), dtype=torch.float32, device=device
@@ -136,6 +138,7 @@ def benchmark_flash_decoding_attention(
             max_seq_len_across_batch,
             alibi_slopes,
             "auto",
+            kv_scale,
         )
     elif provider == "triton_flash_decoding_attention":
         fn = lambda: flash_decoding_attention(
@@ -164,6 +167,7 @@ def benchmark_flash_decoding_attention(
             max_seq_len_across_batch,
             mid_output,
             mid_output_lse,
+            alibi_slopes,
             sm_scale,
         )
     else:

examples/inference/benchmark_ops/benchmark_fused_rotary_embdding_unpad.py

@@ -2,7 +2,11 @@ import torch
 
 from colossalai.kernel.kernel_loader import InferenceOpsLoader
 from colossalai.kernel.triton import copy_kv_to_blocked_cache, decoding_fused_rotary_embedding, rotary_embedding
-from tests.test_infer.test_ops.triton.kernel_utils import mock_alloc_block_table_and_kvcache_v2, mock_alloc_single_token
+from tests.test_infer.test_ops.triton.kernel_utils import (
+    mock_alloc_block_table_and_kvcache_v2,
+    mock_alloc_block_table_and_kvcache_v3,
+    mock_alloc_single_token,
+)
 
 inference_ops = InferenceOpsLoader().load()
 
@@ -68,11 +72,17 @@ def benchmark_rotary_emb(
     cache_shape = (BATCH_SIZE * max_num_blocks_per_seq, num_kv_heads, block_size, head_dim)
     k_cache = torch.zeros(size=cache_shape, dtype=dtype, device="cuda")
     v_cache = torch.zeros(size=cache_shape, dtype=dtype, device="cuda")
+    x = 16 // torch.tensor([], dtype=dtype).element_size()
+    new_cache_shape = (BATCH_SIZE * max_num_blocks_per_seq, num_kv_heads, head_dim // x, block_size, x)
+    new_k_cache = torch.zeros(size=new_cache_shape, dtype=dtype, device="cuda")
 
     past_kv_seq_lengths = torch.tensor([SEQ_LEN - 1 for _ in range(BATCH_SIZE)], dtype=torch.int32, device="cuda")
     block_tables = mock_alloc_block_table_and_kvcache_v2(
         k, v, k_cache, v_cache, past_kv_seq_lengths, BATCH_SIZE, max_num_blocks_per_seq, block_size
     )
+    _ = mock_alloc_block_table_and_kvcache_v3(
+        k, v, new_k_cache, v_cache, past_kv_seq_lengths, BATCH_SIZE, max_num_blocks_per_seq, block_size
+    )
     new_k = torch.randn((BATCH_SIZE, num_kv_heads, head_dim), dtype=dtype, device="cuda")
     new_q = torch.randn_like(new_k)
     new_v = torch.randn_like(new_k)
@@ -94,12 +104,12 @@ def benchmark_rotary_emb(
         )
     elif provider == "no_fused_cuda_rotary_emb_func":
         fn = lambda: [
-            inference_ops.rotary_embedding(new_q, new_k, cos, sin),
-            inference_ops.decode_kv_cache_memcpy(new_k, new_v, k_cache, v_cache, kv_seq_lengths, block_tables),
+            inference_ops.rotary_embedding(new_q, new_k, cos, sin, True),
+            inference_ops.decode_kv_cache_memcpy(new_k, new_v, new_k_cache, v_cache, kv_seq_lengths, block_tables),
         ]
     elif provider == "fused_cuda_rotary_emb_func":
         fn = lambda: inference_ops.rotary_embedding_and_cache_copy(
-            new_q, new_k, new_v, cos, sin, k_cache, v_cache, kv_seq_lengths, block_tables
+            new_q, new_k, new_v, cos, sin, new_k_cache, v_cache, kv_seq_lengths, block_tables, True
         )
     else:
         raise ValueError("Undefined provider")

examples/inference/benchmark_ops/benchmark_kv_cache_memcopy.py

@@ -4,6 +4,7 @@ from colossalai.inference.modeling.layers.attention import copy_to_cache
 from colossalai.kernel.kernel_loader import InferenceOpsLoader
 from colossalai.kernel.triton import copy_kv_to_blocked_cache
 from colossalai.utils import get_current_device
+from tests.test_infer.test_ops.cuda.test_kv_cache_memcpy import prepare_data as prepare_data_new_kcache_layout
 from tests.test_infer.test_ops.triton.test_kvcache_copy import prepare_data
 
 try:
@@ -68,6 +69,9 @@ def benchmark_kvcache_copy(
     elif provider == "triton_copy_func":
         fn = lambda: copy_kv_to_blocked_cache(new_k, new_v, k_cache, v_cache, context_lengths, block_tables)
     elif provider == "cuda_copy_func":
+        _, _, k_cache, _, _, _, _, _, _ = prepare_data_new_kcache_layout(
+            bsz, num_kv_heads, block_size, max_seq_len // block_size, context_lengths - 1, device, dtype
+        )
         new_k = new_k.squeeze(1) if new_k.dim() == 4 else new_k
         new_v = new_v.squeeze(1) if new_v.dim() == 4 else new_v
         fn = lambda: inference_ops.decode_kv_cache_memcpy(new_k, new_v, k_cache, v_cache, context_lengths, block_tables)