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

fix fused_set_kv_buffer for rope with Ling-v2 (#20316)

Browse Source
This commit is contained in:
strgrb
2026-03-23 19:20:40 +08:00
committed by GitHub
parent 4641e5a3d2
commit 80d4a0753a
1 changed files with 6 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/models/bailing_moe.py
+6 -2
View File
@@ -532,6 +532,10 @@ class BailingMoEAttention(nn.Module):
head_dim=self.head_dim, head_dim=self.head_dim,
alt_stream=self.alt_stream, alt_stream=self.alt_stream,
) )
can_fuse_set_kv = (
self.head_dim == self.rotary_emb.rotary_dim
and enable_fused_set_kv_buffer(forward_batch)
)
q, k = self.rotary_emb( q, k = self.rotary_emb(
positions, positions,
q, q,
@@ -542,7 +546,7 @@ class BailingMoEAttention(nn.Module):
layer=self.attn, layer=self.attn,
forward_batch=forward_batch, forward_batch=forward_batch,
) )
if enable_fused_set_kv_buffer(forward_batch) if can_fuse_set_kv
else None else None
), ),
) )
@@ -551,7 +555,7 @@ class BailingMoEAttention(nn.Module):
k, k,
v, v,
forward_batch, forward_batch,
save_kv_cache=not enable_fused_set_kv_buffer(forward_batch), save_kv_cache=not can_fuse_set_kv,
) )
attn_output, _ = self.dense(context_layer) attn_output, _ = self.dense(context_layer)
return attn_output return attn_output