Skip to content

feat(cache): keep recurrent states where prompts diverge #2

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

Merged
mzinner merged 1 commit into from
Oct 23, 2026
Merged

feat(cache): keep recurrent states where prompts diverge #2

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
112 changes: 96 additions & 16 deletions dev/tests/engine/kv_first_cache_test.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
#include "TestKvPool.hpp"
#include "engine/Cache.hpp"

#include <cstdlib>
#include <array>
#include <cstdint>
#include <iostream>
Expand Down Expand Up @@ -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<uint32_t> prompt;
std::vector<uint64_t> 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);
Expand All @@ -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<TestState>(bytes));
std::make_shared<TestState>(bytes), role);
}
void publishBoundary(uint32_t block) {
publish(block, 100, engine::StateRole::Boundary);
}

engine::CacheLookup lookup(uint32_t tokens) {
Expand Down Expand Up @@ -342,7 +349,7 @@ void testCheckpointLookupProbeDoesNotPromote() {
CacheFixture fixture;
const uint64_t block = fixture.blocks[0];
fixture.cache.publishCompositeState(block, std::make_shared<TestState>(100),
true);
engine::StateRole::Checkpoint);
const StateCheckpoint checkpoint = fixture.cache.checkpointState(block);
{
auto probe = fixture.lookup(33);
Expand All @@ -360,11 +367,11 @@ void testCheckpointRetirementRespectsUseAndPublicationIdentity() {
const uint64_t block = fixture.blocks[1];
const auto publish = [&] {
fixture.cache.publishCompositeState(block, std::make_shared<TestState>(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,
Expand Down Expand Up @@ -408,9 +415,9 @@ void testRestoredCheckpointsKeepTheirEvictionPriority() {
fixture.publish(0, 150);
fixture.publish(3, 400);
fixture.cache.publishCompositeState(fixture.blocks[1],
std::make_shared<TestState>(200), true);
std::make_shared<TestState>(200), engine::StateRole::Checkpoint);
fixture.cache.publishCompositeState(fixture.blocks[2],
std::make_shared<TestState>(300), true);
std::make_shared<TestState>(300), engine::StateRole::Checkpoint);
{
auto lookup = fixture.lookup(65);
fixture.cache.beginRequest(2);
Expand All @@ -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<TestState>(200), true);
std::make_shared<TestState>(200), engine::StateRole::Checkpoint);
const auto reclaimed = fixture.cache.reclaimOne();
require(reclaimed.madeProgress && reclaimed.reclaimedBytes == 200 &&
fixture.cache.snapshot().kvCache.blocks == 4 &&
Expand All @@ -453,7 +530,7 @@ void testOptionalReclaimLeavesOrdinaryStateIntact() {
CacheFixture fixture;
fixture.publish(0, 150);
fixture.cache.publishCompositeState(fixture.blocks[2],
std::make_shared<TestState>(200), true);
std::make_shared<TestState>(200), engine::StateRole::Checkpoint);
require(fixture.cache.reclaimOneState(true) &&
fixture.cache.snapshot().stateCache.checkpointEntries == 0 &&
fixture.lookup(33).resumeBoundary() == 32,
Expand All @@ -468,9 +545,9 @@ void testCheckpointPinsAndBoundaryUpgrade() {
fixture.publish(0, 100);
fixture.publish(3, 400);
fixture.cache.publishCompositeState(fixture.blocks[1],
std::make_shared<TestState>(200), true);
std::make_shared<TestState>(200), engine::StateRole::Checkpoint);
fixture.cache.publishCompositeState(fixture.blocks[2],
std::make_shared<TestState>(300), true);
std::make_shared<TestState>(300), engine::StateRole::Checkpoint);
const auto checkpoint = fixture.cache.checkpointState(fixture.blocks[1]);
auto first = fixture.lookup(65);
auto second = fixture.lookup(97);
Expand All @@ -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");
}
Expand Down Expand Up @@ -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<BudgetState>(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);
Expand All @@ -540,7 +618,7 @@ void testCheckpointPressurePreservesHotPrefix() {
void testLogicalKvPressureStillReclaimsPages() {
CacheFixture fixture;
fixture.cache.publishCompositeState(fixture.blocks[0],
std::make_shared<TestState>(200), true);
std::make_shared<TestState>(200), engine::StateRole::Checkpoint);
fixture.cache.beginRequest(2);
require(fixture.cache.ensureTokens(2, 32).granted() &&
fixture.cache.snapshot().kvCache.blocks == 3 &&
Expand Down Expand Up @@ -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) {
Expand Down
24 changes: 19 additions & 5 deletions dev/tests/engine/protocol_test.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -123,6 +123,7 @@ RequestFrame exampleRequest() {
request.seed = 0xfedcba9876543210ULL;
request.cohort = Cohort::Constrained;
request.constraint = ConstraintMode::TokenMask;
request.stateHints = {1, 3};
return request;
}

Expand All @@ -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<uint16_t>(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<uint32_t> &hints :
{std::vector<uint32_t>{3, 1}, std::vector<uint32_t>{2, 2},
std::vector<uint32_t>{0}, std::vector<uint32_t>{5}}) {
RequestFrame invalid = exampleRequest();
invalid.stateHints = hints;
CHECK(test, !serializeMessage(Message{invalid}));
}

}

std::vector<Message> everyOtherMessage() {
Expand Down
28 changes: 20 additions & 8 deletions dev/tests/engine/test_protocol_python.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"
)

Expand Down Expand Up @@ -340,14 +340,26 @@ def test_request_header_and_binary_prompt(self):
self.assertEqual(wire[:4], b"SPLH")
self.assertEqual(
struct.unpack_from("<HHHHQI", wire, 4),
(p.PROTOCOL_VERSION, 24, int(p.FrameType.REQUEST), 0, 76, 0),
(p.PROTOCOL_VERSION, 24, int(p.FrameType.REQUEST), 0, 80, 0),
)
self.assertEqual(struct.unpack_from("<Q", wire, 24)[0], request.request_id)
self.assertEqual(struct.unpack_from("<I", wire, 24 + 31)[0], 5)
self.assertEqual(struct.unpack_from("<I", wire, 24 + 56)[0], 0)
self.assertEqual(
struct.unpack_from("<5I", wire, 24 + 56), request.prompt_tokens
struct.unpack_from("<5I", wire, 24 + 60), request.prompt_tokens
)

def test_request_state_hints_follow_the_prompt(self):
request = replace(example_request(), state_hints=(1, 3))
wire = p.serialize_message(request)
self.assertEqual(struct.unpack_from("<I", wire, 24 + 56)[0], 2)
self.assertEqual(struct.unpack_from("<2I", wire, 24 + 60 + 20), (1, 3))
frame = p.encode_message(request)
self.assertEqual(p.decode_frame(frame), request)
for hints in ((3, 1), (2, 2), (0,), (5,)):
with self.subTest(hints=hints), self.assertRaises(p.ProtocolError):
p.serialize_message(replace(request, state_hints=hints))

def test_refresh_request_deadline_changes_only_the_absolute_deadline(self):
request = example_request()
wire = p.serialize_message(request)
Expand Down
Loading