Commits on Feb 1, 2024

1. Support RoPE position info in batch prefill/decode kernels …

   This PR adds q/k position information to batch prefill/decode
   kernels. More specifically, the kernel now accepts two
   additional arrays:
   * `q_rope_position` with shape `(total_q_len,)`, denoting the
   in-sequence position of each position in the input q.
   * `k_rope_pos_offset` with shape `(num_sequence,)`, denoting
   the start position of each sequence in k.
   
   These two arrays helps on-the-fly calculate RoPE in multi-level
   cases.
   
   Tests `test_batch_prefill` and `test_batch_decode` can pass.
   Performance is not validated yet. Per discussion with Zihao,
   this change is not very likely to incur significant perf
   regression.

   

   MasterJH5574 committed Feb 1, 2024
   Configuration menu

   • View commit details
   • Copy full SHA for d89fe51
   • Browse repository at this point

   Copy the full SHA

   d89fe51 View commit details

   Browse the repository at this point in the history