fix synchronize error in _calc_spec_decode_metadata

Signed-off-by: Ronald1995 <[email protected]>

Author: Ronald1995

vllm_ascend/sample/rejection_sampler.py

@@ -383,7 +383,7 @@ def rejection_greedy_sample_pytorch(
         target_argmax[global_idx].to(output_token_ids.dtype),
         output_token_ids
     )
-    output_token_ids_.copy_(output_token_ids)
+    output_token_ids.copy_(output_token_ids_)
     # Fill bonus token.
     needs_bonus = is_greedy & (first_mismatch_pos_per_req
                                >= draft_tokens_per_req)

vllm_ascend/worker/model_runner_v1.py

@@ -2144,16 +2144,31 @@ def _calc_spec_decode_metadata(
         target_logits_indices += arange
 
         # TODO: Optimize the CPU -> NPU copy.
-        cu_num_draft_tokens = torch.from_numpy(cu_num_draft_tokens).to(
-            self.device, non_blocking=True)
-        cu_num_sampled_tokens = torch.from_numpy(cu_num_sampled_tokens).to(
-            self.device, non_blocking=True)
-        logits_indices = torch.from_numpy(logits_indices).to(self.device,
-                                                             non_blocking=True)
-        target_logits_indices = torch.from_numpy(target_logits_indices).to(
-            self.device, non_blocking=True)
-        bonus_logits_indices = torch.from_numpy(bonus_logits_indices).to(
-            self.device, non_blocking=True)
+        cu_num_draft_tokens = (
+            torch.from_numpy(cu_num_draft_tokens)
+            .pin_memory()
+            .to(self.device, non_blocking=True)
+        )
+        cu_num_sampled_tokens = (
+            torch.from_numpy(cu_num_sampled_tokens)
+            .pin_memory()
+            .to(self.device, non_blocking=True)
+        )
+        logits_indices = (
+            torch.from_numpy(logits_indices)
+            .pin_memory()
+            .to(self.device, non_blocking=True)
+        )
+        target_logits_indices = (
+            torch.from_numpy(target_logits_indices)
+            .pin_memory()
+            .to(self.device, non_blocking=True)
+        )
+        bonus_logits_indices = (
+            torch.from_numpy(bonus_logits_indices)
+            .pin_memory()
+            .to(self.device, non_blocking=True)
+        )
 
         # Compute the draft token ids.
         # draft_token_indices:      [  1,   2,   3, 105, 106, 208]