Unnecessary further optimizations 3

src/exo/worker/engines/vllm/growable_cache.py

@@ -56,6 +56,13 @@ def _patch_determine_available_memory() -> None:
 
     @torch.inference_mode()
     def patched(self: "Worker") -> int:
+        import pathlib
+        import shutil
+
+        compile_cache = pathlib.Path.home() / ".cache" / "vllm" / "torch_compile_cache"
+        if compile_cache.exists():
+            shutil.rmtree(compile_cache, ignore_errors=True)
+
         real_empty_cache = torch.cuda.empty_cache
         torch.cuda.empty_cache = lambda: None  # type: ignore
         try:
@@ -243,7 +250,6 @@ def _grow_tensors(
         model_runner.compilation_config.static_forward_context
     )  # type: ignore
     runner_kv_caches: list[torch.Tensor] = model_runner.kv_caches  # type: ignore
-    runner_kv_caches.clear()
 
     from collections import defaultdict
 
@@ -258,12 +264,37 @@ def _grow_tensors(
     for ln in new_kv_caches:
         index2name[extract_layer_index(ln, num_attn_module)].append(ln)
 
+    new_ordered: list[torch.Tensor | list[torch.Tensor]] = []
     for layer_index in sorted(index2name.keys()):
         for ln in index2name[layer_index]:
-            runner_kv_caches.append(new_kv_caches[ln])
+            new_ordered.append(new_kv_caches[ln])
 
-    for layer_name, kv_cache in new_kv_caches.items():
-        forward_context[layer_name].kv_cache = [kv_cache]  # type: ignore
+    for i, new_kv in enumerate(new_ordered):
+        if i < len(runner_kv_caches):
+            old_kv = runner_kv_caches[i]
+            if isinstance(old_kv, list) and isinstance(new_kv, list):
+                for j, (old_t, new_t) in enumerate(zip(old_kv, new_kv)):
+                    old_t.set_(new_t.storage(), new_t.storage_offset(), new_t.shape, new_t.stride())  # type: ignore
+            elif isinstance(old_kv, torch.Tensor) and isinstance(new_kv, torch.Tensor):
+                old_kv.set_(new_kv.storage(), new_kv.storage_offset(), new_kv.shape, new_kv.stride())  # type: ignore
+            else:
+                runner_kv_caches[i] = new_kv
+        else:
+            runner_kv_caches.append(new_kv)
+
+    for layer_name, new_kv in new_kv_caches.items():
+        old_kv_list = forward_context[layer_name].kv_cache  # type: ignore
+        if old_kv_list and len(old_kv_list) > 0:
+            old_entry = old_kv_list[0]
+            if isinstance(old_entry, list) and isinstance(new_kv, list):
+                for j, (old_t, new_t) in enumerate(zip(old_entry, new_kv)):
+                    old_t.set_(new_t.storage(), new_t.storage_offset(), new_t.shape, new_t.stride())  # type: ignore
+            elif isinstance(old_entry, torch.Tensor) and isinstance(new_kv, torch.Tensor):
+                old_entry.set_(new_kv.storage(), new_kv.storage_offset(), new_kv.shape, new_kv.stride())  # type: ignore
+            else:
+                forward_context[layer_name].kv_cache = [new_kv]  # type: ignore
+        else:
+            forward_context[layer_name].kv_cache = [new_kv]  # type: ignore
 
 
 def _grow_block_pool(

src/exo/worker/engines/vllm/vllm_generator.py

@@ -3,6 +3,7 @@ import math
 import os
 import re
 import sys
+import tempfile
 import time
 from collections.abc import Callable, Generator
 from dataclasses import dataclass, field

src/exo/worker/runner/bootstrap.py

@@ -2,6 +2,7 @@ import ctypes
 import os
 import resource
 import sys
+import urllib.request
 from pathlib import Path
 
 import loguru
@@ -14,6 +15,9 @@ from exo.utils.channels import ClosedResourceError, MpReceiver, MpSender
 
 logger: "loguru.Logger" = loguru.logger
 
+_TIKTOKEN_BASE_URL = "https://openaipublic.blob.core.windows.net/encodings"
+_TIKTOKEN_FILES = ["o200k_base.tiktoken", "cl100k_base.tiktoken"]
+
 _CUDA_HOST_LIBS = ["libcuda.so.1", "libnvidia-ml.so.1", "libnvidia-ptxjitcompiler.so.1"]
 _CUDA_HOST_SEARCH_DIRS = [
     Path("/usr/lib/aarch64-linux-gnu"),
@@ -25,6 +29,28 @@ _CUDA_HOST_SEARCH_DIRS = [
 ]
 
 
+def _ensure_tiktoken_encodings() -> None:
+    if os.environ.get("TIKTOKEN_ENCODINGS_BASE"):
+        return
+    from exo.shared.constants import EXO_CACHE_HOME
+
+    enc_dir = EXO_CACHE_HOME / "encodings"
+    enc_dir.mkdir(parents=True, exist_ok=True)
+    for fname in _TIKTOKEN_FILES:
+        dest = enc_dir / fname
+        if dest.exists():
+            continue
+        url = f"{_TIKTOKEN_BASE_URL}/{fname}"
+        logger.info(f"Downloading {url} -> {dest}")
+        try:
+            urllib.request.urlretrieve(url, dest)
+        except Exception:
+            logger.warning(f"Failed to download {fname}, harmony encoding may fail")
+            return
+    os.environ["TIKTOKEN_ENCODINGS_BASE"] = str(enc_dir)
+    logger.info(f"Set TIKTOKEN_ENCODINGS_BASE={enc_dir}")
+
+
 def _ensure_cuda_libs() -> None:
     if sys.platform != "linux":
         return
@@ -80,6 +106,7 @@ def entrypoint(
             os.environ["VLLM_KV_CACHE_LAYOUT"] = "NHD"
             # os.environ["VLLM_BATCH_INVARIANT"] = "1"
             _ensure_cuda_libs()
+            _ensure_tiktoken_encodings()
             from exo.shared.constants import EXO_MODELS_DIR
             from exo.worker.runner.llm_inference.runner import Runner, VllmBuilder
 

src/exo/worker/runner/runner_supervisor.py

@@ -81,7 +81,14 @@ class RunnerSupervisor:
         task_sender, task_recv = mp_channel[Task]()
         cancel_sender, cancel_recv = mp_channel[TaskId]()
 
-        runner_process = mp.Process(
+        from exo.shared.types.worker.instances import VllmInstance
+
+        # vLLM runners use "spawn" to avoid inheriting the parent's CUDA state.
+        # With "fork", the parent's partial CUDA init (from device detection) is
+        # inherited by the child, which conflicts with torch.compile's inductor
+        # backend (cudagraph_mode=none) and causes CUDA illegal instruction errors.
+        ctx = mp.get_context("spawn") if isinstance(bound_instance.instance, VllmInstance) else mp
+        runner_process = ctx.Process(
             target=entrypoint,
             args=(
                 bound_instance,