poly, HIP: the expert cache gives VRAM back and grows again on AMD (#13)

include/strata/hip_compat/cuda.h

@@ -4,6 +4,13 @@
 // (cuMemCreate, cuMemMap, ...) onto HIP's memory management API, which the HIP runtime provides.
 #include <hip/hip_runtime.h>
 
+// ROCm gives the VRAM of a released handle back only when the addresses it was mapped at are freed, and frees only
+// a whole reservation (#13): Engine::vmm_unmap() reserves the addresses again for the pieces it keeps.  Measured on
+// Linux; Windows' HIP runtime is another implementation and keeps the plain release until it is measured there.
+#if !defined(_WIN32)
+#define STRATA_VMM_RELEASE_WITH_ADDRESSES 1
+#endif
+
 typedef hipError_t CUresult;
 typedef int CUdevice;
 // HIP's own CUdeviceptr/handle equivalents are pointer-typed; CUDA's are integers, and strata-poly's Mapped

include/strata/hip_compat/cuda_runtime.h

@@ -92,7 +92,49 @@ hipError_t mem_get_info(size_t* free_bytes, size_t* total_bytes);
 }
 #define cudaMemGetInfo ::strata::hip_compat::mem_get_info
 #else
-#define cudaMemGetInfo hipMemGetInfo
+// Linux: hipMemGetInfo's free figure is the kernel driver's KFD accounting, which on an RX 9070 XT (16 GB, Fedora,
+// ROCm 7 wheels) read 924 MiB while amdgpu's own counters had 2,128 MiB free, and a 1,500 MiB hipMalloc beside the
+// engine succeeded.  The expert cache, sized and grown by that figure, left more than a GiB of the card unused and
+// did not grow back after another program freed VRAM.  The free figure is amdgpu's (mem_info_vram_total less
+// mem_info_vram_used, for the card's PCI address) when the files are there (STRATA_KERNEL_VRAM_FREE=0: as it is).
+// Not on an integrated GPU: its files count only the BIOS carve-out (a Ryzen 7 9800X3D's Radeon: 512 MiB), while HIP
+// gives it system RAM (15,536 MiB free by hipMemGetInfo).
+#include <cctype>
+#include <cstdio>
+#include <cstdlib>
+#include <cstring>
+namespace strata::hip_compat {
+inline bool read_sysfs_u64(const char* path, unsigned long long* v) {
+    FILE* f = std::fopen(path, "r");
+    if (!f) return false;
+    const bool ok = std::fscanf(f, "%llu", v) == 1;
+    std::fclose(f);
+    return ok;
+}
+inline hipError_t mem_get_info(size_t* free_bytes, size_t* total_bytes) {
+    const hipError_t e = hipMemGetInfo(free_bytes, total_bytes);
+    static const bool off = [] {
+        const char* v = std::getenv("STRATA_KERNEL_VRAM_FREE");
+        return v && std::strcmp(v, "0") == 0;
+    }();
+    int dev = 0, integrated = 1;
+    char bus[64] = {};
+    if (e != hipSuccess || off || hipGetDevice(&dev) != hipSuccess ||
+        hipDeviceGetAttribute(&integrated, hipDeviceAttributeIntegrated, dev) != hipSuccess || integrated ||
+        hipDeviceGetPCIBusId(bus, (int) sizeof(bus), dev) != hipSuccess)
+        return e;
+    for (char* c = bus; *c; ++c) *c = (char) std::tolower((unsigned char) *c);
+    char path[160];
+    unsigned long long total = 0, used = 0;
+    std::snprintf(path, sizeof(path), "/sys/bus/pci/devices/%s/mem_info_vram_total", bus);
+    if (!read_sysfs_u64(path, &total)) return e;
+    std::snprintf(path, sizeof(path), "/sys/bus/pci/devices/%s/mem_info_vram_used", bus);
+    if (!read_sysfs_u64(path, &used) || used > total || total == 0) return e;
+    *free_bytes = (size_t) (total - used);
+    return e;
+}
+}  // namespace strata::hip_compat
+#define cudaMemGetInfo ::strata::hip_compat::mem_get_info
 #endif
 #define cudaMemcpy hipMemcpy
 #define cudaMemcpy2D hipMemcpy2D

src/poly/engine.cpp

@@ -23,6 +23,7 @@
 #include <chrono>
 #include <cmath>
 #include <cstdio>
+#include <cstdlib>
 #include <cstring>
 #include <exception>
 #include <filesystem>
