Optimize expert selection

Simplify the case when norm_topk_prob is set. Avoid extra normalization and smaller softmax computation.

Results (Measured on M3 max):

Before:

📊 Benchmark Summary:
==================================================
Prompt:     1066.413 tokens/sec
Generation: 62.131 tokens/sec
==================================================

After:

📊 Benchmark Summary:
==================================================
Prompt:     1103.736 tokens/sec
Generation: 69.197 tokens/sec
==================================================

Testing: presubmit

Author: karimnosseir

python/src/coreai_models/models/macos/qwen3_moe.py

@@ -125,16 +125,19 @@ def __init__(
         self.norm_topk_prob = norm_topk_prob
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        gates = self.gate(x)
-        gates = torch.softmax(gates, dim=-1, dtype=torch.float32)
+        router_logits = self.gate(x).to(torch.float32)
 
-        active_experts_scores, active_experts_indices = torch.topk(
-            gates, self.top_k, dim=-1, largest=True
-        )
-        active_experts_indices = active_experts_indices.to(torch.uint16)
         if self.norm_topk_prob:
-            partition = torch.sum(active_experts_scores, axis=-1, keepdims=True)
-            active_experts_scores = active_experts_scores / partition
+            top_logits, active_experts_indices = torch.topk(
+                router_logits, self.top_k, dim=-1, largest=True
+            )
+            active_experts_scores = torch.softmax(top_logits, dim=-1)
+        else:
+            gates = torch.softmax(router_logits, dim=-1)
+            active_experts_scores, active_experts_indices = torch.topk(
+                gates, self.top_k, dim=-1, largest=True
+            )
+        active_experts_indices = active_experts_indices.to(torch.uint16)
 
         y_active_experts = self.switch_mlp(x, active_experts_indices)
         active_experts_scores = active_experts_scores.unsqueeze(-1)