[ROCm][Perf] Fix Qwen3-vLLM audio encoder TP when heads are not divisible by TP size (#45900)

Signed-off-by: abrahamzewoudie <azewoudi@amd.com>

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)