[BugFix] Disable TP for Qwen3-Omni audio encoder when heads % TP != 0 (#50858)
Signed-off-by: CalvinXKY <xky@users.noreply.github.com> Co-authored-by: CalvinXKY <xky@users.noreply.github.com>
kaiyuanxie committed
Aug 29, 2026 at 11:18 UTC
738bc8811ae31b733a7806adbc7a580d55729329
1 file changed
+8
-1
vllm/model_executor/models/qwen3_omni_moe_thinker.py
+8
-1
@@ -196,8 +196,13 @@ class Qwen3OmniMoeAudioAttention(nn.Module):
196
self.embed_dim = config.d_model
197
self.num_heads = config.encoder_attention_heads
198
self.head_dim = self.embed_dim // self.num_heads
199
+ # Audio encoder uses 20 heads. Shard across TP when divisible
200
+ # (e.g. TP=2/4); otherwise keep unreplicated (e.g. TP=8).
201
tp_size = get_tensor_model_parallel_world_size()
200
- self.num_local_heads = self.num_heads // tp_size
202
+ self.disable_tp = self.num_heads % tp_size != 0
203
+ self.num_local_heads = (
204
+ self.num_heads if self.disable_tp else self.num_heads // tp_size
205
+ )
206
207
if (self.head_dim * self.num_heads) != self.embed_dim:
208
raise ValueError(
@@ -214,6 +219,7 @@ class Qwen3OmniMoeAudioAttention(nn.Module):
219
total_num_kv_heads=self.num_heads,
220
bias=True,
221
prefix=f"{prefix}.qkv_proj",
222
+ disable_tp=self.disable_tp,
223
)
224
225
self.out_proj = RowParallelLinear(
@@ -221,6 +227,7 @@ class Qwen3OmniMoeAudioAttention(nn.Module):
227
output_size=self.embed_dim,
228
bias=True,
229
prefix=f"{prefix}.out_proj",
230
+ disable_tp=self.disable_tp,
231
)
232
233
self.attn = MMEncoderAttention(