🐛 fix tensor parallel

[joerunde]

Description

Fixes a bug introduced in #283 where the non-driver workers did not cache the output tokens for the next decode iteration.

This also allows TP tests with TP=2 to run on cpu, so that we can catch these bugs on GHA runs.

[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 🚀

[prashantgupta24]

bot:test
MARKERS="spyre and not quantized and not multi and not embedding"

[prashantgupta24]

bot:test
MARKERS="spyre and multi"

[prashantgupta24]

LGTM!

[prashantgupta24]

actually let me try to see if CB works with TP first before merging - edit: doesn't seem to work