Skip to content

[Engine] serve: a zero drafter ring (Metal --mtp) must not void the conversation cache or kill the engine - #4384

Draft
catalinb13 wants to merge 2 commits into
Niko1221:mainfrom
catalinb13:pr-metal-mtp-draft-gate
Draft

catalinb13 wants to merge 2 commits into
Niko1221:mainfrom
catalinb13:pr-metal-mtp-draft-gate

Conversation

@catalinb13

Copy link
Copy Markdown

The title must start with its area: [Engine], [GUI], [TUI], [Benchmark], [Docs] or [Other], for example "[Engine] fix decode stop on long prompts" or "[GUI] keep the model list after a tab switch".

Problem and result

On an Apple Silicon box that runs --mtp, the decode window performs the draft rounds itself
(metal/src/engine/window.cpp); the session-side MtpDrafter is loaded but never allocates cells, so
mtp.kv_state() is a default state: max_cells == 0 and every readback through it is zero. The serve
paths that pass &mtp.kv_state() whenever use_mtp is set then describe a draft part that does not
exist, and the conversation cache turns against itself:

  • every parked image fails conversation_snapshot_validate -> discard invalid snapshot -> the whole
    conversation re-prefills on the next turn (observed on our M1 Max: 100772 tokens re-read, 382.7 s);
  • a disk slot file written with the same drafter cannot be read back: the limits reserve a draft part the
    file never carried, so the read refuses its own file;
  • STRATA_SNAPSHOT_VERIFY compares the last session K/V part against the zero drafter, prints
    ERR verifying... to the protocol stdout and return 1 stops the engine;
  • when the K/V byte estimate refuses (the disk-save route of the same problem), the 500 carries no reason
    at all - the estimate collapses to "invalid or overflowing K/V byte estimate".

