Add prefill/decode from seq lens in BaseCausalLMModel #383
+133
−82
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
We do not have a clearly defined interface for LMs. Decode and prefill have different signature when exporting to IREE. Here is added a new ABC CausalLMModelABC that makes a distinction between the two variants. The BaseCausalLMModel provides a default implementation for the new prefill_from_seq_lens and decode_from_seq_lens methods. The export script export_paged_llm_v1 does too much in its exported functions. It computes the attention mask then. It shards its arguments and unshards its result. This change lets it be a thiner wrapper around the new interface functions. Make paged_llm_v1.TorchGenerator use the new interface methods.
rsuderman
requested changes
Nov 6, 2024
| @@ -27,7 +27,7 @@ | |||
| ################################################################################ | |||
|
|
|||
|
|
|||
| class PagedLlamaModelV1(BaseCausalLMModel): | |||
| class PagedLlamaModelV1(BaseCausalLMModel, CausalLMModelABC): | |||
There was a problem hiding this comment.
Double inheritance almost always a bad idea especially considering these are both LLM interfaces. Can we have an alternative option? Why can't we expand this functionality in the existing BaseCausalLMModel?
There was a problem hiding this comment.
I removed CausalLMModelABC. I added the prefill and decode unimplemented methods to BaseCausalLMModel.
sogartar
changed the title
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
We do not have a clearly defined interface for LMs. Decode and prefill have different signature when exporting to IREE. One of them uses an attention mask the other sequence lenghts.
This change adds to BaseCausalLMModel default implementations for the new prefill_from_seq_lens and decode_from_seq_lens methods.
The export script export_paged_llm_v1 does too much in its exported functions. It computes the attention mask then it shards its arguments and unshards its result. This change lets it be a thinner wrapper around the new functions.
Make paged_llm_v1.TorchGenerator use the new interface methods.