From 08b3e67b669d18d5bbd14e456b8bb4794916e402 Mon Sep 17 00:00:00 2001 From: abrahamzewoudie Date: Wed, 9 Sep 2026 16:28:24 +0300 Subject: [PATCH] [ROCm][Perf] Fix Qwen3-vLLM audio encoder TP when heads are not divisible by TP size (#45900) Signed-off-by: abrahamzewoudie --- .../models/qwen3_omni_moe_thinker.py | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/vllm/model_executor/models/qwen3_omni_moe_thinker.py b/vllm/model_executor/models/qwen3_omni_moe_thinker.py index 77a463205e..acc09f5038 100755 --- a/vllm/model_executor/models/qwen3_omni_moe_thinker.py +++ b/vllm/model_executor/models/qwen3_omni_moe_thinker.py @@ -199,10 +199,9 @@ class Qwen3OmniMoeAudioAttention(nn.Module): # Audio encoder uses 20 heads. Shard across TP when divisible # (e.g. TP=2/4); otherwise keep unreplicated (e.g. TP=8). tp_size = get_tensor_model_parallel_world_size() - self.disable_tp = self.num_heads % tp_size != 0 - self.num_local_heads = ( - self.num_heads if self.disable_tp else self.num_heads // tp_size - ) + disable_tp = self.num_heads % tp_size != 0 + effective_tp = 1 if disable_tp else tp_size + self.num_local_heads = self.num_heads // effective_tp if (self.head_dim * self.num_heads) != self.embed_dim: raise ValueError( @@ -218,16 +217,16 @@ class Qwen3OmniMoeAudioAttention(nn.Module): total_num_heads=self.num_heads, total_num_kv_heads=self.num_heads, bias=True, + disable_tp=disable_tp, prefix=f"{prefix}.qkv_proj", - disable_tp=self.disable_tp, ) self.out_proj = RowParallelLinear( input_size=self.embed_dim, output_size=self.embed_dim, bias=True, + disable_tp=disable_tp, prefix=f"{prefix}.out_proj", - disable_tp=self.disable_tp, ) self.attn = MMEncoderAttention( @@ -278,16 +277,20 @@ class Qwen3OmniMoeAudioEncoderLayer(nn.Module): ) self.self_attn_layer_norm = nn.LayerNorm(self.embed_dim) self.activation_fn = _ACTIVATION_REGISTRY[config.activation_function] + tp_size = get_tensor_model_parallel_world_size() + disable_tp = config.encoder_attention_heads % tp_size != 0 self.fc1 = ColumnParallelLinear( self.embed_dim, config.encoder_ffn_dim, bias=True, + disable_tp=disable_tp, prefix=f"{prefix}.fc1", ) self.fc2 = RowParallelLinear( config.encoder_ffn_dim, self.embed_dim, bias=True, + disable_tp=disable_tp, prefix=f"{prefix}.fc2", ) self.final_layer_norm = nn.LayerNorm(self.embed_dim)