[feat] added support for kv_cache with different strides. (#143)

Author: guocuimi

src/kernels/kv_cache_kernels.cu

@@ -12,10 +12,11 @@ __global__ void set_kv_cache_kernel(
     const T* __restrict__ values,      // [n_tokens, n_heads, head_dim]
     T* __restrict__ key_cache,
     T* __restrict__ value_cache,
-    int kv_stride,
-    int n_kv_heads,
-    int head_dim,
-    int block_size) {
+    int64_t k_stride,
+    int64_t v_stride,
+    int64_t n_kv_heads,
+    int64_t head_dim,
+    int64_t block_size) {
   // block/token index
   const int64_t bid = blockIdx.x;
   // which slot to write to
@@ -29,8 +30,9 @@ __global__ void set_kv_cache_kernel(
   const int64_t block_base_idx = block_idx * block_size * n_kv_heads * head_dim;
 
   // copy value one by one for the token
-  for (int i = threadIdx.x; i < n_kv_heads * head_dim; i += blockDim.x) {
-    const int64_t src_idx = bid * kv_stride + i;
+  for (int64_t i = threadIdx.x; i < n_kv_heads * head_dim; i += blockDim.x) {
+    const int64_t k_src_idx = bid * k_stride + i;
+    const int64_t v_src_idx = bid * v_stride + i;
 
     // cache: [n_blocks, block_size, n_heads, head_dim]
     const int64_t head_base_idx =
@@ -42,8 +44,8 @@ __global__ void set_kv_cache_kernel(
     const int head_offset = i % head_dim;
     const int64_t dst_idx = head_base_idx + head_idx * head_dim + head_offset;
 
-    key_cache[dst_idx] = keys[src_idx];
-    value_cache[dst_idx] = values[src_idx];
+    key_cache[dst_idx] = keys[k_src_idx];
+    value_cache[dst_idx] = values[v_src_idx];
   }
 }
 
@@ -53,15 +55,21 @@ void set_kv_cache(
     const torch::Tensor& values,    // [n_tokens, n_kv_heads, head_dim]
     torch::Tensor& key_cache,       // [n_blocks, block_size, n_heads, head_dim]
     torch::Tensor& value_cache) {
-  const int n_tokens = keys.size(0);
-  const int n_kv_heads = keys.size(-2);
-  const int head_dim = keys.size(-1);
-  const int block_size = key_cache.size(-3);
-  const int kv_stride = keys.stride(0);
-  const int n = n_kv_heads * head_dim;
+  // keys and values should be continuous at n_kv_heads and head_dim dims
+  CHECK(keys.stride(-1) == 1 && keys.stride(-2) == keys.size(-1));
+  CHECK(values.stride(-1) == 1 && values.stride(-2) == values.size(-1));
+
+  const int64_t n_tokens = keys.size(-3);
+  const int64_t n_kv_heads = keys.size(-2);
+  const int64_t head_dim = keys.size(-1);
+  const int64_t block_size = key_cache.size(-3);
+  // it is possible that keys and values have different strides
+  const int64_t k_stride = keys.stride(-3);
+  const int64_t v_stride = values.stride(-3);
+  const int64_t n = n_kv_heads * head_dim;
 
   dim3 grid(n_tokens);
-  dim3 block(std::min(n, 1024));
+  dim3 block(std::min<int>(n, 1024));
   DISPATCH_FLOATING_TYPES(keys.scalar_type(), "set_kv_cache_kernel", [&] {
     set_kv_cache_kernel<scalar_t>
         <<<grid, block, 0, at::cuda::getCurrentCUDAStream()>>>(
@@ -70,7 +78,8 @@ void set_kv_cache(
             values.data_ptr<scalar_t>(),
             key_cache.data_ptr<scalar_t>(),
             value_cache.data_ptr<scalar_t>(),
-            kv_stride,
+            k_stride,
+            v_stride,
             n_kv_heads,
             head_dim,
             block_size);

src/kernels/layernorm_kernels.cu

@@ -16,14 +16,14 @@ __global__ void rms_norm_kernel(T* __restrict__ out,
                                 const T* __restrict__ input,
                                 const T* __restrict__ weight,
                                 const float epsilon,
-                                int n) {
-  const int tidx = threadIdx.x;
-  const int bidx = blockIdx.x;
+                                int64_t n) {
+  const auto tidx = threadIdx.x;
+  const auto bidx = blockIdx.x;
 
   __shared__ float s_variance;
   float variance = 0.0f;
 
-  for (int i = tidx; i < n; i += blockDim.x) {
+  for (int64_t i = tidx; i < n; i += blockDim.x) {
     const float x = input[bidx * n + i];
     variance += x * x;
   }
@@ -33,8 +33,8 @@ __global__ void rms_norm_kernel(T* __restrict__ out,
   }
   __syncthreads();
 
-  for (int i = tidx; i < n; i += blockDim.x) {
-    const int idx = bidx * n + i;
+  for (int64_t i = tidx; i < n; i += blockDim.x) {
+    const int64_t idx = bidx * n + i;
     const float x = input[idx];
     out[idx] = (T)(x * s_variance) * weight[i];
   }
@@ -47,10 +47,10 @@ void rms_norm(torch::Tensor& out,
   DCHECK(input.is_contiguous()) << "input tensor must be contiguous";
   DCHECK(out.is_contiguous()) << "output tensor must be contiguous";
 
-  const int n = input.size(1);
+  const int64_t n = input.size(1);
 
   dim3 grid(input.size(0));
-  dim3 block(std::min(n, 1024));
+  dim3 block(std::min<int>(n, 1024));
   DISPATCH_FLOATING_TYPES(input.scalar_type(), "rms_norm_kernel", [&] {
     rms_norm_kernel<scalar_t>
         <<<grid, block, 0, at::cuda::getCurrentCUDAStream()>>>(
@@ -73,15 +73,15 @@ __global__ void rms_norm_residual_kernel(T* __restrict__ out,
                                          const T* __restrict__ input,
                                          const T* __restrict__ weight,
                                          const float epsilon,
-                                         int n) {
-  const int tidx = threadIdx.x;
-  const int bidx = blockIdx.x;
+                                         int64_t n) {
+  const auto tidx = threadIdx.x;
+  const auto bidx = blockIdx.x;
 
   __shared__ float s_variance;
   float variance = 0.0f;
 
-  for (int i = tidx; i < n; i += blockDim.x) {
-    const int idx = bidx * n + i;
+  for (int64_t i = tidx; i < n; i += blockDim.x) {
+    const int64_t idx = bidx * n + i;
     const float r = residual[idx];
     const float x = r + input[idx];
     residual[idx] = x;
@@ -93,8 +93,8 @@ __global__ void rms_norm_residual_kernel(T* __restrict__ out,
   }
   __syncthreads();
 
-  for (int i = tidx; i < n; i += blockDim.x) {
-    const int idx = bidx * n + i;
+  for (int64_t i = tidx; i < n; i += blockDim.x) {
+    const int64_t idx = bidx * n + i;
     const float x = residual[idx];
     out[idx] = (T)(x * s_variance) * weight[i];
   }
@@ -109,10 +109,10 @@ void rms_norm_residual(torch::Tensor& out,
   DCHECK(out.is_contiguous()) << "output tensor must be contiguous";
   DCHECK(residual.is_contiguous()) << "residual tensor must be contiguous";
 
-  const int n = input.size(1);
+  const int64_t n = input.size(1);
 
   dim3 grid(input.size(0));
-  dim3 block(std::min(n, 1024));
+  dim3 block(std::min<int>(n, 1024));
   DISPATCH_FLOATING_TYPES(input.scalar_type(), "rms_norm_residual_kernel", [&] {
     rms_norm_residual_kernel<scalar_t>
         <<<grid, block, 0, at::cuda::getCurrentCUDAStream()>>>(
@@ -133,17 +133,17 @@ __global__ void layer_norm_kernel(T* __restrict__ out,
                                   const T* __restrict__ weight,
                                   const T* __restrict__ bias,
                                   const float epsilon,
-                                  int n) {
-  const int tidx = threadIdx.x;
-  const int bidx = blockIdx.x;
+                                  int64_t n) {
+  const auto tidx = threadIdx.x;
+  const auto bidx = blockIdx.x;
 
   __shared__ float s_mean;
   __shared__ float s_variance;
   float mean = 0.0f;
   float variance = 0.0f;
 
   // calculate mean of the input.
-  for (int i = tidx; i < n; i += blockDim.x) {
+  for (int64_t i = tidx; i < n; i += blockDim.x) {
     mean += input[bidx * n + i];
   }
   mean = block_reduce_sum<float>(mean);
@@ -153,7 +153,7 @@ __global__ void layer_norm_kernel(T* __restrict__ out,
   __syncthreads();
 
   // calculate variance of the input.
-  for (int i = tidx; i < n; i += blockDim.x) {
+  for (int64_t i = tidx; i < n; i += blockDim.x) {
     const float x = input[bidx * n + i] - s_mean;
     variance += x * x;
   }
@@ -163,8 +163,8 @@ __global__ void layer_norm_kernel(T* __restrict__ out,
   }
   __syncthreads();
 
-  for (int i = tidx; i < n; i += blockDim.x) {
-    const int idx = bidx * n + i;
+  for (int64_t i = tidx; i < n; i += blockDim.x) {
+    const int64_t idx = bidx * n + i;
     float local_out = (input[idx] - s_mean) * s_variance * weight[i];
     if (bias != nullptr) {
       local_out += bias[i];
@@ -181,10 +181,10 @@ void layer_norm(torch::Tensor& out,
   DCHECK(input.is_contiguous()) << "input tensor must be contiguous";
   DCHECK(out.is_contiguous()) << "output tensor must be contiguous";
 
-  const int n = input.size(1);
+  const int64_t n = input.size(1);
 
   dim3 grid(input.size(0));
-  dim3 block(std::min(n, 1024));
+  dim3 block(std::min<int>(n, 1024));
   DISPATCH_FLOATING_TYPES(input.scalar_type(), "layer_norm_kernel", [&] {
     layer_norm_kernel<scalar_t>
         <<<grid, block, 0, at::cuda::getCurrentCUDAStream()>>>(

src/kernels/pos_embedding_kernels.cu

@@ -33,47 +33,47 @@ struct RotaryEmbedding {
 // inplace update query and key
 template <typename T>
 __global__ void rotary_embedding_kernel(
-    T* __restrict__ query,              // [n_tokens, n_heads, head_dim]
-    T* __restrict__ key,                // [n_tokens, n_kv_heads, head_dim]
+    T* __restrict__ querys,             // [n_tokens, n_heads, head_dim]
+    T* __restrict__ keys,               // [n_tokens, n_kv_heads, head_dim]
     const int* __restrict__ positions,  // [n_tokens]
     const T* __restrict__ cos_sin,      // [max_positions, 2, rotary_dim/2]
-    int head_dim,
-    int rotary_dim,
-    int n_heads,
-    int n_kv_heads,
-    int q_stride,
-    int k_stride,
+    int64_t head_dim,
+    int64_t rotary_dim,
+    int64_t n_heads,
+    int64_t n_kv_heads,
+    int64_t q_stride,
+    int64_t k_stride,
     bool interleaved) {
   const int tidx = threadIdx.x;
   const int bidx = blockIdx.x;
 
   // figure out cos sin base ptr for the token
-  const int n = rotary_dim / 2;
+  const int64_t n = rotary_dim / 2;
   const T* cos_sin_base = cos_sin + positions[bidx] * rotary_dim;
   const T* cos = cos_sin_base;
   const T* sin = cos_sin_base + n;
 
   // apply rotary embedding to query head by head
   // q base ptr for the token
-  T* q_base = query + bidx * q_stride;
-  for (int i = tidx; i < n_heads * n; i += blockDim.x) {
+  T* q_base = querys + bidx * q_stride;
+  for (int64_t i = tidx; i < n_heads * n; i += blockDim.x) {
     // head idx
-    const int h_idx = i / n;
+    const int64_t h_idx = i / n;
     // rotary idx within head
-    const int r_idx = i % n;
+    const int64_t r_idx = i % n;
     // q ptr for the head
     T* q = q_base + h_idx * head_dim;
     RotaryEmbedding<T>::apply(q, cos, sin, r_idx, n, interleaved);
   }
 
   // apply rotary embedding to key head by head
   // k base ptr for the token
-  T* k_base = key + bidx * k_stride;
-  for (int i = tidx; i < n_kv_heads * n; i += blockDim.x) {
+  T* k_base = keys + bidx * k_stride;
+  for (int64_t i = tidx; i < n_kv_heads * n; i += blockDim.x) {
     // head idx
-    const int h_idx = i / n;
+    const int64_t h_idx = i / n;
     // rotary idx within head
-    const int r_idx = i % n;
+    const int64_t r_idx = i % n;
     // k ptr for the head
     T* k = k_base + h_idx * head_dim;
     RotaryEmbedding<T>::apply(k, cos, sin, r_idx, n, interleaved);
@@ -82,31 +82,30 @@ __global__ void rotary_embedding_kernel(
 
 // apply rotary embedding to query and key inplace
 void apply_rotary_pos_emb(
-    torch::Tensor& query,            // [n_tokens, n_heads, head_dim]
-    torch::Tensor& key,              // [n_tokens, n_kv_heads, head_dim]
+    torch::Tensor& querys,           // [n_tokens, n_heads, head_dim]
+    torch::Tensor& keys,             // [n_tokens, n_kv_heads, head_dim]
     const torch::Tensor& positions,  // [n_tokens]
     const torch::Tensor& cos_sin,    // [max_positions, 2, rotary_dim/2]
     int rotary_dim,
     bool interleaved) {
-  DCHECK(query.is_cuda()) << "query must be on gpu";
-  DCHECK(key.is_cuda()) << "key must be on gpu";
-  DCHECK(query.dim() == 3) << "query must be 3d";
-  DCHECK(key.dim() == 3) << "key must be 3d";
+  // keys and values should be continuous at n_kv_heads and head_dim dims
+  CHECK(querys.stride(-1) == 1 && querys.stride(-2) == querys.size(-1));
+  CHECK(keys.stride(-1) == 1 && keys.stride(-2) == keys.size(-1));
 
-  const int n_tokens = query.size(0);
-  const int n_heads = query.size(1);
-  const int n_kv_heads = key.size(1);
-  const int head_dim = query.size(2);
-  const int q_stride = query.stride(0);
-  const int k_stride = key.stride(0);
+  const int64_t n_tokens = querys.size(-3);
+  const int64_t n_heads = querys.size(-2);
+  const int64_t n_kv_heads = keys.size(-2);
+  const int64_t head_dim = querys.size(-1);
+  const int64_t q_stride = querys.stride(-3);
+  const int64_t k_stride = keys.stride(-3);
 
   const dim3 grid(n_tokens);
-  const dim3 block(std::min(1024, n_heads * rotary_dim) / 2);
-  DISPATCH_FLOATING_TYPES(query.scalar_type(), "rotary_embedding_kernel", [&] {
+  const dim3 block(std::min<int>(1024, n_heads * rotary_dim) / 2);
+  DISPATCH_FLOATING_TYPES(querys.scalar_type(), "rotary_embedding_kernel", [&] {
     rotary_embedding_kernel<scalar_t>
         <<<grid, block, 0, at::cuda::getCurrentCUDAStream()>>>(
-            query.data_ptr<scalar_t>(),
-            key.data_ptr<scalar_t>(),
+            querys.data_ptr<scalar_t>(),
+            keys.data_ptr<scalar_t>(),
             positions.data_ptr<int>(),
             cos_sin.data_ptr<scalar_t>(),
             head_dim,

src/layers/attention/handler.h

@@ -53,6 +53,7 @@ class AttentionHandler {
       const torch::Tensor& key,    // [n_tokens, n_kv_heads, head_dim]
       const torch::Tensor& value,  // [n_tokens, n_kv_heads, head_dim]
       const InputParameters& input_params) = 0;
+
   // create an attention handler
   static std::unique_ptr<AttentionHandler> create_handler(
       const ModelArgs& args,

src/memory/kv_cache_test.cpp

@@ -57,11 +57,10 @@ TEST(KVCacheTest, Basic) {
 }
 
 TEST(KVCacheTest, Random) {
-  const int num_kv_heads = 12;
-  const int head_dim = 128;
-  const int block_size = 4;
-  const int x = 8;
-  const int num_blocks = 2;
+  const int64_t num_kv_heads = 12;
+  const int64_t head_dim = 128;
+  const int64_t block_size = 4;
+  const int64_t num_blocks = 2;
 
   // auto dtype = torch::kFloat16;
   torch::set_default_dtype(
@@ -82,17 +81,21 @@ TEST(KVCacheTest, Random) {
   for (int32_t i = 0; i < 10000; ++i) {
     using ISlice = torch::indexing::Slice;
 
-    const int sample_size = std::min(num_blocks * block_size, 10);
-    const int num_slots = i % sample_size + 1;
+    const int64_t sample_size = std::min<int64_t>(num_blocks * block_size, 10);
+    const int64_t num_slots = i % sample_size + 1;
     torch::Tensor slot_ids =
         torch::randperm(num_blocks * block_size,
                         torch::dtype(torch::kInt).device(device))
             .index({ISlice(0, num_slots)});
 
+    // construct keys and values with different strides
     torch::Tensor keys =
-        torch::rand({num_slots, num_kv_heads, head_dim}, torch::device(device));
+        torch::rand({num_slots, num_kv_heads * 2, head_dim},
+                    torch::device(device))
+            .slice(/*dim=*/1, /*start=*/0, /*end=*/num_kv_heads);
     torch::Tensor values =
         torch::rand({num_slots, num_kv_heads, head_dim}, torch::device(device));
+    EXPECT_NE(keys.stride(0), values.stride(0));
 
     kv_cache.set_kv_cache_cuda(slot_ids, keys, values);