[tx] Right align prompts for decoding #932
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request aims to right-align prompts for the decoding phase after a left-aligned prefill. This is a good change that should allow for using the same efficient attention mechanism for both prefill and decoding. The approach of removing the compute_positions function and introducing a batch_roll function to shift the KV cache and attention mask is sound. However, I've identified a critical issue in the implementation of the new batch_roll function that would lead to incorrect behavior for multi-dimensional tensors like the KV cache. My review includes a suggested fix for this issue.
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request introduces a change to right-align prompts for the decoding phase, following a prefill phase with left-aligned prompts. This is a clever approach to enable the use of optimized cuDNN attention kernels for both stages. The changes are well-contained and logical. The new batch_roll utility is correctly implemented for shifting the KV cache and attention mask. The modifications in the model files to simplify position ID generation for the prefill stage are also appropriate given the new alignment strategy. Overall, the implementation appears correct and robust.
…o tx-right-align-decode
In NovaSky-AI#927, we left aligned the prompts for sampling prefill in preparation for NovaSky-AI#879. Here we shift them after the prefill so they will again be right aligned for the decoding. This will allow us to use the cudnn attention NovaSky-AI#927 for both prefill and decoding. --------- Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
1 participant
In #927, we left aligned the prompts for sampling prefill in preparation for #879. Here we shift them after the prefill so they will again be right aligned for the decoding. This will allow us to use the cudnn attention #927 for both prefill and decoding.