From 36f3476d3fb2935f9f4437f6314989a31000d40a 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/5] 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 81c14a26f3..7dd890811f 100644 --- a/zig/build/qwen27_host.zig +++ b/zig/build/qwen27_host.zig @@ -66,6 +66,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 7d8b32840f..072d2d4b93 100644 --- a/zig/src/core/lane_host.zig +++ b/zig/src/core/lane_host.zig @@ -300,7 +300,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 3511561e3a80315d7e6b45c73d43337ffb7c1f9a 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/5] 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 2ef64e9798..1d8e969671 100644 --- a/build.zig +++ b/build.zig @@ -174,8 +174,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 660926ad06..f08eb6308b 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 80bb5763e7..b44f6e46e3 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 of it 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 87e2ae35f414ba471b961d05134115e7f4885a99 Mon Sep 17 00:00:00 2001 From: ashhart <71890694+ashhart@users.noreply.github.com> Date: Wed, 4 Nov 2026 15:01:47 +0000 Subject: [PATCH 3/5] Qwen: widen canonical prompt chains before M5 --- docs/recipes/qwen3.8-27b.md | 6 + zig/build/qwen27_host.zig | 5 + zig/kernels/metal/qwen27/attention_mpp.metal | 4 +- zig/src/core/contract/q27_chain_records.zig | 10 + zig/src/core/contract/registry.zig | 2 +- zig/src/families/qwen27/attention.zig | 2 +- .../families/qwen27/attention_chain_test.zig | 100 +++++ zig/src/families/qwen27/attention_plan.zig | 26 +- zig/src/qwen27_speed_test.zig | 2 +- zig/tests/contract/metal_limits.zig | 5 + zig/tests/contract/metal_main.zig | 3 + zig/tests/contract/metal_q27_chain.zig | 280 +++++++++++++ zig/tests/contract/q27_chain_ref.zig | 103 +++++ zig/tests/contract/q27_tail_baseline.metal | 395 ++++++++++++++++++ 14 files changed, 931 insertions(+), 12 deletions(-) create mode 100644 zig/src/core/contract/q27_chain_records.zig create mode 100644 zig/src/families/qwen27/attention_chain_test.zig create mode 100644 zig/tests/contract/metal_q27_chain.zig create mode 100644 zig/tests/contract/q27_chain_ref.zig create mode 100644 zig/tests/contract/q27_tail_baseline.metal diff --git a/docs/recipes/qwen3.8-27b.md b/docs/recipes/qwen3.8-27b.md index 9798f40243..33d18ddd3e 100644 --- a/docs/recipes/qwen3.8-27b.md +++ b/docs/recipes/qwen3.8-27b.md @@ -64,6 +64,12 @@ Larger or unbounded cuts keep the general draw. Sampling time has its own pass-p On Metal, M5 prompt passes request 2,048 rows when no other request is decoding; earlier chips keep the 512-row default. The larger frame must fit the configured context and all streams beside loaded weights and the 8 GiB margin, otherwise M5 keeps 512 rows. Startup reserves the actual prompt scratch sizes as well as stream caches. +Qwen3.8-27B's pre-M5 GQA6 prompt chains follow the admitted frame, up to 2,048 rows, using the +one-row tail's key chunks, half probabilities and merge. Branching and mixed-stream rounds retain 128 rows. +The older simdgroup prompt kernel uses different probability and tile arithmetic, so this path uses the tail at +short widths too. The M5 path and other attention geometries keep their existing selection. +Host planning and address checks pass; Metal byte replacement, checkpoint equality and timing remain to qualify. +No speedup over an earlier release is claimed. A larger frame must still fit the existing admission budget. Wide prompt frames keep separate 512-row activation scratch for small rounds, so decoding binds the small buffers. Admission counts both banks. `TF_QWEN27_WORKING_SET_BYTES` can lower the device working-set budget for admission checks. `TF_PROMPT_ROWS` selects another multiple of 16 from 16 to 2,048. Each chunk takes at most 512 rows beside diff --git a/zig/build/qwen27_host.zig b/zig/build/qwen27_host.zig index 7dd890811f..3abc120ccf 100644 --- a/zig/build/qwen27_host.zig +++ b/zig/build/qwen27_host.zig @@ -19,6 +19,10 @@ pub fn tests(b: *std.Build, test_step: *std.Build.Step) void { const allocated = b.createModule(.{ .root_source_file = b.path("zig/qwen27_allocation_host.zig"), .target = b.graph.host, .link_libc = true, .imports = &.{.{ .name = "lanes", .module = pricing_lanes }} }); const window_step = b.step("test-qwen27-window-host", "Device width, copied trees and counted zero-grant drafting without shaders"); for ([_]*std.Build.Module{ window, taps, allocated }) |m| window_step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = m })).step); + const chain_ref = b.createModule(.{ .root_source_file = b.path("zig/tests/contract/q27_chain_ref.zig"), .target = b.graph.host, .imports = &.{.{ .name = "contract", .module = @import("kernel_contract.zig").module(b, b.graph.host, .debug) }} }); + const chain_run = b.addRunArtifact(b.addTest(.{ .root_module = chain_ref })); + test_step.dependOn(&chain_run.step); + b.step("test-qwen27-chain-reference-host", "FP64 chain attention and reduced-sum witness without a GPU").dependOn(&chain_run.step); const rows = b.createModule(.{ .root_source_file = b.path("zig/src/qwen27_rows_test.zig"), .target = b.graph.host }); b.step("test-qwen27-prompt-rows-host", "Device-aware prompt rows beside owned decoding streams without a GPU").dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = rows, .filters = &.{"qwen27 rows:"} })).step); const ane_policy = b.createModule(.{ .root_source_file = b.path("zig/src/families/qwen27/ane_policy.zig"), .target = b.graph.host }); @@ -51,6 +55,7 @@ pub fn tests(b: *std.Build, test_step: *std.Build.Step) void { var names: std.ArrayList([]const u8) = .empty; tryNames(b, &names, @embedFile("../src/families/qwen27/gdn_test.zig")); tryNames(b, &names, @embedFile("../src/families/qwen27/attention_plan.zig")); + tryNames(b, &names, @embedFile("../src/families/qwen27/attention_chain_test.zig")); const plans = b.addRunArtifact(b.addTest(.{ .root_module = keep_module, .filters = names.items })); plans.has_side_effects = true; b.step("test-qwen27-plans-host", "Shared GDN and attention plans with no shaders or GPU").dependOn(&plans.step); diff --git a/zig/kernels/metal/qwen27/attention_mpp.metal b/zig/kernels/metal/qwen27/attention_mpp.metal index 622adf5595..630fc8c31b 100644 --- a/zig/kernels/metal/qwen27/attention_mpp.metal +++ b/zig/kernels/metal/qwen27/attention_mpp.metal @@ -239,7 +239,7 @@ kernel void q27_attn_tail(device const bfloat* QB [[buffer(0)]], device const Q2 const int q = kt + h * 32 + row; int phys = -1; if (q < P) phys = q; - else if (q < nmax) phys = P + paths[node * MAXD + (q - P)]; + else if (q < nmax) phys = meta[MG + st * MS + 10] ? q : P + paths[node * MAXD + (q - P)]; ((threadgroup vec*)KV)[e] = phys >= 0 ? *(const device vec*)(kbase + phys * kstep + col) : vec(0); } threadgroup_barrier(mem_flags::mem_threadgroup); @@ -287,7 +287,7 @@ kernel void q27_attn_tail(device const bfloat* QB [[buffer(0)]], device const Q2 const int q = kt + row; int phys = -1; if (q < P) phys = q; - else if (q < nmax) phys = P + paths[node * MAXD + (q - P)]; + else if (q < nmax) phys = meta[MG + st * MS + 10] ? q : P + paths[node * MAXD + (q - P)]; ((threadgroup vec*)KV)[e] = phys >= 0 ? *(const device vec*)(vbase + phys * vstep + col) : vec(0); } threadgroup_barrier(mem_flags::mem_threadgroup); diff --git a/zig/src/core/contract/q27_chain_records.zig b/zig/src/core/contract/q27_chain_records.zig new file mode 100644 index 0000000000..f936c5619d --- /dev/null +++ b/zig/src/core/contract/q27_chain_records.zig @@ -0,0 +1,10 @@ +//! Canonical chain attention keeps the tail's 128 threads and the ordered merge's 32 threads. +const limits = @import("limits.zig"); + +// Metal 4 tensor descriptors compile at runtime; M1/M2 ahead-of-time limits are not claimed. +pub const records = [_]limits.MetalLaunch{ + .{ .name = "qwen27-chain one q27_attn_tail", .threads = 128, .tg_memory = 18432, .offline = false }, + .{ .name = "qwen27-chain one q27_attn_merge_single", .threads = 32, .tg_memory = 0, .offline = false }, + .{ .name = "qwen27-chain rank q27_attn_tail", .threads = 128, .tg_memory = 18432, .offline = false }, + .{ .name = "qwen27-chain rank q27_attn_merge_single", .threads = 32, .tg_memory = 0, .offline = false }, +}; diff --git a/zig/src/core/contract/registry.zig b/zig/src/core/contract/registry.zig index d23f166d84..9377b1b2df 100644 --- a/zig/src/core/contract/registry.zig +++ b/zig/src/core/contract/registry.zig @@ -568,7 +568,7 @@ test { } const paging = @import("paging_pending.zig"); /// The Metal collector keeps measured slots and the restored vision declarations. -pub const all_metal = @import("attention_pv_records.zig").records ++ metal ++ @import("delta_slots_records.zig").records ++ @import("flashnext_pending.zig").records ++ @import("simd_pairs_records.zig").records ++ @import("metal_pair_records.zig").records ++ @import("glm_rows_records.zig").records ++ @import("weight_only_records.zig").metal ++ @import("hc_activation_records.zig").metal ++ @import("qwen4_records.zig").metal ++ @import("lfm2_records.zig").records ++ @import("nemotron_experts_records.zig").records ++ @import("nemotron_pages_records.zig").records ++ @import("nemotron_attention_records.zig").records ++ +pub const all_metal = @import("q27_chain_records.zig").records ++ @import("attention_pv_records.zig").records ++ metal ++ @import("delta_slots_records.zig").records ++ @import("flashnext_pending.zig").records ++ @import("simd_pairs_records.zig").records ++ @import("metal_pair_records.zig").records ++ @import("glm_rows_records.zig").records ++ @import("weight_only_records.zig").metal ++ @import("hc_activation_records.zig").metal ++ @import("qwen4_records.zig").metal ++ @import("lfm2_records.zig").records ++ @import("nemotron_experts_records.zig").records ++ @import("nemotron_pages_records.zig").records ++ @import("nemotron_attention_records.zig").records ++ @import("nemotron_ssd_records.zig").records ++ @import("flashnext_attention_records.zig").records ++ @import("spark25_records.zig").records ++ @import("gemma_attention_records.zig").records ++ @import("moss_image_records.zig").records ++ @import("dsv41_metal_records.zig").records ++ @import("minimax_records.zig").records ++ @import("ling_route_records.zig").records ++ @import("qwen27_moe_records.zig").records ++ @import("affine_gather_records.zig").records; const learning = @import("registry_learning.zig"); pub const gradient_pending = learning.gradient_pending; diff --git a/zig/src/families/qwen27/attention.zig b/zig/src/families/qwen27/attention.zig index 7ccd2be6b0..732b4f329a 100644 --- a/zig/src/families/qwen27/attention.zig +++ b/zig/src/families/qwen27/attention.zig @@ -334,7 +334,7 @@ pub const Attention = struct { bind(e, .{ q, scratch.get(.q_rows), scratch.get(.qa) }); e.setValue(args, 3); e.dispatchThreads(mtl.Size.of(g.dim, args.packed_rows, g.kv_heads), mtl.Size.of(32, 4, 1)); - if (prompt and a.prompt != null and (plan.geometry.shared_prefix or !a.units) and plan.streams == 1) { + if (prompt and a.prompt != null and (plan.geometry.shared_prefix or (!a.units and g.group() != 6)) and plan.streams == 1) { e.barrier(); try a.promptRows(e, scratch, binding, plan, out, split); if (raw_gate) |gate| { diff --git a/zig/src/families/qwen27/attention_chain_test.zig b/zig/src/families/qwen27/attention_chain_test.zig new file mode 100644 index 0000000000..a8417717ad --- /dev/null +++ b/zig/src/families/qwen27/attention_chain_test.zig @@ -0,0 +1,100 @@ +//! Long prompt chains use absolute keys while branching windows retain their bounded ancestor maps. +const std = @import("std"); +const p = @import("attention_plan.zig"); +const frame_rows = @import("gpu_frame.zig").max_rows; +const a = std.testing.allocator; +const pre_m5 = p.Geometry{ .shared_prefix = false }; + +fn parents() [frame_rows]i32 { + var chain: [frame_rows]i32 = undefined; + for (&chain, 0..) |*parent, r| parent.* = @as(i32, @intCast(r)) - 1; + return chain; +} + +test "attention chain: pre-M5 prompt choices follow every admitted frame and request" { + const requested = @import("prompt_rows.zig"); + for ([_]u32{ 13, 14, 15, 16 }) |generation| { + try std.testing.expectEqual(@as(u32, 512), pre_m5.promptRows(requested.requested(generation, false, null))); + for ([_][]const u8{ "256", "512", "1024", "2048" }, [_]u32{ 256, 512, 1024, 2048 }) |text, want| { + try std.testing.expectEqual(want, pre_m5.promptRows(requested.requested(generation, false, text))); + } + } + var frame: u32 = 16; + while (frame <= frame_rows) : (frame += 16) { + for ([_]u32{ 1, 16, 127, 128, 129, 256, 512, 1024, 2048, 4096 }) |request| { + try std.testing.expectEqual(@min(request, frame), pre_m5.promptRows(@min(request, frame))); + } + } + try std.testing.expectEqual(@as(u32, frame_rows), pre_m5.promptRows(4096)); + try std.testing.expectEqual(@as(u32, 128), (p.Geometry{ .heads = 16, .kv_heads = 4, .shared_prefix = false }).promptRows(512)); + try std.testing.expectEqual(@as(u32, 128), (p.Geometry{ .heads = 32, .kv_heads = 2, .shared_prefix = false }).promptRows(512)); + try std.testing.expectEqual(@as(u32, 128), (p.Geometry{ .shared_prefix = false, .fuse_chain = false }).promptRows(512)); +} + +test "attention chain: long tail positions and causal physical keys equal one-row decode" { + const chain = parents(); + for ([_]usize{ 127, 128, 129, 255, 256, 511, 512, 1023, 1024, 2048 }) |rows| { + for ([_]u32{ 0, 63, 64, 127, 128, 511, 512, 8063, 8191, 8192, 8193 }) |start| { + var plan = try p.Plan.init(a, &.{.{ .parents = chain[0..rows], .start = start, .capacity = 16384 }}, pre_m5); + defer plan.deinit(); + try std.testing.expect(!plan.fused); + try std.testing.expectEqual(@as(i32, @intFromBool(rows > p.path_width)), plan.meta[p.global_words + 10]); + for (0..rows) |r| { + const position = start + @as(u32, @intCast(r)); + try std.testing.expectEqual(@as(i32, @intCast(position)), plan.positions[r]); + try std.testing.expectEqual(@as(i32, @intCast(p.chunkFor(position))), plan.nodes[r * p.node_words + 2]); + for ([_]u32{ 0, start, position, position + 1 }) |key| { + const want: ?u32 = if (key <= position) key else null; + try std.testing.expectEqual(want, try plan.physical(r, key)); + } + } + try std.testing.expectEqual(@as(u32, 0), plan.prefix_chunks); + try std.testing.expect(plan.tail_chunks <= 65); + } + } +} + +test "attention chain: keep validates long chain identity without reading past the ancestor map" { + const chain = parents(); + var plan = try p.Plan.init(a, &.{.{ .parents = &chain, .start = 127, .capacity = 4096 }}, pre_m5); + defer plan.deinit(); + var kept: [frame_rows]u32 = undefined; + for (&kept, 0..) |*row, r| row.* = @intCast(r); + for ([_]usize{ 0, 127, 128, 129, 512, 1024, 2048 }) |n| { + const map = try plan.keepRows(a, &.{kept[0..n]}); + defer a.free(map); + for (map, 0..) |row, r| try std.testing.expectEqual(p.KeepRow{ .stream = 0, .source = @intCast(r), .destination = @intCast(r), .row = @intCast(r) }, row); + } + kept[129] = 0; + try std.testing.expectError(error.AttentionKeep, plan.keepRows(a, &.{kept[0..512]})); +} + +test "attention chain: branching, multiple streams and disabled chains keep their old limits" { + var chain = parents(); + const window = p.Window{ .parents = chain[0..129], .start = 0, .capacity = 4096 }; + try std.testing.expectError(error.AttentionCapacity, p.Plan.init(a, &.{window}, .{ .shared_prefix = false, .fuse_chain = false })); + try std.testing.expectError(error.AttentionCapacity, p.Plan.init(a, &.{ window, window }, pre_m5)); + const small = p.Window{ .parents = chain[0..65], .start = 0, .capacity = 4096 }; + try std.testing.expectError(error.AttentionRows, p.Plan.init(a, &.{ small, small }, pre_m5)); + chain[128] = 0; + try std.testing.expectError(error.AttentionCapacity, p.Plan.init(a, &.{window}, pre_m5)); + var short = try p.Plan.init(a, &.{.{ .parents = &.{ -1, 0, 0 }, .start = 511, .capacity = 4096 }}, pre_m5); + defer short.deinit(); + try std.testing.expectEqual(@as(i32, 0), short.meta[p.global_words + 10]); + try std.testing.expectEqual(@as(?u32, null), try short.physical(2, 513)); + try std.testing.expectEqual(@as(?u32, 513), try short.physical(2, 512)); +} + +test "attention chain: every ragged width through the frame limit has the same final causal row" { + const chain = parents(); + const limit = @import("attention_memory.zig").Layout.init(.{ .rows = @as(u32, frame_rows), .streams = @as(u32, 1), .keys = @as(u32, 16384), .shared_prefix = false, .geometry = pre_m5 }); + for (1..frame_rows + 1) |rows| { + var plan = try p.Plan.init(a, &.{.{ .parents = chain[0..rows], .start = 8191, .capacity = 16384 }}, pre_m5); + defer plan.deinit(); + const last: u32 = @intCast(8191 + rows - 1); + try std.testing.expectEqual(@as(?u32, last), try plan.physical(rows - 1, last)); + try std.testing.expectEqual(@as(?u32, 8191), try plan.physical(rows - 1, 8191)); + try std.testing.expectEqual(@as(?u32, null), try plan.physical(rows - 1, last + 1)); + try std.testing.expect(plan.tiles <= limit.tiles and plan.tail_chunks <= limit.tail_chunks); + } +} diff --git a/zig/src/families/qwen27/attention_plan.zig b/zig/src/families/qwen27/attention_plan.zig index b303eb5238..75a950fb94 100644 --- a/zig/src/families/qwen27/attention_plan.zig +++ b/zig/src/families/qwen27/attention_plan.zig @@ -52,7 +52,7 @@ pub const Geometry = struct { if (g.heads / g.kv_heads != 6 and !one_mac) return error.TargetAttentionShape; } pub fn promptRows(g: Geometry, requested: u32) u32 { - return @min(requested, @as(u32, if (g.shared_prefix and g.fuse_chain) chain_rows else max_rows)); + return @min(requested, @as(u32, if (g.fuse_chain and (g.shared_prefix or g.group() == 6)) chain_rows else max_rows)); } pub fn group(g: Geometry) u32 { @@ -83,12 +83,15 @@ pub const Plan = struct { var tile_count: u32 = 0; var ca_max: u32 = 0; var cb_max: u32 = 0; - var fused = geometry.shared_prefix and geometry.fuse_chain; + var chains = geometry.fuse_chain; for (windows) |w| for (w.parents, 0..) |p, r| if (r > 0 and p != @as(i32, @intCast(r)) - 1) { - fused = false; + chains = false; }; + const fused = geometry.shared_prefix and chains; + const wide_tail = !geometry.shared_prefix and geometry.group() == 6 and chains and windows.len == 1; + const row_limit: usize = if (fused or wide_tail) chain_rows else max_rows; for (windows, 0..) |w, st| { - if (w.parents.len == 0 or w.parents.len > @as(usize, if (fused) chain_rows else max_rows) or w.start > w.capacity or @as(u64, w.start) + w.parents.len > w.capacity or w.capacity > std.math.maxInt(i32) / geometry.dim) return error.AttentionCapacity; + if (w.parents.len == 0 or w.parents.len > row_limit or w.start > w.capacity or @as(u64, w.start) + w.parents.len > w.capacity or w.capacity > std.math.maxInt(i32) / geometry.dim) return error.AttentionCapacity; if (w.parents[0] != -1) return error.AttentionParents; var depths: [chain_rows]u32 = undefined; for (w.parents, 0..) |p, r| { @@ -118,13 +121,14 @@ pub const Plan = struct { meta[at + 7] = @intCast(w.capacity * geometry.dim); meta[at + 8] = meta[at + 7]; meta[at + 9] = @intCast(limit); - // words 10 and 11 are spare: a row's chunk and the prefix chunks it merges are its node words 2 and 3 + // Long non-prefix chains address their own keys by position; trees retain the bounded ancestor map. + meta[at + 10] = @intFromBool(wide_tail and w.parents.len > path_width); row_count = try std.math.add(u32, row_count, @intCast(w.parents.len)); tile_count += count; ca_max = @max(ca_max, ca); cb_max = @max(cb_max, cb); } - if (row_count > @as(u32, if (fused) chain_rows else max_rows)) return error.AttentionRows; + if (row_count > row_limit) return error.AttentionRows; meta[0] = @intCast(windows.len); meta[1] = @intCast(ca_max); meta[2] = @intCast(tile_count); @@ -184,6 +188,7 @@ pub const Plan = struct { const depth: u32 = @intCast(p.nodes[node * node_words]); if (logical < start) return logical; if (logical > start + depth) return null; + if (p.meta[global_words + stream * stream_words + 10] != 0) return logical; return start + @as(u32, @intCast(p.paths[node * path_width + logical - start])); } pub fn workspaceBytes(p: Plan) !usize { @@ -205,7 +210,10 @@ pub const Plan = struct { if (path.len > 0) { const last = path[path.len - 1]; if (last >= width or p.nodes[(first + last) * node_words] + 1 != path.len) return error.AttentionKeep; - for (path, 0..) |r, depth| if (r >= width or p.paths[(first + last) * path_width + depth] != r) return error.AttentionKeep; + const chain = p.meta[at + 10] != 0; + for (path, 0..) |r, depth| { + if (r >= width or (if (chain) @as(i32, @intCast(depth)) else p.paths[(first + last) * path_width + depth]) != r) return error.AttentionKeep; + } } count += path.len; } @@ -521,3 +529,7 @@ test "target attention reads every config dimension without widening kernel qual c.rope_dims = 32; try std.testing.expectError(error.TargetAttentionShape, Geometry.fromConfig(c, true)); } + +test { + _ = @import("attention_chain_test.zig"); +} diff --git a/zig/src/qwen27_speed_test.zig b/zig/src/qwen27_speed_test.zig index dd9c9ffccf..aed9ca6fdf 100644 --- a/zig/src/qwen27_speed_test.zig +++ b/zig/src/qwen27_speed_test.zig @@ -45,7 +45,7 @@ test "qwen27 speed: prompt parts fit the actual fused or unfused attention plan" const geometry = planning.Geometry{ .shared_prefix = tensor, .fuse_chain = fuse }; for ([_]u32{ 16, 128, 129, 512, 2048 }) |requested| { const part = geometry.promptRows(requested); - const want = @min(requested, @as(u32, if (tensor and fuse) @import("families/qwen27/gpu_frame.zig").max_rows else planning.max_rows)); + const want = @min(requested, @as(u32, if (fuse) @import("families/qwen27/gpu_frame.zig").max_rows else planning.max_rows)); try std.testing.expectEqual(want, part); for ([_]u32{ 128, 129, 512, 2051 }) |length| { var at: u32 = 0; diff --git a/zig/tests/contract/metal_limits.zig b/zig/tests/contract/metal_limits.zig index b44a632839..ef32c1b7fc 100644 --- a/zig/tests/contract/metal_limits.zig +++ b/zig/tests/contract/metal_limits.zig @@ -31,6 +31,11 @@ fn launchOf(a: std.mem.Allocator, device: mtl.Device, d: limits.MetalLaunch) !La return std.mem.indexOf(u8, s, p) != null; } }.f; + if (std.mem.startsWith(u8, d.name, "qwen27-chain ")) { + const chain = @import("metal_q27_chain.zig"); + const h: @import("metal_q27_attn.zig").Heads = if (has(d.name, " rank ")) .rank else .q27b; + return .{ .decl = d, .function = fun, .here = try chain.source(a, device.tensorUnits(), false, false, h), .pre_m5 = try chain.source(a, false, false, false, h) }; + } if (std.mem.startsWith(u8, d.name, "attention pv ")) { const full = std.mem.indexOf(u8, d.name, " fp32 ") != null; const probe = @import("metal_attention_pv.zig"); diff --git a/zig/tests/contract/metal_main.zig b/zig/tests/contract/metal_main.zig index c053e62ac7..88b152bd2b 100644 --- a/zig/tests/contract/metal_main.zig +++ b/zig/tests/contract/metal_main.zig @@ -493,6 +493,9 @@ pub fn main(init: std.process.Init) !u8 { var nemotron_unroll: @import("metal_nemotron_unroll.zig").Suite = .{ .a = a, .device = device, .queue = queue, .kernels = &kernels }; defer nemotron_unroll.deinit(); try nemotron_unroll.collect(o.only); + var q27_chains: @import("metal_q27_chain.zig").Suite = .{}; + defer q27_chains.deinit(a); + try q27_chains.collect(a, device, queue, &kernels, o.only); // qwen27 attention_io at the 27B's head counts on one Mac (24 over 4) and on a rank of two (12 over 2) var ios: std.ArrayList(*io_k.Io) = .empty; defer { diff --git a/zig/tests/contract/metal_q27_chain.zig b/zig/tests/contract/metal_q27_chain.zig new file mode 100644 index 0000000000..fb06107806 --- /dev/null +++ b/zig/tests/contract/metal_q27_chain.zig @@ -0,0 +1,280 @@ +//! Long canonical prompt chains replace main's one-row tail on identical heads, cache rows and gates. +const std = @import("std"); +const mtl = @import("metal"); +const tf = @import("tensorfold"); +const contract = @import("contract"); +const rig_ = @import("rig.zig"); +const common = @import("metal_q27_attn.zig"); +const oracle = @import("q27_chain_ref.zig"); +const c = contract.spec; +const g = contract.gates; +const q = tf.qwen27; +const attn = q.attention; +const planning = q.attention_plan; +const D = 256; +pub const Reg = struct { heads: common.Heads, gated: bool = false, fp64: bool = false, control: bool = false }; + +fn depth(shape: c.Shape) usize { + return shape.k + shape.rows + D; +} +fn shapes(comptime h: common.Heads, comptime fp64: bool) []const c.Shape { + const n = comptime common.geometry(h, false).heads; + if (fp64) return &.{ .{ .rows = 1, .n = n, .k = 127 }, .{ .rows = 129, .n = n, .k = 0 }, .{ .rows = 3, .n = n, .k = 511 }, .{ .rows = 3, .n = n, .k = 8193 } }; + return &.{ .{ .rows = 1, .n = n, .k = 0 }, .{ .rows = 16, .n = n, .k = 63 }, .{ .rows = 127, .n = n, .k = 64 }, .{ .rows = 128, .n = n, .k = 127 }, .{ .rows = 129, .n = n, .k = 128 }, .{ .rows = 255, .n = n, .k = 511 }, .{ .rows = 256, .n = n, .k = 512 }, .{ .rows = 511, .n = n, .k = 8063 }, .{ .rows = 512, .n = n, .k = 8191 }, .{ .rows = 1023, .n = n, .k = 8192 }, .{ .rows = 1024, .n = n, .k = 8193 }, .{ .rows = 2048, .n = n, .k = 0 } }; +} +pub fn spec(r: Reg) c.Spec { + return switch (r.heads) { + inline else => |h| switch (r.gated) { + inline else => |gated| switch (r.fp64) { + inline else => |fp| switch (r.control) { + inline else => |control| .{ + .name = std.fmt.comptimePrint("{s}metal qwen27 chain {s} {s} {s}", .{ if (control) "control: " else "", comptime common.label(h), if (gated) "gated" else "plain", if (fp) "float64" else "main row bytes" }), + .accum = .f32, + .depth = depth, + .ops = 12, + .mids = if (gated) 1 else 0, + .out = .bf16, + .replaces = if (fp) null else "frozen main one-row tail and merge", + .nan = .{ .stated = "same output classes as main's one-row tail on identical nonfinite inputs" }, + .tiles = .{ .{ .size = 16 }, .{}, .{ .size = 128 } }, + .served = shapes(h, fp), + }, + }, + }, + }, + }; +} + +pub fn source(a: std.mem.Allocator, tensor: bool, baseline: bool, control: bool, heads: common.Heads) ![]u8 { + const src = q.runtime_sources; + const body = if (baseline) try std.mem.concat(a, u8, &.{ src.glue, @embedFile("q27_tail_baseline.metal"), src.attention_io, src.attention_prompt }) else try std.mem.concat(a, u8, &.{ src.glue, src.attention_mpp, src.attention_io, src.attention_prompt }); + defer a.free(body); + const rewritten = if (tensor) try a.dupe(u8, body) else try q.simd_attention.rewrite(a, body); + defer a.free(rewritten); + const from = "? q : P + paths[node * MAXD + (q - P)]"; + if (control and std.mem.count(u8, rewritten, from) != 2) return error.ControlEdit; + const edited = if (control) try std.mem.replaceOwned(u8, a, rewritten, from, "? P : P + paths[node * MAXD + (q - P)]") else try a.dupe(u8, rewritten); + defer a.free(edited); + const gm = common.geometry(heads, false); + return std.fmt.allocPrint(a, "#define Q27_G 6\n#define Q27_KVH {d}\n{s}", .{ gm.kv_heads, edited }); +} + +pub const Case = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + reg: Reg, + op: attn.Attention, + old: attn.Attention, + state: ?State = null, + shape: c.Shape = .{ .rows = 0, .n = 0, .k = 0 }, + fill: g.Fill = .random, + host: [5][]u16 = .{ &.{}, &.{}, &.{}, &.{}, &.{} }, + refs: [6]rig_.Ref = undefined, + own_refs: [2]rig_.Ref = undefined, + caches: [1]attn.Cache = undefined, + baseline: []u8 = &.{}, + const State = struct { + plan: attn.Plan, + scratch: attn.Scratch, + binding: attn.Binding, + fn deinit(s: *State) void { + s.binding.deinit(); + s.scratch.deinit(); + s.plan.deinit(); + } + }; + pub fn init(a: std.mem.Allocator, device: mtl.Device, queue: mtl.Queue, r: Reg) !Case { + var ops: [2]attn.Attention = undefined; + var made: usize = 0; + errdefer for (ops[0..made]) |*op| op.deinit(); + for (&ops, 0..) |*op, i| { + const text = try source(a, device.tensorUnits(), i == 1, r.control and i == 0, r.heads); + defer a.free(text); + const lib = try mtl.Library.fromSource(device, text, mtl.CompileOptions.mlx()); + defer lib.deinit(); + op.* = try attn.Attention.init(device, lib); + made += 1; + } + return .{ .a = a, .rig = .{ .a = a, .device = device, .queue = queue }, .reg = r, .op = ops[0], .old = ops[1] }; + } + pub fn deinit(t: *Case) void { + t.release(); + t.op.deinit(); + t.old.deinit(); + t.rig.deinit(); + } + fn release(t: *Case) void { + if (t.state) |*s| s.deinit(); + t.state = null; + for (&t.host) |*h| { + t.a.free(h.*); + h.* = &.{}; + } + t.a.free(t.baseline); + t.baseline = &.{}; + } + pub fn kernel(t: *Case) g.Kernel { + return .{ .spec = spec(t.reg), .ctx = t, .vt = &g.hooks(Case), .control = if (t.reg.control) .shapes else null }; + } + fn gm(t: *Case) planning.Geometry { + return common.geometry(t.reg.heads, false); + } + fn stateFor(t: *Case, parents: []const i32, start: usize, frame: usize) !State { + var plan = try attn.Plan.init(t.a, &.{.{ .parents = parents, .start = @intCast(start), .capacity = @intCast(t.shape.k + t.shape.rows + 64) }}, t.gm()); + errdefer plan.deinit(); + const limits = attn.Limits{ .rows = @intCast(frame), .streams = 1, .keys = @intCast(t.shape.k + t.shape.rows + 64), .shared_prefix = false, .geometry = t.gm() }; + var scratch = try attn.Scratch.init(t.rig.device, limits); + errdefer scratch.deinit(); + var binding = try attn.Binding.init(t.rig.device, limits); + errdefer binding.deinit(); + try scratch.upload(plan); + try binding.upload(plan, &t.caches); + return .{ .plan = plan, .scratch = scratch, .binding = binding }; + } + pub fn prepare(t: *Case, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + const gm_ = t.gm(); + if (shape.rows == 0 or shape.rows > 2048 or shape.n != gm_.heads or shape.k > 8193) return false; + t.release(); + t.rig.reset(); + t.shape = shape; + t.fill = fill; + const cap = shape.k + shape.rows + 64; + const hk = gm_.kv_heads; + const heads = gm_.heads; + const sizes = [_]usize{ shape.rows * heads * D, hk * cap * D, hk * cap * D, shape.rows * heads * 2 * D, shape.rows * hk * D }; + var rng_ = std.Random.DefaultPrng.init(seed); + const rng = rng_.random(); + for (&t.host, sizes) |*h, n| { + h.* = try t.a.alloc(u16, n); + for (h.*) |*x| x.* = c.toBf16((rng.float(f32) * 2 - 1) * 0.5); + } + if (fill == .stress) { + @memset(t.host[0], 0); + @memset(t.host[1], 0); + @memset(t.host[3], 0); + for (0..hk) |h| for (0..cap) |key| @memset(t.host[2][(h * cap + key) * D ..][0..D], c.toBf16(oracle.stress(key))); + } + if (fill == .nonfinite) { + for ([_]f32{ std.math.nan(f32), std.math.inf(f32), -std.math.inf(f32) }, 0..) |v, i| t.host[0][(i * heads + i) * D + 3] = c.toBf16(v); + } + for (t.host[0..4], 0..) |h, i| t.refs[i] = try t.rig.input(std.mem.sliceAsBytes(h)); + t.refs[4] = try t.rig.output(shape.rows * heads * D * 2); + t.refs[5] = try t.rig.scratch(shape.rows * heads * D * 2); + t.caches[0] = .{ .keys = common.ref(t.refs[1]), .values = common.ref(t.refs[2]), .capacity = @intCast(cap) }; + var parents: [2048]i32 = undefined; + for (parents[0..shape.rows], 0..) |*p, r| p.* = @as(i32, @intCast(r)) - 1; + t.own_refs = .{ try t.own(1), try t.own(2) }; + t.state = try t.stateFor(parents[0..shape.rows], shape.k, shape.rows); + return true; + } + fn own(t: *Case, side: usize) !rig_.Ref { + const gm_ = t.gm(); + const cap = t.shape.k + t.shape.rows + 64; + for (0..t.shape.rows) |r| for (0..gm_.kv_heads) |h| @memcpy(t.host[4][(r * gm_.kv_heads + h) * D ..][0..D], t.host[side][(h * cap + t.shape.k + r) * D ..][0..D]); + return t.rig.input(std.mem.sliceAsBytes(t.host[4])); + } + fn execute(t: *Case, op: *const attn.Attention, state: *State, query: rig_.Ref, key: rig_.Ref, value: rig_.Ref, gate: rig_.Ref, output: rig_.Ref, prompt: bool) !void { + const pass = t.rig.begin(); + op.encode(pass.e, &state.scratch, &state.binding, state.plan, common.ref(query), common.ref(key), common.ref(value), common.ref(output), &t.caches, if (t.reg.gated) common.ref(gate) else null, prompt, false) catch |err| { + try t.rig.finish(pass); + return err; + }; + try t.rig.finish(pass); + } + pub fn run(t: *Case) !void { + const key = t.own_refs[0]; + const value = t.own_refs[1]; + try t.execute(&t.op, &t.state.?, t.refs[0], key, value, t.refs[3], t.refs[4], true); + } + pub fn replaced(t: *Case) ![]const u8 { + t.a.free(t.baseline); + t.baseline = try t.a.alloc(u8, t.shape.rows * t.gm().heads * D * 2); + const key = t.own_refs[0]; + const value = t.own_refs[1]; + const qw = t.gm().heads * D; + const kw = t.gm().kv_heads * D; + for (0..t.shape.rows) |r| { + var state = try t.stateFor(&.{-1}, t.shape.k + r, 1); + defer state.deinit(); + var query = t.refs[0]; + query.off += r * qw * 2; + var k = key; + k.off += r * kw * 2; + var v = value; + v.off += r * kw * 2; + var gate = t.refs[3]; + gate.off += r * qw * 4; + var output = t.refs[5]; + output.off += r * qw * 2; + try t.execute(&t.old, &state, query, k, v, gate, output, false); + } + @memcpy(t.baseline, t.rig.bytes(t.refs[5], t.baseline.len)); + return t.baseline; + } + pub fn count(t: *Case) usize { + return t.shape.rows * t.gm().heads * D; + } + pub fn reference(t: *Case, want: []f64, mag: []f64) !void { + if (t.fill == .stress) return oracle.stressReference(t.shape.k, t.shape.rows, t.gm().heads * D, t.reg.gated, want, mag); + try oracle.reference(t.a, t.host[0], t.host[1], t.host[2], if (t.reg.gated) t.host[3] else null, t.shape.rows, t.gm().heads, t.gm().kv_heads, t.shape.k, t.shape.k + t.shape.rows + 64, want, mag); + } + pub fn bounds(t: *Case, want: []const f64, mag: []const f64, allow: []f64) !void { + for (want, mag, allow) |w, m, *a| a.* = if (t.fill == .stress) @abs(w) * 0x1p-8 + 0x1p-133 else c.bound(spec(t.reg), t.shape, w, m) + 0x1p-11 * m; + } + pub fn simulate(t: *Case, acc: c.Accum, order: g.Order, output: []f64) !void { + try oracle.simulate(t.a, t.shape.k, t.shape.rows, t.gm().heads * D, t.reg.gated, acc, order, output); + } + pub fn check(t: *Case, at: c.Gate) !?[]const u8 { + if (at == .shapes and t.reg.fp64) { + const want = try t.a.alloc(f64, 3 * t.count()); + defer t.a.free(want); + const mag = want[t.count() .. 2 * t.count()]; + const allow = want[2 * t.count() ..]; + try t.reference(want[0..t.count()], mag); + try t.bounds(want[0..t.count()], mag, allow); + const now = try t.rig.outputBytes(); + for (want[0..t.count()], allow, 0..) |w, bound, i| { + if (!(@abs(c.read(.bf16, now, i) - w) <= bound)) return "output leaves independent FP64 attention bound"; + } + } + if (at != .nonfinite) return null; + const old = try t.replaced(); + const now = try t.rig.outputBytes(); + var bad: usize = 0; + for (0..t.count()) |i| { + const cls = c.classOf(c.read(.bf16, old, i)); + bad += @intFromBool(cls != .finite); + if (cls != c.classOf(c.read(.bf16, now, i))) return "output class differs from main's row"; + } + if (bad == 0 or bad == t.count()) return "nonfinite fixture must produce both finite and nonfinite output classes"; + return null; + } +}; + +pub const Suite = struct { + cases: std.ArrayList(*Case) = .empty, + pub fn deinit(s: *Suite, a: std.mem.Allocator) void { + for (s.cases.items) |x| { + x.deinit(); + a.destroy(x); + } + s.cases.deinit(a); + } + pub fn collect(s: *Suite, a: std.mem.Allocator, device: mtl.Device, queue: mtl.Queue, kernels: *std.ArrayList(g.Kernel), only: ?[]const u8) !void { + for (std.enums.values(common.Heads)) |heads| for ([_]Reg{ + .{ .heads = heads }, .{ .heads = heads, .gated = true }, + .{ .heads = heads, .fp64 = true }, .{ .heads = heads, .fp64 = true, .gated = true }, + .{ .heads = heads, .control = true }, + }) |r| { + if (only) |needle| if (std.mem.indexOf(u8, spec(r).name, needle) == null) continue; + const x = try a.create(Case); + errdefer a.destroy(x); + x.* = try Case.init(a, device, queue, r); + s.cases.append(a, x) catch |err| { + x.deinit(); + return err; + }; + try kernels.append(a, x.kernel()); + }; + } +}; diff --git a/zig/tests/contract/q27_chain_ref.zig b/zig/tests/contract/q27_chain_ref.zig new file mode 100644 index 0000000000..70f4302cf8 --- /dev/null +++ b/zig/tests/contract/q27_chain_ref.zig @@ -0,0 +1,103 @@ +//! Independent FP64 causal attention and a cancellation witness for its FP32 value sums. +const std = @import("std"); +const contract = @import("contract"); +const c = contract.spec; +const D = 256; + +pub fn stress(key: usize) f32 { + return if (key < 32) 32768 else if (key < 96) 1 else if (key < 128) -32768 else 0; +} + +pub fn reference(a: std.mem.Allocator, query: []const u16, keys: []const u16, values: []const u16, gate: ?[]const u16, rows: usize, heads: usize, kv: usize, start: usize, cap: usize, want: []f64, mag: []f64) !void { + const scores = try a.alloc(f64, cap); + defer a.free(scores); + const output = try a.alloc(f64, D); + defer a.free(output); + const absolute = try a.alloc(f64, D); + defer a.free(absolute); + for (0..rows) |r| for (0..heads) |h| { + const kh = h / (heads / kv); + const count = start + r + 1; + var top = -std.math.inf(f64); + for (scores[0..count], 0..) |*s, key| { + var sum: f64 = 0; + for (0..D) |d| sum += @as(f64, c.fromBf16(query[(r * heads + h) * D + d])) * c.fromBf16(keys[(kh * cap + key) * D + d]); + s.* = sum / 16; + top = @max(top, s.*); + } + var total: f64 = 0; + @memset(output, 0); + @memset(absolute, 0); + for (scores[0..count], 0..) |s, key| { + const p = @exp(s - top); + total += p; + for (output, absolute, 0..) |*o, *abs, d| { + const term = p * c.fromBf16(values[(kh * cap + key) * D + d]); + o.* += term; + abs.* += @abs(term); + } + } + for (output, 0..) |o, d| { + const at = (r * heads + h) * D + d; + const factor: f64 = if (gate) |g| c.fromBf16(c.toBf16(@floatCast(1 / (1 + @exp(-@as(f64, c.fromBf16(g[(r * heads + h) * 2 * D + D + d]))))))) else 1; + want[at] = o / total * factor; + mag[at] = absolute[d] / total * @abs(factor); + } + }; +} + +pub fn stressReference(start: usize, rows: usize, width: usize, gated: bool, want: []f64, mag: []f64) void { + for (0..rows) |r| { + var sum: f64 = 0; + for (0..start + r + 1) |key| sum += stress(key); + const value = sum / @as(f64, @floatFromInt(start + r + 1)) * (if (gated) @as(f64, 0.5) else 1); + @memset(want[r * width ..][0..width], value); + @memset(mag[r * width ..][0..width], @abs(value)); + } +} + +pub fn simulate(a: std.mem.Allocator, start: usize, rows: usize, width: usize, gated: bool, acc: c.Accum, order: contract.gates.Order, out: []f64) !void { + const terms = try a.alloc(f64, start + rows); + defer a.free(terms); + for (terms, 0..) |*v, key| v.* = stress(key); + for (0..rows) |r| { + const value = contract.rows.sumIn(acc, order, terms[0 .. start + r + 1]) / @as(f64, @floatFromInt(start + r + 1)) * (if (gated) @as(f64, 0.5) else 1); + @memset(out[r * width ..][0..width], value); + } +} + +test "long attention stress separates reduced sums in every contract order" { + var want: [4]f64 = undefined; + var mag: [4]f64 = undefined; + var got: [4]f64 = undefined; + stressReference(8193, 4, 1, false, &want, &mag); + for ([_]contract.gates.Order{ .sequential, .pairwise, .strided }) |order| { + try simulate(std.testing.allocator, 8193, 4, 1, false, .f32, order, &got); + try std.testing.expectEqualSlices(f64, &want, &got); + for ([_]c.Accum{ .f16, .bf16 }) |acc| { + try simulate(std.testing.allocator, 8193, 4, 1, false, acc, order, &got); + try std.testing.expect(!std.math.isFinite(got[0]) or @abs(got[0] - want[0]) > 0.75 * @abs(want[0])); + } + } +} + +test "FP64 chain formula is causal and applies the rounded gate at every output column" { + const rows = 2; + const heads = 6; + const cap = 3; + const query: [rows * heads * D]u16 = @splat(0); + const keys: [cap * D]u16 = @splat(0); + var values: [cap * D]u16 = undefined; + @memset(values[0..D], c.toBf16(2)); + @memset(values[D .. 2 * D], c.toBf16(-2)); + @memset(values[2 * D ..], c.toBf16(std.math.nan(f32))); + const gate: [rows * heads * 2 * D]u16 = @splat(0); + var want: [rows * heads * D]f64 = undefined; + var mag: [rows * heads * D]f64 = undefined; + for ([_]bool{ false, true }) |gated| { + try reference(std.testing.allocator, &query, &keys, &values, if (gated) &gate else null, rows, heads, 1, 0, cap, &want, &mag); + for (want[0 .. heads * D]) |v| try std.testing.expectEqual(@as(f64, if (gated) 1 else 2), v); + for (want[heads * D ..]) |v| try std.testing.expectEqual(@as(f64, 0), v); + for (mag) |v| try std.testing.expectEqual(@as(f64, if (gated) 1 else 2), v); + } +} diff --git a/zig/tests/contract/q27_tail_baseline.metal b/zig/tests/contract/q27_tail_baseline.metal new file mode 100644 index 0000000000..1f5265f0c8 --- /dev/null +++ b/zig/tests/contract/q27_tail_baseline.metal @@ -0,0 +1,395 @@ +// Frozen from zig/kernels/metal/qwen27/attention_mpp.metal at f23603801e5b9bef9a40dd72ba7b783c8e15bda2. + +#include +#include +using namespace metal; +using namespace mpp::tensor_ops; +#ifndef Q27_G +#define Q27_G 6 +#define Q27_KVH 4 +#endif +constant constexpr int G = Q27_G, KVH = Q27_KVH, D = 256, TK = 64, MAXD = 128, MG = 8, MS = 12, NW = 4; +struct Q27AttnReadPtrs { device const bfloat* keys; device const bfloat* values; }; + + +template +kernel void q27_attn_prefix(device const bfloat* Qp [[buffer(0)]], device const Q27AttnReadPtrs* caches [[buffer(1)]], + constant float* scale [[buffer(2)]], device const int* meta [[buffer(3)]], device const int* tile_stream [[buffer(4)]], + device float* PO [[buffer(5)]], device float* PM [[buffer(6)]], device float* PL [[buffer(7)]], device const int* nodes [[buffer(8)]], + uint3 threadgroup_position_in_grid [[threadgroup_position_in_grid]], ushort thread_index_in_simdgroup [[thread_index_in_simdgroup]], + ushort simdgroup_index_in_threadgroup [[simdgroup_index_in_threadgroup]]) { + const ushort lane = thread_index_in_simdgroup; + const ushort sg = simdgroup_index_in_threadgroup; + const int tile = int(threadgroup_position_in_grid.z) * SG + (sg >> 1); + const bool hi = sg & 1; + const uint hk = threadgroup_position_in_grid.x; + const uint c = threadgroup_position_in_grid.y; + const int NCH = meta[1], SGA = meta[2]; + const int st = tile < SGA ? tile_stream[tile] : 0; + const int L = meta[MG + st * MS + 9], NQ = meta[MG + st * MS + 2], S0 = meta[MG + st * MS + 4]; + const int lt = tile - meta[MG + st * MS + 3]; + const int RP = 16 * SGA; + const short qid = lane >> 2; + const short fm = (qid & 4) | ((lane >> 1) & 3); + const short fn = ((qid & 2) | (lane & 1)) * 4; + const int r0 = tile * 16 + fm, r1 = r0 + 8; + const int q0 = lt * 16 + fm, q1 = q0 + 8; + + const int n0 = (q0 < G * NQ) ? min(L, S0 + q0 / G + 1) : 0; + const int n1 = (q1 < G * NQ) ? min(L, S0 + q1 / G + 1) : 0; + threadgroup half Ps[16 * TK]; + threadgroup half* myP = Ps; + if (tile >= SGA || int(c) >= meta[MG + st * MS + 1]) return; + + const device int* own = nodes + NW * meta[MG + st * MS + 6]; + const int ck0 = own[NW * (min(q0, G * NQ - 1) / G) + 2], ck1 = own[NW * (min(q1, G * NQ - 1) / G) + 2]; + const int wa = simd_min(min(ck0, ck1)), wb = simd_max(max(ck0, ck1)); + const device bfloat* Kb = caches[st].keys; + const device bfloat* Vb = caches[st].values; + tensor, tensor_inline> tQ((device bfloat*)Qp + (int64_t)hk * RP * D, dextents(D, RP)); + tensor, tensor_inline> tK((device bfloat*)Kb + (int64_t)hk * meta[MG + st * MS + 7], dextents(D, L), array({1, D})); + tensor, tensor_inline> tV((device bfloat*)Vb + (int64_t)hk * meta[MG + st * MS + 8], dextents(D, L), array({1, D})); + tensor, tensor_inline> tP(myP, dextents(TK, 16)); + constexpr auto dS = matmul2d_descriptor(16, TK, D, false, true, false, matmul2d_descriptor::mode::multiply); + constexpr auto dO = matmul2d_descriptor(16, 128, TK, false, false, false, matmul2d_descriptor::mode::multiply_accumulate); + matmul2d opS; + matmul2d opO; + auto aQ = tQ.slice(0, tile * 16); + auto bV0 = tV.slice(0, 0); + auto O = opO.template get_destination_cooperative_tensor(); + for (int i = 0; i < 64; i++) O[i] = 0.0f; + float m0 = -INFINITY, m1 = -INFINITY, l0 = 0.0f, l1 = 0.0f; + + for (int w = wa;; w = wb) { + const int e0 = ck0 == w ? n0 : 0, e1 = ck1 == w ? n1 : 0; + const int kbeg = int(c) * w; + const int kend = min(kbeg + w, L); + for (int kt = kbeg; kt < kend; kt += TK) { + auto bK = tK.slice(0, kt); + auto S = opS.template get_destination_cooperative_tensor(); + opS.run(aQ, bK, S); + float s[TK / 2]; + for (int i = 0; i < TK / 2; i++) { + const int key = kt + (i >> 3) * 16 + fn + (i & 3); + s[i] = key < ((i & 4) ? e1 : e0) ? S[i] * scale[0] : -INFINITY; + } + float x0 = -INFINITY, x1 = -INFINITY; + for (int i = 0; i < TK / 2; i++) { if (i & 4) x1 = max(x1, s[i]); else x0 = max(x0, s[i]); } + x0 = max(x0, simd_shuffle_xor(x0, 1)); x0 = max(x0, simd_shuffle_xor(x0, 8)); + x1 = max(x1, simd_shuffle_xor(x1, 1)); x1 = max(x1, simd_shuffle_xor(x1, 8)); + const float nm0 = max(m0, x0), nm1 = max(m1, x1); + const float f0 = (x0 == -INFINITY) ? 1.0f : fast::exp(m0 - nm0); + const float f1 = (x1 == -INFINITY) ? 1.0f : fast::exp(m1 - nm1); + float p[TK / 2]; + for (int i = 0; i < TK / 2; i++) p[i] = (s[i] == -INFINITY) ? 0.0f : fast::exp(s[i] - ((i & 4) ? nm1 : nm0)); + float y0 = 0.0f, y1 = 0.0f; + for (int b = 0; b < TK / 16; b++) { + y0 += (p[b * 8] + p[b * 8 + 1]) + (p[b * 8 + 2] + p[b * 8 + 3]); + y1 += (p[b * 8 + 4] + p[b * 8 + 5]) + (p[b * 8 + 6] + p[b * 8 + 7]); + } + y0 += simd_shuffle_xor(y0, 1); y0 += simd_shuffle_xor(y0, 8); + y1 += simd_shuffle_xor(y1, 1); y1 += simd_shuffle_xor(y1, 8); + if (x0 != -INFINITY) { l0 = l0 * f0 + y0; m0 = nm0; } + if (x1 != -INFINITY) { l1 = l1 * f1 + y1; m1 = nm1; } + auto Pc = opS.template get_destination_cooperative_tensor(); + for (int i = 0; i < TK / 2; i++) Pc[i] = half(p[i]); + auto Pin = opO.template get_left_input_cooperative_tensor(Pc); + for (int i = 0; i < 64; i++) O[i] *= (i & 4) ? f1 : f0; + auto bV = tV.slice(hi ? 128 : 0, kt); + opO.run(Pin, bV, O); + } + if (w == wb) break; + } + const int64_t base = ((int64_t)hk * NCH + c) * RP; + + for (int q = 0; q < 16; q++) { + if ((q & 1) ? q1 >= G * NQ : q0 >= G * NQ) continue; + device float* dst = PO + (base + tile * 16 + fm + (q & 1) * 8) * D + (hi ? 128 : 0) + (q >> 1) * 16 + fn; + *(device float4*)dst = float4(O[4 * q], O[4 * q + 1], O[4 * q + 2], O[4 * q + 3]); + } + if (!hi && (lane & 9) == 0) { + if (q0 < G * NQ) { PM[base + r0] = m0; PL[base + r0] = l0; } + if (q1 < G * NQ) { PM[base + r1] = m1; PL[base + r1] = l1; } + } +} +#define Q27_PREFIX(SG) template [[host_name("q27_attn_prefix_" #SG)]] kernel void q27_attn_prefix(device const bfloat*, device const Q27AttnReadPtrs*, constant float*, device const int*, device const int*, device float*, device float*, device float*, device const int*, uint3, ushort, ushort); +Q27_PREFIX(1) +Q27_PREFIX(2) +Q27_PREFIX(3) +Q27_PREFIX(4) +Q27_PREFIX(5) +Q27_PREFIX(6) +Q27_PREFIX(7) +Q27_PREFIX(8) +Q27_PREFIX(9) +Q27_PREFIX(10) +Q27_PREFIX(11) +Q27_PREFIX(12) +Q27_PREFIX(13) +Q27_PREFIX(14) +Q27_PREFIX(15) +Q27_PREFIX(16) + +kernel void q27_attn_tail(device const bfloat* QB [[buffer(0)]], device const Q27AttnReadPtrs* caches [[buffer(1)]], + constant float* scale [[buffer(2)]], device const int* meta [[buffer(3)]], device const int* paths [[buffer(4)]], device const int* nodes [[buffer(5)]], + device const float* POA [[buffer(6)]], device const float* PMA [[buffer(7)]], device const float* PLA [[buffer(8)]], + device float* PO [[buffer(9)]], device float* PM [[buffer(10)]], device float* PL [[buffer(11)]], + uint3 threadgroup_position_in_grid [[threadgroup_position_in_grid]], ushort thread_index_in_simdgroup [[thread_index_in_simdgroup]], + ushort sg [[simdgroup_index_in_threadgroup]], ushort tid [[thread_index_in_threadgroup]]) { + + const ushort lane = thread_index_in_simdgroup; + const uint hk = threadgroup_position_in_grid.x; + const uint cb = threadgroup_position_in_grid.y; + const uint node = threadgroup_position_in_grid.z; + const int NCB = meta[3], W = meta[4], RPA = meta[5], CA = meta[6]; + const int st = nodes[NW * node + 1]; + const int P = meta[MG + st * MS + 4], PT = meta[MG + st * MS + 0], CK = nodes[NW * node + 2]; + if (int(cb) >= meta[MG + st * MS + 5]) return; + const int la = meta[MG + st * MS + 3] * 16 + (int(node) - meta[MG + st * MS + 6]) * G; + const int depth = nodes[NW * node]; + const device bfloat* Kb = caches[st].keys; + const device bfloat* Vb = caches[st].values; + const int nmax = P + depth + 1; + const short qid = lane >> 2; + const short fm = (qid & 4) | ((lane >> 1) & 3); + const short fn = ((qid & 2) | (lane & 1)) * 4; + const int r0 = fm, r1 = fm + 8; + const int n0 = r0 < G ? nmax : 0; + const int n1 = r1 < G ? nmax : 0; + threadgroup half myP[16 * TK]; + threadgroup bfloat KV[32 * D]; + const device bfloat* kbase = (const device bfloat*)Kb + (int64_t)hk * meta[MG + st * MS + 7]; + const device bfloat* vbase = (const device bfloat*)Vb + (int64_t)hk * meta[MG + st * MS + 8]; + const int64_t kstep = D, vstep = D; + tensor, tensor_inline> tQ((device bfloat*)QB + ((int64_t)hk * W + node) * 16 * D, dextents(D, 16)); + tensor, tensor_inline> tK32(KV, dextents(D, 32)); + tensor, tensor_inline> tVh(KV, dextents(128, TK)); + tensor, tensor_inline> tP(myP, dextents(TK, 16)); + constexpr auto dS = matmul2d_descriptor(16, 32, D, false, true, false, matmul2d_descriptor::mode::multiply); + constexpr auto dO = matmul2d_descriptor(16, 128, TK, false, false, false, matmul2d_descriptor::mode::multiply_accumulate); + matmul2d opS; + matmul2d opO; + auto Olo = opO.template get_destination_cooperative_tensor(); + auto Ohi = opO.template get_destination_cooperative_tensor(); + for (int i = 0; i < 64; i++) { Olo[i] = 0.0f; Ohi[i] = 0.0f; } + float m0 = -INFINITY, m1 = -INFINITY, l0 = 0.0f, l1 = 0.0f; + const int c0 = PT / CK; + const int c = c0 + int(cb); + const int kbeg = max(c * CK, PT); + const int kend = min((c + 1) * CK, nmax); + if (sg == 0 && cb == 0 && PT > c0 * CK) { + const int64_t baseA = ((int64_t)hk * CA + c0) * RPA + la; + for (int q = 0; q < 16; q++) { + const int row = fm + (q & 1) * 8; + if (row >= G) continue; + const auto src = POA + (baseA + row) * D + (q >> 1) * 16 + fn; + for (int j = 0; j < 4; j++) { Olo[4 * q + j] = src[j]; Ohi[4 * q + j] = src[128 + j]; } + } + if (r0 < G) { m0 = PMA[baseA + r0]; l0 = PLA[baseA + r0]; } + if (r1 < G) { m1 = PMA[baseA + r1]; l1 = PLA[baseA + r1]; } + } + for (int kt = kbeg; kt < kend; kt += TK) { + float sraw[TK / 2]; + for (int h = 0; h < TK / 32; h++) { + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint e = tid; e < 32 * D / 8; e += 128) { + const int row = int(e) / (D / 8), col = (int(e) % (D / 8)) * 8; + const int q = kt + h * 32 + row; + int phys = -1; + if (q < P) phys = q; + else if (q < nmax) phys = P + paths[node * MAXD + (q - P)]; + ((threadgroup vec*)KV)[e] = phys >= 0 ? *(const device vec*)(kbase + phys * kstep + col) : vec(0); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (sg == 0) { + auto S = opS.template get_destination_cooperative_tensor(); + opS.run(tQ, tK32, S); + for (int i = 0; i < 16; i++) sraw[h * 16 + i] = S[i]; + } + } + if (sg == 0) { + float s[TK / 2]; + for (int i = 0; i < TK / 2; i++) { + const int key = kt + (i >> 3) * 16 + fn + (i & 3); + s[i] = key < ((i & 4) ? n1 : n0) ? sraw[i] * scale[0] : -INFINITY; + } + float x0 = -INFINITY, x1 = -INFINITY; + for (int i = 0; i < TK / 2; i++) { if (i & 4) x1 = max(x1, s[i]); else x0 = max(x0, s[i]); } + x0 = max(x0, simd_shuffle_xor(x0, 1)); x0 = max(x0, simd_shuffle_xor(x0, 8)); + x1 = max(x1, simd_shuffle_xor(x1, 1)); x1 = max(x1, simd_shuffle_xor(x1, 8)); + const float nm0 = max(m0, x0), nm1 = max(m1, x1); + const float f0 = (x0 == -INFINITY) ? 1.0f : fast::exp(m0 - nm0); + const float f1 = (x1 == -INFINITY) ? 1.0f : fast::exp(m1 - nm1); + float p[TK / 2]; + for (int i = 0; i < TK / 2; i++) p[i] = (s[i] == -INFINITY) ? 0.0f : fast::exp(s[i] - ((i & 4) ? nm1 : nm0)); + float y0 = 0.0f, y1 = 0.0f; + for (int b = 0; b < TK / 16; b++) { + y0 += (p[b * 8] + p[b * 8 + 1]) + (p[b * 8 + 2] + p[b * 8 + 3]); + y1 += (p[b * 8 + 4] + p[b * 8 + 5]) + (p[b * 8 + 6] + p[b * 8 + 7]); + } + y0 += simd_shuffle_xor(y0, 1); y0 += simd_shuffle_xor(y0, 8); + y1 += simd_shuffle_xor(y1, 1); y1 += simd_shuffle_xor(y1, 8); + if (x0 != -INFINITY) { l0 = l0 * f0 + y0; m0 = nm0; } + if (x1 != -INFINITY) { l1 = l1 * f1 + y1; m1 = nm1; } + for (int f = 0; f < TK / 16; f++) + for (int i = 0; i < 4; i++) { + myP[fm * TK + f * 16 + fn + i] = half(p[f * 8 + i]); + myP[(fm + 8) * TK + f * 16 + fn + i] = half(p[f * 8 + 4 + i]); + } + for (int i = 0; i < 64; i++) { const float f = (i & 4) ? f1 : f0; Olo[i] *= f; Ohi[i] *= f; } + } + for (int hv = 0; hv < 2; hv++) { + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint e = tid; e < TK * 128 / 8; e += 128) { + const int row = int(e) / 16, col = hv * 128 + (int(e) % 16) * 8; + const int q = kt + row; + int phys = -1; + if (q < P) phys = q; + else if (q < nmax) phys = P + paths[node * MAXD + (q - P)]; + ((threadgroup vec*)KV)[e] = phys >= 0 ? *(const device vec*)(vbase + phys * vstep + col) : vec(0); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (sg == 0) { + if (hv == 0) opO.run(tP, tVh, Olo); + else opO.run(tP, tVh, Ohi); + } + } + } + if (sg != 0) return; + const int64_t base = (((int64_t)hk * NCB + cb) * W + node) * 16; + for (int q = 0; q < 16; q++) { + device float* dst = PO + (base + fm + (q & 1) * 8) * D + (q >> 1) * 16 + fn; + *(device float4*)dst = float4(Olo[4 * q], Olo[4 * q + 1], Olo[4 * q + 2], Olo[4 * q + 3]); + *(device float4*)(dst + 128) = float4(Ohi[4 * q], Ohi[4 * q + 1], Ohi[4 * q + 2], Ohi[4 * q + 3]); + } + if ((lane & 9) == 0) { + PM[base + r0] = m0; PL[base + r0] = l0; + PM[base + r1] = m1; PL[base + r1] = l1; + } +} + + +kernel void q27_attn_merge(device const float* POA [[buffer(0)]], device const float* PMA [[buffer(1)]], device const float* PLA [[buffer(2)]], + device const float* POB [[buffer(3)]], device const float* PMB [[buffer(4)]], device const float* PLB [[buffer(5)]], + device const int* meta [[buffer(6)]], device const int* nodes [[buffer(7)]], device bfloat* OUT [[buffer(8)]], + device const bfloat* QG [[buffer(9)]], constant uint& gated [[buffer(10)]], constant uint& heads [[buffer(11)]], + uint3 threadgroup_position_in_grid [[threadgroup_position_in_grid]], uint thread_index_in_simdgroup [[thread_index_in_simdgroup]]) { + const uint lane = thread_index_in_simdgroup; + const uint hk = threadgroup_position_in_grid.x; + const uint r = threadgroup_position_in_grid.y; + const int NCB = meta[3], W = meta[4], RPA = meta[5], CA = meta[6]; + const int st = nodes[NW * (int(r) / G) + 1]; + const int NCBS = meta[MG + st * MS + 5]; + const int ra = meta[MG + st * MS + 3] * 16 + (int(r) / G - meta[MG + st * MS + 6]) * G + int(r) % G; + const int CT = nodes[NW * (int(r) / G) + 3]; + constexpr int DP = D / 32, MB = 8; + const int node = r / G, g = r % G; + float m = -INFINITY, l = 0.0f, o[DP]; + for (int i = 0; i < DP; i++) o[i] = 0.0f; + for (int c0 = 0; c0 < CT; c0 += MB) { + float mcs[MB], lcs[MB]; + float4 lo[MB], hi[MB]; + for (int j = 0; j < MB; j++) { + const int64_t row = ((int64_t)hk * CA + min(c0 + j, CT - 1)) * RPA + ra; + mcs[j] = PMA[row]; + lcs[j] = PLA[row]; + lo[j] = *(device const float4*)(POA + row * D + lane * DP); + hi[j] = *(device const float4*)(POA + row * D + lane * DP + 4); + } + for (int j = 0; j < MB; j++) { + if (c0 + j >= CT) break; + const float mc = mcs[j]; + if (mc == -INFINITY) continue; + const float lc = lcs[j]; + const float nm = max(m, mc); + const float f1 = fast::exp(m - nm), f2 = fast::exp(mc - nm); + l = l * f1 + lc * f2; + for (int i = 0; i < 4; i++) o[i] = o[i] * f1 + lo[j][i] * f2; + for (int i = 0; i < 4; i++) o[4 + i] = o[4 + i] * f1 + hi[j][i] * f2; + m = nm; + } + } + for (int c = 0; c < NCBS; c++) { + const int64_t row = (((int64_t)hk * NCB + c) * W + node) * 16 + g; + const float mc = PMB[row]; + if (mc == -INFINITY) continue; + const float lc = PLB[row]; + const float nm = max(m, mc); + const float f1 = fast::exp(m - nm), f2 = fast::exp(mc - nm); + l = l * f1 + lc * f2; + for (int i = 0; i < DP; i++) o[i] = o[i] * f1 + POB[row * D + lane * DP + i] * f2; + m = nm; + } + const int h = hk * G + g; + for (int i = 0; i < DP; i++) { + const int64_t index = ((int64_t)node * (int64_t)heads + h) * D + lane * DP + i; + const bfloat v = static_cast(o[i] / l); + if (gated) { + const float gate = float(QG[((int64_t)node * (int64_t)heads + h) * 2 * D + D + lane * DP + i]); + OUT[index] = static_cast(float(v) * q27_round(1.0f / (1.0f + exp(-gate)), 0)); + } else OUT[index] = v; + } +} + + +kernel void q27_attn_merge_single(device const float* POA [[buffer(0)]], device const float* PMA [[buffer(1)]], device const float* PLA [[buffer(2)]], + device const float* POB [[buffer(3)]], device const float* PMB [[buffer(4)]], device const float* PLB [[buffer(5)]], + device const int* meta [[buffer(6)]], device const int* nodes [[buffer(7)]], device bfloat* OUT [[buffer(8)]], + device const bfloat* QG [[buffer(9)]], constant uint& gated [[buffer(10)]], + uint3 threadgroup_position_in_grid [[threadgroup_position_in_grid]], uint thread_index_in_simdgroup [[thread_index_in_simdgroup]]) { + const uint lane = thread_index_in_simdgroup; + const uint hk = threadgroup_position_in_grid.x; + const uint r = threadgroup_position_in_grid.y; + const int NCB = meta[3], W = meta[4], RPA = meta[5], CA = meta[6]; + const int st = nodes[NW * (int(r) / G) + 1]; + const int NCBS = meta[MG + st * MS + 5]; + const int ra = meta[MG + st * MS + 3] * 16 + (int(r) / G - meta[MG + st * MS + 6]) * G + int(r) % G; + const int CT = nodes[NW * (int(r) / G) + 3]; + constexpr int DP = D / 32, MB = 8; + const int node = r / G, g = r % G; + float m = -INFINITY, l = 0.0f, o[DP]; + for (int i = 0; i < DP; i++) o[i] = 0.0f; + for (int c0 = 0; c0 < CT; c0 += MB) { + float mcs[MB], lcs[MB]; + float4 lo[MB], hi[MB]; + for (int j = 0; j < MB; j++) { + const int64_t row = ((int64_t)hk * CA + min(c0 + j, CT - 1)) * RPA + ra; + mcs[j] = PMA[row]; + lcs[j] = PLA[row]; + lo[j] = *(device const float4*)(POA + row * D + lane * DP); + hi[j] = *(device const float4*)(POA + row * D + lane * DP + 4); + } + for (int j = 0; j < MB; j++) { + if (c0 + j >= CT) break; + const float mc = mcs[j]; + if (mc == -INFINITY) continue; + const float lc = lcs[j]; + const float nm = max(m, mc); + const float f1 = fast::exp(m - nm), f2 = fast::exp(mc - nm); + l = l * f1 + lc * f2; + for (int i = 0; i < 4; i++) o[i] = o[i] * f1 + lo[j][i] * f2; + for (int i = 0; i < 4; i++) o[4 + i] = o[4 + i] * f1 + hi[j][i] * f2; + m = nm; + } + } + for (int c = 0; c < NCBS; c++) { + const int64_t row = (((int64_t)hk * NCB + c) * W + node) * 16 + g; + const float mc = PMB[row]; + if (mc == -INFINITY) continue; + const float lc = PLB[row]; + const float nm = max(m, mc); + const float f1 = fast::exp(m - nm), f2 = fast::exp(mc - nm); + l = l * f1 + lc * f2; + for (int i = 0; i < DP; i++) o[i] = o[i] * f1 + POB[row * D + lane * DP + i] * f2; + m = nm; + } + const int h = hk * G + g; + for (int i = 0; i < DP; i++) { + const int64_t index = ((int64_t)node * (G * KVH) + h) * D + lane * DP + i; + const bfloat v = static_cast(o[i] / l); + if (gated) { + const float gate = float(QG[((int64_t)node * (G * KVH) + h) * 2 * D + D + lane * DP + i]); + OUT[index] = static_cast(float(v) * q27_round(1.0f / (1.0f + exp(-gate)), 0)); + } else OUT[index] = v; + } +} From 76074faf04eb528548e02f12f07ecafa67eae98b Mon Sep 17 00:00:00 2001 From: ashhart <71890694+ashhart@users.noreply.github.com> Date: Wed, 4 Nov 2026 19:11:21 +0000 Subject: [PATCH 4/5] Fix pre-M5 chain contract library and hook lifetime --- zig/build/gemma4_host.zig | 5 +++ zig/q27_contract_test.zig | 49 ++++++++++++++++++++++++++ zig/tests/contract/metal_q27_chain.zig | 6 ++-- 3 files changed, 58 insertions(+), 2 deletions(-) create mode 100644 zig/q27_contract_test.zig diff --git a/zig/build/gemma4_host.zig b/zig/build/gemma4_host.zig index 36958cb157..86e68aad49 100644 --- a/zig/build/gemma4_host.zig +++ b/zig/build/gemma4_host.zig @@ -40,6 +40,11 @@ pub fn compileHost(b: *std.Build, metal: *std.Build.Module, engine: *std.Build.M runner.addImport("h3_kernels", b.createModule(.{ .root_source_file = b.path("zig/kernels/metal/h3/kernels.zig") })); runner.addAnonymousImport("ling_metal", .{ .root_source_file = b.path("zig/kernels/metal/ling/sources.zig") }); runner.addAnonymousImport("ling_route", .{ .root_source_file = b.path("zig/src/families/ling/route.zig") }); + const chain_root = b.createModule(.{ .root_source_file = b.path("zig/q27_contract_test.zig"), .target = b.graph.host, .link_libc = true, .imports = &.{ .{ .name = "metal", .module = metal }, .{ .name = "tensorfold", .module = engine }, .{ .name = "contract", .module = @import("kernel_contract.zig").module(b, b.graph.host, .debug) } } }); + const chain_filter = b.option([]const u8, "q27-contract-host-filter", "One chain contract host regression"); + const chain_run = b.addRunArtifact(b.addTest(.{ .root_module = chain_root, .filters = if (chain_filter) |f| &.{f} else &.{"q27 contract:"} })); + chain_run.has_side_effects = true; + b.step("test-q27-contract-host", "Actual chain source, stable hooks and named pipeline refusals without GPU calls").dependOn(&chain_run.step); const contract_root = b.createModule(.{ .root_source_file = b.path("zig/gemma4_contract_compile.zig"), .target = b.graph.host, .link_libc = true, .imports = &.{.{ .name = "runner", .module = runner }} }); const contract_object = b.addObject(.{ .name = "gemma4-contract-compile", .root_module = contract_root }); _ = contract_object.getEmittedBin(); diff --git a/zig/q27_contract_test.zig b/zig/q27_contract_test.zig new file mode 100644 index 0000000000..14f28de2be --- /dev/null +++ b/zig/q27_contract_test.zig @@ -0,0 +1,49 @@ +//! Host-only source and lifecycle checks for the real chain contract; no Metal operation is called. +const std = @import("std"); +const tf = @import("tensorfold"); +const g = @import("contract").gates; +const chain = @import("tests/contract/metal_q27_chain.zig"); +const t = std.testing; + +test "q27 contract: both chain libraries contain the served pre-M5 prompt shader" { + for ([_]bool{ false, true }) |baseline| for ([_]bool{ false, true }) |units| { + const text = try chain.source(t.allocator, units, baseline, false, .q27b); + defer t.allocator.free(text); + const symbol = if (units) "Q27_PROMPT(8)" else "Q27_PROMPT_SG(8)"; + try t.expect(std.mem.indexOf(u8, text, symbol) != null); + try t.expect(std.mem.indexOf(u8, text, "q27_attn_prompt_merge_single") != null); + }; + const text = try tf.qwen27.simd_attention.library(t.allocator, false, true); + defer t.allocator.free(text); + try t.expect(std.mem.indexOf(u8, text, "Q27_PROMPT_SG(8)") != null); +} + +noinline fn fromAnotherFrame(x: *chain.Case) g.Kernel { + return x.kernel(); +} + +test "q27 contract: returned hooks survive the registering stack frame" { + var x: chain.Case = undefined; + x.reg = .{ .heads = .q27b }; + x.shape = .{ .rows = 3, .n = 24, .k = 0 }; + const k = fromAnotherFrame(&x); + const expected = comptime g.hooks(chain.Case); + try t.expectEqualDeep(expected, k.vt.*); + try t.expectEqual(@as(usize, 3 * 24 * 256), k.vt.count(k.ctx)); +} + +test "q27 contract: missing prompt pipelines become named gate failures without Metal calls" { + var x: chain.Case = undefined; + x.rig = .{ .a = t.allocator, .device = undefined, .queue = undefined }; + defer x.rig.deinit(); + x.reg = .{ .heads = .q27b }; + x.op.prompt = null; + x.old.prompt = null; + const shape = @import("contract").spec.Shape{ .rows = 16, .n = 24, .k = 0 }; + try t.expectError(error.Qwen27ChainPromptUnavailable, x.prepare(shape, .random, 1)); + var tally: g.Tally = .{}; + try g.runAll(t.allocator, &.{x.kernel()}, 1, &tally); + try t.expectEqual(@as(usize, 1), tally.kernels); + try t.expectEqual(@as(usize, 1), tally.failed); + try t.expectEqual(@as(usize, 0), tally.passed); +} diff --git a/zig/tests/contract/metal_q27_chain.zig b/zig/tests/contract/metal_q27_chain.zig index fb06107806..89e471437c 100644 --- a/zig/tests/contract/metal_q27_chain.zig +++ b/zig/tests/contract/metal_q27_chain.zig @@ -47,7 +47,7 @@ pub fn spec(r: Reg) c.Spec { pub fn source(a: std.mem.Allocator, tensor: bool, baseline: bool, control: bool, heads: common.Heads) ![]u8 { const src = q.runtime_sources; - const body = if (baseline) try std.mem.concat(a, u8, &.{ src.glue, @embedFile("q27_tail_baseline.metal"), src.attention_io, src.attention_prompt }) else try std.mem.concat(a, u8, &.{ src.glue, src.attention_mpp, src.attention_io, src.attention_prompt }); + const body = if (baseline) try std.mem.concat(a, u8, &.{ src.glue, @embedFile("q27_tail_baseline.metal"), src.attention_io, src.attention_prompt, if (tensor) "" else src.attention_prompt_sg }) else try std.mem.concat(a, u8, &.{ src.glue, src.attention_mpp, src.attention_io, src.attention_prompt, if (tensor) "" else src.attention_prompt_sg }); defer a.free(body); const rewritten = if (tensor) try a.dupe(u8, body) else try q.simd_attention.rewrite(a, body); defer a.free(rewritten); @@ -114,8 +114,9 @@ pub const Case = struct { t.baseline = &.{}; } pub fn kernel(t: *Case) g.Kernel { - return .{ .spec = spec(t.reg), .ctx = t, .vt = &g.hooks(Case), .control = if (t.reg.control) .shapes else null }; + return .{ .spec = spec(t.reg), .ctx = t, .vt = &vt, .control = if (t.reg.control) .shapes else null }; } + const vt = g.hooks(Case); fn gm(t: *Case) planning.Geometry { return common.geometry(t.reg.heads, false); } @@ -134,6 +135,7 @@ pub const Case = struct { pub fn prepare(t: *Case, shape: c.Shape, fill: g.Fill, seed: u64) !bool { const gm_ = t.gm(); if (shape.rows == 0 or shape.rows > 2048 or shape.n != gm_.heads or shape.k > 8193) return false; + if (t.op.prompt == null or t.old.prompt == null) return error.Qwen27ChainPromptUnavailable; t.release(); t.rig.reset(); t.shape = shape; From f57d2a4320bae3d5564bf91298aad5ea097f1782 Mon Sep 17 00:00:00 2001 From: ashhart <71890694+ashhart@users.noreply.github.com> Date: Fri, 6 Nov 2026 00:21:35 +0000 Subject: [PATCH 5/5] tests: the chain contract finds the prompt shaders by their new instantiations, since main's masked variants --- zig/q27_contract_test.zig | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/zig/q27_contract_test.zig b/zig/q27_contract_test.zig index 14f28de2be..47baf6410a 100644 --- a/zig/q27_contract_test.zig +++ b/zig/q27_contract_test.zig @@ -9,13 +9,13 @@ test "q27 contract: both chain libraries contain the served pre-M5 prompt shader for ([_]bool{ false, true }) |baseline| for ([_]bool{ false, true }) |units| { const text = try chain.source(t.allocator, units, baseline, false, .q27b); defer t.allocator.free(text); - const symbol = if (units) "Q27_PROMPT(8)" else "Q27_PROMPT_SG(8)"; + const symbol = if (units) "Q27_PROMPT(8, false," else "Q27_PROMPT_SG(8, false,"; try t.expect(std.mem.indexOf(u8, text, symbol) != null); try t.expect(std.mem.indexOf(u8, text, "q27_attn_prompt_merge_single") != null); }; const text = try tf.qwen27.simd_attention.library(t.allocator, false, true); defer t.allocator.free(text); - try t.expect(std.mem.indexOf(u8, text, "Q27_PROMPT_SG(8)") != null); + try t.expect(std.mem.indexOf(u8, text, "Q27_PROMPT_SG(8, false,") != null); } noinline fn fromAnotherFrame(x: *chain.Case) g.Kernel {