[Bugfix] Pad hidden_states to avoid cross-ring AllGatherV

Signed-off-by: ApsarasX <apsarax@outlook.com>

Author: ApsarasX

vllm_ascend/models/deepseek_v2.py

@@ -229,14 +229,13 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
 
         if self.tp_size > 1:
             # pass
-            num_tokens, hidden_size = hidden_states.shape
-            if num_tokens < self.tp_size:
-                target_size = self.tp_size
-                new_hidden_states = torch.empty([target_size, hidden_size],
-                                                dtype=hidden_states.dtype,
-                                                device=hidden_states.device)
-                new_hidden_states[:num_tokens] = hidden_states
-                hidden_states = new_hidden_states
+            num_tokens, _ = hidden_states.shape
+            padded_num_tokens = (self.tp_size -
+                                 num_tokens % self.tp_size) % self.tp_size
+            # Pad hidden_states to make it divisible by tp_size to avoid cross-ring AllGatherV on 910B2C
+            if padded_num_tokens > 0:
+                hidden_states = nn.functional.pad(hidden_states,
+                                                  (0, 0, 0, padded_num_tokens))
             chunk_hidden_states = torch.tensor_split(hidden_states,
                                                      self.tp_size,
                                                      dim=0)
@@ -259,8 +258,8 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
             dist.all_gather(list(chunk_hidden_states), router_hidden_states,
                             self.tp_group)
             final_hidden_states = torch.cat(chunk_hidden_states, dim=0)
-            if num_tokens < self.tp_size:
-                final_hidden_states = final_hidden_states[:num_tokens]
+            if padded_num_tokens > 0:
+                final_hidden_states = final_hidden_states[:-padded_num_tokens]
         else:
             final_hidden_states = router_hidden_states