clef: Cloudflare Clef decisions on /v1/systemone through its joint schema head (--clef DIR)

THIRD_PARTY_NOTICES.md

@@ -193,6 +193,16 @@ The Zig tensor-parallel fabric's RoCE path (`zig/src/fabric/tp/mailbox.cu`, `roc
 the same one-shot all-gather protocol of b12x's RoCEnante (`b12x/comm/roce/`, PRs #295 and #315) for any family that
 runs tensor parallelism over it, under the same copyright and Apache License 2.0 terms.
 
+## Cloudflare Clef
+
+`zig/src/server/clef_encode.zig`, `clef_head.zig` and `clef_serve.zig` port Cloudflare's Clef release code,
+`joint_schema_model.py` in [Cloudflare/clef](https://huggingface.co/Cloudflare/clef) and
+[Cloudflare/clef-flash](https://huggingface.co/Cloudflare/clef-flash): the record rendering and encoding
+(`render`, `question_options`, `encode_record`), the verbatim system prompt, the joint schema head
+(`JointSchemaHead`, `EvidenceRoutingLayer`) and the SystemOne answer shape (`systemone_answer`). Copyright Cloudflare,
+Inc., Apache License 2.0; see [the license text](LICENSES/Apache-2.0.txt). `tools/clef_tiny_golden.py` runs that
+file to write the test fixture `zig/src/server/fixtures/clef_tiny.json` (random weights, no Cloudflare data).
+
 ## Vendored code and weights
 
 In the Python engine, frozen at 0.6.6, `src/tensorfold/drafters/vendor/z_lab_dflash/model_mlx.py` is the unmodified

tools/clef_tiny_golden.py

@@ -0,0 +1,53 @@
+"""Write zig/src/server/fixtures/clef_tiny.json: a tiny random Clef joint head and its logits from Cloudflare's own code.
+
+usage: python3 tools/clef_tiny_golden.py JOINT_SCHEMA_MODEL_PY OUT_JSON   (needs torch; JOINT_SCHEMA_MODEL_PY is Cloudflare's
+joint_schema_model.py from the Clef release). The weights are rounded to bf16 so the Zig head reads the same values.
+"""
+import importlib.util
+import json
+import sys
+
+import torch
+
+spec = importlib.util.spec_from_file_location("jsm", sys.argv[1])
+jsm = importlib.util.module_from_spec(spec)
+sys.modules["jsm"] = jsm
+spec.loader.exec_module(jsm)
+
+torch.manual_seed(1234)
+H, W, ROUTING, LAYERS, HEADS, FF, VOCAB, N = 16, 8, 1, 1, 2, 16, 20, 12
+head = jsm.JointSchemaHead(H, W, ROUTING, LAYERS, HEADS, FF).eval()
+with torch.no_grad():
+    for p in head.parameters():
+        p.copy_((torch.randn_like(p) * 0.5).bfloat16().float())
+    head.prior_logit_scale.fill_(0.5)
+    head.joint_logit_scale.fill_(1.0)
+    head.residual_gate.fill_(0.25)
+    hidden = torch.randn(1, N, H) * 2
+    rows = (torch.randn(VOCAB, H)).bfloat16().float()
+ids = torch.randint(0, VOCAB, (1, N))
+questions = [
+    jsm.EncodedQuestion("q0", 1, (0, 3), ((3, 5), (5, 7)), ("a", "b")),
+    jsm.EncodedQuestion("q1", 2, (7, 9), ((9, 10), (10, 11), (11, 12)), ("0", "1", "2")),
+]
+record = jsm.EncodedRecord(tuple(ids[0].tolist()), tuple(questions), "tiny")
+with torch.no_grad():
+    logits = head(hidden, ids, torch.ones(1, N, dtype=torch.long), [record], rows)[0]
+
+
+def bits(t):
+    return t.contiguous().view(-1).bfloat16().view(torch.int16).to(torch.int32).bitwise_and(0xFFFF).tolist()
+
+
+out = {
+    "config": {"hidden": H, "width": W, "routing": ROUTING, "layers": LAYERS, "heads": HEADS, "ff": FF},
+    "tensors": {name: bits(p.detach()) for name, p in head.state_dict().items()},
+    "lm_head": bits(rows),
+    "hidden": hidden.view(-1).tolist(),
+    "ids": ids[0].tolist(),
+    "questions": [{"kind": q.question_type, "span": list(q.question_span), "options": [list(s) for s in q.option_spans]} for q in questions],
+    "logits": [l.tolist() for l in logits],
+}
+with open(sys.argv[2], "w") as f:
+    json.dump(out, f, separators=(",", ":"))
+print("wrote", sys.argv[2], "logits", out["logits"])

zig/build/cuda.zig

@@ -454,6 +454,7 @@ pub fn hostTests(b: *std.Build, draft_ids: *std.Build.Module, build_options: *st
     media_gate.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = server_tests, .filters = &.{ "image", "video" } })).step);
     media_gate.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = mods.core, .filters = &.{ "qwen_vl", "video" } })).step);
     b.step("test-http-limits-host", "Total HTTP read deadlines and connection admission without shaders").dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = server_tests, .filters = &.{"HTTP limits:"} })).step);
