RYDE-WORK/ktransformers
1
0
Fork 0
mirror of https://github.com/RYDE-WORK/ktransformers.git synced 2026-01-19 21:03:18 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Merge pull request #315 from kvcache-ai/Atream-add-adapted

Browse Source 
Atream add adapted
This commit is contained in:
Atream 2025-02-15 15:45:34 +08:00 committed by GitHub
commit f9f9f746c0
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 9 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
ktransformers/operators/attention.py
View File

@ -262,7 +262,7 @@ class KDeepseekV2Attention(BaseInjectedModule, DeepseekV2Attention):
"""
# flash attn doesn't support head_dim bigger than 256
# use vLLM triton attention kernel for MQA
# use triton attention kernel adapted from vLLM and SGLang for MQA
decode_attention_fwd_grouped(query_states, compressed_kv_with_k_pe, compressed_kv, attn_output,
page_table,
position_ids.squeeze(0).to(torch.int32), attn_logits,
@ -551,4 +551,4 @@ class KLlamaAttention(BaseInjectedModule):
if not output_attentions:
attn_weights = None
return attn_output, attn_weights, past_key_value
return attn_output, attn_weights, past_key_value

8
ktransformers/operators/triton_attention.py
View File

@ -1,3 +1,9 @@
# Adapted from
# https://github.com/sgl-project/sglang/blob/9f635ea50de920aa507f486daafba26a5b837574/python/sglang/srt/layers/attention/triton_ops/decode_attention.py
# which was originally adapted from
# https://github.com/ModelTC/lightllm/blob/96353e868a840db4d103138caf15ed9dbea8c186/lightllm/models/deepseek2/triton_kernel/gqa_flash_decoding_stage1.py
# https://github.com/ModelTC/lightllm/blob/96353e868a840db4d103138caf15ed9dbea8c186/lightllm/models/deepseek2/triton_kernel/gqa_flash_decoding_stage2.py
import triton
import triton.language as tl
@ -376,4 +382,4 @@ def decode_attention_fwd_grouped(
)
_decode_softmax_reducev_fwd(attn_logits, q, o, v_buffer, b_seq_len,
num_kv_splits)
num_kv_splits)