diff --git a/src/exo/worker/engines/vllm/vllm_generator.py b/src/exo/worker/engines/vllm/vllm_generator.py
index 5f06450d..ab9bb947 100644
--- a/src/exo/worker/engines/vllm/vllm_generator.py
+++ b/src/exo/worker/engines/vllm/vllm_generator.py
@@ -1,11 +1,9 @@
import gc
-import itertools
import json
import re
import sys
import time
-from collections import deque
-from collections.abc import Callable, Generator, Iterable
+from collections.abc import Callable, Generator
from dataclasses import dataclass, field
from pathlib import Path
@@ -20,16 +18,10 @@ from exo.shared.types.api import (
PromptTokensDetails,
Usage,
)
-from exo.shared.types.chunks import PrefillProgressChunk
from exo.shared.types.common import ModelId
-from exo.shared.types.events import ChunkGenerated, Event
from exo.shared.types.memory import Memory
-from exo.shared.types.tasks import CANCEL_ALL_TASKS, TaskId, TextGeneration
-from exo.shared.types.worker.runner_response import (
- GenerationResponse,
- ToolCallResponse,
-)
-from exo.utils.channels import MpReceiver, MpSender
+from exo.shared.types.text_generation import TextGenerationTaskParams
+from exo.shared.types.worker.runner_response import GenerationResponse
from exo.worker.engines.kv_cache import TorchKVCache
from exo.worker.engines.mlx.cache import KVPrefixCache
from exo.worker.engines.vllm.growable_cache import (
@@ -42,13 +34,6 @@ from exo.worker.engines.vllm.prompt_format import (
make_vllm_sampling_params,
)
from exo.worker.runner.bootstrap import logger
-from exo.worker.runner.llm_inference.batch_generator import (
- Cancelled,
- Finished,
- GeneratorQueue,
- InferenceGenerator,
-)
-from exo.worker.runner.llm_inference.model_output_parsers import apply_vllm_parsers
from exo.worker.runner.llm_inference.tool_parsers import ToolParser, infer_tool_parser
@@ -66,628 +51,317 @@ def _build_layer_groups(kv_cache_config: object) -> list[int]:
@dataclass
-class _ActiveRequest:
- task: TextGeneration
+class _EngineRequest:
+ uid: int
request_id: str
prompt_token_count: int
prompt_token_ids: list[int]
- queue: GeneratorQueue[GenerationResponse]
- parsed_gen: Generator[GenerationResponse | ToolCallResponse | None]
prefill_done: bool = False
prefill_steps: int = 0
prev_text: str = ""
prev_token_count: int = 0
start_time: float = field(default_factory=time.perf_counter)
first_token_time: float | None = None
+ on_generation_token: Callable[[], None] | None = None
+ on_prefill_progress: Callable[[int, int], None] | None = None
-@dataclass(eq=False)
-class VllmSequentialGenerator(InferenceGenerator):
- engine: LLMEngine
- model_id: ModelId
- tool_parser: ToolParser | None
- cancel_receiver: MpReceiver[TaskId]
- event_sender: MpSender[Event]
- prefix_cache: KVPrefixCache
+def _save_prefix_cache(
+ engine: LLMEngine,
+ prefix_cache: KVPrefixCache,
+ request_id: str,
+ prompt_token_ids: list[int],
+ prompt_token_count: int,
+) -> None:
+ try:
+ coordinator = None
+ model_runner = _growable_model_runner_ref[0]
+ kv_cache_config = None
+ try:
+ engine_core = engine.engine_core.engine_core # type: ignore
+ coordinator = engine_core.scheduler.kv_cache_manager.coordinator # type: ignore
+ kv_cache_config = engine_core.scheduler.kv_cache_manager.kv_cache_config # type: ignore
+ except Exception:
+ pass
+ if coordinator is None or model_runner is None or kv_cache_config is None:
+ return
- _cancelled_tasks: set[TaskId] = field(default_factory=set, init=False)
- _all_tasks: dict[TaskId, TextGeneration] = field(default_factory=dict, init=False)
- _maybe_queue: list[TextGeneration] = field(default_factory=list, init=False)
- _queue: deque[TextGeneration] = field(default_factory=deque, init=False)
- _maybe_cancel: list[TextGeneration] = field(default_factory=list, init=False)
- _active: _ActiveRequest | None = field(default=None, init=False)
-
- def warmup(self) -> None:
- tokenizer = self.engine.get_tokenizer()
- messages = [{"role": "user", "content": "Prompt to warm up the inference engine. Repeat this."}]
- prompt_text: str = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) # type: ignore
- token_ids: list[int] = tokenizer.encode(prompt_text, add_special_tokens=False) # type: ignore
- params = SamplingParams(max_tokens=50, detokenize=False)
- self.engine.add_request("warmup", {"prompt_token_ids": token_ids}, params)
- tokens_generated = 0
- while self.engine.has_unfinished_requests():
- self.engine.step()
- tokens_generated += 1
- logger.info(f"vLLM warmup complete, generated {tokens_generated} tokens")
-
- def submit(self, task: TextGeneration) -> None:
- self._cancelled_tasks.discard(CANCEL_ALL_TASKS)
- self._all_tasks[task.task_id] = task
- self._maybe_queue.append(task)
-
- def agree_on_tasks(self) -> None:
- self._queue.extend(self._maybe_queue)
- self._maybe_queue.clear()
-
- def agree_on_cancellations(self) -> None:
- for task_id in self.cancel_receiver.collect():
- if task_id == CANCEL_ALL_TASKS:
- self._cancelled_tasks.add(CANCEL_ALL_TASKS)
- else:
- if task_id in self._all_tasks:
- self._maybe_cancel.append(self._all_tasks[task_id])
- self._cancelled_tasks.update(task.task_id for task in self._maybe_cancel)
- self._maybe_cancel.clear()
-
- def step(
- self,
- ) -> Iterable[
- tuple[TaskId, GenerationResponse | ToolCallResponse | Cancelled | Finished]
- ]:
- self.agree_on_cancellations()
-
- if self._active is None and not self._queue:
- self.agree_on_tasks()
-
- if self._active is None and not self._queue:
- return []
-
- tokenizer = self.engine.get_tokenizer()
- think_start: str | None = getattr(tokenizer, "think_start", None)
- think_end: str | None = getattr(tokenizer, "think_end", None)
-
- if self._active is None:
- task = self._queue.popleft()
- if self.should_cancel(task.task_id):
- self._cancelled_tasks.discard(task.task_id)
- return [(task.task_id, Cancelled())]
- token_ids, prompt_text, prompt_token_count = format_vllm_prompt(
- self.engine, task.task_params
- )
- logger.info(prompt_text)
- request_id = str(task.task_id)
- sampling_params = make_vllm_sampling_params(self.engine, task.task_params, self.model_id)
- self.engine.add_request(
- request_id,
- {"prompt_token_ids": token_ids},
- sampling_params,
- )
-
- queue: GeneratorQueue[GenerationResponse] = GeneratorQueue()
- parsed_gen = apply_vllm_parsers(
- queue.gen(),
- self.model_id,
- prompt_text,
- self.tool_parser,
- task.task_params.tools,
- think_start=think_start,
- think_end=think_end,
- )
-
- self._active = _ActiveRequest(
- task=task,
- request_id=request_id,
- prompt_token_count=prompt_token_count,
- prompt_token_ids=token_ids,
- queue=queue,
- parsed_gen=parsed_gen,
- )
-
- active = self._active
-
- if self.should_cancel(active.task.task_id):
- self._active = None
- self._cancelled_tasks.discard(active.task.task_id)
- return [(active.task.task_id, Cancelled())]
-
- if not self.engine.has_unfinished_requests():
- self._active = None
- return [(active.task.task_id, Finished())]
-
- if not active.prefill_done:
- max_batch_tokens: int = getattr(self.engine.model_config, "max_num_batched_tokens", 2048) or 2048 # type: ignore[reportUnknownMemberType]
- prefill_steps = 0
- while self.engine.has_unfinished_requests():
- self.agree_on_cancellations()
- if self.should_cancel(active.task.task_id):
- self.engine.abort_request([active.request_id])
- self._active = None
- self._cancelled_tasks.discard(active.task.task_id)
- return [(active.task.task_id, Cancelled())]
- outputs = self.engine.step()
- prefill_steps += 1
- for output in outputs:
- if output.request_id != active.request_id:
- continue
- if len(output.outputs[0].token_ids) > 0:
- active.first_token_time = time.perf_counter()
- active.prefill_done = True
- self._save_prefix_cache(active)
- break
- if not active.prefill_done:
- self.event_sender.send(ChunkGenerated(
- command_id=active.task.command_id,
- chunk=PrefillProgressChunk(
- model=self.model_id,
- processed_tokens=min(prefill_steps * max_batch_tokens, active.prompt_token_count),
- total_tokens=active.prompt_token_count,
- ),
- ))
- if active.prefill_done:
+ internal_id: str | None = None
+ for mgr in coordinator.single_type_managers: # type: ignore
+ for key in mgr.req_to_blocks: # type: ignore
+ if str(key).startswith(request_id): # type: ignore
+ internal_id = str(key) # type: ignore
break
- if not active.prefill_done:
- self._active = None
- return [(active.task.task_id, Finished())]
- # Fall through to process the outputs from the final prefill step
- else:
- outputs = self.engine.step()
- finished = False
- results: list[
- tuple[
- TaskId,
- GenerationResponse | ToolCallResponse | Cancelled | Finished,
- ]
- ] = []
+ if internal_id:
+ break
+ if internal_id is None:
+ return
+
+ null_block = coordinator.block_pool.null_block # type: ignore
+ block_ids_per_group: list[list[int]] = []
+ token_offset_per_group: list[int] = []
+ for mgr in coordinator.single_type_managers: # type: ignore
+ blocks = mgr.req_to_blocks.get(internal_id) # type: ignore
+ if not blocks:
+ block_ids_per_group.append([])
+ token_offset_per_group.append(0)
+ continue
+ block_size: int = mgr.block_size # type: ignore
+ num_leading_nulls = 0
+ for b in blocks: # type: ignore
+ if b is null_block or b.is_null: # type: ignore
+ num_leading_nulls += 1
+ else:
+ break
+ real_blocks = [b for b in blocks if b is not null_block and not b.is_null] # type: ignore
+ block_ids_per_group.append([b.block_id for b in real_blocks]) # type: ignore
+ token_offset_per_group.append(num_leading_nulls * block_size)
+
+ layer_to_group = _build_layer_groups(kv_cache_config)
+ torch_cache = TorchKVCache.from_vllm_cache(
+ model_runner.kv_caches, # type: ignore
+ block_ids_per_group,
+ layer_to_group,
+ prompt_token_count,
+ token_offset_per_group,
+ )
+ prefix_cache.add_from_torch(prompt_token_ids, torch_cache)
+ except Exception:
+ logger.opt(exception=True).warning("Failed to save prefix cache")
+
+
+def _build_generation_response(
+ tokenizer: object,
+ token_id: int,
+ finish_reason: str | None,
+ prompt_token_count: int,
+ completion_tokens: int,
+ start_time: float,
+ first_token_time: float | None,
+) -> GenerationResponse:
+ token_text: str = tokenizer.decode([token_id]) # type: ignore[reportUnknownMemberType]
+ finish_usage: Usage | None = None
+ finish_stats: GenerationStats | None = None
+ mapped_finish_reason: str | None = None
+ if finish_reason:
+ now = time.perf_counter()
+ prefill_elapsed = (first_token_time or now) - start_time
+ decode_elapsed = now - (first_token_time or now)
+ finish_usage = Usage(
+ prompt_tokens=prompt_token_count,
+ completion_tokens=completion_tokens,
+ total_tokens=prompt_token_count + completion_tokens,
+ prompt_tokens_details=PromptTokensDetails(),
+ completion_tokens_details=CompletionTokensDetails(),
+ )
+ finish_stats = GenerationStats(
+ prompt_tps=prompt_token_count / prefill_elapsed if prefill_elapsed > 0 else 0.0,
+ generation_tps=completion_tokens / decode_elapsed if decode_elapsed > 0 else 0.0,
+ prompt_tokens=prompt_token_count,
+ generation_tokens=completion_tokens,
+ peak_memory_usage=Memory.from_bytes(
+ torch.cuda.max_memory_allocated() # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType, reportAttributeAccessIssue]
+ ),
+ )
+ mapped_finish_reason = finish_reason if finish_reason in ("stop", "length", "content_filter") else "stop"
+ return GenerationResponse(
+ text=token_text,
+ token=token_id,
+ finish_reason=mapped_finish_reason,
+ usage=finish_usage,
+ stats=finish_stats,
+ )
+
+
+def vllm_generate(
+ engine: LLMEngine,
+ model_id: ModelId,
+ task: TextGenerationTaskParams,
+ prompt: str,
+ prefix_cache: KVPrefixCache,
+ on_prefill_progress: Callable[[int, int], None] | None = None,
+ distributed_prompt_progress_callback: Callable[[], None] | None = None,
+ on_generation_token: Callable[[], None] | None = None,
+) -> Generator[GenerationResponse, None, None]:
+ token_ids, prompt_text, prompt_token_count = format_vllm_prompt(engine, task)
+ logger.info(prompt_text)
+ request_id = f"vllm-seq-{time.monotonic_ns()}"
+ sampling_params = make_vllm_sampling_params(engine, task, model_id)
+ engine.add_request(request_id, {"prompt_token_ids": token_ids}, sampling_params)
+
+ tokenizer = engine.get_tokenizer()
+ max_batch_tokens: int = getattr(engine.model_config, "max_num_batched_tokens", 2048) or 2048 # type: ignore[reportUnknownMemberType]
+ start_time = time.perf_counter()
+ first_token_time: float | None = None
+ prev_token_count = 0
+ prefill_done = False
+ prefill_steps = 0
+
+ while engine.has_unfinished_requests():
+ if distributed_prompt_progress_callback and not prefill_done:
+ distributed_prompt_progress_callback()
+ outputs = engine.step()
for output in outputs:
- if output.request_id != active.request_id:
+ if output.request_id != request_id:
continue
completion = output.outputs[0]
new_token_count = len(completion.token_ids)
- new_tokens = completion.token_ids[active.prev_token_count :]
-
- new_text = completion.text[len(active.prev_text) :]
+ new_tokens = completion.token_ids[prev_token_count:]
finish_reason = completion.finish_reason
+ prev_token_count = new_token_count
- active.prev_text = completion.text
- active.prev_token_count = new_token_count
- if active.first_token_time is None and new_text:
- active.first_token_time = time.perf_counter()
+ if not prefill_done and not new_tokens:
+ prefill_steps += 1
+ if on_prefill_progress:
+ on_prefill_progress(
+ min(prefill_steps * max_batch_tokens, prompt_token_count),
+ prompt_token_count,
+ )
+ continue
- finish_usage: Usage | None = None
- finish_stats: GenerationStats | None = None
- mapped_finish_reason: str | None = None
- if finish_reason:
- now = time.perf_counter()
- prefill_elapsed = (active.first_token_time or now) - active.start_time
- decode_elapsed = now - (active.first_token_time or now)
- finish_usage = Usage(
- prompt_tokens=active.prompt_token_count,
- completion_tokens=new_token_count,
- total_tokens=active.prompt_token_count + new_token_count,
- prompt_tokens_details=PromptTokensDetails(),
- completion_tokens_details=CompletionTokensDetails(),
- )
- finish_stats = GenerationStats(
- prompt_tps=active.prompt_token_count / prefill_elapsed
- if prefill_elapsed > 0
- else 0.0,
- generation_tps=new_token_count / decode_elapsed
- if decode_elapsed > 0
- else 0.0,
- prompt_tokens=active.prompt_token_count,
- generation_tokens=new_token_count,
- peak_memory_usage=Memory.from_bytes(
- torch.cuda.max_memory_allocated() # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType, reportAttributeAccessIssue]
- ),
- )
- mapped_finish_reason = (
- finish_reason
- if finish_reason in ("stop", "length", "content_filter")
- else "stop"
- )
- finished = True
+ if not prefill_done and new_tokens:
+ first_token_time = time.perf_counter()
+ prefill_done = True
+ _save_prefix_cache(engine, prefix_cache, request_id, token_ids, prompt_token_count)
- tokenizer = self.engine.get_tokenizer()
for i, token_id in enumerate(new_tokens):
is_last = i == len(new_tokens) - 1
- token_text: str = tokenizer.decode([token_id]) # type: ignore[reportUnknownMemberType]
- active.queue.push(
- GenerationResponse(
- text=token_text,
- token=token_id,
- finish_reason=mapped_finish_reason if is_last and finished else None,
- usage=finish_usage if is_last and finished else None,
- stats=finish_stats if is_last and finished else None,
- )
+ if on_generation_token:
+ on_generation_token()
+ yield _build_generation_response(
+ tokenizer, token_id,
+ finish_reason if is_last and finish_reason else None,
+ prompt_token_count, new_token_count,
+ start_time, first_token_time,
)
- try:
- parsed = next(active.parsed_gen)
- except StopIteration:
- self.engine.abort_request([active.request_id])
- results.append((active.task.task_id, Finished()))
- self._active = None
- return results
- if parsed is not None:
- results.append((active.task.task_id, parsed))
-
- if finished:
- logger.info(f"vLLM generation done for request {active.request_id}")
- results.append((active.task.task_id, Finished()))
- self._active = None
-
- return results
-
- def _get_coordinator(self) -> object | None:
- if not hasattr(self, "_coordinator_cached"):
- try:
- engine_core = self.engine.engine_core.engine_core # type: ignore
- self._coordinator_cached: object | None = engine_core.scheduler.kv_cache_manager.coordinator # type: ignore
- except Exception:
- self._coordinator_cached = None
- return self._coordinator_cached
-
- def _get_kv_cache_config(self) -> object | None:
- if not hasattr(self, "_kv_cache_config_cached"):
- try:
- engine_core = self.engine.engine_core.engine_core # type: ignore
- self._kv_cache_config_cached: object | None = engine_core.scheduler.kv_cache_manager.kv_cache_config # type: ignore
- except Exception:
- self._kv_cache_config_cached = None
- return self._kv_cache_config_cached
-
- def _save_prefix_cache(self, active: _ActiveRequest) -> None:
- try:
- coordinator = self._get_coordinator()
- model_runner = _growable_model_runner_ref[0]
- kv_cache_config = self._get_kv_cache_config()
- if coordinator is None or model_runner is None or kv_cache_config is None:
- return
-
- internal_id: str | None = None
- for mgr in coordinator.single_type_managers: # type: ignore
- for key in mgr.req_to_blocks: # type: ignore
- if str(key).startswith(active.request_id): # type: ignore
- internal_id = str(key) # type: ignore
- break
- if internal_id:
- break
- if internal_id is None:
- return
-
- null_block = coordinator.block_pool.null_block # type: ignore
-
- block_ids_per_group: list[list[int]] = []
- token_offset_per_group: list[int] = []
- for mgr in coordinator.single_type_managers: # type: ignore
- blocks = mgr.req_to_blocks.get(internal_id) # type: ignore
- if not blocks:
- block_ids_per_group.append([])
- token_offset_per_group.append(0)
- continue
- block_size: int = mgr.block_size # type: ignore
- num_leading_nulls = 0
- for b in blocks: # type: ignore
- if b is null_block or b.is_null: # type: ignore
- num_leading_nulls += 1
- else:
- break
- real_blocks = [b for b in blocks if b is not null_block and not b.is_null] # type: ignore
- block_ids_per_group.append([b.block_id for b in real_blocks]) # type: ignore
- token_offset_per_group.append(num_leading_nulls * block_size)
-
- layer_to_group = _build_layer_groups(kv_cache_config)
- torch_cache = TorchKVCache.from_vllm_cache(
- model_runner.kv_caches, # type: ignore
- block_ids_per_group,
- layer_to_group,
- active.prompt_token_count,
- token_offset_per_group,
- )
- self.prefix_cache.add_from_torch(active.prompt_token_ids, torch_cache)
- except Exception:
- logger.opt(exception=True).warning("Failed to save prefix cache")
-
- def close(self) -> None:
- del self.engine
- gc.collect()
- torch.cuda.empty_cache()
- if torch.distributed.is_initialized():
- torch.distributed.destroy_process_group()
-_EXO_MAX_CONCURRENT_VLLM_REQUESTS = 8
+def warmup_vllm_engine(engine: LLMEngine) -> None:
+ tokenizer = engine.get_tokenizer()
+ messages = [{"role": "user", "content": "Prompt to warm up the inference engine. Repeat this."}]
+ prompt_text: str = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) # type: ignore
+ token_ids: list[int] = tokenizer.encode(prompt_text, add_special_tokens=False) # type: ignore
+ params = SamplingParams(max_tokens=50, detokenize=False)
+ engine.add_request("warmup", {"prompt_token_ids": token_ids}, params)
+ while engine.has_unfinished_requests():
+ engine.step()
+ logger.info("vLLM warmup complete")
@dataclass(eq=False)
-class VllmBatchGenerator(InferenceGenerator):
+class VllmBatchEngine:
engine: LLMEngine
model_id: ModelId
- tool_parser: ToolParser | None
- cancel_receiver: MpReceiver[TaskId]
- event_sender: MpSender[Event]
prefix_cache: KVPrefixCache
- _cancelled_tasks: set[TaskId] = field(default_factory=set, init=False)
- _all_tasks: dict[TaskId, TextGeneration] = field(default_factory=dict, init=False)
- _maybe_queue: list[TextGeneration] = field(default_factory=list, init=False)
- _queue: deque[TextGeneration] = field(default_factory=deque, init=False)
- _maybe_cancel: list[TextGeneration] = field(default_factory=list, init=False)
- _active: dict[str, _ActiveRequest] = field(default_factory=dict, init=False)
+ _active: dict[int, _EngineRequest] = field(default_factory=dict, init=False)
+ _next_uid: int = field(default=0, init=False)
- def warmup(self) -> None:
- tokenizer = self.engine.get_tokenizer()
- messages = [{"role": "user", "content": "Prompt to warm up the inference engine. Repeat this."}]
- prompt_text: str = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) # type: ignore
- token_ids: list[int] = tokenizer.encode(prompt_text, add_special_tokens=False) # type: ignore
- params = SamplingParams(max_tokens=50, detokenize=False)
- self.engine.add_request("warmup", {"prompt_token_ids": token_ids}, params)
- while self.engine.has_unfinished_requests():
- self.engine.step()
- logger.info("vLLM batch warmup complete")
+ @property
+ def has_work(self) -> bool:
+ return bool(self._active) or self.engine.has_unfinished_requests()
- def submit(self, task: TextGeneration) -> None:
- self._cancelled_tasks.discard(CANCEL_ALL_TASKS)
- self._all_tasks[task.task_id] = task
- self._maybe_queue.append(task)
-
- def agree_on_tasks(self) -> None:
- self._queue.extend(self._maybe_queue)
- self._maybe_queue.clear()
-
- def agree_on_cancellations(self) -> None:
- for task_id in self.cancel_receiver.collect():
- if task_id == CANCEL_ALL_TASKS:
- self._cancelled_tasks.add(CANCEL_ALL_TASKS)
- else:
- if task_id in self._all_tasks:
- self._maybe_cancel.append(self._all_tasks[task_id])
- self._cancelled_tasks.update(task.task_id for task in self._maybe_cancel)
- self._maybe_cancel.clear()
-
- def _start_request(self, task: TextGeneration) -> _ActiveRequest:
- token_ids, prompt_text, prompt_token_count = format_vllm_prompt(
- self.engine, task.task_params
- )
+ def submit(
+ self,
+ task_params: TextGenerationTaskParams,
+ prompt: str,
+ on_prefill_progress: Callable[[int, int], None] | None = None,
+ distributed_prompt_progress_callback: Callable[[], None] | None = None,
+ on_generation_token: Callable[[], None] | None = None,
+ ) -> int:
+ token_ids, prompt_text, prompt_token_count = format_vllm_prompt(self.engine, task_params)
logger.info(prompt_text)
- request_id = str(task.task_id)
- sampling_params = make_vllm_sampling_params(self.engine, task.task_params, self.model_id)
- self.engine.add_request(
- request_id,
- {"prompt_token_ids": token_ids},
- sampling_params,
- )
- tokenizer = self.engine.get_tokenizer()
- think_start: str | None = getattr(tokenizer, "think_start", None)
- think_end: str | None = getattr(tokenizer, "think_end", None)
- queue: GeneratorQueue[GenerationResponse] = GeneratorQueue()
- parsed_gen = apply_vllm_parsers(
- queue.gen(),
- self.model_id,
- prompt_text,
- self.tool_parser,
- task.task_params.tools,
- think_start=think_start,
- think_end=think_end,
- )
- return _ActiveRequest(
- task=task,
+ uid = self._next_uid
+ self._next_uid += 1
+ request_id = f"vllm-batch-{uid}"
+ sampling_params = make_vllm_sampling_params(self.engine, task_params, self.model_id)
+ self.engine.add_request(request_id, {"prompt_token_ids": token_ids}, sampling_params)
+ self._active[uid] = _EngineRequest(
+ uid=uid,
request_id=request_id,
prompt_token_count=prompt_token_count,
prompt_token_ids=token_ids,
- queue=queue,
- parsed_gen=parsed_gen,
+ on_generation_token=on_generation_token,
+ on_prefill_progress=on_prefill_progress,
)
+ return uid
- def _apply_cancellations(
- self,
- ) -> list[tuple[TaskId, Cancelled]]:
- if not self._cancelled_tasks:
+ def step(self) -> list[tuple[int, GenerationResponse]]:
+ if not self.has_work:
return []
- cancel_all = CANCEL_ALL_TASKS in self._cancelled_tasks
- rids_to_abort: list[str] = []
- results: list[tuple[TaskId, Cancelled]] = []
- for rid, active in list(self._active.items()):
- if active.task.task_id in self._cancelled_tasks or cancel_all:
- rids_to_abort.append(rid)
- results.append((active.task.task_id, Cancelled()))
- del self._active[rid]
- if rids_to_abort:
- self.engine.abort_request(rids_to_abort)
- already_cancelled = {tid for tid, _ in results}
- for tid in self._cancelled_tasks:
- if tid != CANCEL_ALL_TASKS and tid not in already_cancelled:
- results.append((tid, Cancelled()))
- self._cancelled_tasks.clear()
- return results
-
- def step(
- self,
- ) -> Iterable[
- tuple[TaskId, GenerationResponse | ToolCallResponse | Cancelled | Finished]
- ]:
- self.agree_on_cancellations()
-
- if not self._queue:
- self.agree_on_tasks()
-
- results: list[
- tuple[TaskId, GenerationResponse | ToolCallResponse | Cancelled | Finished]
- ] = []
-
- while self._queue and len(self._active) < _EXO_MAX_CONCURRENT_VLLM_REQUESTS:
- task = self._queue.popleft()
- if self.should_cancel(task.task_id):
- self._cancelled_tasks.discard(task.task_id)
- results.append((task.task_id, Cancelled()))
- continue
- active = self._start_request(task)
- self._active[active.request_id] = active
-
- if not self._active:
- return itertools.chain(results, self._apply_cancellations())
outputs = self.engine.step()
tokenizer = self.engine.get_tokenizer()
+ max_batch_tokens: int = getattr(self.engine.model_config, "max_num_batched_tokens", 2048) or 2048 # type: ignore[reportUnknownMemberType]
+ results: list[tuple[int, GenerationResponse]] = []
+
+ rid_to_uid = {req.request_id: uid for uid, req in self._active.items()}
for output in outputs:
- rid = output.request_id
- if rid not in self._active:
+ uid = rid_to_uid.get(output.request_id)
+ if uid is None:
continue
- active = self._active[rid]
+ req = self._active[uid]
completion = output.outputs[0]
new_token_count = len(completion.token_ids)
- new_tokens = completion.token_ids[active.prev_token_count:]
+ new_tokens = completion.token_ids[req.prev_token_count:]
finish_reason = completion.finish_reason
+ req.prev_token_count = new_token_count
- active.prev_text = completion.text
- active.prev_token_count = new_token_count
- if active.first_token_time is None and new_tokens:
- active.first_token_time = time.perf_counter()
- if not active.prefill_done:
- active.prefill_done = True
- self._save_prefix_cache(active)
+ if not req.prefill_done and not new_tokens:
+ req.prefill_steps += 1
+ if req.on_prefill_progress:
+ req.on_prefill_progress(
+ min(req.prefill_steps * max_batch_tokens, req.prompt_token_count),
+ req.prompt_token_count,
+ )
+ continue
- finish_usage: Usage | None = None
- finish_stats: GenerationStats | None = None
- mapped_finish_reason: str | None = None
- finished = False
- if finish_reason:
- now = time.perf_counter()
- prefill_elapsed = (active.first_token_time or now) - active.start_time
- decode_elapsed = now - (active.first_token_time or now)
- finish_usage = Usage(
- prompt_tokens=active.prompt_token_count,
- completion_tokens=new_token_count,
- total_tokens=active.prompt_token_count + new_token_count,
- prompt_tokens_details=PromptTokensDetails(),
- completion_tokens_details=CompletionTokensDetails(),
+ if not req.prefill_done and new_tokens:
+ req.first_token_time = time.perf_counter()
+ req.prefill_done = True
+ _save_prefix_cache(
+ self.engine, self.prefix_cache,
+ req.request_id, req.prompt_token_ids, req.prompt_token_count,
)
- finish_stats = GenerationStats(
- prompt_tps=active.prompt_token_count / prefill_elapsed if prefill_elapsed > 0 else 0.0,
- generation_tps=new_token_count / decode_elapsed if decode_elapsed > 0 else 0.0,
- prompt_tokens=active.prompt_token_count,
- generation_tokens=new_token_count,
- peak_memory_usage=Memory.from_bytes(
- torch.cuda.max_memory_allocated() # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType, reportAttributeAccessIssue]
- ),
- )
- mapped_finish_reason = (
- finish_reason if finish_reason in ("stop", "length", "content_filter") else "stop"
- )
- finished = True
for i, token_id in enumerate(new_tokens):
is_last = i == len(new_tokens) - 1
- token_text: str = tokenizer.decode([token_id]) # type: ignore[reportUnknownMemberType]
- active.queue.push(
- GenerationResponse(
- text=token_text,
- token=token_id,
- finish_reason=mapped_finish_reason if is_last and finished else None,
- usage=finish_usage if is_last and finished else None,
- stats=finish_stats if is_last and finished else None,
+ if req.on_generation_token:
+ req.on_generation_token()
+ results.append((uid, _build_generation_response(
+ tokenizer, token_id,
+ finish_reason if is_last and finish_reason else None,
+ req.prompt_token_count, new_token_count,
+ req.start_time, req.first_token_time,
+ )))
+
+ if finish_reason:
+ del self._active[uid]
+
+ for req in self._active.values():
+ if not req.prefill_done:
+ req.prefill_steps += 1
+ if req.on_prefill_progress:
+ req.on_prefill_progress(
+ min(req.prefill_steps * max_batch_tokens, req.prompt_token_count),
+ req.prompt_token_count,
)
- )
- try:
- parsed = next(active.parsed_gen)
- except StopIteration:
- self.engine.abort_request([rid])
- results.append((active.task.task_id, Finished()))
- del self._active[rid]
- break
- if parsed is not None:
- results.append((active.task.task_id, parsed))
- else:
- if finished:
- logger.info(f"vLLM generation done for request {rid}")
- results.append((active.task.task_id, Finished()))
- del self._active[rid]
- max_batch_tokens: int = getattr(self.engine.model_config, "max_num_batched_tokens", 2048) or 2048 # type: ignore[reportUnknownMemberType]
- for active in self._active.values():
- if not active.prefill_done:
- active.prefill_steps += 1
- self.event_sender.send(ChunkGenerated(
- command_id=active.task.command_id,
- chunk=PrefillProgressChunk(
- model=self.model_id,
- processed_tokens=min(active.prefill_steps * max_batch_tokens, active.prompt_token_count),
- total_tokens=active.prompt_token_count,
- ),
- ))
+ return results
- return itertools.chain(results, self._apply_cancellations())
-
- def _get_coordinator(self) -> object | None:
- if not hasattr(self, "_coordinator_cached"):
- try:
- engine_core = self.engine.engine_core.engine_core # type: ignore
- self._coordinator_cached: object | None = engine_core.scheduler.kv_cache_manager.coordinator # type: ignore
- except Exception:
- self._coordinator_cached = None
- return self._coordinator_cached
-
- def _get_kv_cache_config(self) -> object | None:
- if not hasattr(self, "_kv_cache_config_cached"):
- try:
- engine_core = self.engine.engine_core.engine_core # type: ignore
- self._kv_cache_config_cached: object | None = engine_core.scheduler.kv_cache_manager.kv_cache_config # type: ignore
- except Exception:
- self._kv_cache_config_cached = None
- return self._kv_cache_config_cached
-
- def _save_prefix_cache(self, active: _ActiveRequest) -> None:
- try:
- coordinator = self._get_coordinator()
- model_runner = _growable_model_runner_ref[0]
- kv_cache_config = self._get_kv_cache_config()
- if coordinator is None or model_runner is None or kv_cache_config is None:
- return
- internal_id: str | None = None
- for mgr in coordinator.single_type_managers: # type: ignore
- for key in mgr.req_to_blocks: # type: ignore
- if str(key).startswith(active.request_id): # type: ignore
- internal_id = str(key) # type: ignore
- break
- if internal_id:
- break
- if internal_id is None:
- return
- null_block = coordinator.block_pool.null_block # type: ignore
- block_ids_per_group: list[list[int]] = []
- token_offset_per_group: list[int] = []
- for mgr in coordinator.single_type_managers: # type: ignore
- blocks = mgr.req_to_blocks.get(internal_id) # type: ignore
- if not blocks:
- block_ids_per_group.append([])
- token_offset_per_group.append(0)
- continue
- block_size: int = mgr.block_size # type: ignore
- num_leading_nulls = 0
- for b in blocks: # type: ignore
- if b is null_block or b.is_null: # type: ignore
- num_leading_nulls += 1
- else:
- break
- real_blocks = [b for b in blocks if b is not null_block and not b.is_null] # type: ignore
- block_ids_per_group.append([b.block_id for b in real_blocks]) # type: ignore
- token_offset_per_group.append(num_leading_nulls * block_size)
- layer_to_group = _build_layer_groups(kv_cache_config)
- torch_cache = TorchKVCache.from_vllm_cache(
- model_runner.kv_caches, # type: ignore
- block_ids_per_group,
- layer_to_group,
- active.prompt_token_count,
- token_offset_per_group,
- )
- self.prefix_cache.add_from_torch(active.prompt_token_ids, torch_cache)
- except Exception:
- logger.opt(exception=True).warning("Failed to save prefix cache")
+ def cancel(self, uids: list[int]) -> None:
+ rids = [self._active[uid].request_id for uid in uids if uid in self._active]
+ if rids:
+ self.engine.abort_request(rids)
+ for uid in uids:
+ self._active.pop(uid, None)
def close(self) -> None:
- for rid in list(self._active):
- self.engine.abort_request([rid])
+ rids = [req.request_id for req in self._active.values()]
+ if rids:
+ self.engine.abort_request(rids)
self._active.clear()
del self.engine
gc.collect()
@@ -796,11 +470,6 @@ def load_vllm_engine(
load_format="fastsafetensors",
enable_prefix_caching=False,
attention_backend="TRITON_ATTN",
- # torch.compile's piecewise compilation produces kernels that assume
- # KV cache blocks are written by previous compiled forward passes.
- # Our prefix cache writes KV externally via tensor indexing and skips
- # the forward pass, which leaves the compiled kernel's internal state
- # inconsistent with the actual block data, producing garbage output.
enforce_eager=True,
disable_log_stats=True,
)
diff --git a/src/exo/worker/runner/llm_inference/batch_generator.py b/src/exo/worker/runner/llm_inference/batch_generator.py
index e0382ca4..5dea6da4 100644
--- a/src/exo/worker/runner/llm_inference/batch_generator.py
+++ b/src/exo/worker/runner/llm_inference/batch_generator.py
@@ -140,13 +140,14 @@ class SequentialGenerator(InferenceGenerator):
| None
) = field(default=None, init=False)
- def warmup(self):
- self.check_for_cancel_every = warmup_inference(
- model=self.model,
- tokenizer=self.tokenizer,
- group=self.group,
- model_id=self.model_id,
- )
+ def warmup(self) -> None:
+ if self.model is not None:
+ self.check_for_cancel_every = warmup_inference(
+ model=self.model,
+ tokenizer=self.tokenizer,
+ group=self.group,
+ model_id=self.model_id,
+ )
def submit(
self,
@@ -304,7 +305,7 @@ class SequentialGenerator(InferenceGenerator):
@dataclass(eq=False)
class BatchGenerator(InferenceGenerator):
- model: Model
+ model: Model | None
tokenizer: TokenizerWrapper
group: mx.distributed.Group | None
kv_prefix_cache: KVPrefixCache | None
@@ -313,6 +314,8 @@ class BatchGenerator(InferenceGenerator):
device_rank: int
cancel_receiver: MpReceiver[TaskId]
event_sender: MpSender[Event]
+ _gen: ExoBatchGenerator # ExoBatchGenerator or VllmBatchEngine
+ max_concurrent_requests: int = EXO_MAX_CONCURRENT_REQUESTS
check_for_cancel_every: int = 50
_cancelled_tasks: set[TaskId] = field(default_factory=set, init=False)
@@ -320,7 +323,6 @@ class BatchGenerator(InferenceGenerator):
_maybe_cancel: list[TextGeneration] = field(default_factory=list, init=False)
_all_tasks: dict[TaskId, TextGeneration] = field(default_factory=dict, init=False)
_queue: deque[TextGeneration] = field(default_factory=deque, init=False)
- _mlx_gen: ExoBatchGenerator = field(init=False)
_active_tasks: dict[
int,
tuple[
@@ -330,21 +332,14 @@ class BatchGenerator(InferenceGenerator):
],
] = field(default_factory=dict, init=False)
- def __post_init__(self) -> None:
- self._mlx_gen = ExoBatchGenerator(
- model=self.model,
- tokenizer=self.tokenizer,
- group=self.group,
- kv_prefix_cache=self.kv_prefix_cache,
- )
-
- def warmup(self):
- self.check_for_cancel_every = warmup_inference(
- model=self.model,
- tokenizer=self.tokenizer,
- group=self.group,
- model_id=self.model_id,
- )
+ def warmup(self) -> None:
+ if self.model is not None:
+ self.check_for_cancel_every = warmup_inference(
+ model=self.model,
+ tokenizer=self.tokenizer,
+ group=self.group,
+ model_id=self.model_id,
+ )
def submit(
self,
@@ -386,7 +381,7 @@ class BatchGenerator(InferenceGenerator):
self.agree_on_tasks()
# Submit any queued tasks to the engine
- while self._queue and len(self._active_tasks) < EXO_MAX_CONCURRENT_REQUESTS:
+ while self._queue and len(self._active_tasks) < self.max_concurrent_requests:
task = self._queue.popleft()
try:
uid = self._start_task(task)
@@ -411,10 +406,10 @@ class BatchGenerator(InferenceGenerator):
)
self._active_tasks[uid] = (task, queue, output_generator)
- if not self._mlx_gen.has_work:
+ if not self._gen.has_work:
return self._apply_cancellations()
- results = self._mlx_gen.step()
+ results = self._gen.step()
output: list[
tuple[TaskId, GenerationResponse | ToolCallResponse | Cancelled | Finished]
@@ -456,7 +451,7 @@ class BatchGenerator(InferenceGenerator):
del self._active_tasks[uid]
if uids_to_cancel:
- self._mlx_gen.cancel(uids_to_cancel)
+ self._gen.cancel(uids_to_cancel)
already_cancelled = {tid for tid, _ in results}
for tid in self._cancelled_tasks:
@@ -516,7 +511,7 @@ class BatchGenerator(InferenceGenerator):
self.agree_on_tasks()
- return self._mlx_gen.submit(
+ return self._gen.submit(
task_params=task.task_params,
prompt=prompt,
on_prefill_progress=on_prefill_progress,
@@ -525,5 +520,7 @@ class BatchGenerator(InferenceGenerator):
)
def close(self) -> None:
- self._mlx_gen.close()
- del self.model, self.tokenizer, self.group
+ self._gen.close()
+ if self.model is not None:
+ del self.model
+ del self.tokenizer, self.group
diff --git a/src/exo/worker/runner/llm_inference/model_output_parsers.py b/src/exo/worker/runner/llm_inference/model_output_parsers.py
index 00393565..0e5a4cf2 100644
--- a/src/exo/worker/runner/llm_inference/model_output_parsers.py
+++ b/src/exo/worker/runner/llm_inference/model_output_parsers.py
@@ -49,12 +49,10 @@ def apply_all_parsers(
starts_in_thinking=detect_thinking_prompt_suffix(prompt, tokenizer),
)
- if issubclass(model_type, GptOssModel):
+ lower = model_id.normalize().lower()
+ if issubclass(model_type, GptOssModel) or "gpt-oss" in lower or "gpt_oss" in lower:
mlx_generator = parse_gpt_oss(mlx_generator)
- elif (
- issubclass(model_type, DeepseekV32Model)
- and "deepseek" in model_id.normalize().lower()
- ):
+ elif issubclass(model_type, DeepseekV32Model) or "deepseek" in lower:
mlx_generator = parse_deepseek_v32(mlx_generator)
elif tool_parser:
mlx_generator = parse_tool_calls(mlx_generator, tool_parser, tools)
@@ -62,42 +60,6 @@ def apply_all_parsers(
return mlx_generator
-def apply_vllm_parsers(
- receiver: Generator[GenerationResponse | None],
- model_id: ModelId,
- prompt: str,
- tool_parser: ToolParser | None,
- tools: list[dict[str, Any]] | None,
- think_start: str | None = None,
- think_end: str | None = None,
-) -> Generator[GenerationResponse | ToolCallResponse | None]:
- gen = receiver
- lower = model_id.normalize().lower()
-
- if "gpt-oss" in lower or "gpt_oss" in lower:
- return parse_gpt_oss(gen)
-
- if "deepseek" in lower:
- gen = parse_thinking_models(
- gen,
- think_start or "",
- think_end or "",
- starts_in_thinking=prompt.rstrip().endswith(think_start or ""),
- )
- return parse_deepseek_v32(gen)
-
- if think_start is not None:
- gen = parse_thinking_models(
- gen,
- think_start,
- think_end,
- starts_in_thinking=prompt.rstrip().endswith(think_start),
- )
- if tool_parser:
- gen = parse_tool_calls(gen, tool_parser, tools)
- return gen
-
-
_GPT_OSS_CHANNEL_TOKEN = 200005
_GPT_OSS_MESSAGE_TOKEN = 200008
diff --git a/src/exo/worker/runner/llm_inference/runner.py b/src/exo/worker/runner/llm_inference/runner.py
index 19a66909..3de42709 100644
--- a/src/exo/worker/runner/llm_inference/runner.py
+++ b/src/exo/worker/runner/llm_inference/runner.py
@@ -462,7 +462,15 @@ class MlxBuilder(Builder):
cancel_receiver=self.cancel_receiver,
event_sender=self.event_sender,
)
+ from exo.worker.runner.llm_inference.batch_generator import ExoBatchGenerator
+
logger.info("using BatchGenerator")
+ gen = ExoBatchGenerator(
+ model=self.inference_model,
+ tokenizer=self.tokenizer,
+ group=self.group,
+ kv_prefix_cache=kv_prefix_cache,
+ )
return BatchGenerator(
model=self.inference_model,
tokenizer=self.tokenizer,
@@ -473,6 +481,7 @@ class MlxBuilder(Builder):
device_rank=device_rank,
cancel_receiver=self.cancel_receiver,
event_sender=self.event_sender,
+ _gen=gen,
)
def shutdown_cleanup(self) -> None:
@@ -509,28 +518,35 @@ class VllmBuilder(Builder):
)
def build(self) -> InferenceGenerator:
- if os.environ.get("EXO_NO_BATCH"):
- from exo.worker.engines.vllm.vllm_generator import VllmSequentialGenerator
+ from mlx_lm.tokenizer_utils import TokenizerWrapper
- logger.info("using VllmSequentialGenerator (batching disabled)")
- return VllmSequentialGenerator(
- engine=self._engine,
- model_id=self.model_id,
- tool_parser=self._tool_parser,
- cancel_receiver=self.cancel_receiver,
- event_sender=self.event_sender,
- prefix_cache=self._prefix_cache,
- )
- from exo.worker.engines.vllm.vllm_generator import VllmBatchGenerator
+ from exo.worker.engines.vllm.vllm_generator import (
+ VllmBatchEngine,
+ warmup_vllm_engine,
+ )
- logger.info("using VllmBatchGenerator")
- return VllmBatchGenerator(
+ warmup_vllm_engine(self._engine)
+ gen = VllmBatchEngine(
engine=self._engine,
model_id=self.model_id,
+ prefix_cache=self._prefix_cache,
+ )
+ tokenizer = TokenizerWrapper(self._engine.get_tokenizer())
+ max_concurrent = 1 if os.environ.get("EXO_NO_BATCH") else 8
+
+ logger.info(f"using BatchGenerator (vLLM, max_concurrent={max_concurrent})")
+ return BatchGenerator(
+ model=None,
+ tokenizer=tokenizer,
+ group=None,
tool_parser=self._tool_parser,
+ kv_prefix_cache=None,
+ model_id=self.model_id,
+ device_rank=0,
cancel_receiver=self.cancel_receiver,
event_sender=self.event_sender,
- prefix_cache=self._prefix_cache,
+ _gen=gen,
+ max_concurrent_requests=max_concurrent,
)
def shutdown_cleanup(self) -> None: