   CUDA_KERNEL_ASSERT(start <= N);
   CUDA_KERNEL_ASSERT(end <= N);
 
-  for (int i = threadIdx.x; i < post; i += blockDim.x) {
-    for (int line = start; line < end; ++line) {
-      grad_out[line * post + i] = grad_in[group * post + i];
+  if (ExactBlock) {
+    grad_out += threadIdx.x;
+    grad_in += threadIdx.x;
+
+    for (int line = start + threadIdx.y; line < end; line += blockDim.y) {
+      grad_out[line * post] = grad_in[group * post];
+    }
+  } else {
+    for (int i = threadIdx.x; i < post; i += blockDim.x) {
+      for (int line = start; line < end; ++line) {
+        grad_out[line * post + i] = grad_in[group * post + i];
+      }
     }
   }
 }

void inclusive_scan_wrapper(
    const int* length_data,
    int len_length,
    Tensor<CUDAContext>* prefix_sum_buffer,
    Tensor<CUDAContext>* prefix_sum_length_buffer,
    CUDAContext* context_) {
  // inclusive scan
  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>(),
      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>());
  // Run inclusive prefix sum
  cub::DeviceScan::InclusiveSum(
      d_temp_storage,
      temp_storage_bytes,
      length_data,
      prefix_sum_length_buffer->mutable_data<int>(),
      len_length,
      context_->cuda_stream());
}

template <typename T>
__global__ void length_sum_kernel(
    const T* in,
    T* out,
    const int* prefix_sum_length_data,
    int N,
    int post,
    int len_length) {

  }
}

template <typename T, typename IndexType, bool ExactBlock = false>
__global__ void sparse_length_sum_kernel(
    const T* __restrict__ in,
    T* __restrict__ out,
    const int* __restrict__ prefix_sum_length_data,
    const IndexType* __restrict__ indices,
    int N,
    int post,
    int len_length,
