Skip to content

Fix tiled KV cache allocation and support prefix cloning. #1015

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

Open
copybara-service wants to merge 1 commit into
base: dev
Choose a base branch
from
Open

Fix tiled KV cache allocation and support prefix cloning. #1015

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
3 changes: 3 additions & 0 deletions BUILD.bazel
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -530,6 +530,9 @@ cc_library(

cc_library(
name = "ops",
srcs = [
"ops/ops.cc",
],
hdrs = [
"ops/ops.h",
],
Expand Down
17 changes: 13 additions & 4 deletions gemma/gemma.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1188,7 +1188,8 @@ static HWY_NOINLINE void PrefillTBatch(const ModelConfig& config,
max_tbatch_size >= prefill_this_query);

// For each batch of tokens in the query:
for (size_t tbatch_start = 0; tbatch_start < prefill_this_query;
for (size_t tbatch_start = qbatch_1.InitialPos(0);
tbatch_start < prefill_this_query;
tbatch_start += max_tbatch_size) {
const size_t tbatch_size =
HWY_MIN(max_tbatch_size, prefill_this_query - tbatch_start);
Expand Down Expand Up @@ -1528,6 +1529,7 @@ size_t PrefillTBatchOrQBatch(const ModelConfig& config,
MatMulEnv& env, TimingInfo& timing_info) {
size_t max_prompt_size = 0;
bool all_prefix_end_are_zero = true;
bool all_initial_pos_are_zero = true;
size_t total_prefill_tokens = 0; // only for throughput stats.
const size_t seq_len = qbatch.KV(0).SeqLen();
for (size_t qi = 0; qi < qbatch.Size(); ++qi) {
Expand All @@ -1540,9 +1542,13 @@ size_t PrefillTBatchOrQBatch(const ModelConfig& config,

// Prefill stops before size - 1 because the last prompt token is the
// first input token for generation.
total_prefill_tokens += prompt.size() - 1;
const size_t initial_pos = qbatch.InitialPos(qi);
if (prompt.size() - 1 > initial_pos) {
total_prefill_tokens += (prompt.size() - 1 - initial_pos);
}

all_prefix_end_are_zero &= qbatch.PrefixEnd(qi) == 0;
all_initial_pos_are_zero &= initial_pos == 0;

// We use a single divisor, so all sequence lengths must be the same.
HWY_ASSERT(qbatch.KV(qi).SeqLen() == seq_len);
Expand All @@ -1559,7 +1565,8 @@ size_t PrefillTBatchOrQBatch(const ModelConfig& config,

timing_info.prefill_start = hwy::platform::Now();
// Batch over the larger of prompt length, or queries.
if ((qbatch.Size() > max_prompt_size) && all_prefix_end_are_zero) {
if ((qbatch.Size() > max_prompt_size) && all_prefix_end_are_zero &&
all_initial_pos_are_zero) {
activations.SetBatchSize(qbatch.Size()); // required before PrefillQBatch
PrefillQBatch(max_prompt_size, config, runtime_config, weights, activations,
qbatch, env, non_eos);
Expand All @@ -1586,7 +1593,9 @@ void StreamAndUpdateEOSAfterPrefill(const ModelConfig& config,
const RuntimeConfig& runtime_config,
QBatch& qbatch, hwy::BitSet4096<>& non_eos,
size_t qi) {
const size_t last_pos_in_prompt = qbatch.Pos(qi) - qbatch.InitialPos(qi);
const size_t prompt_size = qbatch.Prompt(qi).size();
const size_t last_pos_in_prompt =
prompt_size == 0 ? 0 : std::min(qbatch.Pos(qi), prompt_size - 1);

const size_t pos = qbatch.Pos(qi); // during prefill, pos is still correct.
// In autoregressive mode, we have not prefilled the last token, so do
Expand Down
Loading
Loading