feat(llm): adaptive num_ctx tiers + fix KV cache priming num_ctx mismatch

Adds pick_num_ctx() which selects the smallest context window tier
(8192, 16384, 32768) that fits the current messages with 25% headroom,
capped at OLLAMA_NUM_CTX. Threads num_ctx through generation_task.py so
every chat request uses the computed tier rather than a fixed 16384.

Fixes a critical cache miss bug: KV cache priming in app.py and
settings.py was sending requests without num_ctx, so Ollama sized the
cache at its model default (different from the 16384 real requests used),
forcing a full model reload on the first real user message. Both priming
sites now call pick_num_ctx() and pass the matching value.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
2026-04-03 11:47:39 -04:00
parent a6888953dc
commit ef55bcb560
4 changed files with 48 additions and 13 deletions
+6 -4
View File
@@ -183,16 +183,18 @@ def create_app() -> Quart:
"""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.
rather than the full ~4,650-token system prompt, cutting TTFT from ~25s to <1s.
The num_ctx must match what real requests will use so Ollama doesn't reload.
"""
try:
from fabledassistant.services.llm import build_context
from fabledassistant.services.llm import build_context, pick_num_ctx
messages, _ = await build_context(
user_id=user_id,
history=[],
current_note_id=None,
user_message=" ",
)
num_ctx = pick_num_ctx(messages)
async with httpx.AsyncClient(timeout=120.0) as client:
await client.post(
f"{Config.OLLAMA_URL}/api/chat",
@@ -200,11 +202,11 @@ def create_app() -> Quart:
"model": model,
"messages": messages,
"stream": False,
"options": {"num_predict": 1},
"options": {"num_predict": 1, "num_ctx": num_ctx},
"keep_alive": "2h",
},
)
logger.info("Primed KV cache for user %d with model '%s'", user_id, model)
logger.info("Primed KV cache for user %d with model '%s' (num_ctx=%d)", user_id, model, num_ctx)
except Exception:
logger.warning("Failed to prime KV cache for user %d", user_id, exc_info=True)
+4 -3
View File
@@ -15,7 +15,7 @@ logger = logging.getLogger(__name__)
async def _prime_kv_cache_bg(user_id: int, model: str) -> None:
"""Fire-and-forget: prime Ollama's KV cache with the user's system prompt."""
import httpx
from fabledassistant.services.llm import build_context
from fabledassistant.services.llm import build_context, pick_num_ctx
try:
messages, _ = await build_context(
user_id=user_id,
@@ -23,6 +23,7 @@ async def _prime_kv_cache_bg(user_id: int, model: str) -> None:
current_note_id=None,
user_message=" ",
)
num_ctx = pick_num_ctx(messages)
async with httpx.AsyncClient(timeout=120.0) as client:
await client.post(
f"{Config.OLLAMA_URL}/api/chat",
@@ -30,11 +31,11 @@ async def _prime_kv_cache_bg(user_id: int, model: str) -> None:
"model": model,
"messages": messages,
"stream": False,
"options": {"num_predict": 1},
"options": {"num_predict": 1, "num_ctx": num_ctx},
"keep_alive": "2h",
},
)
logger.info("Primed KV cache for user %d with model '%s'", user_id, model)
logger.info("Primed KV cache for user %d with model '%s' (num_ctx=%d)", user_id, model, num_ctx)
except Exception:
logger.warning("Failed to prime KV cache for user %d", user_id, exc_info=True)
@@ -22,7 +22,7 @@ from fabledassistant.services.generation_buffer import (
GenerationBuffer,
GenerationState,
)
from fabledassistant.services.llm import ChatChunk, build_context, generate_completion, stream_chat, stream_chat_with_tools, summarize_history_for_context
from fabledassistant.services.llm import ChatChunk, build_context, generate_completion, pick_num_ctx, stream_chat, stream_chat_with_tools, summarize_history_for_context
from fabledassistant.services.chat import update_conversation_title
from fabledassistant.services.settings import get_setting
from fabledassistant.services.logging import log_generation
@@ -168,6 +168,7 @@ async def _stream_with_retry(
model: str,
tools: list[dict],
think: bool,
num_ctx: int | None = None,
) -> AsyncGenerator[ChatChunk, None]:
"""stream_chat_with_tools with automatic retry on Ollama 500 errors.
@@ -184,7 +185,7 @@ async def _stream_with_retry(
)
await asyncio.sleep(delay)
try:
async for chunk in stream_chat_with_tools(messages, model, tools=tools, think=think):
async for chunk in stream_chat_with_tools(messages, model, tools=tools, think=think, num_ctx=num_ctx):
yield chunk
return
except httpx.HTTPStatusError as exc:
@@ -260,6 +261,11 @@ async def run_generation(
voice_speech_style=voice_speech_style,
)
# Pick the smallest context tier that fits the current messages.
# Using the minimum needed tier reduces KV cache size and speeds up prefill.
num_ctx = pick_num_ctx(messages)
logger.debug("Adaptive num_ctx=%d for conv %d", num_ctx, conv_id)
# Emit context event
buf.append_event("context", {"context": context_meta})
@@ -269,6 +275,7 @@ async def run_generation(
t_start = time.monotonic()
timing: dict = {
"think": think,
"num_ctx": num_ctx,
"tools": [],
"rounds": 0,
"prompt_tokens": None,
@@ -298,7 +305,7 @@ async def run_generation(
buf.append_event("status", {"status": "Generating response..." if _round == 0 else "Composing response..."})
t_stream = time.monotonic()
async for chunk in _stream_with_retry(messages, model, tools, think):
async for chunk in _stream_with_retry(messages, model, tools, think, num_ctx=num_ctx):
if buf.cancel_event.is_set():
cancelled = True
break
@@ -535,7 +542,7 @@ async def run_assist_generation(
await asyncio.sleep(delay)
try:
buf.content_so_far = ""
async for chunk in stream_chat(messages, model, options={"num_predict": Config.OLLAMA_NUM_CTX}):
async for chunk in stream_chat(messages, model, options={"num_predict": num_ctx}, num_ctx=num_ctx):
buf.content_so_far += chunk
buf.append_event("chunk", {"chunk": chunk})
+27 -2
View File
@@ -19,6 +19,28 @@ from fabledassistant.services.settings import get_setting
logger = logging.getLogger(__name__)
# Context window tiers. The smallest tier that fits the current input is used
# so Ollama allocates a smaller KV cache, reducing prefill time and VRAM usage.
# Requests using the same tier hit Ollama's prefix cache; a tier upgrade causes
# a one-time model reload but then the larger cache stays warm.
_CTX_TIERS = (8192, 16384, 32768)
def pick_num_ctx(messages: list[dict]) -> int:
"""Return the smallest context tier that fits *messages* with 25% headroom.
Stays at or below Config.OLLAMA_NUM_CTX (the configured ceiling).
"""
total_chars = sum(len(m.get("content") or "") for m in messages)
estimated_tokens = int(total_chars / 3.5)
needed = int(estimated_tokens * 1.25) + 256 # 25% headroom + output buffer
cap = Config.OLLAMA_NUM_CTX
for tier in _CTX_TIERS:
if tier >= needed and tier <= cap:
return tier
return cap
STOP_WORDS = frozenset({
"a", "an", "the", "is", "it", "to", "in", "for", "of", "and", "or",
"on", "at", "by", "with", "from", "as", "be", "was", "were", "been",
@@ -112,6 +134,7 @@ async def stream_chat(
model: str,
options: dict | None = None,
think: bool = False,
num_ctx: int | None = None,
) -> AsyncGenerator[str, None]:
"""Stream chat completion from Ollama, yielding content chunks.
@@ -119,7 +142,7 @@ async def stream_chat(
Thinking tokens are silently discarded anyway, but disabling avoids the
multi-minute delay before the first content token arrives.
"""
merged_options = {"num_ctx": Config.OLLAMA_NUM_CTX}
merged_options = {"num_ctx": num_ctx or Config.OLLAMA_NUM_CTX}
if options:
merged_options.update(options)
payload: dict = {"model": model, "messages": messages, "stream": True, "options": merged_options, "think": think, "keep_alive": "2h"}
@@ -159,6 +182,7 @@ async def stream_chat_with_tools(
model: str,
tools: list[dict] | None = None,
think: bool = False,
num_ctx: int | None = None,
) -> AsyncGenerator[ChatChunk, None]:
"""Stream chat completion from Ollama with tool support.
@@ -170,7 +194,8 @@ async def stream_chat_with_tools(
Thinking tokens are consumed by Ollama and not forwarded to the caller;
only the final response content is yielded. Expect higher TTFT when enabled.
"""
options: dict = {"num_ctx": Config.OLLAMA_NUM_CTX}
resolved_ctx = num_ctx or Config.OLLAMA_NUM_CTX
options: dict = {"num_ctx": resolved_ctx}
if tools:
options["num_predict"] = 8192
payload: dict = {