electron-rare/exo
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Distributed callbacks

Browse Source
This commit is contained in:
Ryuichi Leo Takashige
2026-03-16 21:17:14 +00:00
parent dc68ddbac0
commit e96f084051
8 changed files with 118 additions and 679 deletions
Download Patch File Download Diff File Expand all files Collapse all files
scripts/investigate_mlx_kv.py
-121
View File
@@ -1,121 +0,0 @@
#!/usr/bin/env python3
"""Investigate MLX KV cache shapes and test NHD conversion.
Run on Mac:
uv run python scripts/investigate_mlx_kv.py --model-path ~/.exo/models/mlx-community--gpt-oss-20b-MXFP4-Q8
"""
import argparse
import mlx.core as mx
from mlx_lm import load
from mlx_lm.generate import stream_generate
from mlx_lm.models.cache import ArraysCache, KVCache, RotatingKVCache
from mlx_lm.sample_utils import make_sampler
from exo.worker.engines.kv_cache import TorchKVCache
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--model-path", required=True)
parser.add_argument("--max-tokens", type=int, default=10)
args = parser.parse_args()
print(f"Loading model from {args.model_path}")
model, tokenizer = load(args.model_path)
print(f"Model loaded. {len(model.layers)} layers")
# Create cache and run generation to populate it
prompt = "Hello, how are you?"
prompt_tokens = mx.array(tokenizer.encode(prompt))
print(f"Prompt: {prompt!r} ({len(prompt_tokens)} tokens)")
cache = model.make_cache()
print(f"\n=== Raw MLX Cache ({len(cache)} layers) ===")
for i, c in enumerate(cache):
print(f" [{i}] {type(c).__name__}", end="")
if isinstance(c, RotatingKVCache):
print(f" max_size={c.max_size} keep={c.keep}")
elif isinstance(c, ArraysCache):
print(f" size={len(c.state)}")
else:
print()
# Generate tokens to fill the cache
print(f"\nGenerating {args.max_tokens} tokens...")
generated = []
for out in stream_generate(
model=model,
tokenizer=tokenizer,
prompt=prompt_tokens,
max_tokens=args.max_tokens,
sampler=make_sampler(temp=0.0),
prompt_cache=cache,
):
generated.append(out.text)
if out.finish_reason:
break
print(f"Generated: {''.join(generated)!r}")
# Inspect populated cache
print("\n=== Populated MLX Cache ===")
for i, c in enumerate(cache):
if isinstance(c, (KVCache, RotatingKVCache)):
if c.keys is not None:
k, v = c.state
mx.eval(k)
mx.eval(v)
print(f" [{i}] {type(c).__name__}: keys={k.shape} values={v.shape} dtype={k.dtype} offset={c.offset}", end="")
if isinstance(c, RotatingKVCache):
print(f" _idx={c._idx} keep={c.keep} max_size={c.max_size} meta_state={c.meta_state}")
else:
print()
else:
print(f" [{i}] {type(c).__name__}: empty")
elif isinstance(c, ArraysCache):
shapes = [list(a.shape) if a is not None else None for a in c.state]
print(f" [{i}] ArraysCache: {shapes}")
# Convert to TorchKVCache
print("\n=== Converting to TorchKVCache (NHD format) ===")
torch_cache = TorchKVCache.from_mlx_cache(cache)
print(torch_cache)
# Value statistics for KV layers
print("\n=== Value Statistics ===")
for idx, layer in torch_cache.kv_layers():
k, v = layer.keys, layer.values
print(f" [{idx}] keys: min={k.min():.6f} max={k.max():.6f} mean={k.mean():.6f} std={k.std():.6f}")
print(f" [{idx}] vals: min={v.min():.6f} max={v.max():.6f} mean={v.mean():.6f} std={v.std():.6f}")
# Round-trip test
print("\n=== Round-trip: TorchKVCache -> MLX Cache ===")
restored_cache = torch_cache.to_mlx_cache()
for i, (orig, restored) in enumerate(zip(cache, restored_cache, strict=True)):
if isinstance(orig, (KVCache, RotatingKVCache)) and orig.keys is not None:
ok, ov = orig.state
rk, rv = restored.state
mx.eval(ok, ov, rk, rv)
k_diff = mx.max(mx.abs(ok - rk)).item()
v_diff = mx.max(mx.abs(ov - rv)).item()
print(f" [{i}] {type(orig).__name__}: key_diff={k_diff:.2e} val_diff={v_diff:.2e}", end="")
if isinstance(orig, RotatingKVCache):
meta_match = orig.meta_state == restored.meta_state
print(f" meta_match={meta_match} orig={orig.meta_state} restored={restored.meta_state}")
else:
offset_match = orig.offset == restored.offset
print(f" offset_match={offset_match} ({orig.offset} vs {restored.offset})")
elif isinstance(orig, ArraysCache):
diffs = []
for _j, (oa, ra) in enumerate(zip(orig.state, restored.state, strict=True)):
if oa is not None and ra is not None:
mx.eval(oa, ra)
diffs.append(mx.max(mx.abs(oa - ra)).item())
print(f" [{i}] ArraysCache: max_diffs={[f'{d:.2e}' for d in diffs]}")
print("\nDone.")
if __name__ == "__main__":
main()
scripts/investigate_roundtrip_mlx.py
-88
View File
@@ -1,88 +0,0 @@
#!/usr/bin/env python3
"""Verify TorchKVCache round-trip produces identical model outputs.
Run on Mac:
uv run python scripts/investigate_roundtrip_mlx.py --model-path ~/.exo/models/mlx-community--gpt-oss-20b-MXFP4-Q8
"""
import argparse
from copy import deepcopy
import mlx.core as mx
from mlx_lm import load
from mlx_lm.generate import stream_generate
from mlx_lm.sample_utils import make_sampler
from exo.worker.engines.kv_cache import TorchKVCache
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--model-path", required=True)
args = parser.parse_args()
print(f"Loading model from {args.model_path}")
model, tokenizer = load(args.model_path)
prompt = "The capital of France is"
prompt_tokens = mx.array(tokenizer.encode(prompt))
print(f"Prompt: {prompt!r} ({len(prompt_tokens)} tokens)")
# Prefill: generate 1 token to populate cache
cache_orig = model.make_cache()
for out in stream_generate(
model=model,
tokenizer=tokenizer,
prompt=prompt_tokens,
max_tokens=1,
sampler=make_sampler(temp=0.0),
prompt_cache=cache_orig,
):
first_token = out.text
break
print(f"First generated token: {first_token!r}")
# Snapshot the original cache (deep copy for comparison after model forward)
cache_snapshot = deepcopy(cache_orig)
# Convert original cache → TorchKVCache → back to MLX
print("\nConverting: MLX -> TorchKVCache (NHD) -> MLX")
torch_cache = TorchKVCache.from_mlx_cache(cache_snapshot)
print(torch_cache)
cache_roundtrip = torch_cache.to_mlx_cache()
# Run model forward with original cache
next_token = mx.array([[tokenizer.encode(first_token)[-1]]])
logits_orig = model(next_token, cache=cache_orig)
mx.eval(logits_orig)
# Run model forward with round-tripped cache
logits_roundtrip = model(next_token, cache=cache_roundtrip)
mx.eval(logits_roundtrip)
# Compare
diff = mx.abs(logits_orig - logits_roundtrip)
max_diff = mx.max(diff).item()
mean_diff = mx.mean(diff).item()
print("\n=== Logit Comparison ===")
print(f" logits_orig shape: {logits_orig.shape}")
print(f" max abs diff: {max_diff:.2e}")
print(f" mean abs diff: {mean_diff:.2e}")
# Check if top token is the same
top_orig = mx.argmax(logits_orig[0, -1]).item()
top_roundtrip = mx.argmax(logits_roundtrip[0, -1]).item()
print(f" top token orig: {top_orig} ({tokenizer.decode([top_orig])!r})")
print(f" top token roundtrip: {top_roundtrip} ({tokenizer.decode([top_roundtrip])!r})")
print(f" top tokens match: {top_orig == top_roundtrip}")
if max_diff == 0.0:
print("\nBIT-EXACT round-trip confirmed.")
elif max_diff < 1e-5:
print(f"\nNear-exact round-trip (max diff {max_diff:.2e}, likely bfloat16 precision).")
else:
print(f"\nWARNING: significant round-trip diff {max_diff:.2e}")
if __name__ == "__main__":
main()
scripts/test_gpt_oss_mlx.py
-184
View File
@@ -1,184 +0,0 @@
#!/usr/bin/env python3
"""Test GPT OSS output parsing with MLX backend.
Run on Mac:
uv run python scripts/test_gpt_oss_mlx.py
Replicates exo's exact flow:
mlx_lm.stream_generate -> GenerationResponse -> parse_gpt_oss -> output
"""
import json
import mlx.core as mx
from mlx_lm import load
from mlx_lm.generate import stream_generate
from mlx_lm.sample_utils import make_sampler
from openai_harmony import (
HarmonyEncodingName,
HarmonyError,
Role,
StreamableParser,
load_harmony_encoding,
)
from exo.shared.constants import EXO_MODELS_DIR
MODEL_PATH = str(EXO_MODELS_DIR / "mlx-community--gpt-oss-20b-MXFP4-Q8")
SYSTEM_PROMPT = "You are a helpful AI assistant. Respond directly and concisely. Do not show your reasoning or thought process. When files are shared with you, analyze them and respond helpfully."
MESSAGES_SINGLE = [
{"role": "system", "content": SYSTEM_PROMPT},
{"role": "user", "content": "Hi!"},
]
MESSAGES_MULTI = [
{"role": "system", "content": SYSTEM_PROMPT},
{"role": "user", "content": "Hi!"},
{"role": "assistant", "content": "Hello! How can I help you today?"},
{"role": "user", "content": "A lot."},
]
MESSAGES_LONG = [
{"role": "system", "content": SYSTEM_PROMPT},
{"role": "user", "content": "Hi!"},
{"role": "assistant", "content": "Hello! How can I help you today?"},
{"role": "user", "content": "What's the capital of France?"},
{"role": "assistant", "content": "Paris."},
{"role": "user", "content": "And Germany?"},
{"role": "assistant", "content": "Berlin."},
{"role": "user", "content": "Tell me a fun fact about Berlin."},
]
MAX_TOKENS = 200
def run_generation(model, tokenizer, messages: list[dict], label: str):
prompt_text: str = tokenizer.apply_chat_template(
messages,
tokenize=False,
add_generation_prompt=True,
)
prompt_tokens: list[int] = tokenizer.encode(prompt_text, add_special_tokens=False)
print(f"\n{'=' * 120}")
print(f"=== {label} ===")
print(f"{'=' * 120}")
print(f"PROMPT TEXT:\n{prompt_text}\n")
print(f"PROMPT TOKENS ({len(prompt_tokens)}): {prompt_tokens}\n")
encoding = load_harmony_encoding(HarmonyEncodingName.HARMONY_GPT_OSS)
parser = StreamableParser(encoding, role=Role.ASSISTANT)
thinking = False
current_tool_name: str | None = None
tool_arg_parts: list[str] = []
print(f"{'idx':>4} | {'token_id':>8} | {'mlx_text':>20} | {'harm_delta':>25} | {'channel':>12} | {'recipient':>25} | {'yield':>40}")
print("-" * 160)
prompt_array = mx.array(prompt_tokens)
_CHANNEL_TOKEN = 200005
_MESSAGE_TOKEN = 200008
_IDLE, _EXPECT_NAME, _EXPECT_MSG = 0, 1, 2
header_state = _IDLE
all_tokens: list[int] = []
all_yielded: list[str] = []
for i, out in enumerate(stream_generate(
model=model,
tokenizer=tokenizer,
prompt=prompt_array,
max_tokens=MAX_TOKENS,
sampler=make_sampler(temp=0.0),
)):
token_id = int(out.token)
all_tokens.append(token_id)
if header_state == _EXPECT_MSG and token_id != _MESSAGE_TOKEN:
parser.process(_MESSAGE_TOKEN)
header_state = _IDLE
elif header_state == _EXPECT_MSG:
header_state = _IDLE
elif header_state == _EXPECT_NAME:
header_state = _EXPECT_MSG
if token_id == _CHANNEL_TOKEN:
header_state = _EXPECT_NAME
try:
parser.process(token_id)
except HarmonyError as e:
print(f" !! HarmonyError at token {i}: {e}")
break
delta = parser.last_content_delta
ch = parser.current_channel
recipient = parser.current_recipient
yielded = ""
effective_recipient = recipient if (recipient is not None and recipient.startswith("functions.")) else None
if effective_recipient != current_tool_name:
if current_tool_name is not None:
tool_name = current_tool_name.removeprefix("functions.")
args = "".join(tool_arg_parts).strip()
yielded = f"TOOL_CALL({tool_name}, {args!r})"
tool_arg_parts = []
current_tool_name = effective_recipient
if current_tool_name is not None:
if delta:
tool_arg_parts.append(delta)
if out.finish_reason is not None:
yielded = f"TOOL_FINISH({json.dumps(''.join(tool_arg_parts))})"
tool_arg_parts = []
else:
is_suppressed = ch == "analysis" or (recipient is not None and recipient.startswith("!"))
if is_suppressed and not thinking:
thinking = True
if not is_suppressed and thinking:
thinking = False
if delta:
prefix = "[THINK] " if thinking else ""
yielded = f"{prefix}{delta!r}"
if out.finish_reason is not None:
yielded += f" [FINISH={out.finish_reason}]"
all_yielded.append(yielded)
print(f"{i:4d} | {token_id:8d} | {out.text!r:>20} | {delta!r:>25} | {str(ch):>12} | {str(recipient):>25} | {yielded:>40}")
if out.finish_reason is not None:
break
print(f"\n--- RAW TOKEN IDS ({len(all_tokens)}) ---")
print(all_tokens)
print("\n--- FINAL TEXT ---")
text_parts = [y for y in all_yielded if y and not y.startswith("TOOL_") and not y.endswith("]") or "[THINK]" in y]
final_text = ""
for y in all_yielded:
if y.startswith("[THINK] "):
continue
if y.startswith("TOOL_"):
continue
clean = y.replace(" [FINISH=stop]", "").replace(" [FINISH=length]", "")
if clean.startswith("'") and clean.endswith("'") or clean.startswith('"') and clean.endswith('"'):
final_text += clean[1:-1]
print(repr(final_text))
def main():
print(f"Loading model from {MODEL_PATH}")
model, tokenizer = load(MODEL_PATH)
print(f"Model loaded. EOS token IDs: {tokenizer.eos_token_ids}")
run_generation(model, tokenizer, MESSAGES_SINGLE, "SINGLE TURN: Hi!")
run_generation(model, tokenizer, MESSAGES_MULTI, "MULTI TURN: Hi! -> A lot.")
run_generation(model, tokenizer, MESSAGES_LONG, "LONG CONVO: 5 turns")
if __name__ == "__main__":
main()
scripts/test_gpt_oss_vllm.py
-245
View File
@@ -1,245 +0,0 @@
#!/usr/bin/env python3
"""Test GPT OSS output parsing with vLLM backend.
Run on Spark:
uv run python scripts/test_gpt_oss_vllm.py
uv run python scripts/test_gpt_oss_vllm.py --no-prefix-cache
uv run python scripts/test_gpt_oss_vllm.py --multi-first
"""
import json
import os
import sys
os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0"
from openai_harmony import (
HarmonyEncodingName,
HarmonyError,
Role,
StreamableParser,
load_harmony_encoding,
)
from vllm.engine.arg_utils import EngineArgs
from vllm.sampling_params import SamplingParams
from vllm.v1.engine.llm_engine import LLMEngine
from exo.shared.constants import EXO_MODELS_DIR
MODEL_ID = str(EXO_MODELS_DIR / "openai--gpt-oss-20b")
SYSTEM_PROMPT = "You are a helpful AI assistant. Respond directly and concisely. Do not show your reasoning or thought process. When files are shared with you, analyze them and respond helpfully."
MESSAGES_SINGLE = [
{"role": "system", "content": SYSTEM_PROMPT},
{"role": "user", "content": "Hi!"},
]
MESSAGES_MULTI = [
{"role": "system", "content": SYSTEM_PROMPT},
{"role": "user", "content": "Hi!"},
{"role": "assistant", "content": "Hello! How can I help you today?"},
{"role": "user", "content": "A lot."},
]
MESSAGES_LONG = [
{"role": "system", "content": SYSTEM_PROMPT},
{"role": "user", "content": "Hi!"},
{"role": "assistant", "content": "Hello! How can I help you today?"},
{"role": "user", "content": "What's the capital of France?"},
{"role": "assistant", "content": "Paris."},
{"role": "user", "content": "And Germany?"},
{"role": "assistant", "content": "Berlin."},
{"role": "user", "content": "Tell me a fun fact about Berlin."},
]
MAX_TOKENS = 200
_CHANNEL_TOKEN = 200005
_MESSAGE_TOKEN = 200008
def run_generation(engine, messages: list[dict], label: str, request_id: str):
tokenizer = engine.get_tokenizer()
prompt_text: str = tokenizer.apply_chat_template(
messages,
tokenize=False,
add_generation_prompt=True,
)
tokens_from_encode: list[int] = tokenizer.encode(prompt_text, add_special_tokens=False)
tokens_from_template: list[int] = tokenizer.apply_chat_template(
messages,
tokenize=True,
add_generation_prompt=True,
)
print(f"\n{'=' * 120}")
print(f"=== {label} ===")
print(f"{'=' * 120}")
print(f"PROMPT TEXT:\n{prompt_text}\n")
print(f"TOKENS FROM encode() ({len(tokens_from_encode)}): {tokens_from_encode}")
print(f"TOKENS FROM template ({len(tokens_from_template)}): {tokens_from_template}")
if tokens_from_encode != tokens_from_template:
print(" !! MISMATCH! Diffs:")
for i, (a, b) in enumerate(zip(tokens_from_encode, tokens_from_template)):
if a != b:
print(f" pos {i}: encode={a} ({tokenizer.decode([a])!r}) vs template={b} ({tokenizer.decode([b])!r})")
if len(tokens_from_encode) != len(tokens_from_template):
print(f" length diff: encode={len(tokens_from_encode)} vs template={len(tokens_from_template)}")
else:
print(" (tokens match)")
prompt_tokens = tokens_from_template
encoding = load_harmony_encoding(HarmonyEncodingName.HARMONY_GPT_OSS)
parser = StreamableParser(encoding, role=Role.ASSISTANT)
thinking = False
current_tool_name: str | None = None
tool_arg_parts: list[str] = []
print(f"{'idx':>4} | {'token_id':>8} | {'vllm_text':>20} | {'harm_delta':>25} | {'channel':>12} | {'recipient':>25} | {'yield':>40}")
print("-" * 160)
sampling_params = SamplingParams(max_tokens=MAX_TOKENS, temperature=0.0)
engine.add_request(request_id, {"prompt_token_ids": prompt_tokens}, sampling_params)
_IDLE, _EXPECT_NAME, _EXPECT_MSG = 0, 1, 2
header_state = _IDLE
all_tokens: list[int] = []
all_yielded: list[str] = []
token_index = 0
prev_token_count = 0
prev_text = ""
while engine.has_unfinished_requests():
outputs = engine.step()
for output in outputs:
if output.request_id != request_id:
continue
completion = output.outputs[0]
new_tokens = completion.token_ids[prev_token_count:]
finish_reason = completion.finish_reason
prev_token_count = len(completion.token_ids)
prev_text = completion.text
for j, token_id in enumerate(new_tokens):
is_last = j == len(new_tokens) - 1
finished = is_last and finish_reason is not None
vllm_text = tokenizer.decode([token_id])
all_tokens.append(token_id)
if header_state == _EXPECT_MSG and token_id != _MESSAGE_TOKEN:
parser.process(_MESSAGE_TOKEN)
header_state = _IDLE
elif header_state == _EXPECT_MSG:
header_state = _IDLE
elif header_state == _EXPECT_NAME:
header_state = _EXPECT_MSG
if token_id == _CHANNEL_TOKEN:
header_state = _EXPECT_NAME
try:
parser.process(token_id)
except HarmonyError as e:
print(f" !! HarmonyError at token {token_index}: {e}")
engine.abort_request([request_id])
return
delta = parser.last_content_delta
ch = parser.current_channel
recipient = parser.current_recipient
yielded = ""
effective_recipient = recipient if (recipient is not None and recipient.startswith("functions.")) else None
if effective_recipient != current_tool_name:
if current_tool_name is not None:
tool_name = current_tool_name.removeprefix("functions.")
args = "".join(tool_arg_parts).strip()
yielded = f"TOOL_CALL({tool_name}, {args!r})"
tool_arg_parts = []
current_tool_name = effective_recipient
if current_tool_name is not None:
if delta:
tool_arg_parts.append(delta)
if finished:
yielded = f"TOOL_FINISH({json.dumps(''.join(tool_arg_parts))})"
tool_arg_parts = []
else:
is_suppressed = ch == "analysis" or (recipient is not None and recipient.startswith("!"))
if is_suppressed and not thinking:
thinking = True
if not is_suppressed and thinking:
thinking = False
if delta:
prefix = "[THINK] " if thinking else ""
yielded = f"{prefix}{delta!r}"
if finished:
yielded += f" [FINISH={finish_reason}]"
all_yielded.append(yielded)
print(f"{token_index:4d} | {token_id:8d} | {vllm_text!r:>20} | {delta!r:>25} | {str(ch):>12} | {str(recipient):>25} | {yielded:>40}")
token_index += 1
if finish_reason is not None:
print(f"\nFinish reason: {finish_reason}")
print(f"\n--- RAW TOKEN IDS ({len(all_tokens)}) ---")
print(all_tokens)
print("\n--- FINAL TEXT ---")
final_text = ""
for y in all_yielded:
if y.startswith("[THINK] "):
continue
if y.startswith("TOOL_"):
continue
clean = y.replace(" [FINISH=stop]", "").replace(" [FINISH=length]", "")
if clean.startswith("'") and clean.endswith("'") or clean.startswith('"') and clean.endswith('"'):
final_text += clean[1:-1]
print(repr(final_text))
def main():
no_prefix_cache = "--no-prefix-cache" in sys.argv
multi_first = "--multi-first" in sys.argv
print(f"Loading vLLM engine from {MODEL_ID}")
print(f" prefix_caching: {'DISABLED' if no_prefix_cache else 'DEFAULT'}")
print(f" order: {'multi-first' if multi_first else 'single-first'}")
kwargs = {}
if no_prefix_cache:
kwargs["enable_prefix_caching"] = False
engine_args = EngineArgs(
model=MODEL_ID,
gpu_memory_utilization=0.9,
trust_remote_code=True,
load_format="fastsafetensors",
**kwargs,
)
engine = LLMEngine.from_engine_args(engine_args)
tokenizer = engine.get_tokenizer()
eos_ids = getattr(tokenizer, "eos_token_id", None)
print(f"VLLM: Engine loaded. eos_token_id from tokenizer: {eos_ids}")
if multi_first:
run_generation(engine, MESSAGES_MULTI, "MULTI TURN: Hi! -> A lot.", "test-multi")
run_generation(engine, MESSAGES_SINGLE, "SINGLE TURN: Hi!", "test-single")
run_generation(engine, MESSAGES_LONG, "LONG CONVO: 5 turns", "test-long")
else:
run_generation(engine, MESSAGES_SINGLE, "SINGLE TURN: Hi!", "test-single")
run_generation(engine, MESSAGES_MULTI, "MULTI TURN: Hi! -> A lot.", "test-multi")
run_generation(engine, MESSAGES_LONG, "LONG CONVO: 5 turns", "test-long")
print("\n=== ALL DONE ===")
if __name__ == "__main__":
main()
src/exo/worker/engines/mlx/cache.py
+5 -6
View File
@@ -208,9 +208,8 @@ class KVPrefixCache:
return prompt_cache, remaining, best_index
def lookup(self, prompt_token_ids: list[int]) -> tuple["TorchKVCache | None", int, int | None]:
"""Prefix cache lookup returning TorchKVCache directly. For vLLM restore path."""
from exo.worker.engines.kv_cache import TorchKVCache as _TKV
def lookup(self, prompt_token_ids: list[int]) -> tuple["object | None", int, int | None]:
from exo.worker.engines.kv_cache import TorchKVCache
prompt_mx = mx.array(prompt_token_ids)
max_length = len(prompt_token_ids)
@@ -234,13 +233,13 @@ class KVPrefixCache:
self._last_used[best_index] = self._access_counter
cached = self.caches[best_index]
if isinstance(cached, _TKV):
if isinstance(cached, TorchKVCache):
return cached.trim_to(best_length), best_length, best_index
torch_cache = _TKV.from_mlx_cache(cached)
torch_cache = TorchKVCache.from_mlx_cache(cached)
return torch_cache.trim_to(best_length), best_length, best_index
def add_from_torch(self, prompt_token_ids: list[int], cache: "TorchKVCache") -> None:
def add_from_torch(self, prompt_token_ids: list[int], cache: "object") -> None:
"""Store a TorchKVCache directly. For vLLM save path — no MLX conversion."""
self._evict_if_needed()
self.prompts.append(mx.array(prompt_token_ids))
src/exo/worker/engines/vllm/vllm_generator.py
+109 -35
View File
@@ -1,4 +1,5 @@
import gc
import itertools
import json
import re
import sys
@@ -19,11 +20,16 @@ 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
from exo.shared.types.worker.runner_response import (
GenerationResponse,
ToolCallResponse,
)
from exo.utils.channels import MpReceiver, MpSender
from exo.worker.engines.kv_cache import TorchKVCache
from exo.worker.engines.mlx.cache import KVPrefixCache
from exo.worker.engines.vllm.growable_cache import (
@@ -68,6 +74,7 @@ class _ActiveRequest:
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)
@@ -80,10 +87,14 @@ class VllmSequentialGenerator(InferenceGenerator):
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)
_pending: deque[TextGeneration] = field(default_factory=deque, 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:
@@ -101,23 +112,34 @@ class VllmSequentialGenerator(InferenceGenerator):
def submit(self, task: TextGeneration) -> None:
self._cancelled_tasks.discard(CANCEL_ALL_TASKS)
self._pending.append(task)
self._all_tasks[task.task_id] = task
self._maybe_queue.append(task)
def _check_cancellations(self) -> None:
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:
self._cancelled_tasks.add(task_id)
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._check_cancellations()
self.agree_on_cancellations()
if self._active is None and not self._pending:
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()
@@ -125,7 +147,7 @@ class VllmSequentialGenerator(InferenceGenerator):
think_end: str | None = getattr(tokenizer, "think_end", None)
if self._active is None:
task = self._pending.popleft()
task = self._queue.popleft()
if self.should_cancel(task.task_id):
self._cancelled_tasks.discard(task.task_id)
return [(task.task_id, Cancelled())]
@@ -172,10 +194,18 @@ class VllmSequentialGenerator(InferenceGenerator):
self._active = None
return [(active.task.task_id, Finished())]
# --- Prefill phase: step engine until first token arrives ---
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
@@ -184,6 +214,15 @@ class VllmSequentialGenerator(InferenceGenerator):
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:
break
if not active.prefill_done:
@@ -366,10 +405,14 @@ class VllmBatchGenerator(InferenceGenerator):
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)
_pending: deque[TextGeneration] = field(default_factory=deque, 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)
def warmup(self) -> None:
@@ -385,14 +428,22 @@ class VllmBatchGenerator(InferenceGenerator):
def submit(self, task: TextGeneration) -> None:
self._cancelled_tasks.discard(CANCEL_ALL_TASKS)
self._pending.append(task)
self._all_tasks[task.task_id] = task
self._maybe_queue.append(task)
def _check_cancellations(self) -> None:
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:
self._cancelled_tasks.add(task_id)
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(
@@ -428,34 +479,44 @@ class VllmBatchGenerator(InferenceGenerator):
parsed_gen=parsed_gen,
)
def _apply_cancellations(
self,
) -> list[tuple[TaskId, Cancelled]]:
if not self._cancelled_tasks:
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._check_cancellations()
self.agree_on_cancellations()
if not self._queue:
self.agree_on_tasks()
results: list[
tuple[TaskId, GenerationResponse | ToolCallResponse | Cancelled | Finished]
] = []
for task_id in list(self._cancelled_tasks):
rid = str(task_id)
if rid in self._active:
self.engine.abort_request([rid])
del self._active[rid]
self._cancelled_tasks.discard(task_id)
results.append((task_id, Cancelled()))
elif CANCEL_ALL_TASKS in self._cancelled_tasks:
for rid, active in list(self._active.items()):
self.engine.abort_request([rid])
results.append((active.task.task_id, Cancelled()))
self._active.clear()
self._cancelled_tasks.discard(CANCEL_ALL_TASKS)
break
while self._pending and len(self._active) < _EXO_MAX_CONCURRENT_VLLM_REQUESTS:
task = self._pending.popleft()
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()))
@@ -464,7 +525,7 @@ class VllmBatchGenerator(InferenceGenerator):
self._active[active.request_id] = active
if not self._active:
return results
return itertools.chain(results, self._apply_cancellations())
outputs = self.engine.step()
tokenizer = self.engine.get_tokenizer()
@@ -543,7 +604,20 @@ class VllmBatchGenerator(InferenceGenerator):
results.append((active.task.task_id, Finished()))
del self._active[rid]
return results
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 itertools.chain(results, self._apply_cancellations())
def _get_coordinator(self) -> object | None:
if not hasattr(self, "_coordinator_cached"):
@@ -690,7 +764,7 @@ def _patch_weight_loading_progress() -> None:
_monkey_patch_iterator(weight_utils, "fastsafetensors_weights_iterator")
import huggingface_hub # pyright: ignore[reportMissingImports]
_noop_metadata = lambda *_a, **_kw: None # pyright: ignore[reportUnknownLambdaType]
def _noop_metadata(*_a: object, **_kw: object) -> None: pass # pyright: ignore[reportUnknownParameterType]
original_metadata = huggingface_hub.get_safetensors_metadata # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]
huggingface_hub.get_safetensors_metadata = _noop_metadata # pyright: ignore[reportAttributeAccessIssue]
for mod in list(sys.modules.values()):
src/exo/worker/runner/bootstrap.py
+1
View File
@@ -80,6 +80,7 @@ def entrypoint(
model_path=str(EXO_MODELS_DIR / model_id.normalize()),
trust_remote_code=bound_instance.bound_shard.model_card.trust_remote_code,
cancel_receiver=cancel_receiver,
event_sender=event_sender,
)
runner = Runner(
bound_instance, event_sender, task_receiver, cancel_receiver, builder
src/exo/worker/runner/llm_inference/runner.py
+3
View File
@@ -485,6 +485,7 @@ class VllmBuilder(Builder):
model_path: str
trust_remote_code: bool
cancel_receiver: MpReceiver[TaskId]
event_sender: MpSender[Event]
group: mx.distributed.Group | None = None
def connect(self, bound_instance: BoundInstance) -> None:
@@ -517,6 +518,7 @@ class VllmBuilder(Builder):
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
@@ -527,6 +529,7 @@ class VllmBuilder(Builder):
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,
)