vrr/kvcache-ai-ktransformers
2
0
Fork 0
mirror of https://github.com/kvcache-ai/ktransformers.git synced 2025-09-09 13:55:27 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Fix ktransformers-server flashinfer wrapper position arg issue;

Browse source 
Fix db position issue
This commit is contained in:
Azure-Tang 2025-04-01 07:30:23 +00:00
parent 203b853c75
commit 31677181c3
3 changed files with 7 additions and 2 deletions
Download patch file Download diff file Expand all files Collapse all files

5
ktransformers/operators/flashinfer_wrapper.py
View file

@ -128,6 +128,9 @@ class MLAWrapper():
if kv_indices is None:
assert self.max_batch_size == 1
kv_indices = self.kv_indices_buf
if bsz_tensor is None:
assert self.max_batch_size == 1
bsz_tensor = self.batch_size_tensor_buf
self.wrapper.plan(
qo_indptr,
@ -166,6 +169,7 @@ class MLAWrapperSingleton():
kv_indptr,
kv_indices,
kv_len_arr,
bsz_tensor,
num_heads,
head_dim_ckv,
head_dim_kpe,
@ -179,6 +183,7 @@ class MLAWrapperSingleton():
kv_indptr,
kv_indices,
kv_len_arr_cur_device,
bsz_tensor,
num_heads,
head_dim_ckv,
head_dim_kpe,