[WIP]Chunked Prefill #188
Open
[WIP]Chunked Prefill #188
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
vipannalla
reviewed
Feb 18, 2025
| @@ -518,6 +518,7 @@ async def send_request( | |||
| """Send the request to JetStream server.""" | |||
| # Tokenize on client side following MLPerf standard. | |||
| token_ids = tokenizer.encode(input_request.prompt) | |||
| print("len token_ids ", len(token_ids)) | |||
There was a problem hiding this comment.
delete? or use log.debug()
Comment on lines
+38
to
+39
| # 64, | ||
| # 128, |
There was a problem hiding this comment.
We should jit compile 128 bucket size
| ) | ||
| else: | ||
| jax.debug.print("calling chunked_prefill for {chunk_num}", chunk_num=chunk_num) | ||
| prefill_result, first_token = prefill_engine.prefill(params=prefill_params | {"cache": prefill_result["cache"]}, |
There was a problem hiding this comment.
Is "cache" supposed to represent KV cache from previous chunks so far? Can we rename it to "cache_so_far"?
Comment on lines
+589
to
+606
| if prefill_result is None: | ||
| jax.debug.print("calling chunked_prefill for {chunk_num}", chunk_num=chunk_num) | ||
| prefill_result, first_token = prefill_engine.prefill(params=prefill_params, | ||
| padded_tokens=padded_tokens[chunk_num], | ||
| true_length=true_lengths[chunk_num], | ||
| positions=positions[chunk_num], | ||
| all_true_length=true_length, | ||
| previous_chunk=prefill_result, | ||
| ) | ||
| else: | ||
| jax.debug.print("calling chunked_prefill for {chunk_num}", chunk_num=chunk_num) | ||
| prefill_result, first_token = prefill_engine.prefill(params=prefill_params | {"cache": prefill_result["cache"]}, | ||
| padded_tokens=padded_tokens[chunk_num], | ||
| true_length=true_lengths[chunk_num], | ||
| positions=positions[chunk_num], | ||
| all_true_length=true_length, | ||
| previous_chunk=prefill_result, | ||
| ) |
There was a problem hiding this comment.
You can get rid of forking to make the code more readable:
cache_so_far = {} if prefill_result is None else {"cache_so_far": prefill_result["cache"]}
prefill_result, first_token = prefill_engine.prefill(params=prefill_params | cache_so_far, ....)
...
| if total token size is 520 and chunk size is 256, | ||
| the function will return 3 chunks and return tuple is as follows- | ||
| [[t0,..t255][t256,..t511][t512,..t519...(padding)]], | ||
| [256, 256, 7+padding], |
There was a problem hiding this comment.
nit: the true lengths returned should be [256, 256, 7] (no padding)
| the function will return 3 chunks and return tuple is as follows- | ||
| [[t0,..t255][t256,..t511][t512,..t519...(padding)]], | ||
| [256, 256, 7+padding], | ||
| [[0,..255],[256,..511],[512..518..]] |
There was a problem hiding this comment.
nit: [512..519..] (ends at 519)
vipannalla
reviewed
Feb 18, 2025
Comment on lines
+607
to
+608
| t_l_array = jnp.expand_dims(jnp.arange(0, chunk_num*prefill_engine.chunk_size + true_lengths[chunk_num]), 0) | ||
| prefill_result['t_l_array'] = t_l_array |
There was a problem hiding this comment.
Isn't t_l_array same as positions array? where is this used?
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.
If chunked prefill is set to True,
Chunk and appropriately pad the tokens
Call prefill for each chunk