KV streaming under WSL: pin the K/V host copy with cudaHostRegister

Under WSL2 the NVIDIA driver caps what cudaHostAlloc(Mapped) can pin at about
1 GiB in all, so the KV streaming host copy (layer.cpp) cannot be allocated
and setup turns streaming off there: the whole KV cache must sit in VRAM,
beside the expert cache. cudaHostRegister of ordinary memory is not capped the
same way; the expert arena already registers tens of GiB with it on the same
machines (pinned.cu).

kv_host_alloc() maps anonymous memory and registers it Portable|Mapped, before
any page is touched, when /dev/dxg exists (WSL2, detected as generate.cpp's
under_wddm() does). Elsewhere cudaHostAlloc stays first. Either way the other
method is the fallback when the first is refused, and
STRATA_KV_HOST_REGISTER=1/0 forces the order. On native Linux and Windows the
behaviour is unchanged unless cudaHostAlloc is refused (by inspection; not run
there). The failure message no longer says WSL pins only 1 GiB.

Measured on one RTX 5090, Ryzen 9 9950X3D, WSL2 (157 GiB), Windows driver
591.86, IQ3_S, --max-context 1048576 --rope-scaling yarn --rope-scale 4
--kv int8 --kv-resident 32768, with this commit:
- start: "the K/V in 12.38 GiB of pinned RAM"; expert cache 11,609 slots
- tools/needle_bench.py: 4/4 found, 128K and 1M (1,015,070 prompt tokens) at
  depths 10 and 90
- decode 99-124 tok/s with 0.5-1M tokens in context (4 runs); 150-191 tok/s
  medians from 55-token prompts to 128K (single runs 113-196); prefill about
  4,200 tok/s for a cold 1M prompt
kv_stream_parity still passes; it allocates its own host pools, so it checks
the streaming kernels, which this change does not touch.

Not tested: layer split or remote experts under WSL, Docker Desktop.
setup.py still never writes --kv-resident on WSL and removes it there
(upgrade_config); until that changes it has to be added by hand.

src/core/layer.cpp

@@ -33,11 +33,15 @@
 #include <algorithm>
 #include <cstdlib>
 #include <exception>
+#include <filesystem>
 #include <cuda_runtime.h>
 #include <cmath>
 #include <cstdio>
 #include <cstring>
 #include <limits>
+#ifndef _WIN32
+#include <sys/mman.h>
+#endif
 #include <memory>
 #include <vector>
 namespace strata::core {
@@ -517,6 +521,53 @@ namespace {
 int64_t g_kv_resident = 0;
 uint64_t g_kv_host_bytes = 0;
 
+// The KV host copy: pinned and device-mapped. cudaHostAlloc is the default, but under WSL2 the driver caps what
+// cudaHostAlloc(Mapped) may pin at about 1 GiB in all, while cudaHostRegister of ordinary memory is not capped there:
+// the expert arena registers tens of GiB the same way (pinned.cu). So under WSL2, detected as generate.cpp's
+// under_wddm() does (/dev/dxg, the device its GPU goes through to the Windows driver), the copy is anonymous memory
+// registered Portable|Mapped before any page is touched, which also leaves WSL's small cudaHostAlloc budget to the
+// staging buffers. Elsewhere cudaHostAlloc stays first. Either way the other method is the fallback when the first
+// is refused. STRATA_KV_HOST_REGISTER=1/0 forces registering first or cudaHostAlloc first. Like the cudaHostAlloc
+// copy it lives as long as the process.
+uint8_t* kv_host_alloc(uint64_t bytes) {
+    auto host_alloc = [](uint64_t n) -> uint8_t* {
+        uint8_t* h = nullptr;
+        if (cudaHostAlloc((void**) &h, n, cudaHostAllocMapped | cudaHostAllocPortable) == cudaSuccess) return h;
+        (void) cudaGetLastError();
+        return nullptr;
+    };
+#ifdef _WIN32
+    return host_alloc(bytes);
+#else
+    static const bool register_first = [] {
+        const char* e = std::getenv("STRATA_KV_HOST_REGISTER");
+        if (e != nullptr && (e[0] == '0' || e[0] == '1')) return e[0] == '1';
+        return std::filesystem::exists("/dev/dxg");
+    }();
+    auto register_anon = [](uint64_t n) -> uint8_t* {
+        void* m = mmap(nullptr, (size_t) n, PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
+        if (m == MAP_FAILED) return nullptr;
+        if (cudaHostRegister(m, (size_t) n, cudaHostRegisterPortable | cudaHostRegisterMapped) != cudaSuccess) {
+            (void) cudaGetLastError();
+            munmap(m, (size_t) n);
+            return nullptr;
+        }
+        static bool said = false;
+        if (!said) {
+            std::fprintf(stderr, "strata: KV streaming: the host copy is pinned with cudaHostRegister\n");
+            said = true;
+        }
+        return (uint8_t*) m;
+    };
+    if (register_first) {
+        if (uint8_t* m = register_anon(bytes)) return m;
+        return host_alloc(bytes);
+    }
+    if (uint8_t* h = host_alloc(bytes)) return h;
+    return register_anon(bytes);
+#endif
+}
+
 // ---- the elastic K/V (--kv-grow): one VMM range per state, each pool array at a chunk boundary in it
 bool g_kv_elastic = false;
 int64_t g_kv_elastic_init = 16384;
@@ -803,12 +854,12 @@ uint64_t qsa_state_init(const ModelGeometry& g, int64_t max_cells, void* base, Q
         const uint64_t bytes = (uint64_t) pages * strata::kernels::kv_block_bytes(s, qsa_kv_format(st)) + 4 * 256;
         uint8_t* h = nullptr;
         uint8_t* d = nullptr;
-        if (cudaHostAlloc((void**) &h, bytes, cudaHostAllocMapped | cudaHostAllocPortable) != cudaSuccess ||
+        if ((h = kv_host_alloc(bytes)) == nullptr ||
             cudaHostGetDevicePointer((void**) &d, h, 0) != cudaSuccess) {
-            // under WSL the NVIDIA driver pins only ~1 GiB in all, which is less than 128K of 8-bit KV needs
+            // both cudaHostRegister and cudaHostAlloc were refused (kv_host_alloc): usually not enough free RAM
             if (p.mode == 1) std::fprintf(stderr, "strata: KV streaming: cannot pin %.2f GiB of RAM for a layer's KV copy "
-                                 "(%.2f GiB pinned so far) - lower the context, or run without --kv-resident (under "
-                                 "WSL the driver pins only about 1 GiB in all)\n", (double) bytes / 1073741824.0,
+                                 "(%.2f GiB pinned so far): RAM or the driver refused to pin it - lower the context, "
+                                 "or run without --kv-resident\n", (double) bytes / 1073741824.0,
                                  (double) g_kv_host_bytes / 1073741824.0);
             return 0;
         }