   feature_buffer_ = dh::ToSpan(d_feature_buffer_storage_);
 
   // --- Initialize result buffers.
-  output_buffer_bits_storage_.resize(n_features);
+  output_buffer_bits_storage_.resize(BitField::ComputeStorageSize(n_features));
   output_buffer_bits_ = BitField(dh::ToSpan(output_buffer_bits_storage_));
-  input_buffer_bits_storage_.resize(n_features);
+  input_buffer_bits_storage_.resize(BitField::ComputeStorageSize(n_features));
   input_buffer_bits_ = BitField(dh::ToSpan(input_buffer_bits_storage_));
   result_buffer_.resize(n_features);
   s_result_buffer_ = dh::ToSpan(result_buffer_);

  return {s_result_buffer_.data(), s_result_buffer_.data() + n_available};
}

__global__ void QueryFeatureListKernel(common::Span<int32_t> feature_list_input,
                                       common::Span<int32_t> node_feature_list,
                                       BitField result_buffer_input,
                                       BitField result_buffer_output) {
  uint32_t tid = threadIdx.x + blockIdx.x * blockDim.x;
  if (tid < feature_list_input.size()) {
    result_buffer_input.Set(feature_list_input[tid]);
  }

  if (tid < node_feature_list.size()) {
    result_buffer_output.Set(node_feature_list[tid]);
  }
  result_buffer_output &= result_buffer_input;
}


}

__global__ void ClearBuffersKernel(
    BitField result_buffer_output, BitField result_buffer_input) {
  auto tid = blockIdx.x * blockDim.x + threadIdx.x;
  if (tid < result_buffer_output.Size()) {
    result_buffer_output.Clear(tid);
  }
  if (tid < result_buffer_input.Size()) {
    result_buffer_input.Clear(tid);
