Metal: a request's stats list its prompt pass widths, on every text family

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).
 

build.zig

@@ -179,8 +179,11 @@ 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/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);
     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);

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"),

zig/build/prefill_widths_metal.zig

@@ -0,0 +1,33 @@
+//! 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);
+}
+
+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);
+}

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);
 }

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);
 }

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");

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");
+}

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 } });

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");

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;
+    };
+}

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);

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));
         }

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);

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,12 +132,16 @@ 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 => {},
         }
     }
     pub fn deinit(b: *Box) void {
         for (&b.tokens) |*ts| ts.deinit(a);
+        for (&b.widths) |*ws| ws.deinit(a);
     }
 };

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());
+}

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);

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);

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));

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));
         }

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));
         }