Skip to content
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

Recompute KV cache for Phi3 when switching from short to long factor #1161

Open
wants to merge 4 commits into
base: main
Choose a base branch
from
Open

Recompute KV cache for Phi3 when switching from short to long factor #1161

Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
6833733
recompute KV cache for Phi3 when switching from short to long factor
ajindal1 Dec 20, 2024
d9622de
fix typo
ajindal1 Dec 20, 2024
24f0e11
remove first switch
ajindal1 Dec 20, 2024
f96ea94
add different Phi3 models
ajindal1 Jan 3, 2025
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
17 changes: 17 additions & 0 deletions src/generators.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -352,6 +352,23 @@ void Generator::GenerateNextToken() {
ThrowErrorIfSessionTerminated(state_->session_terminated_);
if (search_->GetSequenceLength() == 0 && !computed_logits_)
throw std::runtime_error("GenerateNextToken called with no prior state. Please call AppendTokens, SetLogits, or params.SetInputs before calling GenerateNextToken.");

// TODO: Extend the solution to make it work for batch size > 1 and num beams > 1
// Phi3 model switches from short factor to long factor at 4097 (original_max_position_embeddings+1) token, needs Recomputation of Position IDs and KV Cache
// at this stage which is achieved by rewinding to zero and appending the current sequence
if (search_->params_->search.batch_size == 1 && search_->params_->search.num_beams == 1) {
if ((search_->GetSequenceLength() == 4097) && (model_->config_->model.type == "phi3" || model_->config_->model.type == "phimoe")) {
auto current_seq = cpu_span<int32_t>(GetSequence(0).CpuSpan());
RewindToLength(0);
AppendTokens(current_seq);
}
if ((search_->GetSequenceLength() == 8197) && (model_->config_->model.type == "phi3small")) {
auto current_seq = cpu_span<int32_t>(GetSequence(0).CpuSpan());
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do you intend to copy the span over to CPU? You might want to use CopyDeviceToCpu instead of CpuSpan.

RewindToLength(0);
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do you think a more general purpose solution might be to have a new API that resets the generator?

Reset()

Then we don't need to limit ourselves to only supporting scenarios that support RewindToLength(0) (i.e. continuous decoding).

Maybe we can do it in another PR.

AppendTokens(current_seq);
Comment on lines +367 to +368
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could you also add a comment highlighting that this solution works only for scenarios that support continuous decoding.

  • Batch size == 1
  • Num beams == 1
  • Model must be a decoder model (not a multimodal model)
  • EP in {CPU, CUDA}. DML won't work.

}
}

if (!computed_logits_) {
auto next_tokens = search_->GetNextTokens();
if (last_action_ == Action::rewound)
Expand Down
Loading