✅ add assertions for warmup mode context

[prashantgupta24]

Description

This PR adds assertions for operations done within the warmup context. Earlier we saw a bug in which the decode warmup was moved outside the warmup context because of one simple indentation issue and it caused a segmentation fault with no clue as to why it was happening.

These assertions should catch the error when running even on CPU.

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

[sducouedic]

Interesting idea

[joerunde]

I would prefer if we addressed this with a refactor to put everything that should be inside the warmup context into its own method. As is the full warmup method here is already ~150 lines and should be broken up anyway

# setup dummy_requests

with _maybe_warmup_context():
    self._dynamic_warmup(dummy_requests)

# the rest of the post-compile warmup

[yannicks1]

lgtm