[Inference/Kernel] Optimize paged attention: Refactor key cache layout#5643

Merged

SunflowerAries merged 2 commits intohpcaitech:feature/colossal-inferfrom

SunflowerAries:optimize-paged-attn

Apr 25, 2024

Commits on Apr 24, 2024

optimize flashdecodingattention: refactor code with different key cache layout(from [num_blocks, num_kv_heads, block_size, head_size] to [num_blocks, num_kv_heads, head_size/x, block_size, x])
SunflowerAries
committed
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot]
committed