From 80d4a0753a0c7a20a5d9b8d546196c649a0ae44b Mon Sep 17 00:00:00 2001 From: strgrb Date: Mon, 23 Mar 2026 19:20:40 +0800 Subject: [PATCH] fix fused_set_kv_buffer for rope with Ling-v2 (#20316) --- python/sglang/srt/models/bailing_moe.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/python/sglang/srt/models/bailing_moe.py b/python/sglang/srt/models/bailing_moe.py index 4b9e39b82..faaadc90a 100644 --- a/python/sglang/srt/models/bailing_moe.py +++ b/python/sglang/srt/models/bailing_moe.py @@ -532,6 +532,10 @@ class BailingMoEAttention(nn.Module): head_dim=self.head_dim, 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( positions, q, @@ -542,7 +546,7 @@ class BailingMoEAttention(nn.Module): layer=self.attn, forward_batch=forward_batch, ) - if enable_fused_set_kv_buffer(forward_batch) + if can_fuse_set_kv else None ), ) @@ -551,7 +555,7 @@ class BailingMoEAttention(nn.Module): k, v, 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) return attn_output