Flash Next CUDA: /v1/decisions on one GPU, label scores from the prompt's last row

docs/api.md

@@ -99,7 +99,7 @@ Unsupported image input, audio, video and non-text output requests receive HTTP
 
 ## Decisions
 
-`POST /v1/decisions` is served by the MLX server and by the CUDA GLM engine.
+`POST /v1/decisions` is served by the MLX server, by the CUDA GLM engine and by Flash Next on one CUDA GPU.
 Another CUDA engine, one without label scoring, returns HTTP 400.
 The prompt wording is SGLang's decision prompt format version 1: the input, a blank line, the question, one line per
 option, level, or described yes or no answer, and a closing instruction to answer with one label. Choice labels are

zig/src/core/lane_host.zig

@@ -24,6 +24,7 @@ const Stats = api.Stats;
 pub const Rounds = rounds.Rounds;
 pub const Proposers = rounds.Proposers;
 const rounds = @import("lane_rounds.zig");
+const scoring = @import("lane_host_score.zig");
 const lane_stream = @import("lane_stream.zig");
 const lane_halt = @import("lane_halt.zig");
 
@@ -69,6 +70,7 @@ pub const LaneHost = struct {
     memory_lanes: ?usize = null, // lanes the backend had memory for at its last LaneMemory, until a job ends
     warm_quiet: u64 = 0, // ns a warm pass waits for other work before its first chunk (0: none), set by its family
     disk: rows_host.Disk = .{}, // the Mac app's questions about kept states, and the disk writer's quiet
+    scores: scoring.Queue = .{}, // decisions' label scores waiting for the engine thread (lane_host_score.zig)
     /// Runs first on the engine thread: a backend whose API state is per thread (a CUDA context) binds it there.
     on_thread: ?ThreadInit = null,
     /// opt-in: a family round planner (prompt pieces, decode rows in one pass); null: whole or `prefill_step` prompts
@@ -97,6 +99,7 @@ pub const LaneHost = struct {
     pub fn stop(h: *LaneHost) void {
         rows_host.drainThread(h);
         rows_host.close(h);
+        scoring.close(h);
         const closed: api.LearnEvent = .{ .done = .{ .message = "the engine closed" } };
         for (h.lessons.items) |l| l.sink.event(l.sink.ctx, &closed);
         h.lessons.deinit(h.gpa);
@@ -114,19 +117,28 @@ pub const LaneHost = struct {
     pub const halt = lane_halt.halt;
 
     pub fn engine(h: *LaneHost) Engine {
-        return .{ .ctx = h, .vtable = &.{
-            .info = infoFn,
-            .submit = submitFn,
-            .cancel = cancelFn,
-            .status = status.read,
-            .memory = memoryFn,
-            .keepalive = keepaliveFn,
-            .learn = learn.submit,
-            .recall = recalls.attach,
-            .imprint = rows_host.ask,
-            .drain = rows_host.drain,
-        } };
-    }
+        // decisions' label scores only when the backend gives them, so the route refuses as before otherwise
+        return .{ .ctx = h, .vtable = if (h.core.backend.vtable.score != null) &scored_vtable else &plain_vtable };
+    }
+
+    const plain_vtable: Engine.VTable = .{
+        .info = infoFn,
+        .submit = submitFn,
+        .cancel = cancelFn,
+        .status = status.read,
+        .memory = memoryFn,
+        .keepalive = keepaliveFn,
+        .learn = learn.submit,
+        .recall = recalls.attach,
+        .imprint = rows_host.ask,
+        .drain = rows_host.drain,
+    };
+
+    const scored_vtable: Engine.VTable = blk: {
+        var v = plain_vtable;
+        v.score = scoring.score;
+        break :blk v;
+    };
 
     /// The family's queue as a keepalive target, when it set one.
     fn keepaliveFn(ctx: *anyopaque) ?api.keepalive.Target {
@@ -506,6 +518,7 @@ pub const LaneHost = struct {
         while (true) {
             ends.takeCancels(h);
             rows_host.answer(h);
+            scoring.answer(h);
             learn.refuse(h);
             while (h.admitOne()) {}
             h.fillOne();

zig/src/core/lane_host_score.zig

@@ -0,0 +1,75 @@
+//! Engine.score for a lanes backend that scores labels: a request thread asks, the engine thread answers.
+const std = @import("std");
+const LaneHost = @import("lane_host.zig").LaneHost;
+
+/// One scoring question, on the asking thread's stack until `done`.
+pub const Ask = struct {
+    prompt: []const u32,
+    labels: []const u32,
+    logits: []f64,
+    mass: bool,
+    logsumexp: f64 = 0,
+    failed: bool = false,
+    done: bool = false,
+};
+
+/// Questions waiting for the engine thread (under the host's lock).
+pub const Queue = struct {
+    asks: std.ArrayList(*Ask) = .empty,
+};
+
+/// Engine.score: queued for the engine thread, which runs it between two rounds; Failed when the host is closing.
+pub fn score(ctx: *anyopaque, prompt: []const u32, labels: []const u32, logits: []f64, mass: bool) error{Failed}!f64 {
+    const h: *LaneHost = @ptrCast(@alignCast(ctx));
+    var ask: Ask = .{ .prompt = prompt, .labels = labels, .logits = logits, .mass = mass };
+    h.lock();
+    defer h.unlock();
+    if (h.closing) return error.Failed;
+    h.scores.asks.append(h.gpa, &ask) catch return error.Failed;
+    h.wake.broadcast(h.io);
+    while (!ask.done) h.wake.waitUncancelable(h.io, &h.mutex);
+    if (ask.failed) return error.Failed;
+    return ask.logsumexp;
+}
+
+/// Every queued question answered by the backend (engine thread, between rounds).
+pub fn answer(h: *LaneHost) void {
+    h.lock();
+    if (h.scores.asks.items.len == 0) {
+        h.unlock();
+        return;
+    }
+    const list = h.scores.asks.toOwnedSlice(h.gpa) catch {
+        h.unlock();
+        return;
+    };
+    h.unlock();
+    defer h.gpa.free(list);
+    const backend = h.core.backend;
+    for (list) |a| {
+        const run = backend.vtable.score orelse {
+            a.failed = true;
+            continue;
+        };
+        if (run(backend.ptr, a.prompt, a.labels, a.logits, a.mass)) |lse| a.logsumexp = lse else |err| {
+            std.log.warn("lane host: a decision score failed ({s})", .{@errorName(err)});
+            a.failed = true;
+        }
+    }
+    h.lock();
+    for (list) |a| a.done = true;
+    h.wake.broadcast(h.io);
+    h.unlock();
+}
+
+/// Questions left when the engine stops: each fails.
+pub fn close(h: *LaneHost) void {
+    h.lock();
+    for (h.scores.asks.items) |a| {
+        a.failed = true;
+        a.done = true;
+    }
+    h.scores.asks.clearAndFree(h.gpa);
+    h.wake.broadcast(h.io);
+    h.unlock();
+}

zig/src/core/lane_host_score_test.zig

@@ -0,0 +1,61 @@
+//! Decisions' label scores reach a lanes backend that gives them, on the engine thread, and nowhere else.
+const std = @import("std");
+const lanes = @import("lanes");
+const api = @import("engine_api.zig");
+const LaneHost = @import("lane_host.zig").LaneHost;
+
+/// The fake backend's entries with a score of its own: label i scores prompt.len + i, the mass 7.5.
+const Scorer = struct {
+    var engine_thread: ?std.Thread.Id = null;
+    var calls: usize = 0;
+
+    fn score(_: *anyopaque, prompt: []const u32, labels: []const u32, logits: []f64, mass: bool) anyerror!f64 {
+        engine_thread = std.Thread.getCurrentId();
+        calls += 1;
+        if (labels.len == 0) return error.EmptyLabels;
+        for (logits, 0..) |*v, i| v.* = @floatFromInt(prompt.len + i);
+        return if (mass) 7.5 else 0;
+    }
+};
+
+test "an engine whose backend gives no score refuses decisions as unsupported" {
+    const gpa = std.testing.allocator;
+    var cfg = try lanes.Config.init(gpa, .{}, 1, 0);
+    defer cfg.deinit(gpa);
+    var target: lanes.fake.Fake = .{ .gpa = gpa };
+    defer target.deinit();
+    var clock: lanes.fake.FixedClock = .{};
+    var core = lanes.Engine.init(gpa, &cfg, target.backend(), clock.clock());
+    defer core.deinit();
+    var host = LaneHost.init(gpa, std.testing.io, &core, .{ .lanes = 1 });
+    var logits: [2]f64 = undefined;
+    try std.testing.expectError(error.Unsupported, host.engine().score(&.{ 1, 2 }, &.{ 3, 4 }, &logits, true));
+}
+
+test "a backend's score runs on the engine thread, a failure comes back as Failed, and a stopped host fails" {
+    const gpa = std.testing.allocator;
+    var cfg = try lanes.Config.init(gpa, .{}, 1, 0);
+    defer cfg.deinit(gpa);
+    var target: lanes.fake.Fake = .{ .gpa = gpa };
+    defer target.deinit();
+    const plain = target.backend();
+    var vt = plain.vtable.*;
+    vt.score = Scorer.score;
+    var clock: lanes.fake.FixedClock = .{};
+    var core = lanes.Engine.init(gpa, &cfg, .{ .ptr = plain.ptr, .vtable = &vt }, clock.clock());
+    defer core.deinit();
+    var host = LaneHost.init(gpa, std.testing.io, &core, .{ .lanes = 1 });
+    try host.start();
+    const engine = host.engine();
+    Scorer.calls = 0;
+    var logits: [3]f64 = undefined;
+    const lse = try engine.score(&.{ 1, 2, 3, 4 }, &.{ 9, 8, 7 }, &logits, true);
+    try std.testing.expectEqual(@as(f64, 7.5), lse);
+    try std.testing.expectEqualSlices(f64, &.{ 4, 5, 6 }, &logits);
+    try std.testing.expect(Scorer.engine_thread.? != std.Thread.getCurrentId());
+    try std.testing.expectError(error.Failed, engine.score(&.{1}, &.{}, logits[0..0], true));
+    try std.testing.expectEqual(@as(usize, 2), Scorer.calls);
+    host.stop();
+    try std.testing.expectError(error.Failed, engine.score(&.{1}, &.{2}, logits[0..1], false));
+    try std.testing.expectEqual(@as(usize, 2), Scorer.calls);
+}

zig/src/core/lane_host_tests.zig

@@ -5,6 +5,7 @@ test {
     _ = @import("lane_rounds_test.zig");
     _ = @import("lane_halt_test.zig");
     _ = @import("lane_host_end_test.zig");
+    _ = @import("lane_host_score_test.zig");
     _ = @import("lane_review_test.zig");
     _ = @import("lane_init_test.zig");
     _ = @import("lane_failure_review_test.zig");

zig/src/core/lane_rounds.zig

@@ -264,6 +264,7 @@ pub fn run(h: *LaneHost, rd: Rounds) void {
     while (true) {
         ends.takeCancels(h);
         rows_host.answer(h);
+        @import("lane_host_score.zig").answer(h);
         learn.refuse(h);
         h.lock();
         if (h.closing) {

zig/src/core/lanes/backend.zig

@@ -109,6 +109,8 @@ pub const Backend = struct {
         prefill_features: ?*const fn (ptr: *anyopaque, s: *Stream, taps: []const u32, sink: FeatureSink) anyerror!void = null,
         /// The first token's logprob row (the prompt's last row) for a stream with `logprobs`; null: not given.
         first_row: ?*const fn (ptr: *anyopaque, s: *Stream) anyerror!Row = null,
+        /// One logit per label at the row after `prompt` and, when `mass`, the vocabulary's logsumexp (decisions).
+        score: ?*const fn (ptr: *anyopaque, prompt: []const u32, labels: []const u32, logits: []f64, mass: bool) anyerror!f64 = null,
         /// `first` with the first token drawn among `mask`'s allowed tokens (structured output).
         first_masked: ?*const fn (ptr: *anyopaque, s: *Stream, position: u64, mask: []const u32) anyerror!u64 = null,
         /// The stream's held drafts as host tokens in window order, waiting for the draft (a grammar walks them).

zig/src/core/lanes/recorder.zig

@@ -238,7 +238,7 @@ pub const Recorded = struct {
 
     pub fn backend(x: *Recorded) be.Backend {
         const v = x.inner.vtable;
-        x.table = .{ .prefill = prefill, .prefill_step = if (v.prefill_step != null) prefillStep else null, .prefill_many = if (v.prefill_many != null) prefillMany else null, .first = first, .queue = queue, .read = read, .verify = verify, .keep = keep, .draft = draft, .unspeculate = if (v.unspeculate != null) unspeculate else null, .probabilities = if (v.probabilities != null) probabilities else null, .tree = if (v.tree != null) tree else null, .alternatives = if (v.alternatives != null) alternatives else null, .prepare_features = if (v.prepare_features != null) prepareFeatures else null, .features = features, .first_row = if (v.first_row != null) firstRow else null, .first_masked = if (v.first_masked != null) firstMasked else null, .held = if (v.held != null) held else null, .prompt_rows = v.prompt_rows, .release = release };
+        x.table = .{ .prefill = prefill, .prefill_step = if (v.prefill_step != null) prefillStep else null, .prefill_many = if (v.prefill_many != null) prefillMany else null, .first = first, .queue = queue, .read = read, .verify = verify, .keep = keep, .draft = draft, .unspeculate = if (v.unspeculate != null) unspeculate else null, .probabilities = if (v.probabilities != null) probabilities else null, .tree = if (v.tree != null) tree else null, .alternatives = if (v.alternatives != null) alternatives else null, .prepare_features = if (v.prepare_features != null) prepareFeatures else null, .features = features, .first_row = if (v.first_row != null) firstRow else null, .first_masked = if (v.first_masked != null) firstMasked else null, .score = if (v.score != null) score else null, .held = if (v.held != null) held else null, .prompt_rows = v.prompt_rows, .release = release };
         return .{ .ptr = x, .vtable = &x.table, .first_before_draft = x.inner.first_before_draft };
     }
 
@@ -315,6 +315,10 @@ pub const Recorded = struct {
         const i = of(p).inner;
         return i.vtable.first_masked.?(i.ptr, s, position, mask);
     }
+    fn score(p: *anyopaque, prompt: []const u32, labels: []const u32, logits: []f64, mass: bool) anyerror!f64 {
+        const i = of(p).inner;
+        return i.vtable.score.?(i.ptr, prompt, labels, logits, mass);
+    }
     fn held(p: *anyopaque, s: *Stream, out: []u32) anyerror!void {
         const i = of(p).inner;
         return i.vtable.held.?(i.ptr, s, out);

zig/src/families/flashnext/cuda_lanes.zig

@@ -211,6 +211,7 @@ pub fn Namespace(comptime Tri: type) type {
                     .release = releaseFn,
                     .tree = treeFn,
                     .probabilities = probabilitiesFn,
+                    .score = scoreFn,
                 } };
             }
 
@@ -406,6 +407,13 @@ pub fn Namespace(comptime Tri: type) type {
             return self.take(try self.tokenAt(&l.seq, s, position));
         }
 
+        /// Decisions' label scores on one rank (cuda_score.zig): a second rank would miss the pass's collectives.
+        fn scoreFn(ptr: *anyopaque, ids: []const u32, wanted: []const u32, logits: []f64, mass: bool) anyerror!f64 {
+            const self = of(ptr);
+            if (self.ranks != null) return error.UnsupportedRankOperation;
+            return @import("cuda_score.zig").labels(Tri, self.engine, self.tri(), self.context, ids, wanted, logits, mass);
+        }
+
         fn queueFn(ptr: *anyopaque, s: *lanes.Stream, feed: be.Feed, position: u64) anyerror!u64 {
             _ = ptr;
             _ = s;
@@ -529,6 +537,7 @@ pub fn Namespace(comptime Tri: type) type {
             _ = @import("cuda_images_test.zig").Cases(Cuda, Cuda.open, dropLane);
             _ = @import("cuda_lanes_restore_test.zig").Cases(Cuda, Cuda.open, snapRestore, dropLane);
             _ = @import("cuda_lanes_admission_test.zig").Cases(Cuda, Cuda.open, dropLane);
+            _ = @import("cuda_score.zig");
         }
 
         // -- the prompt cache's copies ------------------------------------------------------------------------

zig/src/families/flashnext/cuda_launch.zig

@@ -29,6 +29,8 @@ pub const mtp = @import("cuda_mtp.zig");
 pub const pocket = @import("cuda_pocket.zig");
 /// Living Weights' trainer over the batched pass
 pub const train = @import("cuda_train.zig");
+/// Decisions' label scores (the lanes backend's `score`)
+pub const score = @import("cuda_score.zig");
 pub const native = @import("cuda_lanes.zig").Namespace(Tri);
 /// Our own glue kernels, which every launch takes when no captured set is loaded
 pub const Glue = glue_.Glue;

zig/src/families/flashnext/cuda_score.zig

@@ -0,0 +1,43 @@
+//! Decisions' label scores on CUDA: the prompt through a fresh cache, then its last row's label logits and logsumexp.
+const std = @import("std");
+const prompt = @import("cuda_prompt.zig");
+const mtp = @import("cuda_mtp.zig");
+
+/// One logit per label at the row after `ids` (bf16 widened), and the row's logsumexp in f64 when `mass` (else 0).
+pub fn labels(comptime Tri: type, e: *prompt.Engine, tri: Tri, context: usize, ids: []const u32, wanted: []const u32, logits: []f64, mass: bool) !f64 {
+    // every refusal before a cache is made
+    if (ids.len == 0 or ids.len > context or ids.len > prompt.max_tokens) return error.PromptOutsideContext;
+    if (wanted.len == 0 or wanted.len != logits.len) return error.UnexpectedTensor;
+    for (wanted) |t| if (t >= prompt.vocabulary) return error.LabelOutsideVocabulary;
+    var seq = try e.newSeq(ids.len);
+    defer seq.deinit(e.gpa);
+    try e.forward(Tri, tri, &seq, ids, .{});
+    const row = std.mem.bytesAsSlice(u16, seq.logits[0 .. prompt.vocabulary * 2]);
+    for (wanted, logits) |t, *v| {
+        v.* = mtp.promote(std.mem.littleToNative(u16, row[t]));
+        if (!std.math.isFinite(v.*)) return error.NonFiniteScore;
+    }
+    return if (mass) logsumexp(row) else 0;
+}
+
+/// The bf16 row's logsumexp in f64; a non-finite row is refused.
+fn logsumexp(row: []align(1) const u16) !f64 {
+    var top: f64 = -std.math.inf(f64);
+    for (row) |b| {
+        const v: f64 = mtp.promote(std.mem.littleToNative(u16, b));
+        if (!std.math.isFinite(v)) return error.NonFiniteScore;
+        top = @max(top, v);
+    }
+    var sum: f64 = 0;
+    for (row) |b| sum += @exp(@as(f64, mtp.promote(std.mem.littleToNative(u16, b))) - top);
+    return top + @log(sum);
+}
+
+test "the logsumexp of a bf16 row in f64, and a non-finite row refused" {
+    const one: u16 = 0x3f80; // 1.0
+    const two: u16 = 0x4000; // 2.0
+    const row = [_]u16{ one, two, one };
+    const want = 2.0 + @log(1.0 + 2.0 * @exp(@as(f64, -1.0)));
+    try std.testing.expectApproxEqRel(want, try logsumexp(&row), 1e-15);
+    try std.testing.expectError(error.NonFiniteScore, logsumexp(&[_]u16{ one, 0x7f80 }));
+}

zig/src/native/cuda_bound.zig

@@ -109,6 +109,12 @@ pub fn vtable(comptime H: type, v: be.Backend.VTable) be.Backend.VTable {
                 return x.inner.vtable.first_masked.?(x.inner.ptr, s, position, mask);
             }
         }.f else null,
+        .score = if (v.score != null) struct {
+            fn f(p: *anyopaque, prompt: []const u32, labels: []const u32, logits: []f64, mass: bool) anyerror!f64 {
+                const x = H.bind(p);
+                return x.inner.vtable.score.?(x.inner.ptr, prompt, labels, logits, mass);
+            }
+        }.f else null,
         .held = if (v.held != null) struct {
             fn f(p: *anyopaque, s: *lanes.Stream, out: []u32) anyerror!void {
                 const x = H.bind(p);
@@ -253,6 +259,12 @@ test "each forwarded entry calls the family's entry of the same name, once" {
                 return 0;
             }
         }.f,
+        .score = struct {
+            fn f(p: *anyopaque, _: []const u32, _: []const u32, _: []f64, _: bool) anyerror!f64 {
+                Rec.hit(p, "score");
+                return 0;
+            }
+        }.f,
         .held = struct {
             fn f(p: *anyopaque, _: *S, _: []u32) anyerror!void {
                 Rec.hit(p, "held");
@@ -312,6 +324,9 @@ test "each forwarded entry calls the family's entry of the same name, once" {
     try rec.once("first_row");
     _ = try w.first_masked.?(p, &s, 0, &.{});
     try rec.once("first_masked");
+    var label_logits: [0]f64 = .{};
+    _ = try w.score.?(p, &.{}, &.{}, &label_logits, true);
+    try rec.once("score");
     try w.held.?(p, &s, &out);
     try rec.once("held");
     w.release(p, &s);