[Feature] Support variable length merge states (#71)

Also, reorder the input dimension of `MergeStates` to be consistent with
variable length version.

Author: yzh119

include/flashinfer/cascade.cuh

@@ -137,8 +137,8 @@ __device__ __forceinline__ void threadblock_sync_state(state_t<vec_size>& st, DT
  * \param vec_size The vector size used in the kernel.
  * \tparam DTypeIn The data type of v.
  * \tparam DTypeOut The data type of v_merged.
- * \param v The partial v of index sets. (num_index_sets, n, h, d)
- * \param s The logsumexp value of index sets. (num_index_sets, n, h)
+ * \param v The partial v of index sets. (n, num_index_sets, h, d)
+ * \param s The logsumexp value of index sets. (n, num_index_sets, h)
  * \param v_merged The merged v of index sets union. (n, h, d)
  * \param s_merged The merged logsumexp value of index sets union. (n, h)
  * \param num_heads The number of heads of v.
@@ -150,7 +150,6 @@ __global__ void MergeStatesKernel(DTypeIn* __restrict__ V, float* __restrict__ S
                                   DTypeOut* __restrict__ v_merged, float* __restrict__ s_merged,
                                   uint32_t num_index_sets, uint32_t num_heads, uint32_t head_dim) {
   uint32_t tx = threadIdx.x, ty = threadIdx.y;
-  uint32_t seq_len = gridDim.x;
   uint32_t pos = blockIdx.x;
   uint32_t head_idx = ty;
   state_t<vec_size> st;
@@ -159,9 +158,10 @@ __global__ void MergeStatesKernel(DTypeIn* __restrict__ V, float* __restrict__ S
   v_merged_vec.fill(0.f);
 #pragma unroll 2
   for (uint32_t iter = 0; iter < num_index_sets; ++iter) {
-    float s = S[(iter * seq_len + pos) * num_heads + head_idx];
+    float s = S[(pos * num_index_sets + iter) * num_heads + head_idx];
     vec_t<float, vec_size> v;
-    v.cast_load(V + ((iter * seq_len + pos) * num_heads + head_idx) * head_dim + tx * vec_size);
+    v.cast_load(V + ((pos * num_index_sets + iter) * num_heads + head_idx) * head_dim +
+                tx * vec_size);
     st.merge(v, s, 1);
   }
 
@@ -175,14 +175,14 @@ __global__ void MergeStatesKernel(DTypeIn* __restrict__ V, float* __restrict__ S
 /*!
  * \brief The CUDA kernel that merges self-attention states of a list of index sets,
  *   accelerated for larget number of index sets.
- * \param vec_size The vector size used in the kernel.
- * \param bdx The blockDim.x used in the kernel.
- * \param bdy The blockDim.y used in the kernel.
- * \param num_smem_stages The number of stages of shared memory used in the kernel.
+ * \tparam vec_size The vector size used in the kernel.
+ * \tparam bdx The blockDim.x used in the kernel.
+ * \tparam bdy The blockDim.y used in the kernel.
+ * \tparam num_smem_stages The number of stages of shared memory used in the kernel.
  * \tparam DTypeIn The data type of v.
  * \tparam DTypeOut The data type of v_merged.
- * \param v The partial v of index sets. (num_index_sets, n, h, d)
- * \param s The logsumexp value of index sets. (num_index_sets, n, h)
+ * \param V The partial v of index sets. (n, num_index_sets, h, d)
+ * \param S The logsumexp value of index sets. (n, num_index_sets, h)
  * \param v_merged The merged v of index sets union. (n, h, d)
  * \param s_merged The merged logsumexp value of index sets union. (n, h)
  * \param num_heads The number of heads of v.
@@ -196,9 +196,8 @@ __global__ void MergeStatesLargeNumIndexSetsKernel(DTypeIn* __restrict__ V, floa
                                                    float* __restrict__ s_merged,
                                                    uint32_t num_index_sets, uint32_t num_heads) {
   uint32_t tx = threadIdx.x, ty = threadIdx.y;
-  uint32_t seq_len = gridDim.y;
-  uint32_t pos = blockIdx.y;
-  uint32_t head_idx = blockIdx.x;
+  uint32_t pos = blockIdx.x;
+  uint32_t head_idx = blockIdx.y;
   state_t<vec_size> st;
   constexpr uint32_t vec_bits = sizeof(DTypeIn) * vec_size * 8;
   constexpr uint32_t head_dim = vec_size * bdx;
@@ -211,7 +210,8 @@ __global__ void MergeStatesLargeNumIndexSetsKernel(DTypeIn* __restrict__ V, floa
   for (uint32_t iter = 0; iter < num_smem_stages; ++iter) {
     cp_async::pred_load<vec_bits, PrefetchMode::kPrefetch, SharedMemFillMode::kNoFill>(
         v_smem + (iter * bdy + ty) * head_dim + tx * vec_size,
-        V + (((iter * bdy + ty) * seq_len + pos) * num_heads + head_idx) * head_dim + tx * vec_size,
+        V + ((pos * num_index_sets + (iter * bdy + ty)) * num_heads + head_idx) * head_dim +
+            tx * vec_size,
         (iter * bdy + ty) < num_index_sets);
     cp_async::commit_group();
   }
@@ -220,27 +220,111 @@ __global__ void MergeStatesLargeNumIndexSetsKernel(DTypeIn* __restrict__ V, floa
     if (iter % bdx == 0) {
       s_smem[ty * bdx + tx] =
           iter * bdy + (ty * bdx + tx) < num_index_sets
-              ? S[((iter * bdy + ty * bdx + tx) * seq_len + pos) * num_heads + head_idx]
+              ? S[(pos * num_index_sets + (iter * bdy + ty * bdx + tx)) * num_heads + head_idx]
               : 0.f;
       __syncthreads();
     }
     cp_async::wait_group<num_smem_stages - 1>();
     __syncthreads();
     vec_t<float, vec_size> v;
     v.cast_load(v_smem + ((iter % num_smem_stages) * bdy + ty) * head_dim + tx * vec_size);
+    if (iter * bdy + ty < num_index_sets) {
+      float s = s_smem[(iter % bdx) * bdy + ty];
+      st.merge(v, s, 1);
+    }
     __syncthreads();
     cp_async::pred_load<vec_bits, PrefetchMode::kPrefetch, SharedMemFillMode::kNoFill>(
         v_smem + ((iter % num_smem_stages) * bdy + ty) * head_dim + tx * vec_size,
         V +
-            ((((iter + num_smem_stages) * bdy + ty) * seq_len + pos) * num_heads + head_idx) *
+            ((pos * num_index_sets + ((iter + num_smem_stages) * bdy + ty)) * num_heads +
+             head_idx) *
                 head_dim +
             tx * vec_size,
         (iter + num_smem_stages) * bdy + ty < num_index_sets);
     cp_async::commit_group();
+  }
+  cp_async::wait_group<0>();
+  __syncthreads();
+
+  st.normalize();
+  threadblock_sync_state<bdx, bdy, vec_size>(st, v_smem, s_smem);
+  st.normalize();
+
+  st.o.cast_store(v_merged + (pos * num_heads + head_idx) * head_dim + tx * vec_size);
+  if (s_merged != nullptr) {
+    s_merged[pos * num_heads + head_idx] = st.get_lse();
+  }
+}
+
+/*!
+ * \brief The CUDA kernel to merge self-attention states of multiple index sets, the number of index
+ *   sets at each position might vary.
+ * \tparam vec_size The vector size used in the kernel.
+ * \tparam bdx The blockDim.x used in the kernel.
+ * \tparam bdy The blockDim.y used in the kernel.
+ * \tparam num_smem_stages The number of stages of shared memory used in the kernel.
+ * \tparam DTypeIn The data type of v.
+ * \tparam DTypeOut The data type of v_merged.
+ * \param V The partial v of index sets. (nnz, h, d)
+ * \param S The logsumexp value of index sets. (nnz, h)
+ * \param indptr The start offsets of each position in the variable length array.
+ * \param v_merged The merged v of index sets union. (n, h, d)
+ * \param s_merged The merged logsumexp value of index sets union. (n, h)
+ * \param num_heads The number of heads of v.
+ * \param head_dim The dimension of each head.
+ * \note s are logsumexp values with base 2.
+ */
+template <uint32_t vec_size, uint32_t bdx, uint32_t bdy, uint32_t num_smem_stages, typename DTypeIn,
+          typename DTypeOut, typename IdType>
+__global__ void VariableLengthMergeStatesKernel(DTypeIn* __restrict__ V, float* __restrict__ S,
+                                                IdType* indptr, DTypeOut* __restrict__ v_merged,
+                                                float* __restrict__ s_merged, uint32_t num_heads) {
+  uint32_t tx = threadIdx.x, ty = threadIdx.y;
+  uint32_t pos = blockIdx.x;
+  uint32_t head_idx = blockIdx.y;
+  state_t<vec_size> st;
+  constexpr uint32_t vec_bits = sizeof(DTypeIn) * vec_size * 8;
+  constexpr uint32_t head_dim = vec_size * bdx;
+
+  extern __shared__ uint8_t smem[];
+  DTypeIn* v_smem = (DTypeIn*)smem;
+  float* s_smem = (float*)(smem + num_smem_stages * bdy * head_dim * sizeof(DTypeIn));
+  const uint32_t num_index_sets = indptr[pos + 1] - indptr[pos];
+
+#pragma unroll
+  for (uint32_t iter = 0; iter < num_smem_stages; ++iter) {
+    cp_async::pred_load<vec_bits, PrefetchMode::kPrefetch, SharedMemFillMode::kNoFill>(
+        v_smem + (iter * bdy + ty) * head_dim + tx * vec_size,
+        V + ((indptr[pos] + (iter * bdy + ty)) * num_heads + head_idx) * head_dim + tx * vec_size,
+        (iter * bdy + ty) < num_index_sets);
+    cp_async::commit_group();
+  }
+#pragma unroll 4
+  for (uint32_t iter = 0; iter < ceil_div(num_index_sets, bdy); ++iter) {
+    if (iter % bdx == 0) {
+      s_smem[ty * bdx + tx] =
+          iter * bdy + (ty * bdx + tx) < num_index_sets
+              ? S[(indptr[pos] + (iter * bdy + ty * bdx + tx)) * num_heads + head_idx]
+              : 0.f;
+      __syncthreads();
+    }
+    cp_async::wait_group<num_smem_stages - 1>();
+    __syncthreads();
+    vec_t<float, vec_size> v;
+    v.cast_load(v_smem + ((iter % num_smem_stages) * bdy + ty) * head_dim + tx * vec_size);
     if (iter * bdy + ty < num_index_sets) {
       float s = s_smem[(iter % bdx) * bdy + ty];
       st.merge(v, s, 1);
     }
+    __syncthreads();
+    cp_async::pred_load<vec_bits, PrefetchMode::kPrefetch, SharedMemFillMode::kNoFill>(
+        v_smem + ((iter % num_smem_stages) * bdy + ty) * head_dim + tx * vec_size,
+        V +
+            ((indptr[pos] + ((iter + num_smem_stages) * bdy + ty)) * num_heads + head_idx) *
+                head_dim +
+            tx * vec_size,
+        (iter + num_smem_stages) * bdy + ty < num_index_sets);
+    cp_async::commit_group();
   }
   cp_async::wait_group<0>();
   __syncthreads();
@@ -346,7 +430,7 @@ cudaError_t MergeStates(DTypeIn* v, float* s, DTypeOut* v_merged, float* s_merge
     if (num_index_sets > 2 * (128 / bdx)) {
       constexpr uint32_t num_threads = 128;
       constexpr uint32_t bdy = num_threads / bdx;
-      dim3 nblks(num_heads, seq_len);
+      dim3 nblks(seq_len, num_heads);
       dim3 nthrs(bdx, bdy);
       constexpr uint32_t num_smem_stages = 4;
       auto kernel = MergeStatesLargeNumIndexSetsKernel<vec_size, bdx, bdy, num_smem_stages, DTypeIn,
@@ -369,6 +453,30 @@ cudaError_t MergeStates(DTypeIn* v, float* s, DTypeOut* v_merged, float* s_merge
   return cudaSuccess;
 }
 
+template <typename DTypeIn, typename DTypeOut, typename IdType>
+cudaError_t VariableLengthMergeStates(DTypeIn* v, float* s, IdType* indptr, DTypeOut* v_merged,
+                                      float* s_merged, uint32_t seq_len, uint32_t num_heads,
+                                      uint32_t head_dim, cudaStream_t stream = nullptr) {
+  SWITCH_HEAD_DIM(head_dim, HEAD_DIM, {
+    constexpr uint32_t vec_size = std::max(16U / sizeof(DTypeIn), HEAD_DIM / 32U);
+    constexpr uint32_t bdx = HEAD_DIM / vec_size;
+    constexpr uint32_t num_threads = 128;
+    constexpr uint32_t bdy = num_threads / bdx;
+    dim3 nblks(seq_len, num_heads);
+    dim3 nthrs(bdx, bdy);
+    constexpr uint32_t num_smem_stages = 4;
+    auto kernel = VariableLengthMergeStatesKernel<vec_size, bdx, bdy, num_smem_stages, DTypeIn,
+                                                  DTypeOut, IdType>;
+    void* args[] = {&v, &s, &indptr, &v_merged, &s_merged, &num_heads};
+    uint32_t smem_size =
+        num_smem_stages * bdy * head_dim * sizeof(DTypeIn) + num_threads * sizeof(float);
+    FLASHINFER_CUDA_CALL(
+        cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size));
+    FLASHINFER_CUDA_CALL(cudaLaunchKernel((void*)kernel, nblks, nthrs, args, smem_size, stream));
+  });
+  return cudaSuccess;
+}
+
 }  // namespace flashinfer
 
 #endif  // FLASHINFER_CASCADE_CUH_

include/flashinfer/prefill.cuh

@@ -803,8 +803,8 @@ __global__ void SinglePrefillWithKVCacheKernel(
   DTypeIn* q_ptr_base = q + qkv_info.get_qo_elem_offset(qo_idx_base, kv_head_idx * group_size,
                                                         (tx % 8) * num_elems_per_128b<DTypeIn>());
   DTypeOut* o_ptr_base =
-      split_kv ? ((DTypeOut*)tmp) + chunk_idx * qo_len * qkv_info.get_num_qo_heads() * head_dim +
-                     qkv_info.get_qo_elem_offset(qo_idx_base, kv_head_idx * group_size,
+      split_kv ? ((DTypeOut*)tmp) + chunk_idx * qkv_info.get_num_qo_heads() * head_dim +
+                     qkv_info.get_qo_elem_offset(qo_idx_base * num_chunks, kv_head_idx * group_size,
                                                  (tx % 8) * num_elems_per_128b<DTypeOut>())
                : o + qkv_info.get_qo_elem_offset(qo_idx_base, kv_head_idx * group_size,
                                                  (tx % 8) * num_elems_per_128b<DTypeOut>());
@@ -910,8 +910,9 @@ __global__ void SinglePrefillWithKVCacheKernel(
   normalize_d<num_frags_x, num_frags_y>(o_frag, d);
 
   // write back
-  write_o_reg_gmem<group_size, num_frags_x, num_frags_y>(o_frag, &qo_smem, o_ptr_base, qo_idx_base,
-                                                         qo_len, qo_n_stride, qo_h_stride);
+  write_o_reg_gmem<group_size, num_frags_x, num_frags_y>(
+      o_frag, &qo_smem, o_ptr_base, qo_idx_base, qo_len,
+      split_kv ? qo_n_stride * num_chunks : qo_n_stride, qo_h_stride);
 
   // write lse
   if (lse != nullptr || split_kv) {
@@ -926,8 +927,8 @@ __global__ void SinglePrefillWithKVCacheKernel(
         if (qo_idx < qo_len) {
           if constexpr (split_kv) {
             float* tmp_lse =
-                (float*)(((DTypeOut*)tmp) + num_chunks * qo_len * num_qo_heads * head_dim);
-            tmp_lse[(chunk_idx * qo_len + qo_idx) * num_qo_heads + qo_head_idx] =
+                (float*)(((DTypeOut*)tmp) + qo_len * num_chunks * num_qo_heads * head_dim);
+            tmp_lse[(qo_idx * num_chunks + chunk_idx) * num_qo_heads + qo_head_idx] =
                 math::ptx_log2(d[fx][j]) + float(m[fx][j]);
           } else {
             lse[qo_idx * num_qo_heads + qo_head_idx] = math::ptx_log2(d[fx][j]) + float(m[fx][j]);

python/csrc/cascade.cu

@@ -65,8 +65,8 @@ std::vector<torch::Tensor> merge_states(torch::Tensor v, torch::Tensor s) {
   CHECK_EQ(v.size(0), s.size(0));
   CHECK_EQ(v.size(1), s.size(1));
   CHECK_EQ(v.size(2), s.size(2));
-  unsigned int num_index_sets = v.size(0);
-  unsigned int seq_len = v.size(1);
+  unsigned int seq_len = v.size(0);
+  unsigned int num_index_sets = v.size(1);
   unsigned int num_heads = v.size(2);
   unsigned int head_dim = v.size(3);
   s = s.to(torch::kFloat32);

python/flashinfer/ops/__init__.py

@@ -265,10 +265,10 @@ def merge_states(v: torch.Tensor, s: torch.Tensor):
     ----------
     v : torch.Tensor
         The attention output from the KV segments.
-        Shape: [num_kv_segments, seq_len, num_heads, head_dim]
+        Shape: [seq_len, num_kv_segments, num_heads, head_dim]
     s : torch.Tensor
         The logsumexp value from the KV segments.
-        Shape: [num_kv_segments, seq_len, num_heads]
+        Shape: [seq_len, num_kv_segments, num_heads]
 
     Returns
     -------

src/bench_cascade.cu

@@ -28,20 +28,20 @@ using namespace flashinfer;
 template <typename T>
 void bench_merge_states(nvbench::state& state) {
   const auto num_index_sets = state.get_int64("num_index_sets");
-  const auto batch_size = state.get_int64("batch_size");
+  const auto seq_len = state.get_int64("seq_len");
   const auto num_heads = state.get_int64("num_heads");
   const auto head_dim = state.get_int64("head_dim");
 
-  std::vector<T> V_host(num_index_sets * batch_size * num_heads * head_dim);
-  std::vector<float> S_host(num_index_sets * batch_size * num_heads);
+  std::vector<T> V_host(seq_len * num_index_sets * num_heads * head_dim);
+  std::vector<float> S_host(seq_len * num_index_sets * num_heads);
 
   utils::vec_normal_(V_host);
   utils::vec_uniform_(S_host, 5, 10);
 
   thrust::device_vector<T> V_device(V_host);
   thrust::device_vector<float> S_device(S_host);
-  thrust::device_vector<T> V_merged(batch_size * num_heads * head_dim);
-  thrust::device_vector<float> S_merged(batch_size * num_heads);
+  thrust::device_vector<T> V_merged(seq_len * num_heads * head_dim);
+  thrust::device_vector<float> S_merged(seq_len * num_heads);
 
   state.add_global_memory_reads<T>(V_host.size(), "Read");
   state.add_global_memory_writes<T>(V_merged.size(), "Write");
@@ -51,7 +51,7 @@ void bench_merge_states(nvbench::state& state) {
     cudaError_t status = MergeStates(
         thrust::raw_pointer_cast(V_device.data()), thrust::raw_pointer_cast(S_device.data()),
         thrust::raw_pointer_cast(V_merged.data()), thrust::raw_pointer_cast(S_merged.data()),
-        num_index_sets, batch_size, num_heads, head_dim);
+        num_index_sets, seq_len, num_heads, head_dim);
     timer.stop();
   });
 }