From 2f6af1a3defe256179d8b386c835ff2f2c942832 Mon Sep 17 00:00:00 2001 From: Yuhong Guo Date: Fri, 31 Oct 2025 19:32:49 +0800 Subject: [PATCH] Enable bailing_moe to support TP=16 (#12369) --- python/sglang/srt/models/bailing_moe.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/python/sglang/srt/models/bailing_moe.py b/python/sglang/srt/models/bailing_moe.py index a4af1e43e..541eb94ef 100644 --- a/python/sglang/srt/models/bailing_moe.py +++ b/python/sglang/srt/models/bailing_moe.py @@ -420,14 +420,21 @@ class BailingMoEAttention(nn.Module): attn_tp_size = get_attention_tp_size() assert self.total_num_heads % attn_tp_size == 0 - assert self.total_kv_heads % attn_tp_size == 0 + if self.total_kv_heads >= attn_tp_size: + # Number of KV heads is greater than TP size, so we partition + # the KV heads across multiple tensor parallel GPUs. + assert self.total_kv_heads % attn_tp_size == 0 + else: + # Number of KV heads is less than TP size, so we replicate + # the KV heads across multiple tensor parallel GPUs. + assert attn_tp_size % self.total_kv_heads == 0 assert self.total_num_heads >= self.total_kv_heads self.num_heads = self.total_num_heads // attn_tp_size self.head_dim = config.head_dim or (self.hidden_size // self.total_num_heads) self.q_size = self.head_dim * self.num_heads - self.num_kv_heads = self.total_kv_heads // attn_tp_size + self.num_kv_heads = max(1, self.total_kv_heads // attn_tp_size) self.kv_size = max(1, self.num_kv_heads * self.head_dim) self.scale = self.head_dim**-0.5