[FIX] Fix shape mismatch for swapped sequences when logprobs > 0

[derange-alembic]

When the logprobs of a sequence is set to be larger than 0 and it was swapped out using recomputation policy, the _get_logprobs() in sampler.py will encounter IndexError: shape mismatch: indexing tensors could not be broadcast together with shapes due to the obliviousness of the output_token_ids.

[derange-alembic]

This issue is also reported in #1847.

[Yard1]

This looks good to me. cc @WoosukKwon @zhuohan123

[WoosukKwon]

BTW, I found another bug when running a batch with n = 2.

  File "/home/wskwon/workspace/vllm/vllm/core/scheduler.py", line 284, in schedule
    scheduler_outputs = self._schedule()
  File "/home/wskwon/workspace/vllm/vllm/core/scheduler.py", line 142, in _schedule
    assert seq_group.num_seqs() == 1, (
AssertionError: Waiting sequence group should have only one prompt sequence.

I think this happens when one of the two sequences finish earlier and the other one is preempted and resumed using recomputation.

This was fixed by #2186

[WoosukKwon]

Why is output_tokens used for prompt logprobs?

[richardzhuang0412]

For prompt_logprobs I think the issue still remains

[AetherPrior]

Hi, I've raised a new issue (#3032), as the codebase has changed since then, and the issue still persists.
CC: @derange-alembic , @Yard1 , @WoosukKwon

[github-actions]

This pull request has been automatically marked as stale because it has not had any activity within 90 days. It will be automatically closed if no further activity occurs within 30 days. Leave a comment if you feel this pull request should remain open. Thank you!

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. @derange-alembic please rebase it. https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork