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

fix-update-flashinfer_wrapper_local_chat

Browse source
This commit is contained in:
Atream 2025-02-25 12:47:31 +00:00
parent 5474be5299
commit 477ac28a9c
4 changed files with 15 additions and 4 deletions
Download patch file Download diff file Expand all files Collapse all files

1
ktransformers/util/utils.py
View file

@ -177,6 +177,7 @@ def prefill_and_generate(model, tokenizer, inputs, max_new_tokens=10000, use_cud
else:
inputs_embeds = model.model.embed_tokens(inputs.to("cpu")).to(torch_device)
if use_flashinfer_mla:
MLAWrapperSingleton.update_buffer(past_key_values.max_pages)
MLAWrapperSingleton.need_plan_all()
logits = model(