Skip to content

Implement start_pos per query for batch interface #344

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 3 commits into from
Closed

Implement start_pos per query for batch interface #344

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
208917a
Implement `start_pos` per query for batch interface
ufownl Aug 10, 2024
477ffc1
Fix build issues when tests are enabled
ufownl Aug 10, 2024
5c98189
Fix the position calculation issue in the generation phase
ufownl Aug 11, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions BUILD.bazel
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -58,10 +58,10 @@ cc_test(
)

cc_test(
name = "matvec_test",
name = "gemma_matvec_test",
size = "small",
timeout = "long",
srcs = ["ops/matvec_test.cc"],
srcs = ["ops/gemma_matvec_test.cc"],
local_defines = ["HWY_IS_TEST"],
# for test_suite.
tags = ["hwy_ops_test"],
Expand Down
2 changes: 1 addition & 1 deletion CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -154,7 +154,7 @@ set(GEMMA_TEST_FILES
backprop/optimize_test.cc
ops/ops_test.cc
ops/matmul_test.cc
ops/matvec_test.cc
ops/gemma_matvec_test.cc
evals/gemma_test.cc
)

Expand Down
3 changes: 2 additions & 1 deletion evals/benchmark_helper.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -171,7 +171,8 @@ std::vector<std::pair<std::string, size_t>> GemmaEnv::BatchQueryModel2(
gcpp::TimingInfo timing_info = {.verbosity = app_.verbosity};
runtime_config_.batch_stream_token = batch_stream_token;
inference_args_.CopyTo(runtime_config_);
model_->GenerateBatch(runtime_config_, prompts, /*start_pos=*/0,
model_->GenerateBatch(runtime_config_, prompts,
std::vector<size_t>(num_queries, 0),
KVCaches(&kv_caches_[0], num_queries), timing_info);
return res;
}
Expand Down
136 changes: 85 additions & 51 deletions gemma/gemma-inl.h
View file
Open in desktop

Large diffs are not rendered by default.

13 changes: 7 additions & 6 deletions gemma/gemma.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -66,7 +66,8 @@ Gemma::~Gemma() {
TimingInfo& timing_info); \
extern void GenerateBatch(CONFIGT<TWEIGHT>, const ByteStorageT& weights_u8, \
const RuntimeConfig& runtime_config, \
const MultiplePromptsTokens& prompts, size_t pos, \
const MultiplePromptsTokens& prompts, \
const MultiplePositions& pos, \
const KVCaches& kv_caches, PerClusterPools& pools, \
TimingInfo& timing_info);
GEMMA_FOREACH_CONFIG_AND_WEIGHT(GEMMA_DECLARE);
Expand All @@ -87,9 +88,9 @@ template <class TConfig>
struct GenerateBatchT {
void operator()(const ByteStorageT& weights_u8,
const RuntimeConfig& runtime_config,
const MultiplePromptsTokens& prompts, size_t pos,
const KVCaches& kv_caches, PerClusterPools& pools,
TimingInfo& timing_info) const {
const MultiplePromptsTokens& prompts,
const MultiplePositions& pos, const KVCaches& kv_caches,
PerClusterPools& pools, TimingInfo& timing_info) const {
GenerateBatch(TConfig(), weights_u8, runtime_config, prompts, pos,
kv_caches, pools, timing_info);
}
Expand All @@ -109,8 +110,8 @@ void Gemma::Generate(const RuntimeConfig& runtime_config,

void Gemma::GenerateBatch(const RuntimeConfig& runtime_config,
const MultiplePromptsTokens& prompts,
size_t start_pos, const KVCaches& kv_caches,
TimingInfo& timing_info) {
const MultiplePositions& start_pos,
const KVCaches& kv_caches, TimingInfo& timing_info) {
pools_.StartSpinning();

CallForModelAndWeight<GenerateBatchT>(info_.model, info_.weight, weights_u8_,
Expand Down
4 changes: 3 additions & 1 deletion gemma/gemma.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -143,6 +143,7 @@ struct TimingInfo {

using PromptTokens = hwy::Span<const int>;
using MultiplePromptsTokens = hwy::Span<const PromptTokens>;
using MultiplePositions = hwy::Span<const size_t>;
using KVCaches = hwy::Span<KVCache>;

class Gemma {
Expand All @@ -164,7 +165,8 @@ class Gemma {
size_t start_pos, KVCache& kv_cache, TimingInfo& timing_info);

void GenerateBatch(const RuntimeConfig& runtime_config,
const MultiplePromptsTokens& prompts, size_t start_pos,
const MultiplePromptsTokens& prompts,
const MultiplePositions& start_pos,
const KVCaches& kv_caches, TimingInfo& timing_info);

private:
Expand Down
2 changes: 1 addition & 1 deletion ops/matvec_test.cc → ops/gemma_matvec_test.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -33,7 +33,7 @@

// clang-format off
#undef HWY_TARGET_INCLUDE
#define HWY_TARGET_INCLUDE "ops/matvec_test.cc" // NOLINT
#define HWY_TARGET_INCLUDE "ops/gemma_matvec_test.cc" // NOLINT
// clang-format on
#include "hwy/foreach_target.h" // IWYU pragma: keep
#include "hwy/highway.h"
Expand Down