Support both paged and non-paged attention

[yannicks1]

Final FMS API supporting paged and non-paged attention

This final Fms API will support both the static and continuous batching APIs. Currently consuming this fms branch. Only minimal changes required on vLLM Spyre side.

Note: Do not merge until in fms main.

[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

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 🚀

[joerunde]

Fixing DCO and getting this ready to merge with fms 1.1

[joerunde]

bot:test
MARKERS="cb and spyre"

[joerunde]

bot:test
MARKERS="cb and spyre"

[joerunde]

bot:test
TEST_FILE=tests/e2e/test_spyre_online.py

[prashantgupta24]

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

[prashantgupta24]

new test incoming - test_cb_max_tokens - We will need to uncomment the v1 marker on this test.

[ckadner]

I had to go over the test matrix several times and look at what tests get run in which grouping. Maybe we can simplify those in the future.

[sducouedic]

LGTM

[sducouedic]

# TO DO: implement warmup for continuous batching

Is this comment still valid? beside warming up with dynamic sizes, what has to be done?

[yannicks1]

no, should be fine, I will remove it

[sducouedic]

maybe the with _maybe_warmup_context() can used one level up before the for loop, instead of importing warmup_mode from torch_sendnn for each warmup shape. Maybe not that important since we don't warmup for many shapes usually.

[joerunde]

I think it needs to be done on each shape here so that the context closes and runs update_lazyhandle() after each shape, right?

[sducouedic]

aah ok you are probably right, I don't really understand what is happening with that context to be honest.
btw, I was looking for update_lazyhandle in the code to understand, and couldn't find it, even with a general "find" in the directory. When did that disappear?

[tjohnson31415]

The call to update_lazyhandle is now hidden in the warmup_mode context provided by torch_sendnn such that it gets called when the context exits. Use of warmup_mode was part of the change to remove the deprecated sendnn_decoder backend in #186.

So you'd have to look in the torch_sendnn code, but the warmup_mode context is basically just:

class warmup_mode:
    def __init__(self, enabled: bool = True):
        global _warmup_mode
        self._old_mode = _warmup_mode
        _warmup_mode = enabled

    def __enter__(self):
        return None

    def __exit__(self, exc_type, exc_val, exc_tb):
        global _warmup_mode
        update_lazyhandle()
        _warmup_mode = self._old_mode

[joerunde]

Merging now that we have driver changes to support this 🎉

[yannicks1]

nice to have that finally in main !