Support both paged and non-paged attention

.github/workflows/test.yml

@@ -43,17 +43,14 @@ jobs:
             markers: "v0 and cpu and e2e"
             flags: "--timeout=300"
           - name: "V1-e2e"
-            markers: "v1 and cpu and e2e"
+            markers: "v1 and cpu and e2e and not cb"
             flags: "--timeout=300 --forked"
-          - name: "V1-worker"
-            markers: "v1 and not e2e"
-            flags: "--timeout=300"
-          - name: "utils"
-            markers: "utils"
-            flags: "--timeout=300"
-          - name: "cb"
-            markers: "cb"
+          - name: "V1-cb"
+            markers: "v1 and cpu and cb"
             flags: "--timeout=300 --forked"
+          - name: "V1-worker and utils"
+            markers: "v1 and not e2e or utils"
+            flags: "--timeout=300"
 
     name: "${{ matrix.test_suite.name }} (${{ matrix.vllm_version.name }})"
 
@@ -163,10 +160,6 @@ jobs:
           # `uv run`, to avoid having `uv run` re-sync any dependencies or 
           # re-install the vllm_sypre package from source
           source .venv/bin/activate
-          if [ ${{ matrix.test_suite.markers }} == "cb" ]; then
-          # install custom fms branch
-              uv pip install git+https://github.com/foundation-model-stack/foundation-model-stack@paged_attn_mock --force-reinstall
-          fi
           # commands to run if condition is true
           python3 -m pytest ${{ matrix.test_suite.flags }} \
             tests -v -m "${{ matrix.test_suite.markers }}"

pyproject.toml

@@ -12,7 +12,7 @@ readme = "README.md"
 license = {text = "Apache 2"}
 dependencies = [
     "fms-model-optimizer>=0.2.0",
-    "ibm-fms==1.0.0",
+    "ibm-fms==1.1.0",
     "vllm>=0.9.0,!=0.9.1",
 ]
 requires-python = ">=3.9"

tests/e2e/test_spyre_cb.py

@@ -9,7 +9,8 @@
 from typing import Any
 
 import pytest
-from spyre_util import generate_cb_spyre_vllm_output, get_spyre_model_list
+from spyre_util import (generate_cb_spyre_vllm_output, get_spyre_backend_list,
+                        get_spyre_model_list)
 from vllm import EngineArgs, SamplingParams
 from vllm.v1.engine import EngineCoreRequest
 from vllm.v1.engine.core import EngineCore
@@ -18,16 +19,12 @@
 from vllm_spyre.v1.core.scheduler import ContinuousBatchingSpyreScheduler
 
 
+@pytest.mark.cb
+@pytest.mark.v1
 @pytest.mark.parametrize("max_num_seqs", [2, 3, 4],
                          ids=lambda val: f"max_num_seqs({val})")
 @pytest.mark.parametrize("model", get_spyre_model_list())
-@pytest.mark.parametrize(
-    "backend", [pytest.param("eager", marks=pytest.mark.cpu, id="eager")])
-@pytest.mark.parametrize("cb",
-                         [pytest.param(1, marks=pytest.mark.cb, id="cb")])
-# commenting v1 since we don't want this test to run with v1 marker yet
-# @pytest.mark.parametrize("vllm_version",
-#                          [pytest.param("V1", marks=pytest.mark.v1, id="v1")])
+@pytest.mark.parametrize("backend", get_spyre_backend_list())
 @pytest.mark.parametrize(
     "prompts",
     [
@@ -53,9 +50,7 @@ def test_cb_handling(
     model: str,
     backend: str,
     max_num_seqs: int,
-    cb: int,
     prompts: list[str],
-    # vllm_version: str,
     monkeypatch: pytest.MonkeyPatch,
 ):
     """Test that the spyre worker correctly handles
@@ -80,7 +75,7 @@ def test_cb_handling(
         tensor_parallel_size=1,
         backend=backend,
         max_num_seqs=max_num_seqs,
-        use_cb=cb,
+        use_cb=1,
         monkeypatch=monkeypatch,
     )
 
@@ -647,9 +642,9 @@ def augment_checked_steps(
 
 
 @pytest.mark.cb
+@pytest.mark.v1
 @pytest.mark.parametrize("model", get_spyre_model_list())
-@pytest.mark.parametrize(
-    "backend", [pytest.param("eager", marks=pytest.mark.cpu, id="eager")])
+@pytest.mark.parametrize("backend", get_spyre_backend_list())
 @pytest.mark.parametrize("max_num_seqs", [2])
 @pytest.mark.parametrize(
     "seqs_max_tokens,prompts_lengths,steps_add_reqs,checked_steps,"

vllm_spyre/model_executor/model_loader/spyre.py

@@ -353,17 +353,25 @@ def forward(
         **extra_kwargs,
     ) -> torch.Tensor:
 
+        # import will be not be needed/ handled by FMS soon
+        import fms.utils.spyre.paged  # noqa # pylint: disable=unused-import
+
+        # specify attention type for continuous batching
+        extra_kwargs['attn_name'] = "spyre_paged_attn"
+
+        # additional (paged) attention arguments
+        extra_kwargs['current_tkv_mask'] = current_tkv_mask
+        extra_kwargs['left_padded_prompt_mask'] = left_padded_prompt_mask
+        extra_kwargs['block_table'] = block_table
+        extra_kwargs['slot_mapping'] = slot_mapping
+
         output = self.model(
             input_ids,
             position_ids=position_ids,
             mask=mask,
             past_key_value_states=self.past_key_value_states,
             use_cache=use_cache,
             only_last_token=only_last_token,
-            current_tkv_mask=current_tkv_mask,
-            left_padded_prompt_mask=left_padded_prompt_mask,
-            block_table=block_table,
-            slot_mapping=slot_mapping,
             **extra_kwargs,
         )
 
@@ -401,6 +409,9 @@ def forward(
         **extra_kwargs,
     ) -> torch.Tensor:
 
+        # specify attention type for static batching
+        extra_kwargs['attn_name'] = "sdpa_bidirectional"
+
         output = self.model(
             input_ids,
             position_ids=position_ids,