Fix transformers MoE compatibility with VLLM

src/transformers/models/flex_olmo/modeling_flex_olmo.py

@@ -350,7 +350,7 @@ def forward(self, hidden_states):
         if self.norm_topk_prob:
             router_top_value /= router_top_value.sum(dim=-1, keepdim=True)
         router_top_value = router_top_value.to(router_logits.dtype)
-        router_scores = torch.zeros_like(router_logits).scatter_(1, router_indices, router_top_value)
+        router_scores = router_top_value
         return router_scores, router_indices
 
 

src/transformers/models/gpt_oss/modeling_gpt_oss.py

@@ -155,7 +155,7 @@ def forward(self, hidden_states):
         router_logits = F.linear(hidden_states, self.weight, self.bias)  # (seq_len, num_experts)
         router_top_value, router_indices = torch.topk(router_logits, self.top_k, dim=-1)  # (seq_len, top_k)
         router_top_value = torch.nn.functional.softmax(router_top_value, dim=1, dtype=router_top_value.dtype)
-        router_scores = torch.zeros_like(router_logits).scatter_(1, router_indices, router_top_value)
+        router_scores = router_top_value
         return router_scores, router_indices
 
 

src/transformers/models/gpt_oss/modular_gpt_oss.py

@@ -153,7 +153,7 @@ def forward(self, hidden_states):
         router_logits = F.linear(hidden_states, self.weight, self.bias)  # (seq_len, num_experts)
         router_top_value, router_indices = torch.topk(router_logits, self.top_k, dim=-1)  # (seq_len, top_k)
         router_top_value = torch.nn.functional.softmax(router_top_value, dim=1, dtype=router_top_value.dtype)
-        router_scores = torch.zeros_like(router_logits).scatter_(1, router_indices, router_top_value)
+        router_scores = router_top_value
         return router_scores, router_indices
 
 

src/transformers/models/minimax/modeling_minimax.py

@@ -467,7 +467,7 @@ def forward(self, hidden_states):
         router_logits = torch.nn.functional.softmax(router_logits.float(), dim=-1)
         router_top_value, router_indices = torch.topk(router_logits, self.top_k, dim=-1)  # (seq_len, top_k)
         router_top_value /= router_top_value.sum(dim=-1, keepdim=True)
-        router_scores = torch.zeros_like(router_logits).scatter_(1, router_indices, router_top_value)
+        router_scores = router_top_value
         return router_scores, router_indices
 
 

src/transformers/models/mixtral/modeling_mixtral.py

@@ -109,7 +109,7 @@ def forward(self, hidden_states):
         router_logits = torch.nn.functional.softmax(router_logits.float(), dim=-1)
         router_top_value, router_indices = torch.topk(router_logits, self.top_k, dim=-1)  # (seq_len, top_k)
         router_top_value /= router_top_value.sum(dim=-1, keepdim=True)
-        router_scores = torch.zeros_like(router_logits).scatter_(1, router_indices, router_top_value)
+        router_scores = router_top_value
         return router_scores, router_indices
 
 

src/transformers/models/mixtral/modular_mixtral.py

@@ -188,7 +188,7 @@ def forward(self, hidden_states):
         router_logits = torch.nn.functional.softmax(router_logits.float(), dim=-1)
         router_top_value, router_indices = torch.topk(router_logits, self.top_k, dim=-1)  # (seq_len, top_k)
         router_top_value /= router_top_value.sum(dim=-1, keepdim=True)
-        router_scores = torch.zeros_like(router_logits).scatter_(1, router_indices, router_top_value)
+        router_scores = router_top_value
         return router_scores, router_indices
 
 

src/transformers/models/olmoe/modeling_olmoe.py

@@ -353,7 +353,7 @@ def forward(self, hidden_states):
         if self.norm_topk_prob:
             router_top_value /= router_top_value.sum(dim=-1, keepdim=True)
         router_top_value = router_top_value.to(router_logits.dtype)
-        router_scores = torch.zeros_like(router_logits).scatter_(1, router_indices, router_top_value)
+        router_scores = router_top_value
         return router_scores, router_indices
 
 

src/transformers/models/qwen2_moe/modeling_qwen2_moe.py

@@ -347,7 +347,7 @@ def forward(self, hidden_states):
         if self.norm_topk_prob:
             router_top_value /= router_top_value.sum(dim=-1, keepdim=True)
         router_top_value = router_top_value.to(router_logits.dtype)
-        router_scores = torch.zeros_like(router_logits).scatter_(1, router_indices, router_top_value)
+        router_scores = router_top_value
         return router_scores, router_indices
 
 

src/transformers/models/qwen2_moe/modular_qwen2_moe.py

@@ -106,7 +106,7 @@ def forward(self, hidden_states):
         if self.norm_topk_prob:
             router_top_value /= router_top_value.sum(dim=-1, keepdim=True)
         router_top_value = router_top_value.to(router_logits.dtype)
-        router_scores = torch.zeros_like(router_logits).scatter_(1, router_indices, router_top_value)
+        router_scores = router_top_value
         return router_scores, router_indices
 
 

src/transformers/models/qwen3_moe/modeling_qwen3_moe.py

@@ -267,7 +267,7 @@ def forward(self, hidden_states):
         if self.norm_topk_prob:
             router_top_value /= router_top_value.sum(dim=-1, keepdim=True)
         router_top_value = router_top_value.to(router_logits.dtype)
