Add PagedAttention support (experimental, CUDA only) #17579
Conversation
ericcurtin
force-pushed
the
add-pagedattention
branch
3 times, most recently
from
2a33486 to
14ad291
Compare
github-actions
bot
added
Nvidia GPU
ericcurtin
force-pushed
the
add-pagedattention
branch
from
14ad291 to
06254d1
Compare
Implement PagedAttention algorithm from for memory-efficient KV cache management. This feature reduces memory fragmentation by storing KV cache in fixed-size blocks (similar to virtual memory paging) and enables efficient memory sharing between sequences through copy-on-write semantics. The implementation is experimental and disabled by default. Enable with the --pagedattention flag Signed-off-by: Eric Curtin <eric.curtin@docker.com>
ericcurtin
force-pushed
the
add-pagedattention
branch
from
06254d1 to
1745418
Compare
| const int token_idx = block_idx * BLOCK_SIZE + i; | ||
| if (token_idx >= seq_len) break; | ||
|
|
||
| // TODO: Vectorized K loading and Q·K computation |
There was a problem hiding this comment.
some TODOs look quite sus, I'm wondering if the code is AI-generated and/or this function actually works
beside, probably give some credits to the original kernel: https://github.com/vllm-project/vllm/blob/main/csrc/attention/attention_kernels.cuh
There was a problem hiding this comment.
I mark it experimental for good reason 🙂
There was a problem hiding this comment.
I think it's important to explicitly state if you're using AI to generate this PR or not. the numerous TODOs though out the PR does make it look sus. there will be a human who spend real time and efforts reviewing this PR afterall.
There was a problem hiding this comment.
I mark it experimental for good reason 🙂
I think this PR should be marked as a draft, until it is no longer experimental
3 participants
Implement PagedAttention algorithm for memory-efficient KV cache management. This feature reduces memory fragmentation by storing KV cache in fixed-size blocks (similar to virtual memory paging) and enables efficient memory sharing between sequences through copy-on-write semantics.
The implementation is experimental and disabled by default. Enable with the --pagedattention flag