Kolibri-1 on CUDA: prompts in the shared KV store with TF_KV_SHARE (off by default)

CHANGELOG.md

@@ -4,6 +4,8 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel
 
 ## Unreleased
 
+- `TF_KV_SHARE=N` now keeps Kolibri-1's prompts in CUDA's shared paged KV store too, as for Gemma 4. It is off by default (#1657).
+
 - GLM, GLM-DSA and Flash Next on Metal drop their checkpoint from the file cache once loaded, keeping the SSD expert and PLE files that serving still reads (#1868).
 
 - On Macs, the q27, Nemotron, LFM2 and Qwen3.5 loaders drop their checkpoint from the file cache once the weights are loaded, so unified memory doesn't hold them twice: on an M5 Ultra a loaded Qwen3.8-27B frees about 14 GiB and Nemotron 3.5 Lightning about 17 GiB, with the same replies (part of #1868).

docs/recipes/kolibri-1.md

@@ -48,6 +48,11 @@ rounds run eager. Both give a row the same bits.
 - Refused at load: other quantizations, `norm_topk_prob`, other layer kinds, a config with a vision tower, shapes the
   shared kernels do not take, `--vision`, `--precision` and a captured kernel set (`TENSORFOLD_CUDA_KERNELS`).
 
+`TF_KV_SHARE=N` selects the experimental shared prompt store, with at most N kept states; unset or zero keeps the
+existing linear-cache path. Full-layer K/V is held as core-pool pages, while each snapshot keeps the sliding layers'
+last 512 K/V rows in pages of the same pool. Shared keeps and resumes use 256-token page edges. Copy drafts use
+the existing lane core; this does not add EAGLE-3 or a draft head. The new paged prompt kernels and complete shared-cache checkpoint path still require CUDA qualification.
+
 ## Checks
 
 ```bash

zig/build/kolibri.zig

@@ -9,6 +9,7 @@ pub fn host(b: *std.Build, m: *std.Build.Module, api: *std.Build.Module, all: *s
     const run = b.addRunArtifact(b.addTest(.{ .root_module = m }));
     b.step("test-kolibri-host", "Kolibri config, ownership and copy-draft metadata without a GPU").dependOn(&run.step);
     all.dependOn(&run.step);
+    @import("kolibri_shared_host.zig").host(b, m.import_table.get("cuda").?, m.import_table.get("core").?, m.import_table.get("lanes").?, api, m, all);
     b.step("test-kolibri-context-host", "Logical Kolibri context and physical verification reserve").dependOn(
         &b.addRunArtifact(b.addTest(.{ .root_module = m, .filters = &.{"Kolibri context:"} })).step,
     );

zig/build/kolibri_shared_host.zig

@@ -0,0 +1,9 @@
+//! Real native prompt adapter, pool and Kolibri state copies under an inert CUDA driver.
+const std = @import("std");
+pub fn host(b: *std.Build, cuda: *std.Build.Module, core: *std.Build.Module, lanes: *std.Build.Module, api: *std.Build.Module, kolibri: *std.Build.Module, all: *std.Build.Step) void {
+    const adapter = b.createModule(.{ .root_source_file = b.path("zig/src/native/cuda_prompt_pages.zig"), .target = b.graph.host, .link_libc = true, .imports = &.{ .{ .name = "core", .module = core }, .{ .name = "engine_api", .module = api } } });
+    const root = b.createModule(.{ .root_source_file = b.path("zig/kolibri_shared_host.zig"), .target = b.graph.host, .link_libc = true, .imports = &.{ .{ .name = "cuda", .module = cuda }, .{ .name = "core", .module = core }, .{ .name = "lanes", .module = lanes }, .{ .name = "kolibri1", .module = kolibri }, .{ .name = "prompt_pages", .module = adapter } } });
+    const run = b.addRunArtifact(b.addTest(.{ .root_module = root, .filters = &.{"Kolibri shared:"} }));
+    b.step("test-kolibri-shared-host", "Kolibri paged resumes, shared streams and copy-draft cache bindings without a GPU").dependOn(&run.step);
+    all.dependOn(&run.step);
+}

zig/build/spark25.zig

@@ -9,4 +9,13 @@ pub fn host(b: *std.Build, m: *std.Build.Module, all: *std.Build.Step) void {
     const run = b.addRunArtifact(b.addTest(.{ .root_module = m }));
     b.step("test-spark25-host", "Spark2.5 config admission, RoPE tables, cache sizes, snapshot runs and refusals without a GPU").dependOn(&run.step);
     all.dependOn(&run.step);
+    const cb = @import("cuda.zig");
+    const inert = b.addWriteFiles().add("inert-load-only.fatbin", "INERT DRIVER FIXTURE, NEVER A GPU IMAGE");
+    const images: [cb.kernels.len]?std.Build.LazyPath = @splat(inert);
+    const fake_cuda = cb.runtime(b, b.graph.host, .debug, &images);
+    const fake = module(b, b.graph.host, .debug, fake_cuda, m.import_table.get("core").?, m.import_table.get("lanes").?);
+    const root = b.createModule(.{ .root_source_file = b.path("zig/spark25_load_host.zig"), .target = b.graph.host, .link_libc = true, .imports = &.{ .{ .name = "cuda", .module = fake_cuda }, .{ .name = "spark25", .module = fake } } });
+    const load = b.addRunArtifact(b.addTest(.{ .root_module = root, .filters = &.{"Spark load:"} }));
+    b.step("test-spark25-load-host", "Actual Spark kernel loading over inert images with every driver lookup observed").dependOn(&load.step);
+    all.dependOn(&load.step);
 }

zig/kernels/cuda/window_attention.cu

@@ -8,6 +8,7 @@
 #include "attn_odd.cuh"
 #include "common/qmm_frag.cuh"
 #include "common/cp_async.cuh"
+#include "common/kv_pages.cuh"
 #include "window_split.cuh"
 
 // Prompt rows: prefill_attention.cu's tensor-core kernel with a window (keys (pos - win, pos]) and a ring of keys.
@@ -57,10 +58,10 @@ __device__ __forceinline__ float rowsum(const float (&c)[8]) {
 }
 
 // ODD: every NaN and infinite value leaves the mma, and the rows that see them add them back at float64's weights.
-template <int D, int WARPS, int HPC, int NS, bool ODD>
+template <int D, int WARPS, int HPC, int NS, bool ODD, typename... P>
 __device__ __forceinline__ bool wattn_rows(const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restrict__ kc,
                                            const __nv_bfloat16* __restrict__ vc, __nv_bfloat16* __restrict__ out,
-                                           int p0, int W, int H, int HK, int G, float scale, int win, int ring) {
+                                           int p0, int W, int H, int HK, int G, float scale, int win, int ring, const P&... pg) {
     constexpr int C = D / 8, K16 = D / 16, NB = D / 8, RB = WARPS / HPC, THREADS = 32 * WARPS;
     constexpr bool QREG = D <= 128;               // the queries' fragments stay in registers, else in shared memory
     extern __shared__ uint4 wsmem[];
@@ -81,7 +82,13 @@ __device__ __forceinline__ bool wattn_rows(const __nv_bfloat16* __restrict__ q,
             const int key = key0 + i / C;
             const bool in = key < limit;             // past the cache's filled keys: zeros, never read
             const int row = in ? (ring > 0 ? key % ring : key) : 0; // a ring keeps key p at row p % ring
-            cp16z(dst + sw<D>(i / C, i % C), src + ((int64_t)row * HK + hk) * D + (i % C) * 8, in);
+            if constexpr (sizeof...(P) != 0) {
+                const tf_common::Pg& t = (pg, ...);
+                const int k = in ? key : 0;
+                const int64_t at = (int64_t)t.tab[k / t.tok] * t.page + ((int64_t)(k % t.tok) * HK + hk) * D;
+                cp16z(dst + sw<D>(i / C, i % C), src + at + (i % C) * 8, in);
+            } else
+                cp16z(dst + sw<D>(i / C, i % C), src + ((int64_t)row * HK + hk) * D + (i % C) * 8, in);
         }
     };
     uint32_t qa[QREG ? K16 : 1][4];
@@ -234,8 +241,14 @@ __device__ __forceinline__ bool wattn_rows(const __nv_bfloat16* __restrict__ q,
     if constexpr (!ODD) { // a NaN or infinite output: the block runs again as ODD (finite values: the same bits)
         if (__syncthreads_or(attn_odd::nonfinite_out<NB>(o, r0 < W, r1 < W))) return true;
     } else {
-        const auto row = [&](int k) { return vc + ((int64_t)(ring > 0 ? k % ring : k) * HK + hk) * D; };
-        const auto krow = [&](int k) { return kc + ((int64_t)(ring > 0 ? k % ring : k) * HK + hk) * D; };
+        const auto offset = [&](int k) -> int64_t {
+            if constexpr (sizeof...(P) != 0) {
+                const tf_common::Pg& t = (pg, ...);
+                return (int64_t)t.tab[k / t.tok] * t.page + ((int64_t)(k % t.tok) * HK + hk) * D;
+            } else return ((int64_t)(ring > 0 ? k % ring : k) * HK + hk) * D;
+        };
+        const auto row = [&](int k) { return vc + offset(k); };
+        const auto krow = [&](int k) { return kc + offset(k); };
         const auto lo = [&](int r) { return win > 0 ? max(pw + r - win + 1, 0) : 0; };
         const auto vis = [&](int r, int k) { return wrow + r < W && k >= lo(r) && k <= pw + r; };
         const auto each = [&](int r, auto f) { for (int k = lo(r); k <= pw + r; ++k) f(krow(k)); }; // row r's keys
@@ -256,11 +269,11 @@ __device__ __forceinline__ bool wattn_rows(const __nv_bfloat16* __restrict__ q,
 }
 
 // A block with a NaN or infinite output, run again with ODD on (out of the common path).
-template <int D, int WARPS, int HPC, int NS>
+template <int D, int WARPS, int HPC, int NS, typename... P>
 __device__ __noinline__ void wattn_odd(const __nv_bfloat16* q, const __nv_bfloat16* kc, const __nv_bfloat16* vc,
                                        __nv_bfloat16* out, int p0, int W, int H, int HK, int G, float scale, int win,
-                                       int ring) {
-    wattn_rows<D, WARPS, HPC, NS, true>(q, kc, vc, out, p0, W, H, HK, G, scale, win, ring);
+                                       int ring, const P&... pg) {
+    wattn_rows<D, WARPS, HPC, NS, true>(q, kc, vc, out, p0, W, H, HK, G, scale, win, ring, pg...);
 }
 
 template <int D, int WARPS, int HPC, int NS>
@@ -533,3 +546,20 @@ template __global__ void tf_window_prompt::wattn_kernel<128, 8, 4, 8>(const __nv
 template __global__ void tf_window_prompt::wattn_kernel<128, 8, 8, 8>(const __nv_bfloat16* __restrict__,
     const __nv_bfloat16* __restrict__, const __nv_bfloat16* __restrict__, __nv_bfloat16* __restrict__, int, int, int,
     int, int, float, int, int);
+
+// The same prompt rows and nonfinite replay, with only their cache addresses taken from pages.
+extern "C" __global__ void __launch_bounds__(256, 1) tf_window_prompt_128_h4_pg(
+    const __nv_bfloat16* q, const __nv_bfloat16* kc, const __nv_bfloat16* vc, __nv_bfloat16* out,
+    int p0, int W, int H, int HK, int G, float scale, int win, int ring, const int* tab, long long page, int tok) {
+    const tf_common::Pg pg{tab, page, tok, 0};
+    if (tf_window_prompt::wattn_rows<128, 8, 4, 8, false>(q, kc, vc, out, p0, W, H, HK, G, scale, win, ring, pg))
+        tf_window_prompt::wattn_odd<128, 8, 4, 8>(q, kc, vc, out, p0, W, H, HK, G, scale, win, ring, pg);
+}
+
+extern "C" __global__ void __launch_bounds__(256, 1) tf_window_prompt_128_h8_pg(
+    const __nv_bfloat16* q, const __nv_bfloat16* kc, const __nv_bfloat16* vc, __nv_bfloat16* out,
+    int p0, int W, int H, int HK, int G, float scale, int win, int ring, const int* tab, long long page, int tok) {
+    const tf_common::Pg pg{tab, page, tok, 0};
+    if (tf_window_prompt::wattn_rows<128, 8, 8, 8, false>(q, kc, vc, out, p0, W, H, HK, G, scale, win, ring, pg))
+        tf_window_prompt::wattn_odd<128, 8, 8, 8>(q, kc, vc, out, p0, W, H, HK, G, scale, win, ring, pg);
+}

zig/kolibri_page_driver.zig

@@ -0,0 +1,90 @@
+//! Inert CUDA allocations and ordered copies; no GPU or kernel arithmetic is executed.
+const std = @import("std");
+const cuda = @import("cuda");
+const a = std.testing.allocator;
+const Entry = struct { ptr: u64 = 0, len: usize = 0 };
+var blocks: [512]Entry = @splat(.{});
+pub var skip_copy = false;
+pub var allocations: usize = 0;
+pub var graphs_dropped: usize = 0;
+
+fn alloc(out: *u64, n: usize) callconv(.c) cuda.abi.Result {
+    const mem = a.alignedAlloc(u8, .fromByteUnits(256), n) catch return 2;
+    @memset(mem, 0);
+    for (&blocks) |*b| if (b.ptr == 0) {
+        b.* = .{ .ptr = @intFromPtr(mem.ptr), .len = n };
+        allocations += 1;
+        out.* = b.ptr;
+        return 0;
+    };
+    a.free(mem);
+    return 2;
+}
+fn free(ptr: u64) callconv(.c) cuda.abi.Result {
+    for (&blocks) |*b| if (b.ptr == ptr) {
+        const mem: [*]align(256) u8 = @ptrFromInt(ptr);
+        a.free(mem[0..b.len]);
+        b.* = .{};
+        return 0;
+    };
+    return 1;
+}
+fn h2d(dst: u64, src: ?*const anyopaque, n: usize, _: cuda.abi.Stream) callconv(.c) cuda.abi.Result {
+    const out: [*]u8 = @ptrFromInt(dst);
+    const in: [*]const u8 = @ptrCast(src.?);
+    @memcpy(out[0..n], in[0..n]);
+    return 0;
+}
+fn d2d(dst: u64, src: u64, n: usize, _: cuda.abi.Stream) callconv(.c) cuda.abi.Result {
+    if (!skip_copy) {
+        const out: [*]u8 = @ptrFromInt(dst);
+        const in: [*]const u8 = @ptrFromInt(src);
+        @memcpy(out[0..n], in[0..n]);
+    }
+    return 0;
+}
+fn fill(dst: u64, value: u8, n: usize, _: cuda.abi.Stream) callconv(.c) cuda.abi.Result {
+    const out: [*]u8 = @ptrFromInt(dst);
+    @memset(out[0..n], value);
+    return 0;
+}
+fn sync(_: cuda.abi.Stream) callconv(.c) cuda.abi.Result {
+    return 0;
+}
+fn contextSync() callconv(.c) cuda.abi.Result {
+    return 0;
+}
+fn dropGraph(_: cuda.abi.GraphExec) callconv(.c) cuda.abi.Result {
+    graphs_dropped += 1;
+    return 0;
+}
+
+pub fn driver() cuda.Driver {
+    var api: cuda.abi.Api = undefined;
+    api.cuMemAlloc_v2 = alloc;
+    api.cuMemFree_v2 = free;
+    api.cuMemcpyHtoDAsync_v2 = h2d;
+    api.cuMemcpyDtoDAsync_v2 = d2d;
+    api.cuMemsetD8Async = fill;
+    api.cuStreamSynchronize = sync;
+    api.cuCtxSynchronize = contextSync;
+    api.cuGraphExecDestroy = dropGraph;
+    api.cuLaunchKernel = launch;
+    return .{ .lib = undefined, .api = api };
+}
+
+pub fn empty() bool {
+    for (blocks) |b| if (b.ptr != 0) return false;
+    return true;
+}
+
+pub var last_function: usize = 0;
+pub var last_table: u64 = 0;
+fn launch(fun: cuda.abi.Function, _: c_uint, _: c_uint, _: c_uint, _: c_uint, _: c_uint, _: c_uint, _: c_uint, _: cuda.abi.Stream, args: ?[*]?*anyopaque, _: ?[*]?*anyopaque) callconv(.c) cuda.abi.Result {
+    last_function = @intFromPtr(fun);
+    if (last_function >= 3) {
+        const ptr: *const u64 = @ptrCast(@alignCast(args.?[12].?));
+        last_table = ptr.*;
+    } else last_table = 0;
+    return 0;
+}