maxtext
maxtext copied to clipboard
Ragged attention and test, not actively used.
Adds ragged attention kernel in Pallas as well as a unit test for the new code. Note that the ragged attention kernel is not actively being used in this code and will be incorporated into the rest of the attention code in a future PR.