Repository: vllm-project/vllm-spyre
Pull request #79: [CB ] e2e continuous batching tests
Author: login: prashantgupta24
Created: 2025-04-09T20:30:24Z | Merged: 2025-05-12T20:38:22Z

Adding some continuous batching unit tests, for https://github.com/vllm-project/vllm-spyre/issues/75

cc @yannicks1 @maxdebayser 

Also fixes https://github.com/vllm-project/vllm-spyre/issues/151 - basically installs fms from custom branch before running the `cb` tests.

Once the fms custom branch is merged to main, we can delete the duplicated `test_batch_handling` test within `tests/e2e/test_spyre_basic.py` and have the test within `test_spyre_cb` test both CB and non-cb conditions. We can also remove the `cb` flag and have the test run with other `e2e` tests.


Diff:
diff --git a/.github/workflows/test-spyre.yml b/.github/workflows/test-spyre.yml
index 9c02504bd..1ebc1df1d 100644
--- a/.github/workflows/test-spyre.yml
+++ b/.github/workflows/test-spyre.yml
@@ -55,6 +55,9 @@ jobs:
           - name: "utils"
             markers: "utils"
             flags: "--timeout=300"
+          - name: "cb"
+            markers: "cb"
+            flags: "--timeout=300 --forked"
 
     name: "${{ matrix.test_suite.name }} (${{ matrix.vllm_version.name }})"
 
