EngineX-Hygon/sglang
5
0
Fork 0
Code Issues Pull Requests Actions 7 Projects Releases Wiki Activity

Unify forward mode (#1360)

Browse Source
This commit is contained in:
Liangsheng Yin
2024-09-09 13:49:29 -07:00
committed by GitHub
parent 689ff588ec
commit 69b3bb9ae1
9 changed files with 54 additions and 58 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
python/sglang/srt/layers/radix_attention.py
View File

@@ -197,9 +197,9 @@ class RadixAttention(nn.Module):
k = k.view(-1, self.tp_k_head_num, self.qk_head_dim)
v = v.view(-1, self.tp_v_head_num, self.v_head_dim)
if input_metadata.forward_mode == ForwardMode.EXTEND:
if input_metadata.forward_mode.is_extend():
return self.extend_forward(q, k, v, input_metadata)
elif input_metadata.forward_mode == ForwardMode.DECODE:
elif input_metadata.forward_mode.is_decode():
return self.decode_forward(q, k, v, input_metadata)
def store_kv_cache(self, cache_k, cache_v, input_metadata: InputMetadata):