youngkingdom/vllm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

simplify - get rid of tokenshape

Browse Source 
Signed-off-by: Tyler Michael Smith <tyler@neuralmagic.com>
This commit is contained in:
Tyler Michael Smith
2025-02-03 22:42:09 +00:00
parent 230730c34d
commit c33aeecf24
5 changed files with 10 additions and 90 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
csrc/torch_bindings.cpp
View File

@ -470,18 +470,6 @@ TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _cache_ops), cache_ops) {
cache_ops.impl("reshape_and_cache_flash", torch::kCUDA,
&reshape_and_cache_flash);
// Reshape the key and value tensors and cache them.
cache_ops.def(
"reshape_and_cache_flash_full_cuda(Tensor tensorshape,"
" Tensor key, Tensor value,"
" Tensor! key_cache,"
" Tensor! value_cache,"
" Tensor slot_mapping,"
" str kv_cache_dtype,"
" Tensor k_scale, Tensor v_scale) -> ()");
cache_ops.impl("reshape_and_cache_flash_full_cuda", torch::kCUDA,
&reshape_and_cache_flash_full_cuda);
// Concat kv_c and k_pe and cache them.
cache_ops.def(
"concat_and_cache_mla(Tensor kv_c, Tensor k_pe,"