@@ -147,5 +150,10 @@ 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 }}"
diff --git a/pyproject.toml b/pyproject.toml
index 49f438201..d04288b8a 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -121,6 +121,7 @@ markers = [
     "e2e: Tests using end-to-end engine spin-up",
     "v0: Tests using vLLM v0 engine",
     "v1: Tests using vLLM v1 engine",
+    "cb: Continuous batching tests",
     "cpu: Tests using CPU (i.e. eager) backend",
     "spyre: Tests using Spyre hardware backend",
     "decoder: Tests for decoder models",
diff --git a/tests/e2e/test_spyre_cb.py b/tests/e2e/test_spyre_cb.py
new file mode 100644
index 000000000..a31277d27
--- /dev/null
+++ b/tests/e2e/test_spyre_cb.py
@@ -0,0 +1,235 @@
+"""Verification of continuous batching
+
+Run `python -m pytest tests/test_spyre_cb.py`.
+"""
+
+import pytest
+from spyre_util import generate_cb_spyre_vllm_output, get_spyre_model_list
+from transformers import AutoTokenizer
+from vllm import EngineArgs, SamplingParams
+from vllm.v1.engine.llm_engine import LLMEngine as V1LLMEngine
+
+
+@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(
+    "prompts",
+    [
+        [
+            "7 6 5 4",
+            "10 9 8 7",
+        ],
+        [
+            "7 6 5 4",
+            "10 9 8 7",
+            "8 7 6 5",
+        ],
+        [
+            "7 6 5 4",
+            "10 9 8 7",
+            "8 7 6 5",
+            "9 8 7 6",
+        ],
+    ],
+    ids=lambda val: f"num_prompts({len(val)})",
+)
+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
+    continuous batches of requests that
+    finish after different numbers of forward passes"""
+
+    vllm_sampling_params = SamplingParams(max_tokens=20,
+                                          temperature=0,
+                                          stop="1",
+                                          ignore_eos=True,
+                                          logprobs=0)
+
+    # Ensure that both:
+    # - The model doesn't crash
+    # - The output sequences are correct
+    vllm_results = generate_cb_spyre_vllm_output(
+        model=model,
+        prompts=prompts,
+        max_model_len=2048,
+        block_size=2048,
+        sampling_params=vllm_sampling_params,
+        tensor_parallel_size=1,
+        backend=backend,
+        max_num_seqs=max_num_seqs,
+        use_cb=cb,
+        monkeypatch=monkeypatch,
+    )
+
+    for i, prompt in enumerate(prompts):
+        assert (vllm_results[i]["text"] == [
+            " " + " ".join(
+                str(i)
+                for i in range(int(prompt.split()[-1]) - 1, 1, -1)) + " "
+        ][0])
+
+
+@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")])
+# @pytest.mark.parametrize("vllm_version",
+#                          [pytest.param("V1", marks=pytest.mark.v1, id="v1")])
+def test_cb_with_steps(model: str, backend: str, cb: int,
+                       monkeypatch: pytest.MonkeyPatch):
+    """Test that the spyre worker correctly handles
+    continuous batches of requests and one sequence
+    an exit the batch while a new ones gets prefilled
+    and appended to the batch"""
+    with monkeypatch.context() as m:
+
+        max_tokens = 20
+        max_num_seqs = 2  # defines max batch size
+
+        prompt1 = "7 6 5 4"
+        prompt2 = "10 9 8 7"
+        prompt3 = "8 7 6 5"
+
+        tokenizer = AutoTokenizer.from_pretrained(model)
+        stop_token = "1"
+        # issue with JackFram/llama-160m tokenizer
+        # it returns 2 tokens for "1",
+        # the first is quotes, the second is the
+        # actual number 1.
+        stop_token_ids = [tokenizer.encode(stop_token,
+                                          add_special_tokens=False)[1]] \
+        if "llama-194m" in model \
+        else tokenizer.encode(stop_token,
+                                          add_special_tokens=False)
+
+        sampling_params = SamplingParams(
+            max_tokens=max_tokens,
+            temperature=0,
+            stop_token_ids=stop_token_ids,
+            ignore_eos=True,
+        )
+
+        # set env vars
+        m.setenv("VLLM_SPYRE_USE_CB", "1")
+        m.setenv("VLLM_USE_V1", "1")
+        m.setenv("VLLM_SPYRE_DYNAMO_BACKEND", backend)
+
+        # To get deterministic execution in V1
+        # and to enable InprocClient
+        m.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
+
+        # start the engine
+        engine_args = EngineArgs(model=model,
+                                 max_model_len=2048,
+                                 max_num_seqs=max_num_seqs)
+
+        engine = V1LLMEngine.from_engine_args(engine_args)
+        engine_core = engine.engine_core.engine_core
+
+        # add first request
+        engine.add_request("1", prompt1, sampling_params)
+
+        assert len(engine_core.scheduler.waiting) == 1
+        assert len(engine_core.scheduler.running) == 0
+
+        request_outputs = engine.step()
+        assert len(request_outputs) == 1  # only 1 request
+        assert request_outputs[0].request_id == "1"  # req 1 is decoding
+
+        assert len(engine_core.scheduler.waiting) == 0
+        assert len(engine_core.scheduler.running) == 1
+
+        # add another request
+        engine.add_request("2", prompt2, sampling_params)
+
+        assert len(engine_core.scheduler.waiting) == 1  # req 2 is waiting
+        assert len(engine_core.scheduler.running) == 1
+
+        request_outputs = engine.step()
+        assert len(request_outputs) == 1  # still only 1 request
+        assert request_outputs[0].request_id == "2"  # req 2 is decoding
+
+        assert len(engine_core.scheduler.waiting) == 0
+        assert len(engine_core.scheduler.running) == 2
+
+        request_outputs = engine.step()
+        assert len(request_outputs) == 2
+        assert request_outputs[0].request_id == "2"  # req 2 is decoding
+        assert request_outputs[1].request_id == "1"  # req 1 is decoding
+
+        assert len(engine_core.scheduler.waiting) == 0
+        assert len(engine_core.scheduler.running) == 2
+
+        request_outputs = engine.step()
+        assert len(request_outputs) == 2
+        assert request_outputs[0].request_id == "2"  # req 2 is decoding
+        assert request_outputs[1].request_id == "1"  # req 1 is decoding
+
+        assert len(engine_core.scheduler.waiting) == 0
+        assert len(engine_core.scheduler.running) == 2
+
+        request_outputs = engine.step()
+        assert len(request_outputs) == 2
+        assert request_outputs[0].request_id == "2"  # req 2 is decoding
+        assert request_outputs[1].request_id == "1"  # req 1 is decoding
+
+        assert len(engine_core.scheduler.waiting) == 0
+        assert len(engine_core.scheduler.running) == 2
+
+        # add the third request
+        # but since VLLM_SPYRE_MAX_BATCH_SIZE=2
+        # this request is waiting
+        engine.add_request("3", prompt3, sampling_params)
+
+        assert len(engine_core.scheduler.waiting) == 1
+        assert len(engine_core.scheduler.running) == 2
+
+        request_outputs = engine.step()
+        assert len(request_outputs) == 2  # both requests decoding now
+        assert request_outputs[0].request_id == "2"  # req 2 is decoding
+        assert request_outputs[1].request_id == "1"  # req 1 is decoding
+
+        assert len(engine_core.scheduler.waiting) == 1
+        assert len(engine_core.scheduler.running) == 2
+
+        request_outputs = engine.step()
+        assert len(request_outputs) == 2  # both requests decoding now
+        assert request_outputs[0].request_id == "2"  # req 2 is decoding
+        assert request_outputs[1].request_id == "1"  # req 1 is decoding
+        assert request_outputs[1].finished  # request 1 is done
+        assert request_outputs[1].outputs[0].text == " 3 2 "
+
+        assert len(engine_core.scheduler.waiting) == 1
+        assert len(engine_core.scheduler.running) == 1
+
+        # req 3 is scheduled now
+        request_outputs = engine.step()
+        assert len(request_outputs) == 1
+        assert request_outputs[0].request_id == "3"  # req 3 is decoding
+
+        assert len(engine_core.scheduler.waiting) == 0
+        assert len(engine_core.scheduler.running) == 2
+
+        request_outputs = engine.step()
+        assert len(request_outputs) == 2  # requests 2 and 3 decoding now
+        assert request_outputs[0].request_id == "3"  # req 3 is decoding
+        assert request_outputs[1].request_id == "2"  # req 2 is decoding
+
+        assert len(engine_core.scheduler.waiting) == 0
+        assert len(engine_core.scheduler.running) == 2
diff --git a/tests/spyre_util.py b/tests/spyre_util.py
index fd9dde26b..feb7089de 100644
--- a/tests/spyre_util.py
+++ b/tests/spyre_util.py
@@ -205,6 +205,45 @@ def generate_spyre_vllm_output(model: str, prompts: list[str],
     return results
 
 
+# Support for continuous batching
+def generate_cb_spyre_vllm_output(
+    model: str,
+    prompts: list[str],
+    max_model_len: int,
+    block_size: int,
+    sampling_params: Union[SamplingParams, list[SamplingParams]],
+    tensor_parallel_size: int,
+    backend: str,
+    max_num_seqs: int,
+    use_cb: int,
+    monkeypatch: pytest.MonkeyPatch,
+) -> list[dict[str, Any]]:
+    with monkeypatch.context() as m:
+
+        m.setenv("VLLM_SPYRE_USE_CB", str(use_cb))
+        m.setenv("VLLM_USE_V1", "1")
+        m.setenv("VLLM_SPYRE_DYNAMO_BACKEND", backend)
+
+        vllm_model = LLM(
+            model=model,
+            tokenizer=model,
+            max_model_len=max_model_len,
+            max_num_seqs=max_num_seqs,
+            block_size=block_size,
+            tensor_parallel_size=tensor_parallel_size,
+        )
+
+        vllm_outputs = vllm_model.generate(prompts, sampling_params)
+        results = []
+
+        for req_output in vllm_outputs:
+            result = {}
+            result["text"] = req_output.outputs[0].text
+            results.append(result)
+
+        return results
+
+
 # Hugging Face
 def generate_hf_output(
         model: str, prompts: list[str],


Reviews:
user: login: prashantgupta24
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-09T21:54:25Z
user: login: prashantgupta24
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-09T23:07:27Z
user: login: maxdebayser
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-10T13:14:45Z
user: login: maxdebayser
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-10T13:17:41Z
user: login: maxdebayser
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-10T13:28:32Z
user: login: prashantgupta24
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-10T16:09:24Z
user: login: prashantgupta24
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-10T16:09:45Z
user: login: prashantgupta24
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-10T16:10:42Z
user: login: prashantgupta24
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-11T20:28:13Z
user: login: prashantgupta24
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-14T22:12:25Z
user: login: prashantgupta24
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-14T22:33:09Z
user: login: prashantgupta24
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-14T22:33:37Z
user: login: prashantgupta24
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-14T23:07:25Z
user: login: prashantgupta24
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-02T18:39:37Z
user: login: prashantgupta24
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-12T20:07:55Z
user: login: wallashss
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-12T20:09:57Z
user: login: prashantgupta24
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-12T20:15:57Z
user: login: wallashss
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-12T20:21:19Z
user: login: prashantgupta24
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-12T20:26:46Z
user: login: wallashss
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-12T20:27:14Z
user: login: wallashss
body: LGTM! Thanks a lot for this @prashantgupta24 
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-05-12T20:28:33Z

Review comments:
diff_hunk: @@ -0,0 +1,153 @@
+"""Verification of continuous batching
+
+Run `python -m pytest tests/test_spyre_cb.py`.
+"""
+
+import pytest
+from spyre_util import generate_cb_spyre_vllm_output, get_spyre_model_list
+from vllm import EngineArgs, SamplingParams
+from vllm.v1.engine.llm_engine import LLMEngine as V1LLMEngine
+
+
+@pytest.mark.parametrize("model", get_spyre_model_list())
+@pytest.mark.parametrize("backend", ["eager", "inductor"])
+def test_cb_handling(
+    model: str,
+    backend: str,
+):
+    """Test that the spyre worker correctly handles 
+    continuous batches of requests that
+    finish after different numbers of forward passes"""
+
+    max_tokens1 = 20
+    max_tokens2 = 17
+    max_tokens3 = 11
+
+    sampling_params1 = SamplingParams(max_tokens=max_tokens1,
+                                      temperature=0.0,
+                                      stop="1",
+                                      ignore_eos=True)
+
+    sampling_params2 = SamplingParams(max_tokens=max_tokens2,
+                                      temperature=0.0,
+                                      stop="1",
+                                      ignore_eos=True)
+
+    sampling_params3 = SamplingParams(max_tokens=max_tokens3,
+                                      temperature=0.0,
+                                      stop="1",
+                                      ignore_eos=True)
+
+    vllm_sampling_params = [
+        sampling_params1,
+        sampling_params2,
+        sampling_params3,
+    ]
+
+    # These prompts are ordered so that they don't finish in the
+    # order given
+    prompts = [
+        "7 6 5 4",
+        "10 9 8 7",
+        "8 7 6 5",
+    ]
+
+    # Ensure that both:
+    # - The model doesn't crash
+    # - The output sequences are correct
+    vllm_results = generate_cb_spyre_vllm_output(
+        model=model,
+        prompts=prompts,
+        max_model_len=2048,
+        block_size=2048,
+        sampling_params=vllm_sampling_params,
+        tensor_parallel_size=1,
+        backend=backend,
+    )
+
+    assert vllm_results[0]["text"] == " 3 2 "
+    assert vllm_results[1]["text"] == " 6 5 4 3 2 "
+    assert vllm_results[2]["text"] == " 4 3 2 "
+
+
+@pytest.mark.parametrize("model", get_spyre_model_list())
+@pytest.mark.parametrize("backend", ["eager", "inductor"])
+def test_cb_with_engine_steps(model: str, backend: str,
+                              monkeypatch: pytest.MonkeyPatch):
+    """Test that the spyre worker correctly handles 
+    continuous batches of requests and one sequence 
+    an exit the batch while a new ones gets prefilled 
+    and appended to the batch"""
+    with monkeypatch.context() as m:
+        max_tokens1 = 20
+        max_tokens2 = 17
+
+        sampling_params1 = SamplingParams(max_tokens=max_tokens1,
+                                          temperature=0.0,
+                                          stop="1",
+                                          ignore_eos=True)
+
+        sampling_params2 = SamplingParams(max_tokens=max_tokens2,
+                                          temperature=0.0,
+                                          stop="1",
+                                          ignore_eos=True)
+
+        prompt1 = "7 6 5 4"
+        prompt2 = "10 9 8 7"
+
+        # set env vars
+        m.setenv("VLLM_SPYRE_WARMUP_PROMPT_LENS", "64")
+        m.setenv("VLLM_SPYRE_WARMUP_NEW_TOKENS",
+                 str(max([max_tokens1, max_tokens2])))
+
+        m.setenv("VLLM_SPYRE_USE_CB", "1")
+        m.setenv("VLLM_USE_V1", "1")
+
+        m.setenv("VLLM_SPYRE_MAX_CONTEXT_LENGTH", "2048")
+        m.setenv("VLLM_SPYRE_MAX_BATCH_SIZE", str(2))
+        m.setenv("VLLM_SPYRE_DYNAMO_BACKEND", backend)
+
+        # start the engine
+        engine_args = EngineArgs(
+            model=model,
+            enforce_eager=True,  # reduce test time
+        )
+
+        engine = V1LLMEngine.from_engine_args(engine_args)
path: tests/test_spyre_cb.py
user: login: prashantgupta24
body: might be easier to set up `engine_core` and loop through the `step` function for it, that way we can directly get the `waiting` and `running` queues. Will give it a shot next
created_at: 2025-04-09T21:54:25Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,282 @@
+"""Verification of continuous batching
+
+Run `python -m pytest tests/test_spyre_cb.py`.
+"""
+
+import time
+import uuid
+
+import pytest
+from spyre_util import generate_cb_spyre_vllm_output, get_spyre_model_list
+from transformers import AutoTokenizer
+from vllm import EngineArgs, SamplingParams
+from vllm.v1.engine import EngineCoreRequest
+from vllm.v1.engine.core import EngineCore
+from vllm.v1.engine.llm_engine import LLMEngine as V1LLMEngine
+from vllm.v1.executor.abstract import Executor
+
+
+@pytest.mark.parametrize("model", get_spyre_model_list())
+@pytest.mark.parametrize("backend", ["eager", "inductor"])
+def test_cb_handling(
+    model: str,
+    backend: str,
+):
+    """Test that the spyre worker correctly handles 
+    continuous batches of requests that
+    finish after different numbers of forward passes"""
+
+    max_tokens1 = 20
+    max_tokens2 = 17
+    max_tokens3 = 11
+
+    sampling_params1 = SamplingParams(max_tokens=max_tokens1,
+                                      temperature=0.0,
+                                      stop="1",
+                                      ignore_eos=True)
+
+    sampling_params2 = SamplingParams(max_tokens=max_tokens2,
+                                      temperature=0.0,
+                                      stop="1",
+                                      ignore_eos=True)
+
+    sampling_params3 = SamplingParams(max_tokens=max_tokens3,
+                                      temperature=0.0,
+                                      stop="1",
+                                      ignore_eos=True)
+
+    vllm_sampling_params = [
+        sampling_params1,
+        sampling_params2,
+        sampling_params3,
+    ]
+
+    # These prompts are ordered so that they don't finish in the
+    # order given
+    prompts = [
+        "7 6 5 4",
+        "10 9 8 7",
+        "8 7 6 5",
+    ]
+
+    # Ensure that both:
+    # - The model doesn't crash
+    # - The output sequences are correct
+    vllm_results = generate_cb_spyre_vllm_output(
+        model=model,
+        prompts=prompts,
+        max_model_len=2048,
+        block_size=2048,
+        sampling_params=vllm_sampling_params,
+        tensor_parallel_size=1,
+        backend=backend,
+    )
+
+    assert vllm_results[0]["text"] == " 3 2 "
+    assert vllm_results[1]["text"] == " 6 5 4 3 2 "
+    assert vllm_results[2]["text"] == " 4 3 2 "
+
+
+@pytest.mark.parametrize("model", get_spyre_model_list())
+@pytest.mark.parametrize("backend", ["eager", "inductor"])
+def test_cb_with_engine_steps(model: str, backend: str,
+                              monkeypatch: pytest.MonkeyPatch):
+    """Test that the spyre worker correctly handles 
+    continuous batches of requests and one sequence 
+    an exit the batch while a new ones gets prefilled 
+    and appended to the batch"""
+    with monkeypatch.context() as m:
+        max_tokens1 = 20
+        max_tokens2 = 17
+
+        sampling_params1 = SamplingParams(max_tokens=max_tokens1,
+                                          temperature=0.0,
+                                          stop="1",
+                                          ignore_eos=True)
+
+        sampling_params2 = SamplingParams(max_tokens=max_tokens2,
+                                          temperature=0.0,
+                                          stop="1",
+                                          ignore_eos=True)
+
+        prompt1 = "7 6 5 4"
+        prompt2 = "10 9 8 7"
+
+        # set env vars
+        m.setenv("VLLM_SPYRE_WARMUP_PROMPT_LENS", "64")
+        m.setenv("VLLM_SPYRE_WARMUP_NEW_TOKENS",
+                 str(max([max_tokens1, max_tokens2])))
+
+        m.setenv("VLLM_SPYRE_USE_CB", "1")
+        m.setenv("VLLM_USE_V1", "1")
+
+        m.setenv("VLLM_SPYRE_MAX_CONTEXT_LENGTH", "2048")
+        m.setenv("VLLM_SPYRE_MAX_BATCH_SIZE", str(2))
+        m.setenv("VLLM_SPYRE_DYNAMO_BACKEND", backend)
+
+        # start the engine
+        engine_args = EngineArgs(
+            model=model,
+            enforce_eager=True,  # reduce test time
+        )
+
+        engine = V1LLMEngine.from_engine_args(engine_args)
+
+        # add first request
+        engine.add_request("0", prompt1, sampling_params1)
+
+        request_outputs = engine.step()
+        assert len(request_outputs) == 1  # only 1 request
+        assert request_outputs[0].request_id == "0"
+
+        request_outputs = engine.step()
+        assert len(request_outputs) == 1  # still only 1 request
+        assert request_outputs[0].request_id == "0"
+
+        # add another request
+        engine.add_request("1", prompt2, sampling_params2)
+
+        request_outputs = engine.step()
+        assert len(request_outputs) == 1  # still only 1 request
+        assert request_outputs[0].request_id == "0"
+
+        request_outputs = engine.step()
+        assert len(request_outputs) == 1  # still only 1 request
+        assert request_outputs[0].request_id == "1"  # req 1 is decoding
+
+        request_outputs = engine.step()
+        assert len(request_outputs) == 2  # both requests decoding now
+
+        request_outputs = engine.step()
+        assert len(request_outputs) == 2  # both requests decoding now
+
+        request_outputs = engine.step()
+        assert len(request_outputs) == 2  # both requests decoding now
+        assert request_outputs[0].finished  # request 1 is done
+        assert request_outputs[0].outputs[0].text == " 3 2 "
+
+        request_outputs = engine.step()
+        assert len(request_outputs) == 1  # only req #2 is left
+        assert request_outputs[0].request_id == "1"  # req 1 is decoding
+
+
+@pytest.mark.parametrize("model", get_spyre_model_list())
+@pytest.mark.parametrize("backend", ["eager", "inductor"])
+def test_cb_with_core_steps(model: str, backend: str,
+                            monkeypatch: pytest.MonkeyPatch):
+    """Test that the spyre worker correctly handles
+    continuous batches of requests and one sequence
+    an exit the batch while a new ones gets prefilled
+    and appended to the batch"""
+    with monkeypatch.context() as m:
+        max_tokens1 = 20
+        max_tokens2 = 17
+
+        sampling_params1 = SamplingParams(max_tokens=max_tokens1,
+                                          temperature=0.0,
+                                          stop="1",
+                                          ignore_eos=True)
+
+        sampling_params2 = SamplingParams(max_tokens=max_tokens2,
+                                          temperature=0.0,
+                                          stop="1",
+                                          ignore_eos=True)
+
+        prompt1 = "7 6 5 4"
+        prompt2 = "10 9 8 7"
+
+        # set env vars
+        m.setenv("VLLM_SPYRE_WARMUP_PROMPT_LENS", "64")
+        m.setenv("VLLM_SPYRE_WARMUP_NEW_TOKENS",
+                 str(max([max_tokens1, max_tokens2])))
+
+        m.setenv("VLLM_SPYRE_USE_CB", "1")
+        m.setenv("VLLM_USE_V1", "1")
+
+        m.setenv("VLLM_SPYRE_MAX_CONTEXT_LENGTH", "2048")
+        m.setenv("VLLM_SPYRE_MAX_BATCH_SIZE", str(2))
+        m.setenv("VLLM_SPYRE_DYNAMO_BACKEND", backend)
+
+        # start the engine
+        engine_args = EngineArgs(model=model, )
+
+        vllm_config = engine_args.create_engine_config()
+        executor_class = Executor.get_class(vllm_config)
+
+        engine_core = EngineCore(vllm_config=vllm_config,
+                                 executor_class=executor_class,
+                                 log_stats=True)
+        # engine = V1LLMEngine.from_engine_args(engine_args)
+
+        tokenizer = AutoTokenizer.from_pretrained(model)
+        prompt_tokens1 = tokenizer(prompt1).input_ids
+        prompt_tokens2 = tokenizer(prompt2).input_ids
+
+        # # add first request
+        engine_core.add_request(
+            EngineCoreRequest(
+                request_id=uuid.uuid4(),
+                prompt=prompt1,
+                prompt_token_ids=prompt_tokens1,
+                mm_inputs=None,
+                mm_hashes=None,
+                mm_placeholders=None,
+                sampling_params=sampling_params1,
+                eos_token_id=None,
+                arrival_time=time.time(),
+                lora_request=None,
+            ))
+        assert len(engine_core.scheduler.waiting) == 1
+        assert len(engine_core.scheduler.running) == 0
+
+        _ = engine_core.step()
+        assert len(engine_core.scheduler.waiting) == 0
+        assert len(engine_core.scheduler.running) == 1
+
+        engine_core.add_request(
+            EngineCoreRequest(
+                request_id=uuid.uuid4(),
+                prompt=prompt2,
+                prompt_token_ids=prompt_tokens2,
+                mm_inputs=None,
+                mm_hashes=None,
+                mm_placeholders=None,
+                sampling_params=sampling_params2,
+                eos_token_id=None,
+                arrival_time=time.time(),
+                lora_request=None,
+            ))
+        assert len(engine_core.scheduler.waiting) == 1
+        assert len(engine_core.scheduler.running) == 1
+
+        _ = engine_core.step()
+        assert len(engine_core.scheduler.waiting) == 0
+        assert len(engine_core.scheduler.running) == 1
+
+        _ = engine_core.step()
+        assert len(engine_core.scheduler.waiting) == 0
+        assert len(engine_core.scheduler.running) == 2
+
+        _ = engine_core.step()
+        assert len(engine_core.scheduler.waiting) == 0
+        assert len(engine_core.scheduler.running) == 2
+
+        _ = engine_core.step()
+        assert len(engine_core.scheduler.waiting) == 0
+        assert len(engine_core.scheduler.running) == 2
+
+        _ = engine_core.step()
+        assert len(engine_core.scheduler.waiting) == 0
+        assert len(engine_core.scheduler.running) == 2
+
+        _ = engine_core.step()
+        assert len(engine_core.scheduler.waiting) == 0
+        assert len(engine_core.scheduler.running) == 2
+
+        _ = engine_core.step()
+        assert len(engine_core.scheduler.waiting) == 0
+        assert len(engine_core.scheduler.running) == 2
+
+        _ = engine_core.step()
+        assert len(engine_core.scheduler.waiting) == 0
+        assert len(engine_core.scheduler.running) == 2
path: tests/test_spyre_cb.py
user: login: prashantgupta24
body: Hmm, it seems like this goes on for more than I anticipated, request[0] should have hit the `stop_reason` already and ejected from `running` queue, but it stays here until it hits its `max_tokens` 🤔 Perhaps I'm understanding this wrong.
created_at: 2025-04-09T23:07:27Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,282 @@
+"""Verification of continuous batching
+
+Run `python -m pytest tests/test_spyre_cb.py`.
+"""
+
+import time
+import uuid
+
+import pytest
+from spyre_util import generate_cb_spyre_vllm_output, get_spyre_model_list
+from transformers import AutoTokenizer
+from vllm import EngineArgs, SamplingParams
+from vllm.v1.engine import EngineCoreRequest
+from vllm.v1.engine.core import EngineCore
+from vllm.v1.engine.llm_engine import LLMEngine as V1LLMEngine
+from vllm.v1.executor.abstract import Executor
+
+
+@pytest.mark.parametrize("model", get_spyre_model_list())
+@pytest.mark.parametrize("backend", ["eager", "inductor"])
+def test_cb_handling(
+    model: str,
+    backend: str,
+):
+    """Test that the spyre worker correctly handles 
+    continuous batches of requests that
+    finish after different numbers of forward passes"""
+
+    max_tokens1 = 20
+    max_tokens2 = 17
+    max_tokens3 = 11
+
+    sampling_params1 = SamplingParams(max_tokens=max_tokens1,
+                                      temperature=0.0,
+                                      stop="1",
+                                      ignore_eos=True)
+
+    sampling_params2 = SamplingParams(max_tokens=max_tokens2,
+                                      temperature=0.0,
+                                      stop="1",
+                                      ignore_eos=True)
+
+    sampling_params3 = SamplingParams(max_tokens=max_tokens3,
+                                      temperature=0.0,
+                                      stop="1",
+                                      ignore_eos=True)
+
+    vllm_sampling_params = [
+        sampling_params1,
+        sampling_params2,
+        sampling_params3,
+    ]
+
+    # These prompts are ordered so that they don't finish in the
+    # order given
+    prompts = [
+        "7 6 5 4",
+        "10 9 8 7",
+        "8 7 6 5",
+    ]
+
+    # Ensure that both:
+    # - The model doesn't crash
+    # - The output sequences are correct
+    vllm_results = generate_cb_spyre_vllm_output(
+        model=model,
+        prompts=prompts,
+        max_model_len=2048,
+        block_size=2048,
+        sampling_params=vllm_sampling_params,
+        tensor_parallel_size=1,
+        backend=backend,
+    )
+
+    assert vllm_results[0]["text"] == " 3 2 "
+    assert vllm_results[1]["text"] == " 6 5 4 3 2 "
+    assert vllm_results[2]["text"] == " 4 3 2 "
+
+
+@pytest.mark.parametrize("model", get_spyre_model_list())
+@pytest.mark.parametrize("backend", ["eager", "inductor"])
+def test_cb_with_engine_steps(model: str, backend: str,
+                              monkeypatch: pytest.MonkeyPatch):
+    """Test that the spyre worker correctly handles 
+    continuous batches of requests and one sequence 
+    an exit the batch while a new ones gets prefilled 
+    and appended to the batch"""
+    with monkeypatch.context() as m:
+        max_tokens1 = 20
+        max_tokens2 = 17
+
+        sampling_params1 = SamplingParams(max_tokens=max_tokens1,
+                                          temperature=0.0,
+                                          stop="1",
+                                          ignore_eos=True)
+
+        sampling_params2 = SamplingParams(max_tokens=max_tokens2,
+                                          temperature=0.0,
+                                          stop="1",
+                                          ignore_eos=True)
+
+        prompt1 = "7 6 5 4"
+        prompt2 = "10 9 8 7"
path: tests/test_spyre_cb.py
user: login: maxdebayser
body: Can you perhaps test with 3 prompts and a max batch size of 2 so that we can verify that the third prompt stays in the waiting state until on of the two others leaves the batch?
created_at: 2025-04-10T13:14:45Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,282 @@
+"""Verification of continuous batching
+
+Run `python -m pytest tests/test_spyre_cb.py`.
+"""
+
+import time
+import uuid
+
+import pytest
+from spyre_util import generate_cb_spyre_vllm_output, get_spyre_model_list
+from transformers import AutoTokenizer
+from vllm import EngineArgs, SamplingParams
+from vllm.v1.engine import EngineCoreRequest
+from vllm.v1.engine.core import EngineCore
+from vllm.v1.engine.llm_engine import LLMEngine as V1LLMEngine
+from vllm.v1.executor.abstract import Executor
+
+
+@pytest.mark.parametrize("model", get_spyre_model_list())
+@pytest.mark.parametrize("backend", ["eager", "inductor"])
+def test_cb_handling(
+    model: str,
+    backend: str,
+):
+    """Test that the spyre worker correctly handles 
+    continuous batches of requests that
+    finish after different numbers of forward passes"""
+
+    max_tokens1 = 20
+    max_tokens2 = 17
+    max_tokens3 = 11
+
+    sampling_params1 = SamplingParams(max_tokens=max_tokens1,
+                                      temperature=0.0,
+                                      stop="1",
+                                      ignore_eos=True)
+
+    sampling_params2 = SamplingParams(max_tokens=max_tokens2,
+                                      temperature=0.0,
+                                      stop="1",
+                                      ignore_eos=True)
+
+    sampling_params3 = SamplingParams(max_tokens=max_tokens3,
+                                      temperature=0.0,
+                                      stop="1",
+                                      ignore_eos=True)
+
+    vllm_sampling_params = [
+        sampling_params1,
+        sampling_params2,
+        sampling_params3,
+    ]
+
+    # These prompts are ordered so that they don't finish in the
+    # order given
+    prompts = [
+        "7 6 5 4",
+        "10 9 8 7",
+        "8 7 6 5",
+    ]
+
+    # Ensure that both:
+    # - The model doesn't crash
+    # - The output sequences are correct
path: tests/test_spyre_cb.py
user: login: maxdebayser
body: Perhaps I misunderstood the purpose of this particular test but I think it duplicates this one: https://github.com/vllm-project/vllm-spyre/blob/7e52a8870d78b627ac1852633c6ec4d00621f269/tests/test_spyre_basic.py#L81

This other test can be parameterized to run with and without CB
created_at: 2025-04-10T13:17:41Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,282 @@
+"""Verification of continuous batching
+
+Run `python -m pytest tests/test_spyre_cb.py`.
+"""
+
+import time
+import uuid
+
+import pytest
+from spyre_util import generate_cb_spyre_vllm_output, get_spyre_model_list
+from transformers import AutoTokenizer
+from vllm import EngineArgs, SamplingParams
+from vllm.v1.engine import EngineCoreRequest
+from vllm.v1.engine.core import EngineCore
+from vllm.v1.engine.llm_engine import LLMEngine as V1LLMEngine
+from vllm.v1.executor.abstract import Executor
+
+
+@pytest.mark.parametrize("model", get_spyre_model_list())
+@pytest.mark.parametrize("backend", ["eager", "inductor"])
+def test_cb_handling(
+    model: str,
+    backend: str,
+):
+    """Test that the spyre worker correctly handles 
+    continuous batches of requests that
+    finish after different numbers of forward passes"""
+
+    max_tokens1 = 20
+    max_tokens2 = 17
+    max_tokens3 = 11
+
+    sampling_params1 = SamplingParams(max_tokens=max_tokens1,
+                                      temperature=0.0,
+                                      stop="1",
+                                      ignore_eos=True)
+
+    sampling_params2 = SamplingParams(max_tokens=max_tokens2,
+                                      temperature=0.0,
+                                      stop="1",
+                                      ignore_eos=True)
+
+    sampling_params3 = SamplingParams(max_tokens=max_tokens3,
+                                      temperature=0.0,
+                                      stop="1",
+                                      ignore_eos=True)
+
+    vllm_sampling_params = [
+        sampling_params1,
+        sampling_params2,
+        sampling_params3,
+    ]
+
+    # These prompts are ordered so that they don't finish in the
+    # order given
+    prompts = [
+        "7 6 5 4",
+        "10 9 8 7",
+        "8 7 6 5",
+    ]
+
+    # Ensure that both:
+    # - The model doesn't crash
+    # - The output sequences are correct
+    vllm_results = generate_cb_spyre_vllm_output(
+        model=model,
+        prompts=prompts,
+        max_model_len=2048,
+        block_size=2048,
+        sampling_params=vllm_sampling_params,
+        tensor_parallel_size=1,
+        backend=backend,
+    )
+
+    assert vllm_results[0]["text"] == " 3 2 "
+    assert vllm_results[1]["text"] == " 6 5 4 3 2 "
+    assert vllm_results[2]["text"] == " 4 3 2 "
+
+
+@pytest.mark.parametrize("model", get_spyre_model_list())
+@pytest.mark.parametrize("backend", ["eager", "inductor"])
+def test_cb_with_engine_steps(model: str, backend: str,
+                              monkeypatch: pytest.MonkeyPatch):
+    """Test that the spyre worker correctly handles 
+    continuous batches of requests and one sequence 
+    an exit the batch while a new ones gets prefilled 
+    and appended to the batch"""
+    with monkeypatch.context() as m:
+        max_tokens1 = 20
+        max_tokens2 = 17
+
+        sampling_params1 = SamplingParams(max_tokens=max_tokens1,
+                                          temperature=0.0,
+                                          stop="1",
+                                          ignore_eos=True)
+
+        sampling_params2 = SamplingParams(max_tokens=max_tokens2,
+                                          temperature=0.0,
+                                          stop="1",
+                                          ignore_eos=True)
+
+        prompt1 = "7 6 5 4"
+        prompt2 = "10 9 8 7"
+
+        # set env vars
+        m.setenv("VLLM_SPYRE_WARMUP_PROMPT_LENS", "64")
+        m.setenv("VLLM_SPYRE_WARMUP_NEW_TOKENS",
+                 str(max([max_tokens1, max_tokens2])))
+
+        m.setenv("VLLM_SPYRE_USE_CB", "1")
+        m.setenv("VLLM_USE_V1", "1")
+
+        m.setenv("VLLM_SPYRE_MAX_CONTEXT_LENGTH", "2048")
+        m.setenv("VLLM_SPYRE_MAX_BATCH_SIZE", str(2))
+        m.setenv("VLLM_SPYRE_DYNAMO_BACKEND", backend)
+
+        # start the engine
+        engine_args = EngineArgs(
+            model=model,
+            enforce_eager=True,  # reduce test time
+        )
+
+        engine = V1LLMEngine.from_engine_args(engine_args)
path: tests/test_spyre_cb.py
user: login: maxdebayser
body: Just thinking aloud: if you set `envs.VLLM_ENABLE_V1_MULTIPROCESSING = False`, `from_engine_args()` should create an engine that has an `InprocClient` for the engine core and then the engine step will be in lockstep with the engine core step. So you could combine both tests if you want. Although perhaps it would make the combined test a little harder to read.
created_at: 2025-04-10T13:28:32Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,282 @@
+"""Verification of continuous batching
+
+Run `python -m pytest tests/test_spyre_cb.py`.
+"""
+
+import time
+import uuid
+
+import pytest
+from spyre_util import generate_cb_spyre_vllm_output, get_spyre_model_list
+from transformers import AutoTokenizer
+from vllm import EngineArgs, SamplingParams
+from vllm.v1.engine import EngineCoreRequest
+from vllm.v1.engine.core import EngineCore
+from vllm.v1.engine.llm_engine import LLMEngine as V1LLMEngine
+from vllm.v1.executor.abstract import Executor
+
+
+@pytest.mark.parametrize("model", get_spyre_model_list())
+@pytest.mark.parametrize("backend", ["eager", "inductor"])
+def test_cb_handling(
+    model: str,
+    backend: str,
+):
+    """Test that the spyre worker correctly handles 
+    continuous batches of requests that
+    finish after different numbers of forward passes"""
+
+    max_tokens1 = 20
+    max_tokens2 = 17
+    max_tokens3 = 11
+
+    sampling_params1 = SamplingParams(max_tokens=max_tokens1,
+                                      temperature=0.0,
+                                      stop="1",
+                                      ignore_eos=True)
+
+    sampling_params2 = SamplingParams(max_tokens=max_tokens2,
+                                      temperature=0.0,
+                                      stop="1",
+                                      ignore_eos=True)
+
+    sampling_params3 = SamplingParams(max_tokens=max_tokens3,
+                                      temperature=0.0,
+                                      stop="1",
+                                      ignore_eos=True)
+
+    vllm_sampling_params = [
+        sampling_params1,
+        sampling_params2,
+        sampling_params3,
+    ]
+
+    # These prompts are ordered so that they don't finish in the
+    # order given
+    prompts = [
+        "7 6 5 4",
+        "10 9 8 7",
+        "8 7 6 5",
+    ]
+
+    # Ensure that both:
+    # - The model doesn't crash
+    # - The output sequences are correct
path: tests/test_spyre_cb.py
user: login: prashantgupta24
body: Yeah good point! 
created_at: 2025-04-10T16:09:24Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,282 @@
+"""Verification of continuous batching
+
+Run `python -m pytest tests/test_spyre_cb.py`.
+"""
+
+import time
+import uuid
+
+import pytest
+from spyre_util import generate_cb_spyre_vllm_output, get_spyre_model_list
+from transformers import AutoTokenizer
+from vllm import EngineArgs, SamplingParams
+from vllm.v1.engine import EngineCoreRequest
+from vllm.v1.engine.core import EngineCore
+from vllm.v1.engine.llm_engine import LLMEngine as V1LLMEngine
+from vllm.v1.executor.abstract import Executor
+
+
+@pytest.mark.parametrize("model", get_spyre_model_list())
+@pytest.mark.parametrize("backend", ["eager", "inductor"])
+def test_cb_handling(
+    model: str,
+    backend: str,
+):
+    """Test that the spyre worker correctly handles 
+    continuous batches of requests that
+    finish after different numbers of forward passes"""
+
+    max_tokens1 = 20
+    max_tokens2 = 17
+    max_tokens3 = 11
+
+    sampling_params1 = SamplingParams(max_tokens=max_tokens1,
+                                      temperature=0.0,
+                                      stop="1",
+                                      ignore_eos=True)
+
+    sampling_params2 = SamplingParams(max_tokens=max_tokens2,
+                                      temperature=0.0,
+                                      stop="1",
+                                      ignore_eos=True)
+
+    sampling_params3 = SamplingParams(max_tokens=max_tokens3,
+                                      temperature=0.0,
+                                      stop="1",
+                                      ignore_eos=True)
+
+    vllm_sampling_params = [
+        sampling_params1,
+        sampling_params2,
+        sampling_params3,
+    ]
+
+    # These prompts are ordered so that they don't finish in the
+    # order given
+    prompts = [
+        "7 6 5 4",
+        "10 9 8 7",
+        "8 7 6 5",
+    ]
+
+    # Ensure that both:
+    # - The model doesn't crash
+    # - The output sequences are correct
+    vllm_results = generate_cb_spyre_vllm_output(
+        model=model,
+        prompts=prompts,
+        max_model_len=2048,
+        block_size=2048,
+        sampling_params=vllm_sampling_params,
+        tensor_parallel_size=1,
+        backend=backend,
+    )
+
+    assert vllm_results[0]["text"] == " 3 2 "
+    assert vllm_results[1]["text"] == " 6 5 4 3 2 "
+    assert vllm_results[2]["text"] == " 4 3 2 "
+
+
+@pytest.mark.parametrize("model", get_spyre_model_list())
+@pytest.mark.parametrize("backend", ["eager", "inductor"])
+def test_cb_with_engine_steps(model: str, backend: str,
+                              monkeypatch: pytest.MonkeyPatch):
+    """Test that the spyre worker correctly handles 
+    continuous batches of requests and one sequence 
+    an exit the batch while a new ones gets prefilled 
+    and appended to the batch"""
+    with monkeypatch.context() as m:
+        max_tokens1 = 20
+        max_tokens2 = 17
+
+        sampling_params1 = SamplingParams(max_tokens=max_tokens1,
+                                          temperature=0.0,
+                                          stop="1",
+                                          ignore_eos=True)
+
+        sampling_params2 = SamplingParams(max_tokens=max_tokens2,
+                                          temperature=0.0,
+                                          stop="1",
+                                          ignore_eos=True)
+
+        prompt1 = "7 6 5 4"
+        prompt2 = "10 9 8 7"
path: tests/test_spyre_cb.py
user: login: prashantgupta24
body: Cool!
created_at: 2025-04-10T16:09:45Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,282 @@
+"""Verification of continuous batching
+
+Run `python -m pytest tests/test_spyre_cb.py`.
+"""
+
+import time
+import uuid
+
+import pytest
+from spyre_util import generate_cb_spyre_vllm_output, get_spyre_model_list
+from transformers import AutoTokenizer
+from vllm import EngineArgs, SamplingParams
+from vllm.v1.engine import EngineCoreRequest
+from vllm.v1.engine.core import EngineCore
+from vllm.v1.engine.llm_engine import LLMEngine as V1LLMEngine
+from vllm.v1.executor.abstract import Executor
+
+
+@pytest.mark.parametrize("model", get_spyre_model_list())
+@pytest.mark.parametrize("backend", ["eager", "inductor"])
+def test_cb_handling(
+    model: str,
+    backend: str,
+):
+    """Test that the spyre worker correctly handles 
+    continuous batches of requests that
+    finish after different numbers of forward passes"""
+
+    max_tokens1 = 20
+    max_tokens2 = 17
+    max_tokens3 = 11
+
+    sampling_params1 = SamplingParams(max_tokens=max_tokens1,
+                                      temperature=0.0,
+                                      stop="1",
+                                      ignore_eos=True)
+
+    sampling_params2 = SamplingParams(max_tokens=max_tokens2,
+                                      temperature=0.0,
+                                      stop="1",
+                                      ignore_eos=True)
+
+    sampling_params3 = SamplingParams(max_tokens=max_tokens3,
+                                      temperature=0.0,
+                                      stop="1",
+                                      ignore_eos=True)
+
+    vllm_sampling_params = [
+        sampling_params1,
+        sampling_params2,
+        sampling_params3,
+    ]
+
+    # These prompts are ordered so that they don't finish in the
+    # order given
+    prompts = [
+        "7 6 5 4",
+        "10 9 8 7",
+        "8 7 6 5",
+    ]
+
+    # Ensure that both:
+    # - The model doesn't crash
+    # - The output sequences are correct
+    vllm_results = generate_cb_spyre_vllm_output(
+        model=model,
+        prompts=prompts,
+        max_model_len=2048,
+        block_size=2048,
+        sampling_params=vllm_sampling_params,
+        tensor_parallel_size=1,
+        backend=backend,
+    )
+
+    assert vllm_results[0]["text"] == " 3 2 "
+    assert vllm_results[1]["text"] == " 6 5 4 3 2 "
+    assert vllm_results[2]["text"] == " 4 3 2 "
+
+
+@pytest.mark.parametrize("model", get_spyre_model_list())
+@pytest.mark.parametrize("backend", ["eager", "inductor"])
+def test_cb_with_engine_steps(model: str, backend: str,
+                              monkeypatch: pytest.MonkeyPatch):
+    """Test that the spyre worker correctly handles 
+    continuous batches of requests and one sequence 
+    an exit the batch while a new ones gets prefilled 
+    and appended to the batch"""
+    with monkeypatch.context() as m:
+        max_tokens1 = 20
+        max_tokens2 = 17
+
+        sampling_params1 = SamplingParams(max_tokens=max_tokens1,
+                                          temperature=0.0,
+                                          stop="1",
+                                          ignore_eos=True)
+
+        sampling_params2 = SamplingParams(max_tokens=max_tokens2,
+                                          temperature=0.0,
+                                          stop="1",
+                                          ignore_eos=True)
+
+        prompt1 = "7 6 5 4"
+        prompt2 = "10 9 8 7"
+
+        # set env vars
+        m.setenv("VLLM_SPYRE_WARMUP_PROMPT_LENS", "64")
+        m.setenv("VLLM_SPYRE_WARMUP_NEW_TOKENS",
+                 str(max([max_tokens1, max_tokens2])))
+
+        m.setenv("VLLM_SPYRE_USE_CB", "1")
+        m.setenv("VLLM_USE_V1", "1")
+
+        m.setenv("VLLM_SPYRE_MAX_CONTEXT_LENGTH", "2048")
+        m.setenv("VLLM_SPYRE_MAX_BATCH_SIZE", str(2))
+        m.setenv("VLLM_SPYRE_DYNAMO_BACKEND", backend)
+
+        # start the engine
+        engine_args = EngineArgs(
+            model=model,
+            enforce_eager=True,  # reduce test time
+        )
+
+        engine = V1LLMEngine.from_engine_args(engine_args)
path: tests/test_spyre_cb.py
user: login: prashantgupta24
body: Hmm, yeah it's hard for me already to wrap my head around that :) Let me wrap up the basic tests first then I can try to take a look at this if that works
created_at: 2025-04-10T16:10:41Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,282 @@
+"""Verification of continuous batching
+
+Run `python -m pytest tests/test_spyre_cb.py`.
+"""
+
+import time
+import uuid
+
+import pytest
+from spyre_util import generate_cb_spyre_vllm_output, get_spyre_model_list
+from transformers import AutoTokenizer
+from vllm import EngineArgs, SamplingParams
+from vllm.v1.engine import EngineCoreRequest
+from vllm.v1.engine.core import EngineCore
+from vllm.v1.engine.llm_engine import LLMEngine as V1LLMEngine
+from vllm.v1.executor.abstract import Executor
+
+
+@pytest.mark.parametrize("model", get_spyre_model_list())
+@pytest.mark.parametrize("backend", ["eager", "inductor"])
+def test_cb_handling(
+    model: str,
+    backend: str,
+):
+    """Test that the spyre worker correctly handles 
+    continuous batches of requests that
+    finish after different numbers of forward passes"""
+
+    max_tokens1 = 20
+    max_tokens2 = 17
+    max_tokens3 = 11
+
+    sampling_params1 = SamplingParams(max_tokens=max_tokens1,
+                                      temperature=0.0,
+                                      stop="1",
+                                      ignore_eos=True)
+
+    sampling_params2 = SamplingParams(max_tokens=max_tokens2,
+                                      temperature=0.0,
+                                      stop="1",
+                                      ignore_eos=True)
+
+    sampling_params3 = SamplingParams(max_tokens=max_tokens3,
+                                      temperature=0.0,
+                                      stop="1",
+                                      ignore_eos=True)
+
+    vllm_sampling_params = [
+        sampling_params1,
+        sampling_params2,
+        sampling_params3,
+    ]
+
+    # These prompts are ordered so that they don't finish in the
+    # order given
+    prompts = [
+        "7 6 5 4",
+        "10 9 8 7",
+        "8 7 6 5",
+    ]
+
+    # Ensure that both:
+    # - The model doesn't crash
+    # - The output sequences are correct
+    vllm_results = generate_cb_spyre_vllm_output(
+        model=model,
+        prompts=prompts,
+        max_model_len=2048,
+        block_size=2048,
+        sampling_params=vllm_sampling_params,
+        tensor_parallel_size=1,
+        backend=backend,
+    )
+
+    assert vllm_results[0]["text"] == " 3 2 "
+    assert vllm_results[1]["text"] == " 6 5 4 3 2 "
+    assert vllm_results[2]["text"] == " 4 3 2 "
+
+
+@pytest.mark.parametrize("model", get_spyre_model_list())
+@pytest.mark.parametrize("backend", ["eager", "inductor"])
+def test_cb_with_engine_steps(model: str, backend: str,
+                              monkeypatch: pytest.MonkeyPatch):
+    """Test that the spyre worker correctly handles 
+    continuous batches of requests and one sequence 
+    an exit the batch while a new ones gets prefilled 
+    and appended to the batch"""
+    with monkeypatch.context() as m:
+        max_tokens1 = 20
+        max_tokens2 = 17
+
+        sampling_params1 = SamplingParams(max_tokens=max_tokens1,
+                                          temperature=0.0,
+                                          stop="1",
+                                          ignore_eos=True)
+
+        sampling_params2 = SamplingParams(max_tokens=max_tokens2,
+                                          temperature=0.0,
+                                          stop="1",
+                                          ignore_eos=True)
+
+        prompt1 = "7 6 5 4"
+        prompt2 = "10 9 8 7"
path: tests/test_spyre_cb.py
user: login: prashantgupta24
body: Done
created_at: 2025-04-11T20:28:13Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,282 @@
+"""Verification of continuous batching
+
+Run `python -m pytest tests/test_spyre_cb.py`.
+"""
+
+import time
+import uuid
+
+import pytest
+from spyre_util import generate_cb_spyre_vllm_output, get_spyre_model_list
+from transformers import AutoTokenizer
+from vllm import EngineArgs, SamplingParams
+from vllm.v1.engine import EngineCoreRequest
+from vllm.v1.engine.core import EngineCore
+from vllm.v1.engine.llm_engine import LLMEngine as V1LLMEngine
+from vllm.v1.executor.abstract import Executor
+
+
+@pytest.mark.parametrize("model", get_spyre_model_list())
+@pytest.mark.parametrize("backend", ["eager", "inductor"])
+def test_cb_handling(
+    model: str,
+    backend: str,
+):
+    """Test that the spyre worker correctly handles 
+    continuous batches of requests that
+    finish after different numbers of forward passes"""
+
+    max_tokens1 = 20
+    max_tokens2 = 17
+    max_tokens3 = 11
+
+    sampling_params1 = SamplingParams(max_tokens=max_tokens1,
+                                      temperature=0.0,
+                                      stop="1",
+                                      ignore_eos=True)
+
+    sampling_params2 = SamplingParams(max_tokens=max_tokens2,
+                                      temperature=0.0,
+                                      stop="1",
+                                      ignore_eos=True)
+
+    sampling_params3 = SamplingParams(max_tokens=max_tokens3,
+                                      temperature=0.0,
+                                      stop="1",
+                                      ignore_eos=True)
+
+    vllm_sampling_params = [
+        sampling_params1,
+        sampling_params2,
+        sampling_params3,
+    ]
+
+    # These prompts are ordered so that they don't finish in the
+    # order given
+    prompts = [
+        "7 6 5 4",
+        "10 9 8 7",
+        "8 7 6 5",
+    ]
+
+    # Ensure that both:
+    # - The model doesn't crash
+    # - The output sequences are correct
path: tests/test_spyre_cb.py
user: login: prashantgupta24
body: Uh combining this test with the `test_spyre_basic` is proving to be tricky. There are a number of differences between the 2 tests:
1. Different backends (cb still doesn't support AIU)
1. CB doesn't support v0
1. `logprobs` within `SamplingParams` doesn't work with cb (still investigating)
created_at: 2025-04-14T22:12:25Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,282 @@
+"""Verification of continuous batching
+
+Run `python -m pytest tests/test_spyre_cb.py`.
+"""
+
+import time
+import uuid
+
+import pytest
+from spyre_util import generate_cb_spyre_vllm_output, get_spyre_model_list
+from transformers import AutoTokenizer
+from vllm import EngineArgs, SamplingParams
+from vllm.v1.engine import EngineCoreRequest
+from vllm.v1.engine.core import EngineCore
+from vllm.v1.engine.llm_engine import LLMEngine as V1LLMEngine
+from vllm.v1.executor.abstract import Executor
+
+
+@pytest.mark.parametrize("model", get_spyre_model_list())
+@pytest.mark.parametrize("backend", ["eager", "inductor"])
+def test_cb_handling(
+    model: str,
+    backend: str,
+):
+    """Test that the spyre worker correctly handles 
+    continuous batches of requests that
+    finish after different numbers of forward passes"""
+
+    max_tokens1 = 20
+    max_tokens2 = 17
+    max_tokens3 = 11
+
+    sampling_params1 = SamplingParams(max_tokens=max_tokens1,
+                                      temperature=0.0,
+                                      stop="1",
+                                      ignore_eos=True)
+
+    sampling_params2 = SamplingParams(max_tokens=max_tokens2,
+                                      temperature=0.0,
+                                      stop="1",
+                                      ignore_eos=True)
+
+    sampling_params3 = SamplingParams(max_tokens=max_tokens3,
+                                      temperature=0.0,
+                                      stop="1",
+                                      ignore_eos=True)
+
+    vllm_sampling_params = [
+        sampling_params1,
+        sampling_params2,
+        sampling_params3,
+    ]
+
+    # These prompts are ordered so that they don't finish in the
+    # order given
+    prompts = [
+        "7 6 5 4",
+        "10 9 8 7",
+        "8 7 6 5",
+    ]
+
+    # Ensure that both:
+    # - The model doesn't crash
+    # - The output sequences are correct
+    vllm_results = generate_cb_spyre_vllm_output(
+        model=model,
+        prompts=prompts,
+        max_model_len=2048,
+        block_size=2048,
+        sampling_params=vllm_sampling_params,
+        tensor_parallel_size=1,
+        backend=backend,
+    )
+
+    assert vllm_results[0]["text"] == " 3 2 "
+    assert vllm_results[1]["text"] == " 6 5 4 3 2 "
+    assert vllm_results[2]["text"] == " 4 3 2 "
+
+
+@pytest.mark.parametrize("model", get_spyre_model_list())
+@pytest.mark.parametrize("backend", ["eager", "inductor"])
+def test_cb_with_engine_steps(model: str, backend: str,
+                              monkeypatch: pytest.MonkeyPatch):
+    """Test that the spyre worker correctly handles 
+    continuous batches of requests and one sequence 
+    an exit the batch while a new ones gets prefilled 
+    and appended to the batch"""
+    with monkeypatch.context() as m:
+        max_tokens1 = 20
+        max_tokens2 = 17
+
+        sampling_params1 = SamplingParams(max_tokens=max_tokens1,
+                                          temperature=0.0,
+                                          stop="1",
+                                          ignore_eos=True)
+
+        sampling_params2 = SamplingParams(max_tokens=max_tokens2,
+                                          temperature=0.0,
+                                          stop="1",
+                                          ignore_eos=True)
+
+        prompt1 = "7 6 5 4"
+        prompt2 = "10 9 8 7"
+
+        # set env vars
+        m.setenv("VLLM_SPYRE_WARMUP_PROMPT_LENS", "64")
+        m.setenv("VLLM_SPYRE_WARMUP_NEW_TOKENS",
+                 str(max([max_tokens1, max_tokens2])))
+
+        m.setenv("VLLM_SPYRE_USE_CB", "1")
+        m.setenv("VLLM_USE_V1", "1")
+
+        m.setenv("VLLM_SPYRE_MAX_CONTEXT_LENGTH", "2048")
+        m.setenv("VLLM_SPYRE_MAX_BATCH_SIZE", str(2))
+        m.setenv("VLLM_SPYRE_DYNAMO_BACKEND", backend)
+
+        # start the engine
+        engine_args = EngineArgs(
+            model=model,
+            enforce_eager=True,  # reduce test time
+        )
+
+        engine = V1LLMEngine.from_engine_args(engine_args)
path: tests/test_spyre_cb.py
user: login: prashantgupta24
body: I understand it now. I will go ahead and do it using `InprocClient`. 
created_at: 2025-04-14T22:33:08Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,282 @@
+"""Verification of continuous batching
+
+Run `python -m pytest tests/test_spyre_cb.py`.
+"""
+
+import time
+import uuid
+
+import pytest
+from spyre_util import generate_cb_spyre_vllm_output, get_spyre_model_list
+from transformers import AutoTokenizer
+from vllm import EngineArgs, SamplingParams
+from vllm.v1.engine import EngineCoreRequest
+from vllm.v1.engine.core import EngineCore
+from vllm.v1.engine.llm_engine import LLMEngine as V1LLMEngine
+from vllm.v1.executor.abstract import Executor
+
+
+@pytest.mark.parametrize("model", get_spyre_model_list())
+@pytest.mark.parametrize("backend", ["eager", "inductor"])
+def test_cb_handling(
+    model: str,
+    backend: str,
+):
+    """Test that the spyre worker correctly handles 
+    continuous batches of requests that
+    finish after different numbers of forward passes"""
+
+    max_tokens1 = 20
+    max_tokens2 = 17
+    max_tokens3 = 11
+
+    sampling_params1 = SamplingParams(max_tokens=max_tokens1,
+                                      temperature=0.0,
+                                      stop="1",
+                                      ignore_eos=True)
+
+    sampling_params2 = SamplingParams(max_tokens=max_tokens2,
+                                      temperature=0.0,
+                                      stop="1",
+                                      ignore_eos=True)
+
+    sampling_params3 = SamplingParams(max_tokens=max_tokens3,
+                                      temperature=0.0,
+                                      stop="1",
+                                      ignore_eos=True)
+
+    vllm_sampling_params = [
+        sampling_params1,
+        sampling_params2,
+        sampling_params3,
+    ]
+
+    # These prompts are ordered so that they don't finish in the
+    # order given
+    prompts = [
+        "7 6 5 4",
+        "10 9 8 7",
+        "8 7 6 5",
+    ]
+
+    # Ensure that both:
+    # - The model doesn't crash
+    # - The output sequences are correct
+    vllm_results = generate_cb_spyre_vllm_output(
+        model=model,
+        prompts=prompts,
+        max_model_len=2048,
+        block_size=2048,
+        sampling_params=vllm_sampling_params,
+        tensor_parallel_size=1,
+        backend=backend,
+    )
+
+    assert vllm_results[0]["text"] == " 3 2 "
+    assert vllm_results[1]["text"] == " 6 5 4 3 2 "
+    assert vllm_results[2]["text"] == " 4 3 2 "
+
+
+@pytest.mark.parametrize("model", get_spyre_model_list())
+@pytest.mark.parametrize("backend", ["eager", "inductor"])
+def test_cb_with_engine_steps(model: str, backend: str,
+                              monkeypatch: pytest.MonkeyPatch):
+    """Test that the spyre worker correctly handles 
+    continuous batches of requests and one sequence 
+    an exit the batch while a new ones gets prefilled 
+    and appended to the batch"""
+    with monkeypatch.context() as m:
+        max_tokens1 = 20
+        max_tokens2 = 17
+
+        sampling_params1 = SamplingParams(max_tokens=max_tokens1,
+                                          temperature=0.0,
+                                          stop="1",
+                                          ignore_eos=True)
+
+        sampling_params2 = SamplingParams(max_tokens=max_tokens2,
+                                          temperature=0.0,
+                                          stop="1",
+                                          ignore_eos=True)
+
+        prompt1 = "7 6 5 4"
+        prompt2 = "10 9 8 7"
+
+        # set env vars
+        m.setenv("VLLM_SPYRE_WARMUP_PROMPT_LENS", "64")
+        m.setenv("VLLM_SPYRE_WARMUP_NEW_TOKENS",
+                 str(max([max_tokens1, max_tokens2])))
+
+        m.setenv("VLLM_SPYRE_USE_CB", "1")
+        m.setenv("VLLM_USE_V1", "1")
+
+        m.setenv("VLLM_SPYRE_MAX_CONTEXT_LENGTH", "2048")
+        m.setenv("VLLM_SPYRE_MAX_BATCH_SIZE", str(2))
+        m.setenv("VLLM_SPYRE_DYNAMO_BACKEND", backend)
+
+        # start the engine
+        engine_args = EngineArgs(
+            model=model,
+            enforce_eager=True,  # reduce test time
+        )
+
+        engine = V1LLMEngine.from_engine_args(engine_args)
+
+        # add first request
+        engine.add_request("0", prompt1, sampling_params1)
+
+        request_outputs = engine.step()
+        assert len(request_outputs) == 1  # only 1 request
+        assert request_outputs[0].request_id == "0"
+
+        request_outputs = engine.step()
+        assert len(request_outputs) == 1  # still only 1 request
+        assert request_outputs[0].request_id == "0"
+
+        # add another request
+        engine.add_request("1", prompt2, sampling_params2)
+
+        request_outputs = engine.step()
+        assert len(request_outputs) == 1  # still only 1 request
+        assert request_outputs[0].request_id == "0"
+
+        request_outputs = engine.step()
+        assert len(request_outputs) == 1  # still only 1 request
+        assert request_outputs[0].request_id == "1"  # req 1 is decoding
+
+        request_outputs = engine.step()
+        assert len(request_outputs) == 2  # both requests decoding now
+
+        request_outputs = engine.step()
+        assert len(request_outputs) == 2  # both requests decoding now
+
+        request_outputs = engine.step()
+        assert len(request_outputs) == 2  # both requests decoding now
+        assert request_outputs[0].finished  # request 1 is done
+        assert request_outputs[0].outputs[0].text == " 3 2 "
+
+        request_outputs = engine.step()
+        assert len(request_outputs) == 1  # only req #2 is left
+        assert request_outputs[0].request_id == "1"  # req 1 is decoding
+
+
+@pytest.mark.parametrize("model", get_spyre_model_list())
+@pytest.mark.parametrize("backend", ["eager", "inductor"])
+def test_cb_with_core_steps(model: str, backend: str,
+                            monkeypatch: pytest.MonkeyPatch):
+    """Test that the spyre worker correctly handles
+    continuous batches of requests and one sequence
+    an exit the batch while a new ones gets prefilled
+    and appended to the batch"""
+    with monkeypatch.context() as m:
+        max_tokens1 = 20
+        max_tokens2 = 17
+
+        sampling_params1 = SamplingParams(max_tokens=max_tokens1,
+                                          temperature=0.0,
+                                          stop="1",
+                                          ignore_eos=True)
+
+        sampling_params2 = SamplingParams(max_tokens=max_tokens2,
+                                          temperature=0.0,
+                                          stop="1",
+                                          ignore_eos=True)
+
+        prompt1 = "7 6 5 4"
+        prompt2 = "10 9 8 7"
+
+        # set env vars
+        m.setenv("VLLM_SPYRE_WARMUP_PROMPT_LENS", "64")
+        m.setenv("VLLM_SPYRE_WARMUP_NEW_TOKENS",
+                 str(max([max_tokens1, max_tokens2])))
+
+        m.setenv("VLLM_SPYRE_USE_CB", "1")
+        m.setenv("VLLM_USE_V1", "1")
+
+        m.setenv("VLLM_SPYRE_MAX_CONTEXT_LENGTH", "2048")
+        m.setenv("VLLM_SPYRE_MAX_BATCH_SIZE", str(2))
+        m.setenv("VLLM_SPYRE_DYNAMO_BACKEND", backend)
+
+        # start the engine
+        engine_args = EngineArgs(model=model, )
+
+        vllm_config = engine_args.create_engine_config()
+        executor_class = Executor.get_class(vllm_config)
+
+        engine_core = EngineCore(vllm_config=vllm_config,
+                                 executor_class=executor_class,
+                                 log_stats=True)
+        # engine = V1LLMEngine.from_engine_args(engine_args)
+
+        tokenizer = AutoTokenizer.from_pretrained(model)
+        prompt_tokens1 = tokenizer(prompt1).input_ids
+        prompt_tokens2 = tokenizer(prompt2).input_ids
+
+        # # add first request
+        engine_core.add_request(
+            EngineCoreRequest(
+                request_id=uuid.uuid4(),
+                prompt=prompt1,
+                prompt_token_ids=prompt_tokens1,
+                mm_inputs=None,
+                mm_hashes=None,
+                mm_placeholders=None,
+                sampling_params=sampling_params1,
+                eos_token_id=None,
+                arrival_time=time.time(),
+                lora_request=None,
+            ))
+        assert len(engine_core.scheduler.waiting) == 1
+        assert len(engine_core.scheduler.running) == 0
+
+        _ = engine_core.step()
+        assert len(engine_core.scheduler.waiting) == 0
+        assert len(engine_core.scheduler.running) == 1
+
+        engine_core.add_request(
+            EngineCoreRequest(
+                request_id=uuid.uuid4(),
+                prompt=prompt2,
+                prompt_token_ids=prompt_tokens2,
+                mm_inputs=None,
+                mm_hashes=None,
+                mm_placeholders=None,
+                sampling_params=sampling_params2,
+                eos_token_id=None,
+                arrival_time=time.time(),
+                lora_request=None,
+            ))
+        assert len(engine_core.scheduler.waiting) == 1
+        assert len(engine_core.scheduler.running) == 1
+
+        _ = engine_core.step()
+        assert len(engine_core.scheduler.waiting) == 0
+        assert len(engine_core.scheduler.running) == 1
+
+        _ = engine_core.step()
+        assert len(engine_core.scheduler.waiting) == 0
+        assert len(engine_core.scheduler.running) == 2
+
+        _ = engine_core.step()
+        assert len(engine_core.scheduler.waiting) == 0
+        assert len(engine_core.scheduler.running) == 2
+
+        _ = engine_core.step()
+        assert len(engine_core.scheduler.waiting) == 0
+        assert len(engine_core.scheduler.running) == 2
+
+        _ = engine_core.step()
+        assert len(engine_core.scheduler.waiting) == 0
+        assert len(engine_core.scheduler.running) == 2
+
+        _ = engine_core.step()
+        assert len(engine_core.scheduler.waiting) == 0
+        assert len(engine_core.scheduler.running) == 2
+
+        _ = engine_core.step()
+        assert len(engine_core.scheduler.waiting) == 0
+        assert len(engine_core.scheduler.running) == 2
+
+        _ = engine_core.step()
+        assert len(engine_core.scheduler.waiting) == 0
+        assert len(engine_core.scheduler.running) == 2
path: tests/test_spyre_cb.py
user: login: prashantgupta24
body: " Stop strings require detokenizing the output. Detokenization is handled in the output processor that runs at the AsyncLLM and LLMEngine level
So, we stream tokens from EngineCore -> AsyncLlm. AsyncLLM checks for stop strings in the OutputProcessor class. If a stop string is detected, we stop generating and send an abort message back to EngineCore to terminate the request in the EngineCore " 

This is why the request wasn't stopping using the `stop` string
created_at: 2025-04-14T22:33:37Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,282 @@
+"""Verification of continuous batching
+
+Run `python -m pytest tests/test_spyre_cb.py`.
+"""
+
+import time
+import uuid
+
+import pytest
+from spyre_util import generate_cb_spyre_vllm_output, get_spyre_model_list
+from transformers import AutoTokenizer
+from vllm import EngineArgs, SamplingParams
+from vllm.v1.engine import EngineCoreRequest
+from vllm.v1.engine.core import EngineCore
+from vllm.v1.engine.llm_engine import LLMEngine as V1LLMEngine
+from vllm.v1.executor.abstract import Executor
+
+
+@pytest.mark.parametrize("model", get_spyre_model_list())
+@pytest.mark.parametrize("backend", ["eager", "inductor"])
+def test_cb_handling(
+    model: str,
+    backend: str,
+):
+    """Test that the spyre worker correctly handles 
+    continuous batches of requests that
+    finish after different numbers of forward passes"""
+
+    max_tokens1 = 20
+    max_tokens2 = 17
+    max_tokens3 = 11
+
+    sampling_params1 = SamplingParams(max_tokens=max_tokens1,
+                                      temperature=0.0,
+                                      stop="1",
+                                      ignore_eos=True)
+
+    sampling_params2 = SamplingParams(max_tokens=max_tokens2,
+                                      temperature=0.0,
+                                      stop="1",
+                                      ignore_eos=True)
+
+    sampling_params3 = SamplingParams(max_tokens=max_tokens3,
+                                      temperature=0.0,
+                                      stop="1",
+                                      ignore_eos=True)
+
+    vllm_sampling_params = [
+        sampling_params1,
+        sampling_params2,
+        sampling_params3,
+    ]
+
+    # These prompts are ordered so that they don't finish in the
+    # order given
+    prompts = [
+        "7 6 5 4",
+        "10 9 8 7",
+        "8 7 6 5",
+    ]
+
+    # Ensure that both:
+    # - The model doesn't crash
+    # - The output sequences are correct
path: tests/test_spyre_cb.py
user: login: prashantgupta24
body: If we can remove `logprobs` from the SamplingParams for `tests/test_spyre_basic.py` (which I think is fine since we're not really doing anything with them), we can reuse the tests. But I'm also curious as to why it's not working with cb
created_at: 2025-04-14T23:07:25Z
author_association: COLLABORATOR
path: tests/e2e/test_spyre_basic.py
user: login: prashantgupta24
body: This test has been moved to `test_spyre_cb` since there essentially tests the same things but both with CB and static.
created_at: 2025-05-02T18:39:37Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,238 @@
+"""Verification of continuous batching
+
+Run `python -m pytest tests/test_spyre_cb.py`.
+"""
+
+import pytest
+from spyre_util import generate_cb_spyre_vllm_output, get_spyre_model_list
+from transformers import AutoTokenizer
+from vllm import EngineArgs, SamplingParams
+from vllm.v1.engine.llm_engine import LLMEngine as V1LLMEngine
+
+
+@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(
+    "prompts",
+    [
+        [
+            "7 6 5 4",
+            "10 9 8 7",
+        ],
+        [
+            "7 6 5 4",
+            "10 9 8 7",
+            "8 7 6 5",
+        ],
+        [
+            "7 6 5 4",
+            "10 9 8 7",
+            "8 7 6 5",
+            "9 8 7 6",
+        ],
+    ],
+    ids=lambda val: f"num_prompts({len(val)})",
+)
+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
+    continuous batches of requests that
+    finish after different numbers of forward passes"""
+
+    if cb and len(prompts) < max_num_seqs:
+        pytest.skip("Skipping since this doesn't really test CB")
path: tests/e2e/test_spyre_cb.py
user: login: prashantgupta24
body: Let me know if that doesn't make sense, I believe it could be an overkill to test these scenarios
created_at: 2025-05-12T20:07:55Z
author_association: COLLABORATOR
diff_hunk: @@ -55,6 +55,9 @@ jobs:
           - name: "utils"
             markers: "utils"
             flags: "--timeout=300"
+          - name: "cb"
+            markers: "cb"
+            flags: "--timeout=300 --forked"
path: .github/workflows/test-spyre.yml
user: login: wallashss
body: Do we know if there's order guaranteed here? I had experienced errors when switching back to the static batching with the fms branch installed. 
created_at: 2025-05-12T20:09:57Z
author_association: COLLABORATOR
diff_hunk: @@ -55,6 +55,9 @@ jobs:
           - name: "utils"
             markers: "utils"
             flags: "--timeout=300"
+          - name: "cb"
+            markers: "cb"
+            flags: "--timeout=300 --forked"
path: .github/workflows/test-spyre.yml
user: login: prashantgupta24
body: It's an isolated job on a different runner so hopefully order shouldn't matter?
created_at: 2025-05-12T20:15:57Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,238 @@
+"""Verification of continuous batching
+
+Run `python -m pytest tests/test_spyre_cb.py`.
+"""
+
+import pytest
+from spyre_util import generate_cb_spyre_vllm_output, get_spyre_model_list
+from transformers import AutoTokenizer
+from vllm import EngineArgs, SamplingParams
+from vllm.v1.engine.llm_engine import LLMEngine as V1LLMEngine
+
+
+@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(
+    "prompts",
+    [
+        [
+            "7 6 5 4",
+            "10 9 8 7",
+        ],
+        [
+            "7 6 5 4",
+            "10 9 8 7",
+            "8 7 6 5",
+        ],
+        [
+            "7 6 5 4",
+            "10 9 8 7",
+            "8 7 6 5",
+            "9 8 7 6",
+        ],
+    ],
+    ids=lambda val: f"num_prompts({len(val)})",
+)
+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
+    continuous batches of requests that
+    finish after different numbers of forward passes"""
+
+    if cb and len(prompts) < max_num_seqs:
+        pytest.skip("Skipping since this doesn't really test CB")
path: tests/e2e/test_spyre_cb.py
user: login: wallashss
body: IMHO It does make sense to remove the if and let it run. I found already cases that the data is not properly cleaned because it was expecting the engine has new requests process them. It might be interesting to find those cases if we do not properly set the state of system consistent. 

In fact, I was thinking about to duplicate most of the tests by turn on CB and expect the correct behavior of the system. However, due to the limitation of the branches of fms it will not be not so easy. I think we should consider this later.
created_at: 2025-05-12T20:21:19Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,238 @@
+"""Verification of continuous batching
+
+Run `python -m pytest tests/test_spyre_cb.py`.
+"""
+
+import pytest
+from spyre_util import generate_cb_spyre_vllm_output, get_spyre_model_list
+from transformers import AutoTokenizer
+from vllm import EngineArgs, SamplingParams
+from vllm.v1.engine.llm_engine import LLMEngine as V1LLMEngine
+
+
+@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(
+    "prompts",
+    [
+        [
+            "7 6 5 4",
+            "10 9 8 7",
+        ],
+        [
+            "7 6 5 4",
+            "10 9 8 7",
+            "8 7 6 5",
+        ],
+        [
+            "7 6 5 4",
+            "10 9 8 7",
+            "8 7 6 5",
+            "9 8 7 6",
+        ],
+    ],
+    ids=lambda val: f"num_prompts({len(val)})",
+)
+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
+    continuous batches of requests that
+    finish after different numbers of forward passes"""
+
+    if cb and len(prompts) < max_num_seqs:
+        pytest.skip("Skipping since this doesn't really test CB")
path: tests/e2e/test_spyre_cb.py
user: login: prashantgupta24
body: > I was thinking about to duplicate most of the tests by turn on CB and expect the correct behavior of the system

Interesting point! Yeah we can definitely see about that later. For now I can revert the `if` statement
created_at: 2025-05-12T20:26:46Z
author_association: COLLABORATOR
diff_hunk: @@ -55,6 +55,9 @@ jobs:
           - name: "utils"
             markers: "utils"
             flags: "--timeout=300"
+          - name: "cb"
+            markers: "cb"
+            flags: "--timeout=300 --forked"
path: .github/workflows/test-spyre.yml
user: login: wallashss
body: > It's an isolated job on a different runner so hopefully order shouldn't matter?

If it is isolated env, then it's fine, the order does not matter. Tks.
created_at: 2025-05-12T20:27:14Z
author_association: COLLABORATOR

PR comments:
user: login: github-actions[bot]
created_at: 2025-04-09T20:30:33Z
body: 👋 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 🚀
author_association: NONE
user: login: maxdebayser
created_at: 2025-04-10T13:29:05Z
body: Nice, this is going in the right direction.
author_association: COLLABORATOR
user: login: prashantgupta24
created_at: 2025-04-14T23:00:29Z
body: @maxdebayser there's something funky going on with cb, the basic test with 4 prompts seems to be failing for `max_num_seqs` 2 and 3:
```
test_spyre_cb.py::test_cb_handling[eager-/models/llama-194m-1] PASSED                                                                                                                                                                  [ 25%]
test_spyre_cb.py::test_cb_handling[eager-/models/llama-194m-2] FAILED                                                                                                                                                                  [ 50%]
test_spyre_cb.py::test_cb_handling[eager-/models/llama-194m-3] FAILED                                                                                                                                                                  [ 75%]
test_spyre_cb.py::test_cb_handling[eager-/models/llama-194m-4] PASSED
```
```
FAILED test_spyre_cb.py::test_cb_handling[eager-/models/llama-194m-2] - AssertionError: assert ' 6 5 4 8 7 8 7 6 5 4' == ' 6 5 4 3 2 '

  -  6 5 4 3 2
  +  6 5 4 8 7 8 7 6 5 4
FAILED test_spyre_cb.py::test_cb_handling[eager-/models/llama-194m-3] - AssertionError: assert ' 6 5 4 9 8 7 6 5 4 3' == ' 6 5 4 3 2 '

  -  6 5 4 3 2
  +  6 5 4 9 8 7 6 5 4 3
```

whereas running the same test without cb is successful

```

test_spyre_cb.py::test_cb_handling[eager-/models/llama-194m-1] PASSED                                                                                                                                                                  [ 12%]
test_spyre_cb.py::test_cb_handling[eager-/models/llama-194m-2] PASSED                                                                                                                                                                  [ 25%]
test_spyre_cb.py::test_cb_handling[eager-/models/llama-194m-3] PASSED                                                                                                                                                                  [ 37%]
test_spyre_cb.py::test_cb_handling[eager-/models/llama-194m-4] PASSED
```
author_association: COLLABORATOR
user: login: maxdebayser
created_at: 2025-04-15T13:55:45Z
body: Since there are 4 requests, max_seq_lens 1 and 4 basically run without exercising the CB logic. So 2 and 3 seem to point to a problem in this logic. Perhaps it could be a mask that is wrong. cc: @yannicks1 
author_association: COLLABORATOR
user: login: prashantgupta24
created_at: 2025-04-15T18:33:21Z
body: Yeah this is a pattern, 3 prompts with 2 `max_num_seqs` and 4 prompts with 2 or 3 `max_num_seqs` are failing

```
FAILED test_spyre_cb.py::test_cb_handling[prompts2-1-eager-/models/llama-194m-2] - AssertionError: assert ' 6 5 4 8 7 6 5 4 3 2' == ' 6 5 4 3 2 '
FAILED test_spyre_cb.py::test_cb_handling[prompts3-1-eager-/models/llama-194m-2] - AssertionError: assert ' 6 5 4 8 7 8 7 6 5 4' == ' 6 5 4 3 2 '
FAILED test_spyre_cb.py::test_cb_handling[prompts3-1-eager-/models/llama-194m-3] - AssertionError: assert ' 6 5 4 9 8 7 6 5 4 3' == ' 6 5 4 3 2 '
```
author_association: COLLABORATOR
user: login: prashantgupta24
created_at: 2025-04-16T22:29:23Z
body: Spent almost the whole day today unsuccessfully trying to set up a local working env, ugh! Hard to do on an M1. Unfortunately debugging the lines one by one gives me the correct result for some reason. 

Update: got a local env working on my M1!
author_association: COLLABORATOR
user: login: prashantgupta24
created_at: 2025-04-29T22:46:15Z
body: Alright all tests pass when run against the new CB code with fms branch [paged_attn_mock](https://github.com/foundation-model-stack/foundation-model-stack/tree/paged_attn_mock)
```
tests/e2e/test_spyre_cb.py::test_cb_handling[num_prompts(1)-v1-use_cb(1)-eager-llama-194m-max_num_seqs(1)] PASSED
tests/e2e/test_spyre_cb.py::test_cb_handling[num_prompts(1)-v1-use_cb(1)-eager-llama-194m-max_num_seqs(2)] PASSED
tests/e2e/test_spyre_cb.py::test_cb_handling[num_prompts(1)-v1-use_cb(1)-eager-llama-194m-max_num_seqs(3)] PASSED
tests/e2e/test_spyre_cb.py::test_cb_handling[num_prompts(1)-v1-use_cb(1)-eager-llama-194m-max_num_seqs(4)] PASSED
tests/e2e/test_spyre_cb.py::test_cb_handling[num_prompts(2)-v1-use_cb(1)-eager-llama-194m-max_num_seqs(1)] PASSED
tests/e2e/test_spyre_cb.py::test_cb_handling[num_prompts(2)-v1-use_cb(1)-eager-llama-194m-max_num_seqs(2)] PASSED
tests/e2e/test_spyre_cb.py::test_cb_handling[num_prompts(2)-v1-use_cb(1)-eager-llama-194m-max_num_seqs(3)] PASSED
tests/e2e/test_spyre_cb.py::test_cb_handling[num_prompts(2)-v1-use_cb(1)-eager-llama-194m-max_num_seqs(4)] PASSED
tests/e2e/test_spyre_cb.py::test_cb_handling[num_prompts(3)-v1-use_cb(1)-eager-llama-194m-max_num_seqs(1)] PASSED
tests/e2e/test_spyre_cb.py::test_cb_handling[num_prompts(3)-v1-use_cb(1)-eager-llama-194m-max_num_seqs(2)] PASSED
tests/e2e/test_spyre_cb.py::test_cb_handling[num_prompts(3)-v1-use_cb(1)-eager-llama-194m-max_num_seqs(3)] PASSED
tests/e2e/test_spyre_cb.py::test_cb_handling[num_prompts(3)-v1-use_cb(1)-eager-llama-194m-max_num_seqs(4)] PASSED
tests/e2e/test_spyre_cb.py::test_cb_handling[num_prompts(4)-v1-use_cb(1)-eager-llama-194m-max_num_seqs(1)] PASSED
tests/e2e/test_spyre_cb.py::test_cb_handling[num_prompts(4)-v1-use_cb(1)-eager-llama-194m-max_num_seqs(2)] PASSED
tests/e2e/test_spyre_cb.py::test_cb_handling[num_prompts(4)-v1-use_cb(1)-eager-llama-194m-max_num_seqs(3)] PASSED
tests/e2e/test_spyre_cb.py::test_cb_handling[num_prompts(4)-v1-use_cb(1)-eager-llama-194m-max_num_seqs(4)] PASSED
```
author_association: COLLABORATOR
user: login: wallashss
created_at: 2025-05-09T16:51:41Z
body: What is the status of this PR? Do we have blockers? These tests would be very helpful for the development of CB.
author_association: COLLABORATOR
user: login: prashantgupta24
created_at: 2025-05-09T17:02:28Z
body: > What is the status of this PR? Do we have blockers? These tests would be very helpful for the development of CB.

The tests on this PR work well they just need the custom branch from FMS
author_association: COLLABORATOR
user: login: wallashss
created_at: 2025-05-09T17:26:20Z
body: > The tests on this PR work well they just need the custom branch from FMS

Yeah, you are right... These tests passing will break the others because of the divergence of FMS version. I had this problem before. 

I was wondering if we could merge this as-is and just skip them on CI for now, and we activate later again once we solve the FMS issue. I see value to have this test merged ASAP and I think we will be able to iterate the CB feature faster. We could make a follow up issue for this. What do you think?

Also cc @joerunde 
author_association: COLLABORATOR
user: login: prashantgupta24
created_at: 2025-05-09T17:43:45Z
body: > > The tests on this PR work well they just need the custom branch from FMS
> 
> Yeah, you are right... These tests passing will break the others because of the divergence of FMS version. I had this problem before.
> 
> I was wondering if we could merge this as-is and just skip them on CI for now, and we activate later again once we solve the FMS issue. I see value to have this test merged ASAP and I think we will be able to iterate the CB feature faster. We could make a follow up issue for this. What do you think?
> 
> Also cc @joerunde

Yeah I think we talked about this in one of the stand-ups. Worth trying it out I guess? Another option is to actually install the custom FMS branch in the CI when we actually run this test, although I'm not sure how tricky that will be. Honestly the tests are well contained so it's easy to copy them over whenever you want to test it out for now (which is what I'm doing right now)
author_association: COLLABORATOR
user: login: wallashss
created_at: 2025-05-09T18:39:14Z
body: > Another option is to actually install the custom FMS branch in the CI when we actually run this test, although I'm not sure how tricky that will be

Good idea! Could be, pretty easy !

```
pip install git+https://github.com/foundation-model-stack/foundation-model-stack.git@paged_attn_mock
```

I'm okay with moving forward like this.


author_association: COLLABORATOR
user: login: prashantgupta24
created_at: 2025-05-09T18:57:16Z
body: > > Another option is to actually install the custom FMS branch in the CI when we actually run this test, although I'm not sure how tricky that will be
> 
> Good idea! Could be, pretty easy !
> 
> ```
> pip install git+https://github.com/foundation-model-stack/foundation-model-stack.git@paged_attn_mock
> ```
> 
> I'm okay with moving forward like this.

Yeah, sounds good! Probably better to open a separate story for that. Let me do that now.
author_association: COLLABORATOR
user: login: prashantgupta24
created_at: 2025-05-09T23:43:28Z
body: Alright this PR installs the custom FMS branch! Let me fix up the tests next week then maybe we can merge it
author_association: COLLABORATOR