diff --git a/zig/src/families/qwen27/decode_round.zig b/zig/src/families/qwen27/decode_round.zig index d558bdee80..ac05d71c7f 100644 --- a/zig/src/families/qwen27/decode_round.zig +++ b/zig/src/families/qwen27/decode_round.zig @@ -17,6 +17,11 @@ const Round = @import("round_plan.zig").Round; const Taps = @import("taps.zig").Taps; const image_rows = @import("image_rows.zig"); const mrope = @import("core").qwen_vl.rope; +const kv_grow = @import("kv_grow.zig"); +const kv_slots = @import("kv_slots.zig"); + +/// The host's hook for buffers a slot's K/V growth made. +pub const Grown = struct { ctx: *anyopaque, added: *const fn (ctx: *anyopaque, buffers: []const mtl.Buffer) void }; /// Whether `cb` failed on the GPU; its Metal error (out of memory, a timeout) is logged with `what` it ran. pub fn gpuFailed(cb: mtl.CommandBuffer, what: []const u8) bool { @@ -43,6 +48,10 @@ pub const Runner = struct { slots: u32, ready: u32, // slots [0, ready) have their K/V rows; a later slot takes its own when first used (ensure) capacity: u32, + /// TF_Q27_KV_GROW: slots hold the rows their streams reached, grown by this step; null: the whole context. + kv_step: ?u32 = null, + /// A grown slot's new buffers, for the host's residency set (committed now, as a new slot's are). + grown: ?Grown = null, active: ?[32]u8 = null, failed: bool = false, capture_taps: bool = false, @@ -72,6 +81,11 @@ pub const Runner = struct { /// A runner of `slots` whose first `ready` take their K/V rows now, the rest when first used (ensure). pub fn initReady(a: std.mem.Allocator, m: *Model, slots: u32, capacity: u32, ready: u32) !Runner { + return initGrow(a, m, slots, capacity, ready, null); + } + + /// initReady whose slots start at `kv_step` rows and grow as their streams do (reserve); null: the whole context. + pub fn initGrow(a: std.mem.Allocator, m: *Model, slots: u32, capacity: u32, ready: u32, kv_step: ?u32) !Runner { if (slots == 0 or slots > attn.planning.max_streams or capacity == 0 or m.frame.capacity % 16 != 0 or ready == 0 or ready > slots) return error.BadRoundCapacity; const geometry = try attn.planning.Geometry.fromConfig(m.config, m.device.tensorUnits()); const linear_index = try a.alloc(u32, m.config.layers); @@ -130,14 +144,15 @@ pub const Runner = struct { for (owned.items) |b| b.deinit(); owned.deinit(a); } + const first = if (kv_step) |s| kv_grow.first(s, capacity) else capacity; for (0..na) |layer| for (0..ready) |slot| { - caches[layer * slots + slot] = try slotCache(a, m, &owned, capacity); + caches[layer * slots + slot] = try slotCache(a, m, &owned, first); }; const offsets = try a.alloc(u32, slots); errdefer a.free(offsets); @memset(offsets, 0); const taps = try Taps.init(m.device, m.frame.capacity, @intCast(m.config.hidden)); - return .{ .allocator = a, .model = m, .geometry = geometry, .gdn = pool, .attention = attention, .scratch = scratch, .bindings = bindings, .caches = caches, .linear_index = linear_index, .attention_index = attention_index, .offsets = offsets, .owned = owned, .taps = taps, .slots = slots, .ready = ready, .capacity = capacity }; + return .{ .allocator = a, .model = m, .geometry = geometry, .gdn = pool, .attention = attention, .scratch = scratch, .bindings = bindings, .caches = caches, .linear_index = linear_index, .attention_index = attention_index, .offsets = offsets, .owned = owned, .taps = taps, .slots = slots, .ready = ready, .capacity = capacity, .kv_step = kv_step }; } /// One slot's K/V rows across the attention layers. @@ -155,12 +170,50 @@ pub const Runner = struct { r.owned.shrinkRetainingCapacity(first); } const layers = r.caches.len / r.slots; - for (0..layers) |layer| r.caches[layer * r.slots + slot] = try slotCache(r.allocator, r.model, &r.owned, r.capacity); + const rows = if (r.kv_step) |s| kv_grow.first(s, r.capacity) else r.capacity; + for (0..layers) |layer| r.caches[layer * r.slots + slot] = try slotCache(r.allocator, r.model, &r.owned, rows); r.ready += 1; if (r.residency) |f| for (r.owned.items[first..]) |b| f.join(f.ctx, b); return r.owned.items[first..]; } + /// The K/V rows slot `slot` holds now: the fewest any attention layer holds. + pub fn heldRows(r: *const Runner, slot: u32) u32 { + var rows: u32 = std.math.maxInt(u32); + for (0..r.caches.len / r.slots) |layer| rows = @min(rows, r.caches[layer * r.slots + slot].capacity); + return rows; + } + + /// Slot `slot` holds `need` K/V rows, grown layer by layer (kv_slots.zig). + pub fn reserve(r: *Runner, slot: u32, need: u32) !void { + return kv_slots.reserve(r, slot, need); + } + + /// Slot `slot` back to its first rows once it is free and unread (kv_slots.zig). + pub fn release(r: *Runner, slot: u32) !void { + return kv_slots.release(r, slot); + } + + /// Emptied slot `slot` holds the rows a prompt of `need` tokens fills, without a copy (kv_slots.zig). + pub fn admit(r: *Runner, slot: u32, need: u32) !void { + return kv_slots.admit(r, slot, need); + } + + /// Whether a blit still reads slot `slot` (without waiting for it). + pub fn busy(r: *const Runner, slot: u32) bool { + const cb = r.reading[slot] orelse return false; + return switch (cb.status()) { + .completed, .@"error" => false, + else => true, + }; + } + + /// settleSlot for a hook that holds the runner as an opaque pointer (a drafter context's ring). + pub fn settleOpaque(ctx: *anyopaque, slot: u32) void { + const r: *Runner = @ptrCast(@alignCast(ctx)); + r.settleSlot(slot); + } + /// The K/V buffers of the ready slots (for a residency set). pub fn readyBuffers(r: *const Runner, a: std.mem.Allocator, out: *std.ArrayList(mtl.Buffer)) !void { for (0..r.caches.len / r.slots) |layer| for (0..r.ready) |slot| { @@ -221,6 +274,7 @@ pub const Runner = struct { // A rank's head writes only this rank's vocabulary columns; the tail's matcher reads the whole vocabulary. if (r.tail != null and r.reduce != null) return error.TailOnRank; for (round.slots, round.windows) |slot, window| if (slot >= r.ready or r.offsets[slot] != window.start) return error.CachePositionDiffers; + for (round.slots, round.windows) |slot, window| try r.reserve(slot, window.start + @as(u32, @intCast(window.parents.len))); const paths = try r.allocator.alloc([]const u32, round.windows.len); defer r.allocator.free(paths); @memset(paths, &.{}); @@ -442,7 +496,8 @@ pub const Runner = struct { fn attentionPlan(r: *Runner, round: *const Round) !attn.Plan { var windows: [attn.planning.max_streams]attn.planning.Window = undefined; if (round.windows.len > windows.len) return error.AttentionStreams; - for (round.windows, 0..) |window, i| windows[i] = .{ .parents = window.parents, .start = window.start, .capacity = r.capacity }; + // each slot's held rows: the kernels' head stride (equal to the context unless K/V grows as used) + for (round.windows, round.slots, 0..) |window, slot, i| windows[i] = .{ .parents = window.parents, .start = window.start, .capacity = r.heldRows(slot) }; return attn.Plan.init(r.allocator, windows[0..round.windows.len], r.geometry); } diff --git a/zig/src/families/qwen27/dflash/caches.zig b/zig/src/families/qwen27/dflash/caches.zig index f6fefc31fc..450ce65594 100644 --- a/zig/src/families/qwen27/dflash/caches.zig +++ b/zig/src/families/qwen27/dflash/caches.zig @@ -4,7 +4,7 @@ const cfg = @import("config.zig"); pub fn Pair(comptime Buffer: type) type { return struct { keys: Buffer, values: Buffer }; } -/// `layers` rolling K/V caches of the window plus one block, as a stream's drafter context holds them. +/// `layers` rolling K/V caches of each layer's window plus one block, as a stream's drafter context holds them. pub fn newCaches(comptime Buffer: type, a: std.mem.Allocator, device: anytype, c: cfg.Config, layers: usize) ![]Pair(Buffer) { const Cache = Pair(Buffer); const caches = try a.alloc(Cache, layers); @@ -14,8 +14,8 @@ pub fn newCaches(comptime Buffer: type, a: std.mem.Allocator, device: anytype, c cache.keys.deinit(); cache.values.deinit(); }; - for (caches) |*cache| { - const bytes = @as(usize, c.window + 16) * c.kvWidth() * 2; + for (caches, 0..) |*cache, i| { + const bytes = @as(usize, c.layerWindow(i) + 16) * c.kvWidth() * 2; const keys = try device.buffer(bytes, @import("metal").ResourceOptions.shared | @import("metal").ResourceOptions.untracked); const values = device.buffer(bytes, @import("metal").ResourceOptions.shared | @import("metal").ResourceOptions.untracked) catch |err| { keys.deinit(); diff --git a/zig/src/families/qwen27/dflash/config.zig b/zig/src/families/qwen27/dflash/config.zig index 19f6a39c5b..5da1567fb3 100644 --- a/zig/src/families/qwen27/dflash/config.zig +++ b/zig/src/families/qwen27/dflash/config.zig @@ -2,8 +2,13 @@ const std = @import("std"); const target_config = @import("../config.zig"); pub const repository = "z-lab/Qwen3.8-27B-DFlash2"; -pub const tap_ids = [5]u32{ 5, 19, 33, 47, 61 }; +/// The stock DFlash2 taps, padded to the eight a DFlash v1 draft may name. +pub const tap_ids = [8]u32{ 5, 19, 33, 47, 61, 0, 0, 0 }; +pub const max_taps = 8; +/// DFlash2 (conv kernels and a candidate selector) or z-lab's DFlash v1 (a plain qwen3 body, taps fused by fc). +pub const Variant = enum { dflash2, v1 }; pub const Config = struct { + variant: Variant = .dflash2, hidden: u32 = 5120, intermediate: u32 = 17408, layers: u32 = 5, @@ -21,7 +26,24 @@ pub const Config = struct { topk: u32 = 16, eps: f32 = 1e-6, theta: f64 = 10000000, - taps: [5]u32 = tap_ids, + taps: [8]u32 = tap_ids, + tap_count: u32 = 5, + /// v1: the layers whose layer_types entry is full_attention (DFlash2 layers are all sliding). + full: [8]bool = @splat(false), + /// v1: the positions a stream's drafter context keeps (its full layers see all of them; 0: the window). + history: u32 = 0, + /// Positions kept: the ring's rows and the full layers' context; sliding layers see the last `window` - 1. + pub fn span(c: Config) u32 { + return if (c.variant == .v1 and c.history > c.window) c.history else c.window; + } + /// Rows layer `i` attends behind its block: its own window, or the whole span for a full layer. + pub fn layerWindow(c: Config, i: usize) u32 { + return if (c.full[i]) c.span() else c.window; + } + /// A ring row: v1 keeps fused, normed taps (fc and hidden_norm ran in the taps pass); DFlash2 keeps raw taps. + pub fn ringWidth(c: Config) u32 { + return if (c.variant == .v1) c.hidden else c.tapWidth(); + } pub fn qWidth(c: Config) u32 { return c.heads * c.head_dim; } @@ -29,7 +51,7 @@ pub const Config = struct { return c.kv_heads * c.head_dim; } pub fn tapWidth(c: Config) u32 { - return c.hidden * 5; + return c.hidden * c.tap_count; } pub fn dynamicWidth(c: Config) u32 { return 4 * (c.hidden / c.group); @@ -37,30 +59,54 @@ pub const Config = struct { pub fn depth(c: Config) u32 { return c.block - 1; } + /// DFlash2: 15 a layer and six globals; v1: 11 a layer beside fc, hidden_norm and norm. pub fn tensorCount(c: Config) usize { - return 15 * @as(usize, c.layers) + 6; + return if (c.variant == .v1) 11 * @as(usize, c.layers) + 3 else 15 * @as(usize, c.layers) + 6; } + /// DFlash2: seven linears and two conv projections a layer, fc and the selector's projection; v1: seven and fc. pub fn linearCount(c: Config) usize { - return 9 * @as(usize, c.layers) + 2; + return if (c.variant == .v1) 7 * @as(usize, c.layers) + 1 else 9 * @as(usize, c.layers) + 2; } pub fn check(c: Config) !void { + if (c.variant == .v1) return c.checkV1(); const dimensions = [_]u32{ c.hidden, c.intermediate, c.layers, c.vocab, c.heads, c.kv_heads, c.head_dim, c.target_layers, c.group, c.rank, c.topk }; for (dimensions) |n| if (n == 0) return error.BadDraftConfig; if (c.hidden % 64 != 0 or c.intermediate % 64 != 0 or c.hidden % c.group != 0 or c.heads % c.kv_heads != 0 or c.head_dim % 2 != 0 or c.block != 8 or c.conv != 2 or c.layers > 8 or c.mask >= c.vocab or c.topk > c.vocab or c.window < c.block) return error.BadDraftConfig; if (!std.math.isFinite(c.eps) or c.eps <= 0 or !std.math.isFinite(c.theta) or c.theta <= 0) return error.BadDraftConfig; if (c.target_layers == 64 and !std.mem.eql(u32, &c.taps, &tap_ids)) return error.UnsupportedDraftTaps; - for (c.taps, 0..) |tap, i| for (c.taps[0..i]) |previous| if (tap == previous) return error.UnsupportedDraftTaps; - for (c.taps) |i| if (i >= c.target_layers) return error.BadDraftConfig; + if (c.tap_count != 5) return error.UnsupportedDraftTaps; + for (c.taps[0..5], 0..) |tap, i| for (c.taps[0..i]) |previous| if (tap == previous) return error.UnsupportedDraftTaps; + for (c.taps[0..5]) |i| if (i >= c.target_layers) return error.BadDraftConfig; const q = try std.math.mul(u32, c.heads, c.head_dim); _ = try std.math.mul(u32, c.kv_heads, c.head_dim); _ = try std.math.mul(u32, c.hidden, 5); _ = try std.math.mul(u32, c.hidden / c.group, 4); if (q % 64 != 0) return error.BadDraftConfig; } + /// The drafters name in the serving and startup lines: DFlash for z-labs v1 layout, DFlash2 for its successor. + pub fn kind(c: Config) []const u8 { + return if (c.variant == .v1) "DFlash" else "DFlash2"; + } /// A draft reads one target's taps and vocabulary: its width, vocabulary and tapped depth. pub fn bind(c: Config, hidden: usize, vocab: usize, layers: usize) !void { if (hidden != c.hidden or vocab != c.vocab or layers != c.target_layers) return error.DraftForAnotherTarget; } + /// v1 has no conv kernels and no selector (group, conv and rank zero); its lattice keeps the head's top-k. + fn checkV1(c: Config) !void { + const dimensions = [_]u32{ c.hidden, c.intermediate, c.layers, c.vocab, c.heads, c.kv_heads, c.head_dim, c.target_layers, c.topk, c.tap_count }; + for (dimensions) |n| if (n == 0) return error.BadDraftConfig; + if (c.group != 0 or c.conv != 0 or c.rank != 0) return error.BadDraftConfig; + if (c.hidden % 64 != 0 or c.intermediate % 64 != 0 or c.heads % c.kv_heads != 0 or c.head_dim % 2 != 0 or c.block < 2 or c.block > 16 or c.layers > 8 or c.tap_count > max_taps or c.mask >= c.vocab or c.topk > 16 or c.topk > c.vocab or c.window < c.block) return error.BadDraftConfig; + if (!std.math.isFinite(c.eps) or c.eps <= 0 or !std.math.isFinite(c.theta) or c.theta <= 0) return error.BadDraftConfig; + for (c.taps[0..c.tap_count]) |i| if (i >= c.target_layers) return error.BadDraftConfig; + for (c.taps[0..c.tap_count], 0..) |tap, i| for (c.taps[0..i]) |previous| if (tap == previous) return error.UnsupportedDraftTaps; + for (c.taps[c.tap_count..]) |unused| if (unused != 0) return error.BadDraftConfig; + for (c.full[c.layers..]) |unused| if (unused) return error.BadDraftConfig; + const q = try std.math.mul(u32, c.heads, c.head_dim); + _ = try std.math.mul(u32, c.kv_heads, c.head_dim); + _ = try std.math.mul(u32, c.hidden, c.tap_count); + if (q % 64 != 0) return error.BadDraftConfig; + } pub fn requireStock(c: Config) !void { try c.check(); if (c.hidden != 5120 or c.intermediate != 17408 or c.layers != 5 or c.vocab != 248320 or c.heads != 32 or c.kv_heads != 8 or c.head_dim != 128 or c.target_layers != 64 or c.mask != 248070 or c.window != 2048 or c.group != 16 or c.rank != 256 or c.topk != 16 or c.eps != 1e-6 or c.theta != 10000000) return error.NotStockDFlash2; @@ -93,6 +139,7 @@ pub fn parse(gpa: std.mem.Allocator, raw: []const u8) !Config { const o = try object(p.value); const d = try object(o.get("dflash_config") orelse return error.BadDraftConfig); const architectures = o.get("architectures") orelse return error.BadDraftConfig; + if (architectures == .array and architectures.array.items.len == 1 and architectures.array.items[0] == .string and std.mem.eql(u8, architectures.array.items[0].string, "DFlashDraftModel")) return parseV1(o, d); if (architectures != .array or architectures.array.items.len != 1 or architectures.array.items[0] != .string or !std.mem.eql(u8, architectures.array.items[0].string, "DFlash2DraftModel")) return error.BadDraftConfig; try text(o, "model_type", "qwen3"); try text(o, "dtype", "bfloat16"); @@ -119,6 +166,41 @@ pub fn parse(gpa: std.mem.Allocator, raw: []const u8) !Config { try c.requireStock(); return c; } +/// z-lab's DFlash v1 layout (DFlashDraftModel): its taps, block, mask and layer kinds from the config. +fn parseV1(o: std.json.ObjectMap, d: std.json.ObjectMap) !Config { + try text(o, "model_type", "qwen3"); + try text(o, "dtype", "bfloat16"); + try text(o, "hidden_act", "silu"); + for ([_][]const u8{ "attention_bias", "tie_word_embeddings" }) |key| { + const v = o.get(key) orelse return error.BadDraftConfig; + if (v != .bool or v.bool) return error.BadDraftConfig; + } + // the draft attends to its whole block both ways; a config that says otherwise is not this reading + if (o.get("is_causal")) |v| if (v != .bool or v.bool) return error.BadDraftConfig; + for ([_][]const u8{ "conv_group_size", "conv_kernel_size", "selector_rank", "selector_top_k" }) |key| if (d.get(key) != null) return error.BadDraftConfig; + const sliding = o.get("use_sliding_window") orelse return error.BadDraftConfig; + if (sliding != .bool or !sliding.bool or try number(o, "attention_dropout") != 0) return error.BadDraftConfig; + const rope = try object(o.get("rope_parameters") orelse return error.BadDraftConfig); + try text(rope, "rope_type", "default"); + if (o.get("rope_scaling")) |v| if (v != .null) return error.BadDraftConfig; + var c = Config{ .variant = .v1, .hidden = try integer(o, "hidden_size"), .intermediate = try integer(o, "intermediate_size"), .layers = try integer(o, "num_hidden_layers"), .vocab = try integer(o, "vocab_size"), .heads = try integer(o, "num_attention_heads"), .kv_heads = try integer(o, "num_key_value_heads"), .head_dim = try integer(o, "head_dim"), .target_layers = try integer(o, "num_target_layers"), .block = try integer(d, "block_size"), .mask = try integer(d, "mask_token_id"), .window = try integer(o, "sliding_window"), .group = 0, .conv = 0, .rank = 0, .topk = 16, .eps = @floatCast(try number(o, "rms_norm_eps")), .theta = try number(rope, "rope_theta"), .taps = @splat(0), .tap_count = 0 }; + const taps = d.get("target_layer_ids") orelse return error.BadDraftConfig; + if (taps != .array or taps.array.items.len == 0 or taps.array.items.len > max_taps) return error.BadDraftConfig; + for (taps.array.items, 0..) |v, i| { + if (v != .integer or v.integer < 0 or v.integer > std.math.maxInt(u32)) return error.BadDraftConfig; + c.taps[i] = @intCast(v.integer); + } + c.tap_count = @intCast(taps.array.items.len); + if (c.layers == 0 or c.layers > 8) return error.BadDraftConfig; + const types = o.get("layer_types") orelse return error.BadDraftConfig; + if (types != .array or types.array.items.len != c.layers) return error.BadDraftConfig; + for (types.array.items, 0..) |v, i| { + if (v != .string) return error.BadDraftConfig; + if (std.mem.eql(u8, v.string, "full_attention")) c.full[i] = true else if (!std.mem.eql(u8, v.string, "sliding_attention")) return error.BadDraftConfig; + } + try c.check(); + return c; +} /// --drafter's config against the target's, from the two folders' config.json: a refusal comes before either loads. pub fn admit(gpa: std.mem.Allocator, io: std.Io, target_dir: []const u8, draft_dir: []const u8) !Config { @@ -236,3 +318,89 @@ test "a 9B target refuses the 27B's DFlash2 and a 9B-shaped DFlash2 by name from } try std.testing.expectError(error.NotStockDFlash2, admitFolders(nine, small)); } +// z-lab/Qwen3.5-9B-DFlash's config.json, as published (Ornith-1.5-9B-DFlash publishes the same file). +const v1_fixture = + \\{"architectures":["DFlashDraftModel"],"attention_bias":false,"attention_dropout":0.0, + \\"auto_map":{"AutoModel":"dflash.DFlashDraftModel"},"bos_token_id":null, + \\"dflash_config":{"block_size":16,"mask_token_id":248077,"target_layer_ids":[1,5,9,13,17,21,25,29]}, + \\"dtype":"bfloat16","eos_token_id":248044,"head_dim":128,"hidden_act":"silu","hidden_size":4096, + \\"initializer_range":0.02,"intermediate_size":12288, + \\"layer_types":["sliding_attention","sliding_attention","sliding_attention","sliding_attention","sliding_attention","full_attention"], + \\"max_position_embeddings":262144,"max_window_layers":6,"model_type":"qwen3","num_attention_heads":32, + \\"num_hidden_layers":6,"num_key_value_heads":8,"num_target_layers":32,"pad_token_id":null,"rms_norm_eps":1e-06, + \\"rope_parameters":{"rope_theta":10000000,"rope_type":"default"},"sliding_window":4096,"tie_word_embeddings":false, + \\"transformers_version":"5.7.0","use_cache":true,"use_sliding_window":true,"vocab_size":248320} +; +test "the 9B target takes its own DFlash v1 drafter, and the 27B target refuses it by name" { + const a = std.testing.allocator; + const nine = try targetJson(a, 4096, 12288, 32, 16, 32); + defer a.free(nine); + const big = try targetJson(a, 5120, 17408, 64, 24, 48); + defer a.free(big); + try std.testing.expectEqual(Variant.v1, (try admitFolders(nine, v1_fixture)).variant); + try std.testing.expectError(error.DraftForAnotherTarget, admitFolders(big, v1_fixture)); +} +test "a DFlash v1 config reads its own taps, layer kinds, block and window" { + const c = try parse(std.testing.allocator, v1_fixture); + try std.testing.expectEqual(Variant.v1, c.variant); + try std.testing.expectEqual(@as(u32, 8), c.tap_count); + try std.testing.expectEqualSlices(u32, &.{ 1, 5, 9, 13, 17, 21, 25, 29 }, c.taps[0..8]); + try std.testing.expectEqual(@as(u32, 16), c.block); + try std.testing.expectEqual(@as(u32, 15), c.depth()); + try std.testing.expectEqual(@as(u32, 248077), c.mask); + try std.testing.expectEqual(@as(u32, 4096), c.window); + try std.testing.expectEqual(@as(u32, 4096), c.hidden); + try std.testing.expectEqual(@as(u32, 12288), c.intermediate); + try std.testing.expectEqual(@as(u32, 6), c.layers); + try std.testing.expectEqual(@as(u32, 32), c.target_layers); + try std.testing.expectEqualSlices(bool, &.{ false, false, false, false, false, true }, c.full[0..6]); + try std.testing.expectEqual(@as(u32, 32768), c.tapWidth()); + try std.testing.expectEqual(@as(u32, 0), c.rank); + try std.testing.expectEqual(@as(u32, 16), c.topk); +} +test "a DFlash v1 draft counts eleven tensors a layer and seven linears a layer beside fc" { + const c = try parse(std.testing.allocator, v1_fixture); + try std.testing.expectEqual(@as(usize, 69), c.tensorCount()); + try std.testing.expectEqual(@as(usize, 43), c.linearCount()); +} +test "a DFlash v1 config whose taps fall past the target or repeat is refused" { + const a = std.testing.allocator; + const past = try std.mem.replaceOwned(u8, a, v1_fixture, "25,29]", "25,32]"); + defer a.free(past); + try std.testing.expectError(error.BadDraftConfig, parse(a, past)); + const repeated = try std.mem.replaceOwned(u8, a, v1_fixture, "25,29]", "25,25]"); + defer a.free(repeated); + try std.testing.expectError(error.UnsupportedDraftTaps, parse(a, repeated)); +} +test "a DFlash v1 config that asks for a causal draft is refused" { + const a = std.testing.allocator; + const causal = try std.mem.replaceOwned(u8, a, v1_fixture, "\"attention_bias\":false", "\"attention_bias\":false,\"is_causal\":true"); + defer a.free(causal); + try std.testing.expectError(error.BadDraftConfig, parse(a, causal)); +} +test "a v1 draft keeps fused ring rows and gives its full layer the whole span" { + var c = try parse(std.testing.allocator, v1_fixture); + try std.testing.expectEqual(@as(u32, 4096), c.ringWidth()); + try std.testing.expectEqual(@as(u32, 4096), c.span()); + c.history = 32768; + try std.testing.expectEqual(@as(u32, 32768), c.span()); + try std.testing.expectEqual(@as(u32, 4096), c.layerWindow(0)); + try std.testing.expectEqual(@as(u32, 32768), c.layerWindow(5)); +} +test "DFlash2 keeps raw tap rows and its window everywhere" { + var c = Config{}; + c.history = 32768; + try std.testing.expectEqual(@as(u32, 25600), c.ringWidth()); + try std.testing.expectEqual(@as(u32, 2048), c.layerWindow(0)); + try std.testing.expectEqual(@as(u32, 2048), c.span()); +} +test "the stock DFlash2 config keeps its five taps and twenty-five-thousand-wide rows" { + const c = try parse(std.testing.allocator, fixture); + try std.testing.expectEqual(Variant.dflash2, c.variant); + try std.testing.expectEqual(@as(u32, 5), c.tap_count); + try std.testing.expectEqual(@as(u32, 25600), c.tapWidth()); +} +test "a v1 config names its drafter DFlash, and DFlash2 stays DFlash2" { + try std.testing.expectEqualStrings("DFlash", (try parse(std.testing.allocator, v1_fixture)).kind()); + try std.testing.expectEqualStrings("DFlash2", (try parse(std.testing.allocator, fixture)).kind()); +} diff --git a/zig/src/families/qwen27/dflash/execution.zig b/zig/src/families/qwen27/dflash/execution.zig index e7cbe2a028..b64901756a 100644 --- a/zig/src/families/qwen27/dflash/execution.zig +++ b/zig/src/families/qwen27/dflash/execution.zig @@ -21,6 +21,13 @@ fn rotate(ops: op.Ops, frame: u64, x: op.View, positions: []const u64, heads: u3 try y.check(x.rows, x.width, .bf16); return y; } +fn add(ops: op.Ops, frame: u64, x: op.View, y: op.View) !op.View { + const sum = ops.vtable.add orelse return error.DraftOperatorShape; + try y.check(x.rows, x.width, .bf16); + const z = try sum(ops.ptr, frame, x, y); + try z.check(x.rows, x.width, .bf16); + return z; +} /// `block`: rows a conv window spans; a batched pass stacks several streams' blocks. fn convolve(ops: op.Ops, frame: u64, x: op.View, dynamic: op.View, weights: w.Conv, branch: u32, residual: ?op.View, c: Config, block: u32) !op.View { try dynamic.check(x.rows, c.dynamicWidth(), .bf16); @@ -34,9 +41,9 @@ pub fn append(ops: op.Ops, g: *const w.Graph, preparation: op.Preparation, conte try preparation.check(); const c = g.config; try c.check(); - if (context.cache == 0 or context.begin > context.end or next.cache != context.cache or positions.len == 0 or next.end != try std.math.add(u64, context.end, positions.len) or next.begin != @max(context.begin, next.end -| (@as(u64, c.window) - 1))) return error.BadDraftContext; + if (context.cache == 0 or context.begin > context.end or next.cache != context.cache or positions.len == 0 or next.end != try std.math.add(u64, context.end, positions.len) or next.begin != @max(context.begin, next.end -| (@as(u64, c.span()) - 1))) return error.BadDraftContext; for (positions, 0..) |p, i| if (p != context.end + i or p >= next.end) return error.BadCommittedTaps; - try taps.check(@intCast(positions.len), c.tapWidth(), .bf16); + try taps.check(@intCast(positions.len), c.ringWidth(), .bf16); const frame = try ops.vtable.begin(ops.ptr); defer ops.vtable.release(ops.ptr, frame); const transaction = try ops.vtable.context_begin(ops.ptr, context); @@ -49,9 +56,9 @@ fn checkAppend(g: *const w.Graph, preparation: op.Preparation, context: op.Conte try preparation.check(); const c = g.config; try c.check(); - if (context.cache == 0 or context.begin > context.end or next.cache != context.cache or positions.len == 0 or next.end != try std.math.add(u64, context.end, positions.len) or next.begin != @max(context.begin, next.end -| (@as(u64, c.window) - 1))) return error.BadDraftContext; + if (context.cache == 0 or context.begin > context.end or next.cache != context.cache or positions.len == 0 or next.end != try std.math.add(u64, context.end, positions.len) or next.begin != @max(context.begin, next.end -| (@as(u64, c.span()) - 1))) return error.BadDraftContext; for (positions, 0..) |p, i| if (p != context.end + i or p >= next.end) return error.BadCommittedTaps; - try taps.check(@intCast(positions.len), c.tapWidth(), .bf16); + try taps.check(@intCast(positions.len), c.ringWidth(), .bf16); } /// Where appended keys and values go: a transaction's context, or a batched pass's stream. @@ -61,8 +68,8 @@ fn encodeAppend(ops: op.Ops, g: *const w.Graph, preparation: op.Preparation, tap } fn encodeAppendTo(ops: op.Ops, g: *const w.Graph, preparation: op.Preparation, taps: op.View, positions: []const u64, frame: u64, sink: Sink) !void { const c = g.config; - const fusion = try project(ops, frame, taps, g.fusion, preparation.mode); - const hidden = try norm(ops, frame, fusion, g.hidden_norm, c.eps, 1, c.hidden); + // a v1 ring holds rows already fused and normed (the target's taps pass ran fc and hidden_norm) + const hidden = if (c.variant == .v1) taps else try norm(ops, frame, try project(ops, frame, taps, g.fusion, preparation.mode), g.hidden_norm, c.eps, 1, c.hidden); for (g.layers, 0..) |layer, i| { const raw_keys = try project(ops, frame, hidden, layer.attention.k, preparation.mode); const keys = try rotate(ops, frame, try norm(ops, frame, raw_keys, layer.attention.k_norm, c.eps, c.kv_heads, c.head_dim), positions, c.kv_heads, c); @@ -78,7 +85,7 @@ fn encodeAppendTo(ops: op.Ops, g: *const w.Graph, preparation: op.Preparation, t pub fn appendBlock(ops: op.Ops, g: *const w.Graph, preparation: op.Preparation, context: op.Context, taps: op.View, positions: []const u64, next: op.Context, anchor: u32, block: u32) !op.Lattice { if (block < 2 or block > 16) return error.BadRuntimeBlock; try checkAppend(g, preparation, context, taps, positions, next); - if (anchor >= g.config.vocab or next.end - next.begin > g.config.window - 1) return error.BadDraftContext; + if (anchor >= g.config.vocab or next.end - next.begin > g.config.span() - 1) return error.BadDraftContext; const frame = try ops.vtable.begin(ops.ptr); defer ops.vtable.release(ops.ptr, frame); const transaction = try ops.vtable.context_begin(ops.ptr, context); @@ -98,7 +105,7 @@ pub const Consumer = struct { ptr: *anyopaque, call: *const fn (*anyopaque, u64, fn start(ops: op.Ops, g: *const w.Graph, preparation: op.Preparation, context: op.Context, anchor: u32) !u64 { try preparation.check(); try g.config.check(); - if (anchor >= g.config.vocab or context.cache == 0 or context.begin > context.end or context.end - context.begin > g.config.window - 1) return error.BadDraftContext; + if (anchor >= g.config.vocab or context.cache == 0 or context.begin > context.end or context.end - context.begin > g.config.span() - 1) return error.BadDraftContext; return ops.vtable.begin(ops.ptr); } fn encode(ops: op.Ops, g: *const w.Graph, preparation: op.Preparation, context: op.Context, anchor: u32, frame: u64, block: u32, given: []const u32) !Output { @@ -126,7 +133,9 @@ fn encodeStreams(ops: op.Ops, g: *const w.Graph, preparation: op.Preparation, st for (given, 1..) |t, i| ids[i] = t; var hidden = try ops.vtable.embed(ops.ptr, frame, g.embed, ids[0..rows]); try hidden.check(rows, c.hidden, .bf16); - for (g.layers, 0..) |layer, i| { + if (c.variant == .v1) { + hidden = try layersV1(ops, g, preparation, frame, hidden, positions[0..rows], streams[0].context, if (n > 1) contexts[0..n] else &.{}, block); + } else for (g.layers, 0..) |layer, i| { var residual = hidden; const xn = try norm(ops, frame, hidden, layer.input_norm, c.eps, 1, c.hidden); const dynamic = try project(ops, frame, xn, layer.attention_conv.projection, preparation.mode); @@ -155,9 +164,32 @@ fn encodeStreams(ops: op.Ops, g: *const w.Graph, preparation: op.Preparation, st const head = try ops.vtable.head(ops.ptr, frame, final); if (head.logits.handle == 0 or head.logits.rows != headed or head.logits.dtype != .bf16 or head.logits.width == 0 or head.logits.width > c.vocab or head.logits.stride < head.logits.width) return error.DraftOperatorShape; if (head.token_ids) |mapping| try mapping.check(1, head.logits.width, .u32) else if (head.logits.width != c.vocab) return error.DraftCandidateId; - const projected = try project(ops, frame, final, g.hidden_projection, preparation.mode); + // a v1 draft has no selector projection: its lattice reads none (rank 0), so the final rows stand in + const projected = if (c.variant == .v1) final else try project(ops, frame, final, g.hidden_projection, preparation.mode); return .{ .head = head, .projected = projected }; } +/// A DFlash v1 body: attention and a SwiGLU MLP a layer, each added to its residual (no conv kernels). +fn layersV1(ops: op.Ops, g: *const w.Graph, preparation: op.Preparation, frame: u64, embedded: op.View, positions: []const u64, context: op.Context, contexts: []const op.Context, block: u32) !op.View { + const c = g.config; + const rows = embedded.rows; + var hidden = embedded; + for (g.layers, 0..) |layer, i| { + const xn = try norm(ops, frame, hidden, layer.input_norm, c.eps, 1, c.hidden); + const query = try rotate(ops, frame, try norm(ops, frame, try project(ops, frame, xn, layer.attention.q, preparation.mode), layer.attention.q_norm, c.eps, c.heads, c.head_dim), positions, c.heads, c); + const keys = try rotate(ops, frame, try norm(ops, frame, try project(ops, frame, xn, layer.attention.k, preparation.mode), layer.attention.k_norm, c.eps, c.kv_heads, c.head_dim), positions, c.kv_heads, c); + const values = try project(ops, frame, xn, layer.attention.v, preparation.mode); + const attended = try ops.vtable.attention(ops.ptr, frame, .{ .context = context, .config = c, .layer = @intCast(i), .query = query, .keys = keys, .values = values, .positions = positions, .streams = contexts, .block = block }); + try attended.check(rows, c.qWidth(), .bf16); + hidden = try add(ops, frame, hidden, try project(ops, frame, attended, layer.attention.o, preparation.mode)); + const post = try norm(ops, frame, hidden, layer.post_norm, c.eps, 1, c.hidden); + const gate = try project(ops, frame, post, layer.mlp.gate, preparation.mode); + const up = try project(ops, frame, post, layer.mlp.up, preparation.mode); + const activated = try ops.vtable.swiglu(ops.ptr, frame, gate, up); + try activated.check(rows, c.intermediate, .bf16); + hidden = try add(ops, frame, hidden, try project(ops, frame, activated, layer.mlp.down, preparation.mode)); + } + return hidden; +} /// A stream's new taps for a batched pass: absorbed into its caches before its block reads them. pub const Absorb = struct { taps: op.View, positions: []const u64 }; @@ -166,9 +198,9 @@ pub fn absorbBlocks(ops: op.Ops, g: *const w.Graph, preparation: op.Preparation, try preparation.check(); try g.config.check(); if (block < 2 or block > 16 or streams.len < 2 or streams.len * block > 128 or out.len != streams.len or absorbs.len != streams.len) return error.BadRuntimeBlock; - for (streams) |s| if (s.anchor >= g.config.vocab or s.context.cache == 0 or s.context.begin > s.context.end or s.context.end - s.context.begin > g.config.window - 1) return error.BadDraftContext; + for (streams) |s| if (s.anchor >= g.config.vocab or s.context.cache == 0 or s.context.begin > s.context.end or s.context.end - s.context.begin > g.config.span() - 1) return error.BadDraftContext; for (absorbs, 0..) |a, k| if (a) |x| { - try x.taps.check(@intCast(x.positions.len), g.config.tapWidth(), .bf16); + try x.taps.check(@intCast(x.positions.len), g.config.ringWidth(), .bf16); try encodeAppendTo(ops, g, preparation, x.taps, x.positions, frame, .{ .stream = @intCast(k) }); }; const output = try encodeStreams(ops, g, preparation, streams, frame, block, &.{}); diff --git a/zig/src/families/qwen27/dflash/execution_v1_test.zig b/zig/src/families/qwen27/dflash/execution_v1_test.zig new file mode 100644 index 0000000000..5619426f7c --- /dev/null +++ b/zig/src/families/qwen27/dflash/execution_v1_test.zig @@ -0,0 +1,183 @@ +//! Fake operators prove the DFlash v1 schedule: a plain qwen3 body, two residual adds a layer, no conv, no selector. +const std = @import("std"); +const ck = @import("../checkpoint.zig"); +const cfg = @import("config.zig"); +const op = @import("operators.zig"); +const execution = @import("execution.zig"); +const fixtures = @import("fixtures.zig"); +const session = @import("session.zig"); +const selector = @import("selector.zig"); +const toy = fixtures.toy_v1; +const Fake = struct { + handle: u64 = 1, + projects: u32 = 0, + adds: u32 = 0, + attends: u32 = 0, + appends: u32 = 0, + commits: u32 = 0, + ids: [16]u32 = undefined, + positions: [16]u64 = undefined, + headed: u32 = 0, + fn self(ptr: *anyopaque) *Fake { + return @ptrCast(@alignCast(ptr)); + } + fn view(f: *Fake, rows: u32, width: u32) op.View { + f.handle += 1; + return .{ .handle = f.handle, .rows = rows, .width = width, .stride = width, .dtype = .bf16 }; + } + fn begin(ptr: *anyopaque) !u64 { + self(ptr).handle += 1; + return self(ptr).handle; + } + fn release(_: *anyopaque, _: u64) void {} + fn project(ptr: *anyopaque, _: u64, x: op.View, t: ck.Tensor, _: op.Mode) !op.View { + if (t.bytes.len == 0) return error.FakeAbsentWeight; + self(ptr).projects += 1; + return self(ptr).view(x.rows, @intCast(t.shape[0])); + } + fn norm(ptr: *anyopaque, _: u64, x: op.View, n: op.Norm) !op.View { + if (x.width != n.heads * n.head_dim) return error.FakeNormShape; + return self(ptr).view(x.rows, x.width); + } + fn embed(ptr: *anyopaque, _: u64, _: ck.Linear, ids: []const u32) !op.View { + @memcpy(self(ptr).ids[0..ids.len], ids); + return self(ptr).view(@intCast(ids.len), toy.hidden); + } + fn conv(_: *anyopaque, _: u64, _: op.View, _: op.Conv) !op.View { + return error.FakeV1HasNoConv; + } + fn rope(ptr: *anyopaque, _: u64, x: op.View, a: op.Rotary) !op.View { + if (a.positions.len <= 16) @memcpy(self(ptr).positions[0..a.positions.len], a.positions); + return self(ptr).view(x.rows, x.width); + } + fn attention(ptr: *anyopaque, _: u64, a: op.Attention) !op.View { + if (a.mask != .all_slots or a.layer != self(ptr).attends or a.query.rows != a.positions.len or a.positions[0] != a.context.end) return error.FakeAttentionOrder; + self(ptr).attends += 1; + return self(ptr).view(a.query.rows, a.config.qWidth()); + } + fn swiglu(ptr: *anyopaque, _: u64, a: op.View, b: op.View) !op.View { + if (a.width != b.width or a.rows != b.rows) return error.FakeMlpShape; + return self(ptr).view(a.rows, a.width); + } + fn add(ptr: *anyopaque, _: u64, x: op.View, y: op.View) !op.View { + if (x.rows != y.rows or x.width != y.width or x.width != toy.hidden) return error.FakeAddShape; + self(ptr).adds += 1; + return self(ptr).view(x.rows, x.width); + } + fn mask(ptr: *anyopaque, _: u64, x: op.View, first: u32, count: u32) !op.View { + if (first != 1 or count != x.rows - 1) return error.FakeAnchorShift; + return self(ptr).view(count, x.width); + } + fn head(ptr: *anyopaque, _: u64, x: op.View) !op.Head { + self(ptr).headed = x.rows; + return .{ .logits = self(ptr).view(x.rows, toy.vocab) }; + } + fn lattices(_: *anyopaque, _: u64, _: op.Head, _: op.View, _: cfg.Config, _: u32, _: []op.Lattice) !void { + return error.FakeBatchedHead; + } + fn streamAppend(_: *anyopaque, _: u64, _: u32, _: u32, _: op.View, _: op.View, _: []const u64) !void { + return error.FakeBatchedHead; + } + /// A chain: depth d's best candidate is token 10 + d, its runner-up far behind. + fn lattice(_: *anyopaque, _: u64, h: op.Head, _: op.View, c: cfg.Config) !op.Lattice { + const a = std.testing.allocator; + const ids = try a.alloc(u32, h.logits.rows * c.topk); + errdefer a.free(ids); + const unary = try a.alloc(f32, ids.len); + errdefer a.free(unary); + for (0..h.logits.rows) |d| { + ids[d * c.topk] = @intCast(10 + d); + unary[d * c.topk] = 0; + ids[d * c.topk + 1] = @intCast(40 + d); + unary[d * c.topk + 1] = -20; + } + return .{ .gpa = a, .depth = h.logits.rows, .topk = c.topk, .rank = c.rank, .candidates = ids, .unary = unary, .projected = try a.alloc(f32, 0) }; + } + fn contextBegin(_: *anyopaque, _: op.Context) !u64 { + return 1; + } + fn append(ptr: *anyopaque, _: u64, _: u64, layer: u32, k: op.View, v: op.View, positions: []const u64) !void { + const f = self(ptr); + if (layer != f.appends or k.rows != positions.len or v.rows != k.rows) return error.FakeAppendOrder; + f.appends += 1; + } + fn commit(ptr: *anyopaque, _: u64, _: u64, _: op.Context) !void { + self(ptr).commits += 1; + } + fn stage(_: *anyopaque, _: u64, _: u64, _: op.Context) !void {} + fn abort(_: *anyopaque, _: u64) void {} + fn ops(f: *Fake) op.Ops { + return .{ .ptr = f, .vtable = &.{ .begin = begin, .release = release, .project = project, .norm = norm, .embed = embed, .conv = conv, .rope = rope, .attention = attention, .swiglu = swiglu, .mask_rows = mask, .head = head, .lattice = lattice, .lattices = lattices, .context_begin = contextBegin, .context_append = append, .context_commit = commit, .context_stage = stage, .context_abort = abort, .stream_append = streamAppend, .add = add } }; + } +}; +fn taps() session.Batch { + return .{ .values = .{ .handle = 99, .rows = 2, .width = toy.ringWidth(), .stride = toy.ringWidth(), .dtype = .bf16 }, .positions = &.{ 0, 1 }, .layers = toy.taps, .target_end = 2 }; +} + +test "a v1 block runs attention and an MLP a layer, two residual adds each, and heads its fifteen mask rows" { + const a = std.testing.allocator; + var store = try fixtures.Store.initConfig(a, toy); + defer store.deinit(); + var g = try store.graph(a); + defer g.deinit(); + var fake = Fake{}; + var state = try session.Session.init(toy, 1, 0); + const p = op.Preparation{ .mode = .bf16_reference }; + try state.absorb(fake.ops(), &g, p, taps()); + try std.testing.expectEqual(@as(u32, 3), fake.appends); + // the context append from fused ring rows: k and v a layer, no fc (the target's taps pass ran it) + try std.testing.expectEqual(@as(u32, 2 * 3), fake.projects); + fake.projects = 0; + var tree = try state.treeBlock(a, fake.ops(), &g, p, 5, try state.firstDraw(), 16, .{ .nodes = 15 }); + defer tree.deinit(); + try std.testing.expectEqualSlices(u32, &.{ 5, 62, 62, 62, 62, 62, 62, 62, 62, 62, 62, 62, 62, 62, 62, 62 }, &fake.ids); + try std.testing.expectEqualSlices(u64, &.{ 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17 }, &fake.positions); + try std.testing.expectEqual(@as(u32, 3), fake.attends); + try std.testing.expectEqual(@as(u32, 6), fake.adds); + // q, k, v, o, gate, up and down a layer, and no selector projection + try std.testing.expectEqual(@as(u32, 7 * 3), fake.projects); + try std.testing.expectEqual(@as(u32, 15), fake.headed); + // with no selector, a confident lattice is a chain of each depth's best + try std.testing.expectEqual(@as(usize, 15), tree.tokens.len); + for (tree.tokens, tree.parents, 0..) |token, parent, d| { + try std.testing.expectEqual(@as(u32, @intCast(10 + d)), token); + try std.testing.expectEqual(@as(i32, @intCast(d)) - 1, parent); + } +} +test "a v1 graph holds no conv or selector tensor" { + const a = std.testing.allocator; + var store = try fixtures.Store.initConfig(a, toy); + defer store.deinit(); + var g = try store.graph(a); + defer g.deinit(); + try std.testing.expectEqual(@as(usize, 11 * 3 + 3), store.keys.len); + try std.testing.expectEqual(@as(usize, 0), g.hidden_projection.bytes.len); + try std.testing.expectEqual(@as(usize, 0), g.predecessor.bytes.len); + for (g.layers) |l| try std.testing.expectEqual(@as(usize, 0), l.attention_conv.projection.bytes.len + l.mlp_conv.base.bytes.len); + try std.testing.expect(try g.nativeBytes() > 0); +} +test "a v1 block without an add operator is refused, not run without its residuals" { + const a = std.testing.allocator; + var store = try fixtures.Store.initConfig(a, toy); + defer store.deinit(); + var g = try store.graph(a); + defer g.deinit(); + var fake = Fake{}; + var vtable = fake.ops().vtable.*; + vtable.add = null; + const ops = op.Ops{ .ptr = &fake, .vtable = &vtable }; + try std.testing.expectError(error.DraftOperatorShape, execution.forwardBlock(ops, &g, .{ .mode = .bf16_reference }, .{ .cache = 1, .begin = 0, .end = 2 }, 5, 16)); +} + +test "a v1 drafter's context keeps its whole span, not just the sliding window" { + var c = toy; + c.history = 64; + const s = try session.Session.init(c, 1, 0); + var positions: [40]u64 = undefined; + for (&positions, 0..) |*p, i| p.* = i; + const next = try s.admit(.{ .values = .{ .handle = 2, .rows = 40, .width = c.ringWidth(), .stride = c.ringWidth(), .dtype = .bf16 }, .positions = &positions, .layers = c.taps, .target_end = 40 }); + try std.testing.expectEqual(@as(u64, 0), next.begin); + const windowed = try session.Session.init(toy, 1, 0); + const cropped = try windowed.admit(.{ .values = .{ .handle = 2, .rows = 40, .width = toy.ringWidth(), .stride = toy.ringWidth(), .dtype = .bf16 }, .positions = &positions, .layers = toy.taps, .target_end = 40 }); + try std.testing.expectEqual(@as(u64, 40 - 31), cropped.begin); +} diff --git a/zig/src/families/qwen27/dflash/fixtures.zig b/zig/src/families/qwen27/dflash/fixtures.zig index 4eb14e967f..85b187f358 100644 --- a/zig/src/families/qwen27/dflash/fixtures.zig +++ b/zig/src/families/qwen27/dflash/fixtures.zig @@ -4,56 +4,64 @@ const ck = @import("../checkpoint.zig"); const Config = @import("config.zig").Config; const weights = @import("weights.zig"); pub const toy = Config{ .hidden = 64, .intermediate = 128, .vocab = 64, .heads = 8, .kv_heads = 4, .head_dim = 8, .mask = 63, .window = 16, .rank = 16, .topk = 2 }; +/// A tiny DFlash v1 draft: three taps of an eight-layer target, a sixteen-row block, its last layer full. +pub const toy_v1 = Config{ .variant = .v1, .hidden = 64, .intermediate = 128, .layers = 3, .vocab = 64, .heads = 8, .kv_heads = 4, .head_dim = 8, .target_layers = 8, .block = 16, .mask = 62, .window = 32, .group = 0, .conv = 0, .rank = 0, .topk = 2, .taps = .{ 1, 3, 5, 0, 0, 0, 0, 0 }, .tap_count = 3, .full = .{ false, false, true, false, false, false, false, false } }; const Dimensions = struct { values: @FieldType(ck.Tensor, "shape"), rank: usize }; pub const Store = struct { arena: std.heap.ArenaAllocator, + config: Config = toy, keys: [][]const u8 = undefined, tensors: []ck.Tensor = undefined, embed: ck.Tensor = undefined, pub fn init(gpa: std.mem.Allocator) !Store { - var s = Store{ .arena = std.heap.ArenaAllocator.init(gpa) }; + return initConfig(gpa, toy); + } + /// Zeroed tensors for every name the draft of config `c` reads. + pub fn initConfig(gpa: std.mem.Allocator, c: Config) !Store { + var s = Store{ .arena = std.heap.ArenaAllocator.init(gpa), .config = c }; errdefer s.deinit(); const a = s.arena.allocator(); - s.keys = try weights.names(a, toy); + s.keys = try weights.names(a, c); s.tensors = try a.alloc(ck.Tensor, s.keys.len); for (s.keys, s.tensors) |name, *t| { - const shape = dims(name); + const shape = dims(c, name); var count: usize = 2; for (shape.values[0..shape.rank]) |n| count = try std.math.mul(usize, count, n); const data = try a.alloc(u8, count); @memset(data, 0); t.* = .{ .dtype = .bf16, .rank = @intCast(shape.rank), .shape = shape.values, .bytes = data }; } - const data = try a.alloc(u8, toy.vocab * toy.hidden * 2); + const data = try a.alloc(u8, c.vocab * c.hidden * 2); @memset(data, 0); - s.embed = .{ .dtype = .bf16, .rank = 2, .shape = .{ toy.vocab, toy.hidden, 1, 1, 1 }, .bytes = data }; + s.embed = .{ .dtype = .bf16, .rank = 2, .shape = .{ c.vocab, c.hidden, 1, 1, 1 }, .bytes = data }; return s; } pub fn deinit(s: *Store) void { s.arena.deinit(); } pub fn graph(s: *Store, gpa: std.mem.Allocator) !weights.Graph { - return weights.load(gpa, .{ .source = .{ .ptr = s, .getFn = get }, .names = s.keys }, toy, .{ .dense = s.embed }); + return weights.load(gpa, .{ .source = .{ .ptr = s, .getFn = get }, .names = s.keys }, s.config, .{ .dense = s.embed }); } fn get(ptr: *anyopaque, name: []const u8) !ck.Tensor { const s: *Store = @ptrCast(@alignCast(ptr)); for (s.keys, s.tensors) |key, t| if (std.mem.eql(u8, key, name)) return t; return error.MissingTensor; } - fn dims(name: []const u8) Dimensions { - if (std.mem.eql(u8, name, "fc.weight")) return .{ .values = .{ toy.hidden, toy.tapWidth(), 1, 1, 1 }, .rank = 2 }; - if (std.mem.eql(u8, name, "candidate_selector.hidden_projection.weight")) return .{ .values = .{ toy.rank, toy.hidden, 1, 1, 1 }, .rank = 2 }; - if (std.mem.startsWith(u8, name, "candidate_selector.")) return .{ .values = .{ toy.vocab, toy.rank, 1, 1, 1 }, .rank = 2 }; - if (std.mem.endsWith(u8, name, "base_kernel")) return .{ .values = .{ 2, 2, toy.hidden, 1, 1 }, .rank = 3 }; - if (std.mem.endsWith(u8, name, "kernel_projection.weight")) return .{ .values = .{ toy.dynamicWidth(), toy.hidden, 1, 1, 1 }, .rank = 2 }; - if (std.mem.endsWith(u8, name, "q_proj.weight")) return .{ .values = .{ toy.qWidth(), toy.hidden, 1, 1, 1 }, .rank = 2 }; - if (std.mem.endsWith(u8, name, "k_proj.weight") or std.mem.endsWith(u8, name, "v_proj.weight")) return .{ .values = .{ toy.kvWidth(), toy.hidden, 1, 1, 1 }, .rank = 2 }; - if (std.mem.endsWith(u8, name, "o_proj.weight")) return .{ .values = .{ toy.hidden, toy.qWidth(), 1, 1, 1 }, .rank = 2 }; - if (std.mem.endsWith(u8, name, "gate_proj.weight") or std.mem.endsWith(u8, name, "up_proj.weight")) return .{ .values = .{ toy.intermediate, toy.hidden, 1, 1, 1 }, .rank = 2 }; - if (std.mem.endsWith(u8, name, "down_proj.weight")) return .{ .values = .{ toy.hidden, toy.intermediate, 1, 1, 1 }, .rank = 2 }; - if (std.mem.endsWith(u8, name, "q_norm.weight") or std.mem.endsWith(u8, name, "k_norm.weight")) return .{ .values = .{ toy.head_dim, 1, 1, 1, 1 }, .rank = 1 }; - return .{ .values = .{ toy.hidden, 1, 1, 1, 1 }, .rank = 1 }; + fn dims(toy_config: Config, name: []const u8) Dimensions { + const t = toy_config; + if (std.mem.eql(u8, name, "fc.weight")) return .{ .values = .{ t.hidden, t.tapWidth(), 1, 1, 1 }, .rank = 2 }; + if (std.mem.eql(u8, name, "candidate_selector.hidden_projection.weight")) return .{ .values = .{ t.rank, t.hidden, 1, 1, 1 }, .rank = 2 }; + if (std.mem.startsWith(u8, name, "candidate_selector.")) return .{ .values = .{ t.vocab, t.rank, 1, 1, 1 }, .rank = 2 }; + if (std.mem.endsWith(u8, name, "base_kernel")) return .{ .values = .{ 2, 2, t.hidden, 1, 1 }, .rank = 3 }; + if (std.mem.endsWith(u8, name, "kernel_projection.weight")) return .{ .values = .{ t.dynamicWidth(), t.hidden, 1, 1, 1 }, .rank = 2 }; + if (std.mem.endsWith(u8, name, "q_proj.weight")) return .{ .values = .{ t.qWidth(), t.hidden, 1, 1, 1 }, .rank = 2 }; + if (std.mem.endsWith(u8, name, "k_proj.weight") or std.mem.endsWith(u8, name, "v_proj.weight")) return .{ .values = .{ t.kvWidth(), t.hidden, 1, 1, 1 }, .rank = 2 }; + if (std.mem.endsWith(u8, name, "o_proj.weight")) return .{ .values = .{ t.hidden, t.qWidth(), 1, 1, 1 }, .rank = 2 }; + if (std.mem.endsWith(u8, name, "gate_proj.weight") or std.mem.endsWith(u8, name, "up_proj.weight")) return .{ .values = .{ t.intermediate, t.hidden, 1, 1, 1 }, .rank = 2 }; + if (std.mem.endsWith(u8, name, "down_proj.weight")) return .{ .values = .{ t.hidden, t.intermediate, 1, 1, 1 }, .rank = 2 }; + if (std.mem.endsWith(u8, name, "q_norm.weight") or std.mem.endsWith(u8, name, "k_norm.weight")) return .{ .values = .{ t.head_dim, 1, 1, 1, 1 }, .rank = 1 }; + return .{ .values = .{ t.hidden, 1, 1, 1, 1 }, .rank = 1 }; } }; diff --git a/zig/src/families/qwen27/dflash/generation.zig b/zig/src/families/qwen27/dflash/generation.zig index f97ae93c33..fe2959f29a 100644 --- a/zig/src/families/qwen27/dflash/generation.zig +++ b/zig/src/families/qwen27/dflash/generation.zig @@ -23,8 +23,6 @@ pub const Generation = struct { pending: ?u32 = null, // runRound's next token as the last accept picked it; anything else that moves the target clears it pub fn init(a: std.mem.Allocator, runner: *Runner, draft: *Draft) !Generation { if (runner.model != draft.backend.target or runner.model.config.hidden != draft.graph.config.hidden or runner.model.config.layers != draft.graph.config.target_layers) return error.TargetBinding; - runner.capture_taps = true; - for (draft.graph.config.taps, &runner.taps.ids) |id, *to| to.* = id; const options = mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked; const kept = try runner.model.device.buffer(128 * 4, options); errdefer kept.deinit(); @@ -33,6 +31,12 @@ pub const Generation = struct { const matcher = try @import("round_gpu.zig").Matcher.init(runner.model.device); errdefer matcher.deinit(); const commit = try @import("core").tree_commit_gpu.Ops.init(runner.model.device); + errdefer commit.deinit(); + try runner.taps.setFuser(runner.model.device, null); + try runner.taps.setLayers(runner.model.device, draft.graph.config.taps[0..draft.graph.config.tap_count]); + // a v1 ring keeps fused rows: the target's taps pass runs the drafter's fc and hidden_norm as they complete + if (draft.graph.config.variant == .v1) try runner.taps.setFuser(runner.model.device, try draft.fuser(runner.taps.capacity)); + runner.capture_taps = true; return .{ .allocator = a, .runner = runner, .draft = draft, .kept = kept, .taps = taps, .matcher = matcher, .commit = commit }; } pub fn deinit(g: Generation) void { @@ -140,7 +144,8 @@ pub const Generation = struct { const width: u32 = @intCast(g.runner.model.config.vocab); const shape = try g.matcher.encode(e, .{ .buf = logits.buffer, .off = logits.offset }, ids, parents, width, budget, eos); const result = @import("core").tree_commit_gpu.Ref{ .buf = g.matcher.buffers[4] }; - try g.commit.taps(e, result, .{ .buf = g.runner.taps.buffer }, .{ .buf = g.taps }, .{ .rows = shape.rows, .width = g.runner.taps.width, .capacity = g.runner.taps.capacity, .planes = 5 }); + const src = g.runner.taps.source(); + try g.commit.taps(e, result, .{ .buf = src.buffer }, .{ .buf = g.taps }, .{ .rows = shape.rows, .width = src.width, .capacity = g.runner.taps.capacity, .planes = src.planes }); try g.commit.head(e, result, .{ .buf = logits.buffer, .off = logits.offset }, .{ .rows = shape.rows, .width = width, .stride = width }); e.end(); ended = true; diff --git a/zig/src/families/qwen27/dflash/operators.zig b/zig/src/families/qwen27/dflash/operators.zig index 45ad1df88a..de90332fdd 100644 --- a/zig/src/families/qwen27/dflash/operators.zig +++ b/zig/src/families/qwen27/dflash/operators.zig @@ -81,6 +81,8 @@ pub const Ops = struct { context_abort: *const fn (*anyopaque, u64) void, // a batched pass's stream k. stream_append: *const fn (*anyopaque, u64, u32, u32, View, View, []const u64) anyerror!void, + // x + y, row for row: a DFlash v1 layer's residuals (DFlash2 adds them inside its conv). + add: ?*const fn (*anyopaque, u64, View, View) anyerror!View = null, }; }; diff --git a/zig/src/families/qwen27/dflash/root.zig b/zig/src/families/qwen27/dflash/root.zig index c81ed141df..5ddc4472dc 100644 --- a/zig/src/families/qwen27/dflash/root.zig +++ b/zig/src/families/qwen27/dflash/root.zig @@ -19,4 +19,5 @@ test { _ = @import("runtime_context.zig"); _ = @import("fixtures.zig"); _ = @import("execution_test.zig"); + _ = @import("execution_v1_test.zig"); } diff --git a/zig/src/families/qwen27/dflash/runtime_backend.zig b/zig/src/families/qwen27/dflash/runtime_backend.zig index 51c2a4d3db..fedbd9df3a 100644 --- a/zig/src/families/qwen27/dflash/runtime_backend.zig +++ b/zig/src/families/qwen27/dflash/runtime_backend.zig @@ -59,6 +59,7 @@ pub const Backend = struct { config: cfg.Config, primitives: core.draft_ops.Ops, attention: core.shared_attention.Attention, + attention_full: ?core.shared_attention.Attention = null, // a v1 full layer's, over the whole span topk: core.bf16_topk_gpu.Ops, head: @import("runtime_head.zig").Head, frame: Frame, @@ -79,6 +80,8 @@ pub const Backend = struct { errdefer primitives.deinit(); const attention = try core.shared_attention.Attention.init(a, target.device, .{ .query_heads = c.heads, .kv_heads = c.kv_heads, .head_dim = c.head_dim, .rotary_dim = c.head_dim, .capacity = c.window + 16, .output = .bf16 }); errdefer attention.deinit(); + const attention_full: ?core.shared_attention.Attention = if (c.span() > c.window) try core.shared_attention.Attention.init(a, target.device, .{ .query_heads = c.heads, .kv_heads = c.kv_heads, .head_dim = c.head_dim, .rotary_dim = c.head_dim, .capacity = c.span() + 16, .output = .bf16 }) else null; + errdefer if (attention_full) |x| x.deinit(); const topk = try core.bf16_topk_gpu.Ops.init(target.device); errdefer topk.deinit(); var head = try @import("runtime_head.zig").Head.init(a, target.device, target.weights.head); @@ -86,7 +89,7 @@ pub const Backend = struct { // the trimmed head's shape is the drafter's own: compile its pipeline here, not in the first draft round try target.kernels.prepareDecode(head.linear); const caches = try newCaches(a, target.device, c, c.layers); - return .{ .allocator = a, .target = target, .checkpoint = checkpoint, .config = c, .primitives = primitives, .attention = attention, .topk = topk, .head = head, .frame = .{ .allocator = a, .device = target.device, .queue = target.queue }, .caches = caches }; + return .{ .allocator = a, .target = target, .checkpoint = checkpoint, .config = c, .primitives = primitives, .attention = attention, .attention_full = attention_full, .topk = topk, .head = head, .frame = .{ .allocator = a, .device = target.device, .queue = target.queue }, .caches = caches }; } pub fn setResident(b: *Backend, hook: ?mtl.ResidencySet.Follow) void { @import("../residency.zig").setDraft(b, hook); @@ -105,6 +108,7 @@ pub const Backend = struct { freeCaches(b.allocator, b.caches); b.head.deinit(); b.topk.deinit(); + if (b.attention_full) |x| x.deinit(); b.attention.deinit(); b.primitives.deinit(); } @@ -147,6 +151,26 @@ pub const Backend = struct { pub fn linear(name: []const u8, rank: usize) bool { return rank == 2 and std.mem.endsWith(u8, name, ".weight"); } + /// `rows` rows of x (contiguous, `w`'s width) through linear `w` into y, outside any drafter frame. + pub fn encodeLinear(b: *Backend, e: mtl.ComputeEncoder, w: ck.Tensor, mode: op.Mode, x: Ref, y: Ref, rows: u32) !void { + const n: u32 = @intCast(w.shape[0]); + const k: u32 = @intCast(w.shape[1]); + if (mode == .bf16_reference) { + try b.primitives.linear(e, x, try b.weight(w), y, .{ .rows = rows, .n = n, .k = k, .x_stride = k, .w_stride = k, .y_stride = n }); + } else if (b.tiled.get(@intFromPtr(w.bytes.ptr))) |t| { + simd.encode(e, try simd.call(&b.simd.?, t.n, t.k, rows, false, false, 0), t.matrix(), x.buf, x.off, y.buf, y.off); + } else { + const adapter = b.prepared.get(@intFromPtr(w.bytes.ptr)) orelse return error.DraftPreparationUnqualified; + var first: u32 = 0; + while (first < rows) : (first += 16) { + try adapter.encode(e, .{ .buf = x.buf, .off = x.off + @as(usize, first) * k * 2 }, .{ .buf = y.buf, .off = y.off + @as(usize, first) * n * 2 }, @min(rows - first, 16)); + } + } + } + /// The attention whose capacity is layer `layer`'s window. + pub fn attentionFor(b: *const Backend, layer: u32) *const core.shared_attention.Attention { + return if (b.config.layerWindow(layer) > b.config.window) &b.attention_full.? else &b.attention; + } pub fn ops(b: *Backend) op.Ops { return .{ .ptr = b, .vtable = &callbacks.vtable }; } diff --git a/zig/src/families/qwen27/dflash/runtime_callbacks.zig b/zig/src/families/qwen27/dflash/runtime_callbacks.zig index bfea3d8552..cf004361b1 100644 --- a/zig/src/families/qwen27/dflash/runtime_callbacks.zig +++ b/zig/src/families/qwen27/dflash/runtime_callbacks.zig @@ -89,6 +89,18 @@ fn swiglu(ptr: *anyopaque, epoch: u64, g: op.View, u: op.View) !op.View { e.barrier(); return y; } +fn add(ptr: *anyopaque, epoch: u64, x: op.View, y: op.View) !op.View { + const b = backend(ptr); + const e = try b.encoder(epoch); + if (x.rows != y.rows or x.width != y.width or x.dtype != .bf16 or y.dtype != .bf16 or x.stride != x.width or y.stride != y.width) return error.DraftOperatorShape; + const z = try b.frame.alloc(x.rows, x.width, .bf16); + const xr = try b.frame.ref(x); + const yr = try b.frame.ref(y); + const zr = try b.frame.ref(z); + b.target.glue.residual(e, .{ .buffer = xr.buf, .offset = xr.off }, .{ .buffer = yr.buf, .offset = yr.off }, .{ .buffer = zr.buf, .offset = zr.off }, x.width, x.rows); + e.barrier(); + return z; +} fn maskRows(ptr: *anyopaque, epoch: u64, x: op.View, first: u32, rows: u32) !op.View { const b = backend(ptr); _ = try b.encoder(epoch); @@ -153,4 +165,4 @@ fn lattices(ptr: *anyopaque, epoch: u64, h: op.Head, projected: op.View, c: Conf made += 1; } } -pub const vtable = op.Ops.VTable{ .begin = begin, .release = release, .project = project, .norm = norm, .embed = embed, .conv = conv, .rope = rope, .attention = context.attention, .swiglu = swiglu, .mask_rows = maskRows, .head = head, .lattice = lattice, .lattices = lattices, .context_begin = context.begin, .context_append = context.append, .context_commit = context.commit, .context_stage = context.stage, .context_abort = context.abort, .stream_append = context.streamAppend }; +pub const vtable = op.Ops.VTable{ .begin = begin, .release = release, .project = project, .norm = norm, .embed = embed, .conv = conv, .rope = rope, .attention = context.attention, .swiglu = swiglu, .mask_rows = maskRows, .head = head, .lattice = lattice, .lattices = lattices, .context_begin = context.begin, .context_append = context.append, .context_commit = context.commit, .context_stage = context.stage, .context_abort = context.abort, .stream_append = context.streamAppend, .add = add }; diff --git a/zig/src/families/qwen27/dflash/runtime_context.zig b/zig/src/families/qwen27/dflash/runtime_context.zig index c76c34028c..f50cd9f453 100644 --- a/zig/src/families/qwen27/dflash/runtime_context.zig +++ b/zig/src/families/qwen27/dflash/runtime_context.zig @@ -23,7 +23,7 @@ pub fn append(ptr: *anyopaque, transaction: u64, epoch: u64, layer: u32, keys: o try values.check(keys.rows, keys.width, .bf16); if (b.appended != 0 and b.appended_rows != positions.len) return error.BadCommittedTaps; for (positions, 0..) |position, i| if (position != b.context.end + i) return error.BadCommittedTaps; - try appendRows(b, e, b.caches[layer], keys, values, positions); + try appendRows(b, e, b.caches[layer], layer, keys, values, positions); b.appended |= @as(u32, 1) << @intCast(layer); b.appended_rows = @intCast(positions.len); } @@ -33,24 +33,25 @@ pub fn streamAppend(ptr: *anyopaque, epoch: u64, stream: u32, layer: u32, keys: if (b.transaction != null or stream >= b.batch.len or layer >= b.config.layers or positions.len == 0) return error.BadDraftContext; try keys.check(@intCast(positions.len), b.config.kvWidth(), .bf16); try values.check(keys.rows, keys.width, .bf16); - try appendRows(b, e, b.batch[stream][layer], keys, values, positions); + try appendRows(b, e, b.batch[stream][layer], layer, keys, values, positions); } /// Rows' keys and values into `cache` at their positions' ring slots, 16 a dispatch. -fn appendRows(b: *Backend, e: @import("metal").ComputeEncoder, cache: @import("runtime_backend.zig").Cache, keys: op.View, values: op.View, positions: []const u64) !void { +fn appendRows(b: *Backend, e: @import("metal").ComputeEncoder, cache: @import("runtime_backend.zig").Cache, layer: u32, keys: op.View, values: op.View, positions: []const u64) !void { + const window = b.config.layerWindow(layer); const k = try b.frame.ref(keys); const v = try b.frame.ref(values); var first: usize = 0; while (first < positions.len) : (first += 16) { const count = @min(positions.len - first, 16); var slots: [16]u32 = undefined; - for (slots[0..count], 0..) |*slot, i| slot.* = @intCast(positions[first + i] % b.config.window); + for (slots[0..count], 0..) |*slot, i| slot.* = @intCast(positions[first + i] % window); const off = first * b.config.kvWidth() * 2; - try b.attention.appendKv(e, slots[0..count], .{ .buf = k.buf, .off = k.off + off }, .{ .buf = v.buf, .off = v.off + off }, .{ .buf = cache.keys }, .{ .buf = cache.values }); + try b.attentionFor(layer).appendKv(e, slots[0..count], .{ .buf = k.buf, .off = k.off + off }, .{ .buf = v.buf, .off = v.off + off }, .{ .buf = cache.keys }, .{ .buf = cache.values }); e.barrier(); } } fn ready(b: *Backend, transaction: u64, epoch: u64, next: op.Context) !void { - if (b.transaction != transaction or transaction != epoch or b.appended != (@as(u32, 1) << @intCast(b.config.layers)) - 1 or next.cache != b.context.cache or next.end != b.context.end + b.appended_rows or next.begin != @max(b.context.begin, next.end -| (@as(u64, b.config.window) - 1))) return error.BadDraftContext; + if (b.transaction != transaction or transaction != epoch or b.appended != (@as(u32, 1) << @intCast(b.config.layers)) - 1 or next.cache != b.context.cache or next.end != b.context.end + b.appended_rows or next.begin != @max(b.context.begin, next.end -| (@as(u64, b.config.span()) - 1))) return error.BadDraftContext; } pub fn commit(ptr: *anyopaque, transaction: u64, epoch: u64, next: op.Context) !void { const b = cast(ptr); @@ -81,11 +82,11 @@ pub fn abort(ptr: *anyopaque, transaction: u64) void { b.staged = null; } } -fn gather(b: *Backend, e: @import("metal").ComputeEncoder, context: op.Context, cache: @import("metal").Buffer, out: Ref) !void { +fn gather(b: *Backend, e: @import("metal").ComputeEncoder, context: op.Context, cache: @import("metal").Buffer, window: u32, out: Ref) !void { const rows: u32 = @intCast(context.end - context.begin); if (rows == 0) return; - const at: u32 = @intCast(context.begin % b.config.window); - const first = @min(rows, b.config.window - at); + const at: u32 = @intCast(context.begin % window); + const first = @min(rows, window - at); const width = b.config.kvWidth(); try b.copy(e, .{ .buf = cache, .off = @as(usize, at) * width * 2 }, out, first, width); if (first < rows) try b.copy(e, .{ .buf = cache }, .{ .buf = out.buf, .off = out.off + @as(usize, first) * width * 2 }, rows - first, width); @@ -107,12 +108,14 @@ pub fn attention(ptr: *anyopaque, epoch: u64, p: op.Attention) !op.View { } /// Query rows [first, first + rows) against `context`'s cached keys and their own, into the same rows of `result`. -fn attendRows(b: *Backend, e: @import("metal").ComputeEncoder, p: op.Attention, context: op.Context, caches: []const @import("runtime_backend.zig").Cache, first: u32, rows: u32, result: op.View) !void { +fn attendRows(b: *Backend, e: @import("metal").ComputeEncoder, p: op.Attention, kept: op.Context, caches: []const @import("runtime_backend.zig").Cache, first: u32, rows: u32, result: op.View) !void { + const window = b.config.layerWindow(p.layer); + const context = layerContext(kept, window); const context_rows: u32 = @intCast(context.end - context.begin); const keys = context_rows + rows; // An unwrapped context is contiguous in the cache: the query rows go behind it and attention reads in place. - const at: u32 = @intCast(context.begin % b.config.window); - const in_place = at + context_rows <= b.config.window and at + keys <= b.config.window + 16; + const at: u32 = @intCast(context.begin % window); + const in_place = at + context_rows <= window and at + keys <= window + 16; var k: Ref = undefined; var v: Ref = undefined; if (in_place) { @@ -123,8 +126,8 @@ fn attendRows(b: *Backend, e: @import("metal").ComputeEncoder, p: op.Attention, const merged_v = try b.frame.alloc(keys, b.config.kvWidth(), .bf16); k = try b.frame.ref(merged_k); v = try b.frame.ref(merged_v); - try gather(b, e, context, caches[p.layer].keys, k); - try gather(b, e, context, caches[p.layer].values, v); + try gather(b, e, context, caches[p.layer].keys, window, k); + try gather(b, e, context, caches[p.layer].values, window, v); } const off = @as(usize, context_rows) * b.config.kvWidth() * 2; const own_k = try b.frame.ref(p.keys); @@ -133,16 +136,27 @@ fn attendRows(b: *Backend, e: @import("metal").ComputeEncoder, p: op.Attention, try b.copy(e, .{ .buf = own_v.buf, .off = own_v.off + @as(usize, first) * p.values.stride * 2 }, .{ .buf = v.buf, .off = v.off + off }, rows, p.values.width); const mask = try b.frame.alloc(rows, keys, .u8); const mask_ref = try b.frame.ref(mask); - try fillMask(mask_ref.buf.contents()[mask_ref.off..][0 .. @as(usize, rows) * keys], context_rows, rows, b.config.window, b.config.block); + try fillMask(mask_ref.buf.contents()[mask_ref.off..][0 .. @as(usize, rows) * keys], context_rows, rows, window, b.config.block); const q = try b.frame.ref(p.query); const y = try b.frame.ref(result); const scale = 1.0 / @sqrt(@as(f32, @floatFromInt(b.config.head_dim))); const R = @import("core").shared_attention.Ref; const views = [_]R{ .{ .buf = q.buf, .off = q.off + @as(usize, first) * p.query.stride * 2 }, .{ .buf = k.buf, .off = k.off }, .{ .buf = v.buf, .off = v.off }, .{ .buf = mask_ref.buf, .off = mask_ref.off }, .{ .buf = y.buf, .off = y.off + @as(usize, first) * result.stride * 2 } }; // Long contexts share each 256-key chunk across the query heads of one KV head. - if (keys > 1024) try b.attention.encodeSplit(e, rows, keys, keys, scale, views[0], views[1], views[2], views[3], views[4]) else try b.attention.encode(e, rows, keys, keys, scale, views[0], views[1], views[2], views[3], views[4]); + const attend = b.attentionFor(p.layer); + if (keys > 1024) try attend.encodeSplit(e, rows, keys, keys, scale, views[0], views[1], views[2], views[3], views[4]) else try attend.encode(e, rows, keys, keys, scale, views[0], views[1], views[2], views[3], views[4]); } +/// The rows a layer of window `window` attends: the kept context's last `window` - 1 (all of it for a full layer). +pub fn layerContext(kept: op.Context, window: u32) op.Context { + return .{ .cache = kept.cache, .begin = @max(kept.begin, kept.end -| (@as(u64, window) - 1)), .end = kept.end }; +} +test "a sliding layer attends its own window of a longer kept context, a full layer all of it" { + const kept = op.Context{ .cache = 1, .begin = 0, .end = 10000 }; + try std.testing.expectEqual(@as(u64, 10000 - 4095), layerContext(kept, 4096).begin); + try std.testing.expectEqual(@as(u64, 0), layerContext(kept, 32768).begin); + try std.testing.expectEqual(@as(u64, 7), layerContext(.{ .cache = 1, .begin = 7, .end = 100 }, 4096).begin); +} /// Slots before `trained` cannot see later slots, retaining their trained-block mask. pub fn fillMask(out: []u8, context_rows: u32, queries: u32, window: u32, trained: u32) !void { if (queries == 0 or window == 0 or context_rows >= window or trained == 0) return error.BadDraftContext; diff --git a/zig/src/families/qwen27/dflash/runtime_lattice.zig b/zig/src/families/qwen27/dflash/runtime_lattice.zig index e55d6f8209..46e8143743 100644 --- a/zig/src/families/qwen27/dflash/runtime_lattice.zig +++ b/zig/src/families/qwen27/dflash/runtime_lattice.zig @@ -10,9 +10,10 @@ fn value(r: @import("runtime_frame.zig").Ref, at: usize) f32 { pub fn read(a: std.mem.Allocator, frame: *const Frame, topk: op.View, hidden: op.View, c: Config, head: op.Head) !op.Lattice { if (!frame.completed) return error.DraftLatticeShape; try topk.check(head.logits.rows, @sizeOf(@import("core").bf16_topk.Row), .u8); - try hidden.check(head.logits.rows, c.rank, .bf16); + // rank 0 (a v1 draft): no selector projection to read, whatever `hidden` holds + if (c.rank > 0) try hidden.check(head.logits.rows, c.rank, .bf16); const picked = try frame.ref(topk); - const projected = try frame.ref(hidden); + const projected = if (c.rank > 0) try frame.ref(hidden) else picked; const mapping = if (head.token_ids) |ids| try frame.ref(ids) else null; const count = @as(usize, head.logits.rows) * c.topk; var out = op.Lattice{ .gpa = a, .depth = head.logits.rows, .topk = c.topk, .rank = c.rank, .candidates = try a.alloc(u32, count), .unary = undefined, .projected = undefined }; diff --git a/zig/src/families/qwen27/dflash/runtime_model.zig b/zig/src/families/qwen27/dflash/runtime_model.zig index 0ed3fe34aa..74e934ee4d 100644 --- a/zig/src/families/qwen27/dflash/runtime_model.zig +++ b/zig/src/families/qwen27/dflash/runtime_model.zig @@ -9,6 +9,8 @@ const Backend = @import("runtime_backend.zig").Backend; const Target = @import("../model.zig").Model; const Session = @import("session.zig").Session; const op = @import("operators.zig"); +const tap_ring = @import("../tap_ring.zig"); +const kv_grow = @import("../kv_grow.zig"); const execution = @import("execution.zig"); const selector = @import("selector.zig"); const backend_mod = @import("runtime_backend.zig"); @@ -25,13 +27,20 @@ pub const Model = struct { tree_block: u32 = 16, pending: mtl.Buffer, committed_end: u64 = 0, + fuse_out: ?mtl.Buffer = null, // fc's rows before hidden_norm, for the target's taps pass + fuse_rows: u32 = 0, pub fn load(a: std.mem.Allocator, io: std.Io, target: *Target, dir: []const u8, mode: op.Mode) !*Model { + return loadSpan(a, io, target, dir, mode, 0); + } + /// `history`: positions a v1 drafter keeps for its full layers (the served span); 0 keeps the window. + pub fn loadSpan(a: std.mem.Allocator, io: std.Io, target: *Target, dir: []const u8, mode: op.Mode, history: u32) !*Model { if (mode == .python_q4g64) return error.DraftPreparationUnqualified; const path = try std.fs.path.join(a, &.{ dir, "config.json" }); defer a.free(path); const text = try std.Io.Dir.cwd().readFileAlloc(io, path, a, .limited(1 << 20)); defer a.free(text); - const c = try cfg.parse(a, text); + var c = try cfg.parse(a, text); + if (c.variant == .v1) c.history = history; try c.bind(target.config.hidden, target.config.vocab, target.config.layers); const shards = try core.checkpoint_host.shardFiles(a, io, dir); defer core.checkpoint_host.freeShardFiles(a, shards); @@ -45,6 +54,8 @@ pub const Model = struct { model.allocator = a; model.tree_block = 16; model.committed_end = 0; + model.fuse_out = null; + model.fuse_rows = 0; model.checkpoint = core.checkpoint_metal.Checkpoint.init(a); errdefer model.checkpoint.deinit(); // q4 drafts read their linears only through prepared copies: those stay in the file until prepare reads them @@ -56,13 +67,14 @@ pub const Model = struct { errdefer model.graph.deinit(); if (mode == .prepared_q4_reference) try model.backend.prepare(); model.checkpoint.closeFiles(); - model.pending = try target.device.buffer(@as(usize, c.window) * c.tapWidth() * 2, mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked); + model.pending = try target.device.buffer(@as(usize, c.span()) * c.ringWidth() * 2, mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked); errdefer model.pending.deinit(); model.session = try Session.init(c, 1, 0); model.preparation = .{ .mode = mode, .quantization_verified = mode == .prepared_q4_reference, .source_sha256 = if (mode == .prepared_q4_reference) core.affine4.signature() else @splat(0) }; return model; } pub fn deinit(m: *Model) void { + if (m.fuse_out) |b| b.deinit(); m.graph.deinit(); m.backend.deinit(); m.pending.deinit(); @@ -84,30 +96,49 @@ pub const Model = struct { m.session = try Session.init(m.graph.config, 1, 0); m.committed_end = 0; } + /// The target's taps pass fuses rows for this v1 drafter: fc then hidden_norm on up to `rows` rows a pass. + pub fn fuser(m: *Model, rows: u32) !@import("../taps.zig").Fuser { + if (m.graph.config.variant != .v1 or rows == 0) return error.BadDraftOptions; + if (m.fuse_out) |b| b.deinit(); + m.fuse_out = null; + m.fuse_out = try m.backend.target.device.buffer(@as(usize, rows) * m.graph.config.hidden * 2, opts); + m.fuse_rows = rows; + return .{ .ptr = m, .width = m.graph.config.hidden, .encode = fuseEncode }; + } + fn fuseEncode(ptr: *anyopaque, e: mtl.ComputeEncoder, x: @import("../projection.zig").Ref, y: @import("../projection.zig").Ref, rows: u32) !void { + const m: *Model = @ptrCast(@alignCast(ptr)); + if (rows == 0 or rows > m.fuse_rows or m.backend.failed) return error.BadCommittedTaps; + const c = m.graph.config; + const out = m.fuse_out.?; + try m.backend.encodeLinear(e, m.graph.fusion, m.preparation.mode, .{ .buf = x.buffer, .off = x.offset }, .{ .buf = out }, rows); + e.barrier(); + try m.backend.primitives.norm(e, .{ .buf = out }, try m.backend.weight(m.graph.hidden_norm), .{ .buf = y.buffer, .off = y.offset }, .{ .rows = rows, .heads = 1, .dim = c.hidden, .x_stride = c.hidden, .y_stride = c.hidden, .eps = c.eps }); + e.barrier(); + } pub fn absorb(m: *Model, taps: @import("runtime_backend.zig").Ref, rows: u32) !void { if (m.backend.failed or m.backend.frame.active or m.backend.transaction != null or rows == 0 or rows > 128) return error.BadCommittedTaps; try m.ring().absorb(taps.buf, taps.off, rows); } /// The committed taps the drafter reads, as a ring a lane's state copies. pub fn ring(m: *Model) @import("../tap_ring.zig").Ring { - return .{ .buf = m.pending, .end = &m.committed_end, .window = m.graph.config.window, .stride = @as(usize, m.graph.config.tapWidth()) * 2 }; + return .{ .buf = m.pending, .end = &m.committed_end, .window = m.graph.config.span(), .stride = @as(usize, m.graph.config.ringWidth()) * 2 }; } pub fn flush(m: *Model) !void { if (m.backend.failed or m.backend.frame.active or m.backend.transaction != null or m.committed_end == 0 or m.session.context.end > m.committed_end) return error.BadDraftContext; - const begin = m.committed_end -| (@as(u64, m.graph.config.window) - 1); + const begin = m.committed_end -| (@as(u64, m.graph.config.span()) - 1); if (m.session.context.end < begin) { m.session.context = .{ .cache = 1, .begin = begin, .end = begin }; m.session.ready = false; m.backend.context = m.session.context; } - const stride = @as(usize, m.graph.config.tapWidth()) * 2; + const stride = @as(usize, m.graph.config.ringWidth()) * 2; while (m.session.context.end < m.committed_end) { const first = m.session.context.end; - const slot: u32 = @intCast(first % m.graph.config.window); - const rows: u32 = @intCast(@min(@min(m.committed_end - first, 128), m.graph.config.window - slot)); + const slot: u32 = @intCast(first % m.graph.config.span()); + const rows: u32 = @intCast(@min(@min(m.committed_end - first, 128), m.graph.config.span() - slot)); var positions: [128]u64 = undefined; for (positions[0..rows], 0..) |*position, i| position.* = first + i; - const view = try m.backend.frame.borrow(.{ .buf = m.pending, .off = @as(usize, slot) * stride }, rows, m.graph.config.tapWidth()); + const view = try m.backend.frame.borrow(.{ .buf = m.pending, .off = @as(usize, slot) * stride }, rows, m.graph.config.ringWidth()); try m.session.absorb(m.backend.ops(), &m.graph, m.preparation, .{ .values = view, .positions = positions[0..rows], .layers = m.graph.config.taps, .target_end = first + rows }); } } @@ -170,10 +201,10 @@ pub const Model = struct { defer ops.vtable.release(ops.ptr, frame); var absorbs: [8]?execution.Absorb = @splat(null); var positions: [8][128]u64 = undefined; - const stride = @as(usize, c.tapWidth()) * 2; + const stride = @as(usize, c.ringWidth()) * 2; for (news[0..n], absorbs[0..n], 0..) |f, *a, k| if (f) |x| { for (positions[k][0..x.rows], 0..) |*p, i| p.* = x.first + i; - a.* = .{ .taps = try m.backend.frame.add(.{ .buf = x.ring, .off = @as(usize, x.slot) * stride }, x.rows, c.tapWidth(), c.tapWidth(), .bf16), .positions = positions[k][0..x.rows] }; + a.* = .{ .taps = try m.backend.frame.add(.{ .buf = x.ring, .off = @as(usize, x.slot) * stride }, x.rows, c.ringWidth(), c.ringWidth(), .bf16), .positions = positions[k][0..x.rows] }; }; var lattices: [8]op.Lattice = undefined; try execution.absorbBlocks(ops, &m.graph, m.preparation, frame, absorbs[0..n], streams[0..n], block, lattices[0..n]); @@ -200,7 +231,7 @@ pub const Model = struct { /// A short run of committed taps the session lacks, for a batched frame to absorb (null: none, or flush first). const Fresh = struct { ring: mtl.Buffer, slot: u32, rows: u32, first: u64, next: op.Context }; fn fresh(m: *const Model) ?Fresh { - const window = m.graph.config.window; + const window = m.graph.config.span(); const span = @import("catchup.zig").fresh(m.session.ready, m.session.context.begin, m.session.context.end, m.committed_end, window) orelse return null; return .{ .ring = m.pending, .slot = span.slot, .rows = span.rows, .first = span.first, .next = .{ .cache = m.session.context.cache, .begin = span.begin, .end = m.committed_end } }; } @@ -216,14 +247,14 @@ pub const Model = struct { if (m.committed_end == 0) return error.DraftContextNotReady; const block = @min(m.tree_block, nodes + 1); if (block < 2) return .{ .gpa = m.allocator, .tokens = &.{}, .parents = &.{}, .scores = &.{} }; - const window = m.graph.config.window; + const window = m.graph.config.span(); if (@import("catchup.zig").fresh(m.session.ready, m.session.context.begin, m.session.context.end, m.committed_end, window)) |span| { const rows = span.rows; const slot = span.slot; var positions: [128]u64 = undefined; for (positions[0..rows], 0..) |*position, i| position.* = m.session.context.end + i; - const stride = @as(usize, m.graph.config.tapWidth()) * 2; - const view = try m.backend.frame.borrow(.{ .buf = m.pending, .off = @as(usize, slot) * stride }, rows, m.graph.config.tapWidth()); + const stride = @as(usize, m.graph.config.ringWidth()) * 2; + const view = try m.backend.frame.borrow(.{ .buf = m.pending, .off = @as(usize, slot) * stride }, rows, m.graph.config.ringWidth()); return m.session.absorbTreeBlock(m.allocator, m.backend.ops(), &m.graph, m.preparation, .{ .values = view, .positions = positions[0..rows], .layers = m.graph.config.taps, .target_end = m.session.context.end + rows }, pending, block, options); } try m.flush(); @@ -239,21 +270,65 @@ pub const Context = struct { caches: []Cache, session: Session, context: op.Context = .{ .cache = 1, .begin = 0, .end = 0 }, + device: mtl.Device, + span: u32, // the ring's window + stride: usize, // a ring row's bytes + held: u32, // the ring's rows: the span unless it grows as used + step: ?u32 = null, // a ring that grows as used (TF_Q27_KV_GROW): its first rows and step + settle: ?Settle = null, // the slot's pending reads land before its ring's buffer changes + + /// The runner's wait for a blit still reading the slot (Runner.settleSlot). + pub const Settle = struct { ctx: *anyopaque, slot: u32, run: *const fn (ctx: *anyopaque, slot: u32) void }; - /// `drafting`: the stream drafts its own trees, so it gets K/V caches; else only its ring. - pub fn init(a: std.mem.Allocator, m: *const Model, drafting: bool) !*Context { + /// `drafting`: the stream drafts its own trees, so it gets K/V caches; its ring grows by `step` (null: the span). + pub fn init(a: std.mem.Allocator, m: *const Model, drafting: bool, step: ?u32) !*Context { const c = m.graph.config; const device = m.backend.target.device; const x = try a.create(Context); errdefer a.destroy(x); - const pending = try device.buffer(@as(usize, c.window) * c.tapWidth() * 2, opts); + const grows = step; // drafter frames wait for their ring reads before closing: a ring grows between them + const held = if (grows) |s| kv_grow.first(s, c.span()) else c.span(); + const pending = try device.buffer(@as(usize, held) * c.ringWidth() * 2, opts); errdefer pending.deinit(); const caches = try backend_mod.newCaches(a, device, c, if (drafting) c.layers else 0); errdefer backend_mod.freeCaches(a, caches); - x.* = .{ .pending = pending, .caches = caches, .session = try Session.init(c, 1, 0) }; + x.* = .{ .pending = pending, .caches = caches, .session = try Session.init(c, 1, 0), .device = device, .span = c.span(), .stride = @as(usize, c.ringWidth()) * 2, .held = held, .step = grows }; x.setResident(m.backend.residency); return x; } + /// The ring's buffer holding `rows` (Ring.Grow): a wider one takes every held row at its place. + fn growRing(ctx: *anyopaque, rows: u32) anyerror!mtl.Buffer { + const x: *Context = @ptrCast(@alignCast(ctx)); + if (rows <= x.held) return x.pending; + const stride = x.stride; + const next = tap_ring.grownRows(x.held, rows, x.step.?, x.span); + const fresh = try x.device.buffer(@as(usize, next) * stride, opts); + if (x.settle) |s| s.run(s.ctx, s.slot); + tap_ring.carry(x.pending, fresh, x.held, stride); + std.log.info("qwen27: slot {d} taps ring grew from {d} to {d} rows", .{ if (x.settle) |h| h.slot else 0, x.held, next }); + x.swapRing(fresh, next); + return fresh; + } + /// The ring back to its first rows once its stream left; nothing reads the old rows once `settle` returns. + fn shrinkRing(x: *Context) !void { + const s = x.step orelse return; + const first = kv_grow.first(s, x.span); + if (x.held <= first) return; + const stride = x.stride; + const fresh = try x.device.buffer(@as(usize, first) * stride, opts); + if (x.settle) |h| h.run(h.ctx, h.slot); + std.log.info("qwen27: slot {d} taps ring gave back {d} to {d} rows", .{ if (x.settle) |h| h.slot else 0, x.held, first }); + x.swapRing(fresh, first); + } + fn swapRing(x: *Context, fresh: mtl.Buffer, rows: u32) void { + if (x.residency) |f| { + f.leave(f.ctx, x.pending); + f.join(f.ctx, fresh); + } + x.pending.deinit(); + x.pending = fresh; + x.held = rows; + } pub fn setResident(x: *Context, hook: ?mtl.ResidencySet.Follow) void { if (x.residency) |f| x.apply(f, false); x.residency = hook; @@ -271,13 +346,14 @@ pub const Context = struct { } /// Empty, as Model.reset leaves the model's own. pub fn reset(x: *Context, c: cfg.Config) !void { + try x.shrinkRing(); x.committed_end = 0; x.session = try Session.init(c, 1, 0); x.context = .{ .cache = 1, .begin = 0, .end = 0 }; } /// Its taps ring, where the lane backend absorbs the stream's kept rows while it is parked. pub fn ring(x: *Context, c: cfg.Config) @import("../tap_ring.zig").Ring { - return .{ .buf = x.pending, .end = &x.committed_end, .window = c.window, .stride = @as(usize, c.tapWidth()) * 2 }; + return .{ .buf = x.pending, .end = &x.committed_end, .window = c.span(), .stride = @as(usize, c.ringWidth()) * 2, .grow = if (x.step != null) .{ .ctx = x, .run = growRing } else null }; } /// Its buffers, for a residency set. pub fn buffers(x: *const Context, list: *std.ArrayList(mtl.Buffer), a: std.mem.Allocator) !void { diff --git a/zig/src/families/qwen27/dflash/selector.zig b/zig/src/families/qwen27/dflash/selector.zig index c15dfb19ab..f6e0e9aaf1 100644 --- a/zig/src/families/qwen27/dflash/selector.zig +++ b/zig/src/families/qwen27/dflash/selector.zig @@ -84,7 +84,8 @@ fn expand(g: *const weights.Graph, lattice: op.Lattice, options: Options, first: } pub fn propose(gpa: std.mem.Allocator, g: *const weights.Graph, lattice: op.Lattice, anchor: u32, first_position: u64, options: Options) !Tree { try g.config.check(); - for ([_]weights.Tensor{ g.predecessor, g.successor }) |t| if (t.dtype != .bf16 or t.rank != 2 or t.shape[0] != g.config.vocab or t.shape[1] != g.config.rank) return error.DraftCodebookShape; + // a v1 draft has no codebooks (rank 0): its edges score nothing and its tree follows each depth's marginals + if (g.config.rank > 0) for ([_]weights.Tensor{ g.predecessor, g.successor }) |t| if (t.dtype != .bf16 or t.rank != 2 or t.shape[0] != g.config.vocab or t.shape[1] != g.config.rank) return error.DraftCodebookShape; try lattice.check(g.config); if (anchor >= g.config.vocab or options.nodes > 63 or options.children == 0 or options.children > 16 or lattice.topk > 16 or !std.math.isFinite(options.tau) or options.tau <= 0 or !std.math.isFinite(options.edge) or !std.math.isFinite(options.noise_weight) or !std.math.isFinite(options.temperature) or options.temperature <= 0) return error.BadDraftSelector; if (!std.math.isFinite(options.overhead_ms) or options.overhead_ms < 0) return error.BadDraftCosts; diff --git a/zig/src/families/qwen27/dflash/session.zig b/zig/src/families/qwen27/dflash/session.zig index 1756feb6b5..896e20abb9 100644 --- a/zig/src/families/qwen27/dflash/session.zig +++ b/zig/src/families/qwen27/dflash/session.zig @@ -5,7 +5,7 @@ const op = @import("operators.zig"); const weights = @import("weights.zig"); const execution = @import("execution.zig"); const selector = @import("selector.zig"); -pub const Batch = struct { values: op.View, positions: []const u64, layers: [5]u32, target_end: u64 }; +pub const Batch = struct { values: op.View, positions: []const u64, layers: [8]u32, target_end: u64 }; pub const Session = struct { config: cfg.Config, context: op.Context, @@ -17,11 +17,11 @@ pub const Session = struct { } pub fn admit(s: Session, batch: Batch) !op.Context { if (!std.mem.eql(u32, &batch.layers, &s.config.taps) or batch.positions.len == 0) return error.BadCommittedTaps; - try batch.values.check(@intCast(batch.positions.len), s.config.tapWidth(), .bf16); + try batch.values.check(@intCast(batch.positions.len), s.config.ringWidth(), .bf16); const end = try std.math.add(u64, s.context.end, batch.positions.len); if (end != batch.target_end) return error.BadCommittedTaps; for (batch.positions, 0..) |position, i| if (position != s.context.end + i or position >= batch.target_end) return error.BadCommittedTaps; - const begin = @max(s.context.begin, end -| (@as(u64, s.config.window) - 1)); + const begin = @max(s.context.begin, end -| (@as(u64, s.config.span()) - 1)); return .{ .cache = s.context.cache, .begin = begin, .end = end }; } pub fn absorb(s: *Session, ops: op.Ops, graph: *const weights.Graph, preparation: op.Preparation, batch: Batch) !void { diff --git a/zig/src/families/qwen27/dflash/weights.zig b/zig/src/families/qwen27/dflash/weights.zig index 570230cce5..de2931eee6 100644 --- a/zig/src/families/qwen27/dflash/weights.zig +++ b/zig/src/families/qwen27/dflash/weights.zig @@ -34,6 +34,11 @@ pub const Graph = struct { }; const suffixes = [_][]const u8{ "input_layernorm.weight", "post_attention_layernorm.weight", "self_attn.q_proj.weight", "self_attn.k_proj.weight", "self_attn.v_proj.weight", "self_attn.o_proj.weight", "self_attn.q_norm.weight", "self_attn.k_norm.weight", "mlp.gate_proj.weight", "mlp.up_proj.weight", "mlp.down_proj.weight", "attention_conv.kernel_projection.weight", "attention_conv.base_kernel", "mlp_conv.kernel_projection.weight", "mlp_conv.base_kernel" }; const globals = [_][]const u8{ "fc.weight", "hidden_norm.weight", "norm.weight", "candidate_selector.hidden_projection.weight", "candidate_selector.predecessor_codebook", "candidate_selector.successor_codebook" }; +/// A DFlash v1 draft's own tensors: the first eleven suffixes a layer, and fc, hidden_norm and norm. +const v1_suffixes = suffixes[0..11]; +const v1_globals = globals[0..3]; +/// A tensor a v1 draft does not have (its conv and selector slots): no bytes, never bound to a kernel. +pub const absent = Tensor{ .dtype = .bf16, .rank = 0, .shape = @splat(1), .bytes = &.{} }; pub fn names(gpa: std.mem.Allocator, c: Config) ![][]const u8 { const result = try gpa.alloc([]const u8, c.tensorCount()); var made: usize = 0; @@ -41,11 +46,12 @@ pub fn names(gpa: std.mem.Allocator, c: Config) ![][]const u8 { for (result[0..made]) |name| gpa.free(name); gpa.free(result); } - for (globals) |name| { + const v1 = c.variant == .v1; + for (if (v1) v1_globals else &globals) |name| { result[made] = try gpa.dupe(u8, name); made += 1; } - for (0..c.layers) |i| for (suffixes) |suffix| { + for (0..c.layers) |i| for (if (v1) v1_suffixes else &suffixes) |suffix| { result[made] = try std.fmt.allocPrint(gpa, "layers.{d}.{s}", .{ i, suffix }); made += 1; }; @@ -99,6 +105,7 @@ pub fn load(gpa: std.mem.Allocator, source: Inventory, c: Config, embed: ck.Line const src = source.source; const layers = try gpa.alloc(Layer, c.layers); errdefer gpa.free(layers); + if (c.variant == .v1) return loadV1(gpa, src, c, layers, embed); for (layers, 0..) |*l, i| l.* = .{ .input_norm = try layerTensor(src, i, "input_layernorm.weight", &.{c.hidden}), .post_norm = try layerTensor(src, i, "post_attention_layernorm.weight", &.{c.hidden}), @@ -110,6 +117,37 @@ pub fn load(gpa: std.mem.Allocator, source: Inventory, c: Config, embed: ck.Line return .{ .gpa = gpa, .config = c, .layers = layers, .fusion = try get(src, "fc.weight", &.{ c.hidden, c.tapWidth() }), .hidden_norm = try get(src, "hidden_norm.weight", &.{c.hidden}), .norm = try get(src, "norm.weight", &.{c.hidden}), .hidden_projection = try get(src, "candidate_selector.hidden_projection.weight", &.{ c.rank, c.hidden }), .predecessor = try get(src, "candidate_selector.predecessor_codebook", &.{ c.vocab, c.rank }), .successor = try get(src, "candidate_selector.successor_codebook", &.{ c.vocab, c.rank }), .embed = embed }; } +/// v1 layers: the qwen3 body alone; the conv and selector slots stay `absent`. +fn loadV1(gpa: std.mem.Allocator, src: ck.Source, c: Config, layers: []Layer, embed: ck.Linear) !Graph { + const none = Conv{ .projection = absent, .base = absent }; + for (layers, 0..) |*l, i| l.* = .{ + .input_norm = try layerTensor(src, i, "input_layernorm.weight", &.{c.hidden}), + .post_norm = try layerTensor(src, i, "post_attention_layernorm.weight", &.{c.hidden}), + .attention = .{ .q = try layerTensor(src, i, "self_attn.q_proj.weight", &.{ c.qWidth(), c.hidden }), .k = try layerTensor(src, i, "self_attn.k_proj.weight", &.{ c.kvWidth(), c.hidden }), .v = try layerTensor(src, i, "self_attn.v_proj.weight", &.{ c.kvWidth(), c.hidden }), .o = try layerTensor(src, i, "self_attn.o_proj.weight", &.{ c.hidden, c.qWidth() }), .q_norm = try layerTensor(src, i, "self_attn.q_norm.weight", &.{c.head_dim}), .k_norm = try layerTensor(src, i, "self_attn.k_norm.weight", &.{c.head_dim}) }, + .mlp = .{ .gate = try layerTensor(src, i, "mlp.gate_proj.weight", &.{ c.intermediate, c.hidden }), .up = try layerTensor(src, i, "mlp.up_proj.weight", &.{ c.intermediate, c.hidden }), .down = try layerTensor(src, i, "mlp.down_proj.weight", &.{ c.hidden, c.intermediate }) }, + .attention_conv = none, + .mlp_conv = none, + }; + return .{ .gpa = gpa, .config = c, .layers = layers, .fusion = try get(src, "fc.weight", &.{ c.hidden, c.tapWidth() }), .hidden_norm = try get(src, "hidden_norm.weight", &.{c.hidden}), .norm = try get(src, "norm.weight", &.{c.hidden}), .hidden_projection = absent, .predecessor = absent, .successor = absent, .embed = embed }; +} + +test "a DFlash v1 inventory is its sixty-nine qwen3 names, no conv or selector" { + const a = std.testing.allocator; + var c = Config{ .variant = .v1, .hidden = 4096, .intermediate = 12288, .layers = 6, .target_layers = 32, .block = 16, .mask = 248077, .window = 4096, .group = 0, .conv = 0, .rank = 0, .topk = 16, .taps = .{ 1, 5, 9, 13, 17, 21, 25, 29 }, .tap_count = 8 }; + c.full[5] = true; + try c.check(); + const list = try names(a, c); + defer freeNames(a, list); + try inventory(a, list, c); + try std.testing.expectEqual(@as(usize, 69), list.len); + for (list) |name| { + try std.testing.expect(std.mem.indexOf(u8, name, "conv") == null); + try std.testing.expect(std.mem.indexOf(u8, name, "selector") == null); + } + try std.testing.expectEqualStrings("fc.weight", list[0]); + try std.testing.expectEqualStrings("layers.5.mlp.down_proj.weight", list[list.len - 1]); +} + test "inventory is81exact names including unsuffixed BF16codebooks" { const a = std.testing.allocator; const list = try names(a, .{}); diff --git a/zig/src/families/qwen27/glue_gpu.zig b/zig/src/families/qwen27/glue_gpu.zig index 5fd13d67b1..6de19546c1 100644 --- a/zig/src/families/qwen27/glue_gpu.zig +++ b/zig/src/families/qwen27/glue_gpu.zig @@ -107,7 +107,7 @@ pub const Kernels = struct { } pub fn tapRows(k: Kernels, e: mtl.ComputeEncoder, x: Ref, kept: Ref, y: Ref, p: Gather) !void { - if (p.rows == 0 or p.rows > p.capacity or p.planes != 5 or p.width == 0) return error.BadTapGather; + if (p.rows == 0 or p.rows > p.capacity or p.planes == 0 or p.planes > 8 or p.width == 0) return error.BadTapGather; if (x.offset > x.buffer.length() or @as(usize, p.capacity) * p.width * p.planes * 2 > x.buffer.length() - x.offset or kept.offset > kept.buffer.length() or @as(usize, p.rows) * 4 > kept.buffer.length() - kept.offset or y.offset > y.buffer.length() or @as(usize, p.rows) * p.width * p.planes * 2 > y.buffer.length() - y.offset) return error.BadTapGather; e.setPipeline(k.gather); bind(e, .{ x, kept }); diff --git a/zig/src/families/qwen27/kv_grow.zig b/zig/src/families/qwen27/kv_grow.zig new file mode 100644 index 0000000000..dd69edc417 --- /dev/null +++ b/zig/src/families/qwen27/kv_grow.zig @@ -0,0 +1,46 @@ +//! Grow-as-used K/V (opt-in): a slot holds the rows its stream reached, not the whole context. +const std = @import("std"); + +/// One blit of a K/V buffer's head into a wider buffer, in bytes. +pub const Move = struct { from: usize, to: usize, len: usize }; + +/// TF_Q27_KV_GROW: the rows a slot starts with and grows by (a multiple of 16); unset or 0 keeps every slot full. +pub fn step(text: ?[]const u8) !?u32 { + const value = text orelse return null; + const n = std.fmt.parseInt(u32, value, 10) catch return error.BadKvGrowStep; + if (n == 0) return null; + if (n % 16 != 0) return error.BadKvGrowStep; + return n; +} + +/// A slot's first rows: one step, at most the context. +pub fn first(s: u32, context: u32) u32 { + return @min(s, context); +} + +/// Rows for `need`: the need rounded up to the step, or a quarter more than `held`, at most the context. +pub fn rows(held: u32, need: u32, s: u32, context: u32) u32 { + if (need <= held) return held; + const rounded = (@as(u64, need) + s - 1) / s * s; + const quarter = @as(u64, held) + held / 4; + return @intCast(@min(@max(rounded, quarter), context)); +} + +/// An empty slot's rows for a prompt of `need`: rounded up to the step, at least one step, at most the context. +pub fn admitted(need: u32, s: u32, context: u32) u32 { + const rounded = (@as(u64, need) + s - 1) / s * s; + return @intCast(@min(@max(rounded, s), context)); +} + +/// Head h of a [heads][old][dim] bf16 buffer to its place in a [heads][new][dim] one. +pub fn moves(out: []Move, heads: usize, dim: usize, old: u32, new: u32) []Move { + std.debug.assert(out.len >= heads and new >= old); + const row = dim * @sizeOf(u16); + for (out[0..heads], 0..) |*m, h| m.* = .{ .from = h * old * row, .to = h * new * row, .len = @as(usize, old) * row }; + return out[0..heads]; +} + +/// A slot's K and V bytes across `layers` attention layers at `r` rows of `kv` bf16 values. +pub fn slotBytes(layers: usize, kv: usize, r: u32) usize { + return layers * 2 * @as(usize, r) * kv * @sizeOf(u16); +} diff --git a/zig/src/families/qwen27/kv_slots.zig b/zig/src/families/qwen27/kv_slots.zig new file mode 100644 index 0000000000..349e92e704 --- /dev/null +++ b/zig/src/families/qwen27/kv_slots.zig @@ -0,0 +1,103 @@ +//! Grow-as-used K/V for a runner's slots (TF_Q27_KV_GROW): grow before rows are written, give back once a slot is free. +const std = @import("std"); +const mtl = @import("metal"); +const Runner = @import("decode_round.zig").Runner; +const gpuFailed = @import("decode_round.zig").gpuFailed; +const attn = @import("attention.zig"); +const kv_grow = @import("kv_grow.zig"); + +/// Slot `slot` holds `need` K/V rows: layer by layer, a wider pair takes every old byte at its place (waited blit). +pub fn reserve(r: *Runner, slot: u32, need: u32) !void { + const s = r.kv_step orelse return; + if (slot >= r.ready) return error.SlotNotReady; + const held = r.heldRows(slot); + if (need <= held) return; + if (need > r.capacity) return error.AttentionCapacity; + r.settleSlot(slot); + if (r.failed or r.active != null) return error.RoundNotReady; + const rows = kv_grow.rows(held, need, s, r.capacity); + const layers = r.caches.len / r.slots; + const bytes = @as(usize, rows) * r.model.config.kvDim() * 2; + // one layer at a time, so growth holds at most one layer's old pair beside the new rows + for (0..layers) |layer| { + const at = layer * r.slots + slot; + const c = r.caches[at]; + if (c.capacity >= rows) continue; // grown before an earlier failure + if (c.keys.offset != 0 or c.values.offset != 0 or c.key_stride != 0 or c.value_stride != 0) return error.KvGrowLayout; + var fresh: [2]mtl.Buffer = undefined; + fresh[0] = try r.model.device.buffer(bytes, mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked); + errdefer fresh[0].deinit(); + fresh[1] = try r.model.device.buffer(bytes, mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked); + errdefer fresh[1].deinit(); + const old = [2]mtl.Buffer{ c.keys.buffer, c.values.buffer }; + { + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const cb = r.model.queue.commandBuffer(); + const e = cb.blit(); + var moves: [attn.planning.max_streams]kv_grow.Move = undefined; + for (old, fresh) |from, to| { + for (kv_grow.moves(&moves, r.model.config.kv_heads, r.model.config.head_dim, c.capacity, rows)) |mv| e.copy(from, mv.from, to, mv.to, mv.len); + } + e.end(); + cb.commit(); + cb.wait(); + if (gpuFailed(cb, "a slot's K/V growth")) return error.RoundGpuFailure; + } + swapPair(r, at, fresh, rows); + } + std.log.info("qwen27: slot {d} K/V grew from {d} to {d} rows ({d} of {d} MiB held)", .{ slot, held, rows, kv_grow.slotBytes(layers, r.model.config.kvDim(), rows) >> 20, r.slotBytes() >> 20 }); +} + +/// Slot `slot` back to its first rows once no stream holds it and no blit reads it: its rows are gone (offset 0). +pub fn release(r: *Runner, slot: u32) !void { + const s = r.kv_step orelse return; + return refit(r, slot, kv_grow.first(s, r.capacity)); +} + +/// An emptied slot takes the rows a prompt of `need` tokens fills, at once and without a copy (a fresh pass). +pub fn admit(r: *Runner, slot: u32, need: u32) !void { + const s = r.kv_step orelse return; + if (need > r.capacity) return error.AttentionCapacity; + return refit(r, slot, kv_grow.admitted(need, s, r.capacity)); +} + +/// Every layer of empty slot `slot` holds exactly `first` rows: new buffers where it holds any other count. +fn refit(r: *Runner, slot: u32, first: u32) !void { + if (slot >= r.ready) return; + var held: u32 = 0; + for (0..r.caches.len / r.slots) |layer| held = @max(held, r.caches[layer * r.slots + slot].capacity); + if (held == first and r.heldRows(slot) == first) return; + r.settleSlot(slot); + if (r.failed) return error.RoundNotReady; + const bytes = @as(usize, first) * r.model.config.kvDim() * 2; + for (0..r.caches.len / r.slots) |layer| { + const at = layer * r.slots + slot; + const c = r.caches[at]; + if (c.capacity == first) continue; + var fresh: [2]mtl.Buffer = undefined; + fresh[0] = try r.model.device.buffer(bytes, mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked); + errdefer fresh[0].deinit(); + fresh[1] = try r.model.device.buffer(bytes, mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked); + swapPair(r, at, fresh, first); + } + r.offsets[slot] = 0; + std.log.info("qwen27: slot {d} K/V refit from {d} to {d} rows", .{ slot, held, first }); +} + +/// Cache `at` takes `fresh` (`rows` each); its old pair leaves the residency set and is freed. +fn swapPair(r: *Runner, at: usize, fresh: [2]mtl.Buffer, rows: u32) void { + const c = r.caches[at]; + for ([_]mtl.Buffer{ c.keys.buffer, c.values.buffer }, fresh) |from, to| { + for (r.owned.items) |*b| if (b.id == from.id) { + b.* = to; + }; + if (r.residency) |f| { + f.leave(f.ctx, from); + f.join(f.ctx, to); + } + from.deinit(); + } + r.caches[at] = .{ .keys = .{ .buffer = fresh[0] }, .values = .{ .buffer = fresh[1] }, .capacity = rows }; + if (r.grown) |g| g.added(g.ctx, &fresh); +} diff --git a/zig/src/families/qwen27/prompt_memory.zig b/zig/src/families/qwen27/prompt_memory.zig index 2a0c8b6752..af5cbe1a60 100644 --- a/zig/src/families/qwen27/prompt_memory.zig +++ b/zig/src/families/qwen27/prompt_memory.zig @@ -15,8 +15,21 @@ pub fn frameBytes(c: Config, rows: u32) usize { return buffers(frame.allocationSizes(c, rows)) + (if (rows > frame.small_rows) buffers(frame.decodeSizes(c, rows)) else 0); } +/// A prompt row's tap bytes: `planes` raw planes, plus a fusing drafter's row-major copy, fc rows and fused row. +pub fn tapRowBytes(hidden: usize, planes: usize, fused: bool) usize { + return (planes + if (fused) planes + 2 else 0) * hidden * 2; +} +test "qwen27 memory: an eight-tap fusing drafter's prompt taps cost what its planes and fusion buffers hold" { + try std.testing.expectEqual(@as(usize, 5 * 5120 * 2), tapRowBytes(5120, 5, false)); + try std.testing.expectEqual(@as(usize, 18 * 4096 * 2), tapRowBytes(4096, 8, true)); +} + /// Frame bytes already belong to the model; these are the runner's future bytes at its peak. pub fn runnerBytes(c: Config, context: u32, slots: u32, rows: u32, shared: bool, ring: usize) !usize { + return runnerBytesTaps(c, context, slots, rows, shared, ring, tapRowBytes(c.hidden, 5, false)); +} +/// runnerBytes with the drafter's own tap row bytes (tapRowBytes). +pub fn runnerBytesTaps(c: Config, context: u32, slots: u32, rows: u32, shared: bool, ring: usize, tap_row: usize) !usize { var attention: usize = 0; for (0..c.layers) |i| attention += @intFromBool(c.kind(i) == .attention); const linear = c.layers - attention; @@ -27,7 +40,7 @@ pub fn runnerBytes(c: Config, context: u32, slots: u32, rows: u32, shared: bool, const kv = attention * @as(usize, slots) * buffers([_]usize{ @as(usize, context) * c.kvDim() * 2, @as(usize, context) * c.kvDim() * 2 }); const bindings = attention * buffers([_]usize{ (planning.global_words + planning.max_streams * planning.stream_words) * 4, planning.max_streams * 16, rows * 16 }); const before = if (slots > 1) buffers([_]usize{budget.state}) else 0; - const taps = buffers([_]usize{@as(usize, rows) * 5 * c.hidden * 2}); + const taps = buffers([_]usize{@as(usize, rows) * tap_row}); return kv + bindings + before + taps + buffers(scratch.sizes) + buffers(budget.allocationSizes()) + slots * (ring + (256 << 20)); } @@ -37,11 +50,12 @@ pub const Fit = struct { slots: u32, shared: bool, ring: usize = 0, + tap_row: usize = 0, // tapRowBytes for the drafter's taps; 0: five raw planes current_rows: u32, left: usize, pub fn cost(f: Fit, rows: u32, slots: u32) !usize { - return (try runnerBytes(f.config, f.context, slots, rows, f.shared, f.ring)) + (frameBytes(f.config, rows) -| frameBytes(f.config, f.current_rows)); + return (try runnerBytesTaps(f.config, f.context, slots, rows, f.shared, f.ring, if (f.tap_row > 0) f.tap_row else tapRowBytes(f.config.hidden, 5, false))) + (frameBytes(f.config, rows) -| frameBytes(f.config, f.current_rows)); } pub fn choose(f: Fit, requested: u32, explicit: bool) !u32 { diff --git a/zig/src/families/qwen27/qwen27.zig b/zig/src/families/qwen27/qwen27.zig index 9475212464..097a283ffb 100644 --- a/zig/src/families/qwen27/qwen27.zig +++ b/zig/src/families/qwen27/qwen27.zig @@ -26,6 +26,7 @@ pub const gpu_frame = @import("gpu_frame.zig"); pub const model = @import("model.zig"); pub const projection = @import("projection.zig"); pub const prompt_memory = @import("prompt_memory.zig"); +pub const kv_grow = @import("kv_grow.zig"); pub const prompt_rows = @import("prompt_rows.zig"); pub const prompt_glue = @import("prompt_glue.zig"); pub const prompt_gpu = @import("prompt_gpu.zig"); diff --git a/zig/src/families/qwen27/residency.zig b/zig/src/families/qwen27/residency.zig index bb14753233..32daac2472 100644 --- a/zig/src/families/qwen27/residency.zig +++ b/zig/src/families/qwen27/residency.zig @@ -27,6 +27,7 @@ pub fn setRunner(r: anytype, hook: ?mtl.ResidencySet.Follow) void { pub fn appendDraft(b: anytype, a: std.mem.Allocator, list: *std.ArrayList(mtl.Buffer)) !void { try fields.append(a, list, &b.attention); + if (b.attention_full) |*x| try fields.append(a, list, x); try fields.append(a, list, &b.topk); try list.appendSlice(a, b.head.owned.items); try list.appendSlice(a, b.frame.pool.items); @@ -42,6 +43,7 @@ pub fn appendDraft(b: anytype, a: std.mem.Allocator, list: *std.ArrayList(mtl.Bu fn draftBuffers(b: anytype, hook: mtl.ResidencySet.Follow, joining: bool) void { fields.apply(&b.attention, hook, joining); + if (b.attention_full) |*x| fields.apply(x, hook, joining); fields.apply(&b.topk, hook, joining); for (b.head.owned.items) |buffer| fields.one(buffer, hook, joining); for (b.caches) |*c| fields.apply(c, hook, joining); diff --git a/zig/src/families/qwen27/snapshot.zig b/zig/src/families/qwen27/snapshot.zig index d99a2c3da4..731bd7c100 100644 --- a/zig/src/families/qwen27/snapshot.zig +++ b/zig/src/families/qwen27/snapshot.zig @@ -243,7 +243,10 @@ pub fn restoreSlot(a: std.mem.Allocator, r: *Runner, slot: u32, ring: ?Ring, st: if (st.bytes != slotBytes(r, ring, st.at) or st.at > r.capacity) return error.SnapshotLayout; try st.landed(); try r.reset(slot); - try copy(a, r, slot, ring, st, false); + try r.reserve(slot, st.at); // a slot whose K/V and ring grow as used holds the state's rows first + var held = ring; + if (held) |*g| try g.hold(st.at); + try copy(a, r, slot, held, st, false); r.offsets[slot] = st.at; if (ring) |g| g.end.* = st.draft_end; } @@ -255,12 +258,15 @@ pub fn moveSlot(a: std.mem.Allocator, r: *Runner, from: u32, from_ring: ?Ring, t if (from == to or r.failed or r.active != null or from >= r.slots or to >= r.slots or (from_ring == null) != (to_ring == null)) return error.BadSlotMove; if (from_ring) |g| if (g.end.* != at) return error.NotAtMark; try r.reset(to); + try r.reserve(to, at); + var to_held = to_ring; + if (to_held) |*g| try g.hold(at); var src: std.ArrayList(Part) = .empty; defer src.deinit(a); var dst: std.ArrayList(Part) = .empty; defer dst.deinit(a); try parts(r, from, from_ring, at, &src, a, false); - try parts(r, to, to_ring, at, &dst, a, false); + try parts(r, to, to_held, at, &dst, a, false); if (src.items.len != dst.items.len) return error.SnapshotLayout; const pool = mtl.objc.Pool.push(); defer pool.pop(); diff --git a/zig/src/families/qwen27/tap_gather.zig b/zig/src/families/qwen27/tap_gather.zig index 03dcd5bc54..73a6668a02 100644 --- a/zig/src/families/qwen27/tap_gather.zig +++ b/zig/src/families/qwen27/tap_gather.zig @@ -1,15 +1,20 @@ //! Completed shared-memory tap planes gathered on the host without an extra GPU submission. const std = @import("std"); pub fn gather(planes: []const u8, capacity: usize, written: usize, width: usize, rows: []const u32, out: []u8) !void { + return gatherPlanes(5, planes, capacity, written, width, rows, out); +} +/// `count` planes of `capacity` rows each, row `rows[i]` of every plane side by side in out's row i. +pub fn gatherPlanes(count: usize, planes: []const u8, capacity: usize, written: usize, width: usize, rows: []const u32, out: []u8) !void { + if (count == 0 or count > 8) return error.BadTapGather; const bytes = try std.math.mul(usize, width, 2); - const stride = try std.math.mul(usize, bytes, 5); + const stride = try std.math.mul(usize, bytes, count); const source_size = try std.math.mul(usize, capacity, stride); const output_size = try std.math.mul(usize, rows.len, stride); if (planes.len < source_size or out.len < output_size) return error.ShortTapGather; if (written > capacity or rows.len > written) return error.BadTapGather; for (rows) |row| if (row >= written) return error.BadTapGather; for (rows, 0..) |row, i| { - for (0..5) |plane| { + for (0..count) |plane| { const src = (plane * capacity + row) * bytes; @memcpy(out[i * stride + plane * bytes ..][0..bytes], planes[src..][0..bytes]); } diff --git a/zig/src/families/qwen27/tap_ring.zig b/zig/src/families/qwen27/tap_ring.zig index 8c61a3faf8..d68e419d4c 100644 --- a/zig/src/families/qwen27/tap_ring.zig +++ b/zig/src/families/qwen27/tap_ring.zig @@ -1,30 +1,52 @@ //! A stream's committed target taps over the drafter's window, by position modulo the window. const std = @import("std"); const mtl = @import("metal"); +const kv_grow = @import("kv_grow.zig"); pub const Ring = struct { buf: mtl.Buffer, end: *u64, // taps committed so far (the stream's fed tokens) window: u32, stride: usize, // one row's bytes + grow: ?Grow = null, // a ring that holds the rows its stream reached: the buffer holding `rows`; null: the window + + /// The buffer holding at least `rows` rows (it may be a new one; the old rows stay at their places). + pub const Grow = struct { ctx: *anyopaque, run: *const fn (ctx: *anyopaque, rows: u32) anyerror!mtl.Buffer }; /// Rows [off, off + rows) of `taps` follow the committed ones. pub fn absorb(r: Ring, taps: mtl.Buffer, off: usize, rows: u32) !void { const size = @as(usize, rows) * r.stride; if (rows == 0 or off > taps.length() or size > taps.length() - off) return error.BadCommittedTaps; const end = try std.math.add(u64, r.end.*, rows); + const buf = if (r.grow) |g| try g.run(g.ctx, @intCast(@min(end, r.window))) else r.buf; var first: usize = 0; while (first < rows) { const slot: usize = @intCast((r.end.* + first) % r.window); const count = @min(rows - first, r.window - slot); - @memcpy(r.buf.contents()[slot * r.stride ..][0 .. count * r.stride], taps.contents()[off + first * r.stride ..][0 .. count * r.stride]); + @memcpy(buf.contents()[slot * r.stride ..][0 .. count * r.stride], taps.contents()[off + first * r.stride ..][0 .. count * r.stride]); first += count; } r.end.* = end; } + /// The ring holds a state of `rows` tokens before a blit writes it (a restore); its buffer may change. + pub fn hold(r: *Ring, rows: u64) !void { + if (r.grow) |g| r.buf = try g.run(g.ctx, @intCast(@min(rows, r.window))); + } + /// Ring bytes a state after `at` tokens needs: slots [0, at) until the ring wraps, then all of it. pub fn bytes(r: Ring, at: u64) usize { return @as(usize, @intCast(@min(at, r.window))) * r.stride; } }; + +/// Rows a growing ring takes for `need`: the K/V growth rule, at most the window. +pub fn grownRows(held: u32, need: u32, step: u32, window: u32) u32 { + return kv_grow.rows(held, need, step, window); +} + +/// A grown ring's first `held` rows into its new buffer, at the same places (the CPU writes rings). +pub fn carry(old: mtl.Buffer, new: mtl.Buffer, held: u32, stride: usize) void { + const n = @as(usize, held) * stride; + @memcpy(new.contents()[0..n], old.contents()[0..n]); +} diff --git a/zig/src/families/qwen27/taps.zig b/zig/src/families/qwen27/taps.zig index 23e04b5751..73a363e0b0 100644 --- a/zig/src/families/qwen27/taps.zig +++ b/zig/src/families/qwen27/taps.zig @@ -5,13 +5,25 @@ const Ref = @import("projection.zig").Ref; const Glue = @import("glue_gpu.zig").Kernels; pub const layers = [_]usize{ 5, 19, 33, 47, 61 }; +/// The most tap planes a drafter may name (a DFlash v1 draft names up to eight). +pub const max_planes = 8; + +/// A drafter's fusion of a row's taps (fc then hidden_norm), encoded where the taps complete. +pub const Fuser = struct { ptr: *anyopaque, width: u32, encode: *const fn (*anyopaque, mtl.ComputeEncoder, Ref, Ref, u32) anyerror!void }; +/// Where accepted rows gather from: the raw planes, or one plane of fused rows. +pub const Source = struct { buffer: mtl.Buffer, width: u32, planes: u32 }; pub const Taps = struct { buffer: mtl.Buffer, + fuser: ?Fuser = null, + fused: ?mtl.Buffer = null, // [capacity][fuser width] + rowmajor: ?mtl.Buffer = null, // [capacity][count * width]: a row's planes side by side, fc's input + identity: ?mtl.Buffer = null, // 0, 1, ... capacity - 1: every row, for the row-major gather capacity: u32, width: u32, seen: u8 = 0, - ids: [5]usize = layers, + ids: [max_planes]usize = .{ 5, 19, 33, 47, 61, 0, 0, 0 }, + count: u32 = layers.len, rows: u32 = 0, pub fn init(device: mtl.Device, capacity: u32, width: u32) !Taps { @@ -21,9 +33,59 @@ pub const Taps = struct { } pub fn deinit(t: Taps) void { + t.dropFused(); t.buffer.deinit(); } + fn dropFused(t: Taps) void { + if (t.fused) |b| b.deinit(); + if (t.rowmajor) |b| b.deinit(); + if (t.identity) |b| b.deinit(); + } + + /// Completed taps fused for a drafter whose ring keeps fused rows; null: planes stay raw. + pub fn setFuser(t: *Taps, device: mtl.Device, f: ?Fuser) !void { + t.dropFused(); + t.fuser = null; + t.fused = null; + t.rowmajor = null; + t.identity = null; + const fuser = f orelse return; + if (fuser.width == 0) return error.BadTapShape; + const options = mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked; + const fused = try device.buffer(@as(usize, t.capacity) * fuser.width * 2, options); + errdefer fused.deinit(); + const rowmajor = try device.buffer(@as(usize, t.capacity) * t.count * t.width * 2, options); + errdefer rowmajor.deinit(); + const identity = try device.buffer(@as(usize, t.capacity) * 4, options); + for (identity.slice(u32, t.capacity), 0..) |*x, i| x.* = @intCast(i); + t.fuser = fuser; + t.fused = fused; + t.rowmajor = rowmajor; + t.identity = identity; + } + + /// The rows accepted taps gather from. + pub fn source(t: Taps) Source { + if (t.fused) |f| return .{ .buffer = f, .width = t.fuser.?.width, .planes = 1 }; + return .{ .buffer = t.buffer, .width = t.width, .planes = t.count }; + } + + /// The drafter's tap layers; planes are re-made only when their count changes. + pub fn setLayers(t: *Taps, device: mtl.Device, ids: []const u32) !void { + if (ids.len == 0 or ids.len > max_planes or (t.fuser != null and ids.len != t.count)) return error.BadTapShape; + if (ids.len != t.count) { + const bytes = try std.math.mul(usize, try std.math.mul(usize, t.capacity, t.width), 2 * ids.len); + const buffer = try device.buffer(bytes, mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked); + t.buffer.deinit(); + t.buffer = buffer; + t.count = @intCast(ids.len); + } + t.ids = @splat(0); + for (ids, t.ids[0..ids.len]) |id, *to| to.* = id; + t.begin(); + } + pub fn begin(t: *Taps) void { t.seen = 0; t.rows = 0; @@ -31,7 +93,7 @@ pub const Taps = struct { pub fn record(t: *Taps, glue: Glue, e: mtl.ComputeEncoder, layer: usize, hidden: Ref, rows: u32) !void { if (rows == 0 or rows > t.capacity) return error.BadTapShape; - const index = std.mem.indexOfScalar(usize, &t.ids, layer) orelse return; + const index = std.mem.indexOfScalar(usize, t.ids[0..t.count], layer) orelse return; if (t.seen != 0 and rows != t.rows) return error.BadTapShape; const bit = @as(u8, 1) << @intCast(index); if (t.seen & bit != 0) return error.DuplicateTap; @@ -39,15 +101,22 @@ pub const Taps = struct { e.barrier(); t.seen |= bit; t.rows = rows; + if (t.fuser) |f| if (t.seen == fullMask(t.count)) { + const kept = Ref{ .buffer = t.identity.? }; + try glue.tapRows(e, .{ .buffer = t.buffer }, kept, .{ .buffer = t.rowmajor.? }, .{ .rows = rows, .width = t.width, .capacity = t.capacity, .planes = t.count }); + e.barrier(); + try f.encode(f.ptr, e, .{ .buffer = t.rowmajor.? }, .{ .buffer = t.fused.? }, rows); + e.barrier(); + }; } pub fn plane(t: Taps, index: usize) !Ref { - if (index >= layers.len) return error.BadTapIndex; + if (index >= t.count) return error.BadTapIndex; return .{ .buffer = t.buffer, .offset = index * @as(usize, t.capacity) * t.width * 2 }; } pub fn complete(t: Taps) !void { - if (t.seen != (1 << layers.len) - 1) return error.IncompleteTaps; + if (t.seen != fullMask(t.count)) return error.IncompleteTaps; } pub fn accepted(t: Taps, glue: Glue, e: mtl.ComputeEncoder, path: []const u32, kept: Ref, out: Ref) !void { @@ -55,11 +124,21 @@ pub const Taps = struct { if (path.len == 0) return; if (path.len > t.rows) return error.BadTapGather; for (path) |row| if (row >= t.rows) return error.BadTapGather; - try glue.tapRows(e, .{ .buffer = t.buffer }, kept, out, .{ .rows = @intCast(path.len), .width = t.width, .capacity = t.capacity, .planes = layers.len }); + const s = t.source(); + try glue.tapRows(e, .{ .buffer = s.buffer }, kept, out, .{ .rows = @intCast(path.len), .width = s.width, .capacity = t.capacity, .planes = s.planes }); e.barrier(); } }; +fn fullMask(count: u32) u8 { + return @intCast((@as(u16, 1) << @intCast(count)) - 1); +} + +test "eight tap planes complete when all eight are seen" { + try std.testing.expectEqual(@as(u8, 0xff), fullMask(8)); + try std.testing.expectEqual(@as(u8, 0x1f), fullMask(5)); +} + test "DFlash tap ids name completed target blocks" { try std.testing.expectEqualSlices(usize, &.{ 5, 19, 33, 47, 61 }, &layers); } diff --git a/zig/src/native/qwen27_draft.zig b/zig/src/native/qwen27_draft.zig index 7c4886529b..ba2830414f 100644 --- a/zig/src/native/qwen27_draft.zig +++ b/zig/src/native/qwen27_draft.zig @@ -10,7 +10,8 @@ pub const Draft = struct { owns_model: bool, generation: df.generation.Generation, previous_capture: bool, - previous_taps: [5]usize, + previous_taps: [q.taps.max_planes]usize, + previous_count: u32, // the planes the buffer holds stay at least this many: restoring fewer is safe pub fn open(gpa: std.mem.Allocator, io: std.Io, target: *q.model.Model, runner: *q.decode_round.Runner, dir: []const u8, mode: df.operators.Mode) !*Draft { if (runner.model != target) return error.TargetBinding; @@ -28,12 +29,14 @@ pub const Draft = struct { errdefer gpa.destroy(d); const previous_capture = runner.capture_taps; const previous_taps = runner.taps.ids; + const previous_count = runner.taps.count; const generation = df.generation.Generation.init(gpa, runner, model) catch |err| { runner.capture_taps = previous_capture; runner.taps.ids = previous_taps; + if (runner.taps.count > previous_count) runner.taps.count = previous_count; return err; }; - d.* = .{ .gpa = gpa, .runner = runner, .model = model, .owns_model = owns_model, .generation = generation, .previous_capture = previous_capture, .previous_taps = previous_taps }; + d.* = .{ .gpa = gpa, .runner = runner, .model = model, .owns_model = owns_model, .generation = generation, .previous_capture = previous_capture, .previous_taps = previous_taps, .previous_count = previous_count }; model.backend.profile(runner.trace); // the target's TF_Q27_PROFILE trace times the drafter's passes too return d; } @@ -67,7 +70,9 @@ pub const Draft = struct { defer pool.pop(); d.generation.deinit(); d.runner.capture_taps = d.previous_capture; + d.runner.taps.setFuser(d.runner.model.device, null) catch unreachable; // dropping a fuser allocates nothing d.runner.taps.ids = d.previous_taps; + if (d.runner.taps.count > d.previous_count) d.runner.taps.count = d.previous_count; d.model.backend.profile(null); // a borrowed drafter outlives the host and its trace if (d.owns_model) d.model.deinit(); d.gpa.destroy(d); diff --git a/zig/src/native/qwen27_keep.zig b/zig/src/native/qwen27_keep.zig index cb335235c1..f99690b2f5 100644 --- a/zig/src/native/qwen27_keep.zig +++ b/zig/src/native/qwen27_keep.zig @@ -42,7 +42,8 @@ const Commit = struct { var ended = false; errdefer if (!ended) e.end(); try r.keepDevice(&b.round.?, e, g.commit, ref); - if (!ready) try g.commit.taps(e, ref, .{ .buf = r.taps.buffer }, .{ .buf = g.taps }, .{ .rows = rows, .width = r.taps.width, .capacity = r.taps.capacity, .planes = 5 }); + const src = r.taps.source(); + if (!ready) try g.commit.taps(e, ref, .{ .buf = src.buffer }, .{ .buf = g.taps }, .{ .rows = rows, .width = src.width, .capacity = r.taps.capacity, .planes = src.planes }); if (!ready) try g.commit.head(e, ref, .{ .buf = logits.buffer, .off = logits.offset }, .{ .rows = rows, .width = vocab, .stride = vocab }); e.end(); ended = true; @@ -84,7 +85,8 @@ pub fn encodeTail(ptr: *anyopaque, e: mtl.ComputeEncoder) !void { const shape = try g.matcher.encode(e, .{ .buf = logits.buffer, .off = logits.offset }, b.round.?.ids, b.round.?.windows[0].parents, vocab, contract.max_rows, &.{}); const ref = tf.tree_commit_gpu.Ref{ .buf = g.matcher.buffers[4] }; try r.taps.complete(); - try g.commit.taps(e, ref, .{ .buf = r.taps.buffer }, .{ .buf = g.taps }, .{ .rows = shape.rows, .width = r.taps.width, .capacity = r.taps.capacity, .planes = 5 }); + const src = r.taps.source(); + try g.commit.taps(e, ref, .{ .buf = src.buffer }, .{ .buf = g.taps }, .{ .rows = shape.rows, .width = src.width, .capacity = r.taps.capacity, .planes = src.planes }); const metric = mtl.pass_metrics.phase(.keep); defer metric.end(); std.debug.assert(b.prepared_keep == null); diff --git a/zig/src/native/qwen27_lanes.zig b/zig/src/native/qwen27_lanes.zig index d55d4e6801..b27183c9c9 100644 --- a/zig/src/native/qwen27_lanes.zig +++ b/zig/src/native/qwen27_lanes.zig @@ -88,6 +88,7 @@ pub const Metal = struct { standing: @import("qwen27_resume.zig").Standing = .{}, // what slot 0 stands at once its stream left (#1197) rewind: ?mtl.Buffer = null, // slot 0's DeltaNet states and taps at its last pass's history end slot_resume: bool = true, // TF_QWEN27_SLOT_RESUME: a later prompt resumes where slot 0 stands (0: off) + give_back: [max_streams]bool = @splat(false), // a free slot's grown K/V and ring wait to go back (slots.giveBack) pub fn init(gpa: std.mem.Allocator, runner: *q.decode_round.Runner, draft: ?*Draft) !*Metal { if (runner.slots == 0 or runner.slots > max_streams or runner.model.frame.capacity < batch_rows) return error.TargetBinding; @@ -109,7 +110,10 @@ pub const Metal = struct { if (draft) |d| { const g = d.model.ring(); b.gathered_taps = try device.buffer(batch_rows * g.stride, opts); - for (b.contexts[1..runner.ready], 1..) |*x, slot| x.* = try Context.init(gpa, d.model, slot < b.draft_slots); + for (b.contexts[1..runner.ready], 1..) |*x, slot| { + x.* = try Context.init(gpa, d.model, slot < b.draft_slots, runner.kv_step); + x.*.?.settle = .{ .ctx = runner, .slot = @intCast(slot), .run = q.decode_round.Runner.settleOpaque }; + } } return b; } @@ -335,7 +339,8 @@ pub const Metal = struct { std.debug.assert(r.active == null); if (r.active != null) return error.RoundNotReady; try r.taps.complete(); - try q.tap_gather.gather(r.taps.buffer.contents()[0..r.taps.buffer.length()], r.taps.capacity, r.taps.rows, r.taps.width, rows[0..n], out.contents()[0..out.length()]); + const src = r.taps.source(); + try q.tap_gather.gatherPlanes(src.planes, src.buffer.contents()[0..src.buffer.length()], r.taps.capacity, r.taps.rows, src.width, rows[0..n], out.contents()[0..out.length()]); var at: usize = 0; for (paths, b.windows[0..b.count]) |path, slot| { const g = b.ringOf(slot).?; diff --git a/zig/src/native/qwen27_served.zig b/zig/src/native/qwen27_served.zig index e8f8111df0..6969b5de1a 100644 --- a/zig/src/native/qwen27_served.zig +++ b/zig/src/native/qwen27_served.zig @@ -78,6 +78,12 @@ pub const Host = struct { std.log.info("Qwen3.8-27B: a slot took its memory ({d} of {d} ready)", .{ h.runner.ready, h.runner.slots }); } + /// A slot's grown K/V buffers join the residency set now, as a new slot's do. + fn grownKv(ptr: *anyopaque, buffers: []const mtl.Buffer) void { + const h = of(ptr); + if (h.resident) |*r| r.addNow(buffers); + } + /// A finished reply's state stands where a prompt pass of prompt and reply would: keep it for the next turn. fn keepReply(ptr: *anyopaque, s: *lanes.Stream, slot: u32) void { const h = of(ptr); @@ -302,7 +308,7 @@ pub fn open(gpa: Allocator, io: std.Io, dir: []const u8, o: Options, a: Allocato /// Why --drafter is refused before the target or the drafter loads. fn drafterRefusal(a: Allocator, dir: []const u8, drafter: []const u8, err: anyerror) ![]const u8 { if (err == error.DraftForAnotherTarget) return std.fmt.allocPrint(a, "--drafter {s} drafts for another target: its width, vocabulary or tapped layers do not fit {s} (DraftForAnotherTarget)", .{ drafter, dir }); - return std.fmt.allocPrint(a, "--drafter {s} is not a drafter the native Qwen engine serves for {s} ({s}): it takes z-lab/Qwen3.8-27B-DFlash2 as released, for Qwen3.8-27B", .{ drafter, dir, @errorName(err) }); + return std.fmt.allocPrint(a, "--drafter {s} is not a drafter the native Qwen engine serves for {s} ({s}): it takes z-lab/Qwen3.8-27B-DFlash2 as released for Qwen3.8-27B, or a DFlash drafter (DFlashDraftModel) made for this target", .{ drafter, dir, @errorName(err) }); } /// A loaded model served as `open` serves one, `drafter` bound in place of `o.drafter`; both stay the caller's. @@ -332,14 +338,14 @@ pub fn attach(parent: Allocator, io: std.Io, model: *q.model.Model, drafter: ?*q const gpa = if (profiling) h.counted.allocator() else parent; var draft_model = drafter; const owns_draft = drafter == null and o.drafter != null; - if (owns_draft) draft_model = q.dflash.runtime_model.Model.load(gpa, io, model, o.drafter.?, if (o.drafter_bits == 4) .prepared_q4_reference else .bf16_reference) catch |err| { + if (owns_draft) draft_model = q.dflash.runtime_model.Model.loadSpan(gpa, io, model, o.drafter.?, if (o.drafter_bits == 4) .prepared_q4_reference else .bf16_reference, capacity + window) catch |err| { why.* = try std.fmt.allocPrint(a, "the native Qwen3.8-27B drafter cannot load {s} ({s})", .{ o.drafter.?, @errorName(err) }); return error.DrafterLoad; }; var draft_owned = owns_draft; errdefer if (draft_owned) draft_model.?.deinit(); const memory = q.prompt_memory; - const fit_prompt = memory.Fit{ .config = model.config, .context = capacity + window, .slots = @max(o.streams, 1), .shared = model.device.tensorUnits(), .ring = if (draft_model) |d| @as(usize, d.ring().window) * d.ring().stride else 0, .current_rows = model.frame.capacity, .left = room(model) }; + const fit_prompt = memory.Fit{ .config = model.config, .context = capacity + window, .slots = @max(o.streams, 1), .shared = model.device.tensorUnits(), .ring = if (draft_model) |d| @as(usize, d.ring().window) * d.ring().stride else 0, .tap_row = if (draft_model) |d| memory.tapRowBytes(model.config.hidden, d.graph.config.tap_count, d.graph.config.variant == .v1) else 0, .current_rows = model.frame.capacity, .left = room(model) }; const value = if (std.c.getenv("TF_PROMPT_ROWS")) |v| std.mem.span(v) else null; const row_policy = q.prompt_rows; const rows = @max(window, try fit_prompt.choose(row_policy.requested(model.device.generation(), value), row_policy.override(value) != null)); @@ -352,7 +358,7 @@ pub fn attach(parent: Allocator, io: std.Io, model: *q.model.Model, drafter: ?*q model.frame.deinit(); model.frame = next; } - std.log.info("qwen27: prompt capacity={d} rows, startup stream and scratch reserve={d} bytes", .{ rows, try memory.runnerBytes(model.config, capacity + window, slots, rows, model.device.tensorUnits(), fit_prompt.ring) }); + std.log.info("qwen27: prompt capacity={d} rows, startup stream and scratch reserve={d} bytes", .{ rows, try memory.runnerBytesTaps(model.config, capacity + window, slots, rows, model.device.tensorUnits(), fit_prompt.ring, if (fit_prompt.tap_row > 0) fit_prompt.tap_row else memory.tapRowBytes(model.config.hidden, 5, false)) }); h.gpa = gpa; h.model = model; h.owns_model = false; @@ -362,7 +368,12 @@ pub fn attach(parent: Allocator, io: std.Io, model: *q.model.Model, drafter: ?*q h.disk_stage = null; h.vision = o.images; // closed with the host; a failed attach leaves it to the caller h.trace = null; - h.runner = try q.decode_round.Runner.initReady(gpa, model, slots, capacity + window, 1); + const kv_step = q.kv_grow.step(if (std.c.getenv("TF_Q27_KV_GROW")) |v| std.mem.span(v) else null) catch |err| { + why.* = "TF_Q27_KV_GROW must be 0 (off) or a row count that is a multiple of 16"; + return err; + }; + h.runner = try q.decode_round.Runner.initGrow(gpa, model, slots, capacity + window, 1, kv_step); + if (kv_step) |s| std.log.info("qwen27: K/V grows as used: a slot starts at {d} rows and grows by {d} or a quarter, up to {d}", .{ q.kv_grow.first(s, capacity + window), s, capacity + window }); errdefer h.runner.deinit(); if (profileEvery()) |every| { h.trace = .{ .queue = model.queue, .report_every = every }; // enabled once measure() has priced the windows @@ -375,6 +386,7 @@ pub fn attach(parent: Allocator, io: std.Io, model: *q.model.Model, drafter: ?*q errdefer h.back.deinit(); h.back.chunk = h.runner.promptRows(@min(rows, row_policy.requested(model.device.generation(), value))); h.back.grow = .{ .ctx = h, .room = Host.roomFor, .added = Host.added }; + h.runner.grown = .{ .ctx = h, .added = Host.grownKv }; h.back.vision = o.images; if (h.draft != null) try h.back.measure(io); h.cfg = try lanes.Config.init(gpa, h.back.facts(), window, h.back.nodes); @@ -383,7 +395,7 @@ pub fn attach(parent: Allocator, io: std.Io, model: *q.model.Model, drafter: ?*q h.core = lanes.Engine.init(gpa, &h.cfg, h.back.backend(), h.clock.clock()); errdefer h.core.deinit(); const offer: ?api.Vision = if (o.images) |v| .{ .ctx = h, .prepare = Host.prepareImage, .image_token = v.c.image_token, .markers = &v.markers, .image_tokens = v.capacity, .max_images = v.max_images, .explain = tf.qwen_vl.geometry.problem } else null; - h.host = api.LaneHost.init(gpa, io, &h.core, .{ .name = "qwen27", .draft_sources = .{ .drafter = if (h.draft != null) .{ .kind = "DFlash2", .name = o.drafter_name orelse o.drafter orelse "" } else null }, .structures = @import("qwen27_masks.zig").supported, .lanes = slots, .context_window = capacity, .prefill_step = @intCast(h.back.chunk), .decode_share = o.decode_share, .vision = offer }); + h.host = api.LaneHost.init(gpa, io, &h.core, .{ .name = "qwen27", .draft_sources = .{ .drafter = if (h.draft) |d| .{ .kind = d.model.graph.config.kind(), .name = o.drafter_name orelse o.drafter orelse "" } else null }, .structures = @import("qwen27_masks.zig").supported, .lanes = slots, .context_window = capacity, .prefill_step = @intCast(h.back.chunk), .decode_share = o.decode_share, .vision = offer }); h.host.memory = .{ .ctx = h, .read = readMemory }; h.host.explain = .{ .text = words }; h.host.decoded_rows = true; // a reply's rows prefill with the bits decoding gave them, with the cache on or off @@ -414,7 +426,9 @@ pub fn attach(parent: Allocator, io: std.Io, model: *q.model.Model, drafter: ?*q h.host.keepalive_target = .{ .ctx = &h.warm, .tick = mtl.keepalive.Target.tick }; if (h.trace) |*t| t.enabled = true; // serving rounds count from here: measure() timed whole, unsplit buffers try h.host.start(); - std.log.info("Qwen3.8-27B: {d} streams of {d} tokens{s}{s}", .{ slots, capacity, if (h.draft != null) ", DFlash2 drafts for one of them" else "", if (o.images != null) ", images read" else "" }); + var drafts_buf: [48]u8 = undefined; + const drafts = if (h.draft) |d| std.fmt.bufPrint(&drafts_buf, ", {s} drafts for one of them", .{d.model.graph.config.kind()}) catch "" else ""; + std.log.info("Qwen3.8-27B: {d} streams of {d} tokens{s}{s}", .{ slots, capacity, drafts, if (o.images != null) ", images read" else "" }); return h; } diff --git a/zig/src/native/qwen27_slots.zig b/zig/src/native/qwen27_slots.zig index 1dedd55bcb..e5230fe36d 100644 --- a/zig/src/native/qwen27_slots.zig +++ b/zig/src/native/qwen27_slots.zig @@ -38,6 +38,7 @@ pub fn opening(b: *Metal, s: *lanes.Stream, slot: u32) !Fill { if (resumed == 0) { if (d) |x| try x.reset() else try b.runner.reset(slot); try clearContext(b, slot); + try b.runner.admit(slot, @intCast(ids.len)); // a fresh pass: the slot takes its prompt's rows at once } s.cached = resumed; var at: usize = resumed; @@ -358,7 +359,8 @@ pub fn ready(b: *Metal, slot: u32, s: *const lanes.Stream) !void { defer added.deinit(b.gpa); try added.appendSlice(b.gpa, try b.runner.ensure(slot)); if (context_missing) { - const x = try Context.init(b.gpa, b.draft.?.model, drafting); + const x = try Context.init(b.gpa, b.draft.?.model, drafting, b.runner.kv_step); + x.settle = .{ .ctx = b.runner, .slot = slot, .run = q.decode_round.Runner.settleOpaque }; b.contexts[slot] = x; try x.buffers(&added, b.gpa); } @@ -367,8 +369,12 @@ pub fn ready(b: *Metal, slot: u32, s: *const lanes.Stream) !void { /// A drafter context's bytes: the taps ring, and K/V caches when the slot drafts. fn contextBytes(c: q.dflash.config.Config, drafting: bool) usize { - const ring = @as(usize, c.window) * c.tapWidth() * 2; - return ring + if (drafting) @as(usize, c.layers) * 2 * (@as(usize, c.window) + 16) * c.kvWidth() * 2 else 0; + const ring = @as(usize, c.span()) * c.ringWidth() * 2; + var caches: usize = 0; + if (drafting) for (0..c.layers) |i| { + caches += 2 * (@as(usize, c.layerWindow(i)) + 16) * c.kvWidth() * 2; + }; + return ring + caches; } /// A free slot for a new stream, slot 0 (the drafter's) first. @@ -392,10 +398,29 @@ pub fn unbind(b: *Metal, s: *lanes.Stream) void { var live: usize = 0; for (b.owners[0..b.runner.slots]) |owner| live += @intFromBool(owner != null); if (live <= 1) b.runner.trimBefore(); + b.give_back[slot] = b.runner.kv_step != null; + giveBack(b); +} + +/// Free, unread slots give back grown K/V and rings without waiting; slot 0 keeps what a next turn resumes from. +pub fn giveBack(b: *Metal) void { + for (0..b.runner.slots) |k| { + if (!b.give_back[k]) continue; + const slot: u32 = @intCast(k); + if (b.owners[k] != null) { + b.give_back[k] = false; + continue; + } + if ((slot == 0 and b.standing.tokens.len > 0) or b.runner.busy(slot)) continue; + b.give_back[k] = false; + b.runner.release(slot) catch |err| std.log.warn("qwen27: slot {d} keeps its K/V ({s})", .{ slot, @errorName(err) }); + clearContext(b, slot) catch {}; + } } /// A promote unbind left for later, once no round is pending: a stream still in slot 0 or none to move ends it. pub fn promoteDue(b: *Metal) void { + giveBack(b); if (!b.promote_due or b.round != null) return; b.promote_due = false; promote(b); diff --git a/zig/src/qwen27_memory_test.zig b/zig/src/qwen27_memory_test.zig index 5c6c5d98a2..74b2d10b8d 100644 --- a/zig/src/qwen27_memory_test.zig +++ b/zig/src/qwen27_memory_test.zig @@ -80,3 +80,127 @@ test "qwen27 memory: working sets account for weights and the unchanged 8 GiB ma try std.testing.expectEqual(@as(usize, 16 << 30), try memory.room(192 << 30, weights, "42949672960")); try std.testing.expectEqual(@as(usize, 0), try memory.room(192 << 30, 200 << 30, null)); } + +const kv_grow = @import("families/qwen27/kv_grow.zig"); + +test "qwen27 memory: grow-as-used K/V is off unless TF_Q27_KV_GROW names a 16-row step" { + try std.testing.expectEqual(@as(?u32, null), try kv_grow.step(null)); + try std.testing.expectEqual(@as(?u32, null), try kv_grow.step("0")); + try std.testing.expectEqual(@as(?u32, 4096), try kv_grow.step("4096")); + try std.testing.expectEqual(@as(?u32, 16), try kv_grow.step("16")); + for ([_][]const u8{ "", "-1", "15", "4097", "x", "99999999999" }) |text| try std.testing.expectError(error.BadKvGrowStep, kv_grow.step(text)); +} + +test "qwen27 memory: a slot starts at one step and grows to the need, a quarter more, or the context" { + // first rows: one step, never past the context + try std.testing.expectEqual(@as(u32, 4096), kv_grow.first(4096, 32784)); + try std.testing.expectEqual(@as(u32, 2000), kv_grow.first(4096, 2000)); + // enough already: unchanged + try std.testing.expectEqual(@as(u32, 4096), kv_grow.rows(4096, 4096, 4096, 32784)); + try std.testing.expectEqual(@as(u32, 4096), kv_grow.rows(4096, 10, 4096, 32784)); + // the need rounded up to the step + try std.testing.expectEqual(@as(u32, 8192), kv_grow.rows(4096, 4097, 4096, 32784)); + try std.testing.expectEqual(@as(u32, 32768), kv_grow.rows(4096, 30001, 4096, 32784)); + // a quarter more than held when that is past the rounded need + try std.testing.expectEqual(@as(u32, 20480), kv_grow.rows(16384, 16385, 4096, 32784)); + try std.testing.expectEqual(@as(u32, 20480), kv_grow.rows(16384, 16385, 16, 32784)); + // never past the context + try std.testing.expectEqual(@as(u32, 32784), kv_grow.rows(32768, 32780, 4096, 32784)); + try std.testing.expectEqual(@as(u32, 32784), kv_grow.rows(28672, 28673, 4096, 32784)); +} + +test "qwen27 memory: an emptied slot takes its prompt's rows at once, at least one step, at most the context" { + try std.testing.expectEqual(@as(u32, 4096), kv_grow.admitted(1, 4096, 131088)); + try std.testing.expectEqual(@as(u32, 4096), kv_grow.admitted(4096, 4096, 131088)); + try std.testing.expectEqual(@as(u32, 8192), kv_grow.admitted(4097, 4096, 131088)); + // one growth for a 65,536-token prompt, where growing from the first rows takes ten + try std.testing.expectEqual(@as(u32, 65536), kv_grow.admitted(65536, 4096, 131088)); + try std.testing.expectEqual(@as(u32, 131088), kv_grow.admitted(131080, 4096, 131088)); + try std.testing.expectEqual(@as(u32, 2000), kv_grow.admitted(64, 4096, 2000)); +} + +test "qwen27 memory: growing moves every K/V head to its place in the wider buffer" { + // 4 heads of dim 256 in bf16: head h starts at h * rows * 256 * 2 bytes + var moves: [4]kv_grow.Move = undefined; + const got = kv_grow.moves(&moves, 4, 256, 4096, 8192); + try std.testing.expectEqual(@as(usize, 4), got.len); + for (got, 0..) |m, h| { + try std.testing.expectEqual(h * 4096 * 256 * 2, m.from); + try std.testing.expectEqual(h * 8192 * 256 * 2, m.to); + try std.testing.expectEqual(@as(usize, 4096 * 256 * 2), m.len); + } + // the old buffer's bytes are exactly covered, the new head regions never overlap + try std.testing.expectEqual(@as(usize, 4 * 4096 * 256 * 2), got[3].from + got[3].len); + try std.testing.expect(got[2].to + got[2].len <= got[3].to); + // a slot's bytes at a row count, for the residency set and the memory receipts + try std.testing.expectEqual(@as(usize, 8 * 2 * 4096 * 1024 * 2), kv_grow.slotBytes(8, 1024, 4096)); +} + +const tap_ring = @import("families/qwen27/tap_ring.zig"); +const Ring = tap_ring.Ring; +const mtl_fake = @import("metal"); + +/// A ring that starts at `held` rows and grows by tap_ring.grow, as a parked slot's ring does. +const Growing = struct { + buf: mtl_fake.Buffer, + held: u32, + window: u32, + stride: usize, + asks: u32 = 0, + fn make(n: u32, stride: usize) !mtl_fake.Buffer { + const data = try std.testing.allocator.alignedAlloc(u8, .@"16", n * stride); + @memset(data, 0); + return .{ .id = n, .data = data }; + } + fn free(b: mtl_fake.Buffer) void { + std.testing.allocator.free(b.data); + } + fn run(ctx: *anyopaque, want: u32) anyerror!mtl_fake.Buffer { + const g: *Growing = @ptrCast(@alignCast(ctx)); + g.asks += 1; + if (want <= g.held) return g.buf; + const next = tap_ring.grownRows(g.held, want, 16, g.window); + const fresh = try make(next, g.stride); + tap_ring.carry(g.buf, fresh, g.held, g.stride); + free(g.buf); + g.buf = fresh; + g.held = next; + return fresh; + } +}; + +test "qwen27 memory: a growing taps ring holds the same rows at the same places as a whole-window ring" { + const window: u32 = 64; + const stride: usize = 4; + var whole_end: u64 = 0; + var grow_end: u64 = 0; + const whole_buf = try Growing.make(window, stride); + defer Growing.free(whole_buf); + var g = Growing{ .buf = try Growing.make(16, stride), .held = 16, .window = window, .stride = stride }; + defer Growing.free(g.buf); + const whole = Ring{ .buf = whole_buf, .end = &whole_end, .window = window, .stride = stride }; + var growing = Ring{ .buf = g.buf, .end = &grow_end, .window = window, .stride = stride, .grow = .{ .ctx = &g, .run = Growing.run } }; + const src = try Growing.make(128, stride); + defer Growing.free(src); + var next: u32 = 1; + for ([_]u32{ 5, 11, 1, 20, 30, 40 }) |n| { // past 16, 32 and the 64-row window: grows, then wraps + for (0..n) |i| std.mem.writeInt(u32, src.contents()[i * stride ..][0..4], next + @as(u32, @intCast(i)), .little); + next += n; + try whole.absorb(src, 0, n); + try growing.absorb(src, 0, n); + try std.testing.expectEqual(whole_end, grow_end); + const first = whole_end -| window; + for (first..whole_end) |pos| { + const at = (pos % window) * stride; + try std.testing.expectEqualSlices(u8, whole_buf.contents()[at..][0..stride], g.buf.contents()[at..][0..stride]); + } + } + try std.testing.expectEqual(window, g.held); // a wrapped ring holds the whole window + // a restore into the ring asks for its rows first; a held ring is not asked to grow + var restored = Ring{ .buf = g.buf, .end = &grow_end, .window = window, .stride = stride, .grow = .{ .ctx = &g, .run = Growing.run } }; + try restored.hold(200); + try std.testing.expectEqual(window, g.held); + try std.testing.expectEqual(@as(u32, 16), tap_ring.grownRows(16, 3, 16, 64)); + try std.testing.expectEqual(@as(u32, 48), tap_ring.grownRows(16, 33, 16, 64)); + try std.testing.expectEqual(@as(u32, 64), tap_ring.grownRows(48, 49, 16, 64)); +} diff --git a/zig/tests/qwen27/dflash_graph.zig b/zig/tests/qwen27/dflash_graph.zig index e382276faf..a99c93ee39 100644 --- a/zig/tests/qwen27/dflash_graph.zig +++ b/zig/tests/qwen27/dflash_graph.zig @@ -7,7 +7,7 @@ const q = tf.qwen27; const df = q.dflash; const Backend = df.runtime_backend.Backend; const Model = df.runtime_model.Model; -pub const config = df.config.Config{ .hidden = 512, .intermediate = 512, .vocab = 96, .heads = 8, .kv_heads = 2, .head_dim = 64, .target_layers = 5, .mask = 95, .window = 64, .rank = 32, .topk = 4, .taps = .{ 0, 1, 2, 3, 4 } }; +pub const config = df.config.Config{ .hidden = 512, .intermediate = 512, .vocab = 96, .heads = 8, .kv_heads = 2, .head_dim = 64, .target_layers = 5, .mask = 95, .window = 64, .rank = 32, .topk = 4, .taps = .{ 0, 1, 2, 3, 4, 0, 0, 0 } }; fn dimensions(name: []const u8) []const usize { if (std.mem.eql(u8, name, "fc.weight")) return &.{ 512, 2560 }; if (std.mem.eql(u8, name, "candidate_selector.hidden_projection.weight")) return &.{ 32, 512 }; diff --git a/zig/tests/qwen27/dflash_runtime.zig b/zig/tests/qwen27/dflash_runtime.zig index 607b7fb9f8..c20a08c419 100644 --- a/zig/tests/qwen27/dflash_runtime.zig +++ b/zig/tests/qwen27/dflash_runtime.zig @@ -40,7 +40,7 @@ pub fn main(init: std.process.Init) !void { var runner = try q.decode_round.Runner.init(a, target, 1, 256); defer runner.deinit(); runner.capture_taps = true; - runner.taps.ids = .{ 0, 1, 2, 3, 4 }; + runner.taps.ids = .{ 0, 1, 2, 3, 4, 0, 0, 0 }; const draft = try draft_graph.create(a, target); defer draft.deinit(); const eager = try draft_graph.create(a, target); diff --git a/zig/tests/qwen27/forward_diff.zig b/zig/tests/qwen27/forward_diff.zig index 7d08092778..cf719345fa 100644 --- a/zig/tests/qwen27/forward_diff.zig +++ b/zig/tests/qwen27/forward_diff.zig @@ -60,7 +60,7 @@ pub fn run(a: std.mem.Allocator, perturb: bool) !void { defer reference.deinit(); candidate.capture_taps = true; reference.capture_taps = true; - candidate.taps.ids = .{ 0, 1, 2, 3, 4 }; + candidate.taps.ids = .{ 0, 1, 2, 3, 4, 0, 0, 0 }; reference.taps.ids = candidate.taps.ids; var ids: [128]u32 = undefined; for (&ids, 0..) |*token, i| token.* = @intCast((i * 17 + 7) % 96);