[high prio][CB] 🐛 fix warmup

[yannicks1]

🐛 fix warmup

Currently, the number of blocks (in particular the kv cache tensor shape) has to be consistent across warmup and inference. PR #362 introduced inconsistency which led to recompilation for a certain model/TP combination.
This PR fixes above and refactors the warmup to make it more readable.

[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
export TORCH_SENDNN_CACHE_ENABLE=0
export VLLM_SPYRE_TEST_MODEL_LIST='ibm-granite/granite-3.3-8b-instruct'

[wallashss]

It's not private anymore.

Suggested change

	def _get_num_blocks_available(self) -> int:
	def _get_num_blocks_available(self) -> int:

[wallashss]

LGTM!