feat(drama3): reuse a voice reference's prefill across requests

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(.{

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 {

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");

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");

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) {

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;

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));
+}