-        router_scores = torch.zeros_like(router_logits).scatter_(1, router_indices, router_top_value)
+        router_scores = router_top_value
         return router_scores, router_indices
 
 

src/transformers/models/qwen3_next/modeling_qwen3_next.py

@@ -877,7 +877,7 @@ def forward(self, hidden_states):
         if self.norm_topk_prob:
             router_top_value /= router_top_value.sum(dim=-1, keepdim=True)
         router_top_value = router_top_value.to(router_logits.dtype)
-        router_scores = torch.zeros_like(router_logits).scatter_(1, router_indices, router_top_value)
+        router_scores = router_top_value
         return router_scores, router_indices
 
 

src/transformers/models/qwen3_omni_moe/modeling_qwen3_omni_moe.py

@@ -1376,7 +1376,7 @@ def forward(self, hidden_states):
         if self.norm_topk_prob:
             router_top_value /= router_top_value.sum(dim=-1, keepdim=True)
         router_top_value = router_top_value.to(router_logits.dtype)
-        router_scores = torch.zeros_like(router_logits).scatter_(1, router_indices, router_top_value)
+        router_scores = router_top_value
         return router_scores, router_indices
 
 
@@ -2802,7 +2802,7 @@ def forward(self, hidden_states):
         if self.norm_topk_prob:
             router_top_value /= router_top_value.sum(dim=-1, keepdim=True)
         router_top_value = router_top_value.to(router_logits.dtype)
-        router_scores = torch.zeros_like(router_logits).scatter_(1, router_indices, router_top_value)
+        router_scores = router_top_value
         return router_scores, router_indices
 
 

src/transformers/models/qwen3_vl_moe/modeling_qwen3_vl_moe.py

@@ -383,7 +383,7 @@ def forward(self, hidden_states):
         if self.norm_topk_prob:
             router_top_value /= router_top_value.sum(dim=-1, keepdim=True)
         router_top_value = router_top_value.to(router_logits.dtype)
-        router_scores = torch.zeros_like(router_logits).scatter_(1, router_indices, router_top_value)
+        router_scores = router_top_value
         return router_scores, router_indices
 
 

test_moe_inference.sh

@@ -0,0 +1,88 @@
+#!/bin/bash
+
+# Script to test vLLM inference with transformers backend for MoE models
+# Each model is tested and results are reported
+
+# Define colors for output
+GREEN='\033[0;32m'
+RED='\033[0;31m'
+YELLOW='\033[1;33m'
+NC='\033[0m' # No Color
+
+# Define models to test (architecture -> HuggingFace checkpoint)
+declare -A MODELS=(
+    ["olmoe"]="allenai/OLMoE-1B-7B-0924"
+    ["mixtral"]="TitanML/tiny-mixtral"
+    ["qwen2_moe"]="Qwen/Qwen1.5-MoE-A2.7B-Chat"
+    ["qwen3_moe"]="tiny-random/qwen3-moe"
+    ["gpt_oss"]="tiny-random/gpt-oss"
+)
+
+# Results tracking
+declare -A RESULTS
+
+echo "=========================================="
+echo "  MoE Models Inference Test Script"
+echo "=========================================="
+echo ""
+
+# Function to run inference test
+run_test() {
+    local model_name=$1
+    local model_checkpoint=$2
+
+    echo -e "${YELLOW}Testing: ${model_name} (${model_checkpoint})${NC}"
+    echo "-------------------------------------------"
+
+    # Set environment and run the command
+    VLLM_ENABLE_V1_MULTIPROCESSING=0 python vllm/examples/offline_inference/basic/generate.py \
+        --model "$model_checkpoint" \
+        --model-impl transformers \
+        --enforce-eager \
+        --no-enable-prefix-caching \
+        2>&1
+
+    local exit_code=$?
+
+    if [ $exit_code -eq 0 ]; then
+        RESULTS[$model_name]="SUCCESS"
+        echo -e "${GREEN}✓ ${model_name}: SUCCESS${NC}"
+    else
+        RESULTS[$model_name]="FAILED (exit code: $exit_code)"
+        echo -e "${RED}✗ ${model_name}: FAILED (exit code: $exit_code)${NC}"
+    fi
+
+    echo ""
+    return $exit_code
+}
+
+# Run tests for each model
+for model_name in "${!MODELS[@]}"; do
+    run_test "$model_name" "${MODELS[$model_name]}"
+done
+
+# Print summary
+echo ""
+echo "=========================================="
+echo "  SUMMARY"
+echo "=========================================="
+echo ""
+
+success_count=0
+fail_count=0
+
+for model_name in "${!RESULTS[@]}"; do
+    result="${RESULTS[$model_name]}"
+    if [[ "$result" == "SUCCESS" ]]; then
+        echo -e "${GREEN}✓ ${model_name}: ${result}${NC}"
+        ((success_count++))
+    else
+        echo -e "${RED}✗ ${model_name}: ${result}${NC}"
+        ((fail_count++))
+    fi
+done
+
+echo ""
+echo "-------------------------------------------"
+echo -e "Total: ${GREEN}${success_count} passed${NC}, ${RED}${fail_count} failed${NC}"
+echo "=========================================="