[CB] Support pseudo batch size 1 for decode, adjust warmup

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

[yannicks1]

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

[yannicks1]

I found the culprit and fixed it here. Tested manually on the card and works now! Ready for review!

[prashantgupta24]

Can we not wait for #312 instead?

[yannicks1]

@prashantgupta24 We don't know when the support for 'true' batch size 1 for decode will be established. It requires some changes lower in the stack. However, the request for batch size 1 is here and with this workaround we are able to support it. While I agree that #312 cleans up the logic and is nicer from a code point, please note that the performance benefit of batch size 1 over batch size 2 is marginal. As the compiler team has more pressing issues, #312 has been de-prioritized and we want to merge this instead...

[yannicks1]

I removed the test case max_num_seqs=1 here to save some runs. Can we please get that approved and into main @maxdebayser @prashantgupta24 @wallashss

[prashantgupta24]

I'll try to take a look at this after lunch!

[yannicks1]

@wallashss @joerunde @prashantgupta24 @maxdebayser can we please get this in...

[maxdebayser]

Sorry, I somehow missed the comment where you had already answered about the re-initialization of the input batch. That was my last remaining question I think.