Memcpy kernel for flash attention

[suquark]

Memcpy kernel for flash attention

num_tokens: 64, num_heads: 40, head_size: 128, block_size: 8, num_blocks: 1024, dtype: torch.float16
[Latency] gather_cached_kv: 0.008 ms
[Throughput] gather_cached_kv: 156.479 GB/s
num_tokens: 128, num_heads: 40, head_size: 128, block_size: 8, num_blocks: 1024, dtype: torch.float16
[Latency] gather_cached_kv: 0.011 ms
[Throughput] gather_cached_kv: 216.171 GB/s
num_tokens: 256, num_heads: 40, head_size: 128, block_size: 8, num_blocks: 1024, dtype: torch.float16
[Latency] gather_cached_kv: 0.032 ms
[Throughput] gather_cached_kv: 152.631 GB/s
num_tokens: 512, num_heads: 40, head_size: 128, block_size: 8, num_blocks: 1024, dtype: torch.float16
[Latency] gather_cached_kv: 0.057 ms
[Throughput] gather_cached_kv: 172.325 GB/s
num_tokens: 1024, num_heads: 40, head_size: 128, block_size: 8, num_blocks: 1024, dtype: torch.float16
[Latency] gather_cached_kv: 0.104 ms
[Throughput] gather_cached_kv: 187.537 GB/s
num_tokens: 2048, num_heads: 40, head_size: 128, block_size: 8, num_blocks: 1024, dtype: torch.float16
[Latency] gather_cached_kv: 0.204 ms
[Throughput] gather_cached_kv: 191.603 GB/s

The performance is pretty good (theoretical optimal throughput is 1.6TB/s for A100-40GB), considering the memory layout is not ideal.

result for unoptimized kernel:

num_tokens: 64, num_heads: 40, head_size: 128, block_size: 8, num_blocks: 1024, dtype: torch.float16
[Latency] gather_cached_kv: 0.010 ms
[Throughput] gather_cached_kv: 125.891 GB/s
num_tokens: 128, num_heads: 40, head_size: 128, block_size: 8, num_blocks: 1024, dtype: torch.float16
[Latency] gather_cached_kv: 0.015 ms
[Throughput] gather_cached_kv: 160.678 GB/s
num_tokens: 256, num_heads: 40, head_size: 128, block_size: 8, num_blocks: 1024, dtype: torch.float16
[Latency] gather_cached_kv: 0.032 ms
[Throughput] gather_cached_kv: 150.732 GB/s
num_tokens: 512, num_heads: 40, head_size: 128, block_size: 8, num_blocks: 1024, dtype: torch.float16
[Latency] gather_cached_kv: 0.060 ms
[Throughput] gather_cached_kv: 162.482 GB/s
num_tokens: 1024, num_heads: 40, head_size: 128, block_size: 8, num_blocks: 1024, dtype: torch.float16
[Latency] gather_cached_kv: 0.108 ms
[Throughput] gather_cached_kv: 180.763 GB/s
num_tokens: 2048, num_heads: 40, head_size: 128, block_size: 8, num_blocks: 1024, dtype: torch.float16
[Latency] gather_cached_kv: 0.206 ms
[Throughput] gather_cached_kv: 189.757 GB/s

the optimized kernel works much better for smaller number of tokens (+20% speedup)

[suquark]

implementation is done. need testing (will do it on Thursday)

the memory saving strategy is orthogonal to this kernel, so I would not include it in this PR

[WoosukKwon]

Hey @suquark thanks for the PR! I have a quick question: have you also measured the performance diff between the two kernels before and after the optimization?

[suquark]

see the PR comment for the optimized kernel performance comparison

[WoosukKwon]

LGTM.