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-hopper-flashinfer

Browse source
This commit is contained in:
Atream 2025-04-29 11:06:34 +08:00
parent 38333cf129
commit b0318fc01c
3 changed files with 6 additions and 3 deletions
Download patch file Download diff file Expand all files Collapse all files

3
ktransformers/models/custom_modeling_deepseek_v2.py
View file

@ -50,7 +50,8 @@ class KDeepseekV2ForCausalLM(DeepseekV2PreTrainedModel):
self.wrapper = flashinfer.mla.BatchMLAPagedAttentionWrapper(
self.workspace_buffer, use_cuda_graph=use_cuda_graph,
qo_indptr=self.qo_indptr_buf,kv_indptr=self.paged_kv_indptr_buf,
kv_indices=self.paged_kv_indices_buf,kv_len_arr=self.paged_kv_len_buf
kv_indices=self.paged_kv_indices_buf,kv_len_arr=self.paged_kv_len_buf,
backend = "fa2",
)
def batch_embeddings(self, batch: ForwardBatchInput, device="cuda:0"):