Flash Next shared sessions: FLASHNEXT_LOAD_LIMIT_GB lets the slots take the server past 70% of RAM

docs/recipes/qwen3.8-flash-next.md

@@ -74,6 +74,11 @@ New requests join between rounds, and cancelling a request releases its slot for
 `auto` and `--parallel 1` retain the existing single-reply driver. Shared sessions currently run on one Mac;
 `--parallel N` above one with `--speed-up` is refused. Use `--parallel 1` for the two-Mac link.
 
+The ceiling is 70% of RAM, or `FLASHNEXT_LOAD_LIMIT_GB` (GB, up to the GPU's recommended working set). It covers the
+server's whole footprint: the model, the sessions and the prompt cache. Left out, `--prompt-cache-gib` takes what the
+sessions leave under it; a given budget must fit under the same ceiling, and only `--prompt-cache-over-cap` lets it
+pass.
+
 Additional slots allocate caches for the selected context. Startup reserves their metadata and the prompt-cache
 budget, then checks the physical footprint before accepting requests. A fixed slot count that does not fit is
 refused; lower the count, context or prompt-cache budget. `--prompt-cache-over-cap` overrides only the prompt-cache

zig/src/native/cache_fit.zig

@@ -15,13 +15,31 @@ pub const Fit = struct { budget: u64, room: u64, cap: u64 };
 /// The budget for `total` bytes of RAM and a footprint of `ready` bytes once loaded: `gib` when given and it fits (or `over`),
 /// else all that fits; a given budget past the room is refused.
 pub fn fit(total: u64, ready: u64, gib: ?f64, over: bool) error{CacheOverCap}!Fit {
-    const cap = total / 100 * SHARE_PERCENT;
+    return fitUnder(total / 100 * SHARE_PERCENT, ready, gib, over);
+}
+
+/// `fit` under a given ceiling `cap` (bytes) instead of 70% of RAM: a load limit covers the cache as well.
+pub fn fitUnder(cap: u64, ready: u64, gib: ?f64, over: bool) error{CacheOverCap}!Fit {
     const room = cap -| ready -| MARGIN;
     const want = if (gib) |g| (if (g > 0) std.math.lossyCast(u64, g * GiB) else 0) else room;
     if (want > room and !over) return error.CacheOverCap;
     return .{ .budget = want, .room = room, .cap = cap };
 }
 
+/// The ceiling a served model may hold: 70% of `total`, or `text` (a load limit in GB, as GLM_LOAD_LIMIT_GB) up to
+/// `most` (the smaller of RAM and the GPU's recommended working set); text that is not a number above 0 is refused.
+pub const Ceiling = struct { bytes: u64, default: u64, asked: bool = false, capped: bool = false };
+
+pub fn ceiling(total: u64, most: u64, text: ?[]const u8) error{BadLoadLimit}!Ceiling {
+    const default = total / 100 * SHARE_PERCENT;
+    const t = std.mem.trim(u8, text orelse "", " \t");
+    if (t.len == 0) return .{ .bytes = default, .default = default };
+    const gb = std.fmt.parseFloat(f64, t) catch return error.BadLoadLimit;
+    if (!std.math.isFinite(gb) or gb <= 0) return error.BadLoadLimit;
+    const top: f64 = @floatFromInt(@min(total, most));
+    return .{ .bytes = @intFromFloat(@min(gb * 1e9, top)), .default = default, .asked = true, .capped = gb * 1e9 > top };
+}
+
 /// This machine's RAM (null where the OS does not say).
 pub fn ram() ?u64 {
     return std.process.totalSystemMemory() catch null;
@@ -40,19 +58,27 @@ pub fn footprint() ?u64 {
 
 /// The cache plan: what 70% of RAM leaves past this process, or `gib` if it fits (else refused unless `over`).
 pub fn budget(gib: ?f64, over: bool, spare: u64, a: std.mem.Allocator, why: *[]const u8, rank: []const u8) !api.PromptCachePlan {
+    return budgetUnder(null, gib, over, spare, a, why, rank);
+}
+
+/// `budget` under a load limit `limit` (bytes; null: 70% of RAM), so the cache and the rest stay under one ceiling.
+pub fn budgetUnder(limit: ?u64, gib: ?f64, over: bool, spare: u64, a: std.mem.Allocator, why: *[]const u8, rank: []const u8) !api.PromptCachePlan {
     const total = ram() orelse 0;
     const ready = footprint() orelse 0;
     if (total == 0 or ready == 0) {
         const b = if (gib) |g| (if (g > 0) std.math.lossyCast(u64, g * GiB) else 0) else spare;
         std.log.info("prompt cache: {d:.1} GiB for kept prompt states{s} (no memory reading)", .{ gibs(b), rank });
         return .{ .source = if (gib != null) .explicit else .metal_working_set, .budget_bytes = b, .explicit_budget = gib != null };
     }
-    const f = fit(total, ready, gib, over) catch |e| {
-        const left = (fit(total, ready, null, false) catch unreachable).room; // without a given budget it never refuses
-        why.* = try std.fmt.allocPrint(a, "--prompt-cache-gib {d} would take this server past 70% of this Mac's memory: it holds {d:.1} GiB once loaded, 70% of {d:.0} GiB is {d:.1} GiB, and {d} GiB stays free for prompt buffers, so {d:.1} GiB is left for kept prompt states. Leave --prompt-cache-gib out to use that, pass a smaller one, or add --prompt-cache-over-cap to keep {d} GiB anyway.", .{ gib.?, gibs(ready), gibs(total), gibs(total / 100 * SHARE_PERCENT), MARGIN >> 30, gibs(left), gib.? });
+    const cap = limit orelse total / 100 * SHARE_PERCENT;
+    const f = fitUnder(cap, ready, gib, over) catch |e| {
+        const left = (fitUnder(cap, ready, null, false) catch unreachable).room; // without a given budget it never refuses
+        if (limit != null) {
+            why.* = try std.fmt.allocPrint(a, "--prompt-cache-gib {d} would take this server past its load limit: it holds {d:.1} GiB once loaded, the limit is {d:.1} GiB, and {d} GiB stays free for prompt buffers, so {d:.1} GiB is left for kept prompt states. Leave --prompt-cache-gib out to use that, pass a smaller one, raise the limit, or add --prompt-cache-over-cap to keep {d} GiB anyway.", .{ gib.?, gibs(ready), gibs(cap), MARGIN >> 30, gibs(left), gib.? });
+        } else why.* = try std.fmt.allocPrint(a, "--prompt-cache-gib {d} would take this server past 70% of this Mac's memory: it holds {d:.1} GiB once loaded, 70% of {d:.0} GiB is {d:.1} GiB, and {d} GiB stays free for prompt buffers, so {d:.1} GiB is left for kept prompt states. Leave --prompt-cache-gib out to use that, pass a smaller one, or add --prompt-cache-over-cap to keep {d} GiB anyway.", .{ gib.?, gibs(ready), gibs(total), gibs(cap), MARGIN >> 30, gibs(left), gib.? });
         return e;
     };
-    std.log.info("prompt cache: {d:.1} GiB from {d:.1} GiB free under the 70% cap ({d:.1} GiB in use once loaded, {d:.0} GiB of RAM, {d} GiB kept for prompt buffers){s}{s}", .{ gibs(f.budget), gibs(f.room), gibs(ready), gibs(total), MARGIN >> 30, if (f.budget > f.room) ", past the cap by --prompt-cache-over-cap" else "", rank });
+    std.log.info("prompt cache: {d:.1} GiB from {d:.1} GiB free under the {s} ({d:.1} GiB in use once loaded, {d:.0} GiB of RAM, {d} GiB kept for prompt buffers){s}{s}", .{ gibs(f.budget), gibs(f.room), if (limit != null) "load limit" else "70% cap", gibs(ready), gibs(total), MARGIN >> 30, if (f.budget > f.room) ", past the cap by --prompt-cache-over-cap" else "", rank });
     return .{
         .source = .physical_footprint,
         .budget_bytes = f.budget,
@@ -91,6 +117,41 @@ test "a given budget that fits is kept, a larger one is refused unless the overr
     try std.testing.expectEqual(@as(u64, 64 * GiB), (try fit(512 * GiB, 172 * GiB, 64, false)).budget); // a 512 GB Mac has room
 }
 
+test "under a load limit the automatic cache gets what the limit leaves, and a given one is held to the same limit" {
+    // 100 GiB of RAM, a 60 GiB load limit, 50 GiB in use once the sessions are allocated (the case in review)
+    const cap: u64 = 60 * GiB;
+    const ready: u64 = 50 * GiB;
+    const auto = try fitUnder(cap, ready, null, false);
+    try std.testing.expectEqual(@as(u64, 8 * GiB), auto.budget); // 60 - 50 - 2, not the 18 GiB 70% of RAM leaves
+    try std.testing.expect(ready + MARGIN + auto.budget <= cap);
+    try std.testing.expectEqual(@as(u64, 18 * GiB), (try fit(100 * GiB, ready, null, false)).budget); // 70% of RAM: past the limit
+    try std.testing.expectEqual(@as(u64, 8 * GiB), (try fitUnder(cap, ready, 8, false)).budget);
+    try std.testing.expectError(error.CacheOverCap, fitUnder(cap, ready, 16, false));
+    try std.testing.expectEqual(@as(u64, 16 * GiB), (try fitUnder(cap, ready, 16, true)).budget); // the override still passes it
+    // a raised limit raises both: past 70% of RAM a given cache that fits under the limit needs no override
+    const raised: u64 = 80 * GiB;
+    try std.testing.expectEqual(@as(u64, 16 * GiB), (try fitUnder(raised, ready, 16, false)).budget);
+    try std.testing.expectError(error.CacheOverCap, fit(100 * GiB, 60 * GiB, 16, false));
+    try std.testing.expectEqual(@as(u64, 18 * GiB), (try fitUnder(raised, 60 * GiB, null, false)).budget);
+}
+
+test "the ceiling: 70% of RAM, or a load limit in GB up to RAM and the GPU's working set; bad text refused" {
+    const ram256: u64 = 256 * GiB;
+    const ws: u64 = 239_000_000_000;
+    const d = try ceiling(ram256, ws, null);
+    try std.testing.expectEqual(ram256 / 100 * 70, d.bytes);
+    try std.testing.expect(!d.asked and !d.capped);
+    const asked = try ceiling(ram256, ws, "215");
+    try std.testing.expectEqual(@as(u64, 215_000_000_000), asked.bytes);
+    try std.testing.expect(asked.asked and !asked.capped and asked.bytes > asked.default);
+    try std.testing.expect((try ceiling(ram256, ws, " 300 ")).capped);
+    try std.testing.expectEqual(ws, (try ceiling(ram256, ws, "300")).bytes);
+    try std.testing.expectEqual(ram256 / 100 * 70, (try ceiling(ram256, ws, "  ")).bytes);
+    try std.testing.expectError(error.BadLoadLimit, ceiling(ram256, ws, "lots"));
+    try std.testing.expectError(error.BadLoadLimit, ceiling(ram256, ws, "0"));
+    try std.testing.expectError(error.BadLoadLimit, ceiling(ram256, ws, "-4"));
+}
+
 test "the OS readings: RAM and this process's footprint" {
     if (builtin.os.tag != .macos) return error.SkipZigTest;
     const r = ram() orelse return error.NoRam;

zig/src/native/flashnext_batch_host.zig

@@ -21,6 +21,19 @@ const Host = struct {
     warm: mtl.keepalive.Target,
 };
 
+/// What the server may hold in all, sessions and prompt cache: 70% of RAM, or FLASHNEXT_LOAD_LIMIT_GB (GB, as
+/// GLM_LOAD_LIMIT_GB) up to the GPU's recommended working set; past the default the log names the risk, as GLM's does.
+fn ceiling(total: u64, working_set: u64) error{BadLoadLimit}!u64 {
+    const text: ?[]const u8 = if (std.c.getenv("FLASHNEXT_LOAD_LIMIT_GB")) |v| std.mem.span(v) else null;
+    const c = fit.ceiling(total, working_set, text) catch |err| {
+        std.log.err("flash next: FLASHNEXT_LOAD_LIMIT_GB={s} is not a load limit: give a number of GB above 0, or unset it for 70% of RAM", .{text.?});
+        return err;
+    };
+    if (c.capped) std.log.info("flash next: FLASHNEXT_LOAD_LIMIT_GB={s} passes this GPU's recommended working set; the limit is {d:.1} GB", .{ text.?, fit.gibs(c.bytes) * 1.073741824 });
+    if (c.bytes > c.default) std.log.warn("flash next: FLASHNEXT_LOAD_LIMIT_GB lets the server (sessions and prompt cache) take this Mac to {d:.1} GiB, past the default {d:.1} GiB (70% of RAM): a load that wires most of the RAM can stall macOS until its watchdog kills the server", .{ fit.gibs(c.bytes), fit.gibs(c.default) });
+    return c.bytes;
+}
+
 pub fn open(gpa: std.mem.Allocator, io: std.Io, dir: []const u8, dump: ?[]const u8, window: i64, streams: u32, cache_gib: ?f64, over: bool, a: std.mem.Allocator, why: *[]const u8) !api.Opened {
     const pool = mtl.objc.Pool.push();
     defer pool.pop();
@@ -32,9 +45,10 @@ pub fn open(gpa: std.mem.Allocator, io: std.Io, dir: []const u8, dump: ?[]const
     eng.hostMode();
     const ready = fit.footprint() orelse return error.NoMemoryReading;
     const total = fit.ram() orelse return error.NoMemoryReading;
+    const cap = try ceiling(total, eng.r.device.maxWorkingSet());
     const reserved = if (!over) (if (cache_gib) |g| (if (g > 0) std.math.lossyCast(u64, g * fit.GiB) else 0) else 0) else 0;
     const fixed = Session.overhead(tf.flashnext_replay.CAP) + tf.flashnext_replay.MAXR * tf.flashnext_replay.WIDE * 2;
-    const available = total / 100 * fit.SHARE_PERCENT -| ready -| fit.MARGIN -| reserved;
+    const available = cap -| ready -| fit.MARGIN -| reserved;
     if (available < fixed) return error.OverMemoryLimit;
     const room = available - fixed;
     const slots = try tf.flashnext_batch_meta.fit(streams, true, room, Session.bytes(capacity));
@@ -45,14 +59,13 @@ pub fn open(gpa: std.mem.Allocator, io: std.Io, dir: []const u8, dump: ?[]const
     h.back = try Backend.init(gpa, eng, slots, capacity);
     errdefer h.back.deinit();
     const allocated = fit.footprint() orelse return error.NoMemoryReading;
-    if (allocated +| fit.MARGIN +| reserved > total / 100 * fit.SHARE_PERCENT)
-        return error.OverMemoryLimit;
+    if (allocated +| fit.MARGIN +| reserved > cap) return error.OverMemoryLimit;
     h.cfg = try lanes.Config.init(gpa, h.back.facts(), tf.flashnext_replay.MAXR, tf.flashnext_replay.MAXR - 1);
     errdefer h.cfg.deinit(gpa);
     h.clock = .{ .io = io };
     h.core = lanes.Engine.init(gpa, &h.cfg, h.back.backend(), h.clock.clock());
     errdefer h.core.deinit();
-    const plan = try fit.budget(cache_gib, over, 0, a, why, "");
+    const plan = try fit.budgetUnder(cap, cache_gib, over, 0, a, why, ""); // the cache under the same limit as the sessions
     eng.snap_pool.max = plan.budget_bytes;
     h.host = api.LaneHost.init(gpa, io, &h.core, .{ .name = "flashnext-zig", .lanes = slots, .greedy_only = true, .context_window = @intCast(window), .prefill_step = @intCast(eng.pr.step), .prompt_cache_plan = plan });
     h.cache = null;

zig/src/native/metal.zig

@@ -307,7 +307,7 @@ fn openFlashNext(a: Allocator, gpa: Allocator, io: std.Io, o: api.Open, problem:
         }
         var why: []const u8 = "";
         return @import("flashnext_batch_host.zig").open(gpa, io, o.dir, dump, if (window > 0) @min(window, tf.flashnext_replay.CAP - tf.flashnext_engine.MARGIN) else tf.flashnext_replay.CAP - tf.flashnext_engine.MARGIN, o.lanes, o.prompt_cache_gib, o.prompt_cache_over_cap, a, &why) catch |err| {
-            problem.* = if (err == error.CacheOverCap) why else try std.fmt.allocPrint(a, "the native Flash Next batching engine cannot open ({s})", .{@errorName(err)});
+            problem.* = if (err == error.CacheOverCap) why else if (err == error.BadLoadLimit) "FLASHNEXT_LOAD_LIMIT_GB is not a load limit: give a number of GB above 0, or unset it for 70% of RAM" else if (err == error.OverMemoryLimit) "the native Flash Next batching engine cannot open (OverMemoryLimit): the sessions and the given prompt cache do not fit under the load limit (70% of RAM, or FLASHNEXT_LOAD_LIMIT_GB) past the loaded model; fewer streams, a smaller --context or --prompt-cache-gib, or a higher FLASHNEXT_LOAD_LIMIT_GB" else try std.fmt.allocPrint(a, "the native Flash Next batching engine cannot open ({s})", .{@errorName(err)});
             return null;
         };
     }