diff --git a/zig/build/server.zig b/zig/build/server.zig index 7bc8e16048..5ed055d82d 100644 --- a/zig/build/server.zig +++ b/zig/build/server.zig @@ -30,6 +30,7 @@ pub fn addTo( if (family) |drama3| { const source = b.addWriteFiles().add("drama_prompt.zig", \\pub const Reference = @import("drama3").prompt.Reference; + \\pub const Built = @import("drama3").prompt.Built; \\pub const build = @import("drama3").prompt.build; ); m.addImport("drama_prompt", b.createModule(.{ diff --git a/zig/src/core/speech_api.zig b/zig/src/core/speech_api.zig index 4839cfee91..d593239a45 100644 --- a/zig/src/core/speech_api.zig +++ b/zig/src/core/speech_api.zig @@ -14,6 +14,8 @@ pub const SpeechRequest = struct { top_k: u32 = 30, /// Publishes decoded prefixes; whole replies also decode incrementally when a speech guard is enabled. stream: bool = false, + /// End of the invariant reference prompt, excluding request text; 0 disables reuse; backend saves no lookahead row. + reference_rows: u32 = 0, }; pub const SpeechInfo = struct { diff --git a/zig/src/families/drama3/drama3.zig b/zig/src/families/drama3/drama3.zig index 227c835701..88e31598ca 100644 --- a/zig/src/families/drama3/drama3.zig +++ b/zig/src/families/drama3/drama3.zig @@ -3,6 +3,7 @@ pub const completion = @import("completion.zig"); pub const speech_decode = @import("speech_decode.zig"); pub const model = @import("model.zig"); pub const kv_state = @import("kv_state.zig"); +pub const reference_cache = @import("reference_cache.zig"); pub const codec_load = @import("codec_load.zig"); pub const speed_probe = @import("cuda/speed_probe.zig"); pub const replay = @import("replay.zig").Replay(@This()); @@ -29,6 +30,7 @@ test { _ = config; _ = speed_probe; _ = model; + _ = reference_cache; _ = kv_state; _ = @import("codec_map.zig"); _ = @import("codec_schema.zig"); diff --git a/zig/src/families/drama3/drama3_metal.zig b/zig/src/families/drama3/drama3_metal.zig index 2e2d550c61..41b2e059e4 100644 --- a/zig/src/families/drama3/drama3_metal.zig +++ b/zig/src/families/drama3/drama3_metal.zig @@ -3,6 +3,7 @@ pub const completion = @import("completion.zig"); pub const speech_decode = @import("speech_decode.zig"); pub const model = @import("model.zig"); pub const kv_state = @import("kv_state.zig"); +pub const reference_cache = @import("reference_cache.zig"); pub const codec_load = @import("codec_load.zig"); pub const replay = @import("replay.zig").Replay(@This()); pub const config = @import("config.zig"); @@ -19,6 +20,7 @@ pub const codec = @import("metal/codec.zig"); pub const tts = @import("metal/tts.zig"); pub const Engine = engine.Engine; +pub const Snapshot = engine.Engine.Snapshot; pub const Sampling = engine.Sampling; pub const Options = engine.Options; @@ -29,6 +31,7 @@ test { _ = config; _ = model; _ = kv_state; + _ = reference_cache; _ = @import("codec_map.zig"); _ = @import("codec_schema.zig"); _ = @import("codec_pth.zig"); diff --git a/zig/src/families/drama3/generate.zig b/zig/src/families/drama3/generate.zig index 2bd0cbb3e0..1e4c7a5060 100644 --- a/zig/src/families/drama3/generate.zig +++ b/zig/src/families/drama3/generate.zig @@ -256,7 +256,7 @@ fn runRequest(gpa: std.mem.Allocator, io: std.Io, a: Args, e: *d3.Engine, tts: ? nref = 1; } if (nref > 0) ref_frames = refs[0].codes.len; - rows = try d3.prompt.build(ar, tok, text, refs[0..nref]); + rows = (try d3.prompt.build(ar, tok, text, refs[0..nref])).frames; } if (a.replay) try d3.replay.verify(gpa, io, e, tts orelse return error.CodecRequired, rows, q.smp); if (a.speed_probe) { diff --git a/zig/src/families/drama3/metal/engine.zig b/zig/src/families/drama3/metal/engine.zig index 00db2d8a1a..54f761c37d 100644 --- a/zig/src/families/drama3/metal/engine.zig +++ b/zig/src/families/drama3/metal/engine.zig @@ -294,6 +294,32 @@ pub const Engine = struct { e.rows_done = n; } + /// Recurrent state only. The caller must invalidate it before its lane's live prefix KV changes. + pub const Snapshot = @import("../reference_cache.zig").Snapshot; + + /// Captures into one reusable, GPU-owned buffer pair (including partial allocation on failure). + pub fn snapshot(e: *Engine, saved: *Snapshot) !void { + try saved.capture(e.g, e.state_bufs, e.state_bytes); + } + + pub fn restore(e: *Engine, saved: *const Snapshot) !void { + try saved.restore(e.g, e.state_bufs, e.state_bytes); + } + + /// Starts a prompt whose first `start` rows are already in the engine's state; writes the rest, no reset. + pub fn beginPromptFrom(e: *Engine, rows: []const [W]i32, start: usize) !void { + try model.validate(rows, e.c.ids, e.o.max_ctx, 1); + const n = rows.len; + if (n < 2 or start == 0 or start >= n) return error.PromptTooShort; + if (n + 2 > e.o.max_ctx) return error.PromptTooLong; + try e.put(e.pframes, std.mem.sliceAsBytes(rows)); + // The suffix prefill/finishPrompt rebuilds all other decode state, including hlast and RAS. + @memset(e.g.host(e.count, 8), 0); + @memset(e.g.host(e.step, 8), 0); + e.rows_done = start; + e.prompt_rows = start; + } + fn tower(e: *Engine, tt: usize, c0: usize, hfin: u64, probe: ?Probe) !void { var hin = hfin; var nb: [64]u8 = undefined; diff --git a/zig/src/families/drama3/prompt.zig b/zig/src/families/drama3/prompt.zig index 6b3a874994..e9e40b6c48 100644 --- a/zig/src/families/drama3/prompt.zig +++ b/zig/src/families/drama3/prompt.zig @@ -58,7 +58,13 @@ fn refTranscript(a: std.mem.Allocator, transcript: []const u8) ![]u8 { return std.mem.concat(a, u8, &.{ body[0..tl], tagged[0..lead], body[tl..] }); } -pub fn build(gpa: std.mem.Allocator, tok: anytype, text: []const u8, refs: []const Reference) ![][W]i32 { +pub const Built = struct { + frames: [][W]i32, + /// End of the invariant reference prompt, before request text (0 without references). + reference_rows: u32, +}; + +pub fn build(gpa: std.mem.Allocator, tok: anytype, text: []const u8, refs: []const Reference) !Built { if (text.len == 0 or text.len > 16384 or !std.unicode.utf8ValidateSlice(text) or refs.len > 4) return error.BadPrompt; for (refs) |r| { if (r.transcript.len == 0 or r.transcript.len > 4096 or !std.unicode.utf8ValidateSlice(r.transcript) or @@ -73,7 +79,7 @@ pub fn build(gpa: std.mem.Allocator, tok: anytype, text: []const u8, refs: []con const sem_begin = tok.specialTokenId("<|semantic:0|>") orelse return error.NoSemanticToken; var rows: std.ArrayList([W]i32) = .empty; errdefer rows.deinit(gpa); - const Seg = union(enum) { text: []const u8, codes: []const [cfg.codebooks]i32 }; + const Seg = union(enum) { text: []const u8, codes: []const [cfg.codebooks]i32, reference_end }; var segs: std.ArrayList(Seg) = .empty; try segs.append(a, .{ .text = "<|im_start|>system\n" }); try segs.append(a, .{ .text = if (refs.len > 0) reference_instruction else tts_instruction }); @@ -84,9 +90,11 @@ pub fn build(gpa: std.mem.Allocator, tok: anytype, text: []const u8, refs: []con } try segs.append(a, .{ .text = message_end }); try segs.append(a, .{ .text = "<|im_start|>user\n" }); + try segs.append(a, .reference_end); try segs.append(a, .{ .text = try withSpeaker(a, text) }); try segs.append(a, .{ .text = message_end }); try segs.append(a, .{ .text = generation_start }); + var reference_rows: u32 = 0; for (segs.items) |seg| switch (seg) { .text => |t| { const ids = try tok.encode(a, t); @@ -102,8 +110,11 @@ pub fn build(gpa: std.mem.Allocator, tok: anytype, text: []const u8, refs: []con for (0..cfg.codebooks) |c| row[1 + c] = f[c]; try rows.append(gpa, row); }, + .reference_end => if (refs.len > 0) { + reference_rows = @intCast(rows.items.len); + }, }; - return rows.toOwnedSlice(gpa); + return .{ .frames = try rows.toOwnedSlice(gpa), .reference_rows = reference_rows }; } test "speaker tags as prompt.py places them" { @@ -115,3 +126,103 @@ test "speaker tags as prompt.py places them" { try std.testing.expectEqualStrings("<|speaker:1|> hi", try refTranscript(a, " <|speaker:1|> hi")); try std.testing.expectEqualStrings("<|speaker:0|> We asked", try refTranscript(a, "We asked")); } + +const ByteTokenizer = struct { + pub fn specialTokenId(_: @This(), _: []const u8) ?u32 { + return 1000; + } + + pub fn encode(_: @This(), a: std.mem.Allocator, text: []const u8) ![]u32 { + const ids = try a.alloc(u32, text.len); + for (text, ids) |byte, *id| id.* = byte; + return ids; + } +}; + +fn expectTextRows(text: []const u8, rows: []const [W]i32) !void { + try std.testing.expectEqual(text.len, rows.len); + for (text, rows) |byte, row| { + var expected: [W]i32 = @splat(0); + expected[0] = byte; + try std.testing.expectEqualSlices(i32, &expected, &row); + } +} + +test "reference prompt boundary preserves every row and ends before request text" { + const t = std.testing; + const codes = [_][cfg.codebooks]i32{.{ 7, 1, 2, 3, 4, 5, 6, 7, 8, 9 }}; + const refs = [_]Reference{.{ .transcript = "Ref.", .codes = &codes }}; + const built = try build(t.allocator, ByteTokenizer{}, "One.", &refs); + defer t.allocator.free(built.frames); + const before = "<|im_start|>system\n" ++ reference_instruction ++ + "\n\n[Reference 1]\nText: <|speaker:0|> Ref.\nSpeech: "; + const after = message_end ++ "<|im_start|>user\n"; + const suffix = "<|speaker:0|> One." ++ message_end ++ generation_start; + const boundary = before.len + codes.len + after.len; + try t.expectEqual(@as(u32, boundary), built.reference_rows); + try t.expectEqual(boundary + suffix.len, built.frames.len); + try expectTextRows(before, built.frames[0..before.len]); + try t.expectEqualSlices(i32, &.{ 1007, 7, 1, 2, 3, 4, 5, 6, 7, 8, 9 }, &built.frames[before.len]); + try expectTextRows(after, built.frames[before.len + 1 .. boundary]); + try expectTextRows(suffix, built.frames[boundary..]); + + const other = try build(t.allocator, ByteTokenizer{}, " <|speaker:1|> Different text.", &refs); + defer t.allocator.free(other.frames); + try t.expectEqual(built.reference_rows, other.reference_rows); + try t.expectEqualSlices(u8, std.mem.sliceAsBytes(built.frames[0..boundary]), std.mem.sliceAsBytes(other.frames[0..boundary])); + try expectTextRows(" <|speaker:1|> Different text." ++ message_end ++ generation_start, other.frames[boundary..]); + + const plain = try build(t.allocator, ByteTokenizer{}, "One.", &.{}); + defer t.allocator.free(plain.frames); + try t.expectEqual(@as(u32, 0), plain.reference_rows); + try expectTextRows("<|im_start|>system\n" ++ tts_instruction ++ message_end ++ "<|im_start|>user\n" ++ suffix, plain.frames); +} + +test "reference prompt boundary includes all transcripts, wrappers and multiple references" { + const t = std.testing; + const codes = [_][cfg.codebooks]i32{ @splat(1), @splat(2) }; + const refs = [_]Reference{ + .{ .transcript = " <|speaker:1|> Ref.", .codes = &codes }, + .{ .transcript = "Second.", .codes = codes[0..1] }, + }; + const built = try build(t.allocator, ByteTokenizer{}, "One.", &refs); + defer t.allocator.free(built.frames); + const before = "<|im_start|>system\n" ++ reference_instruction ++ + "\n\n[Reference 1]\nText: <|speaker:1|> Ref.\nSpeech: "; + const between = "\n\n[Reference 2]\nText: <|speaker:0|> Second.\nSpeech: "; + const after = message_end ++ "<|im_start|>user\n"; + const boundary = before.len + codes.len + between.len + 1 + after.len; + try t.expectEqual(@as(u32, boundary), built.reference_rows); + try expectTextRows(before, built.frames[0..before.len]); + try expectTextRows(between, built.frames[before.len + 2 ..][0..between.len]); + try expectTextRows(after, built.frames[boundary - after.len .. boundary]); + try expectTextRows("<|speaker:0|> One." ++ message_end ++ generation_start, built.frames[boundary..]); +} + +test "reference cache: last reference code and equal-length changed tower lookahead both miss" { + const t = std.testing; + const cache = @import("reference_cache.zig"); + var prefix: cache.Prefix = .{}; + defer prefix.deinit(t.allocator); + var codes = [_][cfg.codebooks]i32{ @splat(1), @splat(2) }; + const refs = [_]Reference{.{ .transcript = "Ref.", .codes = &codes }}; + const first = try build(t.allocator, ByteTokenizer{}, "One.", &refs); + defer t.allocator.free(first.frames); + const saved_rows = cache.prefillRows(first.frames.len, first.reference_rows, 1); + try t.expect(!try prefix.begin(t.allocator, first.frames, saved_rows)); + prefix.saved(); + + codes[codes.len - 1][cfg.codebooks - 1] += 1; + const changed = try build(t.allocator, ByteTokenizer{}, "One.", &refs); + defer t.allocator.free(changed.frames); + try t.expectEqual(first.reference_rows, changed.reference_rows); + try t.expectEqual(saved_rows, cache.prefillRows(changed.frames.len, changed.reference_rows, 1)); + try t.expect(!try prefix.begin(t.allocator, changed.frames, saved_rows)); + prefix.saved(); + try t.expect(!try prefix.begin(t.allocator, first.frames, saved_rows)); + prefix.saved(); + + // Every processed row is unchanged; only the next row read by the tower differs. + first.frames[saved_rows][0] += 1; + try t.expect(!try prefix.begin(t.allocator, first.frames, saved_rows)); +} diff --git a/zig/src/families/drama3/reference_cache.zig b/zig/src/families/drama3/reference_cache.zig new file mode 100644 index 0000000000..7178db144c --- /dev/null +++ b/zig/src/families/drama3/reference_cache.zig @@ -0,0 +1,204 @@ +//! One live KV prefix per lane, identified exactly through the tower's next-row dependency. +const std = @import("std"); +const Frame = @import("model.zig").Frame; + +/// Stop on an existing chunk boundary with the tower lookahead inside the invariant prompt, keeping uncached cuts. +pub fn prefillRows(total: usize, reference_rows: usize, chunk: usize) usize { + if (reference_rows < 2 or reference_rows >= total) return 0; + std.debug.assert(chunk > 0); + return (reference_rows - 1) / chunk * chunk; +} + +pub const Prefix = struct { + rows: std.ArrayList(Frame) = .empty, // saved rows plus one tower lookahead row; never borrows a request + valid: bool = false, + + pub fn deinit(p: *Prefix, a: std.mem.Allocator) void { + p.rows.deinit(a); + p.* = .{}; + } + + pub fn invalidate(p: *Prefix) void { + p.valid = false; + } + + /// Called before reset/prefill overwrites live KV; a miss invalidates even on allocation failure or cancel. + pub fn begin(p: *Prefix, a: std.mem.Allocator, frames: []const Frame, saved_rows: usize) !bool { + if (saved_rows > 0) { + std.debug.assert(saved_rows < frames.len); + const identity = frames[0 .. saved_rows + 1]; + if (p.valid and std.mem.eql(u8, std.mem.sliceAsBytes(p.rows.items), std.mem.sliceAsBytes(identity))) return true; + p.invalidate(); + try p.rows.resize(a, identity.len); + @memcpy(p.rows.items, identity); + } else { + p.invalidate(); + p.rows.clearRetainingCapacity(); + } + return false; + } + + /// Only after both the prefix prefill and recurrent-state copy have completed successfully. + pub fn saved(p: *Prefix) void { + std.debug.assert(p.rows.items.len >= 2); + p.valid = true; + } +}; + +/// Recurrent state only; KV stays in its lane. Buffers are GPU-owned, reused, and freed on close even if partial. +pub const Snapshot = struct { + buffers: [2]u64 = .{ 0, 0 }, + + pub fn capture(s: *Snapshot, g: anytype, live: [2]u64, bytes: [2]usize) !void { + try g.flush(); + for (0..2) |i| if (s.buffers[i] == 0) { + s.buffers[i] = try g.alloc(bytes[i]); + }; + for (0..2) |i| g.copy(live[i], s.buffers[i], bytes[i]); + try g.flush(); + } + + pub fn restore(s: *const Snapshot, g: anytype, live: [2]u64, bytes: [2]usize) !void { + for (0..2) |i| { + std.debug.assert(s.buffers[i] != 0); + g.copy(s.buffers[i], live[i], bytes[i]); + } + try g.flush(); + } +}; + +test "reference cache: boundary excludes tower lookahead and preserves original chunk cuts" { + const t = std.testing; + try t.expectEqual(@as(usize, 9), prefillRows(20, 10, 1)); + try t.expectEqual(@as(usize, 8), prefillRows(20, 10, 4)); + try t.expectEqual(@as(usize, 0), prefillRows(20, 4, 4)); + for ([_]usize{ 0, 1, 20, 21, std.math.maxInt(usize) }) |boundary| + try t.expectEqual(@as(usize, 0), prefillRows(20, boundary, 1)); + for ([_]usize{ 1, 2, 4, 16 }) |chunk| { + for (2..20) |boundary| { + const rows = prefillRows(20, boundary, chunk); + try t.expect(rows < boundary); + try t.expectEqual(@as(usize, 0), rows % chunk); + } + } +} + +test "reference cache: exact owned identity covers every prefix and tower lookahead cell" { + const t = std.testing; + var p: Prefix = .{}; + defer p.deinit(t.allocator); + const frames: [6]Frame = @splat(@splat(7)); + try t.expect(!try p.begin(t.allocator, &frames, 3)); + p.saved(); + var changed = frames; + changed[4][0] = 42; // request text beyond the lookahead does not change the saved state + try t.expect(try p.begin(t.allocator, &changed, 3)); + try t.expect(try p.begin(t.allocator, changed[0..5], 3)); + for (0..4) |row| { + for (0..11) |col| { + changed = frames; + changed[row][col] += 1; + try t.expect(!try p.begin(t.allocator, &changed, 3)); + p.saved(); + changed[row][col] = frames[row][col]; // the identity must not borrow the caller's buffer + try t.expect(!try p.begin(t.allocator, &frames, 3)); + p.saved(); + } + } + try t.expect(!try p.begin(t.allocator, &frames, 2)); // equal contents, different prefix length +} + +test "reference cache: default, partial replacement and allocation errors invalidate before overwrite" { + const t = std.testing; + var p: Prefix = .{}; + defer p.deinit(t.allocator); + const a: [5]Frame = @splat(@splat(1)); + const b: [5]Frame = @splat(@splat(2)); + _ = try p.begin(t.allocator, &a, 3); + p.saved(); + try t.expect(try p.begin(t.allocator, &a, 3)); + // Cancelling a same-prefix suffix leaves prefix KV and the saved recurrence untouched. + try t.expect(try p.begin(t.allocator, &a, 3)); + try t.expect(!try p.begin(t.allocator, &b, 0)); // cancellation during uncached/default prefill + try t.expect(!p.valid); // even before a replacement snapshot can be reached + try t.expect(!try p.begin(t.allocator, &a, 3)); + p.saved(); + try t.expect(!try p.begin(t.allocator, &b, 3)); // cancelled/error before a replacement snapshot + try t.expect(!p.valid); + try t.expect(!try p.begin(t.allocator, &a, 3)); + p.saved(); + const larger: [100]Frame = @splat(@splat(3)); + var failing = std.testing.FailingAllocator.init(t.allocator, .{ .fail_index = 0, .resize_fail_index = 0 }); + try t.expectError(error.OutOfMemory, p.begin(failing.allocator(), &larger, 99)); + try t.expect(!p.valid); + try t.expect(!try p.begin(t.allocator, &a, 3)); + p.saved(); + p.invalidate(); // an execution error after a hit cannot publish a reusable prefix + try t.expect(!try p.begin(t.allocator, &a, 3)); +} + +/// Small CPU buffers exercise the same capture/restore and partial-allocation ownership as Metal. +const TestGpu = struct { + buffers: std.ArrayList([]u8) = .empty, + fail_at: ?usize = null, + + pub fn alloc(g: *TestGpu, bytes: usize) !u64 { + if (g.fail_at == g.buffers.items.len) return error.OutOfMemory; + const buffer = try std.testing.allocator.alloc(u8, bytes); + errdefer std.testing.allocator.free(buffer); + @memset(buffer, 0); + try g.buffers.append(std.testing.allocator, buffer); + return @intCast(g.buffers.items.len); + } + + pub fn flush(_: *TestGpu) !void {} + + pub fn copy(g: *TestGpu, from: u64, to: u64, bytes: usize) void { + @memcpy(g.buffers.items[@intCast(to - 1)][0..bytes], g.buffers.items[@intCast(from - 1)][0..bytes]); + } + + fn deinit(g: *TestGpu) void { + for (g.buffers.items) |b| std.testing.allocator.free(b); + g.buffers.deinit(std.testing.allocator); + } +}; + +test "reference cache: switches reuse one snapshot pair and restore the latest recurrent state" { + const t = std.testing; + var g: TestGpu = .{}; + defer g.deinit(); + const bytes: [2]usize = .{ 4, 8 }; + const live: [2]u64 = .{ try g.alloc(bytes[0]), try g.alloc(bytes[1]) }; + var snapshot: Snapshot = .{}; + for (1..33) |voice| { + const value: u8 = @intCast(voice); + for (g.buffers.items[0..2]) |buffer| @memset(buffer, value); + try snapshot.capture(&g, live, bytes); + try t.expectEqual(@as(usize, 4), g.buffers.items.len); + for (g.buffers.items[0..2]) |buffer| @memset(buffer, 0); + try snapshot.restore(&g, live, bytes); + for (g.buffers.items[0..2]) |buffer| for (buffer) |b| try t.expectEqual(value, b); + } +} + +test "reference cache: partial snapshot allocation stays owned and is reused or freed at close" { + const t = std.testing; + for ([_]bool{ false, true }) |retry| { + var g: TestGpu = .{}; + defer g.deinit(); + const bytes: [2]usize = .{ 4, 8 }; + const live: [2]u64 = .{ try g.alloc(bytes[0]), try g.alloc(bytes[1]) }; + var snapshot: Snapshot = .{}; + g.fail_at = 3; + try t.expectError(error.OutOfMemory, snapshot.capture(&g, live, bytes)); + try t.expectEqual(@as(usize, 3), g.buffers.items.len); + try t.expectEqual(@as(u64, 3), snapshot.buffers[0]); + try t.expectEqual(@as(u64, 0), snapshot.buffers[1]); + if (retry) { + g.fail_at = null; + try snapshot.capture(&g, live, bytes); + try t.expectEqual(@as(usize, 4), g.buffers.items.len); + try t.expectEqual(@as(u64, 3), snapshot.buffers[0]); + } + } +} diff --git a/zig/src/native/drama_host.zig b/zig/src/native/drama_host.zig index 467fb3e77a..df77baaddb 100644 --- a/zig/src/native/drama_host.zig +++ b/zig/src/native/drama_host.zig @@ -5,6 +5,7 @@ const d3 = @import("drama3"); const completion = d3.completion; const decode = d3.speech_decode; const rounds = @import("core").frame_rounds; +const reference = d3.reference_cache; const is_cuda = !@hasDecl(d3, "gpu"); const cuda = if (is_cuda) @import("cuda") else struct {}; const Limit = if (is_cuda) @import("drama_limit.zig").Limit else struct { @@ -12,7 +13,15 @@ const Limit = if (is_cuda) @import("drama_limit.zig").Limit else struct { fn deinit(_: *@This()) void {} }; const Warm = if (is_cuda) struct {} else d3.residency.Warm; -const Session = struct { e: *d3.Engine, t: *d3.tts.Tts }; +const Session = struct { + e: *d3.Engine, + t: *d3.tts.Tts, + /// Valid only while this lane's prefix KV remains untouched; snapshots own no KV copies. + prefix: reference.Prefix = .{}, + cache: Snap = .{}, +}; +/// Metal snapshot buffers belong to the shared GPU owner; CUDA keeps its uncached path. +const Snap = if (is_cuda) struct {} else d3.Snapshot; const Job = struct { id: api.Id, request: *const api.SpeechRequest, @@ -173,18 +182,48 @@ const Host = struct { } fn advance(h: *Host, work: rounds.Work) !void { const i = work.lane; - const s = h.sessions[i]; + const lane = &h.sessions[i]; + errdefer if (!is_cuda) lane.prefix.invalidate(); + const s = lane.*; const j = &h.jobs[i].?; const r = j.request; if (work.stage == .prompt) { - if (!j.state.started) { - try s.e.startPrompt(r.frames); - if (is_cuda) try s.t.c.reset(s.t.cs) else try s.t.c.reset(); - j.state.started = true; + if (is_cuda) { + if (!j.state.started) { + try s.e.startPrompt(r.frames); + try s.t.c.reset(s.t.cs); + j.state.started = true; + } + const count = s.e.o.promptRows(r.frames.len - j.state.prompt); + try s.e.promptChunk(r.frames.len, j.state.prompt, count, null); + j.state.prompt += count; + } else { + const ref_rows = reference.prefillRows(r.frames.len, r.reference_rows, s.e.o.max_rows); + if (!j.state.started) { + // Invalidate before reset, GPU work or any fallible preparation for a miss. + if (try lane.prefix.begin(h.gpa, r.frames, ref_rows)) { + try s.e.restore(&lane.cache); + try s.e.beginPromptFrom(r.frames, ref_rows); + j.state.prompt = ref_rows; + std.log.info("drama3: reference cache hit ({d} rows restored)", .{ref_rows}); + } else { + try s.e.startPrompt(r.frames); + j.state.prompt = 0; + } + try s.t.c.reset(); + j.state.started = true; + } + var count = s.e.o.promptRows(r.frames.len - j.state.prompt); + const want_snapshot = ref_rows > 0 and !lane.prefix.valid; + if (want_snapshot and j.state.prompt < ref_rows) count = @min(count, ref_rows - j.state.prompt); + try s.e.promptChunk(r.frames.len, j.state.prompt, count, null); + j.state.prompt += count; + if (want_snapshot and j.state.prompt == ref_rows) { + try s.e.snapshot(&lane.cache); + lane.prefix.saved(); + std.log.info("drama3: reference cache built ({d} rows)", .{ref_rows}); + } } - const count = s.e.o.promptRows(r.frames.len - j.state.prompt); - try s.e.promptChunk(r.frames.len, j.state.prompt, count, null); - j.state.prompt += count; if (j.state.prompt == r.frames.len) { try s.e.finishPrompt(r.frames); h.lock(); @@ -260,6 +299,7 @@ const Host = struct { fn close(p: *anyopaque) void { const h = self(p); drain(h); + for (h.sessions[0..h.queue.lanes]) |*s| if (!is_cuda) s.prefix.deinit(h.gpa); var safe = true; for (h.sessions[0..h.queue.lanes]) |*s| waitSession(s) catch { safe = false; @@ -301,6 +341,7 @@ pub fn open(gpa: std.mem.Allocator, io: std.Io, dir: []const u8, _: []const u8, var made: usize = 0; errdefer while (made > 0) { made -= 1; + if (!is_cuda) h.sessions[made].prefix.deinit(gpa); h.sessions[made].t.destroy(); h.sessions[made].e.destroy(); }; diff --git a/zig/src/server/drama_voice.zig b/zig/src/server/drama_voice.zig index 9591caa3e6..41f24c969d 100644 --- a/zig/src/server/drama_voice.zig +++ b/zig/src/server/drama_voice.zig @@ -97,19 +97,20 @@ const Tokenizer = struct { } }; -pub fn build(a: Allocator, tokenizer: text_mod.Text, input: []const u8, refs: []const prompt.Reference) ![][11]i32 { +pub fn build(a: Allocator, tokenizer: text_mod.Text, input: []const u8, refs: []const prompt.Reference) !prompt.Built { return prompt.build(a, Tokenizer{ .value = tokenizer }, input, refs); } /// A voice's prompt: its own ids and frame rows. -pub const Prompt = struct { prompt: []const u32, frames: []const [11]i32 }; +pub const Prompt = struct { prompt: []const u32, frames: []const [11]i32, reference_rows: u32 = 0 }; /// A voice's prompt ids and frame rows for an input. pub fn promptFor(a: Allocator, tokenizer: text_mod.Text, input: []const u8, refs: []const prompt.Reference) !Prompt { - const frames = try build(a, tokenizer, input, refs); + const built = try build(a, tokenizer, input, refs); + const frames = built.frames; const ids = try a.alloc(u32, frames.len); for (frames, ids) |row, *token| token.* = @intCast(row[0]); - return .{ .prompt = ids, .frames = frames }; + return .{ .prompt = ids, .frames = frames, .reference_rows = built.reference_rows }; } /// A voice's prompt for a speech request: its ids and frame rows; null when the request named no voice. @@ -231,5 +232,6 @@ test "voices: real tokenizer rows equal the CLI path without server postprocessi const refs = [_]prompt.Reference{.{ .transcript = "A small cloud passes.", .codes = &codes }}; const cli_rows = try prompt.build(a, &tok, "A cloud.", &refs); const http_rows = try build(a, hf.text(), "A cloud.", &refs); - try std.testing.expectEqualSlices(u8, std.mem.sliceAsBytes(cli_rows), std.mem.sliceAsBytes(http_rows)); + try std.testing.expectEqualSlices(u8, std.mem.sliceAsBytes(cli_rows.frames), std.mem.sliceAsBytes(http_rows.frames)); + try std.testing.expectEqual(cli_rows.reference_rows, http_rows.reference_rows); } diff --git a/zig/src/server/speech.zig b/zig/src/server/speech.zig index 541cf21ec1..201e993a05 100644 --- a/zig/src/server/speech.zig +++ b/zig/src/server/speech.zig @@ -43,7 +43,7 @@ pub const Format = enum { } }; -const Plan = struct { input: []const u8, prompt: []const u32, frames: ?[]const [11]i32 = null, refs: []const @import("drama_prompt").Reference = &.{}, seed: u64, format: Format, stream: bool, budget_s: ?f64 = null }; +const Plan = struct { input: []const u8, prompt: []const u32, frames: ?[]const [11]i32 = null, refs: []const @import("drama_prompt").Reference = &.{}, reference_rows: u32 = 0, seed: u64, format: Format, stream: bool, budget_s: ?f64 = null }; /// One whole take: its audio, why it ended, and whether it ended in a mid-utterance silence. const Take = struct { pcm: []f32, reason: api.Reason, silence_mid: bool = false, trip: bool = false, jump: ?f64 = null }; @@ -122,7 +122,7 @@ fn prepare(srv: *Server, cx: *Cx, sp: api.SpeechInfo, raw: Value) errors.Refused return cx.fail(.request, drama_prompt.error_prefix ++ " unknown voice {s}; see /v1/voices", .{name}); return cx.fail(.request, drama_prompt.error_prefix ++ " unknown voice {s}; only the default voice is supported", .{name}); }; - if (try @import("drama_voice.zig").requestPrompt(cx, srv.text, sp, input, refs)) |made| return .{ .input = input, .prompt = made.prompt, .frames = made.frames, .refs = refs, .seed = api.seedFor(made.prompt, srv.config.seed_salt), .format = format, .stream = stream }; + if (try @import("drama_voice.zig").requestPrompt(cx, srv.text, sp, input, refs)) |made| return .{ .input = input, .prompt = made.prompt, .frames = made.frames, .refs = refs, .reference_rows = made.reference_rows, .seed = api.seedFor(made.prompt, srv.config.seed_salt), .format = format, .stream = stream }; const prompt = try promptOf(srv, cx, sp, input); return .{ .input = input, .prompt = prompt, .seed = api.seedFor(prompt, srv.config.seed_salt), .format = format, .stream = stream }; } @@ -178,6 +178,7 @@ fn synthesize(srv: *Server, c: *Conn, cx: *Cx, id: []const u8, sp: api.SpeechInf .top_p = drama_prompt.top_p, .top_k = drama_prompt.top_k, .stream = plan.stream, + .reference_rows = plan.reference_rows, }; var box: Mailbox = .{ .io = srv.io, .gpa = srv.gpa }; defer box.deinit(); @@ -222,6 +223,7 @@ fn redraw(srv: *Server, c: *Conn, cx: *Cx, id: []const u8, sp: api.SpeechInfo, p if (made) |voice| { next.prompt = voice.prompt; next.frames = voice.frames; + next.reference_rows = voice.reference_rows; } else { next.prompt = promptOf(srv, cx, sp, input) catch |e| return unused(id, k, fallback, if (e == error.OutOfMemory) "out of memory" else cx.message); next.frames = rows(cx.a, next.prompt) catch |e| return unused(id, k, fallback, @errorName(e)); diff --git a/zig/src/server/speech_stream_test.zig b/zig/src/server/speech_stream_test.zig index 00264dcfcf..b245715b96 100644 --- a/zig/src/server/speech_stream_test.zig +++ b/zig/src/server/speech_stream_test.zig @@ -32,6 +32,7 @@ pub const Script = struct { cancelled: std.atomic.Value(bool) = .init(false), thread: ?std.Thread = null, seeds: [4]u64 = undefined, + reference_rows: [4]u32 = undefined, takes: usize = 0, voices: []const @import("drama_voice.zig").Voice = &.{}, get_voices: bool = false, @@ -55,6 +56,7 @@ pub const Script = struct { } const take = e.takes; e.seeds[take] = request.seed; + e.reference_rows[take] = request.reference_rows; e.takes += 1; if (e.inline_) return run(e, id, request, sink, take); e.thread = std.Thread.spawn(.{}, run, .{ e, id, request, sink, take }) catch return error.Busy; @@ -518,13 +520,14 @@ test "voices: HTTP listing and reference rows use the family CLI builder" { } }; const expected = try family.build(a, CliTokenizer{}, "hi", &refs); - var script: Script = .{ .steps = &.{frame}, .voices = &voices, .expected_rows = expected }; + var script: Script = .{ .steps = &.{frame}, .voices = &voices, .expected_rows = expected.frames }; var s: Served = undefined; try s.start(&script, "HTTP/1.1", "{\"model\":\"m\",\"input\":\"hi\",\"voice\":\"tone\",\"response_format\":\"pcm\"}", 2000); defer s.deinit(); try expectHeader(try s.readHead(), " 200 "); s.join(&script); try std.testing.expect(script.rows_equal); + try std.testing.expectEqual(expected.reference_rows, script.reference_rows[0]); try std.testing.expectEqualStrings(&pcm(1), try s.readRest()); var listing_script: Script = .{ .steps = &.{}, .voices = &voices, .get_voices = true }; var listing: Served = undefined; @@ -553,5 +556,6 @@ test "voices: without configuration the default prompt rows stay unchanged" { try expectHeader(try s.readHead(), " 200 "); s.join(&script); try std.testing.expect(script.rows_equal); + try std.testing.expectEqual(@as(u32, 0), script.reference_rows[0]); try std.testing.expectEqualStrings(&pcm(1), try s.readRest()); }