[CB] use used block ids for dummy batch size 2

[yannicks1]

[CB] use used block ids for dummy batch size 2

solves #258

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

Status: To be tested on AIU Spyre

[yannicks1]

bot:test
TEST_FILE=tests/e2e/test_spyre_cb.py MARKERS="spyre" VLLM_SPYRE_TEST_MODEL_LIST="tiny-granite-3.2-8b"

[yannicks1]

bot:test
TEST_FILE=tests/e2e/test_spyre_cb.py MARKERS="spyre" VLLM_SPYRE_TEST_MODEL_LIST='tiny-granite-3.2-8b'

[yannicks1]

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"

[yannicks1]

Good news, I spun up a pod and tested this, and it actually works 🎉
ran test_spyre_cb.py and cb_spyre_inference.py and validated with aiu-smi that we actually run on Spyre.

ready to review!

[prashantgupta24]

Wondering if we can add a test for this somehow 🤔

[prashantgupta24]

Worth adding a test where we only send 1 request and assert a bunch of stuff? 🤷

[prashantgupta24]

@gmarinho2 had something in his PR here - https://github.com/vllm-project/vllm-spyre/pull/252/files but he was assuming we used dummy_req_ids2blocks which we don't anymore

[yannicks1]

what about an assert on the shape dimension that corresponds to the batch size for the decode path? might be an overkill to write test for this. Plus an assert is saver...

[prashantgupta24]

Adding only 1 request and asserting scheduler's tkv could be a good start?

[yannicks1]

I think that should be sufficient. Reluctant to add even more tests as they are already taking long and this can be easily asserted in the code

[wallashss]

Agree, I think the base tests already cover this case.

[prashantgupta24]

Ah sorry I worded it incorrectly, yeah we already have test cases that cover the scenario, what I meant was we could add asserts in the tests themselves to cover this specific PR scenario where len(cached_requests) == 1. On the other hand, since we only have access to the model_runner obj during testing, the only thing we could have asserted was model_runner.model.indices.size which is already covered by the assert in the code, so I guess it's okay 🤷

[wallashss]

LGTM

[yannicks1]

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

[yannicks1]

ran the only failing test manually on the card and it succeeded. Thus merging this PR!