   }
 }
 
+/*
+ * Continuation is called after the partitioning stage. It launches kernels
+ * to sort large and small segments using the partitioning results. Separation
+ * of this stage is required to eliminate device-side synchronization in
+ * the CDP mode.
+ */
+template <typename LargeSegmentPolicyT,
+          typename SmallAndMediumPolicyT,
+          typename LargeKernelT,
+          typename SmallKernelT,
+          typename KeyT,
+          typename ValueT,
+          typename BeginOffsetIteratorT,
+          typename EndOffsetIteratorT>
+__host__ __device__ cudaError_t
+DeviceSegmentedSortContinuation(
+    LargeKernelT large_kernel,
+    SmallKernelT small_kernel,
+    int num_segments,
+    KeyT *d_current_keys,
+    KeyT *d_final_keys,
+    detail::device_double_buffer<KeyT> d_keys_double_buffer,
+    ValueT *d_current_values,
+    ValueT *d_final_values,
+    detail::device_double_buffer<ValueT> d_values_double_buffer,
+    BeginOffsetIteratorT d_begin_offsets,
+    EndOffsetIteratorT d_end_offsets,
+    unsigned int *group_sizes,
+    unsigned int *large_and_medium_segments_indices,
+    unsigned int *small_segments_indices,
+    cudaStream_t stream,
+    bool debug_synchronous)
+{
+  cudaError error = cudaSuccess;
+
+  const unsigned int large_segments = group_sizes[0];
+
+  if (large_segments > 0)
+  {
+    // One CTA per segment
+    const unsigned int blocks_in_grid = large_segments;
+
+    if (debug_synchronous)
+    {
+      _CubLog("Invoking "
+              "DeviceSegmentedSortKernelLarge<<<%d, %d, 0, %lld>>>()\n",
+              static_cast<int>(blocks_in_grid),
+              LargeSegmentPolicyT::BLOCK_THREADS,
+              (long long)stream);
+    }
+
+    THRUST_NS_QUALIFIER::cuda_cub::launcher::triple_chevron(
+      blocks_in_grid,
+      LargeSegmentPolicyT::BLOCK_THREADS,
+      0,
+      stream)
+      .doit(large_kernel,
+            large_and_medium_segments_indices,
+            d_current_keys,
+            d_final_keys,
+            d_keys_double_buffer,
+            d_current_values,
+            d_final_values,
+            d_values_double_buffer,
+            d_begin_offsets,
+            d_end_offsets);
+
+    // Check for failure to launch
+    if (CubDebug(error = cudaPeekAtLastError()))
+    {
+      return error;
+    }
+
+    // Sync the stream if specified to flush runtime errors
+    if (debug_synchronous)
+    {
+      if (CubDebug(error = SyncStream(stream)))
+      {
+        return error;
+      }
+    }
+  }
+
+  const unsigned int small_segments  = group_sizes[1];
+  const unsigned int medium_segments =
+    static_cast<unsigned int>(num_segments) -
+    (large_segments + small_segments);
+
+  const unsigned int small_blocks =
+    DivideAndRoundUp(small_segments,
+                     SmallAndMediumPolicyT::SEGMENTS_PER_SMALL_BLOCK);
+
+  const unsigned int medium_blocks =
+    DivideAndRoundUp(medium_segments,
+                     SmallAndMediumPolicyT::SEGMENTS_PER_MEDIUM_BLOCK);
+
+  const unsigned int small_and_medium_blocks_in_grid = small_blocks +
+                                                       medium_blocks;
+
+  if (small_and_medium_blocks_in_grid)
+  {
+    if (debug_synchronous)
+    {
+      _CubLog("Invoking "
+              "DeviceSegmentedSortKernelSmall<<<%d, %d, 0, %lld>>>()\n",
+              static_cast<int>(small_and_medium_blocks_in_grid),
+              SmallAndMediumPolicyT::BLOCK_THREADS,
+              (long long)stream);
+    }
+
+    THRUST_NS_QUALIFIER::cuda_cub::launcher::triple_chevron(
+      small_and_medium_blocks_in_grid,
+      SmallAndMediumPolicyT::BLOCK_THREADS,
+      0,
+      stream)
+      .doit(small_kernel,
+            small_segments,
+            medium_segments,
+            medium_blocks,
+            small_segments_indices,
+            large_and_medium_segments_indices + num_segments - medium_segments,
+            d_current_keys,
+            d_final_keys,
+            d_current_values,
+            d_final_values,
+            d_begin_offsets,
+            d_end_offsets);
+
+    // Check for failure to launch
+    if (CubDebug(error = cudaPeekAtLastError()))
+    {
+      return error;
+    }
+
+    // Sync the stream if specified to flush runtime errors
+    if (debug_synchronous)
+    {
+      if (CubDebug(error = SyncStream(stream)))
+      {
+        return error;
+      }
+    }
+  }
+
+  return error;
+}
+
+/*
+ * Continuation kernel is used only in the CDP mode. It's used to
+ * launch DeviceSegmentedSortContinuation as a separate kernel.
+ */
+template <typename ChainedPolicyT,
+          typename LargeKernelT,
+          typename SmallKernelT,
+          typename KeyT,
+          typename ValueT,
+          typename BeginOffsetIteratorT,
+          typename EndOffsetIteratorT>
+__launch_bounds__(1) __global__ void
+DeviceSegmentedSortContinuationKernel(
+  LargeKernelT large_kernel,
+  SmallKernelT small_kernel,
+  int num_segments,
+  KeyT *d_current_keys,
+  KeyT *d_final_keys,
+  detail::device_double_buffer<KeyT> d_keys_double_buffer,
+  ValueT *d_current_values,
+  ValueT *d_final_values,
+  detail::device_double_buffer<ValueT> d_values_double_buffer,
+  BeginOffsetIteratorT d_begin_offsets,
+  EndOffsetIteratorT d_end_offsets,
+  unsigned int *group_sizes,
+  unsigned int *large_and_medium_segments_indices,
+  unsigned int *small_segments_indices,
+  bool debug_synchronous)
+{
+  using ActivePolicyT = typename ChainedPolicyT::ActivePolicy;
+  using LargeSegmentPolicyT = typename ActivePolicyT::LargeSegmentPolicy;
+  using SmallAndMediumPolicyT =
+    typename ActivePolicyT::SmallAndMediumSegmentedSortPolicyT;
+
+  // In case of CDP:
+  // 1. each CTA has a different main stream
+  // 2. all streams are non-blocking
+  // 3. child grid always completes before the parent grid
+  // 4. streams can be used only from the CTA in which they were created
+  // 5. streams created on the host cannot be used on the device
+  //
+  // Due to (4, 5), we can't pass the user-provided stream in the continuation.
+  // Due to (1, 2, 3) it's safe to pass the main stream.
+  cudaError_t error =
+    DeviceSegmentedSortContinuation<LargeSegmentPolicyT, SmallAndMediumPolicyT>(
+      large_kernel,
+      small_kernel,
+      num_segments,
+      d_current_keys,
+      d_final_keys,
+      d_keys_double_buffer,
+      d_current_values,
+      d_final_values,
+      d_values_double_buffer,
+      d_begin_offsets,
+      d_end_offsets,
+      group_sizes,
+      large_and_medium_segments_indices,
+      small_segments_indices,
+      0, // always launching on the main stream (see motivation above)
+      debug_synchronous);
+
+  CubDebug(error);
+}
 
 template <typename KeyT,
           typename ValueT>

          return error;
        }
      }
    }

    const unsigned int small_segments  = h_group_sizes[1];
    const unsigned int medium_segments =
      static_cast<unsigned int>(num_segments) -
      (large_segments + small_segments);

    const unsigned int small_blocks =
      DivideAndRoundUp(small_segments,
                       SmallAndMediumPolicyT::SEGMENTS_PER_SMALL_BLOCK);

    const unsigned int medium_blocks =
      DivideAndRoundUp(medium_segments,
                       SmallAndMediumPolicyT::SEGMENTS_PER_MEDIUM_BLOCK);

    const unsigned int small_and_medium_blocks_in_grid = small_blocks +
                                                         medium_blocks;

    if (small_and_medium_blocks_in_grid)
    {
      if (debug_synchronous)
      {
        _CubLog("Invoking "
                "DeviceSegmentedSortKernelSmall<<<%d, %d, 0, %lld>>>()\n",
                static_cast<int>(small_and_medium_blocks_in_grid),
                SmallAndMediumPolicyT::BLOCK_THREADS,
                (long long)stream);
      }

      THRUST_NS_QUALIFIER::cuda_cub::launcher::triple_chevron(
        small_and_medium_blocks_in_grid,
        SmallAndMediumPolicyT::BLOCK_THREADS,
        0,
        stream)
        .doit(small_kernel,
              small_segments,
              medium_segments,
              medium_blocks,
              small_segments_indices.get(),
              large_and_medium_segments_indices.get() + num_segments -
                medium_segments,
              d_keys.Current(),
              GetFinalOutput(LargeSegmentPolicyT::RADIX_BITS, d_keys),
              d_values.Current(),
              GetFinalOutput(LargeSegmentPolicyT::RADIX_BITS, d_values),
              d_begin_offsets,
              d_end_offsets);

      // Check for failure to launch
      if (CubDebug(error = cudaPeekAtLastError()))
      {
        return error;
      }

      // Sync the stream if specified to flush runtime errors
      if (debug_synchronous)
      {
        if (CubDebug(error = SyncStream(stream)))
        {
          return error;
        }
      }
    }

    return error;

      return error;
    }

    if (CUB_IS_HOST_CODE)
    {
      #if CUB_INCLUDE_HOST_CODE
      unsigned int h_group_sizes[num_selected_groups];

      if (CubDebug(
            error = cudaMemcpyAsync(h_group_sizes,
                                    group_sizes.get(),
                                    num_selected_groups * sizeof(unsigned int),
                                    cudaMemcpyDeviceToHost,
                                    stream)))
      {
        return error;
      }
