mrope的_get_position修改

vllm/v1/worker/gpu_model_runner.py

@@ -773,17 +773,23 @@ class GPUModelRunner(
     def _get_positions(self, num_tokens: Any):
         if isinstance(num_tokens, int):
             if self.uses_mrope:
+                if self.use_1d_mrope:
+                    return self.mrope_positions.gpu[: 3 * num_tokens].view(
+                        num_tokens, 3
+                    ).T
                 return self.mrope_positions.gpu[:, :num_tokens]
             if self.uses_xdrope_dim > 0:
                 return self.xdrope_positions.gpu[:, :num_tokens]
             return self.positions.gpu[:num_tokens]
         else:
             if self.uses_mrope:
+                if self.use_1d_mrope:
+                    return self.mrope_positions.gpu.view(-1, 3)[num_tokens].T
                 return self.mrope_positions.gpu[:, num_tokens]
             if self.uses_xdrope_dim > 0:
                 return self.xdrope_positions.gpu[:, num_tokens]
             return self.positions.gpu[num_tokens]
-
+        
     def _make_buffer(
         self, *size: int | torch.SymInt, dtype: torch.dtype, numpy: bool = True
     ) -> CpuGpuBuffer: