From 8276320a8a691adf8747d4d93b12e4c2de208e2d Mon Sep 17 00:00:00 2001
From: i-yuanyukun <i-yuanyukun@stepfun.com>
Date: Thu, 18 Dec 2025 17:03:15 +0800
Subject: [PATCH] [Bugfix] compute ffn output param order

---
 vllm/v1/worker/gpu_ffn_model_runner.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/vllm/v1/worker/gpu_ffn_model_runner.py b/vllm/v1/worker/gpu_ffn_model_runner.py
index 3f85267d8536a..921824f408552 100644
--- a/vllm/v1/worker/gpu_ffn_model_runner.py
+++ b/vllm/v1/worker/gpu_ffn_model_runner.py
@@ -351,7 +351,7 @@ class GPUFFNModelRunner(LoRAModelRunnerMixin):
                 hidden_states, dim=0
             )
             ffn_output = self.model.compute_ffn_output(
-                current_layer_idx, gathered_hidden_states
+                gathered_hidden_states, current_layer_idx
             )
 
             # Extract the output corresponding to current rank
@@ -361,7 +361,7 @@ class GPUFFNModelRunner(LoRAModelRunnerMixin):
         else:
             # Single TP case
             rank_ffn_output = self.model.compute_ffn_output(
-                current_layer_idx, hidden_states
+                hidden_states, current_layer_idx
             )
 
         return rank_ffn_output