Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
catalinb13 wants to merge 2 commits into
base: main
Choose a base branch
from
Draft

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
2131bf3
serve: a drafter whose K/V the engine holds elsewhere must not void t…
catalinb13 Oct 28, 2026
e9419c3
docs: the session file's exact-extension rule and the conditional dra…
catalinb13 Oct 28, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 4 additions & 1 deletion docs/DETAILS.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -727,7 +727,10 @@ curl -X POST "http://127.0.0.1:8080/slots/0?action=restore" -H "Content-Type: ap
host caches (retained K/V, oldest unpinned parked conversations, unpinned checkpoints), and the server waits up to 30 s for
RAM accounting to catch up. It trades later prefix-cache hits for room to save and cannot guarantee a save.
- Clients still send their messages and images; the file holds engine state and token identity, not a chat export. A follow-up
that continues the restored conversation reuses its state; a short tail may be read again. Only the deepest checkpoint is saved.
that continues the restored conversation reuses its state; a short tail may be read again. Only the deepest checkpoint is saved,
and the continuation must extend it exactly: a prompt that branches inside the saved conversation is read again from the start,
because a checkpoint's recurrent state exists only at its own boundary. The draft layer's K/V is in the file only when the
drafter owns a K/V ring: on the Apple Silicon engine the decode window runs the draft rounds itself, so its file carries none.
- A file is bound to the model files (sampled by size and first and last MiB, not by path) and to the settings that change what
the cells mean (engine version, CUDA or HIP, `--kv`, `STRATA_KV_ROT`, `--kv-resident`, `--max-context`, `--mtp-window`, the
rope configuration, a loaded control vector and the arithmetic switches). Sampling, seeds, draft tuning and the expert
Expand Down
27 changes: 24 additions & 3 deletions src/core/conversation_state.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -481,7 +481,22 @@ bool conversation_snapshot_bytes(const ConversationView& view, const SessionStat
for (uint64_t i = 0; i < qsa + (draft ? 1 : 0); ++i) {
const auto& st = i == qsa ? *draft : owned(ss, (size_t) i);
const size_t n = conversation_kv_bytes(st, g, upto, i != qsa);
if (!n || !add(bytes, n)) return fail(error, "invalid or overflowing K/V byte estimate");
if (!n || !add(bytes, n)) {
// the estimate hides WHY it went to zero (the reason lives in the capture's own error string);
// a refused save is rare enough to report the reason and the failing state's identity every time
std::string why;
size_t nb = 0;
conversation_kv_capture_bytes(ConversationKv{}, st, g, upto, i != qsa, nb, why);
std::fprintf(stderr, "strata: K/V byte estimate failed: pool %llu of %llu (%s) upto %lld: %s - "
"mode %d int8 %d q4 %d hybrid %d rot %d n_slots %lld n_pages %lld "
"max_cells %lld idx_rows %lld host %d\n",
(unsigned long long) i, (unsigned long long)(qsa + (draft ? 1 : 0)),
i == qsa ? "draft" : "session", (long long) upto, why.c_str(),
(int)st.kv_mode, (int)st.kv_int8, (int)st.kv_q4, (int)st.kv_hybrid, (int)st.kv_rot,
(long long)st.n_slots, (long long)st.n_pages, (long long)st.max_cells,
(long long)st.idx_pooled_rows, (int)st.host.present());
return fail(error, "invalid or overflowing K/V byte estimate");
}
}
return true;
}
Expand Down Expand Up @@ -667,8 +682,14 @@ bool conversation_session_read_limits(SessionReadLimits& limits, const SessionSt
const int64_t upto = (int64_t) std::min<uint64_t>(tokens, (uint64_t) std::max<int64_t>(st.max_cells, 0));
if (!conversation_kv_part_sizes(st, g, upto, true, l.max_kv_bytes[j], error)) return false;
}
const int64_t dupto = (int64_t) std::min<uint64_t>(tokens, (uint64_t) std::max<int64_t>(draft.max_cells, 0));
if (!conversation_kv_part_sizes(draft, g, dupto, false, l.max_kv_bytes.back(), error)) return false;
// an empty drafter (the Apple Silicon decode window runs the draft rounds itself, so the session-side
// MtpDrafter keeps a default state): a file written this way has no draft part, so its bounds stay {} -
// conversation_kv_part_sizes would refuse the empty state and reject the engine's own file
if (draft.max_cells <= 0) l.max_kv_bytes.back() = {};
else {
const int64_t dupto = (int64_t) std::min<uint64_t>(tokens, (uint64_t) std::max<int64_t>(draft.max_cells, 0));
if (!conversation_kv_part_sizes(draft, g, dupto, false, l.max_kv_bytes.back(), error)) return false;
}
limits = std::move(l);
return true;
}
Expand Down
26 changes: 14 additions & 12 deletions src/program/generate.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10646,8 +10646,8 @@ int main(int argc, char** argv) {
if (!conversations.enabled() || !live_ok || live.empty()) return true;
const size_t n_st = stages.size();
ParkSource src{live, live_imgs, checks, cvec_cached, ss, {},
(n_st > 0 && !mtp_first) || !use_mtp ? nullptr : &mtp.kv_state(),
use_mtp && !mtp_first && n_st > 0 ? &mtp.kv_state() : nullptr, true, ""};
(n_st > 0 && !mtp_first) || !use_mtp || mtp.kv_state().max_cells <= 0 ? nullptr : &mtp.kv_state(),
use_mtp && mtp.kv_state().max_cells > 0 && !mtp_first && n_st > 0 ? &mtp.kv_state() : nullptr, true, ""};
for (auto& st : stages) src.stage_sessions.push_back(&st->ss);
return park_source(src, held);
};
Expand Down Expand Up @@ -13622,7 +13622,8 @@ int main(int argc, char** argv) {
const double read_ms = ms();
// the whole image against this engine, still without any device write
blocking("validate", bytes);
if (!strata::core::conversation_snapshot_validate(image, ss, g, mtp.kv_state(), err)) {
if (!strata::core::conversation_snapshot_validate(image, ss, g,
use_mtp && mtp.kv_state().max_cells > 0 ? &mtp.kv_state() : nullptr, err)) {
refuse(err);
continue;
}
Expand All @@ -13635,7 +13636,8 @@ int main(int argc, char** argv) {
live_ok = false;
// host -> device in synchronous copies of the whole state: one bounded allowance
blocking("transfer", bytes);
if (strata::core::conversation_snapshot_restore(image, ss, g, mtp.kv_state(), err) !=
if (strata::core::conversation_snapshot_restore(image, ss, g,
use_mtp && mtp.kv_state().max_cells > 0 ? &mtp.kv_state() : nullptr, err) !=
strata::core::ConversationRestore::restored) {
// validated above: a failure here is a transfer failure, after device writes began - never
// decode from a partial state; the server starts the engine again
Expand Down Expand Up @@ -14113,7 +14115,7 @@ int main(int argc, char** argv) {
if (incoming) for (size_t i = 0; i < stages.size(); ++i) {
const strata::core::OnDevice on(stages[i]->dev);
if (!strata::core::conversation_snapshot_validate(incoming->stage_images[i], stages[i]->ss, g,
use_mtp && !mtp_first && i + 1 == stages.size() ? &mtp.kv_state() : nullptr, err)) { // the draft: the last stage's
use_mtp && mtp.kv_state().max_cells > 0 && !mtp_first && i + 1 == stages.size() ? &mtp.kv_state() : nullptr, err)) { // the draft: the last stage's
std::fprintf(stderr, "strata serve: conversation cache: discard invalid stage snapshot (%s)\n",
err.c_str());
incoming.reset();
Expand All @@ -14122,7 +14124,7 @@ int main(int argc, char** argv) {
}
}
if (incoming && !strata::core::conversation_snapshot_validate(*incoming, ss, g,
use_mtp && (stages.empty() || mtp_first) ? &mtp.kv_state() : nullptr, err)) {
use_mtp && mtp.kv_state().max_cells > 0 && (stages.empty() || mtp_first) ? &mtp.kv_state() : nullptr, err)) {
std::fprintf(stderr, "strata serve: conversation cache: discard invalid snapshot (%s)\n", err.c_str());
incoming.reset();
err.clear();
Expand Down Expand Up @@ -14171,8 +14173,8 @@ int main(int argc, char** argv) {
if (borrow) {
const strata::core::SavedConversation* donor = conversations.find(parked.serial);
if (donor == nullptr ||
!strata::core::conversation_snapshot_validate(*donor, ss, g, use_mtp ? &mtp.kv_state() : nullptr,
err)) {
!strata::core::conversation_snapshot_validate(*donor, ss, g,
use_mtp && mtp.kv_state().max_cells > 0 ? &mtp.kv_state() : nullptr, err)) {
std::fprintf(stderr, "strata serve: conversation cache: cannot borrow from a parked conversation "
"(%s)\n", err.empty() ? "gone" : err.c_str());
borrow = false;
Expand Down Expand Up @@ -14341,7 +14343,7 @@ int main(int argc, char** argv) {
const strata::core::SavedConversation* donor = conversations.find(parked.serial); // still parked
if (donor == nullptr ||
strata::core::conversation_snapshot_restore_prefix(*donor, parked.tokens, ss, g,
use_mtp ? &mtp.kv_state() : nullptr, err) !=
use_mtp && mtp.kv_state().max_cells > 0 ? &mtp.kv_state() : nullptr, err) !=
strata::core::ConversationRestore::restored) {
// validated above and guarded while parking: a failure here is fatal, as for a whole restore
std::printf("ERR restoring a borrowed conversation prefix: %s\n",
Expand All @@ -14367,7 +14369,7 @@ int main(int argc, char** argv) {
if (incoming) {
const auto t0 = Clock::now();
if (strata::core::conversation_snapshot_restore(*incoming, ss, g,
use_mtp && (stages.empty() || mtp_first) ? &mtp.kv_state() : nullptr, err) !=
use_mtp && mtp.kv_state().max_cells > 0 && (stages.empty() || mtp_first) ? &mtp.kv_state() : nullptr, err) !=
strata::core::ConversationRestore::restored) {
// Already prevalidated above: a failure here is fatal, never
// permission to decode from a partially restored session.
Expand All @@ -14377,15 +14379,15 @@ int main(int argc, char** argv) {
for (size_t i = 0; i < stages.size(); ++i) { // the later stages' parts
const strata::core::OnDevice on(stages[i]->dev);
if (strata::core::conversation_snapshot_restore(incoming->stage_images[i], stages[i]->ss, g,
use_mtp && !mtp_first && i + 1 == stages.size() ? &mtp.kv_state() : nullptr, err) !=
use_mtp && mtp.kv_state().max_cells > 0 && !mtp_first && i + 1 == stages.size() ? &mtp.kv_state() : nullptr, err) !=
strata::core::ConversationRestore::restored) {
std::printf("ERR restoring parked conversation (stage CUDA%d): %s\n", stages[i]->dev,
err.c_str());
return 1;
}
if (cudaDeviceSynchronize() != cudaSuccess) { std::printf("ERR stage sync\n"); return 1; }
}
if (use_mtp && std::getenv("STRATA_SNAPSHOT_VERIFY") != nullptr) {
if (use_mtp && mtp.kv_state().max_cells > 0 && std::getenv("STRATA_SNAPSHOT_VERIFY") != nullptr) {
uint64_t draft_hash = 0;
const auto& draft_kv = stages.empty() || mtp_first ? incoming->kv.back() : incoming->stage_images.back().kv.back();
const strata::core::OnDevice on_d(stages.empty() || mtp_first ? 0 : stages.back()->dev);
Expand Down