Commit 3f946cfb authored by zhuwenwen's avatar zhuwenwen
Browse files

fix alibi

parent 0890ebc1
...@@ -290,7 +290,7 @@ __device__ void paged_attention_with_mask_kernel_TC( ...@@ -290,7 +290,7 @@ __device__ void paged_attention_with_mask_kernel_TC(
const int token_idx = block_idx * BLOCK_SIZE+rowid; const int token_idx = block_idx * BLOCK_SIZE+rowid;
if(alibi_slope[i] != 0){ if(alibi_slope[i] != 0){
float alibi=alibi_slope[i]* (token_idx - seq_len + 1); float alibi=alibi_slope[i]* (token_idx - seq_len + 1);
qk_vec[i] = alibi; qk_vec[i] += alibi;
} }
// used for tree-style attention // used for tree-style attention
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment