Rotary Embedding NKI Kernel Tutorial #53
Conversation
specify indicies for other dims when storing tiles
There was a problem hiding this comment.
We cannot merge this at the moment. The source of the tutorial is now hosted in an internal location and is displayed at https://awsdocs-neuron.readthedocs-hosted.com/en/latest/general/nki/tutorials.html. All of the existing tutorial has a RTD file that is in a certain format. We plan to migrate the tutorial source code to this repo soon and then we can add the RST file for this tutorial as well.
aws-qieqingy
dismissed
their stale review
Cannot merge for now due to reason stated in a separate comment.
| """ | ||
| # Rotate Q | ||
| output_tile[0, :, :] = q_tile * cos_tile | ||
| output_tile[0, :, : q_tile.shape[-1] // 2] = output_tile[ |
There was a problem hiding this comment.
do we need to assert q_tile.shape[-1] % 2 == 0
also could we assign q_tile.shape[-1] to a temp variable such that we dont need to type q_tile.shape[-1] repeatedly?
| i_p = nl.arange(128)[:, None] | ||
| i_f = nl.arange(q.shape[-1])[None, :] |
There was a problem hiding this comment.
sorry forgot to mention. We recommend to use mgrid:
i_p, i_f = nl.mgrid[0:128, 0:q.shape[-1]]
There was a problem hiding this comment.
no worries, just made the update
Issue #, if available:
N/A
Description of changes:
NKI kernel tutorial for applying rotary embedding to query, key tensors. Included benchmarking, and correctness tests with script
Testing:
Please see detailed unit test requirements in the CONTRIBUTING.md
nki.baremetalnki.benchmarkPull Request Checklist