[CB] Scheduling constraints regarding number of available blocks/pages

[yannicks1]

[CB] Scheduling constraints regarding number of available blocks/pages

changes:

• moved hard coded BLOCK_SIZE (64) variable to Platform class and import it where needed instead of defining it in multiple different places.
• introduced scheduler constraints regarding number of available blocks/pages in can_schedule() (need to keep track of the reserved block ids per request in model_runner.reserved_blocks)
• wrote unit test for new scheduler constraint asserting n_reserved_blocks and n_used_blocks
• introduce env variable VLLM_SPYRE_N_BLOCKS to set the number of available blocks (needs to be done during initialzation of all classes, if someone knows a better way, please tell me) for the unit tests.
• using math.ceil(n / d) instead of ((n + d - 1) // d) for better readability
• renaming model_runner.free_pages to model_runner.block_pool (as it is called in upstream vLLM)

closes #260

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

FYI

scheduler.n_free_blocks

should give us the number of free blocks at each step and should be assertable :)

We can technically get the exact blocks used by:

engine_core.model_executor.driver_worker.worker.model_runner.free_blocks

but that might be an overkill here?

[yannicks1]

no, I think this is a nice additional check! Will add this

[yannicks1]

done: e5697bc

[sducouedic]

LGTM

[sducouedic]

I am confused why these two variables are prefixed with 'max'

[yannicks1]

good point, was a relict of the past. I just updated the variable names!

[joerunde]

I don't think this comment is correct? From the code it looks like this is an override for the number of kv cache blocks which is used in place of the simple max_model_len * max_num_seqs calculation.

I'm assuming this is for passing in the known good values for the available blocks given a tested model and card combo. Can we also consider using the existing kv cache override instead of setting up a new one? There's --num-gpu-blocks-override which will set the available blocks in the scheduler config

[yannicks1]

yes, the comment is from copy pasting the code from VLLM_SPYRE_USE_CB 😆
good pointer with --num-gpu-blocks-override, will look into that!

[joerunde]

What I'm seeing here is that:

• self._get_num_blocks_available() only uses info from self.model_runner
• The resulting n_blocks_avail is only used to modify the model runner, and the model runner's model

So I think that _get_num_blocks_available needs to move to the model runner class, and it should be responsible for finalizing itself after the warmup is complete

[yannicks1]

although this was not part of this PR, you are absolutely right and I moved the function 😇

[joerunde]

Ah what I meant here is that the model runner should encapsulate this entirely. This whole block can be replaced with model_runner.finish_warmup(), and there should be no access of the model runner's private methods or direct access of the model here in the worker

[joerunde]

can we consolidate some of the if envs_spyre.VLLM_SPYRE_N_BLOCKS > 0: checks into here and return the override from this method?