[ROCm][Perf] Fix Qwen3-vLLM audio encoder TP when heads are not divisible by TP size (#45900)
Signed-off-by: abrahamzewoudie <azewoudi@amd.com>
abrahamzewoudie committed
Sep 9, 2026 at 16:28 UTC
08b3e67b669d18d5bbd14e456b8bb4794916e402
1 file changed
+9
-6
vllm/model_executor/models/qwen3_omni_moe_thinker.py
+9
-6
@@ -199,10 +199,9 @@ class Qwen3OmniMoeAudioAttention(nn.Module):
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()
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
- )
202
+ disable_tp = self.num_heads % tp_size != 0
203
+ effective_tp = 1 if disable_tp else tp_size
204
+ self.num_local_heads = self.num_heads // effective_tp
205
206
if (self.head_dim * self.num_heads) != self.embed_dim:
207
raise ValueError(
@@ -218,16 +217,16 @@ class Qwen3OmniMoeAudioAttention(nn.Module):
217
total_num_heads=self.num_heads,
218
total_num_kv_heads=self.num_heads,
219
bias=True,
220
+ disable_tp=disable_tp,
221
prefix=f"{prefix}.qkv_proj",
222
- disable_tp=self.disable_tp,
222
)
223
224
self.out_proj = RowParallelLinear(
225
input_size=self.embed_dim,
226
output_size=self.embed_dim,
227
bias=True,
228
+ disable_tp=disable_tp,
229
prefix=f"{prefix}.out_proj",
230
- disable_tp=self.disable_tp,
230
)
231
232
self.attn = MMEncoderAttention(
@@ -278,16 +277,20 @@ class Qwen3OmniMoeAudioEncoderLayer(nn.Module):
277
)
278
self.self_attn_layer_norm = nn.LayerNorm(self.embed_dim)
279
self.activation_fn = _ACTIVATION_REGISTRY[config.activation_function]
280
+ tp_size = get_tensor_model_parallel_world_size()
281
+ disable_tp = config.encoder_attention_heads % tp_size != 0
282
self.fc1 = ColumnParallelLinear(
283
self.embed_dim,
284
config.encoder_ffn_dim,
285
bias=True,
286
+ disable_tp=disable_tp,
287
prefix=f"{prefix}.fc1",
288
)
289
self.fc2 = RowParallelLinear(
290
config.encoder_ffn_dim,
291
self.embed_dim,
292
bias=True,
293
+ disable_tp=disable_tp,
294
prefix=f"{prefix}.fc2",
295
)
296
self.final_layer_norm = nn.LayerNorm(self.embed_dim)