From 0125b709fa261c4dbfaa331cb1a5d0f197e91e8f Mon Sep 17 00:00:00 2001 From: ashhart <71890694+ashhart@users.noreply.github.com> Date: Thu, 5 Nov 2026 15:09:00 -0800 Subject: [PATCH 1/4] Report completed Qwen27 prompt pass widths in request stats Count the actual target slices in the shared Metal prompt path. Keep verify and drafter absorb work out of the list and preserve inference when observation allocation fails. Forward the request-owned list through the lane host and existing raw/chat runtime fields. Host regressions cover stepped and serial cuts, the lone scheduler, drafted verify exclusion, observation failure, mailbox ownership, and whole/streamed raw and chat replies. Refs #2089, #2120. --- zig/build/q27_profile_host.zig | 1 + zig/build/qwen27_host.zig | 5 + zig/src/core/lane_host.zig | 2 +- zig/src/core/lanes/stream.zig | 3 + zig/src/native/qwen27_prompt_rows.zig | 10 ++ zig/src/native/qwen27_slots.zig | 1 + zig/src/qwen27_prefill_stats_test.zig | 188 ++++++++++++++++++++++ zig/src/server/reply_gpu_seconds_test.zig | 40 ++++- 8 files changed, 247 insertions(+), 3 deletions(-) create mode 100644 zig/src/qwen27_prefill_stats_test.zig diff --git a/zig/build/q27_profile_host.zig b/zig/build/q27_profile_host.zig index 66cbd67533..bc35b851bc 100644 --- a/zig/build/q27_profile_host.zig +++ b/zig/build/q27_profile_host.zig @@ -11,4 +11,5 @@ pub fn steps(b: *std.Build, server_tests: *std.Build.Module) void { step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = profile })).step); step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = api, .filters = &.{"gpu_seconds:"} })).step); step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = server_tests, .filters = &.{"gpu_seconds:"} })).step); + b.step("test-qwen27-prefill-widths-http-host", "Completed prompt widths on raw and chat HTTP routes without devices").dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = server_tests, .filters = &.{"prefill widths:"} })).step); } diff --git a/zig/build/qwen27_host.zig b/zig/build/qwen27_host.zig index 294a5911bd..19ffb22850 100644 --- a/zig/build/qwen27_host.zig +++ b/zig/build/qwen27_host.zig @@ -68,6 +68,11 @@ pub fn tests(b: *std.Build, test_step: *std.Build.Step) void { b.step("test-qwen27-equality-host", "Shared lanes drafted, resumed and concurrent equality on the fake target").dependOn(&equal.step); const resume_lanes = b.createModule(.{ .root_source_file = b.path("zig/src/core/lanes/lanes.zig"), .target = b.graph.host, .link_libc = true }); const resume_api = b.createModule(.{ .root_source_file = b.path("zig/src/core/engine_api.zig"), .target = b.graph.host, .link_libc = true, .imports = &.{.{ .name = "lanes", .module = resume_lanes }} }); + const prompt_stats = b.createModule(.{ .root_source_file = b.path("zig/src/qwen27_prefill_stats_test.zig"), .target = b.graph.host, .optimize = .debug, .link_libc = true, .imports = &.{ .{ .name = "lanes", .module = resume_lanes }, .{ .name = "engine_api", .module = resume_api } } }); + const prompt_stats_run = b.addRunArtifact(b.addTest(.{ .root_module = prompt_stats, .filters = &.{"qwen27 prefill stats:"} })); + prompt_stats_run.has_side_effects = true; + test_step.dependOn(&prompt_stats_run.step); + b.step("test-qwen27-prefill-stats-host", "Completed prompt passes through the lane host and server mailbox without a GPU").dependOn(&prompt_stats_run.step); const resume_rules = b.createModule(.{ .root_source_file = b.path("zig/src/native/qwen27_resume.zig"), .target = b.graph.host, .optimize = .debug, .imports = &.{.{ .name = "engine_api", .module = resume_api }} }); b.step("test-qwen27-resume-host", "Slot 0's in-place resume and rewind point: the cache's rule, with no GPU").dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = resume_rules, .filters = &.{"qwen27 resume:"} })).step); } diff --git a/zig/src/core/lane_host.zig b/zig/src/core/lane_host.zig index cd1c773ace..c9582e2dbf 100644 --- a/zig/src/core/lane_host.zig +++ b/zig/src/core/lane_host.zig @@ -308,7 +308,7 @@ pub const LaneHost = struct { h.gpa.free(job.marks); h.gpa.free(job.key); const s = &job.stream; - const stats: Stats = if (job.started) .{ .rounds = s.rounds, .grammar = if (s.grammar_time.windows > 0) s.grammar_time else null, .drafted = s.drafted, .accepted = s.accepted, .min_rows = s.min_rows, .loop_period = s.loop_period, .gpu_seconds = s.gpu_seconds, .prefill_seconds = if (job.prefilled) |done| @as(f64, @floatFromInt(@as(i64, @intCast(@max(0, done - job.began))))) / 1e9 else null } else .{}; + const stats: Stats = if (job.started) .{ .rounds = s.rounds, .grammar = if (s.grammar_time.windows > 0) s.grammar_time else null, .drafted = s.drafted, .accepted = s.accepted, .min_rows = s.min_rows, .prefill_widths = s.prefill_widths.items, .loop_period = s.loop_period, .gpu_seconds = s.gpu_seconds, .prefill_seconds = if (job.prefilled) |done| @as(f64, @floatFromInt(@as(i64, @intCast(@max(0, done - job.began))))) / 1e9 else null } else .{}; var told = if (job.started) recalls.finished(h, job, reason, stats) else stats; told.queued_seconds = if (job.admitted_ns) |at| @as(f64, @floatFromInt(@max(0, at - job.submitted))) / 1e9 else 0; emit(job, .{ .finished = .{ .reason = reason, .stats = told, .message = message } }); diff --git a/zig/src/core/lanes/stream.zig b/zig/src/core/lanes/stream.zig index 831babc560..d950d5640d 100644 --- a/zig/src/core/lanes/stream.zig +++ b/zig/src/core/lanes/stream.zig @@ -181,6 +181,8 @@ pub const Stream = struct { rows: std.ArrayList(Row) = .empty, // with `logprobs`: one a token of emitted(), in order prompt_logprobs: ?u8 = null, prompt_rows: std.ArrayList(Row) = .empty, // with `prompt_logprobs`: one a prompt token after the first, in order + prefill_widths: std.ArrayList(u32) = .empty, // completed target prompt passes, reported by an instrumented backend + prefill_widths_failed: bool = false, // an allocation failure invalidates this observation without failing the reply reuse: Reuse = .{}, decode_spans: []const [2]u32 = &.{}, decide: ?u32 = null, @@ -296,6 +298,7 @@ pub const Stream = struct { s.context.deinit(gpa); s.rows.deinit(gpa); s.prompt_rows.deinit(gpa); + s.prefill_widths.deinit(gpa); s.force.deinit(gpa); s.long_loop.deinit(gpa); s.dropRoundState(gpa); diff --git a/zig/src/native/qwen27_prompt_rows.zig b/zig/src/native/qwen27_prompt_rows.zig index 9f885bb3a3..ff7cf679ad 100644 --- a/zig/src/native/qwen27_prompt_rows.zig +++ b/zig/src/native/qwen27_prompt_rows.zig @@ -1,4 +1,5 @@ //! Each prompt chunk chooses its width from the streams that can decode beside it. +const std = @import("std"); pub fn choose(b: anytype, current: anytype) usize { for (b.owners[0..b.runner.slots]) |owner| if (owner) |stream| { if (stream != current and !stream.finished and !stream.paused and !b.fills.contains(stream)) { @@ -7,3 +8,12 @@ pub fn choose(b: anytype, current: anytype) usize { }; return b.chunk; } + +/// Only completed target prompt passes count; an observation failure never fails inference or reports a partial list. +pub fn completed(gpa: std.mem.Allocator, stream: anytype, rows: usize) void { + if (stream.prefill_widths_failed) return; + stream.prefill_widths.append(gpa, @intCast(rows)) catch { + stream.prefill_widths.clearRetainingCapacity(); + stream.prefill_widths_failed = true; + }; +} diff --git a/zig/src/native/qwen27_slots.zig b/zig/src/native/qwen27_slots.zig index 1dedd55bcb..b341a41284 100644 --- a/zig/src/native/qwen27_slots.zig +++ b/zig/src/native/qwen27_slots.zig @@ -132,6 +132,7 @@ pub fn chunk(b: *Metal, s: *lanes.Stream, f: *Fill, _: u32) !bool { if (b.ringOf(slot)) |g| try ringChunk(b, g, part.len); } charge(b, s, before); + @import("qwen27_prompt_rows.zig").completed(b.gpa, s, part.len); f.at = end; s.fill_at.store(@intCast(end), .monotonic); // a prompt chunk waits: `at` is landed progress if (end == f.rewind) saveRewind(b, f.rewind); diff --git a/zig/src/qwen27_prefill_stats_test.zig b/zig/src/qwen27_prefill_stats_test.zig new file mode 100644 index 0000000000..4a47778740 --- /dev/null +++ b/zig/src/qwen27_prefill_stats_test.zig @@ -0,0 +1,188 @@ +//! Completed prompt passes through the real lane host and server mailbox, with a CPU target and no shaders. +const std = @import("std"); +const lanes = @import("lanes"); +const api = @import("engine_api"); +const Mailbox = @import("server/mailbox.zig").Mailbox; +const rows = @import("native/qwen27_prompt_rows.zig"); +const t = std.testing; +const gpa = t.allocator; + +/// The fake supplies actual completed slice boundaries; only the production collector observes them. +const Observe = struct { + allocator: std.mem.Allocator = gpa, + enabled: bool = true, + seed_capacity: bool = false, + last: usize = 0, + lone_calls: usize = 0, + backend: lanes.backend.Backend = undefined, + + fn of(f: *lanes.fake.Fake) *Observe { + return @ptrCast(@alignCast(f.prefill_hook_ctx.?)); + } + + fn whole(ptr: *anyopaque, s: *lanes.Stream) anyerror!void { + const f: *lanes.fake.Fake = @ptrCast(@alignCast(ptr)); + of(f).last = 0; + f.prefill_hook = cut; + defer f.prefill_hook = null; + try f.backend().prefill(s); + } + + fn cut(ctx: *anyopaque, s: *lanes.Stream, _: usize) void { + const o: *Observe = @ptrCast(@alignCast(ctx)); + const at = s.fill_at.load(.monotonic); + if (o.enabled) rows.completed(o.allocator, s, at - o.last); + o.last = at; + } + + fn step(ptr: *anyopaque, s: *lanes.Stream, max_rows: u32) anyerror!bool { + const f: *lanes.fake.Fake = @ptrCast(@alignCast(ptr)); + const before = s.fill_at.load(.monotonic); + const done = try f.stepped().vtable.prefill_step.?(ptr, s, max_rows); + if (of(f).enabled and of(f).seed_capacity and s.prefill_widths.capacity == 0) { + const storage = try gpa.alloc(u32, 1); + s.prefill_widths = .fromOwnedSlice(storage); + s.prefill_widths.clearRetainingCapacity(); + } + if (of(f).enabled) rows.completed(of(f).allocator, s, s.fill_at.load(.monotonic) - before); + return done; + } + + /// Exercise the host's lone scheduler path; this is a CPU driver, not a real DFlash model. + fn lone(ctx: *anyopaque, s: *lanes.Stream, hooks: api.LoneHooks) anyerror!bool { + const o: *Observe = @ptrCast(@alignCast(ctx)); + o.lone_calls += 1; + _ = try o.backend.opening(gpa, s); + hooks.committed(hooks.ctx); + while (!s.finished) { + const token = lanes.fake.next(s.context.items, s.sampling, s.context.items.len); + _ = try s.commit(gpa, &.{token}, &.{}); + hooks.committed(hooks.ctx); + } + o.backend.release(s); + return false; + } +}; + +const Mode = enum { stepped, serial, lone }; + +fn wait(box: *Mailbox) !void { + for (0..10_000) |_| { + box.mutex.lockUncancelable(t.io); + const done = box.finished; + box.mutex.unlock(t.io); + if (done) return; + std.Io.sleep(t.io, .fromMilliseconds(1), .awake) catch {}; + } + return error.Timeout; +} + +/// Two successive requests own distinct lists. The mailbox must copy each before the worker frees its stream. +fn serve(mode: Mode, n: usize, drafts: bool, observe: *Observe, boxes: []Mailbox) !void { + var costs: [8]lanes.config.Cost = undefined; + for (&costs, 1..) |*c, w| c.* = .{ .width = @intCast(w), .ms = 5.0 + 0.8 * @as(f64, @floatFromInt(w)) }; + var cfg = try lanes.Config.init(gpa, .{ .exact_width = 8, .gpu_tokens = true, .hidden_rows = true, .mtp = true, .speculate = true, .speculate_early = false, .drafts = 3, .window_costs = &costs, .mtp_step_ms = 0.1, .batch_rows = 2048, .max_streams = 2, .draft_streams = true }, 8, 7); + defer cfg.deinit(gpa); + cfg.timed = false; + var target: lanes.fake.Fake = .{ .gpa = gpa, .prefill_chunks = 8, .prefill_hook_ctx = observe }; + defer target.deinit(); + var vt = target.backend().vtable.*; + vt.prefill = Observe.whole; + if (mode == .stepped) vt.prefill_step = Observe.step; + const backend: lanes.backend.Backend = .{ .ptr = &target, .vtable = &vt, .chunk_invariant = true }; + observe.backend = backend; + var clock: lanes.fake.FixedClock = .{}; + var core = lanes.Engine.init(gpa, &cfg, backend, clock.clock()); + defer core.deinit(); + var host = api.LaneHost.init(gpa, t.io, &core, .{ .lanes = 2, .prefill_step = 2048 }); + if (mode == .lone) host.lone = .{ .ctx = observe, .run = Observe.lone, .resumes = false, .plain = true }; + const prompt = try gpa.alloc(u32, n); + defer gpa.free(prompt); + for (prompt, 0..) |*p, i| p.* = @intCast(i % lanes.fake.vocab); + const request: api.Request = .{ .prompt = prompt, .max_tokens = 24, .drafts = drafts }; + try host.start(); + defer host.stop(); + for (boxes, 1..) |*box, id| { + try host.engine().submit(id, &request, .{ .ctx = box, .event = Mailbox.onEvent }); + try wait(box); + try t.expectEqual(api.Reason.length, box.reason); + try t.expectEqual(@as(usize, 24), box.tokens.items.len); + var history: std.ArrayList(u32) = .empty; + defer history.deinit(gpa); + try history.appendSlice(gpa, prompt); + for (box.tokens.items) |token| { + try t.expectEqual(lanes.fake.next(history.items, null, history.items.len), token); + try history.append(gpa, token); + } + } +} + +fn newBox() Mailbox { + return .{ .io = t.io, .gpa = gpa }; +} + +fn sameReply(a: *const Mailbox, b: *const Mailbox) !void { + try t.expectEqualSlices(u32, a.tokens.items, b.tokens.items); + try t.expectEqual(a.stats.rounds, b.stats.rounds); + try t.expectEqual(a.stats.drafted, b.stats.drafted); + try t.expectEqual(a.stats.accepted, b.stats.accepted); + try t.expectEqual(a.stats.min_rows, b.stats.min_rows); +} + +test "qwen27 prefill stats: stepped 16k reports eight actual 2048 passes with request-owned lifetime" { + var boxes = [_]Mailbox{ newBox(), newBox() }; + defer for (&boxes) |*b| b.deinit(); + var observe: Observe = .{}; + try serve(.stepped, 16384, false, &observe, &boxes); + for (&boxes) |*b| try t.expectEqualSlices(u32, &.{ 2048, 2048, 2048, 2048, 2048, 2048, 2048, 2048 }, b.widths); + try sameReply(&boxes[0], &boxes[1]); +} + +test "qwen27 prefill stats: serial reports its nonuniform actual cuts rather than capacity" { + var boxes = [_]Mailbox{newBox()}; + defer boxes[0].deinit(); + var observe: Observe = .{}; + try serve(.serial, 16381, false, &observe, &boxes); + try t.expectEqualSlices(u32, &.{ 2047, 2048, 2047, 2048, 2048, 2047, 2048, 2048 }, boxes[0].widths); +} + +test "qwen27 prefill stats: drafted verify rounds do not append prompt widths or change tokens and counters" { + var on = [_]Mailbox{newBox()}; + defer on[0].deinit(); + var off = [_]Mailbox{newBox()}; + defer off[0].deinit(); + var observe: Observe = .{}; + try serve(.stepped, 16384, true, &observe, &on); + observe.enabled = false; + try serve(.stepped, 16384, true, &observe, &off); + try t.expect(on[0].stats.rounds > 0); + try t.expect(on[0].stats.drafted > 0); + try t.expectEqualSlices(u32, &.{ 2048, 2048, 2048, 2048, 2048, 2048, 2048, 2048 }, on[0].widths); + try t.expectEqual(@as(usize, 0), off[0].widths.len); + try sameReply(&on[0], &off[0]); +} + +test "qwen27 prefill stats: lone scheduler forwards serial pass widths" { + var boxes = [_]Mailbox{newBox()}; + defer boxes[0].deinit(); + var observe: Observe = .{}; + try serve(.lone, 16384, true, &observe, &boxes); + try t.expectEqual(@as(usize, 1), observe.lone_calls); + try t.expectEqualSlices(u32, &.{ 2048, 2048, 2048, 2048, 2048, 2048, 2048, 2048 }, boxes[0].widths); +} + +test "qwen27 prefill stats: observation allocation failure keeps inference and emits no partial list" { + var failed = std.testing.FailingAllocator.init(gpa, .{ .fail_index = 0, .resize_fail_index = 0 }); + var bad = [_]Mailbox{newBox()}; + defer bad[0].deinit(); + var off = [_]Mailbox{newBox()}; + defer off[0].deinit(); + var observe: Observe = .{ .allocator = failed.allocator(), .seed_capacity = true }; + // One owned slot admits the first observation; the forced growth failure must discard it on the second. + try serve(.stepped, 16384, false, &observe, &bad); + observe.enabled = false; + try serve(.stepped, 16384, false, &observe, &off); + try t.expect(failed.has_induced_failure); + try t.expectEqual(@as(usize, 0), bad[0].widths.len); + try sameReply(&bad[0], &off[0]); +} diff --git a/zig/src/server/reply_gpu_seconds_test.zig b/zig/src/server/reply_gpu_seconds_test.zig index d120b97a29..e7a95e085a 100644 --- a/zig/src/server/reply_gpu_seconds_test.zig +++ b/zig/src/server/reply_gpu_seconds_test.zig @@ -72,19 +72,30 @@ fn steady(a: Allocator, value: json.Value, drop: ?[]const u8) ![]u8 { /// Every message the handler sent for a completion whose engine ended with `stats`, steadied, `drop` left out. fn served(a: Allocator, stats: api.Stats, streaming: bool, drop: ?[]const u8) ![]const []u8 { + return servedRoute(a, stats, streaming, drop, false); +} + +fn servedRoute(a: Allocator, stats: api.Stats, streaming: bool, drop: ?[]const u8, chat: bool) ![]const []u8 { const gpa = std.testing.allocator; var f = try fixture.Fixture.init(gpa, std.testing.io, null); defer f.deinit(); + if (chat) try f.tmp.dir.writeFile(std.testing.io, .{ .sub_path = "tokenizer_config.json", .data = + \\{"eos_token":"<|im_end|>","pad_token":"<|endoftext|>","chat_template":"{{ messages[0]['content'] }}"} + }); var problem: []const u8 = ""; const text = try HfText.load(gpa, std.testing.io, f.path, null, a, &problem); defer text.deinit(); var script = Script{ .stats = stats }; const srv = try Server.init(gpa, std.testing.io, script.engine(), text.text(), .{ .served_name = "m", .model_ids = &.{"m"}, .enable_thinking = false, .use_drafts = false }, null); defer srv.deinit(); - const body = (try json.parse(a, if (streaming) "{\"model\":\"m\",\"prompt\":[0],\"max_tokens\":3,\"stream\":true}" else "{\"model\":\"m\",\"prompt\":[0],\"max_tokens\":3}")).ok; + const request = if (chat) + (if (streaming) "{\"model\":\"m\",\"messages\":[{\"role\":\"user\",\"content\":\"H\"}],\"max_tokens\":3,\"stream\":true}" else "{\"model\":\"m\",\"messages\":[{\"role\":\"user\",\"content\":\"H\"}],\"max_tokens\":3}") + else + (if (streaming) "{\"model\":\"m\",\"prompt\":[0],\"max_tokens\":3,\"stream\":true}" else "{\"model\":\"m\",\"prompt\":[0],\"max_tokens\":3}"); + const body = (try json.parse(a, request)).ok; var conn = Conn{ .fd = -1, .peer = "", .buf = &.{}, .gpa = a }; var response = Reply{ .a = a }; - openai.run(srv, a, response.out(), .{ .conn = &conn }, false, body); + openai.run(srv, a, response.out(), .{ .conn = &conn }, chat, body); const values = if (streaming) response.events.items else &.{response.whole orelse return error.MissingReply}; const out = try a.alloc([]u8, values.len); for (values, out) |v, *bytes| bytes.* = try steady(a, v, drop); @@ -134,3 +145,28 @@ test "grammar timing: measured totals reach whole and streamed replies and uncon try std.testing.expectEqual(@as(usize, 1), measured); } } + +test "prefill widths: raw and chat replies serialize completed passes, whole and streamed, once" { + var arena = std.heap.ArenaAllocator.init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + for ([_]bool{ false, true }) |chat| for ([_]bool{ false, true }) |streaming| { + for ([_]bool{ false, true }) |observed| { + const expected: []const u32 = if (observed) &.{ 2048, 17 } else &.{}; + const replies = try servedRoute(a, .{ .prefill_widths = expected }, streaming, null, chat); + var finishes: usize = 0; + for (replies) |bytes| { + const value = (try json.parse(a, bytes)).ok; + if (value.object.get("tensorfold")) |runtime| { + finishes += 1; + const got = runtime.object.get("prefill_widths").?.array; + try std.testing.expectEqual(expected.len, got.len); + for (expected, got) |width, v| { + try std.testing.expectEqual(width, try std.fmt.parseInt(u32, v.int, 10)); + } + } + } + try std.testing.expectEqual(@as(usize, 1), finishes); + } + }; +} From 801408562e15e0a7c311fc01fb2056dd28439b74 Mon Sep 17 00:00:00 2001 From: ashhart <71890694+ashhart@users.noreply.github.com> Date: Thu, 5 Nov 2026 15:36:47 -0800 Subject: [PATCH 2/4] Report completed Nemotron and GLM Metal prompt pass widths Share the request-owned observation policy with Qwen27. Count GLM chunks after completion and Nemotron prompt buffers at their existing landing wait, preserving its one-chunk lookahead. Clear abandoned observation owners on a failed release drain. Host cases exercise chunked and one-pass prompts, actual cut planning, completion and ownership, verify exclusion, and request stats. Add shader-free compiles of the real family callbacks. Refs #2089, #2120. --- build.zig | 2 + zig/build/glm_expiry.zig | 4 +- zig/build/prefill_widths_metal.zig | 18 +++ zig/src/core/lanes/lanes.zig | 1 + zig/src/core/lanes/prefill_widths.zig | 11 ++ zig/src/families/glm/prefill.zig | 1 + zig/src/families/glm/prefill_harness.zig | 5 + zig/src/families/glm/prefill_test.zig | 33 +++++ zig/src/families/nemotron/backend.zig | 29 ++-- zig/src/families/nemotron/prompt_flight.zig | 45 ++++++ .../families/nemotron/prompt_flight_test.zig | 135 ++++++++++++++++++ zig/src/metal_prefill_widths_compile.zig | 12 ++ zig/src/native/qwen27_prompt_rows.zig | 10 -- zig/src/native/qwen27_slots.zig | 2 +- zig/src/qwen27_prefill_stats_test.zig | 2 +- 15 files changed, 276 insertions(+), 34 deletions(-) create mode 100644 zig/build/prefill_widths_metal.zig create mode 100644 zig/src/core/lanes/prefill_widths.zig create mode 100644 zig/src/families/nemotron/prompt_flight.zig create mode 100644 zig/src/families/nemotron/prompt_flight_test.zig create mode 100644 zig/src/metal_prefill_widths_compile.zig diff --git a/build.zig b/build.zig index b9c7d32198..0b73f30235 100644 --- a/build.zig +++ b/build.zig @@ -179,8 +179,10 @@ pub fn build(b: *std.Build) void { gemma4_host_mods.sources.addImport("nemotron_packed_metallib", embedded(b, b.addWriteFiles().add("gemma4-compile-empty.metallib", ""), "nemotron_packed.metallib")); @import("zig/build/gemma4_host.zig").compileHost(b, gemma4_host_mods.metal, gemma4_host_mods.engine, gemma4_host_mods.lanes, gemma4_host_mods.sources); @import("zig/build/qwen27_host.zig").compileHost(b, gemma4_host_mods.metal, gemma4_host_mods.engine, gemma4_host_mods.lanes); + @import("zig/build/prefill_widths_metal.zig").compileHost(b, gemma4_host_mods.engine, gemma4_host_mods.lanes); } @import("zig/build/qwen27_host.zig").tests(b, test_step); + @import("zig/build/prefill_widths_metal.zig").tests(b, test_step); const xgrammar = grammar_build.option(b); dist_build.targets(b, draft_ids, build_options, release_version, xgrammar); @import("zig/build/release_paths_step.zig").tests(b, test_step); diff --git a/zig/build/glm_expiry.zig b/zig/build/glm_expiry.zig index 1403d9e012..26f024b057 100644 --- a/zig/build/glm_expiry.zig +++ b/zig/build/glm_expiry.zig @@ -52,7 +52,9 @@ pub fn steps(b: *std.Build, test_step: *std.Build.Step, lanes: *std.Build.Module .imports = &.{ .{ .name = "metal", .module = metal }, .{ .name = "lanes", .module = lanes }, .{ .name = "fabric", .module = fabric }, .{ .name = "kernel_sources", .module = sources } }, }) }); if (b.graph.host.result.os.tag == .macos) schedule.root_module.linkSystemLibrary("objc", .{}); - b.step("test-glm-scheduling-host", "Actual GLM prefill chunks and lane-host decode fairness using a fake engine").dependOn(&b.addRunArtifact(schedule).step); + const schedule_run = b.addRunArtifact(schedule); + b.step("test-glm-scheduling-host", "Actual GLM prefill chunks and lane-host decode fairness using a fake engine").dependOn(&schedule_run.step); + b.step("test-glm-prefill-widths-host", "Actual GLM pass widths, restore cuts and shared lane stats without devices").dependOn(&schedule_run.step); if (b.graph.host.result.os.tag == .macos) { const compile = b.addObject(.{ .name = "glm-restore-host", .root_module = b.createModule(.{ .root_source_file = b.path("zig/src/glm_restore_compile.zig"), diff --git a/zig/build/prefill_widths_metal.zig b/zig/build/prefill_widths_metal.zig new file mode 100644 index 0000000000..5285161019 --- /dev/null +++ b/zig/build/prefill_widths_metal.zig @@ -0,0 +1,18 @@ +//! Metal prompt observations tested at their real completion boundaries; no shaders or model execution. +const std = @import("std"); + +pub fn tests(b: *std.Build, parent: *std.Build.Step) void { + const lanes = b.createModule(.{ .root_source_file = b.path("zig/src/core/lanes/lanes.zig"), .target = b.graph.host, .link_libc = true }); + const root = b.createModule(.{ .root_source_file = b.path("zig/src/families/nemotron/prompt_flight_test.zig"), .target = b.graph.host, .optimize = .debug, .link_libc = true, .imports = &.{.{ .name = "lanes", .module = lanes }} }); + const run = b.addRunArtifact(b.addTest(.{ .root_module = root, .filters = &.{"Nemotron prefill widths:"} })); + run.has_side_effects = true; + b.step("test-nemotron-prefill-widths-host", "Actual Nemotron cut and flight completion with recording buffers").dependOn(&run.step); + parent.dependOn(&run.step); +} + +pub fn compileHost(b: *std.Build, engine: *std.Build.Module, lanes: *std.Build.Module) void { + const root = b.createModule(.{ .root_source_file = b.path("zig/src/metal_prefill_widths_compile.zig"), .target = b.graph.host, .optimize = .debug, .link_libc = true, .imports = &.{ .{ .name = "tensorfold", .module = engine }, .{ .name = "lanes", .module = lanes } } }); + const object = b.addObject(.{ .name = "metal-prefill-widths-compile", .root_module = root }); + _ = object.getEmittedBin(); + b.step("check-nemotron-prefill-widths-host", "Compile actual Metal prefill, read, land and release without shaders").dependOn(&object.step); +} diff --git a/zig/src/core/lanes/lanes.zig b/zig/src/core/lanes/lanes.zig index f7151b96b0..1aa939fe95 100644 --- a/zig/src/core/lanes/lanes.zig +++ b/zig/src/core/lanes/lanes.zig @@ -11,6 +11,7 @@ pub const recall = @import("recall.zig"); pub const recall_store = @import("recall_store.zig"); pub const sampling = @import("sampling.zig"); pub const fill = @import("fill.zig"); +pub const prefill_widths = @import("prefill_widths.zig"); pub const shape = @import("shape.zig"); pub const plan_lanes = @import("plan_lanes.zig"); pub const cost_cache = @import("cost_cache.zig"); diff --git a/zig/src/core/lanes/prefill_widths.zig b/zig/src/core/lanes/prefill_widths.zig new file mode 100644 index 0000000000..2952c0cb93 --- /dev/null +++ b/zig/src/core/lanes/prefill_widths.zig @@ -0,0 +1,11 @@ +//! Optional request-owned observations of completed target prompt passes, independent of backend arithmetic. +const std = @import("std"); + +/// Observation failure never fails inference or reports a partial list. +pub fn completed(gpa: std.mem.Allocator, stream: anytype, rows: usize) void { + if (stream.prefill_widths_failed) return; + stream.prefill_widths.append(gpa, @intCast(rows)) catch { + stream.prefill_widths.clearRetainingCapacity(); + stream.prefill_widths_failed = true; + }; +} diff --git a/zig/src/families/glm/prefill.zig b/zig/src/families/glm/prefill.zig index 30a8237856..9962814fcc 100644 --- a/zig/src/families/glm/prefill.zig +++ b/zig/src/families/glm/prefill.zig @@ -23,6 +23,7 @@ pub fn step(b: anytype, s: *Stream, _: u32) !bool { const n = b.promptRows(p.at, end); if (n == 0 or n > end - p.at) return error.ChunkOutOfStep; try b.chunkPrompt(s, p.slot, p.at, n); + @import("lanes").prefill_widths.completed(b.gpa, s, n); p.at += n; s.fill_at.store(p.at, .monotonic); if (std.mem.indexOfScalar(u32, s.reuse.marks, p.at) != null) if (s.reuse.hook) |h| h.at(h.ptr, s, p.at); diff --git a/zig/src/families/glm/prefill_harness.zig b/zig/src/families/glm/prefill_harness.zig index 421180e9af..e39d9844c7 100644 --- a/zig/src/families/glm/prefill_harness.zig +++ b/zig/src/families/glm/prefill_harness.zig @@ -121,6 +121,7 @@ pub const Fake = struct { pub const Box = struct { tokens: [2]std.ArrayList(u32) = @splat(.empty), + widths: [2]std.ArrayList(u32) = @splat(.empty), done: std.atomic.Value(u32) = .init(0), bad: bool = false, pub fn event(ptr: *anyopaque, id: api.Id, e: *const api.Event) void { @@ -131,6 +132,9 @@ pub const Box = struct { }, .finished => |end| { if (end.reason == .failed) b.bad = true; + b.widths[id - 1].appendSlice(a, end.stats.prefill_widths) catch { + b.bad = true; + }; _ = b.done.fetchAdd(1, .release); }, else => {}, @@ -138,5 +142,6 @@ pub const Box = struct { } pub fn deinit(b: *Box) void { for (&b.tokens) |*ts| ts.deinit(a); + for (&b.widths) |*ws| ws.deinit(a); } }; diff --git a/zig/src/families/glm/prefill_test.zig b/zig/src/families/glm/prefill_test.zig index 5ed1b3cb86..de2653d244 100644 --- a/zig/src/families/glm/prefill_test.zig +++ b/zig/src/families/glm/prefill_test.zig @@ -102,6 +102,10 @@ fn run(drafts: bool, concurrent: bool) !Box { try std.testing.expectEqual(@as(usize, 32), box.tokens[1].items.len); try std.testing.expectEqual(@as(usize, 2), fake.begins); try std.testing.expectEqual(@as(usize, 2), fake.ends); + try std.testing.expectEqualSlices(u32, &.{4}, box.widths[0].items); + var prompt_total: usize = 0; + for (box.widths[1].items) |width| prompt_total += width; + try std.testing.expectEqual(@as(usize, 102400), prompt_total); if (concurrent) { try std.testing.expect(fake.chunks >= 102400 / fill.decode_rows); try std.testing.expect(fake.max_rows <= fill.decode_rows); @@ -515,3 +519,32 @@ test "glm prefill: a family that opts into chunk invariance gets the cap, and GL try std.testing.expect(try fill.step(&f, &s, hint)); // ...so the prompt ends on the cold plan's cuts try std.testing.expectEqualSlices(u32, &.{ 0, 2048, 4096 }, f.starts.items); // every start on 2048: no BadChunk } + +test "glm prefill: widths report actual planned cuts and a single pass" { + for ([_]bool{ false, true }) |chunked| { + var f: Fake = .{}; + defer f.deinit(); + const tokens: [23]u32 = @splat(7); + var s = try Stream.init(a, .{ .id = "widths", .prompt = &tokens, .max_new = 4, .chunks = if (chunked) &.{ 5, 18 } else &.{} }); + defer s.deinit(a); + while (!try fill.step(&f, &s, 0)) {} + const expected: []const u32 = if (chunked) &.{ 5, 13, 5 } else &.{23}; + try std.testing.expectEqualSlices(u32, expected, s.prefill_widths.items); + try std.testing.expectEqualSlices(u32, &tokens, f.target.lanes.get(&s).?.history.items); + f.abortPrompt(&s); + } +} + +test "glm prefill: widths exclude a restored prefix and failed target chunks" { + var f: Fake = .{ .restored_at = 3 }; + defer f.deinit(); + const tokens: [23]u32 = @splat(7); + var s = try Stream.init(a, .{ .id = "resumed-widths", .prompt = &tokens, .max_new = 4, .chunks = &.{ 5, 18 } }); + defer s.deinit(a); + try std.testing.expect(!try fill.step(&f, &s, 0)); + try std.testing.expectEqualSlices(u32, &.{2}, s.prefill_widths.items); + f.fail_chunk = true; + try std.testing.expectError(error.GpuFailed, fill.step(&f, &s, 0)); + try std.testing.expectEqualSlices(u32, &.{2}, s.prefill_widths.items); + try std.testing.expectEqual(@as(usize, 0), f.filling.count()); +} diff --git a/zig/src/families/nemotron/backend.zig b/zig/src/families/nemotron/backend.zig index fafc3698b4..f026411d51 100644 --- a/zig/src/families/nemotron/backend.zig +++ b/zig/src/families/nemotron/backend.zig @@ -37,7 +37,7 @@ pub const fused_rows = 16; pub const Kind = timing.Kind; -const Flight = struct { end: u64, cb: mtl.CommandBuffer, kind: Kind }; +const Flight = struct { end: u64, cb: mtl.CommandBuffer, kind: Kind, prompt: ?@import("prompt_flight.zig").Prompt = null }; /// A window's host rows in the token ring: the pending token at `stage`, then `n` tokens after its `held` drafts. const Fill = struct { stage: u64, held: usize, n: usize }; @@ -185,18 +185,7 @@ pub const Metal = struct { /// Wait for the oldest command buffer in flight and record its GPU time. pub fn land(self: *Metal) !void { - const f = self.flights.orderedRemove(0); - defer f.cb.release(); - f.cb.wait(); - if (f.cb.failure()) |text| { - std.log.err("command buffer failed: {s}", .{text}); - return error.GpuFailed; - } - self.landed = @max(self.landed, f.end); - self.last_ms = f.cb.gpuSeconds() * 1e3; - self.gpu_ms += self.last_ms; - self.timing.ms[@backingInt(f.kind)] += self.last_ms; - self.timing.count[@backingInt(f.kind)] += 1; + return @import("prompt_flight.zig").land(self); } pub fn drain(self: *Metal) !void { @@ -304,15 +293,11 @@ pub const Metal = struct { // a command buffer a chunk, each committed before the one ahead is waited on, so the GPU keeps a chunk queued var k: usize = 0; while (at < ids.len) { - while (k < s.chunks.len and s.chunks[k] <= at) k += 1; - var end = if (k < s.chunks.len) @min(s.chunks[k], ids.len) else ids.len; - // the decision row ends its chunk: the rows after it read the lessons it latched - if (c.decide != lesson_route.none and c.decide >= at and c.decide + 1 < end) end = @intCast(c.decide + 1); - const rows = @min(self.o.chunk, end - at); + const decision: ?usize = if (c.decide == lesson_route.none) null else @intCast(c.decide); + const rows = @import("prompt_flight.zig").rows(s, at, self.o.chunk, decision, &k); if (s.isCancelled()) return error.Cancelled; // release() waits for the chunk still in flight try kv_pages.grow(c, at + rows); - try self.submit(.prefill, self.next, Chunks{ .n = ids.len, .c = c, .at = at, .rows = rows }); - while (self.flights.items.len > 1) try self.land(); + try @import("prompt_flight.zig").submit(self, s, Chunks{ .n = ids.len, .c = c, .at = at, .rows = rows }); at += rows; // landed tokens only: the newest chunk is still queued on the GPU s.fill_at.store(@intCast(at - rows), .monotonic); @@ -590,7 +575,9 @@ pub const Metal = struct { fn releaseFn(ptr: *anyopaque, s: *lanes.Stream) void { const self: *Metal = @ptrCast(@alignCast(ptr)); - self.drain() catch {}; + self.drain() catch { + @import("prompt_flight.zig").forget(self, s); + }; if (self.caches.fetchRemove(s)) |kv| { kv.value.deinit(&self.pool); self.gpa.destroy(kv.value); diff --git a/zig/src/families/nemotron/prompt_flight.zig b/zig/src/families/nemotron/prompt_flight.zig new file mode 100644 index 0000000000..06bbc26613 --- /dev/null +++ b/zig/src/families/nemotron/prompt_flight.zig @@ -0,0 +1,45 @@ +//! Nemotron's existing pipelined completion boundary; a prompt observation lives until its flight lands. +const lanes = @import("lanes"); + +pub const Prompt = struct { stream: *lanes.Stream, rows: usize }; + +/// The original prompt cuts: requested boundaries, the decision row, and the configured row cap. +pub inline fn rows(s: *const lanes.Stream, at: usize, cap: usize, decision: ?usize, next: *usize) usize { + while (next.* < s.chunks.len and s.chunks[next.*] <= at) next.* += 1; + var end = if (next.* < s.chunks.len) @min(s.chunks[next.*], s.prompt_len) else s.prompt_len; + if (decision) |d| if (d >= at and d + 1 < end) { + end = d + 1; + }; + return @min(cap, end - at); +} + +/// Keep the existing one-chunk lookahead. Submission is not completion. +pub inline fn submit(b: anytype, s: *lanes.Stream, body: anytype) !void { + try b.submit(.prefill, b.next, body); + b.flights.items[b.flights.items.len - 1].prompt = .{ .stream = s, .rows = body.rows }; + while (b.flights.items.len > 1) try b.land(); +} + +/// The real backend and recording command buffers use precisely the same wait, error and timing path. +pub inline fn land(b: anytype) !void { + const f = b.flights.orderedRemove(0); + defer f.cb.release(); + f.cb.wait(); + if (f.cb.failure()) |text| { + @import("std").log.err("command buffer failed: {s}", .{text}); + return error.GpuFailed; + } + b.landed = @max(b.landed, f.end); + b.last_ms = f.cb.gpuSeconds() * 1e3; + b.gpu_ms += b.last_ms; + b.timing.ms[@backingInt(f.kind)] += b.last_ms; + b.timing.count[@backingInt(f.kind)] += 1; + if (f.prompt) |p| lanes.prefill_widths.completed(b.gpa, p.stream, p.rows); +} + +/// An existing failed drain may leave flights queued after release; none may keep the freed request pointer. +pub inline fn forget(b: anytype, s: *lanes.Stream) void { + for (b.flights.items) |*f| if (f.prompt) |p| { + if (p.stream == s) f.prompt = null; + }; +} diff --git a/zig/src/families/nemotron/prompt_flight_test.zig b/zig/src/families/nemotron/prompt_flight_test.zig new file mode 100644 index 0000000000..b4676f6375 --- /dev/null +++ b/zig/src/families/nemotron/prompt_flight_test.zig @@ -0,0 +1,135 @@ +//! Actual Nemotron cuts, lookahead and completion accounting, with recording command buffers and no GPU. +const std = @import("std"); +const lanes = @import("lanes"); +const workflow = @import("prompt_flight.zig"); +const a = std.testing.allocator; +const t = std.testing; +const Kind = enum(u2) { prefill, step, verify, draft }; +const State = struct { waits: usize = 0, releases: usize = 0 }; +const Buffer = struct { + state: *State, + pub fn wait(b: Buffer) void { + b.state.waits += 1; + } + pub fn release(b: Buffer) void { + b.state.releases += 1; + } + pub fn failure(_: Buffer) ?[]const u8 { + return null; + } + pub fn gpuSeconds(_: Buffer) f64 { + return 0.002; + } +}; +const Flight = struct { end: u64, cb: Buffer, kind: Kind, prompt: ?workflow.Prompt = null }; +const Fake = struct { + gpa: std.mem.Allocator = a, + next: u64 = 1, + landed: u64 = 0, + last_ms: f64 = 0, + gpu_ms: f64 = 0, + timing: struct { ms: [4]f64 = @splat(0), count: [4]usize = @splat(0) } = .{}, + states: [32]State = @splat(.{}), + submitted: usize = 0, + flights: std.ArrayList(Flight) = .empty, + fail_submit: bool = false, + pub fn submit(f: *Fake, kind: Kind, end: u64, _: anytype) !void { + if (f.fail_submit) return error.SubmitFailed; + try f.flights.append(f.gpa, .{ .end = end, .kind = kind, .cb = .{ .state = &f.states[f.submitted] } }); + f.submitted += 1; + } + pub fn land(f: *Fake) !void { + try workflow.land(f); + } + fn drain(f: *Fake) !void { + while (f.flights.items.len > 0) try f.land(); + } + fn deinit(f: *Fake) void { + f.flights.deinit(f.gpa); + } +}; + +fn fill(f: *Fake, s: *lanes.Stream, cap: usize, decision: ?usize) !void { + var at: usize = 0; + var next: usize = 0; + while (at < s.prompt_len) { + const n = workflow.rows(s, at, cap, decision, &next); + try workflow.submit(f, s, .{ .rows = n }); + at += n; + } +} + +test "Nemotron prefill widths: actual requested cuts, decision row and cap land in order with one chunk ahead" { + const ids: [29]u32 = @splat(7); + var s = try lanes.Stream.init(a, .{ .id = "cut", .prompt = &ids, .max_new = 1, .chunks = &.{ 3, 20 } }); + defer s.deinit(a); + var f: Fake = .{}; + defer f.deinit(); + try fill(&f, &s, 8, 12); + try t.expectEqualSlices(u32, &.{ 3, 8, 2, 7, 8 }, s.prefill_widths.items); + try t.expectEqual(@as(usize, 1), f.flights.items.len); + try t.expectEqual(@as(usize, 0), f.states[5].waits); + try f.drain(); + try t.expectEqualSlices(u32, &.{ 3, 8, 2, 7, 8, 1 }, s.prefill_widths.items); + for (f.states[0..6]) |cb| { + try t.expectEqual(@as(usize, 1), cb.waits); + try t.expectEqual(@as(usize, 1), cb.releases); + } + try t.expectEqual(@as(usize, 6), f.timing.count[0]); + try t.expectEqual(@as(f64, 12), f.gpu_ms); +} + +test "Nemotron prefill widths: one pass waits once; verify and draft buffers never count as prompt passes" { + var s = try lanes.Stream.init(a, .{ .id = "one", .prompt = &.{ 1, 2, 3, 4, 5, 6, 7 }, .max_new = 1 }); + defer s.deinit(a); + var f: Fake = .{}; + defer f.deinit(); + try fill(&f, &s, 2048, null); + try t.expectEqual(@as(usize, 0), s.prefill_widths.items.len); + try f.submit(.verify, 9, .{}); + try f.submit(.draft, 10, .{}); + try f.drain(); + try t.expectEqualSlices(u32, &.{7}, s.prefill_widths.items); + try t.expectEqualSlices(usize, &.{ 1, 0, 1, 1 }, &f.timing.count); + try t.expectEqual(@as(u64, 10), f.landed); + for (f.states[0..3]) |cb| try t.expectEqual(@as(usize, 1), cb.waits); +} + +test "Nemotron prefill widths: interleaved flights retain their request owner and failed submission appends nothing" { + var first = try lanes.Stream.init(a, .{ .id = "first", .prompt = &.{ 1, 2, 3 }, .max_new = 1 }); + defer first.deinit(a); + var second = try lanes.Stream.init(a, .{ .id = "second", .prompt = &.{ 4, 5 }, .max_new = 1 }); + defer second.deinit(a); + var f: Fake = .{}; + defer f.deinit(); + try workflow.submit(&f, &first, .{ .rows = @as(usize, 3) }); + try workflow.submit(&f, &second, .{ .rows = @as(usize, 2) }); + try t.expectEqualSlices(u32, &.{3}, first.prefill_widths.items); + try t.expectEqual(@as(usize, 0), second.prefill_widths.items.len); + f.fail_submit = true; + try t.expectError(error.SubmitFailed, workflow.submit(&f, &second, .{ .rows = @as(usize, 1) })); + try t.expectEqual(@as(usize, 1), f.flights.items.len); + try f.drain(); + try t.expectEqualSlices(u32, &.{3}, first.prefill_widths.items); + try t.expectEqualSlices(u32, &.{2}, second.prefill_widths.items); +} + +test "Nemotron prefill widths: failed release drain forgets only its owner before the stream can be freed" { + var first = try lanes.Stream.init(a, .{ .id = "released", .prompt = &.{ 1, 2, 3 }, .max_new = 1 }); + defer first.deinit(a); + var second = try lanes.Stream.init(a, .{ .id = "retained", .prompt = &.{ 4, 5 }, .max_new = 1 }); + defer second.deinit(a); + var f: Fake = .{}; + defer f.deinit(); + try f.submit(.prefill, 1, .{}); + f.flights.items[0].prompt = .{ .stream = &first, .rows = 3 }; + try f.submit(.prefill, 2, .{}); + f.flights.items[1].prompt = .{ .stream = &second, .rows = 2 }; + workflow.forget(&f, &first); + try t.expect(f.flights.items[0].prompt == null); + try t.expect(f.flights.items[1].prompt.?.stream == &second); + try t.expectEqual(@as(usize, 0), f.states[0].waits); + try f.drain(); + try t.expectEqual(@as(usize, 0), first.prefill_widths.items.len); + try t.expectEqualSlices(u32, &.{2}, second.prefill_widths.items); +} diff --git a/zig/src/metal_prefill_widths_compile.zig b/zig/src/metal_prefill_widths_compile.zig new file mode 100644 index 0000000000..552cc02473 --- /dev/null +++ b/zig/src/metal_prefill_widths_compile.zig @@ -0,0 +1,12 @@ +//! Force native Nemotron's actual prompt and completion bodies to compile; never run this object. +const tf = @import("tensorfold"); +const lanes = @import("lanes"); + +export fn nemotronPrefillWidthsCompile(b: *tf.nemotron.backend.Metal, s: *lanes.Stream) u8 { + const backend = b.backend(); + backend.prefill(s) catch return 1; + const first = backend.first(s, s.prompt_len) catch return 1; + _ = backend.read(first) catch return 1; + backend.release(s); + return 0; +} diff --git a/zig/src/native/qwen27_prompt_rows.zig b/zig/src/native/qwen27_prompt_rows.zig index ff7cf679ad..9f885bb3a3 100644 --- a/zig/src/native/qwen27_prompt_rows.zig +++ b/zig/src/native/qwen27_prompt_rows.zig @@ -1,5 +1,4 @@ //! Each prompt chunk chooses its width from the streams that can decode beside it. -const std = @import("std"); pub fn choose(b: anytype, current: anytype) usize { for (b.owners[0..b.runner.slots]) |owner| if (owner) |stream| { if (stream != current and !stream.finished and !stream.paused and !b.fills.contains(stream)) { @@ -8,12 +7,3 @@ pub fn choose(b: anytype, current: anytype) usize { }; return b.chunk; } - -/// Only completed target prompt passes count; an observation failure never fails inference or reports a partial list. -pub fn completed(gpa: std.mem.Allocator, stream: anytype, rows: usize) void { - if (stream.prefill_widths_failed) return; - stream.prefill_widths.append(gpa, @intCast(rows)) catch { - stream.prefill_widths.clearRetainingCapacity(); - stream.prefill_widths_failed = true; - }; -} diff --git a/zig/src/native/qwen27_slots.zig b/zig/src/native/qwen27_slots.zig index b341a41284..73b80877df 100644 --- a/zig/src/native/qwen27_slots.zig +++ b/zig/src/native/qwen27_slots.zig @@ -132,7 +132,7 @@ pub fn chunk(b: *Metal, s: *lanes.Stream, f: *Fill, _: u32) !bool { if (b.ringOf(slot)) |g| try ringChunk(b, g, part.len); } charge(b, s, before); - @import("qwen27_prompt_rows.zig").completed(b.gpa, s, part.len); + lanes.prefill_widths.completed(b.gpa, s, part.len); f.at = end; s.fill_at.store(@intCast(end), .monotonic); // a prompt chunk waits: `at` is landed progress if (end == f.rewind) saveRewind(b, f.rewind); diff --git a/zig/src/qwen27_prefill_stats_test.zig b/zig/src/qwen27_prefill_stats_test.zig index 4a47778740..c4a6117936 100644 --- a/zig/src/qwen27_prefill_stats_test.zig +++ b/zig/src/qwen27_prefill_stats_test.zig @@ -3,7 +3,7 @@ const std = @import("std"); const lanes = @import("lanes"); const api = @import("engine_api"); const Mailbox = @import("server/mailbox.zig").Mailbox; -const rows = @import("native/qwen27_prompt_rows.zig"); +const rows = lanes.prefill_widths; const t = std.testing; const gpa = t.allocator; From 7897ac3da8e4d993b050d64fee0dda073f11ca6a Mon Sep 17 00:00:00 2001 From: ashhart <71890694+ashhart@users.noreply.github.com> Date: Thu, 5 Nov 2026 16:47:29 -0800 Subject: [PATCH 3/4] Append completed Metal prompt widths without moving prompt loops --- build.zig | 1 + zig/build/prefill_widths_metal.zig | 15 ++++++ zig/build/qwen4.zig | 3 ++ zig/qwen4_prefill_widths_host.zig | 5 ++ zig/src/families/gemma4/backend.zig | 1 + zig/src/families/glm_dsa/backend.zig | 1 + zig/src/families/kimi_k3/backend.zig | 1 + zig/src/families/lfm2_moe/backend.zig | 1 + zig/src/families/llama/backend.zig | 1 + zig/src/families/minimax_m3/backend.zig | 1 + zig/src/families/qwen3_5/prefill.zig | 1 + zig/src/families/qwen4/engine_test.zig | 2 +- zig/src/families/qwen4/lanes_adapter.zig | 1 + zig/src/families/qwen4/lanes_adapter_test.zig | 4 ++ zig/src/families/spark25/metal_backend.zig | 1 + zig/src/families/step5/backend.zig | 1 + zig/src/metal_prefill_rest_compile.zig | 48 +++++++++++++++++++ zig/src/qwen35_cache_test.zig | 23 +++++++++ zig/src/qwen4_prefill_widths_compile.zig | 4 ++ 19 files changed, 114 insertions(+), 1 deletion(-) create mode 100644 zig/qwen4_prefill_widths_host.zig create mode 100644 zig/src/metal_prefill_rest_compile.zig create mode 100644 zig/src/qwen4_prefill_widths_compile.zig diff --git a/build.zig b/build.zig index 0b73f30235..3aa689962a 100644 --- a/build.zig +++ b/build.zig @@ -180,6 +180,7 @@ pub fn build(b: *std.Build) void { @import("zig/build/gemma4_host.zig").compileHost(b, gemma4_host_mods.metal, gemma4_host_mods.engine, gemma4_host_mods.lanes, gemma4_host_mods.sources); @import("zig/build/qwen27_host.zig").compileHost(b, gemma4_host_mods.metal, gemma4_host_mods.engine, gemma4_host_mods.lanes); @import("zig/build/prefill_widths_metal.zig").compileHost(b, gemma4_host_mods.engine, gemma4_host_mods.lanes); + @import("zig/build/prefill_widths_metal.zig").rest(b, gemma4_host_mods.engine, gemma4_host_mods.lanes); } @import("zig/build/qwen27_host.zig").tests(b, test_step); @import("zig/build/prefill_widths_metal.zig").tests(b, test_step); diff --git a/zig/build/prefill_widths_metal.zig b/zig/build/prefill_widths_metal.zig index 5285161019..bbba897dfc 100644 --- a/zig/build/prefill_widths_metal.zig +++ b/zig/build/prefill_widths_metal.zig @@ -16,3 +16,18 @@ pub fn compileHost(b: *std.Build, engine: *std.Build.Module, lanes: *std.Build.M _ = object.getEmittedBin(); b.step("check-nemotron-prefill-widths-host", "Compile actual Metal prefill, read, land and release without shaders").dependOn(&object.step); } + +pub fn rest(b: *std.Build, engine: *std.Build.Module, lanes: *std.Build.Module) void { + const metal = engine.import_table.get("metal").?; + const bodies = b.createModule(.{ .root_source_file = b.path("zig/src/metal_prefill_rest_compile.zig"), .target = b.graph.host, .optimize = .debug, .link_libc = true, .imports = &.{ .{ .name = "tensorfold", .module = engine }, .{ .name = "lanes", .module = lanes }, .{ .name = "kimi", .module = @import("metal_targets.zig").kimiModule(b, b.graph.host, .debug, metal, engine, lanes) } } }); + const object = b.addObject(.{ .name = "metal-prefill-rest-compile", .root_module = bodies }); + _ = object.getEmittedBin(); + const check = b.step("check-metal-prefill-rest-host", "Actual remaining Metal prompt callback bodies without shaders or execution"); + check.dependOn(&object.step); + const api = b.createModule(.{ .root_source_file = b.path("zig/src/core/engine_api.zig"), .target = b.graph.host, .link_libc = true, .imports = &.{.{ .name = "lanes", .module = lanes }} }); + const host = b.createModule(.{ .root_source_file = b.path("zig/src/native/qwen4_metal.zig"), .target = b.graph.host, .link_libc = true, .imports = &.{ .{ .name = "metal", .module = metal }, .{ .name = "tensorfold", .module = engine }, .{ .name = "engine_api", .module = api }, .{ .name = "contract", .module = @import("kernel_contract.zig").module(b, b.graph.host, .debug) } } }); + @import("qwen4.zig").metalNative(b, b.graph.host, host); + const qwen = b.addObject(.{ .name = "qwen4-prefill-widths-compile", .root_module = b.createModule(.{ .root_source_file = b.path("zig/src/qwen4_prefill_widths_compile.zig"), .target = b.graph.host, .optimize = .debug, .link_libc = true, .imports = &.{.{ .name = "host", .module = host }} }) }); + _ = qwen.getEmittedBin(); + check.dependOn(&qwen.step); +} diff --git a/zig/build/qwen4.zig b/zig/build/qwen4.zig index 559c94622f..603d8f9804 100644 --- a/zig/build/qwen4.zig +++ b/zig/build/qwen4.zig @@ -32,6 +32,9 @@ pub fn steps(b: *std.Build, parent: *std.Build.Step) void { @import("qwen4_cuda_attention_host.zig").wire(b, m); const step = b.step("test-qwen4-host", "Qwen4 explicit config admission and synthetic graph contracts"); step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = m })).step); + const widths_root = b.createModule(.{ .root_source_file = b.path("zig/qwen4_prefill_widths_host.zig"), .target = b.graph.host, .optimize = .debug, .link_libc = true, .imports = &.{.{ .name = "qwen4", .module = m }} }); + const widths = b.step("test-qwen4-prefill-widths-host", "Existing adapter fixture counts Metal prompt passes and excludes CUDA and verify"); + for ([_]*std.Build.Module{ m, widths_root }) |root| widths.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = root, .filters = &.{"Qwen4 lanes:"} })).step); const mixer_step = b.step("test-qwen4-native-mixers-host", "Actual native mixer bodies under CPU collective instruments"); mixer_step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = m, .filters = &.{ "Qwen4 GDN native contract", "Qwen4 HC native contract", "Qwen4 rounding bounds", "Qwen4 PLE native contract" } })).step); const ple_step = b.step("test-qwen4-ple-host", "Generated native lexical bodies against binary64 nodes"); diff --git a/zig/qwen4_prefill_widths_host.zig b/zig/qwen4_prefill_widths_host.zig new file mode 100644 index 0000000000..668570a34a --- /dev/null +++ b/zig/qwen4_prefill_widths_host.zig @@ -0,0 +1,5 @@ +//! Select the Metal tag for the existing memory-backed Qwen4 adapter fixture; no device calls. +pub const metal_prefill_widths = true; +test { + _ = @import("qwen4"); +} diff --git a/zig/src/families/gemma4/backend.zig b/zig/src/families/gemma4/backend.zig index 26afd89094..416afe6bd8 100644 --- a/zig/src/families/gemma4/backend.zig +++ b/zig/src/families/gemma4/backend.zig @@ -231,6 +231,7 @@ pub const Metal = struct { cache.tap = .{ .buf = last, .start = ids.len - 1, .rows = 1 }; } } + lanes.prefill_widths.completed(self.gpa, stream, end - at); at = end; if (std.mem.indexOfScalar(u32, stream.reuse.marks, @intCast(cache.len)) != null) if (stream.reuse.hook) |k| k.at(k.ptr, stream, @intCast(cache.len)); } diff --git a/zig/src/families/glm_dsa/backend.zig b/zig/src/families/glm_dsa/backend.zig index 1bda04b6f6..b853f8216f 100644 --- a/zig/src/families/glm_dsa/backend.zig +++ b/zig/src/families/glm_dsa/backend.zig @@ -116,6 +116,7 @@ pub const Backend = struct { if (mark < s.reuse.marks.len) end = @min(end, s.reuse.marks[mark]); // a mark ends a chunk: its state is kept there const n = @min(end - at, gpu.prompt_rows); if (try fl.chunk(e, i, at, n, false)) |t| b.first[i] = t; + lanes.prefill_widths.completed(b.gpa, s, n); at += n; s.fill_at.store(at, .monotonic); // fl.chunk waits on rank 0: `at` is landed progress if (std.mem.indexOfScalar(u32, s.reuse.marks, at) != null) if (s.reuse.hook) |k| k.at(k.ptr, s, at); diff --git a/zig/src/families/kimi_k3/backend.zig b/zig/src/families/kimi_k3/backend.zig index 5e7016b631..75c0393a04 100644 --- a/zig/src/families/kimi_k3/backend.zig +++ b/zig/src/families/kimi_k3/backend.zig @@ -165,6 +165,7 @@ pub const Metal = struct { }; const rows: u32 = @intCast(end - pos); try m.run(&.{.{ .slot = slot, .rows = rows, .commit = true }}, s.prompt()[pos..end]); + lanes.prefill_widths.completed(m.a, s, rows); m.storeLast(slot, rows - 1); m.state.advance(slot, rows, true, rows); s.fill_at.store(@intCast(end), .release); diff --git a/zig/src/families/lfm2_moe/backend.zig b/zig/src/families/lfm2_moe/backend.zig index 3e72e800fd..ffe1ffb88d 100644 --- a/zig/src/families/lfm2_moe/backend.zig +++ b/zig/src/families/lfm2_moe/backend.zig @@ -204,6 +204,7 @@ pub const Metal = struct { try k.take(k.ptr, stream, f, at); if (end == ids.len) try self.keepLastTap(cache, end - 1 - at, end - 1); } + lanes.prefill_widths.completed(self.gpa, stream, end - at); at = end; if (stream.reuse.hook) |hk| if (std.mem.indexOfScalar(u32, stream.reuse.marks, @intCast(at)) != null) hk.at(hk.ptr, stream, @intCast(at)); diff --git a/zig/src/families/llama/backend.zig b/zig/src/families/llama/backend.zig index 6c465fc4ca..cb6d81a345 100644 --- a/zig/src/families/llama/backend.zig +++ b/zig/src/families/llama/backend.zig @@ -190,6 +190,7 @@ pub const Metal = struct { const s = if (end - at > st.batch_rows) &self.prompt else &self.scratch; try fwd.run(self.model, s, cache, ids[at..end], .{}); cache.len += end - at; + if (stream) |sm| lanes.prefill_widths.completed(self.gpa, sm, end - at); at = end; if (stream) |sm| if (std.mem.indexOfScalar(u32, sm.reuse.marks, @intCast(cache.len)) != null) if (sm.reuse.hook) |k| k.at(k.ptr, sm, @intCast(cache.len)); } diff --git a/zig/src/families/minimax_m3/backend.zig b/zig/src/families/minimax_m3/backend.zig index bc8bd3106f..fc68870438 100644 --- a/zig/src/families/minimax_m3/backend.zig +++ b/zig/src/families/minimax_m3/backend.zig @@ -258,6 +258,7 @@ pub const Metal = struct { } } cache.len += end - at; + if (stream) |sm| lanes.prefill_widths.completed(self.gpa, sm, end - at); at = end; if (stream) |sm| if (std.mem.indexOfScalar(u32, sm.reuse.marks, @intCast(cache.len)) != null) if (sm.reuse.hook) |k| k.at(k.ptr, sm, @intCast(cache.len)); } diff --git a/zig/src/families/qwen3_5/prefill.zig b/zig/src/families/qwen3_5/prefill.zig index 46f2610057..0be0f2127e 100644 --- a/zig/src/families/qwen3_5/prefill.zig +++ b/zig/src/families/qwen3_5/prefill.zig @@ -33,6 +33,7 @@ pub fn run(cache: *st.Cache, stream: *lanes.Stream, chunk: usize, ctx: anytype, while (mark < stream.reuse.marks.len and stream.reuse.marks[mark] <= at) mark += 1; if (mark < stream.reuse.marks.len) end = @min(end, stream.reuse.marks[mark]); try advance(ctx, cache, ids[at..end]); + lanes.prefill_widths.completed(ctx.gpa, stream, end - at); at = end; stream.fill_at.store(@intCast(at), .monotonic); if (stream.reuse.hook) |hook| if (std.mem.indexOfScalar(u32, stream.reuse.marks, @intCast(at)) != null) diff --git a/zig/src/families/qwen4/engine_test.zig b/zig/src/families/qwen4/engine_test.zig index d49ee466cd..81571ca07d 100644 --- a/zig/src/families/qwen4/engine_test.zig +++ b/zig/src/families/qwen4/engine_test.zig @@ -44,7 +44,7 @@ const Arena = @import("arena.zig").Arena(B); const Tensor = Arena.Tensor; pub const G = struct { pub const Storage = B; - pub const kind = @import("native_admission.zig").Backend.cuda; + pub const kind: @import("native_admission.zig").Backend = if (@hasDecl(@import("root"), "metal_prefill_widths")) .metal else .cuda; b: *B, project_rows: usize = 0, begins: usize = 0, diff --git a/zig/src/families/qwen4/lanes_adapter.zig b/zig/src/families/qwen4/lanes_adapter.zig index 7b4514abce..355822a130 100644 --- a/zig/src/families/qwen4/lanes_adapter.zig +++ b/zig/src/families/qwen4/lanes_adapter.zig @@ -112,6 +112,7 @@ pub fn Adapter(comptime E: type, comptime D: type, comptime L: type) type { l.e.fail(); return err; }; + if (comptime E.Gpu.kind == .metal) L.prefill_widths.completed(l.a, s, end - lane.filled); lane.filled = end; s.fill_at.store(@intCast(end), .release); for (s.reuse.marks) |cut| if (cut == end) { diff --git a/zig/src/families/qwen4/lanes_adapter_test.zig b/zig/src/families/qwen4/lanes_adapter_test.zig index a2f96a1699..e173dd48aa 100644 --- a/zig/src/families/qwen4/lanes_adapter_test.zig +++ b/zig/src/families/qwen4/lanes_adapter_test.zig @@ -87,6 +87,8 @@ test "Qwen4 lanes: shared verify, prefix rejection, stale handles and cold-cut p try t.expectEqual(@as(u32, 3), left.fill_at.load(.acquire)); try t.expect(try b.vtable.prefill_step.?(b.ptr, &left, 0)); try b.prefill(&right); + try t.expectEqualSlices(u32, if (E.Gpu.kind == .metal) &.{ 3, 3 } else &.{}, left.prefill_widths.items); + try t.expectEqualSlices(u32, if (E.Gpu.kind == .metal) &.{2} else &.{}, right.prefill_widths.items); const lh = try b.first(&left, 6); const rh = try b.first(&right, 2); try t.expect(lh != rh); @@ -104,6 +106,8 @@ test "Qwen4 lanes: shared verify, prefix rejection, stale handles and cold-cut p try t.expectError(error.Qwen4DrawPosition, b.verify(&bad, &verified)); try t.expectEqual(begins, f.e.gpu.begins); try b.verify(&windows, &verified); + try t.expectEqualSlices(u32, if (E.Gpu.kind == .metal) &.{ 3, 3 } else &.{}, left.prefill_widths.items); + try t.expectEqualSlices(u32, if (E.Gpu.kind == .metal) &.{2} else &.{}, right.prefill_widths.items); try t.expectEqual(@as(usize, 4), f.e.gpu.project_rows); try t.expectEqualSlices(u32, &.{ 2, 3 }, &drafts0); try t.expectEqualSlices(u32, &.{ 0, 0, 0 }, &sample0); diff --git a/zig/src/families/spark25/metal_backend.zig b/zig/src/families/spark25/metal_backend.zig index 844ac939a1..a27f9a16d8 100644 --- a/zig/src/families/spark25/metal_backend.zig +++ b/zig/src/families/spark25/metal_backend.zig @@ -180,6 +180,7 @@ pub const Metal = struct { const s = if (end - at > st.batch_rows) &self.prompt else &self.scratch; try fwd.run(self.model, s, cache, ids[at..end], .{}); cache.len += end - at; + if (stream) |sm| lanes.prefill_widths.completed(self.gpa, sm, end - at); at = end; if (stream) |sm| if (std.mem.indexOfScalar(u32, sm.reuse.marks, @intCast(cache.len)) != null) if (sm.reuse.hook) |k| k.at(k.ptr, sm, @intCast(cache.len)); } diff --git a/zig/src/families/step5/backend.zig b/zig/src/families/step5/backend.zig index 3a51f022d7..408075b172 100644 --- a/zig/src/families/step5/backend.zig +++ b/zig/src/families/step5/backend.zig @@ -120,6 +120,7 @@ pub const Metal = struct { @memcpy(cache.hidden.contents()[0 .. c.hidden * 4], last[0 .. c.hidden * 4]); const known = if (end < ids.len) n else n - 1; // rows whose next token is in ids if (c.mtp > 0 and known > 0) try fwd.mtp(self.model, s, cache, 0, .{ .buf = s.hid }, ids[at + 1 .. at + 1 + known], start); + if (stream) |sm| lanes.prefill_widths.completed(self.gpa, sm, end - at); at = end; if (stream) |sm| if (std.mem.indexOfScalar(u32, sm.reuse.marks, @intCast(cache.len)) != null) if (sm.reuse.hook) |k| k.at(k.ptr, sm, @intCast(cache.len)); } diff --git a/zig/src/metal_prefill_rest_compile.zig b/zig/src/metal_prefill_rest_compile.zig new file mode 100644 index 0000000000..e672039eba --- /dev/null +++ b/zig/src/metal_prefill_rest_compile.zig @@ -0,0 +1,48 @@ +//! Compile actual Metal prefill callback bodies; this object is never executed. +const tf = @import("tensorfold"); +const lanes = @import("lanes"); + +export fn gemma4PrefillCompile(b: *tf.gemma4.backend.Metal, s: *lanes.Stream) bool { + b.backend().prefill(s) catch return false; + return true; +} + +export fn lfm2PrefillCompile(b: *tf.lfm2_moe.backend.Metal, s: *lanes.Stream) bool { + b.backend().prefill(s) catch return false; + return true; +} + +export fn step5PrefillCompile(b: *tf.step5.backend.Metal, s: *lanes.Stream) bool { + b.backend().prefill(s) catch return false; + return true; +} + +export fn sparkPrefillCompile(b: *tf.spark25_metal.backend.Metal, s: *lanes.Stream) bool { + b.backend().prefill(s) catch return false; + return true; +} + +export fn minimaxPrefillCompile(b: *tf.minimax_m3.backend.Metal, s: *lanes.Stream) bool { + b.backend().prefill(s) catch return false; + return true; +} + +export fn llamaPrefillCompile(b: *tf.llama.backend.Metal, s: *lanes.Stream) bool { + b.backend().prefill(s) catch return false; + return true; +} + +export fn qwen35PrefillCompile(b: *tf.qwen35.backend.Metal, s: *lanes.Stream) bool { + b.backend().prefill(s) catch return false; + return true; +} + +export fn dsaPrefillCompile(b: *tf.glm_dsa.backend.Backend, s: *lanes.Stream) bool { + b.backend().prefill(s) catch return false; + return true; +} + +export fn kimiPrefillCompile(b: *@import("kimi").backend.Metal, s: *lanes.Stream) bool { + b.backend().prefill(s) catch return false; + return true; +} diff --git a/zig/src/qwen35_cache_test.zig b/zig/src/qwen35_cache_test.zig index ccbbc3ca75..2918c37fc0 100644 --- a/zig/src/qwen35_cache_test.zig +++ b/zig/src/qwen35_cache_test.zig @@ -181,6 +181,7 @@ test "Qwen retained: malformed positions and lengths refuse before changing a ca } const Harness = struct { + gpa: std.mem.Allocator = a, metal: q.Metal = undefined, store: pc.Store = undefined, active: ?*st.Cache = null, @@ -302,3 +303,25 @@ test "Qwen retained: a failed chunk never captures its incomplete state" { try t.expectError(error.FakeForward, h.pass(&failed, ids[0..300], &.{})); try t.expectEqual(@as(usize, 0), h.store.entries.items.len); } + +test "Qwen retained: completed widths record actual chunks and one pass" { + const Context = struct { + gpa: std.mem.Allocator = a, + fn advance(_: @This(), cache: *st.Cache, ids: []const u32) !void { + try step({}, cache, ids); + } + }; + var allocations: mtl.Allocations = .{}; + const device: mtl.Device = .{ .allocations = &allocations }; + const ids: [23]u32 = @splat(7); + for ([_]bool{ false, true }) |chunked| { + var cache = try st.Cache.init(a, device, 128); + defer cache.deinit(); + seed(&cache); + var stream = try lanes.Stream.init(a, .{ .id = "widths", .prompt = &ids, .max_new = 4, .chunks = if (chunked) &.{ 5, 18 } else &.{} }); + defer stream.deinit(a); + try prefill.run(&cache, &stream, 128, Context{}, Context.advance); + try t.expectEqualSlices(u32, if (chunked) &.{ 5, 13, 5 } else &.{23}, stream.prefill_widths.items); + try t.expectEqual(@as(usize, 23), cache.len); + } +} diff --git a/zig/src/qwen4_prefill_widths_compile.zig b/zig/src/qwen4_prefill_widths_compile.zig new file mode 100644 index 0000000000..087edbf253 --- /dev/null +++ b/zig/src/qwen4_prefill_widths_compile.zig @@ -0,0 +1,4 @@ +//! Emit actual Qwen4 Metal startup and its lane callbacks, never execute this object. +export fn qwen4PrefillWidthsCompile() usize { + return @intFromPtr(&@import("host").open); +} From b215807d0cb0ab83a4a52da89c97725f40fa3dac Mon Sep 17 00:00:00 2001 From: ashhart <71890694+ashhart@users.noreply.github.com> Date: Fri, 6 Nov 2026 08:08:32 +0000 Subject: [PATCH 4/4] changelog: #2293 under Unreleased --- CHANGELOG.md | 3 +++ 1 file changed, 3 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 7086a2d5ba..4dcf39c832 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,9 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel ## Unreleased +- Metal: a request's finished stats list its prompt pass widths (`tensorfold.prefill_widths`, one entry per completed + target pass) on every Metal text family. Telemetry only (#2293). + - Internal: 100 lean findings fixed in core, the server, cluster and the model families (comments reworded to fit one line, and Kolibri-1's page geometry in its own file), with no change in behaviour (#2292).