remove cuda graph copy

Signed-off-by: yewentao256 <zhyanwentao@126.com>

vllm/model_executor/models/granitemoe.py

@@ -113,8 +113,11 @@ class GraniteMoeMoE(nn.Module):
         )
 
     def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
-        # NOTE: hidden_states can have either 1D or 2D shape.
-        orig_shape = hidden_states.shape
+        assert hidden_states.dim() in (1, 2), (
+            "GraniteMoeMoE only supports 1D or 2D inputs"
+        )
+        num_tokens = hidden_states.shape[0]
+        is_input_1d = hidden_states.dim() == 1
         hidden_states = hidden_states.view(-1, self.hidden_size)
 
         if self.is_sequence_parallel:
@@ -128,10 +131,10 @@ class GraniteMoeMoE(nn.Module):
             final_hidden_states = tensor_model_parallel_all_gather(
                 final_hidden_states, 0
             )
-            num_tokens = orig_shape[0]
+
             final_hidden_states = final_hidden_states[:num_tokens]
 
-        return final_hidden_states.view(orig_shape)
+        return final_hidden_states.squeeze(0) if is_input_1d else final_hidden_states
 
 
 class GraniteMoeAttention(nn.Module):

vllm/model_executor/models/qwen2_moe.py

@@ -168,8 +168,10 @@ class Qwen2MoeSparseMoeBlock(nn.Module):
         )
 
     def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
-        # NOTE: hidden_states can have either 1D or 2D shape.
-        orig_shape = hidden_states.shape
+        assert hidden_states.dim() in (1, 2), (
+            "Qwen2MoeSparseMoeBlock only supports 1D or 2D inputs"
+        )
+        is_input_1d = hidden_states.dim() == 1
         hidden_dim = hidden_states.shape[-1]
         hidden_states = hidden_states.view(-1, hidden_dim)
 
@@ -185,7 +187,7 @@ class Qwen2MoeSparseMoeBlock(nn.Module):
                 final_hidden_states
             )
 
-        return final_hidden_states.view(orig_shape)
+        return final_hidden_states.squeeze(0) if is_input_1d else final_hidden_states
 
 
 class Qwen2MoeAttention(nn.Module):

vllm/model_executor/models/qwen3_next.py

@@ -178,8 +178,10 @@ class Qwen3NextSparseMoeBlock(nn.Module):
         )
 
     def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
-        # NOTE: hidden_states can have either 1D or 2D shape.
-        orig_shape = hidden_states.shape
+        assert hidden_states.dim() in (1, 2), (
+            "Qwen3NextSparseMoeBlock only supports 1D or 2D inputs"
+        )
+        is_input_1d = hidden_states.dim() == 1
         num_tokens, hidden_dim = hidden_states.shape
         hidden_states = hidden_states.view(-1, hidden_dim)
 
@@ -211,7 +213,7 @@ class Qwen3NextSparseMoeBlock(nn.Module):
                 final_hidden_states
             )
 
-        return final_hidden_states.view(orig_shape)
+        return final_hidden_states.squeeze(0) if is_input_1d else final_hidden_states
 
 
 class Qwen3NextGatedDeltaNet(nn.Module, MambaBase):

vllm/model_executor/models/step3_text.py

@@ -85,7 +85,10 @@ class FusedMoEBlock(nn.Module):
         )
 
     def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
-        orig_shape = hidden_states.shape
+        assert hidden_states.dim() in (1, 2), (
+            "FusedMoEBlock only supports 1D or 2D inputs"
+        )
+        is_input_1d = hidden_states.dim() == 1
         hidden_dim = hidden_states.shape[-1]
         hidden_states = hidden_states.view(-1, hidden_dim)
 
@@ -97,7 +100,7 @@ class FusedMoEBlock(nn.Module):
         if self.tp_size > 1:
             final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states)
 
-        return final_hidden_states.view(orig_shape)
+        return final_hidden_states.squeeze(0) if is_input_1d else final_hidden_states
 
 
 class Step3TextMLP(nn.Module):