Need to handle 1D and 2D position_ids at the same time

Signed-off-by: Jin Li <[email protected]>

Author: liji-nv

tensorrt_llm/_torch/modules/attention.py

@@ -951,7 +951,7 @@ def forward_impl(self, position_ids: Optional[torch.Tensor],
 
         hidden_states = hidden_states[:num_tokens, ...]
         if position_ids is not None:
-            position_ids = position_ids[:, :num_tokens]
+            position_ids = position_ids[..., :num_tokens]
 
         if self.is_lite:
             compressed_kv, k_pe = self.kv_a_proj_with_mqa(hidden_states).split(