Skip to content

GitLab

Unverified Commit ebe4d1db authored by Liang's avatar Liang Committed by GitHub
Browse files

Fix boundary check in paged attention kernel (#1241)

parent b5a10eb0
Show whitespace changes
Inline Side-by-side
Showing with 1 addition and 1 deletion
csrc/attention/attention_kernels.cu
View file @ ebe4d1db
...@@ -269,7 +269,7 @@ __global__ void single_query_cached_kv_attention_kernel( ...@@ -269,7 +269,7 @@ __global__ void single_query_cached_kv_attention_kernel(
// See https://github.com/vllm-project/vllm/issues/641#issuecomment-1682544472 // See https://github.com/vllm-project/vllm/issues/641#issuecomment-1682544472
scalar_t* v_vec_ptr = reinterpret_cast<scalar_t*>(&v_vec); scalar_t* v_vec_ptr = reinterpret_cast<scalar_t*>(&v_vec);
#pragma unroll #pragma unroll
for (int j = 0; j <= V_VEC_SIZE; j++) { for (int j = 0; j < V_VEC_SIZE; j++) {
v_vec_ptr[j] = token_idx + j < context_len ? v_vec_ptr[j] : zero_value; v_vec_ptr[j] = token_idx + j < context_len ? v_vec_ptr[j] : zero_value;
} }
} }
......
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