examples: split vLLM generation across prefill and decode sessions

examples/vllm/README.md

@@ -1,14 +1,22 @@
 # vLLM Example
 
-This example verifies that a Python environment can start a small vLLM model from Flame's example tree.
+This example uses `flamepy.runner.Runner` to run a small vLLM model across two sessions: one prefill engine and one decode engine.
+
+Each `rr.service(...)` call creates a Flame session. Object instances are stateful, so each engine loads the model once on its executor. Prefill returns a `PrefillOutput` object (prompt token ids plus the first generated token). Decode continues from that object via `TokensPrompt`. This is a session-level prefill/decode split, not vLLM KV-connector disaggregation; decode recomputes KV on the same small model.
+
+Requires a Flame cluster and a runtime that can start vLLM. Override the model with `VLLM_MODEL` (default `facebook/opt-125m`).
 
 ## Run
 
-From the Flame console or a local development shell with the example dependencies installed:
+From the Flame console:
 
 ```bash
 cd /opt/examples/vllm
 uv run main.py
 ```
 
-The script loads `facebook/opt-125m` with vLLM. It is intended as a minimal smoke test for vLLM dependencies and GPU/runtime setup, not as a registered Flame service.
+## Files
+
+- `main.py`: Opens a Runner, creates prefill and decode services, prints the completion.
+- `engine.py`: Stateful `PrefillEngine` and `DecodeEngine` used as Runner services.
+- `pyproject.toml`: Package dependencies (`vllm`). The Flame environment provides `flamepy`.

examples/vllm/engine.py

@@ -0,0 +1,52 @@
+import os
+from dataclasses import dataclass, field
+
+from flamepy.runner import SessionContext
+from vllm import LLM, SamplingParams
+from vllm.inputs import TokensPrompt
+
+DEFAULT_MODEL = "facebook/opt-125m"
+
+
+def _model_name() -> str:
+    return os.getenv("VLLM_MODEL", DEFAULT_MODEL)
+
+
+@dataclass
+class PrefillOutput:
+    prompt_token_ids: list[int]
+    first_token_ids: list[int] = field(default_factory=list)
+
+
+class _VllmEngine:
+    def __init__(self):
+        self._llm = None
+
+    def _engine(self) -> LLM:
+        if self._llm is None:
+            self._llm = LLM(model=_model_name())
+        return self._llm
+
+
+class PrefillEngine(_VllmEngine):
+    _session_context = SessionContext(session_id="vllm-prefill")
+
+    def prefill(self, prompt: str) -> PrefillOutput:
+        outputs = self._engine().generate([prompt], SamplingParams(max_tokens=1))
+        output = outputs[0]
+        return PrefillOutput(
+            prompt_token_ids=list(output.prompt_token_ids),
+            first_token_ids=list(output.outputs[0].token_ids),
+        )
+
+
+class DecodeEngine(_VllmEngine):
+    _session_context = SessionContext(session_id="vllm-decode")
+
+    def decode(self, prefill: PrefillOutput, max_tokens: int = 16) -> str:
+        token_ids = list(prefill.prompt_token_ids) + list(prefill.first_token_ids)
+        outputs = self._engine().generate(
+            [TokensPrompt(prompt_token_ids=token_ids)],
+            SamplingParams(max_tokens=max_tokens),
+        )
+        return outputs[0].outputs[0].text

examples/vllm/main.py

@@ -1,10 +1,16 @@
-from vllm import LLM
+from engine import DecodeEngine, PrefillEngine
+from flamepy.runner import Runner
+
 
 def main():
-    print("Hello from vllm!")
+    with Runner("vllm-example") as rr:
+        prefill = rr.service(PrefillEngine())
+        decode = rr.service(DecodeEngine())
+
+        out = prefill.prefill("Once upon a time")
+        text = decode.decode(out, 16).get()
+        print(text)
 
-    # Initialize the vLLM engine.
-    llm = LLM(model="facebook/opt-125m")
 
 if __name__ == "__main__":
     main()

examples/vllm/pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 name = "flame-vllm"
 version = "0.2.0"
-description = "Add your description here"
+description = "vLLM prefill and decode sessions by flamepy.Runner"
 readme = "README.md"
 requires-python = ">=3.13"
 dependencies = [