Integrates ragged attention to JetStream Pytorch #93
Conversation
| mask, | ||
| start, | ||
| input_pos, | ||
| pre_batch, |
There was a problem hiding this comment.
Few comments on what is pre_batch and pre_block?
Also should gemma/model_exportable.py also be modified?
There was a problem hiding this comment.
Thanks for pointing out this! Pushed the commits to update the Gemma model.
| slot, | ||
| ) | ||
|
|
||
| def precompute_ragged_block_indices(self, decode_state: DecodeState): |
There was a problem hiding this comment.
Nit: do you need other attribute from decode_state besides start and input_pos, pass start and input_pos instead of heavy object decode_state.
There was a problem hiding this comment.
Only the start and input_pos is used, but considering we just passing the reference to the decode_state, so it should not affect the performance at all.
| end_ref, | ||
| line_end_ref, | ||
| pre_b_ref, | ||
| pre_i_ref, |
There was a problem hiding this comment.
q, k, v and related are easy to read. What are the b, i, o, m, l, bk and pre? Can you add brief description to describe them?
| slot, | ||
| ) | ||
|
|
||
| def precompute_ragged_block_indices(self, decode_state: DecodeState): |
…al performance. Fix the typo that should use jax.lax.div instead of jnp.div
…ocessing API from JetStream.
dense_attention_quantized and use option to control if it's quantization or not. Use the new torch_xla2 API.
…a ring buffer. Will cause error.
* refactor flags * clean up: * fix run_server * move common flags to global * format * update * udpate readme * update run_interactive
…flags for debugging and performance tuning.
… align with main.
wang2yn84
force-pushed
the
ragged-attention-final2
branch
from
cbb2fe9 to
ab38726
Compare
…nts. The error message is missing positional arguments.
…nput_pos) back to original to avoid unnecessary issues.
…nel. Fix other lint errors.
Currently the performance is on par with dense attention. We can keep improving the performance in the following PRs.