refactor decode_kv_cache_memcpy

Author: SunflowerAries

colossalai/inference/modeling/models/nopadding_llama.py

@@ -98,15 +98,8 @@ def llama_model_forward(
     """
     block_tables = inputmetadata.block_tables
     sequence_lengths = inputmetadata.sequence_lengths
-    batch_size = inputmetadata.batch_size
     kv_seq_len = inputmetadata.kv_seq_len
 
-    # NOTE: After testing, the performance of this configuration is relatively good. With updates
-    # and optimizations to the CUDA kernel implementation, a more detailed analysis of this configuration's
-    # selection should be conducted.
-    if batch_size >= 32 and kv_seq_len > 512:
-        use_cuda_kernel = False
-
     # NOTE (yuanheng-zhao): fow now, only triton kernels support verification process
     # during speculative-decoding (`q_len > 1`)
     # We will expicitly disable `use_cuda_kernel` here when speculative-decoding is enabled

examples/inference/benchmark_ops/benchmark_flash_decoding_attention.py

@@ -20,7 +20,7 @@
 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,7 @@ 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)
+    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 +137,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(

examples/inference/benchmark_ops/benchmark_fused_rotary_embdding_unpad.py

@@ -105,7 +105,7 @@ def benchmark_rotary_emb(
     elif provider == "no_fused_cuda_rotary_emb_func":
         fn = lambda: [
             inference_ops.rotary_embedding(new_q, new_k, cos, sin, True),
-            inference_ops.decode_kv_cache_memcpy(new_k, new_v, k_cache, v_cache, kv_seq_lengths, block_tables),
+            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(

examples/inference/benchmark_ops/benchmark_kv_cache_memcopy.py

@@ -4,6 +4,7 @@
 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)

extensions/csrc/kernel/cuda/decode_kv_cache_memcpy_kernel.cu

@@ -20,7 +20,8 @@ __global__ void decode_kv_cache_memcpy_kernel(
     const int block_size,
     const int64_t key_stride,
     const int64_t value_stride,
-    const int block_table_stride
+    const int block_table_stride,
+    const int x
 )
 {
     const int seq_id = blockIdx.x;
@@ -38,28 +39,42 @@ __global__ void decode_kv_cache_memcpy_kernel(
     for (; i <= (hidden_size - VecSize); i += blockDim.x * VecSize) {
         const int head_id = i / head_dim;
         const int head_offset = i % head_dim;
+        const int x_id = head_offset / x;
+        const int x_offset = head_offset % x;
         const int64_t key_src_id = seq_id * key_stride + i;
         const int64_t value_src_id = seq_id * value_stride + i;
-        const int64_t target_id = block_id * hidden_size * block_size
+        const int64_t target_key_id = block_id * hidden_size * block_size
+                                      + head_id * block_size * head_dim
+                                      + x_id * block_size * x
+                                      + block_offset * x
+                                      + x_offset;
+        const int64_t target_value_id = block_id * hidden_size * block_size
                                       + head_id * block_size * head_dim
                                       + block_offset * head_dim + head_offset;
 
-        copy_vector<scalar_t, VecSize>(key_cache + target_id, key + key_src_id);
-        copy_vector<scalar_t, VecSize>(value_cache + target_id, value + value_src_id);
+        copy_vector<scalar_t, VecSize>(key_cache + target_key_id, key + key_src_id);
+        copy_vector<scalar_t, VecSize>(value_cache + target_value_id, value + value_src_id);
     }
 
     if (!Aligned) {
         for (; i < hidden_size; ++i ) {
             const int head_id = i / head_dim;
             const int head_offset = i % head_dim;
+            const int x_id = head_offset / x;
+            const int x_offset = head_offset % x;
             const int64_t key_src_id = seq_id * key_stride + i;
             const int64_t value_src_id = seq_id * value_stride + i;
-            const int64_t target_id = block_id * hidden_size * block_size
+            const int64_t target_key_id = block_id * hidden_size * block_size
+                                        + head_id * block_size * head_dim
+                                        + x_id * block_size * x
+                                        + block_offset * x
+                                        + x_offset;
+            const int64_t target_value_id = block_id * hidden_size * block_size
                                         + head_id * block_size * head_dim
                                         + block_offset * head_dim + head_offset;
 
-            key_cache[target_id] = key[key_src_id];
-            value_cache[target_id] = value[value_src_id];
+            key_cache[target_key_id] = key[key_src_id];
+            value_cache[target_value_id] = value[value_src_id];
         }
     }
 
@@ -69,15 +84,16 @@ template<typename scalar_t>
 void apply_decode_kv_cache_memcpy(
     at::Tensor& key,                 // [num_tokens, head_num, head_dim]
     at::Tensor& value,               // [num_tokens, head_num, head_dim]
-    at::Tensor& key_cache,           // [num_blocks, head_num, block_size, head_dim]
+    at::Tensor& key_cache,           // [num_blocks, head_num, head_dim/x, block_size, x]
     at::Tensor& value_cache,         // [num_blocks, head_num, block_size, head_dim]
     at::Tensor& sequence_lengths,    // [batch_size]
     at::Tensor& block_tables)        // [batch_size, max_seq_len]
 {
     int num_tokens = key.size(0);
     int head_num = key.size(1);
     int head_dim = key.size(2);
-    int block_size = key_cache.size(2);
+    int block_size = key_cache.size(3);
+    int x = key_cache.size(4);
 
     int64_t key_stride = key.stride(0);
     int64_t value_stride = value.stride(0);
@@ -110,7 +126,8 @@ void apply_decode_kv_cache_memcpy(
                 block_size,                                                                             \
                 key_stride,                                                                             \
                 value_stride,                                                                           \
-                block_table_stride                                                                      \
+                block_table_stride,                                                                     \
+                x                                                                                       \
             );                                                                                          \
     } while(0)
 
@@ -146,7 +163,7 @@ void apply_decode_kv_cache_memcpy(
 void decode_kv_cache_memcpy(
     at::Tensor& key,                 // [num_tokens, head_num, head_dim]
     at::Tensor& value,               // [num_tokens, head_num, head_dim]
-    at::Tensor& key_cache,           // [num_blocks, head_num, block_size, head_dim]
+    at::Tensor& key_cache,           // [num_blocks, head_num, head_dim/x, block_size, x]
     at::Tensor& value_cache,         // [num_blocks, head_num, block_size, head_dim]
     at::Tensor& sequence_lengths,    // [batch_size]
     at::Tensor& block_tables)        // [batch_size, max_seq_len]

extensions/pybind/inference/inference.cpp

@@ -1,9 +1,10 @@
 #include <torch/extension.h>
 
 void decode_kv_cache_memcpy(
-    torch::Tensor& key,        // [num_tokens, num_heads, head_size]
-    torch::Tensor& value,      // [num_tokens, num_heads, head_size]
-    torch::Tensor& key_cache,  // [num_blocks, num_heads, block_size, head_size]
+    torch::Tensor& key,    // [num_tokens, num_heads, head_size]
+    torch::Tensor& value,  // [num_tokens, num_heads, head_size]
+    torch::Tensor&
+        key_cache,  // [num_blocks, head_num, head_dim/x, block_size, x]
     torch::Tensor&
         value_cache,  // [num_blocks, num_heads, block_size, head_size]
     torch::Tensor& sequence_lengths,  // [batch_size]

tests/test_infer/test_ops/cuda/test_flash_decoding_attention.py

@@ -193,6 +193,7 @@ def test_vllm_flash_decoding_attention(
     max_seq_len_across_batch = kv_seq_lengths.max().item()
     output = torch.empty((BATCH_SIZE, NUM_ATTN_HEADS, HEAD_SIZE), dtype=dtype, device=device)
     sm_scale = 1.0 / (HEAD_SIZE**0.5)
+    kv_scale = 1.0
 
     k_torch = convert_kv_unpad_to_padded(k_unpad, kv_seq_lengths, BATCH_SIZE, max_seq_len_across_batch)
     v_torch = convert_kv_unpad_to_padded(v_unpad, kv_seq_lengths, BATCH_SIZE, max_seq_len_across_batch)
@@ -250,6 +251,7 @@ def test_vllm_flash_decoding_attention(
         max_seq_len_across_batch,
         alibi_slopes,
         "auto",
+        kv_scale,
     )
     numpy_allclose(out_ref, output, rtol=rtol, atol=atol)
 

tests/test_infer/test_ops/cuda/test_kv_cache_memcpy.py

@@ -4,14 +4,42 @@
 
 from colossalai.kernel.kernel_loader import InferenceOpsLoader
 from colossalai.utils import get_current_device
-from tests.test_infer.test_ops.triton.kernel_utils import generate_caches_and_block_tables_v3
-from tests.test_infer.test_ops.triton.test_kvcache_copy import prepare_data
+from tests.test_infer.test_ops.triton.kernel_utils import generate_caches_and_block_tables_v3, mock_alloc_single_token
 
 inference_ops = InferenceOpsLoader().load()
 
 HEAD_DIM = 72
 
 
+def prepare_data(
+    bsz,
+    num_kv_heads,
+    block_size,
+    max_num_blocks_per_seq,
+    context_lengths,
+    device="cuda",
+    dtype=torch.float16,
+):
+    num_tokens = torch.sum(context_lengths).item()
+
+    max_seq_len_in_batch = context_lengths.max()
+    cu_seqlens = F.pad(torch.cumsum(context_lengths, dim=0, dtype=torch.torch.int32), (1, 0))
+
+    kv_size = (num_tokens, num_kv_heads, HEAD_DIM)
+    key = torch.empty(size=kv_size, dtype=dtype, device=device).normal_(mean=0.0, std=0.5)
+    value = torch.empty(size=kv_size, dtype=dtype, device=device).normal_(mean=0.0, std=0.5)
+
+    k_cache_ref, v_cache_ref, block_tables = generate_caches_and_block_tables_v3(
+        key, value, context_lengths, bsz, max_num_blocks_per_seq, block_size, dtype, device
+    )
+
+    block_tables = block_tables.to(device=device)
+    k_cache = torch.zeros_like(k_cache_ref)
+    v_cache = torch.zeros_like(v_cache_ref)
+
+    return key, value, k_cache, v_cache, cu_seqlens, block_tables, max_seq_len_in_batch, k_cache_ref, v_cache_ref
+
+
 def run_decode_copy_kv_to_caches(
     bsz: int,
     block_size: int,
@@ -24,32 +52,41 @@ def run_decode_copy_kv_to_caches(
     torch.cuda.synchronize()
     torch.cuda.reset_peak_memory_stats()
 
+    n = 1
+
     max_seq_len = block_size * max_num_blocks_per_seq
     dtype = torch.float32
     device = get_current_device()
 
-    new_k, new_v, k_cache, v_cache, kv_seq_lengths, block_tables = prepare_data(
-        bsz,
-        num_kv_heads,
-        HEAD_DIM,
-        block_size,
-        max_num_blocks_per_seq,
-        same_context_len,
-        max_seq_len,
-        device=device,
-        dtype=dtype,
+    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
+        else torch.randint(low=1, high=max_seq_len - n, size=(bsz,), dtype=torch.int32, device=device)
+    )
+
+    key, value, k_cache, v_cache, _, block_tables, _, _, _ = prepare_data(
+        bsz, num_kv_heads, block_size, max_num_blocks_per_seq, past_kv_seq_lengths, 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
-    inference_ops.decode_kv_cache_memcpy(new_k, new_v, k_cache, v_cache, kv_seq_lengths, block_tables)
+    new_k = torch.randn((bsz, num_kv_heads, HEAD_DIM), dtype=dtype, device=device)
+    new_v = torch.randn((bsz, num_kv_heads, HEAD_DIM), dtype=dtype, device=device)
 
-    past_kv_seq_len = kv_seq_lengths - 1
+    # mock allocating blocks for the new k/v and update block tables
+    for _ in range(n):
+        mock_alloc_single_token(block_tables, past_kv_seq_lengths, block_size)
+        past_kv_seq_lengths += 1
+
+    inference_ops.decode_kv_cache_memcpy(new_k, new_v, k_cache, v_cache, past_kv_seq_lengths, block_tables)
+
+    past_kv_seq_len = past_kv_seq_lengths - 1
     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_target = k_cache[target_block_ids, :, :, offsets_in_block, :]
     k_source = new_k.squeeze()
     v_target = v_cache[target_block_ids, :, offsets_in_block, :]
+    k_target = k_target.reshape(v_target.shape)
     v_source = new_v.squeeze()
 
     assert k_target.shape == k_source.shape
@@ -77,22 +114,17 @@ def run_context_copy_kv_to_cache(
     else:
         context_lengths = torch.randint(low=1, high=max_seq_len, size=(bsz,), dtype=torch.int32, device=device)
 
-    num_tokens = torch.sum(context_lengths).item()
-
-    max_seq_len_in_batch = context_lengths.max()
-    cu_seqlens = F.pad(torch.cumsum(context_lengths, dim=0, dtype=torch.torch.int32), (1, 0))
-
-    kv_size = (num_tokens, num_kv_heads, HEAD_DIM)
-    key = torch.empty(size=kv_size, dtype=dtype, device=device).normal_(mean=0.0, std=0.5)
-    value = torch.empty(size=kv_size, dtype=dtype, device=device).normal_(mean=0.0, std=0.5)
-
-    k_cache_ref, v_cache_ref, block_tables = generate_caches_and_block_tables_v3(
-        key, value, context_lengths, bsz, max_num_blocks_per_seq, block_size, dtype, device
-    )
-
-    block_tables = block_tables.to(device=device)
-    k_cache = torch.zeros_like(k_cache_ref)
-    v_cache = torch.zeros_like(v_cache_ref)
+    (
+        key,
+        value,
+        k_cache,
+        v_cache,
+        cu_seqlens,
+        block_tables,
+        max_seq_len_in_batch,
+        k_cache_ref,
+        v_cache_ref,
+    ) = prepare_data(bsz, num_kv_heads, block_size, max_num_blocks_per_seq, context_lengths, device, dtype)
 
     inference_ops.context_kv_cache_memcpy(
         key, value, k_cache, v_cache, context_lengths, cu_seqlens, block_tables, max_seq_len_in_batch