After this change the same gates pass nullptr for a drafter that owns no ring, the images carry no draft
part, validate, park, borrow and restore round-trip, and the verify readback runs only where there is a
session-side ring to read. Outputs stay byte-identical to the no-park run; speculative decode is untouched
(it rides the window's own buffers).

Related discussion

Approach

One predicate: mtp.kv_state().max_cells > 0 - "this drafter owns a ring" - added to every serve callsite
that passes &mtp.kv_state() into park, borrow, restore, stage validate and the disk slot paths, plus the
STRATA_SNAPSHOT_VERIFY block. conversation_session_read_limits keeps the draft bounds empty for a
drafter that owns no ring, because a file written this way carries no draft part and
conversation_kv_part_sizes would refuse the empty state. The K/V byte estimate now names the failing
pool, its identity (draft/session, index) and the capture's own reason on stderr before returning the
refusal; that dump is what let us find the root cause through a single HTTP 500.

Compatibility: on CUDA, HIP and SYCL the drafter allocates its ring at load, so every gate takes the true
side and those paths are unchanged line-for-line. --mtp is optional; without it use_mtp is false and
nothing changes. Files written before this change by a CUDA engine still restore, and files written by a
Metal engine after it carry no draft part by construction.

Validation

Passed:

  • Metal build of this exact diff: cmake -G Ninja -B build-metal -S metal -DMETAL_CPP_DIR=third_party/metal-cpp -DSTRATA_GGML_DIR=third_party/llama.cpp && ninja -C build-metal strata at 5e8a4cc on macOS arm64 - links
    clean in 27 s (one pre-existing warning, expert_source.hpp:1267 missing-override, untouched here);
    ./build-metal/strata --help runs.
  • Runtime behaviour of the same gates (in an equivalent fork build carrying this diff plus the window-capture
    work of [Engine] metal: the conversation cache and session files keep the decode window's own K/V (#2502) #3541's shape, so the disk path is exercised end to end): M1 Max 64 GB, huihui
    Qwen3.8-Flash-Next UD-Q4_K_XL, --mtp ... --spec 2, conversation cache with 6 slots. Parked
    100882 tokens in 141 ms; disk save of the same conversation 1.67 GB in 1112.9 ms; boot restore 493 ms;
    first follow-up reused 100724 of 100772 tokens in 7.4 s against 382.7 s cold; the recall needle answered
    correctly after restore; a full re-tokenized continuation kept its needle; greedy decode 22.9-23.1 tok/s
    with --spec 2 (21.8 without MTP); the K/V parity check passed at 2.0e-05 relative L2. Without the gates
    the same build logged discard invalid snapshot on every park and re-prefilled, and
    STRATA_SNAPSHOT_VERIFY=1 stopped the engine at the restore with ERR verifying the drafter's K/V.

Failed:

  • none observed with the gates in place.

Not run:

  • CUDA, HIP and SYCL builds: this machine is Apple Silicon with no GPU toolkit. The changed expressions are
    backend-neutral and the true side of every gate reproduces the current code, but a CUDA build and a
    parking_test.py run with --mtp on an NVIDIA box would close this gap - that is the check a maintainer
    or CI should ask for;
  • executing this branch's binary against a live model: the 104 GB pack needs the whole 64 GB machine, and
    the runtime evidence above comes from the equivalent fork build;
  • tools/parking_test.py on this branch as-is (it expects the disk-save capture port of [Engine] metal: the conversation cache and session files keep the decode window's own K/V (#2502) #3541, which this
    PR deliberately does not duplicate).

Limits

  • On Metal the draft ring's contents are still not parked, because the window owns them; a restored
    conversation restarts speculation with a cold draft cache (observed: first decode after a full re-prefill
    turn 18-20 tok/s, next request back at 20.1). [Engine] Metal: working session cache, and the 12 GiB its arena was reserving for nothing #3564 is the proper fix for parking the draft itself.
  • The park path keeps only the deepest checkpoint on disk; a continuation that branches inside the live
    session re-reads from zero by design (a GDN recurrent state exists only at checkpoint boundaries). This
    PR neither changes nor hides that: it only stops a zero ring from discarding the cache wholesale.
  • Marked draft until a CUDA/HIP/SYCL build and a hardware parking_test.py run with --mtp are reported.

…he conversation cache

On Apple Silicon the decode window runs the MTP draft rounds itself (metal/src/engine/window.cpp):
the session-side MtpDrafter is loaded but keeps no cells, so `mtp.kv_state()` is a default state with
max_cells 0 and its readback is all zeros. Every serve path that passed `&mtp.kv_state()` unconditionally
then asked the snapshot code to describe a draft part that does not exist: a parked or saved image failed
`conversation_snapshot_validate` ("discard invalid snapshot" -> the whole conversation re-prefills), a disk
slot read refused its own file, and `STRATA_SNAPSHOT_VERIFY` read the zero state back and exited the engine
with 1.

Gate those callsites on `mtp.kv_state().max_cells > 0` - the same predicate that answers "does this drafter
own a ring" - so the images carry no draft part and validate against a draft-less drafter; keep the draft
bounds empty in `conversation_session_read_limits` for the same state; only run the verify readback when
there is a ring to read; and dump the failing pool, its identity and the capture's own reason when the K/V
byte estimate refuses, because the estimate's 500 carried no diagnosis.

On CUDA, HIP and SYCL the drafter allocates its ring at load, so every gate passes with the true side and
the default path is unchanged. A drafter bound to the window's real buffers (the `mtp_kv_bind` route of
Niko1221#3564) reports max_cells > 0 and its draft part is described again - this is the fallback for the state
that has none, not a competing design.

Measured on an equivalent diff inside a fork Metal build (M1 Max 64 GB, huihui Qwen3.8-Flash-Next
UD-Q4_K_XL, --mtp spec 2): park 100882 tokens in 141 ms, restore in 493 ms, first follow-up reused
100724 of 100772 tokens (7.6 s against 382.7 s cold), the needle answer stayed correct, decode kept
22.9-23.1 tok/s greedy, and the K/V parity check passed at 2.0e-05 relative L2. Without the gates the
same build re-prefilled every turn and `STRATA_SNAPSHOT_VERIFY` stopped the engine.
…ft part

A continuation that branches inside the saved conversation reads again from the start (a checkpoint's recurrent
state exists only at its own boundary), and the draft layer's K/V is in the file only when the drafter owns a
ring - on the Apple Silicon engine the decode window runs the draft rounds itself, so its files carry none.
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