Recompute KV cache for Phi3 when switching from short to long factor #1161
Conversation
src/generators.cpp
Outdated
| // 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 (model_->config_->model.type == "phi3" && search_->params_->search.batch_size == 1 && search_->params_->search.num_beams == 1) { |
There was a problem hiding this comment.
The model types should be phi3, phi3small, phi3v, and phimoe. For phi3v, rewinding is not currently supported.
src/generators.cpp
Outdated
| // 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 (model_->config_->model.type == "phi3" && search_->params_->search.batch_size == 1 && search_->params_->search.num_beams == 1) { | ||
| if (search_->GetSequenceLength() == 4097 && first_switch) { |
There was a problem hiding this comment.
For phi3small, the switch value is at 8K and not 4K.
There was a problem hiding this comment.
Could we record original_max_position_embeddings in the config and avoid hardcoding this value here?
src/generators.cpp
Outdated
| if (model_->config_->model.type == "phi3" && search_->params_->search.batch_size == 1 && search_->params_->search.num_beams == 1) { | ||
| if (search_->GetSequenceLength() == 4097 && first_switch) { | ||
| first_switch = false; | ||
| auto current_seq = cpu_span<int32_t>(GetSequence(0).CpuSpan()); |
There was a problem hiding this comment.
Do you intend to copy the span over to CPU? You might want to use CopyDeviceToCpu instead of CpuSpan.
| RewindToLength(0); | ||
| AppendTokens(current_seq); |
There was a problem hiding this comment.
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 (search_->GetSequenceLength() == 4097 && first_switch) { | ||
| first_switch = false; | ||
| auto current_seq = cpu_span<int32_t>(GetSequence(0).CpuSpan()); | ||
| RewindToLength(0); |
There was a problem hiding this comment.
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.
src/generators.cpp
Outdated
| // 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 (model_->config_->model.type == "phi3" && search_->params_->search.batch_size == 1 && search_->params_->search.num_beams == 1) { | ||
| if (search_->GetSequenceLength() == 4097 && first_switch) { |
There was a problem hiding this comment.
Could we record original_max_position_embeddings in the config and avoid hardcoding this value here?
Recompute KV cache for Phi3 when switching from short to long factor.
Verified that this PR fixes the issue for: