                             common::Span<GradientPair const> gpair,
                             common::Span<const uint32_t> d_ridx,
                             common::Span<GradientSumT> histogram,
-                            HistRounding<GradientSumT> rounding,
-                            bool force_global_memory) {
+                            HistRounding<GradientSumT> rounding, bool force_global_memory) {
   // decide whether to use shared memory
   int device = 0;
   dh::safe_cuda(cudaGetDevice(&device));
   // opt into maximum shared memory for the kernel if necessary
   size_t max_shared_memory = dh::MaxSharedMemoryOptin(device);
 
-  size_t smem_size = sizeof(typename HistRounding<GradientSumT>::SharedSumT) *
-                     feature_groups.max_group_bins;
+  size_t smem_size =
+      sizeof(typename HistRounding<GradientSumT>::SharedSumT) * feature_groups.max_group_bins;
   bool shared = !force_global_memory && smem_size <= max_shared_memory;
   smem_size = shared ? smem_size : 0;
 
+  constexpr int kBlockThreads = 1024;
+  constexpr int kItemsPerThread = 8;
+  constexpr int kItemsPerTile = kBlockThreads * kItemsPerThread;
+
   auto runit = [&](auto kernel) {
     if (shared) {
-      dh::safe_cuda(cudaFuncSetAttribute(
-          kernel, cudaFuncAttributeMaxDynamicSharedMemorySize,
-          max_shared_memory));
+      dh::safe_cuda(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize,
+                                         max_shared_memory));
     }
 
     // determine the launch configuration
-    int min_grid_size;
-    int block_threads = 1024;
-    dh::safe_cuda(cudaOccupancyMaxPotentialBlockSize(
-        &min_grid_size, &block_threads, kernel, smem_size, 0));
-
     int num_groups = feature_groups.NumGroups();
     int n_mps = 0;
-    dh::safe_cuda(
-        cudaDeviceGetAttribute(&n_mps, cudaDevAttrMultiProcessorCount, device));
+    dh::safe_cuda(cudaDeviceGetAttribute(&n_mps, cudaDevAttrMultiProcessorCount, device));
     int n_blocks_per_mp = 0;
-    dh::safe_cuda(cudaOccupancyMaxActiveBlocksPerMultiprocessor(
-        &n_blocks_per_mp, kernel, block_threads, smem_size));
+    dh::safe_cuda(cudaOccupancyMaxActiveBlocksPerMultiprocessor(&n_blocks_per_mp, kernel,
+                                                                kBlockThreads, smem_size));
+    // This gives the number of blocks to keep the device occupied
+    // Use this as the maximum number of blocks
     unsigned grid_size = n_blocks_per_mp * n_mps;
 
-    // TODO(canonizer): This is really a hack, find a better way to distribute
-    // the data among thread blocks. The intention is to generate enough thread
-    // blocks to fill the GPU, but avoid having too many thread blocks, as this
-    // is less efficient when the number of rows is low. At least one thread
-    // block per feature group is required. The number of thread blocks:
-    // - for num_groups <= num_groups_threshold, around  grid_size * num_groups
-    // - for num_groups_threshold <= num_groups <= num_groups_threshold *
-    // grid_size,
-    //     around grid_size * num_groups_threshold
-    // - for num_groups_threshold * grid_size <= num_groups, around num_groups
-    int num_groups_threshold = 4;
-    grid_size = common::DivRoundUp(
-        grid_size, common::DivRoundUp(num_groups, num_groups_threshold));
-
-    using T = typename GradientSumT::ValueT;
+    // Otherwise launch blocks such that each block has a minimum amount of work to do
+    // There are fixed costs to launching each block, e.g. zeroing shared memory
+    // The below amount of minimum work was found by experimentation
+    constexpr int kMinItemsPerBlock = kItemsPerTile;
+    int columns_per_group = common::DivRoundUp(matrix.row_stride, feature_groups.NumGroups());
+    // Average number of matrix elements processed by each group
+    std::size_t items_per_group = d_ridx.size() * columns_per_group;
+
+    // Allocate number of blocks such that each block has about kMinItemsPerBlock work
+    // Up to a maximum where the device is saturated
+    grid_size =
+        min(grid_size,
+            unsigned(common::DivRoundUp(items_per_group, kMinItemsPerBlock)));
+
     dh::LaunchKernel {dim3(grid_size, num_groups),
-          static_cast<uint32_t>(block_threads),
-          smem_size} (kernel, matrix, feature_groups, d_ridx,
-                      histogram.data(), gpair.data(), rounding);
+        static_cast<uint32_t>(kBlockThreads), smem_size}(
+        kernel, matrix, feature_groups, d_ridx, histogram.data(), gpair.data(), rounding);
   };
 
   if (shared) {
-    runit(SharedMemHistKernel<GradientSumT, true>);
+    runit(SharedMemHistKernel<GradientSumT, true, kBlockThreads, kItemsPerThread>);
   } else {
-    runit(SharedMemHistKernel<GradientSumT, false>);
+    runit(SharedMemHistKernel<GradientSumT, false, kBlockThreads, kItemsPerThread>);
   }
 
   dh::safe_cuda(cudaGetLastError());
 }
 
 template void BuildGradientHistogram<GradientPairPrecise>(
-    EllpackDeviceAccessor const& matrix,
-    FeatureGroupsAccessor const& feature_groups,
-    common::Span<GradientPair const> gpair,
-    common::Span<const uint32_t> ridx,
-    common::Span<GradientPairPrecise> histogram,
-    HistRounding<GradientPairPrecise> rounding,
+    EllpackDeviceAccessor const& matrix, FeatureGroupsAccessor const& feature_groups,
+    common::Span<GradientPair const> gpair, common::Span<const uint32_t> ridx,
+    common::Span<GradientPairPrecise> histogram, HistRounding<GradientPairPrecise> rounding,
     bool force_global_memory);
 
 }  // namespace tree

/*!
 * Copyright 2020-2021 by XGBoost Contributors
 */
#include <thrust/reduce.h>
#include <thrust/iterator/transform_iterator.h>
#include <algorithm>
#include <ctgmath>
#include <limits>

#include "xgboost/base.h"
#include "row_partitioner.cuh"

#include "histogram.cuh"

#include "../../data/ellpack_page.cuh"
#include "../../common/device_helpers.cuh"

namespace xgboost {
namespace tree {

}  // anonymous namespace

struct Clip : public thrust::unary_function<GradientPair, Pair> {
  static XGBOOST_DEV_INLINE float Pclip(float v) { return v > 0 ? v : 0; }
  static XGBOOST_DEV_INLINE float Nclip(float v) { return v < 0 ? abs(v) : 0; }

  XGBOOST_DEV_INLINE Pair operator()(GradientPair x) const {
    auto pg = Pclip(x.GetGrad());
