batch_util.h

Go to the documentation of this file.

// SPDX-License-Identifier: Apache-2.0
#ifndef ENTROPIC_INFERENCE_BATCH_UTIL_H
#define ENTROPIC_INFERENCE_BATCH_UTIL_H
 
#include <algorithm>
#include <cstddef>
#include <vector>
 
#include "warm_keep_util.h"  // common_prefix_len
 
namespace entropic {
 
template <typename Tok>
inline std::size_t batch_shared_prefix_len(
    const std::vector<std::vector<Tok>>& seqs) {
    if (seqs.size() < 2) { return 0; }
    std::size_t shortest = seqs[0].size();
    std::size_t shared = seqs[0].size();
    for (std::size_t i = 1; i < seqs.size(); ++i) {
        shortest = std::min(shortest, seqs[i].size());
        shared = std::min(shared, common_prefix_len(seqs[0], seqs[i]));
    }
    if (shortest == 0) { return 0; }
    return std::min(shared, shortest - 1);
}
 
inline bool batch_is_viable(std::size_t n, int n_parallel, std::size_t shared,
                            bool hybrid, std::size_t total_suffix,
                            int n_batch) {
    return n >= 2
        && !hybrid
        && n_parallel >= 1
        && n <= static_cast<std::size_t>(n_parallel)
        && shared > 0
        && total_suffix <= static_cast<std::size_t>(n_batch);
}
 
}  // namespace entropic
 
#endif  // ENTROPIC_INFERENCE_BATCH_UTIL_H