diff --git a/examples/vllm/README.md b/examples/vllm/README.md index b053e781..2b999a9e 100644 --- a/examples/vllm/README.md +++ b/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`. diff --git a/examples/vllm/engine.py b/examples/vllm/engine.py new file mode 100644 index 00000000..910d46de --- /dev/null +++ b/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 diff --git a/examples/vllm/main.py b/examples/vllm/main.py index b75c46f9..780e6208 100644 --- a/examples/vllm/main.py +++ b/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() diff --git a/examples/vllm/pyproject.toml b/examples/vllm/pyproject.toml index fcd7f86d..7dffe5d3 100644 --- a/examples/vllm/pyproject.toml +++ b/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 = [