lane core: several queued prompts in one pass (prefill_many, addStreams, the lane host admitting together)

zig/src/core/lane_host.zig

@@ -274,29 +274,48 @@ pub const LaneHost = struct {
         h.mutex.unlock(h.io);
     }
 
-    /// Prefills the next queued request into a free lane; false when none waits or no lane is free.
+    /// The next queued request into a free lane (or several, when the backend batches prompts); false when none can go.
     fn admitOne(h: *LaneHost) bool {
+        if (admit.batches(h)) return admit.many(h);
+        const job = switch (h.takeJob()) {
+            .none => return false,
+            .failed => return true,
+            .job => |j| j,
+        };
+        const began = h.setup(job) orelse return true;
+        return h.launch(job, began);
+    }
+
+    pub const Taken = union(enum) { none, failed, job: *Job };
+
+    /// The next queued request a lane is free for, moved to the admitted list (`failed`: it was finished instead).
+    pub fn takeJob(h: *LaneHost) Taken {
         h.lock();
         // a stale device takes no reply: `learn.refuse` ends the queued ones instead
         if (h.stale or h.queued.items.len == 0 or h.admitted.items.len >= @min(h.info_.lanes, h.memory_lanes orelse h.info_.lanes)) {
             h.unlock();
-            return false;
+            return .none;
         }
         const at = admit.pick(h) orelse { // a shared cache pool's head waits for room
             h.unlock();
-            return false;
+            return .none;
         };
         const job = h.queued.orderedRemove(at);
         job.seq = h.admissions;
         h.admissions += 1;
         h.admitted.append(h.gpa, job) catch {
             h.unlock();
             h.finish(job, .failed, "out of memory");
-            return true;
+            return .failed;
         };
         h.unlock();
+        return .{ .job = job };
+    }
+
+    /// A taken job's stream, set up to start; null when the job ended instead (dropped or cancelled).
+    pub fn setup(h: *LaneHost, job: *Job) ?i96 {
         const r = job.request;
-        if (r.background and r.max_tokens == 0 and h.warm_quiet > 0 and !h.quiet(job)) return h.cancel(job); // other work came first: it never ran
+        if (r.background and r.max_tokens == 0 and h.warm_quiet > 0 and !h.quiet(job)) return h.ended(h.cancel(job)); // other work came first
         var reuse: lanes.stream.Reuse = .{};
         // the entry stays alive until the backend restores it: nothing keeps between here and this stream's own pass
         if (h.cache) |store| if (r.prompt_logprobs == null and !r.untaught) if (store.lookupRewind(h.gpa, r.cacheKey(), r.bornLen(), r.history_len, r.rewind_len, r.shared_prefixes, r.chunks, h.spans(r))) |l| {
@@ -305,7 +324,7 @@ pub const LaneHost = struct {
             job.marks = l.marks;
             reuse = .{ .saved = if (l.entry) |e| e.saved else null, .at = if (l.entry) |e| e.at else 0, .marks = l.marks, .hook = .{ .ptr = job, .at = Job.kept } };
         } else |_| {};
-        job.proposer = lanes.SuffixLookup.init(h.gpa, .{ .min_match = h.min_match }) catch return h.drop(job, "the drafter could not start");
+        job.proposer = lanes.SuffixLookup.init(h.gpa, .{ .min_match = h.min_match }) catch return h.ended(h.drop(job, "the drafter could not start"));
         const recall = recalls.reader(h, job);
         job.stream = lanes.Stream.init(h.gpa, .{
             .id = "request",
@@ -332,7 +351,7 @@ pub const LaneHost = struct {
             .grammar = if (r.structure) |st| h.constraint(st) catch |e| {
                 job.proposer.deinit();
                 recalls.release(job);
-                return h.drop(job, if (e == error.OutOfMemory) "out of memory" else "the reply's grammar could not start (xgrammar)");
+                return h.ended(h.drop(job, if (e == error.OutOfMemory) "out of memory" else "the reply's grammar could not start (xgrammar)"));
             } else null,
             .decode_spans = h.spans(r),
             .decide = r.decide,
@@ -342,26 +361,46 @@ pub const LaneHost = struct {
         }) catch {
             job.proposer.deinit();
             recalls.release(job);
-            return h.drop(job, "out of memory");
+            return h.ended(h.drop(job, "out of memory"));
         };
         job.started = true;
         job.fill_on.store(true, .release); // status readers count the fill from here until the pass reports
         const began = std.Io.Clock.awake.now(h.io).toNanoseconds();
         job.began = began;
+        return began;
+    }
+
+    fn ended(_: *LaneHost, _: bool) ?i96 {
+        return null;
+    }
+
+    /// Whether a set-up job starts outside a shared prompt pass: a warm pass, or the lone driver's.
+    pub fn startsAlone(h: *LaneHost, job: *Job) bool {
+        return (job.request.max_tokens == 0 and h.info_.warm_turns) or lone.fits(h, job);
+    }
+
+    /// A set-up job into the engine: a warm pass, the lone driver, the fill queue, or its own prompt pass.
+    pub fn launch(h: *LaneHost, job: *Job, began: i96) bool {
+        const r = job.request;
         if (r.max_tokens == 0 and h.info_.warm_turns) return h.warmPass(job, began);
         if (lone.fits(h, job)) return lone.run(h, job, began, false);
         if (h.core.fills()) {
             h.filling.append(h.gpa, job) catch return h.drop(job, "out of memory");
             return true;
         }
-        h.core.addStream(&job.stream) catch |e| return switch (e) {
+        h.core.addStream(&job.stream) catch |e| return h.failedPass(job, e);
+        h.prefilled(job, began);
+        _ = h.deliver(job);
+        return true;
+    }
+
+    /// A prompt pass that failed: cancelled, no lane memory yet (back to the queue), or the backend's words.
+    pub fn failedPass(h: *LaneHost, job: *Job, e: anyerror) bool {
+        return switch (e) {
             error.Cancelled => h.cancel(job),
             error.LaneMemory => h.requeue(job),
             else => h.drop(job, h.words(e)),
         };
-        h.prefilled(job, began);
-        _ = h.deliver(job);
-        return true;
     }
 
     /// The next filling job's next prompt chunk: the other streams' rounds run between chunks, not after the prompt.

zig/src/core/lane_host_admit.zig

@@ -2,6 +2,7 @@
 const api = @import("engine_api.zig");
 const LaneHost = @import("lane_host.zig").LaneHost;
 const Job = LaneHost.Job;
+const lanes = @import("lanes");
 
 /// The queued job to admit (its index), or null to hold; the caller holds the lock and a lane is free.
 pub fn pick(h: *LaneHost) ?usize {
@@ -21,3 +22,53 @@ pub fn pick(h: *LaneHost) ?usize {
 fn fits(a: api.Admit, job: *const Job) bool {
     return a.fits(a.ctx, @intCast(job.request.prompt.len), job.request.max_tokens);
 }
+
+/// Whether to admit several at once: prompts batch, none fill by chunks, and two or more wait with two lanes free.
+pub fn batches(h: *LaneHost) bool {
+    if (!h.core.batchesPrompts() or h.core.fills()) return false;
+    h.lock();
+    defer h.unlock();
+    const lanes_now = @min(h.info_.lanes, h.memory_lanes orelse h.info_.lanes);
+    return h.queued.items.len >= 2 and lanes_now -| h.admitted.items.len >= 2;
+}
+
+/// Every queued request a lane is free for, prompts in one pass (lone and warm ones as before); false: none taken.
+pub fn many(h: *LaneHost) bool {
+    var jobs: [lanes.Engine.max_together]*Job = undefined;
+    var began: [lanes.Engine.max_together]i96 = undefined;
+    var n: usize = 0;
+    var took = false;
+    while (n < jobs.len) {
+        const job = switch (h.takeJob()) {
+            .none => break,
+            .failed => {
+                took = true;
+                continue;
+            },
+            .job => |j| j,
+        };
+        took = true;
+        const at = h.setup(job) orelse continue;
+        if (h.startsAlone(job)) {
+            _ = h.launch(job, at);
+            continue;
+        }
+        jobs[n] = job;
+        began[n] = at;
+        n += 1;
+    }
+    if (n == 1) return h.launch(jobs[0], began[0]);
+    var streams: [lanes.Engine.max_together]*lanes.Stream = undefined;
+    for (jobs[0..n], 0..) |job, i| streams[i] = &job.stream;
+    var errs: [lanes.Engine.max_together]?anyerror = undefined;
+    h.core.addStreams(streams[0..n], errs[0..n]);
+    for (jobs[0..n], began[0..n], errs[0..n]) |job, at, err| {
+        if (err) |e| {
+            _ = h.failedPass(job, e);
+            continue;
+        }
+        h.prefilled(job, at);
+        _ = h.deliver(job);
+    }
+    return took;
+}