diff --git a/docs/api.md b/docs/api.md index 11793d2ef2..ab07c4de6e 100644 --- a/docs/api.md +++ b/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 diff --git a/zig/src/core/lane_host.zig b/zig/src/core/lane_host.zig index 7f24ff457e..9283aa3920 100644 --- a/zig/src/core/lane_host.zig +++ b/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(); diff --git a/zig/src/core/lane_host_score.zig b/zig/src/core/lane_host_score.zig new file mode 100644 index 0000000000..d287e8d618 --- /dev/null +++ b/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(); +} diff --git a/zig/src/core/lane_host_score_test.zig b/zig/src/core/lane_host_score_test.zig new file mode 100644 index 0000000000..36ab2ee118 --- /dev/null +++ b/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); +} diff --git a/zig/src/core/lane_host_tests.zig b/zig/src/core/lane_host_tests.zig index d121279474..8c3272ca72 100644 --- a/zig/src/core/lane_host_tests.zig +++ b/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"); diff --git a/zig/src/core/lane_rounds.zig b/zig/src/core/lane_rounds.zig index 74a476d540..a49203cd68 100644 --- a/zig/src/core/lane_rounds.zig +++ b/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) { diff --git a/zig/src/core/lanes/backend.zig b/zig/src/core/lanes/backend.zig index 03c5e46d12..bdaff62dd2 100644 --- a/zig/src/core/lanes/backend.zig +++ b/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). diff --git a/zig/src/core/lanes/recorder.zig b/zig/src/core/lanes/recorder.zig index cd95b88032..d4e23a2dbd 100644 --- a/zig/src/core/lanes/recorder.zig +++ b/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); diff --git a/zig/src/families/flashnext/cuda_lanes.zig b/zig/src/families/flashnext/cuda_lanes.zig index 226c988b55..33aba813fd 100644 --- a/zig/src/families/flashnext/cuda_lanes.zig +++ b/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 ------------------------------------------------------------------------ diff --git a/zig/src/families/flashnext/cuda_launch.zig b/zig/src/families/flashnext/cuda_launch.zig index 99e25c0c12..1d86b6c5b2 100644 --- a/zig/src/families/flashnext/cuda_launch.zig +++ b/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; diff --git a/zig/src/families/flashnext/cuda_score.zig b/zig/src/families/flashnext/cuda_score.zig new file mode 100644 index 0000000000..b654cf4fbd --- /dev/null +++ b/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 })); +} diff --git a/zig/src/native/cuda_bound.zig b/zig/src/native/cuda_bound.zig index 8992a81b82..0f26153496 100644 --- a/zig/src/native/cuda_bound.zig +++ b/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); diff --git a/zig/tests/cuda/flashnext_score_tests.zig b/zig/tests/cuda/flashnext_score_tests.zig new file mode 100644 index 0000000000..1e12c12e77 --- /dev/null +++ b/zig/tests/cuda/flashnext_score_tests.zig @@ -0,0 +1,51 @@ +//! Decisions' label scores against the prompt's own last row from a plain pass, byte for byte. +const std = @import("std"); +const cuda = @import("cuda"); +const flash = @import("flashnext"); +const check = @import("check.zig"); +const prompt = flash.prompt; + +pub fn run(gpu: check.Gpu, dir: []const u8, ids_path: []const u8) !void { + const source = try std.Io.Dir.cwd().readFileAlloc(gpu.io, ids_path, gpu.gpa, .limited(1 << 20)); + defer gpu.gpa.free(source); + var base: std.ArrayList(u32) = .empty; + defer base.deinit(gpu.gpa); + var words = std.mem.tokenizeAny(u8, source, " \r\n\t"); + while (words.next()) |word| try base.append(gpu.gpa, try std.fmt.parseInt(u32, word, 10)); + if (base.items.len < 16) return error.BadPrompt; + var stream = try cuda.Stream.init(gpu.d, true); + defer stream.deinit(); + var glue = try flash.Glue.load(gpu.d); + defer glue.deinit(); + const tri: flash.Tri = .{ .s = stream, .own = &glue }; + const e = try prompt.Engine.init(gpu.gpa, gpu.io, @constCast(gpu.d), &stream, dir); + defer e.deinit(); + // a short prompt, one past a pass chunk's attention block, and one past 2,048 keys (sparse selection) + for ([_]usize{ 16, 300, 2600 }) |length| { + const n = @min(length, base.items.len); + const ids = base.items[0..n]; + var plain = try e.newSeq(n); + defer plain.deinit(gpu.gpa); + try e.forward(flash.Tri, tri, &plain, ids, .{}); + const best = prompt.argmax(plain.logits); + const wanted = [_]u32{ 13, best, 220, 1, @intCast(prompt.vocabulary - 1) }; + var got: [wanted.len]f64 = undefined; + const lse = try flash.score.labels(flash.Tri, e, tri, n, ids, &wanted, &got, true); + const row = std.mem.bytesAsSlice(u16, plain.logits); + var top: f64 = -std.math.inf(f64); + for (row) |b| top = @max(top, @as(f64, flash.mtp.promote(b))); + var sum: f64 = 0; + for (row) |b| sum += @exp(@as(f64, flash.mtp.promote(b)) - top); + for (wanted, got) |t, v| { + const want: f64 = flash.mtp.promote(row[t]); + try check.expect(@as(u64, @bitCast(v)) == @as(u64, @bitCast(want)), "label {d}'s logit is the plain pass's row value ({d} vs {d})", .{ t, v, want }); + } + try check.expect(lse == top + @log(sum), "the logsumexp is the plain row's ({d} vs {d})", .{ lse, top + @log(sum) }); + var best_label: usize = 0; + for (got, 0..) |v, i| if (v > got[best_label]) { + best_label = i; + }; + try check.expect(wanted[best_label] == best, "the best label is the greedy token", .{}); + check.pass("decision scores at {d} prompt tokens: {d} label logits and the logsumexp equal the plain pass's last row", .{ n, wanted.len }); + } +} diff --git a/zig/tests/cuda/main.zig b/zig/tests/cuda/main.zig index 837f1bb070..70d8a98ffe 100644 --- a/zig/tests/cuda/main.zig +++ b/zig/tests/cuda/main.zig @@ -70,6 +70,7 @@ const route_rows_tests = @import("route_rows_tests.zig"); const flashnext_model_rank_tests = @import("flashnext_model_rank_tests.zig"); const flashnext_fidelity_tests = @import("flashnext_fidelity_tests.zig"); const flashnext_window_tests = @import("flashnext_window_tests.zig"); +const flashnext_score_tests = @import("flashnext_score_tests.zig"); const flashnext_rank_tests = @import("flashnext_rank_tests.zig"); const flashnext_draft_affine_tests = @import("flashnext_draft_affine_tests.zig"); const flashnext_shared_tests = @import("flashnext_shared_tests.zig"); @@ -98,6 +99,7 @@ const usage = \\ flashnext-fidelity-single MODEL FIXTURE OUTPUT single-rank reference isolation \\ flashnext-fidelity-rank MODEL RANK PORT FIXTURE OUTPUT fixed-history full score rows \\ flashnext-window-rank MODEL RANK PORT IDS_FILE two-rank serial/window and rollback parity + \\ flashnext-score MODEL IDS_FILE decisions' label scores against a plain pass's last row \\ flashnext-model-rank MODEL RANK PORT IDS_FILE explicit full-model replay experiment (no drafts) \\ flashnext-rank RANK PORT [sum|eof|timeout|missing] two-process collective and failure probes \\ flashnext-draft-affine proposal-head affine projection against FP64 and row-count invariance @@ -243,6 +245,7 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { if (std.mem.eql(u8, cmd, "flashnext-fidelity-single")) return flashnext_fidelity_tests.runSingle(gpu, try arg(rest, 0), try arg(rest, 1), try arg(rest, 2), false); if (std.mem.eql(u8, cmd, "flashnext-fidelity-single-serial")) return flashnext_fidelity_tests.runSingle(gpu, try arg(rest, 0), try arg(rest, 1), try arg(rest, 2), true); if (std.mem.eql(u8, cmd, "flashnext-fidelity-rank")) return flashnext_fidelity_tests.run(gpu, try arg(rest, 0), try std.fmt.parseInt(usize, try arg(rest, 1), 10), try std.fmt.parseInt(u16, try arg(rest, 2), 10), try arg(rest, 3), try arg(rest, 4)); + if (std.mem.eql(u8, cmd, "flashnext-score")) return flashnext_score_tests.run(gpu, try arg(rest, 0), try arg(rest, 1)); if (std.mem.eql(u8, cmd, "flashnext-window-rank")) return flashnext_window_tests.run(gpu, try arg(rest, 0), try std.fmt.parseInt(usize, try arg(rest, 1), 10), try std.fmt.parseInt(u16, try arg(rest, 2), 10), try arg(rest, 3)); if (std.mem.eql(u8, cmd, "flashnext-model-rank")) return flashnext_model_rank_tests.run(gpu, try arg(rest, 0), try std.fmt.parseInt(usize, try arg(rest, 1), 10), try std.fmt.parseInt(u16, try arg(rest, 2), 10), try arg(rest, 3)); if (std.mem.eql(u8, cmd, "flashnext-rank")) return flashnext_rank_tests.run(gpu, try std.fmt.parseInt(usize, try arg(rest, 0), 10), try std.fmt.parseInt(u16, try arg(rest, 1), 10), if (rest.len > 2) rest[2] else "sum");