From 69addb182ad92f8057ade3205003d1fbbabc06fa Mon Sep 17 00:00:00 2001 From: Grant Overton <65690763+grantoverton@users.noreply.github.com> Date: Sun, 1 Nov 2026 04:07:46 -1000 Subject: [PATCH 1/9] qwen27 dflash: read and run the 9B's DFlash drafter beside DFlash2 A DFlashDraftModel config now parses as variant v1: its taps (up to eight), block, mask token and layer kinds come from the config, with no conv or selector keys. Its graph is the plain qwen3 body (eleven tensors a layer, fc, hidden_norm, norm), each layer adding attention and MLP to its residual through a new add operator on the target's residual kernel. With no codebooks (rank 0) the selector's edge term is zero, so trees come from each depth's marginals. The target's tap planes follow the drafter's tap count (eight here), and the keep, gather and glue paths pass that count instead of five. DFlash2 parsing, its schedule and its tests are unchanged. --- zig/src/families/qwen27/dflash/config.zig | 144 ++++++++++++++- zig/src/families/qwen27/dflash/execution.zig | 36 +++- .../qwen27/dflash/execution_v1_test.zig | 170 ++++++++++++++++++ zig/src/families/qwen27/dflash/fixtures.zig | 46 +++-- zig/src/families/qwen27/dflash/generation.zig | 7 +- zig/src/families/qwen27/dflash/operators.zig | 2 + zig/src/families/qwen27/dflash/root.zig | 1 + .../qwen27/dflash/runtime_callbacks.zig | 14 +- .../qwen27/dflash/runtime_lattice.zig | 5 +- zig/src/families/qwen27/dflash/selector.zig | 3 +- zig/src/families/qwen27/dflash/session.zig | 2 +- zig/src/families/qwen27/dflash/weights.zig | 42 ++++- zig/src/families/qwen27/glue_gpu.zig | 2 +- zig/src/families/qwen27/tap_gather.zig | 9 +- zig/src/families/qwen27/taps.zig | 37 +++- zig/src/native/qwen27_draft.zig | 8 +- zig/src/native/qwen27_keep.zig | 4 +- zig/src/native/qwen27_lanes.zig | 2 +- zig/src/native/qwen27_served.zig | 2 +- 19 files changed, 484 insertions(+), 52 deletions(-) create mode 100644 zig/src/families/qwen27/dflash/execution_v1_test.zig diff --git a/zig/src/families/qwen27/dflash/config.zig b/zig/src/families/qwen27/dflash/config.zig index 19f6a39c5b..273bb12369 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,10 @@ 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), pub fn qWidth(c: Config) u32 { return c.heads * c.head_dim; } @@ -29,7 +37,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,20 +45,24 @@ 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); @@ -61,6 +73,22 @@ pub const Config = struct { 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 +121,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 +148,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 +300,69 @@ 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)); } +// Ornith-1.5-9B-DFlash's config.json, as published. +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 "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()); +} diff --git a/zig/src/families/qwen27/dflash/execution.zig b/zig/src/families/qwen27/dflash/execution.zig index e7cbe2a028..6cc44bc4ca 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); @@ -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 }; 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..7a733ddfb8 --- /dev/null +++ b/zig/src/families/qwen27/dflash/execution_v1_test.zig @@ -0,0 +1,170 @@ +//! 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.tapWidth(), .stride = toy.tapWidth(), .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: fc, then k and v a layer + try std.testing.expectEqual(@as(u32, 1 + 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)); +} 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..2f7a783144 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,9 @@ 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.setLayers(runner.model.device, draft.graph.config.taps[0..draft.graph.config.tap_count]); + 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 +141,7 @@ 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 }); + 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 = g.runner.taps.count }); 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_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_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/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..4cba5a15b0 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, 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/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/taps.zig b/zig/src/families/qwen27/taps.zig index 23e04b5751..d96455723e 100644 --- a/zig/src/families/qwen27/taps.zig +++ b/zig/src/families/qwen27/taps.zig @@ -5,13 +5,16 @@ 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; pub const Taps = struct { buffer: mtl.Buffer, 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 { @@ -24,6 +27,21 @@ pub const Taps = struct { t.buffer.deinit(); } + /// 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) 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 +49,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; @@ -42,12 +60,12 @@ pub const Taps = struct { } 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 +73,20 @@ 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 }); + try glue.tapRows(e, .{ .buffer = t.buffer }, kept, out, .{ .rows = @intCast(path.len), .width = t.width, .capacity = t.capacity, .planes = t.count }); 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..3cad73efd5 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; } @@ -68,6 +71,7 @@ pub const Draft = struct { d.generation.deinit(); d.runner.capture_taps = d.previous_capture; 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..096b245506 100644 --- a/zig/src/native/qwen27_keep.zig +++ b/zig/src/native/qwen27_keep.zig @@ -42,7 +42,7 @@ 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 }); + 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 = r.taps.count }); 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 +84,7 @@ 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 }); + 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 = r.taps.count }); 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..77e6ee8460 100644 --- a/zig/src/native/qwen27_lanes.zig +++ b/zig/src/native/qwen27_lanes.zig @@ -335,7 +335,7 @@ 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()]); + try q.tap_gather.gatherPlanes(r.taps.count, 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()]); 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..7113f50d27 100644 --- a/zig/src/native/qwen27_served.zig +++ b/zig/src/native/qwen27_served.zig @@ -302,7 +302,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. From 0cddf71eee5591b613b2d0160d1dcfd78022091f Mon Sep 17 00:00:00 2001 From: Grant Overton <65690763+grantoverton@users.noreply.github.com> Date: Sun, 1 Nov 2026 04:07:46 -1000 Subject: [PATCH 2/9] qwen27 served: the serving line names a v1 drafter "DFlash", DFlash2 stays "DFlash2" --- zig/src/native/qwen27_served.zig | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/zig/src/native/qwen27_served.zig b/zig/src/native/qwen27_served.zig index 7113f50d27..917956c6be 100644 --- a/zig/src/native/qwen27_served.zig +++ b/zig/src/native/qwen27_served.zig @@ -383,7 +383,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 = if (d.model.graph.config.variant == .v1) "DFlash" else "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.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 From 0d595a666d427b8efa5ad3768ce60aa3bc0ad8bc Mon Sep 17 00:00:00 2001 From: Grant Overton <65690763+grantoverton@users.noreply.github.com> Date: Sun, 1 Nov 2026 04:07:46 -1000 Subject: [PATCH 3/9] qwen27 prompt memory: the planner counts a drafter's own tap planes and fusion rows runnerBytes kept five raw planes a prompt row. tapRowBytes now prices a drafter's planes, and for a fusing (v1) drafter the row-major copy, fc's rows and the fused row: 18 x hidden x 2 bytes for the 9B drafter's eight taps. Fit carries it as tap_row; 0 keeps the five-plane figure, so DFlash2 and drafter-less serves plan exactly as before. --- zig/src/families/qwen27/dflash/config.zig | 2 +- zig/src/families/qwen27/prompt_memory.zig | 18 ++++++++++++++++-- 2 files changed, 17 insertions(+), 3 deletions(-) diff --git a/zig/src/families/qwen27/dflash/config.zig b/zig/src/families/qwen27/dflash/config.zig index 273bb12369..b88dcdf872 100644 --- a/zig/src/families/qwen27/dflash/config.zig +++ b/zig/src/families/qwen27/dflash/config.zig @@ -300,7 +300,7 @@ 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)); } -// Ornith-1.5-9B-DFlash's config.json, as published. +// 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, 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 { From a584a26b0da79315a30fd44760c7ef94d7f09533 Mon Sep 17 00:00:00 2001 From: Grant Overton <65690763+grantoverton@users.noreply.github.com> Date: Sun, 1 Nov 2026 04:07:46 -1000 Subject: [PATCH 4/9] qwen27 dflash v1: the ring keeps fused rows, and the full layer sees the whole context The target's taps pass now runs the drafter's fc and hidden_norm as a v1 drafter's taps complete (Taps.setFuser), so every gather reads one plane of fused rows: a ring row and a prompt-cache state's drafter part are 8 KiB a token instead of 64 KiB. The drafter's context append starts from those rows (k and v alone). DFlash2 keeps its raw planes. The ring and the context now keep the served span (context plus the lane window) instead of the 4096-token window. Sliding layers attend their own window of it; the full-attention layer attends all of it, with its own cache and attention sized to the span. --- zig/src/families/qwen27/dflash/caches.zig | 6 +- zig/src/families/qwen27/dflash/config.zig | 30 +++++++++ zig/src/families/qwen27/dflash/execution.zig | 20 +++--- .../qwen27/dflash/execution_v1_test.zig | 19 +++++- zig/src/families/qwen27/dflash/generation.zig | 6 +- .../qwen27/dflash/runtime_backend.zig | 26 +++++++- .../qwen27/dflash/runtime_context.zig | 46 +++++++++----- .../families/qwen27/dflash/runtime_model.zig | 61 ++++++++++++++----- zig/src/families/qwen27/dflash/session.zig | 4 +- zig/src/families/qwen27/residency.zig | 2 + zig/src/families/qwen27/taps.zig | 56 ++++++++++++++++- zig/src/native/qwen27_draft.zig | 1 + zig/src/native/qwen27_keep.zig | 6 +- zig/src/native/qwen27_lanes.zig | 3 +- zig/src/native/qwen27_served.zig | 6 +- zig/src/native/qwen27_slots.zig | 8 ++- 16 files changed, 238 insertions(+), 62 deletions(-) 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 b88dcdf872..9a6b0e00ab 100644 --- a/zig/src/families/qwen27/dflash/config.zig +++ b/zig/src/families/qwen27/dflash/config.zig @@ -30,6 +30,20 @@ pub const Config = struct { 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; } @@ -360,6 +374,22 @@ test "a DFlash v1 config that asks for a causal draft is refused" { 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); diff --git a/zig/src/families/qwen27/dflash/execution.zig b/zig/src/families/qwen27/dflash/execution.zig index 6cc44bc4ca..b64901756a 100644 --- a/zig/src/families/qwen27/dflash/execution.zig +++ b/zig/src/families/qwen27/dflash/execution.zig @@ -41,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); @@ -56,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. @@ -68,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); @@ -85,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); @@ -105,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 { @@ -198,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 index 7a733ddfb8..5619426f7c 100644 --- a/zig/src/families/qwen27/dflash/execution_v1_test.zig +++ b/zig/src/families/qwen27/dflash/execution_v1_test.zig @@ -111,7 +111,7 @@ const Fake = struct { } }; fn taps() session.Batch { - return .{ .values = .{ .handle = 99, .rows = 2, .width = toy.tapWidth(), .stride = toy.tapWidth(), .dtype = .bf16 }, .positions = &.{ 0, 1 }, .layers = toy.taps, .target_end = 2 }; + 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" { @@ -125,8 +125,8 @@ test "a v1 block runs attention and an MLP a layer, two residual adds each, and 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: fc, then k and v a layer - try std.testing.expectEqual(@as(u32, 1 + 2 * 3), fake.projects); + // 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(); @@ -168,3 +168,16 @@ test "a v1 block without an add operator is refused, not run without its residua 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/generation.zig b/zig/src/families/qwen27/dflash/generation.zig index 2f7a783144..fe2959f29a 100644 --- a/zig/src/families/qwen27/dflash/generation.zig +++ b/zig/src/families/qwen27/dflash/generation.zig @@ -32,7 +32,10 @@ pub const Generation = struct { 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 }; } @@ -141,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 = g.runner.taps.count }); + 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/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_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_model.zig b/zig/src/families/qwen27/dflash/runtime_model.zig index 0ed3fe34aa..6ef687e149 100644 --- a/zig/src/families/qwen27/dflash/runtime_model.zig +++ b/zig/src/families/qwen27/dflash/runtime_model.zig @@ -25,13 +25,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 +52,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 +65,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 +94,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 +199,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 +229,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 +245,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(); @@ -246,7 +275,7 @@ pub const Context = struct { 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 pending = try device.buffer(@as(usize, c.span()) * 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); @@ -277,7 +306,7 @@ pub const Context = struct { } /// 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 }; } /// 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/session.zig b/zig/src/families/qwen27/dflash/session.zig index 4cba5a15b0..896e20abb9 100644 --- a/zig/src/families/qwen27/dflash/session.zig +++ b/zig/src/families/qwen27/dflash/session.zig @@ -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/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/taps.zig b/zig/src/families/qwen27/taps.zig index d96455723e..73a363e0b0 100644 --- a/zig/src/families/qwen27/taps.zig +++ b/zig/src/families/qwen27/taps.zig @@ -8,8 +8,17 @@ 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, @@ -24,12 +33,47 @@ 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) return error.BadTapShape; + 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); @@ -57,6 +101,13 @@ 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 { @@ -73,7 +124,8 @@ 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 = t.count }); + 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(); } }; diff --git a/zig/src/native/qwen27_draft.zig b/zig/src/native/qwen27_draft.zig index 3cad73efd5..ba2830414f 100644 --- a/zig/src/native/qwen27_draft.zig +++ b/zig/src/native/qwen27_draft.zig @@ -70,6 +70,7 @@ 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 diff --git a/zig/src/native/qwen27_keep.zig b/zig/src/native/qwen27_keep.zig index 096b245506..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 = r.taps.count }); + 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 = r.taps.count }); + 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 77e6ee8460..57e457329c 100644 --- a/zig/src/native/qwen27_lanes.zig +++ b/zig/src/native/qwen27_lanes.zig @@ -335,7 +335,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.gatherPlanes(r.taps.count, 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 917956c6be..c03a74c912 100644 --- a/zig/src/native/qwen27_served.zig +++ b/zig/src/native/qwen27_served.zig @@ -332,14 +332,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 +352,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; diff --git a/zig/src/native/qwen27_slots.zig b/zig/src/native/qwen27_slots.zig index 1dedd55bcb..06e4579c2c 100644 --- a/zig/src/native/qwen27_slots.zig +++ b/zig/src/native/qwen27_slots.zig @@ -367,8 +367,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. From 3a9e1efbd65cb99b933ed53d0a46c861e67ac16f Mon Sep 17 00:00:00 2001 From: Grant Overton <65690763+grantoverton@users.noreply.github.com> Date: Sun, 1 Nov 2026 06:18:44 -1000 Subject: [PATCH 5/9] qwen27 dflash: the developer tools' five-tap literals fill the eight tap slots Commit "read and run the 9B's DFlash drafter beside DFlash2" made the tap ids and a draft config's taps eight-slot arrays, and three tf-qwen27 tools still set five: zig build dev-tools failed on dflash_graph.zig, dflash_runtime.zig and forward_diff.zig. They now name the same five taps with three empty slots (the plane count stays five). --- zig/tests/qwen27/dflash_graph.zig | 2 +- zig/tests/qwen27/dflash_runtime.zig | 2 +- zig/tests/qwen27/forward_diff.zig | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) 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); From ae0d7595700f8eb6e81e08ff1c2356199563e76c Mon Sep 17 00:00:00 2001 From: Grant Overton <65690763+grantoverton@users.noreply.github.com> Date: Sun, 1 Nov 2026 11:45:24 -1000 Subject: [PATCH 6/9] qwen27 served: the startup line names a v1 drafter "DFlash" too The serving line already named a v1 drafter "DFlash", but the startup log line still said "DFlash2 drafts for one of them" for either. Config.kind() gives the name, DFlash for z-lab's v1 layout and DFlash2 for DFlash2, and both lines take it from there. Test: a v1 config is named DFlash and the stock DFlash2 config DFlash2. --- zig/src/families/qwen27/dflash/config.zig | 8 ++++++++ zig/src/native/qwen27_served.zig | 6 ++++-- 2 files changed, 12 insertions(+), 2 deletions(-) diff --git a/zig/src/families/qwen27/dflash/config.zig b/zig/src/families/qwen27/dflash/config.zig index 9a6b0e00ab..5da1567fb3 100644 --- a/zig/src/families/qwen27/dflash/config.zig +++ b/zig/src/families/qwen27/dflash/config.zig @@ -83,6 +83,10 @@ pub const Config = struct { _ = 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; @@ -396,3 +400,7 @@ test "the stock DFlash2 config keeps its five taps and twenty-five-thousand-wide 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/native/qwen27_served.zig b/zig/src/native/qwen27_served.zig index c03a74c912..ba18d6de3f 100644 --- a/zig/src/native/qwen27_served.zig +++ b/zig/src/native/qwen27_served.zig @@ -383,7 +383,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) |d| .{ .kind = if (d.model.graph.config.variant == .v1) "DFlash" else "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 +414,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; } From c86af93659da2ec13a9ff1d085e57a786f41e8d8 Mon Sep 17 00:00:00 2001 From: Grant Overton <65690763+grantoverton@users.noreply.github.com> Date: Sun, 1 Nov 2026 10:54:48 -1000 Subject: [PATCH 7/9] qwen27: grow a slot's K/V as its stream grows (opt-in TF_Q27_KV_GROW) Every slot took K/V rows for the whole --context on its first use, so four live 9B streams at --context 32768 held 4 GiB of K/V while each stream used a few thousand tokens. With TF_Q27_KV_GROW= a slot starts at that many rows and, when a round, a restored state or a moved slot needs more, grows to the need rounded up to the step or a quarter more than it holds, at most the context. Growth goes one attention layer at a time: a wider buffer pair, one waited blit of every old byte to its place (each K/V head from h*old*dim to h*new*dim), the old pair out of the residency set and freed, so growing holds at most one layer's old pair beside the new rows. It runs where moveSlot and reset already touch a slot (settleSlot, no pending keep, no failure). The attention kernels already take each stream's head stride from its cache, so the plan now passes the slot's held rows; arithmetic and key order depend only on positions. Unset, every slot holds the context as before, and the startup fit and slot admission still charge the full context. --- zig/src/families/qwen27/decode_round.zig | 86 ++++++++++++++++++++++-- zig/src/families/qwen27/kv_grow.zig | 40 +++++++++++ zig/src/families/qwen27/qwen27.zig | 1 + zig/src/families/qwen27/snapshot.zig | 2 + zig/src/native/qwen27_served.zig | 14 +++- zig/src/qwen27_memory_test.zig | 45 +++++++++++++ 6 files changed, 183 insertions(+), 5 deletions(-) create mode 100644 zig/src/families/qwen27/kv_grow.zig diff --git a/zig/src/families/qwen27/decode_round.zig b/zig/src/families/qwen27/decode_round.zig index d558bdee80..5e48e8855e 100644 --- a/zig/src/families/qwen27/decode_round.zig +++ b/zig/src/families/qwen27/decode_round.zig @@ -17,6 +17,10 @@ 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"); + +/// 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 +47,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 +80,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 +143,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 +169,74 @@ 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: 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; + } + for (old, 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); + } + 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 }); + } + /// 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 +297,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 +519,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/kv_grow.zig b/zig/src/families/qwen27/kv_grow.zig new file mode 100644 index 0000000000..dae269be52 --- /dev/null +++ b/zig/src/families/qwen27/kv_grow.zig @@ -0,0 +1,40 @@ +//! 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)); +} + +/// 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/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/snapshot.zig b/zig/src/families/qwen27/snapshot.zig index d99a2c3da4..d00ed22f4d 100644 --- a/zig/src/families/qwen27/snapshot.zig +++ b/zig/src/families/qwen27/snapshot.zig @@ -243,6 +243,7 @@ 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 r.reserve(slot, st.at); // a slot whose K/V grows as used holds the state's rows first try copy(a, r, slot, ring, st, false); r.offsets[slot] = st.at; if (ring) |g| g.end.* = st.draft_end; @@ -255,6 +256,7 @@ 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 src: std.ArrayList(Part) = .empty; defer src.deinit(a); var dst: std.ArrayList(Part) = .empty; diff --git a/zig/src/native/qwen27_served.zig b/zig/src/native/qwen27_served.zig index ba18d6de3f..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); @@ -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); diff --git a/zig/src/qwen27_memory_test.zig b/zig/src/qwen27_memory_test.zig index 5c6c5d98a2..b9669bd923 100644 --- a/zig/src/qwen27_memory_test.zig +++ b/zig/src/qwen27_memory_test.zig @@ -80,3 +80,48 @@ 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: 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)); +} From c25538afb774f9b03e87d7ec0df7b9c5e011c989 Mon Sep 17 00:00:00 2001 From: Grant Overton <65690763+grantoverton@users.noreply.github.com> Date: Sun, 1 Nov 2026 10:54:59 -1000 Subject: [PATCH 8/9] qwen27: free slots give back grown K/V; drafter taps rings grow as used With TF_Q27_KV_GROW a slot kept its high-water rows after its stream ended, and every parked stream's drafter taps ring was sized to the whole span, so one long request left its rows behind in every slot it touched. Give-back: when a stream leaves slot k, the slot goes back to its first rows (fresh small buffers, no copy) as soon as no blit reads it. The check is a non-blocking status read of the kept-state save, retried at each backend call, so no stream waits on it. Slot 0 keeps its rows while it stands for a next turn's resume; the next fresh pass there gives them back, as does any fresh pass in any slot. Rings: a parked stream's taps ring starts at the same first rows and grows by the K/V rule when taps are absorbed or a kept state is restored or moved into it (Ring.grow, Ring.hold); the old rows move to their places on the CPU after the slot's pending reads land. Drafter frames wait for their ring reads before closing, so a ring only changes between them. A context's reset gives its ring back. The growth helpers move to kv_slots.zig to keep decode_round.zig under 600 lines. --- zig/src/families/qwen27/decode_round.zig | 72 +++++---------- .../families/qwen27/dflash/runtime_model.zig | 57 +++++++++++- zig/src/families/qwen27/kv_slots.zig | 92 +++++++++++++++++++ zig/src/families/qwen27/snapshot.zig | 10 +- zig/src/families/qwen27/tap_ring.zig | 24 ++++- zig/src/native/qwen27_lanes.zig | 6 +- zig/src/native/qwen27_slots.zig | 23 ++++- zig/src/qwen27_memory_test.zig | 69 ++++++++++++++ 8 files changed, 292 insertions(+), 61 deletions(-) create mode 100644 zig/src/families/qwen27/kv_slots.zig diff --git a/zig/src/families/qwen27/decode_round.zig b/zig/src/families/qwen27/decode_round.zig index 5e48e8855e..46e2fb235e 100644 --- a/zig/src/families/qwen27/decode_round.zig +++ b/zig/src/families/qwen27/decode_round.zig @@ -18,6 +18,7 @@ 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 }; @@ -183,58 +184,29 @@ pub const Runner = struct { return rows; } - /// Slot `slot` holds `need` K/V rows: layer by layer, a wider pair takes every old byte at its place (waited blit). + /// Slot `slot` holds `need` K/V rows, grown layer by layer (kv_slots.zig). 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; + 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); + } + + /// 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); - 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; - } - for (old, 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); - } - 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 }); } /// The K/V buffers of the ready slots (for a residency set). diff --git a/zig/src/families/qwen27/dflash/runtime_model.zig b/zig/src/families/qwen27/dflash/runtime_model.zig index 6ef687e149..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"); @@ -268,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 - /// `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 { + /// 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; 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.span()) * c.ringWidth() * 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; @@ -300,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.span(), .stride = @as(usize, c.ringWidth()) * 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/kv_slots.zig b/zig/src/families/qwen27/kv_slots.zig new file mode 100644 index 0000000000..6a49f26139 --- /dev/null +++ b/zig/src/families/qwen27/kv_slots.zig @@ -0,0 +1,92 @@ +//! 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; + if (slot >= r.ready) return; + const first = kv_grow.first(s, r.capacity); + 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) 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} gave back its K/V ({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/snapshot.zig b/zig/src/families/qwen27/snapshot.zig index d00ed22f4d..731bd7c100 100644 --- a/zig/src/families/qwen27/snapshot.zig +++ b/zig/src/families/qwen27/snapshot.zig @@ -243,8 +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 r.reserve(slot, st.at); // a slot whose K/V grows as used holds the state's rows first - 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; } @@ -257,12 +259,14 @@ pub fn moveSlot(a: std.mem.Allocator, r: *Runner, from: u32, from_ring: ?Ring, t 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_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/native/qwen27_lanes.zig b/zig/src/native/qwen27_lanes.zig index 57e457329c..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; } diff --git a/zig/src/native/qwen27_slots.zig b/zig/src/native/qwen27_slots.zig index 06e4579c2c..1ac85eaea5 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.release(slot); // a fresh pass: the slot's grown rows are no one's now } 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); } @@ -396,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 b9669bd923..675c3c8bca 100644 --- a/zig/src/qwen27_memory_test.zig +++ b/zig/src/qwen27_memory_test.zig @@ -125,3 +125,72 @@ test "qwen27 memory: growing moves every K/V head to its place in the wider buff // 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)); +} From f6ea8ee0f2e2b630a37f992539065de792e39b32 Mon Sep 17 00:00:00 2001 From: Grant Overton <65690763+grantoverton@users.noreply.github.com> Date: Sun, 1 Nov 2026 17:08:04 -1000 Subject: [PATCH 9/9] qwen27: a fresh pass takes its prompt's K/V rows at once (TF_Q27_KV_GROW) With K/V grown as used, a fresh prompt pass gave its slot back to the first rows and then grew it step by step through the prompt: ten growths for a 65,536-token prompt, each a waited blit of every layer into new buffers. The slot is empty at that point, so it now takes the rows the prompt fills (rounded up to the step) in new buffers with no copy, and keeps them when it already holds exactly that many. Decode and restores still grow by the same rule. --- zig/src/families/qwen27/decode_round.zig | 5 +++++ zig/src/families/qwen27/kv_grow.zig | 6 ++++++ zig/src/families/qwen27/kv_slots.zig | 19 +++++++++++++++---- zig/src/native/qwen27_slots.zig | 2 +- zig/src/qwen27_memory_test.zig | 10 ++++++++++ 5 files changed, 37 insertions(+), 5 deletions(-) diff --git a/zig/src/families/qwen27/decode_round.zig b/zig/src/families/qwen27/decode_round.zig index 46e2fb235e..ac05d71c7f 100644 --- a/zig/src/families/qwen27/decode_round.zig +++ b/zig/src/families/qwen27/decode_round.zig @@ -194,6 +194,11 @@ pub const Runner = struct { 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; diff --git a/zig/src/families/qwen27/kv_grow.zig b/zig/src/families/qwen27/kv_grow.zig index dae269be52..dd69edc417 100644 --- a/zig/src/families/qwen27/kv_grow.zig +++ b/zig/src/families/qwen27/kv_grow.zig @@ -26,6 +26,12 @@ pub fn rows(held: u32, need: u32, s: u32, context: u32) u32 { 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); diff --git a/zig/src/families/qwen27/kv_slots.zig b/zig/src/families/qwen27/kv_slots.zig index 6a49f26139..349e92e704 100644 --- a/zig/src/families/qwen27/kv_slots.zig +++ b/zig/src/families/qwen27/kv_slots.zig @@ -52,18 +52,29 @@ pub fn reserve(r: *Runner, slot: u32, need: u32) !void { /// 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; - const first = kv_grow.first(s, r.capacity); 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) return; + 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; + 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(); @@ -71,7 +82,7 @@ pub fn release(r: *Runner, slot: u32) !void { swapPair(r, at, fresh, first); } r.offsets[slot] = 0; - std.log.info("qwen27: slot {d} gave back its K/V ({d} to {d} rows)", .{ slot, held, first }); + 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. diff --git a/zig/src/native/qwen27_slots.zig b/zig/src/native/qwen27_slots.zig index 1ac85eaea5..e5230fe36d 100644 --- a/zig/src/native/qwen27_slots.zig +++ b/zig/src/native/qwen27_slots.zig @@ -38,7 +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.release(slot); // a fresh pass: the slot's grown rows are no one's now + 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; diff --git a/zig/src/qwen27_memory_test.zig b/zig/src/qwen27_memory_test.zig index 675c3c8bca..74b2d10b8d 100644 --- a/zig/src/qwen27_memory_test.zig +++ b/zig/src/qwen27_memory_test.zig @@ -109,6 +109,16 @@ test "qwen27 memory: a slot starts at one step and grows to the need, a quarter 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;