diff --git a/dev/tests/engine/kv_first_cache_test.cpp b/dev/tests/engine/kv_first_cache_test.cpp index f6970d1..916219c 100644 --- a/dev/tests/engine/kv_first_cache_test.cpp +++ b/dev/tests/engine/kv_first_cache_test.cpp @@ -1,6 +1,7 @@ #include "TestKvPool.hpp" #include "engine/Cache.hpp" +#include #include #include #include @@ -105,11 +106,13 @@ CacheNamespace cacheNamespace() { struct CacheFixture { test::TestKvBacking backing{4, 100}; KvPool pool{backing}; - engine::Cache cache{pool, cacheNamespace()}; + engine::Cache cache; std::vector prompt; std::vector blocks; - CacheFixture() { + explicit CacheFixture( + uint32_t boundaryStates = engine::Cache::kDefaultBoundaryStates) + : cache(pool, cacheNamespace(), boundaryStates) { for (uint32_t block = 0; block < 4; ++block) { for (uint32_t row = 0; row < KvCache::pageTokens; ++row) prompt.push_back(1000 + block * 100 + row); @@ -124,9 +127,13 @@ struct CacheFixture { cache.endRequest(1); } - void publish(uint32_t block, uint64_t bytes = 100) { + void publish(uint32_t block, uint64_t bytes = 100, + engine::StateRole role = engine::StateRole::Ordinary) { cache.publishCompositeState(blocks.at(block), - std::make_shared(bytes)); + std::make_shared(bytes), role); + } + void publishBoundary(uint32_t block) { + publish(block, 100, engine::StateRole::Boundary); } engine::CacheLookup lookup(uint32_t tokens) { @@ -342,7 +349,7 @@ void testCheckpointLookupProbeDoesNotPromote() { CacheFixture fixture; const uint64_t block = fixture.blocks[0]; fixture.cache.publishCompositeState(block, std::make_shared(100), - true); + engine::StateRole::Checkpoint); const StateCheckpoint checkpoint = fixture.cache.checkpointState(block); { auto probe = fixture.lookup(33); @@ -360,11 +367,11 @@ void testCheckpointRetirementRespectsUseAndPublicationIdentity() { const uint64_t block = fixture.blocks[1]; const auto publish = [&] { fixture.cache.publishCompositeState(block, std::make_shared(100), - true); + engine::StateRole::Checkpoint); return fixture.cache.checkpointState(block); }; const StateCheckpoint first = publish(); - require(fixture.cache.reuseCompositeState(block, true), + require(fixture.cache.reuseCompositeState(block, engine::StateRole::Checkpoint), "concurrent progress publication did not deduplicate"); require(fixture.cache.retireCheckpointState(first) && fixture.cache.snapshot().stateCache.entries == 0, @@ -408,9 +415,9 @@ void testRestoredCheckpointsKeepTheirEvictionPriority() { fixture.publish(0, 150); fixture.publish(3, 400); fixture.cache.publishCompositeState(fixture.blocks[1], - std::make_shared(200), true); + std::make_shared(200), engine::StateRole::Checkpoint); fixture.cache.publishCompositeState(fixture.blocks[2], - std::make_shared(300), true); + std::make_shared(300), engine::StateRole::Checkpoint); { auto lookup = fixture.lookup(65); fixture.cache.beginRequest(2); @@ -434,11 +441,81 @@ void testRestoredCheckpointsKeepTheirEvictionPriority() { "pressure eviction was counted as rolling retirement"); } +// States where a prompt's messages begin outlive ordinary states, however +// old, up to a capped number; past the cap the least recently used boundary +// becomes ordinary again. +void testBoundaryStatesOutliveOrdinaryStates() { + { + CacheFixture fixture(2); + fixture.publishBoundary(0); + fixture.publish(1); + fixture.publish(2); + require(fixture.cache.reclaimOneState() && fixture.cache.reclaimOneState() && + fixture.lookup(33).resumeBoundary() == 32 && + fixture.lookup(65).resumeBoundary() == 32, + "ordinary states did not go before an older boundary state"); + require(fixture.cache.reclaimOneState() && + fixture.lookup(33).resumeBoundary() == 0, + "the boundary state was not evicted last"); + } + { + CacheFixture fixture(2); + for (uint32_t block = 0; block < 3; ++block) + fixture.publishBoundary(block); + auto snapshot = fixture.cache.snapshot().stateCache; + require(snapshot.boundaryEntries == 2 && snapshot.boundaryCapacity == 2 && + snapshot.boundaryDemotions == 1, + "publishing past the cap did not demote the oldest boundary"); + require(fixture.cache.reclaimOneState() && + fixture.lookup(33).resumeBoundary() == 0 && + fixture.lookup(65).resumeBoundary() == 64, + "the demoted boundary was not evicted first"); + fixture.publish(3); + require(fixture.cache.reuseCompositeState(fixture.blocks[3], + engine::StateRole::Boundary), + "reuse did not find the published state"); + snapshot = fixture.cache.snapshot().stateCache; + require(snapshot.boundaryEntries == 2 && snapshot.boundaryDemotions == 2, + "reuse as a boundary did not upgrade the state within the cap"); + } +} + +// A hint state starts as the oldest ordinary state, so it goes before older +// end-of-turn states unless a later prompt resumes from it. +void testHintStatesStartColdUntilReused() { + { + CacheFixture fixture; + fixture.publish(0); + fixture.publish(1); + fixture.publish(2, 100, engine::StateRole::Hint); + require(fixture.cache.reclaimOneState() && + fixture.lookup(97).resumeBoundary() == 64, + "a new hint state was not evicted before older ordinary states"); + } + { + CacheFixture fixture; + fixture.publish(0); + fixture.publish(2, 100, engine::StateRole::Hint); + fixture.publish(1); + { + auto lookup = fixture.lookup(97); + require(lookup.resumeBoundary() == 96, "the hint state was not found"); + fixture.cache.beginRequest(2); + fixture.cache.restoreRequest(2, lookup); + fixture.cache.endRequest(2); + } + require(fixture.cache.reclaimOneState() && + fixture.lookup(97).resumeBoundary() == 96 && + fixture.lookup(33).resumeBoundary() == 0, + "a reused hint state did not become ordinary"); + } +} + void testCheckpointReclaimPrecedesOlderKv() { CacheFixture fixture; fixture.publish(0, 150); fixture.cache.publishCompositeState(fixture.blocks[2], - std::make_shared(200), true); + std::make_shared(200), engine::StateRole::Checkpoint); const auto reclaimed = fixture.cache.reclaimOne(); require(reclaimed.madeProgress && reclaimed.reclaimedBytes == 200 && fixture.cache.snapshot().kvCache.blocks == 4 && @@ -453,7 +530,7 @@ void testOptionalReclaimLeavesOrdinaryStateIntact() { CacheFixture fixture; fixture.publish(0, 150); fixture.cache.publishCompositeState(fixture.blocks[2], - std::make_shared(200), true); + std::make_shared(200), engine::StateRole::Checkpoint); require(fixture.cache.reclaimOneState(true) && fixture.cache.snapshot().stateCache.checkpointEntries == 0 && fixture.lookup(33).resumeBoundary() == 32, @@ -468,9 +545,9 @@ void testCheckpointPinsAndBoundaryUpgrade() { fixture.publish(0, 100); fixture.publish(3, 400); fixture.cache.publishCompositeState(fixture.blocks[1], - std::make_shared(200), true); + std::make_shared(200), engine::StateRole::Checkpoint); fixture.cache.publishCompositeState(fixture.blocks[2], - std::make_shared(300), true); + std::make_shared(300), engine::StateRole::Checkpoint); const auto checkpoint = fixture.cache.checkpointState(fixture.blocks[1]); auto first = fixture.lookup(65); auto second = fixture.lookup(97); @@ -493,7 +570,8 @@ void testCheckpointPinsAndBoundaryUpgrade() { fixture.cache.lookup(fixture.prompt).resumeBoundary() == 64, "upgraded checkpoint did not join ordinary LRU"); - require(fixture.cache.reuseCompositeState(fixture.blocks[1], true) && + require(fixture.cache.reuseCompositeState(fixture.blocks[1], + engine::StateRole::Checkpoint) && !fixture.cache.checkpointState(fixture.blocks[1]), "optional publication downgraded an ordinary boundary"); } @@ -522,7 +600,7 @@ void testCheckpointPressurePreservesHotPrefix() { require(cache.ensureTokens(3, 32).granted(), "cold KV admission failed"); const uint64_t coldBlock = cache.publishCommittedBlocks(3, cold, 32); cache.publishCompositeState(coldBlock, std::make_shared(budget), - true); + engine::StateRole::Checkpoint); require(budget.used == SharedBudget::capacity, "checkpoint did not fill the shared allocation budget"); const TokenAdmission denied = cache.ensureTokens(3, 64); @@ -540,7 +618,7 @@ void testCheckpointPressurePreservesHotPrefix() { void testLogicalKvPressureStillReclaimsPages() { CacheFixture fixture; fixture.cache.publishCompositeState(fixture.blocks[0], - std::make_shared(200), true); + std::make_shared(200), engine::StateRole::Checkpoint); fixture.cache.beginRequest(2); require(fixture.cache.ensureTokens(2, 32).granted() && fixture.cache.snapshot().kvCache.blocks == 3 && @@ -570,6 +648,8 @@ int main() { testDuplicateProbePromotesStateWithoutLookupAccounting(); testUnifiedRecencyAndPhysicalReclaimAccounting(); testFinishedRequestLeavesTailKvBeforeItsState(); + testBoundaryStatesOutliveOrdinaryStates(); + testHintStatesStartColdUntilReused(); std::cout << "KV-first cache tests passed\n"; return 0; } catch (const std::exception &error) { diff --git a/dev/tests/engine/protocol_test.cpp b/dev/tests/engine/protocol_test.cpp index f2aea2e..73c7981 100644 --- a/dev/tests/engine/protocol_test.cpp +++ b/dev/tests/engine/protocol_test.cpp @@ -123,6 +123,7 @@ RequestFrame exampleRequest() { request.seed = 0xfedcba9876543210ULL; request.cohort = Cohort::Constrained; request.constraint = ConstraintMode::TokenMask; + request.stateHints = {1, 3}; return request; } @@ -135,24 +136,37 @@ void testRequestWireAndRoundTrip() { return; const auto &wire = *serialized.value; - CHECK(test, wire.size() == - kFrameHeaderBytes + 56 + request.promptTokens.size() * 4); + const size_t words = request.promptTokens.size() + request.stateHints.size(); + CHECK(test, wire.size() == kFrameHeaderBytes + 60 + words * 4); CHECK(test, std::string(wire.begin(), wire.begin() + 4) == "SPLH"); CHECK(test, loadU16(wire, 4) == kProtocolVersion); CHECK(test, loadU16(wire, 6) == kFrameHeaderBytes); CHECK(test, loadU16(wire, 8) == static_cast(FrameType::Request)); CHECK(test, loadU16(wire, 10) == 0); - CHECK(test, loadU64(wire, 12) == 56 + request.promptTokens.size() * 4); + CHECK(test, loadU64(wire, 12) == 60 + words * 4); CHECK(test, loadU32(wire, 20) == 0); CHECK(test, loadU64(wire, kFrameHeaderBytes) == request.requestId); CHECK(test, loadU32(wire, kFrameHeaderBytes + 31) == request.promptTokens.size()); - CHECK(test, loadU32(wire, kFrameHeaderBytes + 56) == 0); - CHECK(test, loadU32(wire, kFrameHeaderBytes + 56 + 16) == 0xffffffffU); + CHECK(test, loadU32(wire, kFrameHeaderBytes + 56) == request.stateHints.size()); + CHECK(test, loadU32(wire, kFrameHeaderBytes + 60) == 0); + CHECK(test, loadU32(wire, kFrameHeaderBytes + 60 + 16) == 0xffffffffU); + CHECK(test, loadU32(wire, kFrameHeaderBytes + 60 + 20) == 1); + CHECK(test, loadU32(wire, kFrameHeaderBytes + 60 + 24) == 3); RequestFrame decoded = roundTrip(request); CHECK(test, decoded == request); + // Hints must ascend strictly, start after the first token and stay inside + // the prompt. + for (const std::vector &hints : + {std::vector{3, 1}, std::vector{2, 2}, + std::vector{0}, std::vector{5}}) { + RequestFrame invalid = exampleRequest(); + invalid.stateHints = hints; + CHECK(test, !serializeMessage(Message{invalid})); + } + } std::vector everyOtherMessage() { diff --git a/dev/tests/engine/test_protocol_python.py b/dev/tests/engine/test_protocol_python.py index cae8975..ec5f43d 100644 --- a/dev/tests/engine/test_protocol_python.py +++ b/dev/tests/engine/test_protocol_python.py @@ -15,22 +15,22 @@ REQUEST_GOLDEN = ( - "53504c4806001800010000004c0000000000000000000000efcdab8967452301" + "53504c480700180001000000500000000000000000000000efcdab8967452301" "000201008098281765060040a5ae02000000000080000005000000cdcc4c3f33" - "33733f200000001032547698badcfe0000000000010000002a00000000000080" - "ffffffff" + "33733f200000001032547698badcfe000000000000000000010000002a000000" + "00000080ffffffff" ) ERROR_GOLDEN = ( - "53504c4806001800050100002700000000000000000000000200000000000000" + "53504c4807001800050100002700000000000000000000000200000000000000" "0000090000000c0000006770755f6661756c744d6574616c206661696c6564" ) STATUS_GOLDEN = ( - "53504c4806001800070100002d00000000000000000000002803000000000000" + "53504c4807001800070100002d00000000000000000000002803000000000000" "050000007b22736368656d615f76657273696f6e223a342c227265616479223a" "747275657d" ) INITIAL_MASK_GOLDEN = ( - "53504c4806001800030100001800000000000000000000005b00000000000000" + "53504c4807001800030100001800000000000000000000005b00000000000000" "06000000000000000400000000000000" ) @@ -340,14 +340,26 @@ def test_request_header_and_binary_prompt(self): self.assertEqual(wire[:4], b"SPLH") self.assertEqual( struct.unpack_from("" else super().convert_tokens_to_ids(token) + ) + + def hints(self, tokens, tokenizer=None): + app = request_frontend.Frontend( + tokenizer or self.VocabTokenizer(), None, "test-model", 4096, 16, 5, 1 + ) + return app.state_hints(tokens) + + def test_system_end_and_last_message_start(self): + # system, user, assistant, user, then the reply being generated + tokens = [7, 1, 1, 7, 2, 2, 7, 3, 3, 7, 4, 4, 7, 5] + self.assertEqual(self.hints(tokens), (3, 9)) + + def test_one_message_after_the_system_prompt_gives_one_hint(self): + self.assertEqual(self.hints([7, 1, 7, 2, 2, 7, 3]), (2,)) + + class CharTokenizer(VocabTokenizer): + """One token per character, with <|im_start|> as one special token.""" + + START = "<|im_start|>" + + def encode(self, text): + ids, offsets, at = [], [], 0 + while at < len(text): + if text.startswith(self.START, at): + ids.append(7) + offsets.append((at, at + len(self.START))) + at += len(self.START) + else: + ids.append(1000 + ord(text[at])) + offsets.append((at, at + 1)) + at += 1 + return ids, offsets + + def __call__(self, text, **kwargs): + ids, offsets = self.encode(text) + result = {"input_ids": ids} + if kwargs.get("return_offsets_mapping"): + result["offset_mapping"] = offsets + return result + + def test_last_message_content_parts_are_hinted(self): + tokenizer = self.CharTokenizer() + text = ( + "<|im_start|>system\nrules<|im_start|>user\n" + "first\nWhat is new?\nlast" + "<|im_start|>assistant\n" + ) + tokens, _ = tokenizer.encode(text) + message = { + "role": "user", + "content": [ + {"type": "text", "text": "first"}, + {"type": "text", "text": "What is new?"}, + {"type": "text", "text": "last"}, + ], + } + app = request_frontend.Frontend(tokenizer, None, "test-model", 4096, 16, 5, 1) + starts = [i for i, token in enumerate(tokens) if token == 7] + question = starts[1] + len( + tokenizer.encode(text[text.index("<|im_start|>user") : text.index("What")])[ + 0 + ] + ) + reminder = starts[1] + len( + tokenizer.encode( + text[text.index("<|im_start|>user") : text.index("last")] + )[0] + ) + self.assertEqual( + app.state_hints(tokens, text, message), (starts[1], question, reminder) + ) + + def test_no_hints_without_message_starts(self): + self.assertEqual(self.hints([1, 2, 3]), ()) + self.assertEqual(self.hints([7, 1, 2]), ()) + self.assertEqual(self.hints([7, 1, 7, 2], FakeTokenizer()), ()) + + if __name__ == "__main__": unittest.main() diff --git a/runtime/engine/Cache.cpp b/runtime/engine/Cache.cpp index 0ba5f65..b5187c2 100644 --- a/runtime/engine/Cache.cpp +++ b/runtime/engine/Cache.cpp @@ -6,8 +6,9 @@ namespace splash::engine { -Cache::Cache(KvPool &pool, CacheNamespace cacheNamespace) - : pool_(pool), kv_(pool, cacheNamespace, recency_), states_(kv_, recency_) { +Cache::Cache(KvPool &pool, CacheNamespace cacheNamespace, uint32_t boundaryStates) + : pool_(pool), kv_(pool, cacheNamespace, recency_), + states_(kv_, recency_, boundaryStates) { } void Cache::beginRequest(uint64_t requestId) { @@ -198,14 +199,14 @@ uint64_t Cache::blockAt(uint64_t requestId, uint32_t boundary) const { return active.cachedBlocks[index]; } -bool Cache::reuseCompositeState(uint64_t kvBlock, bool checkpoint) { - return states_.touchIfResident(kvBlock, checkpoint); +bool Cache::reuseCompositeState(uint64_t kvBlock, StateRole role) { + return states_.touchIfResident(kvBlock, role); } void Cache::publishCompositeState(uint64_t kvBlock, std::shared_ptr state, - bool checkpoint) { - states_.publish(kvBlock, std::move(state), checkpoint); + StateRole role) { + states_.publish(kvBlock, std::move(state), role); } StateCheckpoint Cache::checkpointState(uint64_t kvBlock) const noexcept { @@ -248,7 +249,7 @@ CacheReclaimResult Cache::reclaimOne(CacheReclaimMode mode) { checkpoint ? std::nullopt : oldestStateFreeKvBlock(); if (!state && !kv) return {}; - if (state && (!kv || state->lastUsed <= kv->lastUsed)) { + if (state && (!kv || (!state->protectedState && state->lastUsed <= kv->lastUsed))) { const StateEviction eviction = states_.evict(state->id); if (!eviction.evicted) throw std::logic_error("state eviction candidate became pinned"); @@ -260,10 +261,11 @@ CacheReclaimResult Cache::reclaimOne(CacheReclaimMode mode) { : 0}; } -bool Cache::reclaimOneState(bool checkpointsOnly) { +bool Cache::reclaimOneState(bool checkpointsOnly, bool spareProtected) { const std::optional state = states_.evictionCandidate(); return state && (!checkpointsOnly || states_.checkpoint(state->id)) && + (!spareProtected || !state->protectedState) && states_.evict(state->id).evicted; } diff --git a/runtime/engine/Cache.hpp b/runtime/engine/Cache.hpp index 4b1d8df..721074c 100644 --- a/runtime/engine/Cache.hpp +++ b/runtime/engine/Cache.hpp @@ -67,7 +67,11 @@ enum class CacheReclaimMode { ReuseBacking, ReleaseBacking }; // composite states. Physical recurrent-state cells remain model-owned. class Cache final { public: - Cache(KvPool &pool, CacheNamespace cacheNamespace); + // Boundary states kept ahead of all others (StateCache); enough for the + // system prompts of a session and its side agents. + static constexpr uint32_t kDefaultBoundaryStates = 4; + Cache(KvPool &pool, CacheNamespace cacheNamespace, + uint32_t boundaryStates = kDefaultBoundaryStates); Cache(const Cache &) = delete; Cache &operator=(const Cache &) = delete; @@ -91,10 +95,11 @@ class Cache final { uint64_t requestId, std::span exactTokens, uint32_t committedTokens); [[nodiscard]] uint64_t blockAt(uint64_t requestId, uint32_t boundary) const; - [[nodiscard]] bool reuseCompositeState(uint64_t kvBlock, bool checkpoint = false); + [[nodiscard]] bool reuseCompositeState(uint64_t kvBlock, + StateRole role = StateRole::Ordinary); void publishCompositeState(uint64_t kvBlock, std::shared_ptr state, - bool checkpoint = false); + StateRole role = StateRole::Ordinary); [[nodiscard]] StateCheckpoint checkpointState(uint64_t kvBlock) const noexcept; // False only while this exact disposable publication is pinned. bool retireCheckpointState(StateCheckpoint checkpoint) noexcept; @@ -113,8 +118,10 @@ class Cache final { [[nodiscard]] CacheReclaimResult reclaimOne( CacheReclaimMode mode = CacheReclaimMode::ReleaseBacking); // Recycles an unpinned state, preferring checkpoints, for a required state - // publication. KV bytes only return once an extent unmaps. - [[nodiscard]] bool reclaimOneState(bool checkpointsOnly = false); + // publication. KV bytes only return once an extent unmaps. With + // spareProtected, a boundary state is never recycled. + [[nodiscard]] bool reclaimOneState(bool checkpointsOnly = false, + bool spareProtected = false); // Empty resident backing exists but the previous release is still in // flight; more reclaim work becomes possible without evicting anything. [[nodiscard]] bool releaseDeferred() const noexcept; diff --git a/runtime/engine/CacheRecency.hpp b/runtime/engine/CacheRecency.hpp index f1a359c..900e401 100644 --- a/runtime/engine/CacheRecency.hpp +++ b/runtime/engine/CacheRecency.hpp @@ -21,6 +21,8 @@ class CacheRecency final { struct CacheEvictionCandidate final { uint64_t id = 0; uint64_t lastUsed = 0; + // A state many prompts resume from: evict plain KV before it. + bool protectedState = false; }; } // namespace splash::engine diff --git a/runtime/engine/Engine.cpp b/runtime/engine/Engine.cpp index 34835e9..da7d508 100644 --- a/runtime/engine/Engine.cpp +++ b/runtime/engine/Engine.cpp @@ -477,6 +477,38 @@ void Engine::configureStateBoundaries(Request &active, uint32_t stateBoundary, } addCandidate(junctionBoundary, Request::StateBoundary::Purpose::Junction); addCandidate(latestReplayBoundary, Request::StateBoundary::Purpose::Replay); + // The server's hints, at the KV block they fall in. The first ends the + // setup other conversations share, so it is kept as a boundary. The others + // are skipped near another planned state, the end of the prompt included: + // a near-duplicate costs a snapshot, a cache entry and a short prefill + // command, and saves almost nothing. + const std::vector &hints = active.request.stateHints; + // The restored state counts as planned; rolling checkpoints do not, since + // they are retired once the next one is published. + const auto nearPlanned = [&](uint32_t block) { + constexpr uint32_t spacing = Request::StateBoundary::kHintSpacingTokens; + if (latestReplayBoundary - block < spacing || + (block > stateBoundary && block - stateBoundary < spacing)) + return true; + return std::any_of(active.stateBoundaries.begin(), + active.stateBoundaries.end(), + [&](const Request::StateBoundary &planned) { + return planned.purpose != + Request::StateBoundary::Purpose::Checkpoint && + std::max(planned.tokens, block) - + std::min(planned.tokens, block) < + spacing; + }); + }; + for (size_t i = 0; i < hints.size(); ++i) { + const uint32_t block = hints[i] / KvCache::pageTokens * KvCache::pageTokens; + if (block >= latestReplayBoundary) + continue; + if (!i) + addCandidate(block, Request::StateBoundary::Purpose::Boundary); + else if (!nearPlanned(block)) + addCandidate(block, Request::StateBoundary::Purpose::Hint); + } std::sort(active.stateBoundaries.begin(), active.stateBoundaries.end(), [](const Request::StateBoundary &left, const Request::StateBoundary &right) { @@ -531,16 +563,32 @@ void Engine::publishReachedStateBoundaries(Request &active, promptProcessed) { const Request::StateBoundary objective = active.stateBoundaries[active.stateBoundaryCursor++]; - const bool checkpoint = - objective.purpose == Request::StateBoundary::Purpose::Checkpoint; - const bool junction = - objective.purpose == Request::StateBoundary::Purpose::Junction; - uint64_t &failures = checkpoint ? counters_.checkpointPublicationFailures - : junction ? counters_.junctionMaterializationFailures - : counters_.replayStatePublicationFailures; - uint64_t &publications = checkpoint ? counters_.checkpointPublications - : junction ? counters_.junctionMaterializations - : counters_.replayStatePublications; + using Purpose = Request::StateBoundary::Purpose; + const bool checkpoint = objective.purpose == Purpose::Checkpoint; + const StateRole role = checkpoint ? StateRole::Checkpoint + : objective.purpose == Purpose::Boundary ? StateRole::Boundary + : objective.purpose == Purpose::Hint ? StateRole::Hint + : StateRole::Ordinary; + const auto [publicationCounter, failureCounter] = + [&]() -> std::pair { + switch (objective.purpose) { + case Purpose::Checkpoint: + return {&counters_.checkpointPublications, + &counters_.checkpointPublicationFailures}; + case Purpose::Junction: + return {&counters_.junctionMaterializations, + &counters_.junctionMaterializationFailures}; + case Purpose::Hint: + case Purpose::Boundary: + return {&counters_.hintPublications, &counters_.hintPublicationFailures}; + case Purpose::Replay: + break; + } + return {&counters_.replayStatePublications, + &counters_.replayStatePublicationFailures}; + }(); + uint64_t &publications = *publicationCounter; + uint64_t &failures = *failureCounter; // The scheduler ends a command exactly at an armed boundary; a boundary // passed inside a command has no materialized state to copy. if (objective.tokens != promptProcessed) { @@ -549,7 +597,7 @@ void Engine::publishReachedStateBoundaries(Request &active, } try { const uint64_t block = cache_.blockAt(active.request.id, objective.tokens); - if (cache_.reuseCompositeState(block, checkpoint)) { + if (cache_.reuseCompositeState(block, role)) { ++counters_.deduplicatedStatePublications; } else { // Recycle the previous recovery point before allocating its replacement. @@ -560,7 +608,9 @@ void Engine::publishReachedStateBoundaries(Request &active, } std::shared_ptr state = model_.snapshot(active.request.id); - if (!state && cache_.reclaimOneState(checkpoint)) { + // A speculative hint never recycles a boundary state. + if (!state && + cache_.reclaimOneState(checkpoint, role == StateRole::Hint)) { state = model_.snapshot(active.request.id); if (state) ++counters_.recycledStatePublications; @@ -569,7 +619,7 @@ void Engine::publishReachedStateBoundaries(Request &active, ++failures; continue; } - cache_.publishCompositeState(block, std::move(state), checkpoint); + cache_.publishCompositeState(block, std::move(state), role); ++publications; } if (active.latestCheckpoint.kvBlock != block) diff --git a/runtime/engine/Engine.hpp b/runtime/engine/Engine.hpp index 010bbe0..567200c 100644 --- a/runtime/engine/Engine.hpp +++ b/runtime/engine/Engine.hpp @@ -55,6 +55,8 @@ struct EngineSnapshot final { uint64_t replayStatePublicationFailures = 0; uint64_t junctionMaterializations = 0; uint64_t junctionMaterializationFailures = 0; + uint64_t hintPublications = 0; + uint64_t hintPublicationFailures = 0; uint64_t checkpointPublications = 0; uint64_t checkpointPublicationFailures = 0; uint64_t resourceSuspensions = 0; @@ -114,7 +116,13 @@ class Engine final { struct Request final { struct StateBoundary final { - enum class Purpose : uint8_t { Checkpoint, Replay, Junction }; + // Higher purposes win when two land on the same block. A hint is an + // ordinary state where the server says a message part begins; a + // boundary (the first hint, where the shared setup ends) is also kept + // ahead of other states. + enum class Purpose : uint8_t { Checkpoint, Replay, Junction, Hint, Boundary }; + // Hints closer than this to another planned state are skipped. + static constexpr uint32_t kHintSpacingTokens = 4 * KvCache::pageTokens; uint32_t tokens = 0; Purpose purpose = Purpose::Replay; }; diff --git a/runtime/engine/NativeRuntime.cpp b/runtime/engine/NativeRuntime.cpp index 9474d4b..3ecbe6d 100644 --- a/runtime/engine/NativeRuntime.cpp +++ b/runtime/engine/NativeRuntime.cpp @@ -265,6 +265,7 @@ bool NativeRuntime::handleRequest(protocol::RequestFrame &request) { request.seed}; engineRequest.constraint = mapConstraint(request.constraint); engineRequest.returnProgress = request.returnProgress; + engineRequest.stateHints = request.stateHints; engineRequest.deadlineMilliseconds = nowMonotonic + double(remaining) / 1000.0; core_.submit(std::move(engineRequest)); diff --git a/runtime/engine/Protocol.cpp b/runtime/engine/Protocol.cpp index 94009fc..c625bca 100644 --- a/runtime/engine/Protocol.cpp +++ b/runtime/engine/Protocol.cpp @@ -21,7 +21,7 @@ std::string_view failureClassName(FailureClass failureClass); constexpr std::array kMagic{'S', 'P', 'L', 'H'}; -constexpr uint64_t kRequestFixedBytes = 56; +constexpr uint64_t kRequestFixedBytes = 60; constexpr uint64_t kCancelFixedBytes = 8; constexpr uint64_t kMaskResponseFixedBytes = 20; constexpr uint64_t kStatusRequestFixedBytes = 8; @@ -374,6 +374,18 @@ std::optional validateRequest(const RequestFrame &request, return invalid(IssueCode::LimitExceeded, "prompt token count exceeds its limit"); } + if (request.stateHints.size() > limits.maxStateHints) { + return invalid(IssueCode::LimitExceeded, + "state hint count exceeds its limit"); + } + for (size_t i = 0; i < request.stateHints.size(); ++i) { + if (!request.stateHints[i] || + request.stateHints[i] >= request.promptTokens.size() || + (i && request.stateHints[i] <= request.stateHints[i - 1])) { + return invalid(IssueCode::InvalidPayloadLength, + "state hints must ascend within the prompt"); + } + } const SamplingParameters &sampling = request.sampling; if (!std::isfinite(sampling.temperature) || sampling.temperature < 0.0f || !std::isfinite(sampling.topP) || sampling.topP <= 0.0f || @@ -570,8 +582,8 @@ ProtocolResult encodeRequest(const RequestFrame &request, } uint64_t payloadBytes = 0; uint64_t tokenBytes = 0; - if (!checkedMultiply(request.promptTokens.size(), sizeof(uint32_t), - tokenBytes) || + if (!checkedMultiply(request.promptTokens.size() + request.stateHints.size(), + sizeof(uint32_t), tokenBytes) || !checkedAdd(kRequestFixedBytes, tokenBytes, payloadBytes) || payloadBytes > std::numeric_limits::max()) { return failure( @@ -592,8 +604,11 @@ ProtocolResult encodeRequest(const RequestFrame &request, writer.u32(request.sampling.topK); writer.u64(request.seed); writer.u8(request.returnProgress); + writer.u32(static_cast(request.stateHints.size())); for (uint32_t token : request.promptTokens) writer.u32(token); + for (uint32_t hint : request.stateHints) + writer.u32(hint); return success(Frame{FrameType::Request, writer.take()}); } @@ -774,6 +789,7 @@ ProtocolResult decodeRequest(const Frame &frame, uint8_t constraint = 0; uint8_t returnProgress = 0; uint32_t promptCount = 0; + uint32_t hintCount = 0; if (!reader.u64(request.requestId) || !reader.u8(priority) || !reader.u8(cohort) || !reader.u8(constraint) || !reader.u64(request.absoluteDeadlineUnixMicros) || @@ -782,7 +798,7 @@ ProtocolResult decodeRequest(const Frame &frame, !reader.f32(request.sampling.temperature) || !reader.f32(request.sampling.topP) || !reader.u32(request.sampling.topK) || !reader.u64(request.seed) || - !reader.u8(returnProgress)) { + !reader.u8(returnProgress) || !reader.u32(hintCount)) { return failure(makeIssue(FailureClass::ProtocolFatal, IssueCode::InvalidPayloadLength, 0, "request fixed payload is truncated")); @@ -807,6 +823,17 @@ ProtocolResult decodeRequest(const Frame &frame, request.requestId, "prompt count does not match the binary token payload")); } + if (hintCount > limits.maxStateHints) { + return failure( + makeIssue(FailureClass::RequestError, IssueCode::LimitExceeded, + request.requestId, "state hint count exceeds its limit")); + } + if (!reader.words(hintCount, request.stateHints)) { + return failure( + makeIssue(FailureClass::RequestError, IssueCode::InvalidPayloadLength, + request.requestId, + "state hint count does not match the payload")); + } if (reader.remaining()) { return failure( makeIssue(FailureClass::RequestError, IssueCode::InvalidPayloadLength, diff --git a/runtime/engine/Protocol.hpp b/runtime/engine/Protocol.hpp index ee66e02..4a84032 100644 --- a/runtime/engine/Protocol.hpp +++ b/runtime/engine/Protocol.hpp @@ -13,7 +13,7 @@ namespace splash::protocol { -inline constexpr uint16_t kProtocolVersion = 6; +inline constexpr uint16_t kProtocolVersion = 7; inline constexpr size_t kFrameHeaderBytes = 24; inline constexpr uint32_t kStatusSchemaVersion = 5; // The largest prompt (1 Mi tokens of 4 bytes) fits with room to spare. @@ -102,6 +102,7 @@ struct ProtocolLimits { uint64_t maxStatusJsonBytes = 32ULL * 1024 * 1024; uint64_t maxErrorStringBytes = 1ULL * 1024 * 1024; uint32_t maxPromptTokens = 1U << 20; + uint32_t maxStateHints = 16; uint32_t maxLogicalOutputTokens = 1U << 20; uint32_t maxTokenBatch = 4096; uint32_t maxSimulationTokens = 32; @@ -150,6 +151,8 @@ struct RequestFrame { Cohort cohort = Cohort::Greedy; ConstraintMode constraint = ConstraintMode::None; bool returnProgress = false; + // Ascending prompt positions where messages begin (EngineRequest::stateHints). + std::vector stateHints; bool operator==(const RequestFrame &) const = default; }; diff --git a/runtime/engine/RuntimeResources.mm b/runtime/engine/RuntimeResources.mm index a5336ed..a2fe957 100644 --- a/runtime/engine/RuntimeResources.mm +++ b/runtime/engine/RuntimeResources.mm @@ -2,6 +2,7 @@ #include "metal/abi/ExecutionGeometry.h" #import +#include #include #include @@ -497,8 +498,21 @@ throw RuntimeResourcesError(RuntimeResourceStage::ModelLoading, throw std::runtime_error("model factory returned no state storage"); } auto kvPool = std::make_unique(*kvPages); - auto cache = - std::make_unique(*kvPool, cacheIdentity.cacheNamespace); + // SPLASH_BOUNDARY_STATES sets how many boundary states the cache keeps + // ahead of others; 0 keeps them as ordinary states. A value that is not a + // small non-negative integer keeps the default. + const uint32_t boundaryStates = [] { + const char *value = std::getenv("SPLASH_BOUNDARY_STATES"); + if (!value || *value < '0' || *value > '9') + return engine::Cache::kDefaultBoundaryStates; + char *end = nullptr; + const unsigned long parsed = std::strtoul(value, &end, 10); + if (*end || parsed > 1024) + return engine::Cache::kDefaultBoundaryStates; + return static_cast(parsed); + }(); + auto cache = std::make_unique( + *kvPool, cacheIdentity.cacheNamespace, boundaryStates); if (kvPages->declaredBytes() != budget.kvVirtualBytes || kvPages->actualAllocatedBytes() > budget.kvVirtualBytes) { diff --git a/runtime/engine/StateCache.cpp b/runtime/engine/StateCache.cpp index 621061b..bdd7f66 100644 --- a/runtime/engine/StateCache.cpp +++ b/runtime/engine/StateCache.cpp @@ -75,24 +75,72 @@ std::optional StateCache::acquireBlock(uint64_t kvBlock) { return CompositeStateLease(*this, kvBlock, boundary, entry.state); } -bool StateCache::touchIfResident(uint64_t kvBlock, bool checkpoint) { - auto found = entries_.find(kvBlock); - if (found == entries_.end()) - return false; - if (!kv_.contains(kvBlock)) { - throw std::logic_error("composite state outlived its target KV block"); +StateCache::EvictionQueue &StateCache::queueOf(const Entry &entry) noexcept { + switch (entry.role) { + case StateRole::Checkpoint: + return checkpointEviction_; + case StateRole::Boundary: + return boundaryEviction_; + case StateRole::Ordinary: + case StateRole::Hint: + break; } - Entry &entry = found->second; + return ordinaryEviction_; +} + +void StateCache::setRole(uint64_t kvBlock, StateRole role) noexcept { + Entry &entry = entries_.at(kvBlock); + if (entry.role == role) + return; const bool evictable = entry.evictable; if (evictable) removeEvictable(kvBlock); - if (!checkpoint && entry.checkpoint) { + const uint64_t bytes = entry.state->bytes(); + if (entry.role == StateRole::Checkpoint) { --checkpointEntries_; - checkpointBytes_ -= entry.state->bytes(); - entry.checkpoint = false; + checkpointBytes_ -= bytes; + } else if (entry.role == StateRole::Boundary) { + --boundaryEntries_; + } + entry.role = role; + if (role == StateRole::Checkpoint) { + ++checkpointEntries_; + checkpointBytes_ += bytes; + } else if (role == StateRole::Boundary) { + ++boundaryEntries_; } if (evictable) insertEvictable(kvBlock); +} + +bool StateCache::reserveBoundary() noexcept { + if (!boundaryCapacity_) + return false; + if (boundaryEntries_ < boundaryCapacity_) + return true; + const uint64_t oldest = boundaryEviction_.oldest; + if (!oldest) + return false; + setRole(oldest, StateRole::Ordinary); + ++boundaryDemotions_; + return true; +} + +bool StateCache::touchIfResident(uint64_t kvBlock, StateRole role) { + auto found = entries_.find(kvBlock); + if (found == entries_.end()) + return false; + if (!kv_.contains(kvBlock)) { + throw std::logic_error("composite state outlived its target KV block"); + } + const StateRole current = found->second.role; + if (role == StateRole::Boundary && current != StateRole::Boundary && + reserveBoundary()) + setRole(kvBlock, StateRole::Boundary); + else if (current == StateRole::Checkpoint && role != StateRole::Checkpoint) + setRole(kvBlock, StateRole::Ordinary); + else + touch(kvBlock); ++deduplicatedPublications_; return true; } @@ -107,7 +155,7 @@ void StateCache::touch(uint64_t kvBlock) noexcept { void StateCache::publish(uint64_t kvBlock, std::shared_ptr state, - bool checkpoint) { + StateRole role) { if (!state || !state->bytes()) { throw std::invalid_argument("composite state payload is empty"); } @@ -124,29 +172,33 @@ void StateCache::publish(uint64_t kvBlock, Entry entry; entry.publication = publications_ + 1; entry.state = std::move(state); - entry.checkpoint = checkpoint; - auto [_, inserted] = entries_.emplace(kvBlock, std::move(entry)); - if (!inserted) + if (entries_.contains(kvBlock)) throw std::logic_error("duplicate composite state key"); + if (role == StateRole::Boundary && !reserveBoundary()) + role = StateRole::Ordinary; + entry.role = role; + entries_.emplace(kvBlock, std::move(entry)); bytes_ += stateBytes; - if (checkpoint) { + if (role == StateRole::Checkpoint) { ++checkpointEntries_; checkpointBytes_ += stateBytes; + } else if (role == StateRole::Boundary) { + ++boundaryEntries_; } - insertEvictable(kvBlock); + insertEvictable(kvBlock, role == StateRole::Hint); ++publications_; } StateCheckpoint StateCache::checkpoint(uint64_t kvBlock) const noexcept { const auto found = entries_.find(kvBlock); - if (found == entries_.end() || !found->second.checkpoint) + if (found == entries_.end() || found->second.role != StateRole::Checkpoint) return {}; return {kvBlock, found->second.publication}; } bool StateCache::retireCheckpoint(StateCheckpoint checkpoint) noexcept { const auto found = entries_.find(checkpoint.kvBlock); - if (found == entries_.end() || !found->second.checkpoint || + if (found == entries_.end() || found->second.role != StateRole::Checkpoint || found->second.publication != checkpoint.publication) return true; return erase(checkpoint.kvBlock, true).evicted; @@ -158,12 +210,14 @@ bool StateCache::contains(uint64_t kvBlock) const noexcept { std::optional StateCache::evictionCandidate() const noexcept { - const uint64_t oldest = checkpointEviction_.oldest - ? checkpointEviction_.oldest - : ordinaryEviction_.oldest; + const uint64_t oldest = checkpointEviction_.oldest ? checkpointEviction_.oldest + : ordinaryEviction_.oldest ? ordinaryEviction_.oldest + : boundaryEviction_.oldest; if (!oldest) return std::nullopt; - return CacheEvictionCandidate{oldest, entries_.at(oldest).lastUsed}; + const Entry &entry = entries_.at(oldest); + return CacheEvictionCandidate{oldest, entry.lastUsed, + entry.role == StateRole::Boundary}; } StateEviction StateCache::evict(uint64_t kvBlock) noexcept { @@ -176,11 +230,13 @@ StateEviction StateCache::erase(uint64_t kvBlock, bool retirement) noexcept { return {}; const uint64_t reclaimed = found->second.state->bytes(); removeEvictable(kvBlock); - if (found->second.checkpoint) { + if (found->second.role == StateRole::Checkpoint) { --checkpointEntries_; checkpointBytes_ -= reclaimed; if (!retirement) ++checkpointEvictions_; + } else if (found->second.role == StateRole::Boundary) { + --boundaryEntries_; } bytes_ -= reclaimed; entries_.erase(found); @@ -207,6 +263,9 @@ StateCacheSnapshot StateCache::snapshot() const noexcept { result.checkpointBytes = checkpointBytes_; result.checkpointRetirements = checkpointRetirements_; result.checkpointEvictions = checkpointEvictions_; + result.boundaryEntries = boundaryEntries_; + result.boundaryCapacity = boundaryCapacity_; + result.boundaryDemotions = boundaryDemotions_; return result; } @@ -224,16 +283,24 @@ void StateCache::release(uint64_t kvBlock) noexcept { kv_.releaseActive(kvBlock); } -void StateCache::insertEvictable(uint64_t kvBlock) noexcept { +void StateCache::insertEvictable(uint64_t kvBlock, bool coldest) noexcept { Entry &entry = entries_.at(kvBlock); if (entry.evictable || entry.pins) std::terminate(); - EvictionQueue &queue = - entry.checkpoint ? checkpointEviction_ : ordinaryEviction_; + EvictionQueue &queue = queueOf(entry); + entry.evictable = true; + if (coldest && queue.oldest) { + Entry &oldest = entries_.at(queue.oldest); + entry.lastUsed = oldest.lastUsed; + entry.previousEvictable = 0; + entry.nextEvictable = queue.oldest; + oldest.previousEvictable = kvBlock; + queue.oldest = kvBlock; + return; + } entry.previousEvictable = queue.newest; entry.nextEvictable = 0; entry.lastUsed = recency_.next(); - entry.evictable = true; if (queue.newest) entries_.at(queue.newest).nextEvictable = kvBlock; else @@ -245,8 +312,7 @@ void StateCache::removeEvictable(uint64_t kvBlock) noexcept { Entry &entry = entries_.at(kvBlock); if (!entry.evictable) std::terminate(); - EvictionQueue &queue = - entry.checkpoint ? checkpointEviction_ : ordinaryEviction_; + EvictionQueue &queue = queueOf(entry); if (entry.previousEvictable) entries_.at(entry.previousEvictable).nextEvictable = entry.nextEvictable; else diff --git a/runtime/engine/StateCache.hpp b/runtime/engine/StateCache.hpp index d26a9ec..df025ec 100644 --- a/runtime/engine/StateCache.hpp +++ b/runtime/engine/StateCache.hpp @@ -64,8 +64,18 @@ struct StateCacheSnapshot { uint64_t checkpointRetirements = 0; // Pressure and logical eviction; rolling retirements are counted separately. uint64_t checkpointEvictions = 0; + // States at prompt message boundaries, kept ahead of all others. + uint32_t boundaryEntries = 0; + uint32_t boundaryCapacity = 0; + uint64_t boundaryDemotions = 0; }; +// What a published state is for, which decides when it is evicted: +// checkpoints first, then ordinary and hint states by recency (a hint, a +// state the server guessed later prompts may resume from, starts as the +// oldest until one does), then boundary states. +enum class StateRole : uint8_t { Ordinary, Checkpoint, Hint, Boundary }; + struct StateEviction final { bool evicted = false; uint64_t reclaimedBytes = 0; @@ -75,7 +85,10 @@ struct StateEviction final { // target-KV block. The pair is restored atomically. class StateCache final { public: - StateCache(KvCache &kv, CacheRecency &recency) : kv_(kv), recency_(recency) {} + // At most boundaryCapacity boundary states are kept as such; publishing one + // more demotes the least recently used to an ordinary state. + StateCache(KvCache &kv, CacheRecency &recency, uint32_t boundaryCapacity) + : kv_(kv), recency_(recency), boundaryCapacity_(boundaryCapacity) {} StateCache(const StateCache &) = delete; StateCache &operator=(const StateCache &) = delete; @@ -85,12 +98,13 @@ class StateCache final { void recordLookup(bool hit) noexcept { hit ? ++hits_ : ++misses_; } // Reuses a published state without a restore pin or lookup accounting. - // A normal boundary upgrades a checkpoint; a checkpoint cannot downgrade - // an ordinary state. Restoration alone never changes the boundary's role. - [[nodiscard]] bool touchIfResident(uint64_t kvBlock, bool checkpoint = false); + // Any other role upgrades a checkpoint, and a boundary role upgrades any + // state; nothing downgrades. Restoration alone never changes a role. + [[nodiscard]] bool touchIfResident(uint64_t kvBlock, + StateRole role = StateRole::Ordinary); void publish(uint64_t kvBlock, std::shared_ptr state, - bool checkpoint = false); + StateRole role = StateRole::Ordinary); // Publication identity protects replacement states from stale handles. [[nodiscard]] StateCheckpoint checkpoint(uint64_t kvBlock) const noexcept; // Ensures this publication is no longer a disposable checkpoint. Returns @@ -101,7 +115,9 @@ class StateCache final { void touch(uint64_t kvBlock) noexcept; [[nodiscard]] bool contains(uint64_t kvBlock) const noexcept; - // Unpinned checkpoints precede ordinary states regardless of recency. + // Unpinned checkpoints precede ordinary and hint states, which precede + // boundary states, regardless of recency. A boundary candidate is marked + // protectedState. [[nodiscard]] std::optional evictionCandidate() const noexcept; [[nodiscard]] StateEviction evict(uint64_t kvBlock) noexcept; @@ -117,7 +133,7 @@ class StateCache final { uint64_t previousEvictable = 0; uint64_t nextEvictable = 0; bool evictable = false; - bool checkpoint = false; + StateRole role = StateRole::Ordinary; uint64_t publication = 0; }; @@ -130,8 +146,14 @@ class StateCache final { void release(uint64_t kvBlock) noexcept; [[nodiscard]] std::optional acquireBlock(uint64_t kvBlock); - void insertEvictable(uint64_t kvBlock) noexcept; + // Coldest inserts at the oldest end of the queue, as old as its oldest. + void insertEvictable(uint64_t kvBlock, bool coldest = false) noexcept; void removeEvictable(uint64_t kvBlock) noexcept; + [[nodiscard]] EvictionQueue &queueOf(const Entry &entry) noexcept; + // Changes an entry's role, keeping the role counts and queues in step. + void setRole(uint64_t kvBlock, StateRole role) noexcept; + // Makes room for one more boundary state; false when none can be demoted. + [[nodiscard]] bool reserveBoundary() noexcept; KvCache &kv_; CacheRecency &recency_; @@ -139,6 +161,7 @@ class StateCache final { uint64_t bytes_ = 0; EvictionQueue ordinaryEviction_; EvictionQueue checkpointEviction_; + EvictionQueue boundaryEviction_; uint32_t pinnedEntries_ = 0; uint64_t hits_ = 0; uint64_t misses_ = 0; @@ -149,6 +172,9 @@ class StateCache final { uint64_t checkpointBytes_ = 0; uint64_t checkpointRetirements_ = 0; uint64_t checkpointEvictions_ = 0; + uint32_t boundaryEntries_ = 0; + uint32_t boundaryCapacity_ = 0; + uint64_t boundaryDemotions_ = 0; }; } // namespace splash::engine diff --git a/runtime/engine/Status.cpp b/runtime/engine/Status.cpp index bb99632..bba7e54 100644 --- a/runtime/engine/Status.cpp +++ b/runtime/engine/Status.cpp @@ -185,7 +185,10 @@ std::string runtimeStatusJson( << ",\"checkpoint_entries\":" << state.checkpointEntries << ",\"checkpoint_bytes\":" << state.checkpointBytes << ",\"checkpoint_evictions\":" << state.checkpointEvictions - << ",\"checkpoint_retirements\":" << state.checkpointRetirements << "}" + << ",\"checkpoint_retirements\":" << state.checkpointRetirements + << ",\"boundary_entries\":" << state.boundaryEntries + << ",\"boundary_capacity\":" << state.boundaryCapacity + << ",\"boundary_demotions\":" << state.boundaryDemotions << "}" << ",\"cache\":{\"lookups\":" << lookup.lookups << ",\"hits\":" << core.cacheHits << ",\"cold_misses\":" << core.coldMisses << ",\"hit_rate\":" << hitRate @@ -206,6 +209,8 @@ std::string runtimeStatusJson( << ",\"checkpoint_publications\":" << core.checkpointPublications << ",\"checkpoint_publication_failures\":" << core.checkpointPublicationFailures + << ",\"hint_publications\":" << core.hintPublications + << ",\"hint_publication_failures\":" << core.hintPublicationFailures << ",\"resource_suspensions\":" << core.resourceSuspensions << ",\"resource_resumptions\":" << core.resourceResumptions << ",\"resource_replay_tokens\":" << core.resourceReplayTokens << "}" diff --git a/runtime/engine/Types.hpp b/runtime/engine/Types.hpp index a14beef..0b79375 100644 --- a/runtime/engine/Types.hpp +++ b/runtime/engine/Types.hpp @@ -35,6 +35,10 @@ struct EngineRequest final { ConstraintMode constraint = ConstraintMode::None; double deadlineMilliseconds = 0.0; bool returnProgress = false; + // Ascending prompt positions where its messages begin, from the server. + // The state there is kept so later prompts that share the messages before + // it, and differ after, resume instead of prefilling from the start. + std::vector stateHints; [[nodiscard]] ModelRequest modelView() const noexcept { return {id, cohort, prompt, maxNewTokens, sampling, constraint}; diff --git a/server/backend.py b/server/backend.py index cef264e..9df1e32 100644 --- a/server/backend.py +++ b/server/backend.py @@ -109,6 +109,8 @@ class Job: response_previous_id: str | None = None response_history_items: list | None = None return_progress: bool = False + # Prompt positions where the engine keeps the state (Frontend.state_hints). + state_hints: tuple = () class CallbackStreamer: @@ -446,6 +448,7 @@ def _generation_request(self, job): constraint = wire.ConstraintMode.NONE return engine_runtime.GenerationRequest( prompt_tokens=tuple(job.prompt_tokens), + state_hints=tuple(job.state_hints), logical_max_output_tokens=job.max_new_tokens, deadline=self._deadline(job), priority=priority, diff --git a/server/frontend.py b/server/frontend.py index f2f3c75..1065bbd 100644 --- a/server/frontend.py +++ b/server/frontend.py @@ -6,6 +6,7 @@ import secrets import threading import time +from bisect import bisect_left from collections import OrderedDict from contextlib import contextmanager from dataclasses import dataclass @@ -66,12 +67,13 @@ RESPONSE_STORE_BUDGET_BYTES = 64 * 1024 * 1024 +# Each chat message begins with this special token in the model's template. +MESSAGE_START = "<|im_start|>" def _thinking_from_prefix(rendered): - marker = "<|im_start|>" - start = rendered.rfind(marker) - prefix = rendered[start + len(marker) :] if start >= 0 else "" + start = rendered.rfind(MESSAGE_START) + prefix = rendered[start + len(MESSAGE_START) :] if start >= 0 else "" if not prefix.startswith("assistant\n") or "<|im_end|>" in prefix: raise APIError( 400, "chat template must end with an assistant generation prefix" @@ -160,6 +162,24 @@ def stats(self): } +def _special_token_id(tokenizer, token): + """The id of a special token, or None when the tokenizer has none.""" + convert = getattr(tokenizer, "convert_tokens_to_ids", None) + if convert is None: + return None + token_id = convert(token) + if token_id is None or token_id == getattr(tokenizer, "unk_token_id", None): + return None + return token_id + + +def _last_request_message(body): + """The last message as the client sent it: the frontend joins content + parts into one string before rendering, and the hints need the parts.""" + messages = body.get("messages") if isinstance(body, dict) else None + return messages[-1] if isinstance(messages, list) and messages else None + + @dataclass class Prompt: messages: list @@ -194,6 +214,7 @@ def __init__( if not isinstance(preparation_capacity, int) or preparation_capacity <= 0: raise ValueError("frontend preparation capacity must be positive") self.tokenizer = tokenizer + self.message_start_token = _special_token_id(tokenizer, MESSAGE_START) self.backend = backend self.model = model self.max_context = max_context @@ -242,6 +263,75 @@ def prepare( with self._preparation(deadline): return self._prepare(body, tool_namespaces, deadline, clamp_output_budget) + # Content parts of the last message to hint, counted from its end. + HINTED_LAST_PARTS = 2 + + def state_hints(self, tokens, text=None, last_message=None): + """Where the engine should keep the state, so a later prompt that + shares everything before a point and differs after it resumes there. + The first hint ends the system prompt and tools, which a new + conversation or a side agent on the same setup shares; the engine + keeps it ahead of other states. The others mark where the last + message and its final content parts begin: a new question in the same + conversation, often after reminders the client puts first. The end of + each prompt is kept anyway.""" + if self.message_start_token is None: + return () + starts = [ + i for i, token in enumerate(tokens) if token == self.message_start_token + ] + if len(starts) < 2: + return () + hints = {starts[1]} + # The final start opens the reply being generated. + if len(starts) >= 3: + hints.add(starts[-2]) + hints.update(self._part_hints(starts[-2], text, last_message)) + return tuple(sorted(hints)) + + def _part_hints(self, message_token, text, message): + """Token positions where the last message's final content parts begin. + Only that message is tokenized again: it opens with a special token, + so its tokens line up with the whole prompt's.""" + parts = message.get("content") if isinstance(message, dict) else None + if not isinstance(parts, list) or len(parts) < 2 or not isinstance(text, str): + return () + # The last message opens at the second-to-last start; the last opens + # the reply. + opening = text.rfind(MESSAGE_START, 0, text.rfind(MESSAGE_START)) + if opening < 0: + return () + segment = text[opening:] + chars, at = [], 0 + # The first part begins with the message, which is hinted already. + for index, part in enumerate(parts): + needle = part.get("text") if isinstance(part, dict) else None + if not isinstance(needle, str) or not needle.strip(): + continue + found = segment.find(needle[:200].strip(), at) + if found < 0: + continue + if index: + chars.append(found) + at = found + 1 + chars = chars[-self.HINTED_LAST_PARTS :] + if not chars: + return () + try: + offsets = self.tokenizer( + segment, add_special_tokens=False, return_offsets_mapping=True + )["offset_mapping"] + except (TypeError, KeyError, ValueError, NotImplementedError): + # Python (slow) tokenizers raise NotImplementedError for offsets. + return () + begins = [begin for begin, _ in offsets] + hints = [] + for char in chars: + index = bisect_left(begins, char) + if 0 < index < len(begins): + hints.append(message_token + index) + return hints + def count_tokens(self, body, *, deadline=None): if deadline is None: deadline = self.request_deadline(body) @@ -583,6 +673,9 @@ def _prepare(self, body, tool_namespaces, deadline, clamp_output_budget=False): job = Job( request_id=request_id, prompt_tokens=prompt_tokens, + state_hints=self.state_hints( + prompt_tokens, rendered.text, _last_request_message(body) + ), max_new_tokens=max_new, seed=seed, temperature=temperature, diff --git a/server/protocol.py b/server/protocol.py index dc3da1b..8ab5ecf 100644 --- a/server/protocol.py +++ b/server/protocol.py @@ -13,7 +13,7 @@ from enum import IntEnum, IntFlag from typing import TypeAlias -PROTOCOL_VERSION = 6 +PROTOCOL_VERSION = 7 FRAME_HEADER_BYTES = 24 STATUS_SCHEMA_VERSION = 5 # Largest top-k the native sampler keeps as candidates. @@ -24,7 +24,7 @@ _HEADER = struct.Struct("<4sHHHHQI") # Replay can update the integer deadlines without decoding sampling floats. _REQUEST_HEAD = struct.Struct(" limits.max_prompt_tokens: raise ValueError("prompt token count exceeds its limit") + hints = _words(request.state_hints, "state hints") + if len(hints) > limits.max_state_hints: + raise ValueError("state hint count exceeds its limit") except ValueError as error: return _issue( FailureClass.REQUEST_ERROR, @@ -595,6 +602,16 @@ def _request_issue( str(error), request_id, ) + if any( + not hint or hint >= len(prompt) or (i and hint <= hints[i - 1]) + for i, hint in enumerate(hints) + ): + return _issue( + FailureClass.REQUEST_ERROR, + IssueCode.INVALID_PAYLOAD_LENGTH, + "state hints must ascend within the prompt", + request_id, + ) try: temperature = _float32(request.sampling.temperature, "temperature") top_p = _float32(request.sampling.top_p, "top_p") @@ -952,23 +969,29 @@ def _encode_message( if isinstance(message, RequestFrame): _raise_issue(_request_issue(message, limits)) prompt = _words(message.prompt_tokens, "prompt tokens") + hints = _words(message.state_hints, "state hints") temperature = _float32(message.sampling.temperature, "temperature") top_p = _float32(message.sampling.top_p, "top_p") - payload = _REQUEST.pack( - message.request_id, - int(message.priority), - int(message.cohort), - int(message.constraint), - message.absolute_deadline_unix_micros, - message.remaining_deadline_micros, - message.logical_max_output_tokens, - len(prompt), - temperature, - top_p, - message.sampling.top_k, - message.seed, - message.return_progress, - ) + _pack_words(prompt) + payload = ( + _REQUEST.pack( + message.request_id, + int(message.priority), + int(message.cohort), + int(message.constraint), + message.absolute_deadline_unix_micros, + message.remaining_deadline_micros, + message.logical_max_output_tokens, + len(prompt), + temperature, + top_p, + message.sampling.top_k, + message.seed, + message.return_progress, + len(hints), + ) + + _pack_words(prompt) + + _pack_words(hints) + ) frame_type = FrameType.REQUEST elif isinstance(message, CancelFrame): _raise_issue(_cancel_issue(message)) @@ -1229,6 +1252,7 @@ def _decode_request(payload: bytes, limits: ProtocolLimits) -> RequestFrame: top_k, seed, return_progress, + hint_count, ) = _REQUEST.unpack_from(payload) if return_progress > 1: _fail( @@ -1244,9 +1268,18 @@ def _decode_request(payload: bytes, limits: ProtocolLimits) -> RequestFrame: "prompt token count exceeds its limit", request_id, ) + if hint_count > limits.max_state_hints: + _fail( + FailureClass.REQUEST_ERROR, + IssueCode.LIMIT_EXCEEDED, + "state hint count exceeds its limit", + request_id, + ) try: prompt = _unpack_prefix_words(payload, _REQUEST.size, prompt_count) - if len(payload) != _REQUEST.size + 4 * prompt_count: + hints_at = _REQUEST.size + 4 * prompt_count + hints = _unpack_prefix_words(payload, hints_at, hint_count) + if len(payload) != hints_at + 4 * hint_count: raise ValueError("trailing request bytes") except (ValueError, struct.error): _fail( @@ -1285,6 +1318,7 @@ def _decode_request(payload: bytes, limits: ProtocolLimits) -> RequestFrame: request_id, ), bool(return_progress), + hints, ) _raise_issue(_request_issue(request, limits)) return request diff --git a/server/runtime.py b/server/runtime.py index 6ed692c..4245e9f 100644 --- a/server/runtime.py +++ b/server/runtime.py @@ -165,6 +165,7 @@ class GenerationRequest: constraint: wire.ConstraintMode = wire.ConstraintMode.NONE mask_provider: MaskProvider | None = None return_progress: bool = False + state_hints: tuple[int, ...] = () @dataclass(slots=True, frozen=True) @@ -585,6 +586,7 @@ def submit( cohort=request.cohort, constraint=request.constraint, return_progress=request.return_progress, + state_hints=request.state_hints, ) try: encoded = wire.serialize_message(protocol_request)