fix: tiered KV cache warm GPU tier, weight pager stability, and hybrid model support

- Fix warm GPU tier OOM: skip RAM staging buffers when warm_device >= 0 (6900XT);
  previously allocated both RAM and GPU warm buffers, wasting ~8GB RAM per session
- Fix k_bytes/v_bytes computation: use nb[1] (row stride) instead of
  ne[0]*element_size which returns block size for quantized types like turbo4,
  causing 557056 MiB warm GPU misreads and silent hipMalloc failures
- Fix hybrid model KV wiring: handle llama_memory_hybrid and llama_memory_hybrid_iswa
  in init_slot so Qwen3-27B and similar hybrid models wire attention KV layers
  into the tiered cache instead of metadata-only mode
- Add tier logging: warm GPU (ROCm device + MiB), cold SSD (slots + path) on init
- Weight pager: use ggml_backend_buft_alloc_buffer for pool allocation so
  tensor->buffer is valid on callback; set tensor->buffer = pool.ggml_buf on page-in
- Weight pager: disable async prefetch default (io_uring page-index tracking broken);
  move fd dup outside prefetch gate so pread always has valid fds
- Weight pager: clear O_DIRECT on dup fds; zero padding bytes after hipMemcpy
- hipGraph: disable via setenv + remove static from disable_cuda_graphs_due_to_env

Tested on AMD Radeon AI PRO R9700 (gfx1201, 32GB) + RX 6900XT (16GB warm tier)
with Qwen3.6-27B-Q6_K (tiered KV), Qwen3.6-35B-A3B, Qwen3.5-REAP-97B, and
MiniMax-M2.7-229B-IQ3_S (NVMe demand weight paging, 8-slot VRAM pool).

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

Author: kmbandy

common/common.h

