[Core] Support logprobs with spec decode + async scheduling #29223
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request adds support for logprobs with speculative decoding and asynchronous scheduling. The changes involve refactoring how cumulative token counts are calculated and passed to correctly process logprobs in these scenarios. The modifications in vllm/v1/sample/rejection_sampler.py and vllm/v1/worker/gpu_model_runner.py seem correct and well-structured. New tests are added to cover these cases. My main concern is the significant increase in tolerance in tests/v1/sample/test_logprobs.py for comparing logprobs, which might hide numerical precision issues. Please see the specific comment for details.
|
This pull request has merge conflicts that must be resolved before it can be |
Signed-off-by: Nick Hill <nhill@redhat.com>
njhill
force-pushed
the
as-sd-logprobs
branch
from
3caee35 to
cc55c14
Compare
njhill
requested review from
ApostaC,
WoosukKwon,
alexm-redhat,
heheda12345,
robertgshaw2-redhat and
ywang96
as code owners
| cu_num_tokens = None | ||
| if return_cu_num_tokens: | ||
| cu_num_tokens = [0] + valid_mask.sum(axis=1).cumsum().tolist() | ||
| if len(discard_req_indices) > 0: |
There was a problem hiding this comment.
Why is this done after computing cu_num_tokens?
There was a problem hiding this comment.
Because cu_num_tokens is used to index into the logprobs tensors that don't take the discarded indices into account, so doing it beforehand results in incorrect output.
Originally it was done before which was a bug, fixed by #29216.
2 participants
No description provided.