popov/neuron-tai
1
0
Fork 0
Code Issues 42 Pull Requests Actions Packages Projects Releases Wiki Activity

story: DGR-038 Implement isolated shard-local Hot KV State

Browse Source
This commit is contained in:
Dobromir Popov
2026-08-01 01:32:46 +03:00
parent a1df87deb6
commit 49560b396f
6 changed files with 318 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files

31
packages/node/native/worker/llama_shard_engine.h
View File

@@ -25,6 +25,10 @@ struct WorkerIdentity {
uint32_t start_layer = 0;
uint32_t end_layer = 0; // half-open, as on the wire
uint32_t injected_death_after_executions = 0; // opt-in test hook; zero disables
uint32_t hot_kv_max_sessions = 8;
uint32_t hot_kv_context_tokens = 4096;
uint32_t hot_kv_budget_tokens = 32768;
uint32_t hot_kv_ttl_seconds = 300;
};
struct BundleCheck {
@@ -38,15 +42,38 @@ struct EngineHealth {
std::string detail;
};
// This is deliberately expressed in tokens, rather than guessed bytes: llama.cpp
// owns the actual K/V layout for the loaded range and backend. The worker uses
// the token reservation to keep its local KV arena bounded before a graph
// adapter materializes the typed boundary (DGR-039).
struct HotKvStep {
enum class Phase { kPrefill, kDecode };
std::string route_session_id;
uint64_t route_epoch = 0;
Phase phase = Phase::kPrefill;
uint64_t first_position = 0;
uint32_t token_count = 0;
uint64_t expected_past_len = 0;
};
enum class HotKvStatus { kOk, kCacheMiss, kStaleEpoch, kResourceExhausted, kCancelled };
struct HotKvResult {
HotKvStatus status = HotKvStatus::kOk;
uint64_t past_len = 0;
std::string detail;
};
class ShardEngine {
public:
virtual ~ShardEngine() = default;
virtual bool Load(std::string* error) = 0;
virtual BundleCheck Validate(const sp::TensorBundle&, uint64_t max_chunk_bytes) const = 0;
virtual bool Execute(const sp::TensorBundle&, std::string* error) = 0;
virtual HotKvResult OpenSession(const std::string& route_session_id, uint64_t route_epoch) = 0;
virtual HotKvResult Execute(const HotKvStep&, const sp::TensorBundle&, std::string* error) = 0;
virtual const WorkerIdentity& identity() const = 0;
virtual EngineHealth health() const = 0;
virtual void ReleaseSession(const std::string& route_session_id) = 0;
virtual void ReleaseSession(const std::string& route_session_id, uint64_t route_epoch) = 0;
virtual void Shutdown() = 0;
};