fix ci

Author: liuruyan

fastdeploy/worker/gpu_model_runner.py

@@ -2957,9 +2957,8 @@ def _get_prompt_logprobs_list(
                 raw_logprobs = self.sampler.compute_logprobs(logits)
             elif logprobs_mode == "raw_logits":
                 raw_logprobs = logits
-            token_ids, logprobs, ranks = self.sampler.gather_logprobs(
-                raw_logprobs, num_prompt_logprobs, prompt_token_ids_tensor
-            )
+            gathered = self.sampler.gather_logprobs(raw_logprobs, num_prompt_logprobs, prompt_token_ids_tensor)
+            token_ids, logprobs, ranks = gathered.logprob_token_ids, gathered.logprobs, gathered.selected_token_ranks
             # Synchronize before using token_ids, logprobs and ranks to ensure async copy are completed.
             paddle.device.synchronize()
             chunk_slice = slice(start_idx, start_idx + num_logits)

fastdeploy/worker/metax_model_runner.py

@@ -2833,9 +2833,8 @@ def _get_prompt_logprobs_list(
                 raw_logprobs = self.sampler.compute_logprobs(logits)
             elif logprobs_mode == "raw_logits":
                 raw_logprobs = logits
-            token_ids, logprobs, ranks = self.sampler.gather_logprobs(
-                raw_logprobs, num_prompt_logprobs, prompt_token_ids_tensor
-            )
+            gathered = self.sampler.gather_logprobs(raw_logprobs, num_prompt_logprobs, prompt_token_ids_tensor)
+            token_ids, logprobs, ranks = gathered.logprob_token_ids, gathered.logprobs, gathered.selected_token_ranks
             chunk_slice = slice(start_idx, start_idx + num_logits)
             logprobs_tensors.logprob_token_ids[chunk_slice].copy_(token_ids, False)
             logprobs_tensors.logprobs[chunk_slice].copy_(logprobs, False)

fastdeploy/worker/xpu_model_runner.py

@@ -271,9 +271,8 @@ def _get_prompt_logprobs_list(self, hidden_states: paddle.Tensor) -> list[Option
                 raw_logprobs = logits
             else:
                 raw_logprobs = self.sampler.compute_logprobs(logits)
-            token_ids, logprobs, ranks = self.sampler.gather_logprobs(
-                raw_logprobs, num_prompt_logprobs, prompt_token_ids_tensor
-            )
+            gathered = self.sampler.gather_logprobs(raw_logprobs, num_prompt_logprobs, prompt_token_ids_tensor)
+            token_ids, logprobs, ranks = gathered.logprob_token_ids, gathered.logprobs, gathered.selected_token_ranks
             chunk_slice = slice(start_idx, start_idx + num_logits)
             logprobs_tensors.logprob_token_ids[chunk_slice].copy_(token_ids, False)
             logprobs_tensors.logprobs[chunk_slice].copy_(logprobs, False)

tests/worker/test_gpu_prompt_logprobs.py

@@ -206,9 +206,14 @@ def test_prompt_logprobs(self):
             ref_raw_logprobs = model_runner.sampler.compute_logprobs(ref_logits)
             token_is = paddle.to_tensor(req.prompt_token_ids[1:], dtype="int64")
 
-            ref_token_ids, ref_logprobs, ref_ranks = model_runner.sampler.gather_logprobs(
+            gathered = model_runner.sampler.gather_logprobs(
                 ref_raw_logprobs, model_runner.fd_config.model_config.ori_vocab_size, token_is
             )
+            ref_token_ids, ref_logprobs, ref_ranks = (
+                gathered.logprob_token_ids,
+                gathered.logprobs,
+                gathered.selected_token_ranks,
+            )
             prompt_logprobs = model_runner._get_prompt_logprobs_list(hidden_states)[0]
             np.testing.assert_allclose(ref_logprobs.numpy(), prompt_logprobs.logprobs.numpy(), rtol=1e-04, atol=1e-04)
             np.testing.assert_allclose(