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). diff --git a/build.zig b/build.zig index b9c7d32198..3aa689962a 100644 --- a/build.zig +++ b/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); 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..bbba897dfc --- /dev/null +++ b/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); +} 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/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/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/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/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/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/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/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/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/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/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_slots.zig b/zig/src/native/qwen27_slots.zig index 1dedd55bcb..73b80877df 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); + 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 new file mode 100644 index 0000000000..c4a6117936 --- /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 = lanes.prefill_widths; +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/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); +} 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); + } + }; +}