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:
@@ -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.
|
"""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
|
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:
|
try:
|
||||||
from fabledassistant.services.llm import build_context
|
from fabledassistant.services.llm import build_context, pick_num_ctx
|
||||||
messages, _ = await build_context(
|
messages, _ = await build_context(
|
||||||
user_id=user_id,
|
user_id=user_id,
|
||||||
history=[],
|
history=[],
|
||||||
current_note_id=None,
|
current_note_id=None,
|
||||||
user_message=" ",
|
user_message=" ",
|
||||||
)
|
)
|
||||||
|
num_ctx = pick_num_ctx(messages)
|
||||||
async with httpx.AsyncClient(timeout=120.0) as client:
|
async with httpx.AsyncClient(timeout=120.0) as client:
|
||||||
await client.post(
|
await client.post(
|
||||||
f"{Config.OLLAMA_URL}/api/chat",
|
f"{Config.OLLAMA_URL}/api/chat",
|
||||||
@@ -200,11 +202,11 @@ def create_app() -> Quart:
|
|||||||
"model": model,
|
"model": model,
|
||||||
"messages": messages,
|
"messages": messages,
|
||||||
"stream": False,
|
"stream": False,
|
||||||
"options": {"num_predict": 1},
|
"options": {"num_predict": 1, "num_ctx": num_ctx},
|
||||||
"keep_alive": "2h",
|
"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:
|
except Exception:
|
||||||
logger.warning("Failed to prime KV cache for user %d", user_id, exc_info=True)
|
logger.warning("Failed to prime KV cache for user %d", user_id, exc_info=True)
|
||||||
|
|
||||||
|
|||||||
@@ -15,7 +15,7 @@ logger = logging.getLogger(__name__)
|
|||||||
async def _prime_kv_cache_bg(user_id: int, model: str) -> None:
|
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."""
|
"""Fire-and-forget: prime Ollama's KV cache with the user's system prompt."""
|
||||||
import httpx
|
import httpx
|
||||||
from fabledassistant.services.llm import build_context
|
from fabledassistant.services.llm import build_context, pick_num_ctx
|
||||||
try:
|
try:
|
||||||
messages, _ = await build_context(
|
messages, _ = await build_context(
|
||||||
user_id=user_id,
|
user_id=user_id,
|
||||||
@@ -23,6 +23,7 @@ async def _prime_kv_cache_bg(user_id: int, model: str) -> None:
|
|||||||
current_note_id=None,
|
current_note_id=None,
|
||||||
user_message=" ",
|
user_message=" ",
|
||||||
)
|
)
|
||||||
|
num_ctx = pick_num_ctx(messages)
|
||||||
async with httpx.AsyncClient(timeout=120.0) as client:
|
async with httpx.AsyncClient(timeout=120.0) as client:
|
||||||
await client.post(
|
await client.post(
|
||||||
f"{Config.OLLAMA_URL}/api/chat",
|
f"{Config.OLLAMA_URL}/api/chat",
|
||||||
@@ -30,11 +31,11 @@ async def _prime_kv_cache_bg(user_id: int, model: str) -> None:
|
|||||||
"model": model,
|
"model": model,
|
||||||
"messages": messages,
|
"messages": messages,
|
||||||
"stream": False,
|
"stream": False,
|
||||||
"options": {"num_predict": 1},
|
"options": {"num_predict": 1, "num_ctx": num_ctx},
|
||||||
"keep_alive": "2h",
|
"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:
|
except Exception:
|
||||||
logger.warning("Failed to prime KV cache for user %d", user_id, exc_info=True)
|
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,
|
GenerationBuffer,
|
||||||
GenerationState,
|
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.chat import update_conversation_title
|
||||||
from fabledassistant.services.settings import get_setting
|
from fabledassistant.services.settings import get_setting
|
||||||
from fabledassistant.services.logging import log_generation
|
from fabledassistant.services.logging import log_generation
|
||||||
@@ -168,6 +168,7 @@ async def _stream_with_retry(
|
|||||||
model: str,
|
model: str,
|
||||||
tools: list[dict],
|
tools: list[dict],
|
||||||
think: bool,
|
think: bool,
|
||||||
|
num_ctx: int | None = None,
|
||||||
) -> AsyncGenerator[ChatChunk, None]:
|
) -> AsyncGenerator[ChatChunk, None]:
|
||||||
"""stream_chat_with_tools with automatic retry on Ollama 500 errors.
|
"""stream_chat_with_tools with automatic retry on Ollama 500 errors.
|
||||||
|
|
||||||
@@ -184,7 +185,7 @@ async def _stream_with_retry(
|
|||||||
)
|
)
|
||||||
await asyncio.sleep(delay)
|
await asyncio.sleep(delay)
|
||||||
try:
|
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
|
yield chunk
|
||||||
return
|
return
|
||||||
except httpx.HTTPStatusError as exc:
|
except httpx.HTTPStatusError as exc:
|
||||||
@@ -260,6 +261,11 @@ async def run_generation(
|
|||||||
voice_speech_style=voice_speech_style,
|
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
|
# Emit context event
|
||||||
buf.append_event("context", {"context": context_meta})
|
buf.append_event("context", {"context": context_meta})
|
||||||
|
|
||||||
@@ -269,6 +275,7 @@ async def run_generation(
|
|||||||
t_start = time.monotonic()
|
t_start = time.monotonic()
|
||||||
timing: dict = {
|
timing: dict = {
|
||||||
"think": think,
|
"think": think,
|
||||||
|
"num_ctx": num_ctx,
|
||||||
"tools": [],
|
"tools": [],
|
||||||
"rounds": 0,
|
"rounds": 0,
|
||||||
"prompt_tokens": None,
|
"prompt_tokens": None,
|
||||||
@@ -298,7 +305,7 @@ async def run_generation(
|
|||||||
buf.append_event("status", {"status": "Generating response..." if _round == 0 else "Composing response..."})
|
buf.append_event("status", {"status": "Generating response..." if _round == 0 else "Composing response..."})
|
||||||
t_stream = time.monotonic()
|
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():
|
if buf.cancel_event.is_set():
|
||||||
cancelled = True
|
cancelled = True
|
||||||
break
|
break
|
||||||
@@ -535,7 +542,7 @@ async def run_assist_generation(
|
|||||||
await asyncio.sleep(delay)
|
await asyncio.sleep(delay)
|
||||||
try:
|
try:
|
||||||
buf.content_so_far = ""
|
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.content_so_far += chunk
|
||||||
buf.append_event("chunk", {"chunk": chunk})
|
buf.append_event("chunk", {"chunk": chunk})
|
||||||
|
|
||||||
|
|||||||
@@ -19,6 +19,28 @@ from fabledassistant.services.settings import get_setting
|
|||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
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({
|
STOP_WORDS = frozenset({
|
||||||
"a", "an", "the", "is", "it", "to", "in", "for", "of", "and", "or",
|
"a", "an", "the", "is", "it", "to", "in", "for", "of", "and", "or",
|
||||||
"on", "at", "by", "with", "from", "as", "be", "was", "were", "been",
|
"on", "at", "by", "with", "from", "as", "be", "was", "were", "been",
|
||||||
@@ -112,6 +134,7 @@ async def stream_chat(
|
|||||||
model: str,
|
model: str,
|
||||||
options: dict | None = None,
|
options: dict | None = None,
|
||||||
think: bool = False,
|
think: bool = False,
|
||||||
|
num_ctx: int | None = None,
|
||||||
) -> AsyncGenerator[str, None]:
|
) -> AsyncGenerator[str, None]:
|
||||||
"""Stream chat completion from Ollama, yielding content chunks.
|
"""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
|
Thinking tokens are silently discarded anyway, but disabling avoids the
|
||||||
multi-minute delay before the first content token arrives.
|
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:
|
if options:
|
||||||
merged_options.update(options)
|
merged_options.update(options)
|
||||||
payload: dict = {"model": model, "messages": messages, "stream": True, "options": merged_options, "think": think, "keep_alive": "2h"}
|
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,
|
model: str,
|
||||||
tools: list[dict] | None = None,
|
tools: list[dict] | None = None,
|
||||||
think: bool = False,
|
think: bool = False,
|
||||||
|
num_ctx: int | None = None,
|
||||||
) -> AsyncGenerator[ChatChunk, None]:
|
) -> AsyncGenerator[ChatChunk, None]:
|
||||||
"""Stream chat completion from Ollama with tool support.
|
"""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;
|
Thinking tokens are consumed by Ollama and not forwarded to the caller;
|
||||||
only the final response content is yielded. Expect higher TTFT when enabled.
|
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:
|
if tools:
|
||||||
options["num_predict"] = 8192
|
options["num_predict"] = 8192
|
||||||
payload: dict = {
|
payload: dict = {
|
||||||
|
|||||||
Reference in New Issue
Block a user