 void inclusive_scan_wrapper(
     const int* length_data,
     int len_length,
-    Tensor<CUDAContext>* prefix_sum_buffer,
-    Tensor<CUDAContext>* prefix_sum_length_buffer,
+    Tensor<CUDAContext>* temp_buffer,
+    Tensor<CUDAContext>* prefix_sum_out,
     CUDAContext* context_) {
-  // inclusive scan
+  // Retrieve buffer size
   size_t temp_storage_bytes = 0;
-  // Retrieve buffer size
   cub::DeviceScan::InclusiveSum(
       NULL,
       temp_storage_bytes,
       length_data,
-      prefix_sum_length_buffer->mutable_data<int>(),
+      prefix_sum_out->mutable_data<int>(),
       len_length,
       context_->cuda_stream());
   // Allocate temporary storage
-  auto buffer_size = temp_storage_bytes / sizeof(int);
-  buffer_size += temp_storage_bytes % sizeof(int) != 0 ? 1 : 0;
-  prefix_sum_buffer->Resize(buffer_size);
-  void* d_temp_storage =
-      static_cast<void*>(prefix_sum_buffer->mutable_data<int>());
+  auto buffer_size = (temp_storage_bytes + sizeof(int)) / sizeof(int);
+  temp_buffer->Resize(buffer_size);
+  void* d_temp_storage = static_cast<void*>(temp_buffer->mutable_data<int>());
   // Run inclusive prefix sum
   cub::DeviceScan::InclusiveSum(
       d_temp_storage,
       temp_storage_bytes,
       length_data,
-      prefix_sum_length_buffer->mutable_data<int>(),
+      prefix_sum_out->mutable_data<int>(),
       len_length,
       context_->cuda_stream());
 }
 
-template <typename T>
+template <typename T, bool ExactBlock = false>
 __global__ void length_sum_kernel(
-    const T* in,
-    T* out,
-    const int* prefix_sum_length_data,
+    const T* __restrict__ in,
+    T* __restrict__ out,
+    const int* __restrict__ prefix_sum_length_data,
     int N,
     int post,
     int len_length) {

  CUDA_KERNEL_ASSERT(start <= N);
  CUDA_KERNEL_ASSERT(end <= N);

  for (int i = threadIdx.x; i < post; i += blockDim.x) {
    T sum = (T)0;
    for (int line = start; line < end; ++line) {
      sum += in[line * post + i];
    }
    out[group * post + i] = sum;
  }
}

template <typename T>
__global__ void length_sum_gradient_kernel(
    const T* grad_in,
    T* grad_out,
    const int* prefix_sum_length_data,
    int N,
    int post,
    int len_length) {

                len_length,
                dataToReduceSize);
      } else {
        sparse_length_sum_kernel<T, IndexType, false>
            <<<len_length, maxThreads, 0, context_.cuda_stream()>>>(
                in_data,
                out_data,
                prefix_sum_length_data,
