Skip to content

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

Open
plotarmordev wants to merge 1 commit into
ashhart:mainfrom
plotarmordev:drama3/ref-prefill-cache
Open

plotarmordev wants to merge 1 commit into
ashhart:mainfrom
plotarmordev:drama3/ref-prefill-cache

Conversation

@plotarmordev

Copy link
Copy Markdown
Contributor

What this changes

On Metal, a Drama speech request that uses a catalog voice now reuses that voice reference's prompt prefill across requests on the same lane. The first request saves the lane's recurrent state at the end of the reference prompt. A later request with exactly the same reference rows prefills only the rows after them. CUDA stays uncached. The default voice, and requests without a voice, are unchanged.

This builds on the voice-catalog and redraw code now on main (#1120, #1135, #1160, #1162). It is not stacked on any open pull request.

  • Boundary: the prompt builder reports reference_rows, where the invariant reference prompt ends before the request text. SpeechRequest.reference_rows carries it, and 0 turns reuse off. The saved row stops on an existing prefill chunk boundary, and the tower's lookahead row stays inside the invariant prompt. A cached request therefore makes the same chunk cuts as an uncached one.
  • Identity: a hit needs the exact owned prefix frames to match, including the tower's lookahead row. A hash alone is not enough.
  • What is saved: recurrent state only. The KV prefix stays live in its lane. The saved state is invalidated before any default or different-prefix prefill can overwrite that KV, and on any execution error.
  • Ownership: each lane has one reusable GPU-owned snapshot buffer pair, and voice switches reuse it. A partly allocated pair stays owned and is released when the GPU closes.
  • Per-request decode counters are rebuilt after a restore.
  • The server's voice prompt passes reference_rows to the first take and to a plain-speech redraw.
  • Logs: drama3: reference cache built (N rows) and drama3: reference cache hit (N rows restored).

@BobClawblaw wrote the change. @jerryfane reviewed it on a Mac, found two defects and repaired them. Both are credited as co-authors. I ported it onto current main.

Receipt

  • Environment, with the TensorFold commit, MLX or PyTorch version, machine and GPU, checkpoint and revision:
    • My checks: commit b0928cc8 on main 4167b876, and some earlier runs on 0d4592d3, which differs from b0928cc8 only in source comments. Linux aarch64, Zig 0.17.0, CPU only, CUDA_VISIBLE_DEVICES empty. No GPU, no Metal device, no MLX or PyTorch.
    • The Mac evidence below is @jerryfane's, not mine. He ran the same change in a downstream Drama patch series on TensorFold d1018b31, not on this branch. Hardware: Mac Studio, Apple M5 Ultra, 256 GB, native Metal, Zig 0.17.0. Checkpoint: fishaudio/drama-3@e0a9218a. Candidate engine c265ca2f (with this change) against baseline 881151ca (the same series without it). Sources: first review, repair and retest, machine-readable results.
    • Since d1018b31, main has gained the Metal residency set (50ddf20d) and moved the voice prompt into drama_voice.zig. On this branch, snapshot buffers go through Gpu.register, which joins the residency set. That path did not exist in his runs, and I have not run it on Metal.
  • Exactness, from tools/bench_concurrent.py --alone --serial or token_sha against "draft": false: speech returns PCM, so PCM is compared by SHA-256. All results in this item are @jerryfane's Mac runs.
    • Corpus: 24 items, two passes per build. All 48 paired PCM captures are byte-identical to the uncached baseline, and all returned HTTP 200.
    • Voice, default and text transitions and the reference variants (last reference row changed, transcript changed, two references, after cancellation) all passed.
    • Streaming with the built-in and the cached voice: 9/9 checks in each of four runs.
    • Two lanes: 8 requests, all matching the serial baseline.
    • Eight warmed voice switches: 0 KiB process-RSS spread. RSS is not device-memory telemetry.
    • His first review of the earlier head found two defects. A → default → A changed the audio, and each voice switch kept another 61.4 MiB. This change contains his repair, and his retest shows both fixed.
  • Decode speed before and after, from tools/bench_openai.py: not measured. I make no speed claim.
  • Prompt speed before and after: the cold first request's ttft_s from tools/bench_cache.py --require-cold on a long prompt, or "n/a" with the reason when the change cannot reach prompt processing: not measured. I make no speed claim. The uncached prompt path keeps its chunk cuts.
  • Tests run, with pass, fail and skip counts (mine, Linux aarch64, CPU only). b0928cc8 differs from 0d4592d3 only in comments: it joins five multi-line comments into single lines and rewords one of them. No code changed. Each result names the commit it ran on.
    • zig build test --summary all on b0928cc8: 302/302 steps, 3647/3661 tests passed, 14 skipped.
      • One other zig build test run on 0d4592d3, made while a second build ran on the same host, failed in code this change does not touch. fate_test reported rank 1's failure ends both ranks with its reason (expected 2, found 1), and rank 0's clean stop ... aborted on an index out of bounds in fate.zig runHooks. Every other run of the suite on that commit passed, including a fresh-cache run with the same counts as above.
    • CUDA_VISIBLE_DEVICES= zig build test-cuda-host --summary all on b0928cc8: 192/192 steps, 1463/1464 tests passed, 1 skipped; the other steps were cached. On 0d4592d3 with a fresh cache: 192/192 steps, 2796/2801 tests passed, 5 skipped.
    • zig build test-drama-host --summary all on b0928cc8: 38/38 tests passed.
    • zig build test-drama-prefill-host --summary all on b0928cc8: 2/2 tests passed.
    • zig build test-drama-owner --summary all on b0928cc8: 6/6 tests passed.
    • zig build check-drama-cuda-host --summary all on 0d4592d3: 4/4 steps.
    • zig build test-golden on b0928cc8: 306/319 equal to the server goldens, 0 unexpected differences. Additional corpus: 26/26 equal.
    • zig build native -Dversion=1.0.0 on 0d4592d3: built.
    • tools/zig/test_speech_flags.py on 0d4592d3: 1 test, OK, against that binary.
    • tools/zig/test_drama_prefill_compare.py on 0d4592d3: 5 tests, OK.
    • zig fmt --check on the changed files, on b0928cc8: clean.
    • Metal host entry points, on 0d4592d3: I cross-compiled the three objects of check-drama-metal-host for aarch64-macos. To do it, I changed a local copy of the build script to retarget the step and skip framework linking. Nothing was linked or run, and no shaders were compiled. All three objects compiled. A deliberately broken call in the Metal hit path made that compile fail, which shows the Metal branch is analysed.
    • Two mutations on 0d4592d3, each run once, both failed a test:
      • Dropping reference_rows from the server's voice prompt failed voices: HTTP listing and reference rows use the family CLI builder.
      • Changing the snapshot-pair buffer count failed reference cache: switches reuse one snapshot pair and restore the latest recurrent state.
    • python3 tools/zig/lean_check.py on b0928cc8: 437 findings, a list identical to base 4167b876. This change adds no findings.
  • Not run, and why:
    • I did not run Metal at all: no Metal build with the SDK, no shader compile, and no device synthesis. I have no Metal host.
    • I ran nothing on a GPU. CUDA keeps its uncached path.
    • No speed or memory measurements.
    • The new tests use the new API, so they do not compile on base.
    • No GPU kernel changed, and frozen test data was not edited.

Checklist

  • Every token still goes through the lane rounds. No serial path, and nothing that needs drafts off.
  • Drafted output equals "draft": false, a resumed prompt equals a fresh one, and concurrent equals solo.
  • No precision traded for speed. If the bits change, the description says which sums change and why.
  • A new or changed GPU kernel is registered with the kernel contract (docs/kernel-contract.md); each backend's SUMMARY line is in the receipt.
  • Prompt processing is no slower than the last release.
  • The description says which platforms I ran, Metal M1 to M5 and CUDA, and which I could not.
  • New tests fail before the change, pass after it, and skip cleanly without their dependency.
  • Comments and docstrings are one line. No measurements or history in the source.
  • No personal data, machine names, internal hosts or local paths. No AI attribution lines.

Keep one live KV prefix and one reusable recurrent-state snapshot pair per
Metal lane. The prompt builder reports the invariant reference boundary;
cache identity includes the tower lookahead row and compares exact owned
frames rather than a hash alone. Preserve original prefill chunk cuts.

Invalidate before incompatible/default prefill and on execution failure.
Reuse GPU-owned buffers across voice switches, including partial-allocation
failures. Rebuild per-request decode counters and leave CUDA uncached.

Add deterministic boundary, exactness, invalidation, ownership and prompt
regressions.

Co-authored-by: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com>
Co-authored-by: jerryfane <38183696+jerryfane@users.noreply.github.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant