     int N = output_0dim;
     int post = segmentGradsInput.size_from_dim(1);
 
-    if (post > 128) {
-      length_sum_gradient_kernel<T>
-          <<<len_length, 512, 0, context_.cuda_stream()>>>(
-              in_data, out_data, prefix_sum_length_data, N, post, len_length);
-    } else if (post > 64) {
-      length_sum_gradient_kernel<T>
-          <<<len_length, 128, 0, context_.cuda_stream()>>>(
-              in_data, out_data, prefix_sum_length_data, N, post, len_length);
-    } else if (post > 32) {
-      length_sum_gradient_kernel<T>
-          <<<len_length, 64, 0, context_.cuda_stream()>>>(
+    auto maxThreads =
+        GetDeviceProperty(CaffeCudaGetDevice()).maxThreadsPerBlock;
+
+    if (post <= maxThreads) {
+      int multiple = std::min(maxThreads / post, 16);
+      dim3 block(post, multiple);
+
+      length_sum_gradient_kernel<T, true>
+          <<<len_length, block, 0, context_.cuda_stream()>>>(
+
               in_data, out_data, prefix_sum_length_data, N, post, len_length);
     } else {
-      length_sum_gradient_kernel<T>
-          <<<len_length, 32, 0, context_.cuda_stream()>>>(
+      length_sum_gradient_kernel<T, false>
+          <<<len_length, maxThreads, 0, context_.cuda_stream()>>>(
               in_data, out_data, prefix_sum_length_data, N, post, 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) {

void inclusive_scan_wrapper(
    const int* length_data,
    int len_length,
    Tensor<CUDAContext>* temp_buffer,
    Tensor<CUDAContext>* prefix_sum_out,
    CUDAContext* context_) {
  // Retrieve buffer size
  size_t temp_storage_bytes = 0;
  cub::DeviceScan::InclusiveSum(
      NULL,
      temp_storage_bytes,
      length_data,
      prefix_sum_out->mutable_data<int>(),
      len_length,
      context_->cuda_stream());
  // Allocate temporary storage
  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_out->mutable_data<int>(),
      len_length,
      context_->cuda_stream());
}

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