memory-tier: mlock warm KV buffer to prevent kernel swap-out (MAD-179)

Warm tier was a std::vector<uint8_t> sized 3.2 GB at 1M context. On
hosts with low RAM headroom the kernel paged it out during prefill,
causing major-fault thrashing (prefill collapsed to ~13 tok/s while
GPU stayed at 100%).

Fix: introduce LockedBuffer (mmap + best-effort mlock, MAP_LOCKED on
Linux with mlock() fallback, VirtualLock on Windows) and back the
warm tier with it. New CLI flag --kv-tier-warm-mlock (default on)
plumbed through llama_context_params -> llama_memory_params ->
TieredConfig.warm_mlock.

Locking failure is logged with remediation (raise RLIMIT_MEMLOCK)
but does not abort init; the buffer still works, just swappable.

Author: kmbandy

common/arg.cpp

@@ -1418,6 +1418,16 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
             params.kv_warm_device = value;
         }
     ).set_env("LLAMA_ARG_KV_WARM_DEVICE").set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}));
+    add_opt(common_arg(
+        {"--kv-tier-warm-mlock"}, "0|1",
+        "mlock the host-RAM warm tier to prevent kernel swap-out (default: 1). "
+        "Required for predictable performance at long context on RAM-constrained "
+        "hosts. If the lock fails (e.g. RLIMIT_MEMLOCK too small or no "
+        "CAP_IPC_LOCK), allocation still succeeds and a warning is logged.",
+        [](common_params & params, int value) {
+            params.kv_tier_warm_mlock = (value != 0);
+        }
+    ).set_env("LLAMA_ARG_KV_TIER_WARM_MLOCK").set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}));
     add_opt(common_arg(
         {"--kv-tier-total-ctx"}, "N",
         "full ctx budget across all tiers (0 = use n_ctx)",

common/common.cpp

@@ -1541,6 +1541,7 @@ struct llama_context_params common_context_params_to_llama(const common_params &
     cparams.kv_tier_compression         = params.kv_tier_compression;
     cparams.kv_tier_attention_threshold = params.kv_tier_attention_threshold;
     cparams.kv_tier_warm_device         = params.kv_warm_device;
+    cparams.kv_tier_warm_mlock          = params.kv_tier_warm_mlock;
     cparams.kv_tier_total_ctx           = params.kv_tier_total_ctx;
     cparams.kv_tier_semantic_index      = params.kv_semantic_index.empty()
                                               ? nullptr

common/common.h

@@ -602,6 +602,7 @@ struct common_params {
     int kv_tier_compression   = 1;       // 0=none, 1=int4, 2=int8, 3=lz4, 4=quantized
     float kv_tier_attention_threshold = 0.1f;  // attention threshold for eviction
     int   kv_warm_device        = -1;    // HIP device index for warm KV tier (-1 = disabled)
+    bool  kv_tier_warm_mlock    = true;  // mlock the host-RAM warm tier so it's not eligible for kernel swap (see mt-locked-buffer.h)
     int   kv_tier_total_ctx     = 0;     // full ctx budget across all tiers (set at load time)
     std::string kv_semantic_index = "";   // path to embedding model for semantic KV index (empty = disabled)
     float kv_semantic_threshold = 0.65f;  // minimum cosine similarity threshold for prefetch hints

include/llama.h

@@ -402,6 +402,7 @@ extern "C" {
         int32_t      kv_tier_compression;        // mt::Compression
         float        kv_tier_attention_threshold;
         int32_t      kv_tier_warm_device;        // -1 = host RAM
+        bool         kv_tier_warm_mlock;         // mlock the host-RAM warm tier (prevents kernel swap-out)
         int32_t      kv_tier_total_ctx;          // 0 = use n_ctx
         const char * kv_tier_semantic_index;     // nullptr => disabled
         float        kv_tier_semantic_threshold;

src/llama-context.cpp

@@ -295,6 +295,7 @@ llama_context::llama_context(
             /*.kv_tier_compression         =*/ params.kv_tier_compression,
             /*.kv_tier_attention_threshold =*/ params.kv_tier_attention_threshold,
             /*.kv_tier_warm_device         =*/ params.kv_tier_warm_device,
+            /*.kv_tier_warm_mlock          =*/ params.kv_tier_warm_mlock,
             /*.kv_tier_total_ctx           =*/ params.kv_tier_total_ctx,
             /*.kv_tier_semantic_index      =*/ params.kv_tier_semantic_index,
             /*.kv_tier_semantic_threshold  =*/ params.kv_tier_semantic_threshold,
@@ -3354,6 +3355,7 @@ llama_context_params llama_context_default_params() {
         /*.kv_tier_compression         =*/ 1,
         /*.kv_tier_attention_threshold =*/ 0.1f,
         /*.kv_tier_warm_device         =*/ -1,
+        /*.kv_tier_warm_mlock          =*/ true,
         /*.kv_tier_total_ctx           =*/ 0,
         /*.kv_tier_semantic_index      =*/ nullptr,
         /*.kv_tier_semantic_threshold  =*/ 0.65f,

src/llama-memory.h

@@ -33,6 +33,7 @@ struct llama_memory_params {
     int32_t      kv_tier_compression;
     float        kv_tier_attention_threshold;
     int32_t      kv_tier_warm_device;
+    bool         kv_tier_warm_mlock;         // mlock the host-RAM warm tier (prevents kernel swap-out)
     int32_t      kv_tier_total_ctx;
     const char * kv_tier_semantic_index;
     float        kv_tier_semantic_threshold;

src/llama-model.cpp

@@ -2354,6 +2354,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params,
         tcfg.compression         = (mt::Compression)    params.kv_tier_compression;
         tcfg.attention_threshold = params.kv_tier_attention_threshold;
         tcfg.warm_device         = params.kv_tier_warm_device;
+        tcfg.warm_mlock          = params.kv_tier_warm_mlock;
         if (params.kv_tier_semantic_index && params.kv_tier_semantic_index[0]) {
             tcfg.semantic_index = params.kv_tier_semantic_index;
         }

src/memory-tier/mt-config.h

@@ -74,6 +74,16 @@ struct TieredConfig {
     // implement device-warm; the field is preserved for forward compat.
     int warm_device = -1;
 
+    // Pin the host-RAM warm tier in physical memory via mlock() / VirtualLock()
+    // so the kernel cannot page it out to swap under memory pressure. Default
+    // ON because the failure mode is silent and catastrophic: at long context
+    // on RAM-constrained hosts, warm-tier reads turn into swap-disk reads,
+    // collapsing prefill throughput to ~10% of expected. Disable only if the
+    // host cannot grant the lock (RLIMIT_MEMLOCK / CAP_IPC_LOCK) and the lock
+    // failure is acceptable for the workload. When mlock fails, allocation
+    // still succeeds and a warning is logged with the remediation steps.
+    bool warm_mlock = true;
+
     // Optional semantic-similarity prefetch. Non-empty enables it.
     std::string semantic_index;
     float       semantic_threshold = 0.65f;

src/memory-tier/mt-locked-buffer.cpp

@@ -0,0 +1,158 @@
+// mt-locked-buffer.cpp — see header for rationale.
+#include "mt-locked-buffer.h"
+
+#include "llama-impl.h"  // LLAMA_LOG_*
+
+#include <cerrno>
+#include <cstring>
+
+#ifdef _WIN32
+  #include <windows.h>
+#else
+  #include <sys/mman.h>
+  #include <unistd.h>
+#endif
+
+namespace mt {
+
+LockedBuffer::~LockedBuffer() {
+    reset();
+}
+
+LockedBuffer::LockedBuffer(LockedBuffer && other) noexcept
+    : data_(other.data_), size_(other.size_), locked_(other.locked_) {
+    other.data_   = nullptr;
+    other.size_   = 0;
+    other.locked_ = false;
+}
+
+LockedBuffer & LockedBuffer::operator=(LockedBuffer && other) noexcept {
+    if (this != &other) {
+        reset();
+        data_         = other.data_;
+        size_         = other.size_;
+        locked_       = other.locked_;
+        other.data_   = nullptr;
+        other.size_   = 0;
+        other.locked_ = false;
+    }
+    return *this;
+}
+
+void LockedBuffer::reset() {
+    if (!data_) {
+        size_   = 0;
+        locked_ = false;
+        return;
+    }
+#ifdef _WIN32
+    if (locked_) {
+        VirtualUnlock(data_, size_);
+    }
+    VirtualFree(data_, 0, MEM_RELEASE);
+#else
+    if (locked_) {
+        munlock(data_, size_);
+    }
+    munmap(data_, size_);
+#endif
+    data_   = nullptr;
+    size_   = 0;
+    locked_ = false;
+}
+
+bool LockedBuffer::allocate(size_t size_bytes, bool lock) {
+    reset();
+    if (size_bytes == 0) {
+        return true;
+    }
+
+#ifdef _WIN32
+    data_ = static_cast<uint8_t *>(
+        VirtualAlloc(nullptr, size_bytes, MEM_COMMIT | MEM_RESERVE, PAGE_READWRITE));
+    if (!data_) {
+        LLAMA_LOG_ERROR(
+            "mt::LockedBuffer: VirtualAlloc(%zu bytes / %.1f MiB) failed (error %lu)\n",
+            size_bytes, (double) size_bytes / (1024.0 * 1024.0),
+            (unsigned long) GetLastError());
+        return false;
+    }
+    size_ = size_bytes;
+    if (lock) {
+        if (VirtualLock(data_, size_bytes)) {
+            locked_ = true;
+        } else {
+            LLAMA_LOG_WARN(
+                "mt::LockedBuffer: VirtualLock(%.1f MiB) failed (error %lu). "
+                "Buffer is allocated but eligible for page-out under RAM pressure. "
+                "Raise the process working-set quota to fix.\n",
+                (double) size_bytes / (1024.0 * 1024.0),
+                (unsigned long) GetLastError());
+        }
+    }
+#else
+    int flags = MAP_PRIVATE | MAP_ANONYMOUS;
+    bool tried_map_locked = false;
+
+  #ifdef MAP_LOCKED
+    if (lock) {
+        flags |= MAP_LOCKED;
+        tried_map_locked = true;
+    }
+  #endif
+
+    void * p = mmap(nullptr, size_bytes, PROT_READ | PROT_WRITE, flags, -1, 0);
+
+  #ifdef MAP_LOCKED
+    if (p == MAP_FAILED && tried_map_locked) {
+        // MAP_LOCKED often fails on RLIMIT_MEMLOCK; retry without it and try
+        // mlock() after, which gives a clearer error path.
+        int saved_errno = errno;
+        flags &= ~MAP_LOCKED;
+        p = mmap(nullptr, size_bytes, PROT_READ | PROT_WRITE, flags, -1, 0);
+        if (p == MAP_FAILED) {
+            LLAMA_LOG_ERROR(
+                "mt::LockedBuffer: mmap(%.1f MiB) failed: %s "
+                "(also failed with MAP_LOCKED: %s)\n",
+                (double) size_bytes / (1024.0 * 1024.0),
+                strerror(errno), strerror(saved_errno));
+            return false;
+        }
+        // fall through — we have an unlocked mapping; mlock() pass below.
+        tried_map_locked = false;
+    } else if (p != MAP_FAILED && tried_map_locked) {
+        // MAP_LOCKED succeeded — pages are pinned.
+        locked_ = true;
+    }
+  #endif
+
+    if (p == MAP_FAILED) {
+        LLAMA_LOG_ERROR(
+            "mt::LockedBuffer: mmap(%.1f MiB) failed: %s\n",
+            (double) size_bytes / (1024.0 * 1024.0), strerror(errno));
+        return false;
+    }
+
+    data_ = static_cast<uint8_t *>(p);
+    size_ = size_bytes;
+
+    // If we asked for a lock but MAP_LOCKED wasn't used or didn't apply,
+    // do a follow-up mlock() so we get a clear errno path.
+    if (lock && !locked_) {
+        if (mlock(data_, size_bytes) == 0) {
+            locked_ = true;
+        } else {
+            LLAMA_LOG_WARN(
+                "mt::LockedBuffer: mlock(%.1f MiB) failed: %s. "
+                "Buffer is allocated but eligible for swap-out under RAM pressure. "
+                "Raise RLIMIT_MEMLOCK (systemd unit: LimitMEMLOCK=infinity) "
+                "or run with CAP_IPC_LOCK to fix.\n",
+                (double) size_bytes / (1024.0 * 1024.0), strerror(errno));
+        }
+    }
+#endif
+
+    return true;
+}
+
+} // namespace mt

src/memory-tier/mt-locked-buffer.h

@@ -0,0 +1,66 @@
+// mt-locked-buffer.h
+//
+// LockedBuffer: a page-aligned byte buffer that can be optionally pinned in
+// RAM via mlock() / VirtualLock(). Used by the tiered-KV warm staging tier so
+// the kernel cannot page warm-tier KV out to swap under memory pressure (the
+// failure mode that turned warm-tier reads into swap-disk reads, observed on
+// long-context workloads with constrained host RAM).
+//
+// Owns its allocation via mmap (POSIX) or VirtualAlloc (Windows) — does NOT
+// use the C++ heap, so it bypasses any allocator-level placement assumptions
+// the kernel could apply to vector/new memory.
+//
+// Locking is BEST-EFFORT: if the kernel refuses mlock (e.g. RLIMIT_MEMLOCK
+// too small or no CAP_IPC_LOCK), allocation still succeeds and the buffer
+// is unlocked. A warning is logged with a concrete remediation hint.
+//
+// Drop-in replacement for `std::vector<uint8_t>` in the narrow API surface
+// the warm-tier code uses today: data(), size(), and a single bulk allocate.
+//
+// Cross-platform:
+//   Linux : mmap(MAP_PRIVATE|MAP_ANONYMOUS [|MAP_LOCKED]) + fallback mlock()
+//   macOS : mmap(MAP_PRIVATE|MAP_ANONYMOUS) + mlock()  (no MAP_LOCKED)
+//   Win32 : VirtualAlloc(MEM_COMMIT|MEM_RESERVE) + VirtualLock()
+#pragma once
+
+#include <cstddef>
+#include <cstdint>
+
+namespace mt {
+
+class LockedBuffer {
+public:
+    LockedBuffer() = default;
+    ~LockedBuffer();
+
+    // Allocate `size_bytes` of zero-initialized memory.
+    //
+    // If `lock` is true, attempts to pin the memory in RAM. On failure the
+    // buffer remains allocated (operation succeeds) but `is_locked()` will
+    // return false and a warning is logged.
+    //
+    // Returns false only if the underlying allocation itself fails.
+    bool allocate(size_t size_bytes, bool lock = true);
+
+    // Release the buffer (munlock + munmap / VirtualUnlock + VirtualFree).
+    // Safe to call on an empty buffer. Called automatically by the destructor.
+    void reset();
+
+    uint8_t *       data()         { return data_; }
+    const uint8_t * data()   const { return data_; }
+    size_t          size()   const { return size_; }
+    bool            is_locked() const { return locked_; }
+
+    // Disable copy. Allow move so we can reassign in the tiered cache class.
+    LockedBuffer(const LockedBuffer &)             = delete;
+    LockedBuffer & operator=(const LockedBuffer &) = delete;
+    LockedBuffer(LockedBuffer && other) noexcept;
+    LockedBuffer & operator=(LockedBuffer && other) noexcept;
+
+private:
+    uint8_t * data_   = nullptr;
+    size_t    size_   = 0;
+    bool      locked_ = false;
+};
+
+} // namespace mt