fix(server): per-slot byte cap on context checkpoints (closes #67)

common/arg.cpp

@@ -1337,6 +1337,16 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
             params.checkpoint_min_step = value;
         }
     ).set_env("LLAMA_ARG_CHECKPOINT_MIN_SPACING_NT").set_examples({LLAMA_EXAMPLE_SERVER}));
+    add_opt(common_arg(
+        {"--ctx-checkpoints-max-mib"}, "N",
+        string_format("max host-RAM budget per slot for context checkpoints in MiB (default: %d, 0 = no byte cap, count-only). Eviction is FIFO: oldest checkpoint is dropped first to honour either the count cap (--ctx-checkpoints) or this byte cap, whichever bites.", params.ctx_checkpoints_max_mib),
+        [](common_params & params, int value) {
+            if (value < 0) {
+                throw std::invalid_argument("ctx-checkpoints-max-mib must be non-negative");
+            }
+            params.ctx_checkpoints_max_mib = value;
+        }
+    ).set_env("LLAMA_ARG_CTX_CHECKPOINTS_MAX_MIB").set_examples({LLAMA_EXAMPLE_SERVER}));
     add_opt(common_arg(
         {"-cram", "--cache-ram"}, "N",
         string_format("set the maximum cache size in MiB (default: %d, -1 - no limit, 0 - disable)"

common/common.h

@@ -600,8 +600,9 @@ struct common_params {
     int32_t n_cache_reuse       = 0;     // min chunk size to reuse from the cache via KV shifting
     bool    cache_prompt        = true;  // whether to enable prompt caching
     bool    cache_idle_slots    = true;  // save and clear idle slots upon starting a new task
-    int32_t n_ctx_checkpoints   = 32;    // max number of context checkpoints per slot
-    int32_t checkpoint_min_step = 256;   // minimum spacing between context checkpoints
+    int32_t n_ctx_checkpoints        = 32;    // max number of context checkpoints per slot
+    int32_t checkpoint_min_step      = 256;   // minimum spacing between context checkpoints
+    int32_t ctx_checkpoints_max_mib  = 4096;  // max host-RAM budget per slot for context checkpoints (MiB); 0 = no byte cap (count-only)
     int32_t cache_ram_mib       = 8192;  // -1 = no limit, 0 - disable, 1 = 1 MiB, etc.
 
     std::string hostname      = "127.0.0.1";

tools/server/server-context.cpp

@@ -2076,12 +2076,44 @@ struct server_context_impl {
 
     // n_tokens_cur: the number of tokens added to the batch for the current slot
     void create_checkpoint(server_slot & slot, const int64_t n_tokens_cur, llama_pos pos_min, llama_pos pos_max) {
-        while (slot.prompt.checkpoints.size() >= (size_t) params_base.n_ctx_checkpoints) {
-            // make room for the new checkpoint, if needed
-            const auto & cur = slot.prompt.checkpoints.front();
+        // n_ctx_checkpoints <= 0 disables checkpoints entirely. The arg parser
+        // accepts 0 (and any negative value wraps via the size_t cast below to
+        // a huge cap, which is also a no-op), so we short-circuit here rather
+        // than letting create_checkpoint's eviction loop touch an empty list.
+        if (params_base.n_ctx_checkpoints <= 0) {
+            return;
+        }
+
+        // Per-slot footprint accounting (issue #67): the count-only cap let a single slot
+        // accumulate ~20 GB of checkpoints under heierchat's long contexts, driving titan
+        // SystemOOM. The byte cap is FIFO-evicted alongside the count cap; whichever bites
+        // first is the active limit.
+        auto total_bytes = [&]() {
+            size_t b = 0;
+            for (const auto & ckpt : slot.prompt.checkpoints) {
+                b += ckpt.size();
+            }
+            return b;
+        };
+
+        const size_t byte_cap =
+            (params_base.ctx_checkpoints_max_mib > 0)
+                ? (size_t) params_base.ctx_checkpoints_max_mib * 1024 * 1024
+                : 0; // 0 → byte cap disabled, count-only behavior
+
+        auto over_count = [&]() {
+            return slot.prompt.checkpoints.size() >= (size_t) params_base.n_ctx_checkpoints;
+        };
+        auto over_bytes = [&]() {
+            return byte_cap > 0 && !slot.prompt.checkpoints.empty() && total_bytes() >= byte_cap;
+        };
+
+        while (over_count() || over_bytes()) {
+            const auto & front = slot.prompt.checkpoints.front();
+            const char * reason = (over_bytes() && !over_count()) ? "bytes" : "count";
 
-            SLT_WRN(slot, "erasing old context checkpoint (pos_min = %d, pos_max = %d, n_tokens = %" PRId64 ", size = %.3f MiB)\n",
-                    cur.pos_min, cur.pos_max, cur.n_tokens, (float) cur.size() / 1024 / 1024);
+            SLT_WRN(slot, "erasing old context checkpoint (reason=%s, pos_min = %d, pos_max = %d, n_tokens = %" PRId64 ", size = %.3f MiB)\n",
+                    reason, front.pos_min, front.pos_max, front.n_tokens, (float) front.size() / 1024 / 1024);
 
             slot.prompt.checkpoints.erase(slot.prompt.checkpoints.begin());
         }
@@ -2094,9 +2126,11 @@ struct server_context_impl {
         cur.update_dft(ctx_dft.get(), slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);
 
         SLT_INF(slot,
-                "created context checkpoint %d of %d (pos_min = %d, pos_max = %d, n_tokens = %" PRId64 ", size = %.3f MiB)\n",
+                "created context checkpoint %d of %d (pos_min = %d, pos_max = %d, n_tokens = %" PRId64 ", size = %.3f MiB, slot total = %.3f MiB / %d MiB cap)\n",
                 (int) slot.prompt.checkpoints.size(), params_base.n_ctx_checkpoints, cur.pos_min,
-                cur.pos_max, cur.n_tokens, (float) cur.size() / 1024 / 1024);
+                cur.pos_max, cur.n_tokens, (float) cur.size() / 1024 / 1024,
+                (float) total_bytes() / 1024 / 1024,
+                params_base.ctx_checkpoints_max_mib);
     }
 
     void process_single_task(server_task && task) {

tools/server/tests/unit/test_ctx_checkpoints_bytes_cap.py

@@ -0,0 +1,142 @@
+import os
+import tempfile
+import pytest
+from utils import *
+
+server = ServerPreset.tinyllama2()
+
+
+class LogReader:
+    def __init__(self, path):
+        self.path = path
+        self.pos = 0
+
+    def drain(self):
+        with open(self.path) as f:
+            f.seek(self.pos)
+            content = f.read()
+            self.pos = f.tell()
+        return content
+
+
+@pytest.fixture(autouse=True)
+def create_server():
+    global server
+    server = ServerPreset.tinyllama2()
+    server.n_slots = 1
+    server.n_ctx = 512
+    server.n_predict = 16
+    server.temperature = 0.0
+    server.server_slots = True
+    server.kv_unified = True
+    server.debug = True
+    fd, server.log_path = tempfile.mkstemp(suffix='.log')
+    os.close(fd)
+    yield
+
+
+CHAT_TURN_1 = [
+    {"role": "user", "content": "Once upon a time"},
+]
+CHAT_TURN_2 = CHAT_TURN_1 + [
+    {"role": "assistant", "content": "there was a small wooden boat that floated down a calm river under tall pine trees."},
+    {"role": "user", "content": "Continue the story for several more sentences."},
+]
+CHAT_TURN_3 = CHAT_TURN_2 + [
+    {"role": "assistant", "content": "The boat passed a sleeping fox and a curious owl perched on a low branch above the water."},
+    {"role": "user", "content": "Now describe what the fox was dreaming about."},
+]
+
+
+def _post_chat(messages):
+    return server.make_request("POST", "/v1/chat/completions", data={
+        "messages": messages,
+        "max_tokens": 16,
+        "cache_prompt": True,
+    })
+
+
+def test_default_args_dont_break_checkpoints():
+    # No new flag set — defaults must keep the existing count-only behavior.
+    global server
+    server.start()
+    log = LogReader(server.log_path)
+
+    # warm + multi-turn so at least one checkpoint may be created
+    for messages in (CHAT_TURN_1, CHAT_TURN_2, CHAT_TURN_3):
+        res = _post_chat(messages)
+        assert res.status_code == 200, res.body
+
+    # If any checkpoint is created, the success log mentions the per-slot footprint marker
+    # introduced alongside the byte cap. If no checkpoint was created in this short run,
+    # the test still passes — we only assert that the server didn't error.
+    drained = log.drain()
+    if "created context checkpoint" in drained:
+        assert "MiB cap" in drained, "byte-cap footprint marker missing from create_checkpoint log line"
+
+
+def test_byte_cap_disabled_when_zero():
+    # Setting --ctx-checkpoints-max-mib 0 must keep the legacy count-only behavior
+    # without erroring at startup.
+    global server
+    server.ctx_checkpoints_max_mib = 0
+    server.start()
+
+    res = _post_chat(CHAT_TURN_1)
+    assert res.status_code == 200, res.body
+
+
+def test_negative_byte_cap_rejected():
+    # arg parser must reject negative values.
+    global server
+    server.ctx_checkpoints_max_mib = -1
+    with pytest.raises(Exception):
+        server.start()
+
+
+def test_ctx_checkpoints_zero_disables_creation():
+    # --ctx-checkpoints 0 must not crash create_checkpoint (UB on
+    # .front()/.erase(begin()) of an empty list under the previous code).
+    # Expected behavior is "no checkpoints are ever created" — slot still works.
+    global server
+    server.n_ctx_checkpoints = 0
+    server.start()
+    log = LogReader(server.log_path)
+
+    for messages in (CHAT_TURN_1, CHAT_TURN_2, CHAT_TURN_3):
+        res = _post_chat(messages)
+        assert res.status_code == 200, res.body
+
+    drained = log.drain()
+    assert "created context checkpoint" not in drained, (
+        "no checkpoints should be created when --ctx-checkpoints=0"
+    )
+
+
+def test_byte_cap_eviction_reason_bytes():
+    # Force the byte cap to bite first (count cap large, byte cap tiny) and verify
+    # the eviction reason is reported as `bytes` when create_checkpoint fires.
+    # This test is best-effort: if tinyllama2 doesn't accumulate any checkpoints
+    # at the chosen ctx/min-step in this short conversation, we skip rather than
+    # flake.
+    global server
+    server.n_ctx = 1024
+    server.n_ctx_checkpoints = 100      # count cap effectively disabled
+    server.checkpoint_min_step = 16     # easier to trigger multiple checkpoints
+    server.ctx_checkpoints_max_mib = 1  # tiny budget, expect bytes-eviction
+    server.start()
+    log = LogReader(server.log_path)
+
+    for messages in (CHAT_TURN_1, CHAT_TURN_2, CHAT_TURN_3):
+        res = _post_chat(messages)
+        assert res.status_code == 200, res.body
+
+    drained = log.drain()
+    if "context checkpoint" not in drained:
+        pytest.skip("no checkpoints created in this short conversation — cap path not exercised")
+    if "erasing old context checkpoint" not in drained and "created context checkpoint" not in drained:
+        pytest.skip("no eviction happened — single checkpoint stayed under the cap")
+
+    assert "reason=bytes" in drained, (
+        f"eviction occurred but reason!=bytes; full log: {drained[-2000:]}"
+    )

tools/server/tests/utils.py

@@ -106,6 +106,9 @@ class ServerProcess:
     sleep_idle_seconds: int | None = None
     cache_ram: int | None = None
     no_cache_idle_slots: bool = False
+    n_ctx_checkpoints: int | None = None
+    checkpoint_min_step: int | None = None
+    ctx_checkpoints_max_mib: int | None = None
     log_path: str | None = None
     webui_mcp_proxy: bool = False
     backend_sampling: bool = False
@@ -251,6 +254,12 @@ def start(self, timeout_seconds: int = DEFAULT_HTTP_TIMEOUT) -> None:
             server_args.extend(["--cache-ram", self.cache_ram])
         if self.no_cache_idle_slots:
             server_args.append("--no-cache-idle-slots")
+        if self.n_ctx_checkpoints is not None:
+            server_args.extend(["--ctx-checkpoints", self.n_ctx_checkpoints])
+        if self.checkpoint_min_step is not None:
+            server_args.extend(["--checkpoint-min-step", self.checkpoint_min_step])
+        if self.ctx_checkpoints_max_mib is not None:
+            server_args.extend(["--ctx-checkpoints-max-mib", self.ctx_checkpoints_max_mib])
         if self.webui_mcp_proxy:
             server_args.append("--webui-mcp-proxy")
         if self.backend_sampling: