perf(startup): prime Ollama KV cache with system prompt on warm-up

After loading each user's chat model into VRAM, send a minimal chat
request with the real system prompt (num_predict=1) to populate the
KV cache. The first real user message then only needs to process its
own tokens rather than the full ~5,600-token system prompt, reducing
cold-start TTFT from ~25s to <1s.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
2026-04-03 10:47:59 -04:00
parent 3bd0dc6879
commit 4792e6459b
+43 -9
View File
@@ -179,34 +179,64 @@ def create_app() -> Quart:
except Exception:
logger.warning("Failed to warm model '%s'", model, exc_info=True)
async def _prime_kv_cache(user_id: int, model: str) -> None:
"""Send a minimal chat request to prime Ollama's KV cache with the user's system prompt.
This ensures the next real user message only needs to process its own tokens
rather than the full ~5,600-token system prompt, cutting TTFT from ~25s to <1s.
"""
try:
from fabledassistant.services.llm import build_context
messages, _ = await build_context(
user_id=user_id,
history=[],
current_note_id=None,
user_message=" ",
)
async with httpx.AsyncClient(timeout=120.0) as client:
await client.post(
f"{Config.OLLAMA_URL}/api/chat",
json={
"model": model,
"messages": messages,
"stream": False,
"options": {"num_predict": 1},
"keep_alive": "2h",
},
)
logger.info("Primed KV cache for user %d with model '%s'", user_id, model)
except Exception:
logger.warning("Failed to prime KV cache for user %d", user_id, exc_info=True)
async def _warm_user_models() -> None:
"""
Warm whichever chat model(s) users have selected in Settings.
Warm whichever chat model(s) users have selected in Settings, then prime
the KV cache with each user's system prompt so the first real message is fast.
Only warms models that are already installed in Ollama — never auto-pulls.
Falls back silently if no user preferences exist or Ollama is unreachable.
"""
from sqlalchemy import select as sa_select, distinct
from sqlalchemy import select as sa_select
from fabledassistant.models import async_session
from fabledassistant.models.setting import Setting
# 1. Collect all distinct default_model values users have saved.
# 1. Collect (user_id, model) pairs for all users with a saved default_model.
try:
async with async_session() as session:
rows = await session.execute(
sa_select(distinct(Setting.value)).where(
sa_select(Setting.user_id, Setting.value).where(
Setting.key == "default_model",
Setting.value.isnot(None),
Setting.value != "",
)
)
user_models: set[str] = {r for (r,) in rows}
user_model_pairs: list[tuple[int, str]] = list(rows)
except Exception:
logger.debug("Could not read user model preferences from DB", exc_info=True)
return
if not user_models:
if not user_model_pairs:
logger.debug("No user model preferences found; skipping warm-up")
return
@@ -220,11 +250,15 @@ def create_app() -> Quart:
logger.debug("Could not reach Ollama to check installed models", exc_info=True)
return
# 3. Warm only the intersection (installed AND user-preferred).
for model in user_models:
# 3. Warm each unique model, then prime KV cache per user.
warmed: set[str] = set()
for user_id_val, model in user_model_pairs:
base = model.removesuffix(":latest")
if model in installed or f"{base}:latest" in installed or base in installed:
await _warm_model(model)
if model not in warmed:
await _warm_model(model)
warmed.add(model)
await _prime_kv_cache(user_id_val, model)
else:
logger.info(
"User-preferred model '%s' is not installed; skipping warm-up "