fix: logits processors for CB

[wallashss]

Description

This PR introduces a wrapper of the logits processor that are injected to vlllm that will handle the logits processor in a distributed way. The wrapper is initialized with the logits class and the batch_size. So, from the logits processor perspective it will "think" that it's only handling a request per step, while the wrapper receives the batch of logits, slice and redistribute for each request.

[github-actions]

👋 Hi! Thank you for contributing to vLLM support on Spyre.
Just a reminder: Make sure that your code passes all the linting checks, otherwise your PR won't be able to be merged. To do so, first install the linting requirements, then run format.sh and commit the changes. This can be done with uv directly:

uv sync --frozen --group lint --active --inexact

Or this can be done with pip:

uv pip compile --group lint > requirements-lint.txt
pip install -r requirements-lint.txt
bash format.sh

Now you are good to go 🚀

[wallashss]

bot:test

[yannicks1]

can we parametrize the test such that they get executed for SB and CB?

[yannicks1]

or do we already have a similar test for CB?

[wallashss]

We do not have this test for CB. I think the issue here is increase too much the time of CI. Moreover, they do not repro very well the issue of this PR.

[wallashss]

bot:test

[maxdebayser]

I think this is ok as a stopgap solution, but we should think about a more comprehensive solution later when more of the sampling params are implemented with logits processors. There could be performance implications in applying the LPs per requests vs per batch.

[wallashss]

bot:test

[wallashss]

bot:test

[wallashss]

bot:test

[wallashss]

bot:test

[wallashss]

bot:test