chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Support MiniCPM3 (#1371)

Browse Source
This commit is contained in:
William
2024-09-10 17:57:52 +08:00
committed by GitHub
parent fec2d1223c
commit e72275cf7f
5 changed files with 683 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
python/sglang/srt/layers/decode_attention.py
View File

@@ -483,11 +483,14 @@ def _decode_grouped_att_m_fwd(
# shape constraints
Lq, Lk = q.shape[-1], k_buffer.shape[-1]
assert Lq == Lk
assert Lk in {16, 32, 64, 96, 128, 256, 576}
assert Lk in {16, 32, 64, 96, 128, 256, 576, 288}
if Lk == 576:
BLOCK_DMODEL = 512
BLOCK_DPE = 64
elif Lk == 288:
BLOCK_DMODEL = 256
BLOCK_DPE = 32
else:
BLOCK_DMODEL = triton.next_power_of_2(Lk)
BLOCK_DPE = 0

5
python/sglang/srt/layers/extend_attention.py
View File

@@ -280,12 +280,15 @@ def extend_attention_fwd(
assert Lq == Lk and Lv == Lo
# TODO: is the assertion necessary?
assert Lq in {16, 32, 64, 96, 128, 256, 576}
assert Lq in {16, 32, 64, 96, 128, 256, 576, 288}
assert Lv in {16, 32, 64, 96, 128, 256, 512}
if Lq == 576:
BLOCK_DMODEL = 512
BLOCK_DPE = 64
elif Lq == 288:
BLOCK_DMODEL = 256
BLOCK_DPE = 32
else:
BLOCK_DMODEL = triton.next_power_of_2(Lq)
BLOCK_DPE = 0