           in_data, out_data, scan_size, outer_size, inner_size);
     }
   }
+  int64_t grid_size = outer_size * inner_size;
   if (!transpose && !reverse) {
-    BlockScanKernel<T, 128, 4, Op><<<scan_grid, 128, 0, dev_ctx.stream()>>>(
+    BlockScanKernel<T, 128, 4, Op><<<grid_size, 128, 0, dev_ctx.stream()>>>(
         out_data, in_data, outer_size, inner_size, scan_size, exclusive, op);
 
   } else {
     BlockScanKernel<T, 128, 4, Op>
-        <<<scan_grid, 128, 0, dev_ctx.stream()>>>(next_out_data,
+        <<<grid_size, 128, 0, dev_ctx.stream()>>>(next_out_data,
                                                   next_in_data,
                                                   outer_size,
                                                   inner_size,

  int outer_size = height / scan_size;
  int inner_size = width;
  // Consider the size of shared memory, here block size is 128
  dim3 scan_grid(outer_size, inner_size);
  dim3 reverse_grid = scan_grid;
  if (reverse) {

      valid_item = scan_size;
    }

    int offset = block_offset + bx * scan_size;

    T thread_keys[ITEMS_PER_THREAD];
    BlockLoadT(temp_storage.load)