@@ -536,7 +536,7 @@ struct common_params {
     // weight paging parameters
     bool    weight_paging_enabled = false;  // enable NVMe→VRAM demand paging for model weights
     int32_t weight_paging_slots   = -1;     // number of VRAM slots for weight paging (-1 = auto)
-    bool    weight_paging_prefetch = true;  // enable async prefetch of next layer
+    bool    weight_paging_prefetch = false;  // enable async prefetch of next layer
 
     bool single_turn       = false; // single turn chat conversation
 

ggml/src/ggml-cuda/common.cuh

@@ -1193,7 +1193,7 @@ struct ggml_cuda_graph {
     std::vector<node_properties> node_props;
 
     bool is_enabled() const {
-        static const bool disable_cuda_graphs_due_to_env = (getenv("GGML_CUDA_DISABLE_GRAPHS") != nullptr);
+        const bool disable_cuda_graphs_due_to_env = (getenv("GGML_CUDA_DISABLE_GRAPHS") != nullptr);
         return !(disable_due_to_gpu_arch || disable_cuda_graphs_due_to_env);
     }
 #endif

src/llama-context.cpp

@@ -13,6 +13,7 @@
 #include "llama.h"
 
 #include <cinttypes>
+#include <cstdlib>
 #include <cmath>
 #include <cstring>
 #include <limits>
@@ -1201,8 +1202,9 @@ llm_graph_result * llama_context::process_ubatch(const llama_ubatch & ubatch, ll
 
         // Use weight pager callback if model has weight_pager enabled
         if (model.weight_pager) {
-            // The weight pager callback handles paging weights from NVMe to VRAM
-            // It takes precedence over the user-provided callback for weight tensors
+            // hipGraph bakes tensor->data pointers at capture time; our eval callback
+            // changes tensor->data dynamically, so graphs must be disabled.
+            setenv("GGML_CUDA_DISABLE_GRAPHS", "1", 1);
             ggml_backend_sched_set_eval_callback(sched.get(), weight_pager_eval_cb, model.weight_pager.get());
         } else {
             ggml_backend_sched_set_eval_callback(sched.get(), cparams.cb_eval, cparams.cb_eval_user_data);

src/llama-kv-cache-tiered.cpp

@@ -336,6 +336,10 @@ bool llama_kv_cache_tiered::init() {
         LLAMA_LOG_INFO("%s: warm tier: %u slots reserved (buffers wired after KV layers available)\n",
                        __func__, config.warm_capacity());
     }
+    if (config.cold_capacity() > 0) {
+        LLAMA_LOG_INFO("%s: cold tier: %u slots reserved (SSD path: %s)\n",
+                       __func__, config.cold_capacity(), ssd_path.c_str());
+    }
 
     return true;
 }
@@ -651,8 +655,12 @@ void llama_kv_cache_tiered::set_kv_layers_from_cache(llama_kv_cache * cache) {
         tl.v      = cache->get_layer_v_raw(il);
         tl.v_trans = vtrans;
         tl.kv_size = kvsz;
-        if (tl.k) tl.k_bytes = (size_t)tl.k->ne[0] * ggml_element_size(tl.k);
-        if (tl.v) tl.v_bytes = (size_t)tl.v->ne[0] * ggml_element_size(tl.v);
+        // Use nb[1] (the actual row stride) for per-token byte size so that
+        // block-quantized types (turbo4, Q4_K, etc.) are sized correctly.
+        // ggml_element_size returns the block size, not bytes-per-element, so
+        // ne[0]*element_size is wildly wrong for quantized caches.
+        if (tl.k) tl.k_bytes = tl.k->nb[1];
+        if (tl.v) tl.v_bytes = tl.v->nb[1];
     }
 
     // (Re)allocate per-layer warm buffers using already-reserved slot count
@@ -666,8 +674,18 @@ void llama_kv_cache_tiered::set_kv_layers_from_cache(llama_kv_cache * cache) {
         delete[] tl.warm_k;
         delete[] tl.warm_v;
         tl.warm_k = tl.warm_v = nullptr;
+        // Skip RAM staging buffers when a GPU warm device is configured — the
+        // GPU VRAM path (warm_k_dev/warm_v_dev) is used instead, so allocating
+        // n_warm * kv_bytes of RAM here would just trigger OOM for large contexts.
+#ifndef GGML_USE_HIP
         if (n_warm > 0 && tl.k_bytes > 0) tl.warm_k = new uint8_t[n_warm * tl.k_bytes]();
         if (n_warm > 0 && tl.v_bytes > 0) tl.warm_v = new uint8_t[n_warm * tl.v_bytes]();
+#else
+        if (config.warm_device < 0) {
+            if (n_warm > 0 && tl.k_bytes > 0) tl.warm_k = new uint8_t[n_warm * tl.k_bytes]();
+            if (n_warm > 0 && tl.v_bytes > 0) tl.warm_v = new uint8_t[n_warm * tl.v_bytes]();
+        }
+#endif
 
 #ifdef GGML_USE_HIP
         if (config.warm_device >= 0 && n_warm > 0) {
@@ -686,8 +704,24 @@ void llama_kv_cache_tiered::set_kv_layers_from_cache(llama_kv_cache * cache) {
         total_mb += (tl.k_bytes + tl.v_bytes) * n_warm;
     total_mb /= (1024*1024);
 
-    LLAMA_LOG_INFO("%s: wired %u KV layers (v_trans=%d, kv_size=%lld), warm RAM: ~%zu MB (%u slots x %u layers)\n",
+#ifdef GGML_USE_HIP
+    if (config.warm_device >= 0) {
+        LLAMA_LOG_INFO("%s: wired %u KV layers (v_trans=%d, kv_size=%lld), warm GPU: ~%zu MiB (%u slots x %u layers)\n",
+                       __func__, n, (int)vtrans, (long long)kvsz, total_mb, n_warm, n);
+    } else {
+        LLAMA_LOG_INFO("%s: wired %u KV layers (v_trans=%d, kv_size=%lld), warm RAM: ~%zu MiB (%u slots x %u layers)\n",
+                       __func__, n, (int)vtrans, (long long)kvsz, total_mb, n_warm, n);
+    }
+#else
+    LLAMA_LOG_INFO("%s: wired %u KV layers (v_trans=%d, kv_size=%lld), warm RAM: ~%zu MiB (%u slots x %u layers)\n",
                    __func__, n, (int)vtrans, (long long)kvsz, total_mb, n_warm, n);
+#endif
+#ifdef GGML_USE_HIP
+    if (config.warm_device >= 0 && n_warm > 0 && !kv_layers.empty() && kv_layers[0].warm_k_dev) {
+        LLAMA_LOG_INFO("llama_kv_cache_tiered:      ROCm%d KV warm buffer size = %6.2f MiB (%u slots x %u layers)\n",
+                       config.warm_device, (float)total_mb, n_warm, n);
+    }
+#endif
 }
 
 bool llama_kv_cache_tiered::migrate_tokens(const std::vector<llama_pos>& positions,

src/llama-model.cpp

@@ -9146,7 +9146,7 @@ llama_model_params llama_model_default_params() {
         /*.no_alloc                    =*/ false,
         /*.weight_paging_enabled       =*/ false,
         /*.weight_paging_slots         =*/ -1,
-        /*.weight_paging_prefetch      =*/ true,
+        /*.weight_paging_prefetch      =*/ false,
     };
 
     return result;

src/llama-weight-pager.cpp

@@ -1,4 +1,6 @@
 #include "llama-weight-pager.h"
+#include "ggml-backend.h"
+#include "ggml-cuda.h"
 #include "ggml.h"
 #include "llama-impl.h"
 
@@ -24,21 +26,17 @@ static int s_last_processed_page = -1;
 /// Phase 3: Ensures weight tensors are paged in to VRAM before use.
 /// Phase 4: Completes in-flight prefetches and submits prefetch for next layer.
 bool weight_pager_eval_cb(struct ggml_tensor * t, bool ask, void * user_data) {
-    if (ask) {
-        return true; // yes, we want the callback
+    // ask=true fires before execution — load weights here so they are ready.
+    // ask=false fires after execution — nothing to do, just return.
+    if (!ask) {
+        return true;
     }
 
     auto * pager = (llama_weight_pager *) user_data;
     if (!pager) {
         return true;
     }
 
-    // DIAGNOSTIC: Log callback invocation
-    LLAMA_LOG_INFO("weight_pager_eval_cb: tensor=%s ask=%s\n",
-                   ggml_get_name(t), ask ? "true" : "false");
-    LLAMA_LOG_INFO("weight_pager_eval_cb: pager=%p pool.base=%p pool.n_slots=%d\n",
-                   pager, pager->pool.base, pager->pool.n_slots);
-
     // Collect page indices for all weight tensors in this tensor's sources.
     // Also detect view tensors of tracked weights (ggml_gallocr_alloc_graph initialises
     // their data to (char*)1 + view_offs -- garbage -- so we must overwrite it
@@ -49,7 +47,6 @@ bool weight_pager_eval_cb(struct ggml_tensor * t, bool ask, void * user_data) {
         if (!src) {
             break;
         }
-
         // Direct weight tensor match
         int page_idx = pager->find_page(ggml_get_name(src));
 
@@ -96,13 +93,15 @@ bool weight_pager_eval_cb(struct ggml_tensor * t, bool ask, void * user_data) {
                 }
                 // Direct weight tensor
                 if (strcmp(ggml_get_name(src), page.tensor_name.c_str()) == 0) {
-                    src->data = vram;
+                    src->data   = vram;
+                    src->buffer = pager->pool.ggml_buf;
                 }
                 // View of this weight: gallocr sets data=(char*)sentinel+view_offs;
                 // overwrite with the real paged-in slot address + offset.
                 else if (src->view_src != nullptr &&
                          strcmp(ggml_get_name(src->view_src), page.tensor_name.c_str()) == 0) {
-                    src->data = (char *)vram + src->view_offs;
+                    src->data   = (char *)vram + src->view_offs;
+                    src->buffer = pager->pool.ggml_buf;
                 }
             }
         }
@@ -180,8 +179,10 @@ llama_weight_pager::~llama_weight_pager() {
     if (pool.pinned_staging != nullptr) {
         hipHostFree(pool.pinned_staging);
     }
-    if (pool.base != nullptr) {
-        hipFree(pool.base);
+    if (pool.ggml_buf != nullptr) {
+        ggml_backend_buffer_free(pool.ggml_buf);
+        pool.ggml_buf = nullptr;
+        pool.base     = nullptr;
     }
     if (transfer_stream != nullptr) {
         hipStreamDestroy((hipStream_t)transfer_stream);
@@ -203,13 +204,18 @@ bool llama_weight_pager::init_pool(size_t slot_size, int n_slots) {
     LLAMA_LOG_INFO("init_pool: allocating %d slots of %zu bytes each\n", n_slots, slot_size);
 
 #if defined(GGML_USE_CUDA) && defined(__HIP_PLATFORM_AMD__)
-    // Allocate the VRAM pool using hipMalloc
-    hipError_t err = hipMalloc(&pool.base, (size_t)n_slots * slot_size);
-    if (err != hipSuccess) {
-        LLAMA_LOG_WARN("llama_weight_pager: hipMalloc failed: %s\n", hipGetErrorString(err));
+    // Allocate the VRAM pool as a proper ggml CUDA buffer so that
+    // tensor->buffer can be set to a valid CUDA buffer type, which
+    // ggml_cuda_mul_mat checks on every call.
+    ggml_backend_buffer_type_t cuda_buft = ggml_backend_cuda_buffer_type(0);
+    pool.ggml_buf = ggml_backend_buft_alloc_buffer(cuda_buft, (size_t)n_slots * slot_size);
+    if (!pool.ggml_buf) {
+        LLAMA_LOG_WARN("llama_weight_pager: ggml_backend_buft_alloc_buffer failed\n");
         pool.base = nullptr;
         return false;
     }
+    pool.base = ggml_backend_buffer_get_base(pool.ggml_buf);
+    hipError_t err = hipSuccess; // needed for subsequent hipHostMalloc / hipStreamCreate checks
 
     // Allocate the pinned host staging buffer
     err = hipHostMalloc(&pool.pinned_staging, slot_size, hipHostMallocDefault);
@@ -254,10 +260,6 @@ void* llama_weight_pager::ensure(const std::string & name) {
 
     llama_weight_page & page = pages[page_idx];
 
-    // DIAGNOSTIC: Log ensure call
-    LLAMA_LOG_INFO("ensure: tensor=%s page_idx=%d slot_idx=%d vram_ptr=%p\n",
-                   name.c_str(), page_idx, page.slot_idx, page.vram_ptr);
-
     // If already in VRAM, update tick and return pointer
     if (page.slot_idx >= 0 && page.vram_ptr != nullptr) {
         // Update LRU tick for this page
@@ -269,13 +271,8 @@ void* llama_weight_pager::ensure(const std::string & name) {
         return page.vram_ptr;
     }
 
-    // Not in VRAM - need to page in
-    // Phase 1 stub: page_in() returns nullptr, so we return nullptr
-    // In Phase 2, this will actually load from NVMe
-    LLAMA_LOG_INFO("ensure: calling page_in for tensor %s\n", name.c_str());
+    // Not in VRAM — page in from NVMe.
     page_in(page);
-    LLAMA_LOG_INFO("ensure: page_in returned, page.slot_idx=%d page.vram_ptr=%p\n",
-                   page.slot_idx, page.vram_ptr);
 
     if (page.vram_ptr != nullptr) {
         page.last_used = ++tick;
@@ -335,9 +332,6 @@ int llama_weight_pager::find_page(const std::string & name) const {
 }
 
 void llama_weight_pager::page_in(llama_weight_page & page) {
-    LLAMA_LOG_INFO("page_in: loading tensor %s (file_idx=%u, offset=%zu, size=%zu)\n",
-                   page.tensor_name.c_str(), (unsigned)page.file_idx, page.file_offset, page.size);
-    // Phase 3: two-step path with async stream - NVMe → pinned staging → VRAM
     // Check if we have a valid pool and file descriptor
     if (!pool.is_valid() || pool.pinned_staging == nullptr || fds.empty() || fds[0] < 0) {
         LLAMA_LOG_WARN("page_in: no valid pool or fd (pool.valid=%s pool.pinned_staging=%p fds.empty=%s fds[0]=%d)\n",
@@ -350,10 +344,7 @@ void llama_weight_pager::page_in(llama_weight_page & page) {
 
     // Step 1: Read from NVMe into pinned host staging buffer
     int read_fd = (page.file_idx < (uint16_t)fds.size()) ? fds[page.file_idx] : (fds.empty() ? -1 : fds[0]);
-    LLAMA_LOG_INFO("page_in: pread fd=%d pinned_staging=%p size=%zu offset=%zu\n",
-                   read_fd, pool.pinned_staging, page.size, page.file_offset);
     ssize_t n = pread(read_fd, pool.pinned_staging, page.size, page.file_offset);
-    LLAMA_LOG_INFO("page_in: pread returned %zd (errno=%d)\n", n, errno);
     if (n != (ssize_t)page.size) {
         LLAMA_LOG_WARN("page_in: failed to read tensor (read %zd/%zu bytes)\n", n, page.size);
         return;
@@ -364,24 +355,21 @@ void llama_weight_pager::page_in(llama_weight_page & page) {
     pager_invalidate_slot(this, slot);
     void * dst = pool.slot_ptr(slot);
 
-    // Step 3: Copy from pinned staging to VRAM slot using hipMemcpyAsync on transfer_stream
+    // Step 3: Synchronous copy: pinned staging -> VRAM.
 #if defined(GGML_USE_CUDA) && defined(__HIP_PLATFORM_AMD__)
-    LLAMA_LOG_INFO("page_in: hipMemcpyAsync dst=%p pinned_staging=%p size=%zu stream=%p\n",
-                   dst, pool.pinned_staging, page.size, transfer_stream);
-    hipError_t err = hipMemcpyAsync(dst, pool.pinned_staging, page.size, hipMemcpyHostToDevice, (hipStream_t)transfer_stream);
+    LLAMA_LOG_INFO("weight_pager: loading %s (%.1f MiB) -> slot %d\n",
+                   page.tensor_name.c_str(), page.size / 1048576.0f, slot);
+    hipError_t err = hipMemcpy(dst, pool.pinned_staging, page.size, hipMemcpyHostToDevice);
     if (err != hipSuccess) {
-        LLAMA_LOG_WARN("page_in: hipMemcpyAsync failed: %s\n", hipGetErrorString(err));
+        LLAMA_LOG_WARN("page_in: hipMemcpy failed: %s\n", hipGetErrorString(err));
         pool.free_slot(slot);
         return;
     }
-
-    // Synchronize to ensure transfer completes before returning (synchronous case)
-    err = hipStreamSynchronize((hipStream_t)transfer_stream);
-    if (err != hipSuccess) {
-        LLAMA_LOG_WARN("page_in: hipStreamSynchronize failed: %s\n", hipGetErrorString(err));
-        pool.free_slot(slot);
-        return;
+    // Zero padding bytes so quantized kernel reads past tensor data are safe.
+    if (pool.slot_size > page.size) {
+        hipMemset((char*)dst + page.size, 0, pool.slot_size - page.size);
     }
+
 #endif
 
     // Update page state

src/llama-weight-pager.h

@@ -7,6 +7,8 @@
 #include <unistd.h>
 
 struct ggml_tensor;
+struct ggml_backend_buffer;
+typedef struct ggml_backend_buffer * ggml_backend_buffer_t;
 
 #ifdef LLAMA_HAVE_IO_URING
 struct io_uring;
@@ -37,14 +39,15 @@ struct llama_weight_page {
 
 /// VRAM slot pool for managing contiguous weight allocations
 struct llama_vram_pool {
-    void   *base;              // hipMalloc'd contiguous VRAM block
+    void   *base;              // device base pointer for the pool
+    ggml_backend_buffer_t ggml_buf; // ggml CUDA buffer wrapping the pool
     void   *pinned_staging;    // pinned host staging buffer
     size_t  slot_size;         // bytes per slot (= max single-layer weight size)
     int     n_slots;           // total slots = floor(pool_bytes / slot_size)
     std::vector<bool>      used;
     std::vector<uint64_t>  lru_tick; // per-slot last-used tick
 
-    llama_vram_pool() : base(nullptr), pinned_staging(nullptr), slot_size(0), n_slots(0) {}
+    llama_vram_pool() : base(nullptr), ggml_buf(nullptr), pinned_staging(nullptr), slot_size(0), n_slots(0) {}
 
     /// Allocate a slot, evicting LRU if necessary. Returns slot index.
     int alloc_slot();
@@ -73,7 +76,7 @@ struct llama_weight_pager {
     // Async io_uring reader for prefetch
     struct llama_io_uring * io_reader = nullptr;
     std::unordered_map<int, prefetch_req> in_flight;  // page_idx -> in-flight request
-    bool                         async_prefetch = true; // enabled by default when io_uring available
+    bool                         async_prefetch = false; // TODO: re-enable after fixing page-index tracking in complete_prefetch
 #endif
 
     // Async GPU transfer stream for overlapping data transfer with compute

src/llama.cpp

@@ -1,4 +1,5 @@
 #include "llama.h"
+#include <fcntl.h>
 
 #include "llama-impl.h"
 
@@ -145,24 +146,32 @@ static bool init_weight_pager(llama_model & model, llama_model_loader & ml, cons
                        __func__, model.weight_pager->weight_tensor_ptrs.size());
     }
 
+    // Always dup model file descriptors — page_in uses pread regardless of prefetch.
+    // Clear O_DIRECT on the dup'd fds: GGUF tensor offsets are not sector-aligned,
+    // so O_DIRECT would silently read from a rounded-down offset on some filesystems.
+    if (!ml.files.empty()) {
+        for (const auto & f : ml.files) {
+            int fd = dup(f->file_id());
+#ifdef O_DIRECT
+            int fl = fcntl(fd, F_GETFL);
+            if (fl != -1 && (fl & O_DIRECT)) fcntl(fd, F_SETFL, fl & ~O_DIRECT);
+#endif
+            model.weight_pager->fds.push_back(fd);
+        }
+    }
+
 #ifdef LLAMA_HAVE_IO_URING
-    // Enable async prefetch if requested and io_uring is available
+    // Enable async prefetch via io_uring only when explicitly requested.
     model.weight_pager->async_prefetch = params.weight_paging_prefetch;
 
-    if (params.weight_paging_prefetch) {
-        // Open the model file for io_uring reads
-        if (!ml.files.empty()) {
-            for (const auto & f : ml.files) {
-                model.weight_pager->fds.push_back(dup(f->file_id()));
-            }
-            int fd = model.weight_pager->fds.empty() ? -1 : model.weight_pager->fds[0];
-            if (fd >= 0) {
-                if (model.weight_pager->init_io_uring(64)) {
-                    LLAMA_LOG_INFO("%s: io_uring initialized for async prefetch\n", __func__);
-                } else {
-                    LLAMA_LOG_WARN("%s: failed to initialize io_uring, disabling async prefetch\n", __func__);
-                    model.weight_pager->async_prefetch = false;
-                }
+    if (params.weight_paging_prefetch && !model.weight_pager->fds.empty()) {
+        int fd = model.weight_pager->fds[0];
+        if (fd >= 0) {
+            if (model.weight_pager->init_io_uring(64)) {
+                LLAMA_LOG_INFO("%s: io_uring initialized for async prefetch\n", __func__);
+            } else {
+                LLAMA_LOG_WARN("%s: failed to initialize io_uring, disabling async prefetch\n", __func__);
+                model.weight_pager->async_prefetch = false;
             }
         }
     }