feat - support mla kvcache store #888
Conversation
baowendin
changed the title
There was a problem hiding this comment.
Hi @baowendin , thanks for the contribution, the kernel look good to me overall.
Would you mind trying using triton instead? I expect we can get similar performance to cuda. In the future, we hope all elementwise / data-movement kernels can be written in triton, to save maintenance overhead.
baowendin
force-pushed
the
feature/support_mla_store_kvcache
branch
from
db20a81 to
61fb997
Compare
baowendin
force-pushed
the
feature/support_mla_store_kvcache
branch
from
61fb997 to
f7bd89f
Compare
|
hi, I have formatted code with pre-commit, but since I'm not familiar with triton, so this time I can't reformat it with triton, maybe next time ? |
There was a problem hiding this comment.
It would be great to add a benchmark like: https://github.com/flashinfer-ai/flashinfer/blob/main/benchmarks/bench_append_paged_kv_cache.py
| @@ -16,6 +16,7 @@ | |||
| #ifndef FLASHINFER_PAGE_CUH_ | |||
| #define FLASHINFER_PAGE_CUH_ | |||
|
|
|||
| #include <assert.h> | |||
There was a problem hiding this comment.
Would you mind changing them to FLASHINFER_CHECK? (defined in
flashinfer/include/flashinfer/exception.h
Line 41 in 341ae09
assert would only work when you compile the program in debug mode, not release mode.
| import flashinfer | ||
|
|
||
|
|
||
| def test_append_mla_paged_kv_cache(): |
There was a problem hiding this comment.
Can you design some stronger test cases, I apologize that the previous test_page.py is very weak (most of the tests around append page APIs is performed in C++ unittests and in the future we should move them to python).
By "stronger" I mean more cases (nnz/append length), data types, page_size etc.
|
okay, I'll fix these problem and add more test later in this week |
Summary
related to #877
This PR implement MLA cache store，and passed correctness test in the case of ckv_dim=512 and kpe_dim=64，but no further performance test yet. Sincerely hope somebody who familiar with CUDA can help improve performance.