@@ -496,13 +497,64 @@ void Engine::vmm_map(Mapped& m, size_t bytes, size_t piece, const char* what) {
 
 // Gives back the pieces of m that lie past its first `bytes`.  The card must be done with them.
 void Engine::vmm_unmap(Mapped& m, size_t bytes) {
+    const size_t had = m.mapped;
     while (!m.pieces.empty() && m.mapped - m.pieces.back().second >= bytes) {
         const auto [h, n] = m.pieces.back();
         m.mapped -= n;
         m.pieces.pop_back();
         cu_ck(cuMemUnmap(m.base + m.mapped, n), "unmapping VRAM");
         cu_ck(cuMemRelease(h), "unmapping VRAM");
     }
+#ifdef STRATA_VMM_RELEASE_WITH_ADDRESSES
+    // On ROCm the released pieces stay the process's while their addresses are reserved (an RX 9070 XT: 320 MiB
+    // unmapped and released, 32 MiB more free; the addresses freed, all of it).  The kept pieces leave their
+    // addresses, which are freed and reserved again where they were, and come back with their bytes.  Once the first
+    // has left, the slots, the recorded launches and m point at VRAM that is not there: a step that fails from
+    // there on ends the engine (the server starts it again for the next request) rather than reading that VRAM.
+    if (m.mapped == had || !m.base) return;
+    const auto must = [](CUresult r, const char* what) {
+        if (r == CUDA_SUCCESS) return;
+        const char* s = nullptr;
+        (void) cuGetErrorString(r, &s);
+        std::fprintf(stderr, "strata-poly: %s while giving VRAM back: %s; the caches are no longer mapped, so the "
+                             "engine stops\n", what, s ? s : "unknown error");
+        std::abort();
+    };
+    for (size_t at = m.mapped, i = m.pieces.size(); i > 0; --i) {
+        at -= m.pieces[i - 1].second;
+        must(cuMemUnmap(m.base + at, m.pieces[i - 1].second), "unmapping the kept VRAM");
+    }
+    size_t free_b = 0, total_b = 0;
+    (void) cudaMemGetInfo(&free_b, &total_b);
+    const CUdeviceptr was = m.base;
+    must(cuMemAddressFree(m.base, m.reserved), "freeing the addresses");
+    CUdeviceptr p = 0;
+    must(cuMemAddressReserve(&p, m.reserved, 0, was, 0), "reserving the addresses again");
+    if (p != was) {   // the slots and the recorded launches point into the old addresses
+        std::fprintf(stderr, "strata-poly: the VRAM addresses could not be reserved again where they were, so the "
+                             "engine stops\n");
+        std::abort();
+    }
+    CUmemAccessDesc access = {};
+    access.location.type = CU_MEM_LOCATION_TYPE_DEVICE;
+    access.location.id = device_;
+    access.flags = CU_MEM_ACCESS_FLAGS_PROT_READWRITE;
+    size_t at = 0;
+    for (const auto& [h, n] : m.pieces) {
+        must(cuMemMap(m.base + at, n, 0, h, 0), "mapping VRAM again");
+        must(cuMemSetAccess(m.base + at, n, &access, 1), "mapping VRAM again");
+        at += n;
+    }
+    // amdgpu counts the VRAM as free 2 to 5 ms later: a reading before that (mind_vram()'s log line, cover()) is
+    // short by what was given back.  At most 100 ms, as another program may take VRAM meanwhile.
+    const size_t back = (had - m.mapped) / 8 * 7, want = free_b + back;
+    for (int i = 0; i < 100 && free_b < want; ++i) {
+        std::this_thread::sleep_for(std::chrono::milliseconds(1));
+        (void) cudaMemGetInfo(&free_b, &total_b);
+    }
+#else
+    (void) had;
+#endif
 }
 
 // For the destructor: what is left of m, and its addresses, without throwing.
@@ -520,7 +572,7 @@ void* Engine::kv_alloc(size_t bytes, int per) {
     if (again_ != SIZE_MAX) return bufs_.at(again_++).p;
     const size_t all = (size_t) ((opt_.max_context + per - 1) / per) * bytes;
     void* p = nullptr;
-    if (!vmm_) {
+    if (!vmm_ || !opt_.kv_grow) {
         p = dmalloc(all, "the attention cache");
     } else {
         Kv k;
@@ -1525,11 +1577,13 @@ void Engine::load(const Options& o) {
     ev_ = ev_row_[0];
     blas_ = blas_create(st_);
     // The card maps VRAM onto reserved addresses a piece at a time: the attention caches grow with the sequence,
-    // and the expert cache has what a short one leaves (cover()).
+    // and the expert cache has what a short one leaves (cover()).  With --kv-whole the attention caches are one
+    // allocation each (kv_alloc()) and the expert cache still gives VRAM to other programs and takes it back
+    // (mind_vram()): refilled in one allocation instead, it only ever got smaller.
     ck(cudaGetDevice(&device_), "cudaGetDevice");
     CUdevice dev = 0;
     int vmm = 0;
-    if (opt_.kv_grow && cuDeviceGet(&dev, device_) == CUDA_SUCCESS &&
+    if (cuDeviceGet(&dev, device_) == CUDA_SUCCESS &&
         cuDeviceGetAttribute(&vmm, CU_DEVICE_ATTRIBUTE_VIRTUAL_MEMORY_MANAGEMENT_SUPPORTED, dev) == CUDA_SUCCESS && vmm) {
         CUmemAllocationProp prop = {};
         prop.type = CU_MEM_ALLOCATION_TYPE_PINNED;