+    b.step("test-clef-host", "Clef's encoder and joint head; goldens when TF_CLEF_DIR and TF_CLEF_GOLDEN are set").dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = server_tests, .filters = &.{"Clef:"} })).step);
     b.step("test-server-recovery-host", "Incomplete replies, token allocation and bounded SSE storage without shaders").dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = server_tests, .filters = &.{"Server recovery:"} })).step);
     b.step("test-prompt-interior-host", "Last-user metadata excludes later messages without shaders").dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = server_tests, .filters = &.{"interior prompt:"} })).step);
     const server_test = b.addRunArtifact(b.addTest(.{ .root_module = server_tests }));

zig/src/core/engine_api.zig

@@ -292,6 +292,8 @@ pub const Open = struct {
     learn_conversations: bool = false,
     /// --learn-daily-gib: what learned states may write a day under the root, every model together.
     learn_daily_gib: f64 = 32,
+    /// --clef: a decision's longest prompt, whose cache the engine reserves at startup inside its budget (0: none).
+    decision_rows: u32 = 0,
     /// --spill-gib (Macs): states evicted from RAM go to the store `learn_conversations` opens (off: written behind).
     spill: prompt_cache.rows.Spill = .off,
     /// --slide: Sliding Weights learns into the served weights, live (off: learn requests are refused).
@@ -400,6 +402,8 @@ pub const Engine = struct {
         keepalive: ?*const fn (ctx: *anyopaque) ?keepalive.Target = null,
         /// One logit per label at the row after `prompt`, and the vocabulary logsumexp when `mass` (else 0).
         score: ?*const fn (ctx: *anyopaque, prompt: []const u32, labels: []const u32, logits: []f64, mass: bool) error{Failed}!f64 = null,
+        /// Every row's final normed state (bf16 bits, rows x hidden) of `prompt`; null: not given.
+        hidden: ?*const fn (ctx: *anyopaque, prompt: []const u32, out: []u16) error{Failed}!void = null,
         /// Each branch's label logits at its last row after the decision's shared trunk.
         decide: ?*const fn (ctx: *anyopaque, d: *Decision) error{ Failed, Refused }!void = null,
         /// Queue `request` for the family's learner; it stays valid until `done`. Null: this engine does not learn.
@@ -464,6 +468,12 @@ pub const Engine = struct {
         if (e.vtable.drain) |f| f(e.ctx);
     }
 
+    /// The prompt's final normed rows, or Unsupported when this engine does not give them.
+    pub fn hidden(e: Engine, prompt: []const u32, out: []u16) error{ Failed, Unsupported }!void {
+        const f = e.vtable.hidden orelse return error.Unsupported;
+        return f(e.ctx, prompt, out);
+    }
+
     /// The decision hook, or Unsupported when this engine does not score.
     pub fn score(e: Engine, prompt: []const u32, labels: []const u32, logits: []f64, mass: bool) error{ Failed, Unsupported }!f64 {
         const f = e.vtable.score orelse return error.Unsupported;

zig/src/core/lane_host.zig

@@ -119,9 +119,22 @@ pub const LaneHost = struct {
 
     pub fn engine(h: *LaneHost) Engine {
         // 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 s = h.core.backend.vtable.score != null;
+        const rows = h.core.backend.vtable.hidden != null;
+        return .{ .ctx = h, .vtable = if (s and rows) &both_vtable else if (s) &scored_vtable else if (rows) &hidden_vtable else &plain_vtable };
     }
 
+    const hidden_vtable: Engine.VTable = blk: {
+        var v = plain_vtable;
+        v.hidden = scoring.hidden;
+        break :blk v;
+    };
+    const both_vtable: Engine.VTable = blk: {
+        var v = scored_vtable;
+        v.hidden = scoring.hidden;
+        break :blk v;
+    };
+
     const plain_vtable: Engine.VTable = .{
         .info = infoFn,
         .submit = submitFn,
@@ -561,7 +574,7 @@ pub const LaneHost = struct {
                     continue;
                 }
                 if (rows_host.idle(h)) continue;
-                if (h.cancels.items.len == 0 and (h.queued.items.len == 0 or h.admitted.items.len >= h.info_.lanes))
+                if (h.cancels.items.len == 0 and !scoring.busy(h) and (h.queued.items.len == 0 or h.admitted.items.len >= h.info_.lanes))
                     h.wake.waitTimeout(h.io, &h.mutex, .{ .duration = .{ .raw = .fromMilliseconds(100), .clock = .awake } }) catch {};
                 h.unlock();
                 continue;

zig/src/core/lane_host_score.zig

@@ -8,6 +8,8 @@ pub const Ask = struct {
     labels: []const u32,
     logits: []f64,
     mass: bool,
+    hidden: ?[]u16 = null, // a hidden-rows ask: the final normed rows go here instead of label logits
+    at: usize = 0, // a hidden-rows ask's next prompt row: it runs one backend piece per engine-thread pass
     logsumexp: f64 = 0,
     failed: bool = false,
     done: bool = false,
@@ -16,8 +18,14 @@ pub const Ask = struct {
 /// Questions waiting for the engine thread (under the host's lock).
 pub const Queue = struct {
     asks: std.ArrayList(*Ask) = .empty,
+    stepping: std.ArrayList(*Ask) = .empty, // hidden-rows asks part way through their prompt, oldest first
 };
 
+/// A decision's prompt is part way through (under the host's lock): the engine thread should not idle.
+pub fn busy(h: *LaneHost) bool {
+    return h.scores.stepping.items.len > 0;
+}
+
 /// 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));
@@ -32,21 +40,40 @@ pub fn score(ctx: *anyopaque, prompt: []const u32, labels: []const u32, logits:
     return ask.logsumexp;
 }
 
-/// Every queued question answered by the backend (engine thread, between rounds).
+/// Engine.hidden: the prompt's final normed rows, queued for the engine thread as scores are.
+pub fn hidden(ctx: *anyopaque, prompt: []const u32, out: []u16) error{Failed}!void {
+    var logits: [0]f64 = .{};
+    var ask: Ask = .{ .prompt = prompt, .labels = &.{}, .logits = &logits, .mass = false, .hidden = out };
+    const h: *LaneHost = @ptrCast(@alignCast(ctx));
+    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;
+}
+
+/// Queued scores answered, and one prompt piece of the oldest decision (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;
-    };
+    const list = h.scores.asks.toOwnedSlice(h.gpa) catch &.{};
     h.unlock();
     defer h.gpa.free(list);
     const backend = h.core.backend;
     for (list) |a| {
+        if (a.hidden != null) {
+            if (backend.vtable.hidden == null) {
+                a.failed = true;
+                continue;
+            }
+            h.lock();
+            defer h.unlock();
+            h.scores.stepping.append(h.gpa, a) catch {
+                a.failed = true;
+            };
+            continue;
+        }
         const run = backend.vtable.score orelse {
             a.failed = true;
             continue;
@@ -56,8 +83,35 @@ pub fn answer(h: *LaneHost) void {
             a.failed = true;
         }
     }
+    step(h);
     h.lock();
-    for (list) |a| a.done = true;
+    for (list) |a| if (a.hidden == null or a.failed) {
+        a.done = true;
+    };
+    h.wake.broadcast(h.io);
+    h.unlock();
+}
+
+/// One backend piece of the oldest decision's prompt; done (and dropped) at its last row or on a failure.
+fn step(h: *LaneHost) void {
+    h.lock();
+    const a = if (h.scores.stepping.items.len > 0) h.scores.stepping.items[0] else null;
+    h.unlock();
+    const ask = a orelse return;
+    const backend = h.core.backend;
+    var finished = false;
+    if (backend.vtable.hidden.?(backend.ptr, ask.prompt, ask.hidden.?, ask.at)) |next| {
+        ask.at = next;
+        finished = next >= ask.prompt.len;
+    } else |err| {
+        std.log.warn("lane host: a decision's hidden rows failed ({s})", .{@errorName(err)});
+        ask.failed = true;
+        finished = true;
+    }
+    if (!finished) return;
+    h.lock();
+    _ = h.scores.stepping.orderedRemove(0);
+    ask.done = true;
     h.wake.broadcast(h.io);
     h.unlock();
 }
@@ -71,6 +125,11 @@ pub fn close(h: *LaneHost) void {
         a.done = true;
     }
     h.scores.asks.clearAndFree(h.gpa);
+    for (h.scores.stepping.items) |a| {
+        a.failed = true;
+        a.done = true;
+    }
+    h.scores.stepping.clearAndFree(h.gpa);
     h.wake.broadcast(h.io);
     h.unlock();
 }

zig/src/core/lane_host_score_test.zig

@@ -59,3 +59,40 @@ test "a backend's score runs on the engine thread, a failure comes back as Faile
     try std.testing.expectError(error.Failed, engine.score(&.{1}, &.{2}, logits[0..1], false));
     try std.testing.expectEqual(@as(usize, 2), Scorer.calls);
 }
+
+/// A stepped hidden-rows backend: 3 rows a piece, row i's value prompt[i] * 2 + 1; it records each piece's start.
+const Rows = struct {
+    var starts: [8]usize = undefined;
+    var n: usize = 0;
+
+    fn hidden(_: *anyopaque, prompt: []const u32, out: []u16, at: usize) anyerror!usize {
+        starts[n] = at;
+        n += 1;
+        const end = @min(prompt.len, at + 3);
+        for (at..end) |i| out[i] = @intCast(prompt[i] * 2 + 1);
+        return end;
+    }
+};
+
+test "a decision's hidden rows come one backend piece per engine-thread pass, and equal a whole pass" {
+    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.hidden = Rows.hidden;
+    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();
+    defer host.stop();
+    Rows.n = 0;
+    const prompt = [_]u32{ 5, 1, 4, 1, 5, 9, 2, 6 };
+    var out: [8]u16 = undefined;
+    try host.engine().hidden(&prompt, &out);
+    try std.testing.expectEqualSlices(usize, &.{ 0, 3, 6 }, Rows.starts[0..Rows.n]);
+    for (prompt, out) |p, o| try std.testing.expectEqual(@as(u16, @intCast(p * 2 + 1)), o);
+}

zig/src/core/lane_rounds.zig

@@ -286,7 +286,7 @@ pub fn run(h: *LaneHost, rd: Rounds) void {
         }
         if (!work) {
             if (rows_host.idle(h)) continue;
-            if (h.cancels.items.len == 0) h.wake.waitTimeout(h.io, &h.mutex, .{ .duration = .{ .raw = .fromMilliseconds(100), .clock = .awake } }) catch {};
+            if (h.cancels.items.len == 0 and !@import("lane_host_score.zig").busy(h)) h.wake.waitTimeout(h.io, &h.mutex, .{ .duration = .{ .raw = .fromMilliseconds(100), .clock = .awake } }) catch {};
             h.unlock();
             continue;
         }

zig/src/core/lanes/backend.zig

@@ -124,6 +124,8 @@ pub const Backend = struct {
         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,
+        /// A decision prompt piece from row `at` (0: fresh): final normed bf16 rows into `out`; the next piece's row.
+        hidden: ?*const fn (ptr: *anyopaque, prompt: []const u32, out: []u16, at: usize) anyerror!usize = 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/families/qwen3_5/cuda_engine.zig

@@ -538,6 +538,12 @@ pub const Engine = struct {
         recurrent.completed(e);
     }
 
+    /// The last prompt forward's final-norm rows (bf16, `rows` x hidden) into `out`.
+    pub fn finalRows(e: *Engine, rows: usize, out: []u16) !void {
+        if (out.len != rows * e.g.hidden or rows > max_rows) return error.InvalidQwenSegments;
+        try e.d.check(e.d.api.cuMemcpyDtoH_v2(out.ptr, e.x, out.len * 2), "cuMemcpyDtoH");
+    }
+
     pub const Shape = @import("cuda_shapes.zig").Shape;
 
     /// The same scan selected by body, exposed for host launch-contract fixtures.