[CB][do not merge] enable warmup for batch size 1

[yannicks1]

[CB] enable warmup for batch size 1

in #285 we wanted to allow batch size 1 for continuous batching. However, the warmup did not support batch size 1. With this small fix it does work on CPU at least. It does not currently compile on the card. Would need to compare graphs next...

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

[yannicks1]

bot:test
TEST_FILE=tests/e2e/test_spyre_cb.py MARKERS="spyre"

[prashantgupta24]

the warmup did not support batch size 1

naive question - can we add a failing test for that which this PR fixes?

Edit: I guess we had the assert earlier which prevented batch size 1 from running 🤷

[joerunde]

ah interesting- I swear I had tested manually on a dev pod and saw batch size 1 working, but I dunno how that was happening when the test clearly failed 🤦.

Anyway nice fix!

[joerunde]

Alternatively, could the InputBatch construct itself with:

self.max_num_reqs = min(max_num_reqs, 2)

since we know that it'll always need at least 2, and then we avoid reconstructing it in the worker here? That way we have a much smaller diff to back out once we can lift this bs>=2 restriction

[yannicks1]

not sure if I follow here. it has to be >=2 for the warmup. with the min(1,2) we would still fail?

[sducouedic]

would that work if you directly set model_runner.input_batch.max_num_reqs = 2, instead of instantiating a new InputBatch?

[yannicks1]

no, because InputBatch initialization gets model_runner.input_batch.max_num_reqs..

[yannicks1]

max_num_seqs occurs 17 times in the init of the InputBatch. It is not a single attribute, but used to construct several attributes. So re-initializing is simpler...

[yannicks1]

I tried to test this on the card, but it failed. Not clear to me why though. Will have to compare graphs tomorrow.

[yannicks1]

bot:test
TEST_FILE=tests/e2e/test_spyre_cb.py MARKERS="spyre"

[waleedqk]

bot:test
TEST_FILE=tests/e2e/test_spyre_cb.py MARKERS="spyre"

[yannicks1]

bot:test
TEST_FILE=tests/e2e/test_spyre_cb.py MARKERS="spyre"

[JRosenkranz]

I tried to test this on the card, but it failed. Not clear to me why though. Will have to compare graphs tomorrow.

@yannicks1 in order to support this, we will need to add the following which will enable symbolic shapes for size 1 dimensions (if marked dynamic):

from torch.fx.experimental import _config as config
config.backed_size_oblivious = True

[yannicks1]

thanks for the hint @JRosenkranz ! I think though your suggested changes should be tried with #312 which actually uses torch 2.7.1 and applies real batch size 1. This PR (#287) is doing padding under the hood. It is merely supposed to support --max_num_seqs 1 from a scheduling perspective...