branch.hpp

Go to the documentation of this file.

#pragma once
 
// SPDX-License-Identifier: LicenseRef-FSL-1.1-Apache-2.0
// Copyright 2026 Lloyal Labs
 
 
#include "boundaries.hpp"
#include "common.hpp"
#include "decode.hpp"
#include "grammar.hpp"
#include "kv.hpp"
#include "logits.hpp"
#include "metrics.hpp"
#include "sampler.hpp"
 
#include <llama/llama.h>
#include <algorithm>  // std::remove
#include <cassert>    // assert
#include <cmath>      // std::exp, std::log, std::isinf, std::isfinite
#include <cstdint>
#include <cstring>    // std::memcpy
#include <ctime>      // std::time
#include <deque>      // std::deque (pointer stability for BranchStore)
#include <functional> // std::function
#include <limits>     // std::numeric_limits
#include <mutex>
#include <span>       // std::span (C++20)
#include <stdexcept>  // std::runtime_error
#include <string>     // std::to_string
#include <unordered_map>
#include <utility>    // std::pair, std::exchange
#include <vector>
 
namespace lloyal::branch {
 
// ===== HANDLE TYPE =====
 
using BranchHandle = uint32_t;
 
constexpr BranchHandle INVALID_HANDLE = 0;  
constexpr llama_seq_id NO_LEASE = kv::NO_LEASE;  
constexpr int DEFAULT_N_BATCH = 512;        
constexpr uint32_t GEN_SHIFT = 16;          
constexpr uint32_t INDEX_MASK = 0xFFFF;     
 
struct KvPressure {
  uint32_t n_ctx;       
  uint32_t cells_used;  
  uint32_t remaining;   
};
 
inline uint16_t handle_index(BranchHandle h) {
  return static_cast<uint16_t>(h & INDEX_MASK);
}
 
inline uint16_t handle_generation(BranchHandle h) {
  return static_cast<uint16_t>(h >> GEN_SHIFT);
}
 
inline BranchHandle make_handle(uint16_t index, uint16_t generation) {
  return (static_cast<uint32_t>(generation) << GEN_SHIFT) | index;
}
 
// ===== REGISTRY HANDLE TYPES =====
 
using SamplerChainHandle = int32_t;
 
using GrammarHandle = int32_t;
 
using MetricsHandle = int32_t;
 
struct SamplerChainEntry {
  llama_sampler* chain = nullptr;
  bool has_dist = false;   
 
  SamplerChainEntry() = default;
  ~SamplerChainEntry() { if (chain) sampler::free_chain(chain); }
 
  SamplerChainEntry(SamplerChainEntry&& o) noexcept
      : chain(o.chain), has_dist(o.has_dist) { o.chain = nullptr; }
  SamplerChainEntry& operator=(SamplerChainEntry&& o) noexcept {
    if (this != &o) {
      if (chain) sampler::free_chain(chain);
      chain = o.chain; has_dist = o.has_dist; o.chain = nullptr;
    }
    return *this;
  }
  SamplerChainEntry(const SamplerChainEntry&) = delete;
  SamplerChainEntry& operator=(const SamplerChainEntry&) = delete;
};
 
struct GrammarEntry {
  llama_sampler* sampler = nullptr;
 
  GrammarEntry() = default;
  ~GrammarEntry() { if (sampler) grammar::free_sampler(sampler); }
 
  GrammarEntry(GrammarEntry&& o) noexcept : sampler(o.sampler) { o.sampler = nullptr; }
  GrammarEntry& operator=(GrammarEntry&& o) noexcept {
    if (this != &o) {
      if (sampler) grammar::free_sampler(sampler);
      sampler = o.sampler; o.sampler = nullptr;
    }
    return *this;
  }
  GrammarEntry(const GrammarEntry&) = delete;
  GrammarEntry& operator=(const GrammarEntry&) = delete;
};
 
struct CachedSamplingParams {
  float temperature = 0.8f;
  int32_t top_k = 40;
  float top_p = 0.95f;
  float typical_p = 1.0f;
  float min_p = 0.05f;
  float penalty_repeat = 1.0f;
  float penalty_freq = 0.0f;
  float penalty_present = 0.0f;
  int32_t penalty_last_n = 64;
  uint32_t seed = 0;
  bool operator==(const CachedSamplingParams&) const = default;
};
 
template <SamplingParamsLike P>
inline CachedSamplingParams snapshot_params(const P& p, uint32_t absent_seed) {
  using ::lloyal::detail::as_value;
  return CachedSamplingParams{
    as_value(p.temperature, 0.8f),
    as_value(p.top_k, static_cast<int32_t>(40)),
    as_value(p.top_p, 0.95f),
    as_value(p.typical_p, 1.0f),
    as_value(p.min_p, 0.05f),
    as_value(p.penalty_repeat, 1.0f),
    as_value(p.penalty_freq, 0.0f),
    as_value(p.penalty_present, 0.0f),
    as_value(p.penalty_last_n, static_cast<int32_t>(64)),
    as_value(p.seed, absent_seed),
  };
}
 
inline uint32_t clock_seed() { return static_cast<uint32_t>(std::time(nullptr)); }
 
// ===== BRANCH STATE =====
 
struct BranchState {
  llama_context* ctx = nullptr;       
  const llama_model* model = nullptr; 
 
  llama_seq_id seq_id = NO_LEASE;  
  llama_pos position = 0;    
  llama_pos fork_head = 0;   
 
  uint64_t revision = 0;
  uint64_t fork_revision = 0;
  bool poisoned = false;
 
  void poison() {
    poisoned = true;
    ++revision;
  }
 
  uint32_t img_slack_own = 0;
  uint32_t img_slack_total = 0;
 
  SamplerChainHandle sampler_chain = 0;  
  GrammarHandle grammar = 0;             
  bool grammar_sample_first = false;
 
  CachedSamplingParams cached_params;    
 
  boundaries::BoundaryTracker* boundary_tracker = nullptr;  
 
  std::vector<llama_logit_bias> logit_bias;  
  std::function<void(llama_token_data_array&)> steer_fn;  
 
  MetricsHandle metrics = 0;  
 
  llama_token last_token = -1;                   
  std::vector<llama_token_data> last_candidates; 
 
  std::vector<float> logits_snapshot;  
  bool has_logits = false;             
 
  void capture_logits(int32_t index) {
    const float* raw = logits::get(ctx, index);  // throws if ctx is null or no logits at index
    if (n_vocab <= 0 || logits_snapshot.size() < static_cast<size_t>(n_vocab)) {
      throw std::runtime_error("branch: cannot capture logits, branch has no vocab (n_vocab=" +
                               std::to_string(n_vocab) + ")");
    }
    std::memcpy(logits_snapshot.data(), raw, static_cast<size_t>(n_vocab) * sizeof(float));
    has_logits = true;
  }
 
  std::vector<llama_token_data> candidates_buffer;
 
  int n_batch = DEFAULT_N_BATCH;  
  int n_vocab = 0;    
 
  uint16_t generation = 0;  
  bool in_use = false;      
 
  // Topology — maintained by fork/prune/pruneSubtree
  BranchHandle parent = INVALID_HANDLE;     
  std::vector<BranchHandle> children;       
};
 
// ===== BATCHED DECODE ITEM TYPES =====
 
struct DecodeEachItem {
  BranchHandle handle;
  llama_token token;
};
 
struct DecodeScatterItem {
  BranchHandle handle;
  std::span<const llama_token> tokens;
};
 
struct DecodeSegmentsResult {
  int64_t cells = 0;
  llama_pos advance = 0;
};
 
inline void require_distinct_handles(std::span<const BranchHandle> handles, const char* who) {
  std::unordered_map<BranchHandle, size_t> first;
  first.reserve(handles.size());
  for (size_t i = 0; i < handles.size(); ++i) {
    if (handles[i] == INVALID_HANDLE) continue;
    const auto [it, inserted] = first.emplace(handles[i], i);
    if (!inserted) {
      throw std::runtime_error(std::string(who) + " - duplicate handle at indices " +
                               std::to_string(it->second) + " and " + std::to_string(i) +
                               " (sequential calls required for multiple runs per branch)");
    }
  }
}
 
// ===== BRANCH STORE (HANDLE TABLE) =====
 
class BranchStore {
public:
  explicit BranchStore(size_t initial_capacity = 16) {
    // Ensure minimum capacity of 2 (slot 0 reserved + at least 1 usable)
    if (initial_capacity < 2) {
      initial_capacity = 2;
    }
    slots_.resize(initial_capacity);
    // Slot 0 is reserved (handle 0 = invalid)
    slots_[0].in_use = true;
    slots_[0].generation = 0xFFFF;  // Never valid
 
    // Initialize freelist with remaining slots
    // NOTE: Use i-- > 1 pattern to avoid size_t underflow
    for (size_t i = initial_capacity; i-- > 1; ) {
      freelist_.push_back(static_cast<uint16_t>(i));
    }
  }
 
  ~BranchStore() {
    if (tenancy_.ctx != nullptr) {
      LLOYAL_LOG_DEBUG("[BranchStore] WARNING: not drained before destruction");
    }
    for (size_t i = 1; i < slots_.size(); ++i) {
      if (slots_[i].in_use) {
        free_branch_resources(slots_[i]);
      }
    }
  }
 
  struct Allocation { BranchHandle handle; llama_seq_id seq_id; };
 
  Allocation allocate() {
    if (tenancy_.ctx == nullptr) return {INVALID_HANDLE, NO_LEASE};  // drained
    llama_seq_id seq = kv::tenancy::acquire(tenancy_);
    if (seq < 0) return {INVALID_HANDLE, NO_LEASE};
    BranchHandle handle = allocate_slot();
    if (handle == INVALID_HANDLE) {
      // Seq was never used — bookkeeping-only return, no KV calls
      kv::tenancy::release(tenancy_, seq);
      return {INVALID_HANDLE, NO_LEASE};
    }
    // Stamp seq_id on slot so release() can evict properly
    BranchState* st = get(handle);
    st->seq_id = seq;
    return {handle, seq};
  }
 
  void release(BranchHandle handle) {
    if (handle == INVALID_HANDLE) return;
    BranchState* st = get(handle);
    if (!st) return;
    // Eager edge cleanup: remove from parent's children
    if (st->parent != INVALID_HANDLE) {
      BranchState* p = get(st->parent);
      if (p) {
        auto& c = p->children;
        c.erase(std::remove(c.begin(), c.end(), handle), c.end());
      }
    }
    // Subtract unique cells owned by this branch (above fork_head).
    // Position-delta alone undercounts embedding-row prefills (cells grow by
    // n_tokens while position advances by n_pos) — add the branch's own
    // slack so the pressure gauge recovers exactly what was decoded.
    {
      uint32_t unique = st->img_slack_own;
      if (st->position > st->fork_head) {
        unique += static_cast<uint32_t>(st->position - st->fork_head);
      }
      if (unique > 0) {
        cells_used_ = (unique <= cells_used_) ? cells_used_ - unique : 0;
      }
    }
    // Evict lease (KV strip + bookkeeping)
    if (st->seq_id != NO_LEASE)
      kv::tenancy::evict(tenancy_, st->seq_id);
    free_branch_resources(*st);
    reset_slot(*st);
    freelist_.push_back(handle_index(handle));
 
    // All branches released → KV cache empty, reset pressure counter
    if (freelist_.size() == slots_.size() - 1) {
      cells_used_ = 0;
    }
  }
 
  // ===== TENANCY LIFECYCLE =====
 
  void init_tenancy(llama_context* ctx) {
    tenancy_ = kv::tenancy::init(ctx, llama_n_seq_max(ctx));
    cells_used_ = 0;
  }
 
  void drain() {
    if (tenancy_.ctx == nullptr) return;  // idempotent
    kv::tenancy::evict_all(tenancy_);
    for (size_t i = 1; i < slots_.size(); ++i) {
      if (slots_[i].in_use) {
        free_branch_resources(slots_[i]);
        reset_slot(slots_[i]);
      }
    }
    tenancy_.ctx = nullptr;  // marks as drained
    cells_used_ = 0;
  }
 
  void retainOnly(BranchHandle winner) {
    BranchState* w = get(winner);
    if (!w) throw std::runtime_error("retainOnly: invalid winner handle");
    if (w->seq_id == NO_LEASE) throw std::runtime_error("retainOnly: winner has no lease");
    kv::tenancy::retain(tenancy_, w->seq_id);  // nuclear KV pass
    // Collect losers first — don't mutate while iterating
    std::vector<BranchHandle> losers;
    for (size_t i = 1; i < slots_.size(); ++i) {
      if (!slots_[i].in_use) continue;
      BranchHandle h = make_handle(static_cast<uint16_t>(i), slots_[i].generation);
      if (h == winner) continue;
      losers.push_back(h);
    }
    for (auto h : losers)
      release_slot_only(h);  // CPU only, KV already stripped
    w->parent = INVALID_HANDLE;
    w->fork_head = 0;
    w->children.clear();
    // The winner absorbs the whole prefix (fork_head = 0), so its "own"
    // slack becomes the inherited total, and the pressure counter is the
    // prefix's real cell count: position-delta + embedding-row slack.
    w->img_slack_own = w->img_slack_total;
    cells_used_ = static_cast<uint32_t>(w->position) + w->img_slack_total;
    ++w->revision;
  }
 
  void fast_forward(BranchHandle parent, BranchHandle child) {
    BranchState* a = get(parent);
    if (!a) throw std::runtime_error("fast_forward: invalid parent handle");
    BranchState* p = get(child);
    if (!p) throw std::runtime_error("fast_forward: invalid child handle");
    if (p->parent != parent) {
      throw std::runtime_error("fast_forward: the child was not forked from this parent");
    }
    if (a->seq_id == NO_LEASE || p->seq_id == NO_LEASE) {
      throw std::runtime_error("fast_forward: a branch has no lease");
    }
    if (!p->children.empty()) {
      throw std::runtime_error("fast_forward: the child has children of its own");
    }
    if (p->poisoned) {
      throw std::runtime_error("fast_forward: the child is poisoned by a failed decode; prune it");
    }
    if (p->fork_head != a->position) {
      throw std::runtime_error("fast_forward: the parent has decoded since the fork");
    }
    if (p->fork_revision != a->revision) {
      throw std::runtime_error("fast_forward: the parent has changed since the fork");
    }
 
    std::swap(a->seq_id, p->seq_id);
    std::swap(a->logits_snapshot, p->logits_snapshot);
    std::swap(a->has_logits, p->has_logits);
    std::swap(a->sampler_chain, p->sampler_chain);
    std::swap(a->cached_params, p->cached_params);
    std::swap(a->grammar, p->grammar);
    std::swap(a->grammar_sample_first, p->grammar_sample_first);
    std::swap(a->last_token, p->last_token);
    std::swap(a->last_candidates, p->last_candidates);
    std::swap(a->metrics, p->metrics);
    std::swap(a->logit_bias, p->logit_bias);
    a->position = p->position;
    a->img_slack_own += p->img_slack_own;
    a->img_slack_total = p->img_slack_total;
    ++a->revision;
 
    // The child's range is the parent's now, so its release subtracts nothing;
    // it evicts the sequence and frees the resources the parent held before.
    p->position = p->fork_head;
    p->img_slack_own = 0;
    release(child);
  }
 
  // ===== TOPOLOGY QUERIES =====
 
  size_t available() const { return kv::tenancy::available(tenancy_); }
 
  KvPressure kv_pressure() const {
    if (!tenancy_.ctx) return {0, 0, 0};
    uint32_t n_ctx = static_cast<uint32_t>(llama_n_ctx(tenancy_.ctx));
    uint32_t remaining = (n_ctx > cells_used_) ? n_ctx - cells_used_ : 0;
    return { n_ctx, cells_used_, remaining };
  }
 
  void add_cells_used(uint32_t n) { cells_used_ += n; }
 
  BranchHandle parent(BranchHandle h) const {
    const BranchState* st = get(h);
    return st ? st->parent : INVALID_HANDLE;
  }
 
  llama_pos fork_head(BranchHandle h) const {
    const BranchState* st = get(h);
    return st ? st->fork_head : 0;
  }
 
  const std::vector<BranchHandle>& children(BranchHandle h) const {
    static const std::vector<BranchHandle> empty;
    const BranchState* st = get(h);
    return st ? st->children : empty;
  }
 
  bool isLeaf(BranchHandle h) const {
    const BranchState* st = get(h);
    return st ? st->children.empty() : true;
  }
 
  bool isActive(BranchHandle h) const {
    const BranchState* st = get(h);
    return st ? (st->seq_id != NO_LEASE) : false;
  }
 
  BranchState* get(BranchHandle handle) {
    if (handle == INVALID_HANDLE) return nullptr;
 
    uint16_t index = handle_index(handle);
    uint16_t gen = handle_generation(handle);
 
    // Slot 0 is reserved and never valid for external use
    if (index == 0) return nullptr;
 
    if (index >= slots_.size()) return nullptr;
 
    BranchState& slot = slots_[index];
    if (!slot.in_use || slot.generation != gen) {
      return nullptr;
    }
 
    return &slot;
  }
 
  const BranchState* get(BranchHandle handle) const {
    return const_cast<BranchStore*>(this)->get(handle);
  }
 
  // ===== SAMPLER CHAIN REGISTRY =====
 
  template <SamplingParamsLike P>
  SamplerChainHandle create_sampler(const P& params, int32_t n_vocab) {
    SamplerChainHandle h = next_sampler_handle_++;
    SamplerChainEntry entry;
    entry.chain = sampler::create_chain(params, n_vocab);
    float temperature = ::lloyal::detail::as_value(params.temperature, 0.8f);
    entry.has_dist = (temperature > 0.0f);
    sampler_chains_.emplace(h, std::move(entry));
    return h;
  }
 
  SamplerChainHandle clone_sampler(SamplerChainHandle h) {
    if (h == 0) return 0;
    auto it = sampler_chains_.find(h);
    if (it == sampler_chains_.end()) return 0;
    SamplerChainHandle nh = next_sampler_handle_++;
    SamplerChainEntry entry;
    entry.chain = sampler::clone_chain(it->second.chain);
    entry.has_dist = it->second.has_dist;
    sampler_chains_.emplace(nh, std::move(entry));
    return nh;
  }
 
  void free_sampler(SamplerChainHandle h) {
    if (h != 0) sampler_chains_.erase(h);
  }
 
  llama_sampler* get_sampler_chain(SamplerChainHandle h) const {
    if (h == 0) return nullptr;
    auto it = sampler_chains_.find(h);
    return it != sampler_chains_.end() ? it->second.chain : nullptr;
  }
 
  bool sampler_has_dist(SamplerChainHandle h) const {
    if (h == 0) return false;
    auto it = sampler_chains_.find(h);
    return it != sampler_chains_.end() ? it->second.has_dist : false;
  }
 
  // ===== GRAMMAR REGISTRY =====
 
  GrammarHandle create_grammar(const llama_model* model,
                               const char* grammar_str,
                               const char* root = "root") {
    GrammarHandle h = next_grammar_handle_++;
    GrammarEntry entry;
    entry.sampler = grammar::init_sampler(model, grammar_str, root);
    grammars_.emplace(h, std::move(entry));
    return h;
  }
 
  GrammarHandle create_grammar_lazy(
      const llama_model* model,
      const char* grammar_str,
      const std::vector<std::string>& trigger_patterns,
      const std::vector<llama_token>& trigger_tokens,
      const char* root = "root") {
    GrammarHandle h = next_grammar_handle_++;
    GrammarEntry entry;
    entry.sampler = grammar::init_lazy_sampler(
        model, grammar_str, trigger_patterns, trigger_tokens, root);
    if (!entry.sampler) return 0;
    grammars_.emplace(h, std::move(entry));
    return h;
  }
 
  GrammarHandle clone_grammar(GrammarHandle h) {
    if (h == 0) return 0;
    auto it = grammars_.find(h);
    if (it == grammars_.end()) return 0;
    GrammarHandle nh = next_grammar_handle_++;
    GrammarEntry entry;
    entry.sampler = grammar::clone_sampler(it->second.sampler);
    grammars_.emplace(nh, std::move(entry));
    return nh;
  }
 
  void free_grammar(GrammarHandle h) {
    if (h != 0) grammars_.erase(h);
  }
 
  llama_sampler* get_grammar_sampler(GrammarHandle h) const {
    if (h == 0) return nullptr;
    auto it = grammars_.find(h);
    return it != grammars_.end() ? it->second.sampler : nullptr;
  }
 
  // ===== METRICS REGISTRY =====
 
  MetricsHandle create_metrics() {
    MetricsHandle h = next_metrics_handle_++;
    metrics_registry_[h] = metrics::BranchMetricsState{};
    return h;
  }
 
  MetricsHandle clone_metrics(MetricsHandle h) {
    if (h == 0) return 0;
    auto it = metrics_registry_.find(h);
    if (it == metrics_registry_.end()) return 0;
    MetricsHandle nh = next_metrics_handle_++;
    metrics_registry_[nh] = it->second;
    return nh;
  }
 
  void free_metrics(MetricsHandle h) {
    if (h != 0) metrics_registry_.erase(h);
  }
 
  void add_model_surprisal(MetricsHandle h, float surprisal) {
    if (h == 0) return;
    auto it = metrics_registry_.find(h);
    if (it == metrics_registry_.end()) return;
    if (!std::isfinite(surprisal)) return;
    it->second.model.nll_sum_nats += std::max(0.0f, surprisal);
    it->second.model.count++;
  }
 
  void add_sampling_surprisal(MetricsHandle h, float surprisal) {
    if (h == 0) return;
    auto it = metrics_registry_.find(h);
    if (it == metrics_registry_.end()) return;
    if (!std::isfinite(surprisal)) return;
    it->second.sampling.nll_sum_nats += std::max(0.0f, surprisal);
    it->second.sampling.count++;
  }
 
  float get_model_ppl(MetricsHandle h) const {
    if (h == 0) return std::numeric_limits<float>::infinity();
    auto it = metrics_registry_.find(h);
    if (it == metrics_registry_.end() || it->second.model.count == 0)
      return std::numeric_limits<float>::infinity();
    return std::exp(it->second.model.nll_sum_nats /
                    static_cast<float>(it->second.model.count));
  }
 
  float get_sampling_ppl(MetricsHandle h) const {
    if (h == 0) return std::numeric_limits<float>::infinity();
    auto it = metrics_registry_.find(h);
    if (it == metrics_registry_.end() || it->second.sampling.count == 0)
      return std::numeric_limits<float>::infinity();
    return std::exp(it->second.sampling.nll_sum_nats /
                    static_cast<float>(it->second.sampling.count));
  }
 
  // ===== BATCHED DECODE =====
 
  void decode_each(std::span<const DecodeEachItem> items) {
    if (items.empty()) return;
 
    const int32_t n = static_cast<int32_t>(items.size());
 
    // Validate handles and collect states
    std::vector<BranchState*> states(n);
    for (int32_t i = 0; i < n; ++i) {
      states[i] = get(items[i].handle);
      if (!states[i]) {
        throw std::runtime_error("BranchStore::decode_each - invalid handle at index " + std::to_string(i));
      }
      if (i > 0 && states[i]->ctx != states[0]->ctx) {
        throw std::runtime_error("BranchStore::decode_each - all branches must share the same context");
      }
    }
 
    {
      std::vector<BranchHandle> handles(n);
      for (int32_t i = 0; i < n; ++i) handles[i] = items[i].handle;
      require_distinct_handles(handles, "BranchStore::decode_each");
    }
 
    // Build EachItem array from branch states
    std::vector<decode::EachItem> decode_items(n);
    for (int32_t i = 0; i < n; ++i) {
      decode_items[i].token = items[i].token;
      decode_items[i].pos = states[i]->position;
      decode_items[i].seq_id = states[i]->seq_id;
      decode_items[i].output_logits = true;
    }
 
    // Single GPU dispatch
    if (const int32_t rc = decode::each(states[0]->ctx, decode_items.data(), n, scratch_); rc != 0) {
      if (!decode::leaves_branch_intact(rc, /*partial*/ false)) {
        for (auto* st : states) st->poison();
      }
      throw decode::DecodeError(rc, /*partial*/ false, "BranchStore::decode_each - llama_decode failed");
    }
 
    // Capture logits and update positions
    for (int32_t i = 0; i < n; ++i) {
      states[i]->capture_logits(i);
      states[i]->position += 1;
      ++states[i]->revision;
    }
    cells_used_ += static_cast<uint32_t>(items.size());
  }
 
  void decode_scatter(std::span<const DecodeScatterItem> items) {
    if (items.empty()) return;
 
    const int32_t n = static_cast<int32_t>(items.size());
 
    // Validate handles and collect states
    std::vector<BranchState*> states(n);
    for (int32_t i = 0; i < n; ++i) {
      states[i] = get(items[i].handle);
      if (!states[i]) {
        throw std::runtime_error("BranchStore::decode_scatter - invalid handle at index " + std::to_string(i));
      }
      if (i > 0 && states[i]->ctx != states[0]->ctx) {
        throw std::runtime_error("BranchStore::decode_scatter - all branches must share the same context");
      }
    }
    {
      // Only spans that DECODE can collide: an empty item is skipped by
      // bin_pack and occupies no cells, so it is INVALID_HANDLE to the rule.
      std::vector<BranchHandle> material(n);
      for (int32_t i = 0; i < n; ++i) {
        material[i] = items[i].tokens.empty() ? INVALID_HANDLE : items[i].handle;
      }
      require_distinct_handles(material, "BranchStore::decode_scatter");
    }
 
    llama_context* ctx = states[0]->ctx;
    const int32_t batch_limit = static_cast<int32_t>(llama_n_batch(ctx));
 
    // Build flat token spans — bin_pack skips empties internally
    std::vector<std::span<const llama_token>> spans(n);
    for (int32_t i = 0; i < n; ++i) {
      spans[i] = items[i].tokens;
    }
 
    auto chunks = decode::bin_pack(spans.data(), n, batch_limit);
 
    int32_t landed = 0;  // chunks whose llama_decode returned 0
    for (const auto& chunk : chunks) {
      if (chunk.oversized) {
        int32_t idx = chunk.indices[0];
        int32_t tc = static_cast<int32_t>(items[idx].tokens.size());
 
        int32_t committed = 0;
        if (const int32_t rc = decode::many(ctx, items[idx].tokens.data(), tc,
                                            states[idx]->position, batch_limit,
                                            states[idx]->seq_id, &committed); rc != 0) {
          if (!decode::leaves_branch_intact(rc, committed > 0)) states[idx]->poison();
          throw decode::DecodeError(rc, landed > 0 || committed > 0,
              "BranchStore::decode_scatter - decode::many failed for oversized item " + std::to_string(idx));
        }
 
        land_tokens(*states[idx], tc, -1);
        ++landed;
        continue;
      }
 
      // Normal chunk — build ScatterItems and dispatch
      std::vector<decode::ScatterItem> scatter_items(chunk.indices.size());
      for (size_t k = 0; k < chunk.indices.size(); ++k) {
        int32_t idx = chunk.indices[k];
        scatter_items[k].tokens = items[idx].tokens;
        scatter_items[k].start_pos = states[idx]->position;
        scatter_items[k].seq_id = states[idx]->seq_id;
        scatter_items[k].output_logits = true;
      }
 
      if (const int32_t rc = decode::scatter(ctx, scatter_items.data(),
                                             static_cast<int32_t>(scatter_items.size()),
                                             scratch_); rc != 0) {
        // Each branch in the failing chunk lost only this call: intact iff it restored.
        if (!decode::leaves_branch_intact(rc, /*partial*/ false)) {
          for (const int32_t idx : chunk.indices) states[idx]->poison();
        }
        throw decode::DecodeError(rc, landed > 0, "BranchStore::decode_scatter - decode::scatter failed");
      }
 
      // Capture logits and move the books for each item in the chunk
      int32_t cursor = 0;
      for (size_t k = 0; k < scatter_items.size(); ++k) {
        int32_t idx = chunk.indices[k];
        int32_t item_n = static_cast<int32_t>(scatter_items[k].tokens.size());
 
        land_tokens(*states[idx], item_n, cursor + item_n - 1);
        cursor += item_n;
      }
      ++landed;
    }
  }
 
  void decode_embd(BranchHandle handle,
                    const float* rows, int32_t n_tokens, int32_t n_embd_inp,
                    llama_pos n_pos,
                    const llama_pos* pos, int32_t n_pos_per_embd,
                    bool non_causal, bool want_logits) {
    BranchState* state = get(handle);
    if (!state) {
      throw std::runtime_error("BranchStore::decode_embd - invalid handle");
    }
    if (n_pos <= 0 || n_pos > n_tokens) {
      throw std::runtime_error("BranchStore::decode_embd - invalid n_pos");
    }
 
    decode::EmbdItem item;
    item.rows           = rows;
    item.n_rows         = n_tokens;
    item.n_embd_inp     = n_embd_inp;
    item.pos            = pos;
    item.n_pos_per_embd = n_pos_per_embd;
    item.seq_id         = state->seq_id;
    item.non_causal     = non_causal;
    item.output_logits  = want_logits;
 
    // NOT atomic, and deliberately not made so. decode::embd chunks by
    // n_batch, so a failure on a later chunk leaves EARLIER chunks committed
    // while the counters below never move. Rolling those rows back is not
    // available: seq_rm can only rewind a recurrent carrier where the model
    // keeps per-token snapshots, and Gated DeltaNet reports rs_seq = 0, so a
    // rollback there silently restores attention and not the carrier.
    //
    // The branch is therefore poisoned on failure — prune it and replay onto
    // a fresh one, which is the portable correction regardless of layer type.
    // The error's `partial` says whether any chunk landed (see DecodeError).
    // Its own accounting stays CONSISTENT for that prune: neither position
    // nor cells_used_ moved, so release() subtracts what the branch legitimately
    // owned and whole-sequence eviction reclaims the orphaned rows.
    //
    // Admission is the better place to spend effort: SegmentSource::cells()
    // reports this prefill's cost before anything decodes, so a caller can
    // refuse rather than half-commit.
    int32_t committed = 0;
    if (const int32_t rc = decode::embd(state->ctx, item, state->n_batch, scratch_, &committed); rc != 0) {
      if (!decode::leaves_branch_intact(rc, committed > 0)) state->poison();
      throw decode::DecodeError(rc, committed > 0,
          "BranchStore::decode_embd - llama_decode failed; this branch is "
          "poisoned, prune it and replay onto a fresh one");
    }
 
    land_rows(*state, n_tokens, n_pos, want_logits);
  }
 
  DecodeSegmentsResult decode_segments(BranchHandle handle, decode::SegmentSource& source,
                                       const std::function<void(const decode::SegmentProgress&)>& on_progress = {}) {
    BranchState* state = get(handle);
    if (!state) {
      throw std::runtime_error("BranchStore::decode_segments - invalid handle");
    }
 
    const llama_pos start_pos = state->position;
    DecodeSegmentsResult result;
 
    std::vector<size_t> landed_by_input;
    const auto report = [&](const decode::LandedSegment& seg) {
      const int32_t input = source.input_of(seg.index);
      if (!on_progress || input < 0) return;
      if (landed_by_input.size() <= static_cast<size_t>(input)) landed_by_input.resize(input + 1, 0);
      landed_by_input[input] += static_cast<size_t>(seg.cells);
      on_progress(decode::SegmentProgress{input, landed_by_input[input], source.total_of(input)});
    };
 
    try {
      decode::segments(state->ctx, state->seq_id, start_pos, source, state->n_batch,
                       scratch_, segment_pos_, [&](const decode::LandedSegment& seg) {
        if (seg.kind == decode::Segment::Kind::Text) {
          land_tokens(*state, seg.cells, -1);
        } else {
          land_rows(*state, seg.cells, seg.n_pos, seg.has_logits);
        }
        result.cells += seg.cells;
        report(seg);
      });
    } catch (const decode::DecodeError& e) {
      // One operation to the caller. The rail marks its own failures partial
      // when an earlier segment landed or a chunk of this one did; a decode
      // error from the source or from on_progress cannot know what landed. The
      // branch's position does: every segment is material, so a moved position
      // is proof the prefill is part-way through.
      const bool moved = state->position != start_pos;
      if (moved || !decode::leaves_branch_intact(e.rc, e.partial)) state->poison();
      if (e.partial || !moved) throw;
      throw e.as_partial();
    } catch (...) {
      // The source threw after a segment landed: the branch is part-way through.
      if (state->position != start_pos) state->poison();
      throw;
    }
 
    result.advance = state->position - start_pos;
    return result;
  }
 
  void merge_logits(
      BranchHandle dst_handle,
      std::span<const BranchHandle> expert_handles,
      float alpha) {
    BranchState* dst = get(dst_handle);
    if (!dst) {
      throw std::runtime_error("BranchStore::merge_logits - invalid dst handle");
    }
    if (!dst->has_logits) {
      throw std::runtime_error("BranchStore::merge_logits - dst has no captured logits");
    }
    if (dst->n_vocab <= 0) {
      throw std::runtime_error("BranchStore::merge_logits - dst n_vocab invalid");
    }
 
    const int32_t n_vocab = dst->n_vocab;
 
    // Every expert is checked before dst is written, so a refusal changes nothing.
    std::vector<const float*> sources;
    sources.reserve(expert_handles.size());
    for (size_t i = 0; i < expert_handles.size(); ++i) {
      BranchState* src = get(expert_handles[i]);
      if (!src) {
        throw std::runtime_error(
            "BranchStore::merge_logits - invalid expert handle at index " +
            std::to_string(i));
      }
      if (!src->has_logits) {
        throw std::runtime_error(
            "BranchStore::merge_logits - expert has no captured logits at index " +
            std::to_string(i));
      }
      if (src->n_vocab != n_vocab) {
        throw std::runtime_error(
            "BranchStore::merge_logits - n_vocab mismatch at index " +
            std::to_string(i));
      }
      sources.push_back(src->logits_snapshot.data());
    }
    if (sources.empty()) return;  // nothing merged, nothing changed
 
    float* dst_logits = dst->logits_snapshot.data();
    for (const float* src_logits : sources) {
      for (int32_t t = 0; t < n_vocab; ++t) {
        dst_logits[t] += alpha * src_logits[t];
      }
    }
    ++dst->revision;
  }
 
private:
  void land_tokens(BranchState& state, int32_t n_tokens, int32_t logits_index) {
    state.capture_logits(logits_index);
    state.position += n_tokens;
    ++state.revision;
    cells_used_ += static_cast<uint32_t>(n_tokens);
  }
 
  void land_rows(BranchState& state, int32_t n_rows, llama_pos n_pos, bool with_logits) {
    state.position += n_pos;
    cells_used_ += static_cast<uint32_t>(n_rows);
    const uint32_t slack = static_cast<uint32_t>(n_rows - n_pos);
    state.img_slack_own += slack;
    state.img_slack_total += slack;
    ++state.revision;
 
    if (with_logits) {
      state.capture_logits(-1);
    } else {
      // The position advanced but no logits were computed for it. Leaving the
      // previous snapshot in place would let sample() read logits belonging to
      // an EARLIER position — wrong, and silently so. The token rails always
      // recapture, so rows are the one place that can strand them.
      state.has_logits = false;
    }
  }
 
  void free_branch_resources(BranchState& slot) {
    if (slot.sampler_chain != 0) {
      free_sampler(slot.sampler_chain);
      slot.sampler_chain = 0;
    }
    if (slot.grammar != 0) {
      free_grammar(slot.grammar);
      slot.grammar = 0;
    }
    if (slot.boundary_tracker) {
      delete slot.boundary_tracker;
      slot.boundary_tracker = nullptr;
    }
    if (slot.metrics != 0) {
      free_metrics(slot.metrics);
      slot.metrics = 0;
    }
  }
 
  void reset_slot(BranchState& slot) {
    slot.in_use = false;
    slot.generation = static_cast<uint16_t>(slot.generation + 1);  // Prevent ABA
    slot.ctx = nullptr;
    slot.model = nullptr;
    slot.seq_id = NO_LEASE;
    slot.position = 0;
    slot.fork_head = 0;
    slot.revision = 0;
    slot.fork_revision = 0;
    slot.poisoned = false;
    slot.img_slack_own = 0;
    slot.img_slack_total = 0;
    slot.sampler_chain = 0;
    slot.grammar = 0;
    slot.grammar_sample_first = false;
    slot.metrics = 0;
    slot.cached_params = CachedSamplingParams{};
    slot.last_token = -1;
    slot.last_candidates.clear();
    slot.logits_snapshot.clear();
    slot.has_logits = false;
    slot.logit_bias.clear();
    slot.steer_fn = nullptr;
    slot.candidates_buffer.clear();
    slot.n_batch = DEFAULT_N_BATCH;
    slot.n_vocab = 0;
    slot.parent = INVALID_HANDLE;
    slot.children.clear();
  }
 
  BranchHandle allocate_slot() {
    if (freelist_.empty()) {
      size_t old_size = slots_.size();
      size_t new_size = old_size * 2;
      if (new_size > INDEX_MASK) {
        new_size = INDEX_MASK + 1;
      }
      if (old_size >= new_size) {
        LLOYAL_LOG_DEBUG("[branch::allocate_slot] Store full, cannot allocate");
        return INVALID_HANDLE;
      }
      slots_.resize(new_size);
      for (size_t i = new_size; i-- > old_size; ) {
        freelist_.push_back(static_cast<uint16_t>(i));
      }
    }
    uint16_t index = freelist_.back();
    freelist_.pop_back();
    BranchState& slot = slots_[index];
    slot.in_use = true;
    return make_handle(index, slot.generation);
  }
 
  void release_slot_only(BranchHandle handle) {
    if (handle == INVALID_HANDLE) return;
    BranchState* st = get(handle);
    if (!st) return;
    // Eager edge cleanup
    if (st->parent != INVALID_HANDLE) {
      BranchState* p = get(st->parent);
      if (p) {
        auto& c = p->children;
        c.erase(std::remove(c.begin(), c.end(), handle), c.end());
      }
    }
    free_branch_resources(*st);
    reset_slot(*st);
    freelist_.push_back(handle_index(handle));
  }
 
  std::deque<BranchState> slots_;
  std::vector<uint16_t> freelist_;  
 
  kv::tenancy::State tenancy_;
 
  uint32_t cells_used_ = 0;
 
  decode::Scratch scratch_;
 
  std::vector<llama_pos> segment_pos_;
 
  // ===== Handle registries (instance-scoped, not global static) =====
 
  std::unordered_map<SamplerChainHandle, SamplerChainEntry> sampler_chains_;
  SamplerChainHandle next_sampler_handle_ = 1;
 
  std::unordered_map<GrammarHandle, GrammarEntry> grammars_;
  GrammarHandle next_grammar_handle_ = 1;
 
  std::unordered_map<MetricsHandle, metrics::BranchMetricsState> metrics_registry_;
  MetricsHandle next_metrics_handle_ = 1;
};
 
 
// ===== BRANCH API =====
 
template <SamplingParamsLike P>
inline BranchHandle create(
    llama_context* ctx,
    const llama_model* model,
    BranchStore& s,
    llama_pos start_pos,
    const P& params,
    int n_batch = DEFAULT_N_BATCH,
    const char* grammar_str = nullptr,
    boundaries::BoundaryTracker* boundary_tracker = nullptr) {
  if (!ctx || !model) {
    LLOYAL_LOG_DEBUG("[branch::create] NULL ctx or model");
    return INVALID_HANDLE;
  }
 
  // The invariant every logits capture stands on is established here, at
  // birth: a live branch has somewhere to put its logits. Decided BEFORE
  // anything is allocated, so a refusal owns nothing — no slot, no lease,
  // and no range for release() to subtract from the store's books.
  const llama_vocab* vocab = llama_model_get_vocab(model);
  const int n_vocab = llama_vocab_n_tokens(vocab);
  if (n_vocab <= 0) {
    throw std::runtime_error("branch::create - model has no vocab (n_vocab=" +
                             std::to_string(n_vocab) + "); nothing to sample");
  }
 
  auto [handle, seq_id] = s.allocate();
  if (handle == INVALID_HANDLE) {
    return INVALID_HANDLE;
  }
 
  BranchState* state = s.get(handle);
  if (!state) {
    s.release(handle);
    return INVALID_HANDLE;
  }
 
  state->ctx = ctx;
  state->model = model;
  // seq_id already stamped by allocate()
  state->position = start_pos;
  state->n_batch = n_batch;
 
  state->n_vocab = n_vocab;
  state->logits_snapshot.resize(state->n_vocab);
  state->has_logits = false;
  state->candidates_buffer.resize(state->n_vocab);
 
  // The chain is built from the resolved record, so the cache says exactly
  // what the chain draws with — a seedless branch's clock seed included.
  state->cached_params = snapshot_params(params, clock_seed());
  state->sampler_chain = s.create_sampler(state->cached_params, n_vocab);
 
  if (grammar_str && grammar_str[0] != '\0') {
    state->grammar = s.create_grammar(model, grammar_str);
  }
 
  state->boundary_tracker = boundary_tracker;
  state->metrics = s.create_metrics();
 
  LLOYAL_LOG_DEBUG("[branch::create] Created branch handle=%u seq=%d pos=%d",
                   handle, seq_id, start_pos);
 
  return handle;
}
 
 
struct ForkOpts {
  bool clone_logits = true;
};
 
inline BranchHandle fork(BranchHandle source, BranchStore& s, ForkOpts opts = {}) {
  BranchState* src = s.get(source);
  if (!src) {
    LLOYAL_LOG_DEBUG("[branch::fork] Invalid source handle");
    return INVALID_HANDLE;
  }
 
  auto [new_handle, new_seq_id] = s.allocate();
  if (new_handle == INVALID_HANDLE) {
    return INVALID_HANDLE;
  }
 
  BranchState* dst = s.get(new_handle);
  if (!dst) {
    s.release(new_handle);
    return INVALID_HANDLE;
  }
 
  // Copy basic state
  dst->ctx = src->ctx;
  dst->model = src->model;
  dst->seq_id = new_seq_id;
  dst->position = src->position;
  dst->fork_head = src->position;
  dst->fork_revision = src->revision;
  dst->poisoned = src->poisoned;
  // Inherit the prefix's embedding-row slack; the child owns none yet.
  dst->img_slack_total = src->img_slack_total;
  dst->img_slack_own = 0;
  dst->n_batch = src->n_batch;
  dst->n_vocab = src->n_vocab;
 
#ifndef NDEBUG
  assert(kv::pos_max(src->ctx, new_seq_id) < 0 && "tenancy: acquired seq must be clean");
  assert(dst->parent == INVALID_HANDLE && dst->children.empty() && "fresh slot must have no topology");
#endif
 
  // Fork KV cache
  kv::seq_cp(src->ctx, src->seq_id, new_seq_id);
 
  // Record topology
  dst->parent = source;
  src->children.push_back(new_handle);
 
  // Clone sampler chain
  if (src->sampler_chain != 0) {
    dst->sampler_chain = s.clone_sampler(src->sampler_chain);
  }
  dst->cached_params = src->cached_params;
 
  if (src->grammar != 0) {
    dst->grammar = s.clone_grammar(src->grammar);
  }
  dst->grammar_sample_first = src->grammar_sample_first;
 
  if (src->boundary_tracker) {
    dst->boundary_tracker = src->boundary_tracker->clone().release();
  }
 
  if (src->metrics != 0) {
    dst->metrics = s.clone_metrics(src->metrics);
  }
 
  dst->last_token = src->last_token;
  dst->last_candidates = src->last_candidates;
 
  // Logits snapshot: default preserves fork's "sample same distribution" contract.
  // opts.clone_logits=false skips the ~600KB memcpy from src for prefill-overwrite
  // consumers (rerank leaves, embedding probes). Buffer is still kept sized to
  // n_vocab so subsequent decode_scatter/prefill can write into it.
  if (opts.clone_logits) {
    dst->logits_snapshot = src->logits_snapshot;
    dst->has_logits = src->has_logits;
  } else {
    if (static_cast<int32_t>(dst->logits_snapshot.size()) != dst->n_vocab) {
      dst->logits_snapshot.resize(dst->n_vocab);
    }
    dst->has_logits = false;
  }
  dst->logit_bias = src->logit_bias;
 
  dst->candidates_buffer.resize(dst->n_vocab);
 
  LLOYAL_LOG_DEBUG("[branch::fork] Forked handle=%u -> handle=%u seq=%d->%d",
                   source, new_handle, src->seq_id, new_seq_id);
 
  return new_handle;
}
 
inline void set_logit_bias(
    BranchHandle handle,
    const llama_logit_bias* biases,
    size_t n_biases,
    BranchStore& s) {
 
 
  BranchState* state = s.get(handle);
  if (!state) {
    throw std::runtime_error("set_logit_bias: invalid branch handle");
  }
 
  // Replace existing biases with new set
  state->logit_bias.assign(biases, biases + n_biases);
  ++state->revision;
 
  LLOYAL_LOG_DEBUG("[branch::set_logit_bias] Set %zu biases on handle=%u",
                   n_biases, handle);
}
 
inline void clear_logit_bias(
    BranchHandle handle,
    BranchStore& s) {
 
 
  BranchState* state = s.get(handle);
  if (!state) {
    throw std::runtime_error("clear_logit_bias: invalid branch handle");
  }
 
  state->logit_bias.clear();
  ++state->revision;
 
  LLOYAL_LOG_DEBUG("[branch::clear_logit_bias] Cleared biases on handle=%u", handle);
}
 
inline void set_steer(
    BranchHandle handle,
    std::function<void(llama_token_data_array&)> steer_fn,
    BranchStore& s) {
 
 
  BranchState* state = s.get(handle);
  if (!state) {
    throw std::runtime_error("set_steer: invalid branch handle");
  }
 
  state->steer_fn = std::move(steer_fn);
 
  LLOYAL_LOG_DEBUG("[branch::set_steer] Set steer callback on handle=%u", handle);
}
 
inline void clear_steer(
    BranchHandle handle,
    BranchStore& s) {
 
 
  BranchState* state = s.get(handle);
  if (!state) {
    throw std::runtime_error("clear_steer: invalid branch handle");
  }
 
  state->steer_fn = nullptr;
 
  LLOYAL_LOG_DEBUG("[branch::clear_steer] Cleared steer callback on handle=%u", handle);
}
 
template <SamplingParamsLike P>
inline void set_sampler_params(BranchHandle handle, const P& params, BranchStore& s) {
  BranchState* state = s.get(handle);
  if (!state) {
    throw std::runtime_error("set_sampler_params: invalid branch handle");
  }
 
  CachedSamplingParams new_params = snapshot_params(params, state->cached_params.seed);
  if (new_params == state->cached_params && state->sampler_chain != 0) {
    return;  // Memoized — no rebuild needed
  }
 
  // Free old chain
  if (state->sampler_chain != 0) {
    s.free_sampler(state->sampler_chain);
  }
 
  // Create new chain, from the resolved record so the chain draws with the seed the cache says.
  state->sampler_chain = s.create_sampler(new_params, state->n_vocab);
  state->cached_params = new_params;
  ++state->revision;
 
  LLOYAL_LOG_DEBUG("[branch::set_sampler_params] Rebuilt chain on handle=%u temp=%.3f",
                   handle, new_params.temperature);
}
 
inline void set_grammar(
    BranchHandle handle,
    const llama_model* model,
    const char* grammar_str,
    BranchStore& s) {
  BranchState* state = s.get(handle);
  if (!state) {
    throw std::runtime_error("set_grammar: invalid branch handle");
  }
 
  // Free old grammar
  if (state->grammar != 0) {
    s.free_grammar(state->grammar);
    state->grammar = 0;
  }
  state->grammar_sample_first = false;
 
  // Create new grammar if provided
  if (grammar_str && grammar_str[0] != '\0') {
    state->grammar = s.create_grammar(model, grammar_str);
  }
  ++state->revision;
 
  LLOYAL_LOG_DEBUG("[branch::set_grammar] %s grammar on handle=%u",
                   state->grammar != 0 ? "Set" : "Cleared", handle);
}
 
inline void set_grammar_lazy(
    BranchHandle handle,
    const llama_model* model,
    const char* grammar_str,
    const std::vector<std::string>& trigger_patterns,
    const std::vector<llama_token>& trigger_tokens,
    BranchStore& s) {
  BranchState* state = s.get(handle);
  if (!state) {
    throw std::runtime_error("set_grammar_lazy: invalid branch handle");
  }
 
  if (state->grammar != 0) {
    s.free_grammar(state->grammar);
    state->grammar = 0;
  }
  state->grammar_sample_first = false;
 
  if (grammar_str && grammar_str[0] != '\0') {
    state->grammar = s.create_grammar_lazy(
        model, grammar_str, trigger_patterns, trigger_tokens);
  }
  ++state->revision;
 
  LLOYAL_LOG_DEBUG("[branch::set_grammar_lazy] %s grammar on handle=%u",
                   state->grammar != 0 ? "Set" : "Cleared", handle);
}
 
inline void set_grammar_order(BranchHandle handle, bool sample_first, BranchStore& s) {
  BranchState* state = s.get(handle);
  if (!state) {
    throw std::runtime_error("set_grammar_order: invalid branch handle");
  }
  if (state->grammar_sample_first != sample_first) {
    state->grammar_sample_first = sample_first;
    ++state->revision;
  }
}
 
inline void prune(BranchHandle handle, BranchStore& s) {
  BranchState* state = s.get(handle);
  if (!state) return;
  if (!state->children.empty())
    throw std::runtime_error("prune: RESTRICT — branch has children. Use pruneSubtree() for CASCADE.");
  s.release(handle);
}
 
inline void pruneSubtree(BranchHandle h, BranchStore& s) {
  std::vector<BranchHandle> stack{h}, post_order;
  while (!stack.empty()) {
    BranchHandle cur = stack.back(); stack.pop_back();
    post_order.push_back(cur);
    BranchState* st = s.get(cur);
    if (st) for (auto child : st->children) stack.push_back(child);
  }
  for (auto it = post_order.rbegin(); it != post_order.rend(); ++it)
    prune(*it, s);
}
 
inline void reseed_sampler(BranchHandle handle, uint32_t seed, BranchStore& s) {
  BranchState* state = s.get(handle);
  if (!state) {
    throw std::runtime_error("reseed_sampler: invalid branch handle");
  }
  if (!s.sampler_has_dist(state->sampler_chain)) return;
  sampler::reseed_chain(s.get_sampler_chain(state->sampler_chain), seed);
  state->cached_params.seed = seed;
  ++state->revision;
}
 
inline void fast_forward(BranchHandle parent, BranchHandle child, BranchStore& s) {
  s.fast_forward(parent, child);
}
 
inline bool is_poisoned(BranchHandle handle, BranchStore& s) {
  const BranchState* state = s.get(handle);
  return state && state->poisoned;
}
 
inline void force_snapshot_logits(BranchHandle handle, BranchStore& s) {
 
 
  BranchState* state = s.get(handle);
  if (!state) {
    throw std::runtime_error("force_snapshot_logits: invalid branch handle");
  }
 
  state->capture_logits(-1);
  ++state->revision;
}
 
inline void prefill(
    BranchHandle handle,
    const llama_token* tokens,
    size_t n_tokens,
    BranchStore& s) {
 
 
  BranchState* state = s.get(handle);
  if (!state) {
    throw std::runtime_error("prefill: invalid branch handle");
  }
 
  // Pass raw pointer directly - no vector copy needed
  int32_t committed = 0;
  if (const int32_t rc = decode::many(state->ctx, tokens, static_cast<int32_t>(n_tokens),
                                      state->position, state->n_batch, state->seq_id,
                                      &committed); rc != 0) {
    if (!decode::leaves_branch_intact(rc, committed > 0)) state->poison();
    throw decode::DecodeError(rc, committed > 0, "prefill: llama_decode failed");
  }
 
  state->position += static_cast<llama_pos>(n_tokens);
  s.add_cells_used(static_cast<uint32_t>(n_tokens));
  state->capture_logits(-1);
  ++state->revision;
}
 
inline void step(
    BranchHandle handle,
    llama_token token,
    BranchStore& s) {
 
 
  BranchState* state = s.get(handle);
  if (!state) {
    throw std::runtime_error("step: invalid branch handle");
  }
 
  if (const int32_t rc = decode::one(state->ctx, token, state->position, state->seq_id, true); rc != 0) {
    if (!decode::leaves_branch_intact(rc, /*partial*/ false)) state->poison();
    throw decode::DecodeError(rc, /*partial*/ false, "step: llama_decode failed");
  }
  state->position += 1;
  s.add_cells_used(1);
  state->capture_logits(-1);
  ++state->revision;
}
 
inline const float* get_logits(BranchHandle handle, BranchStore& s) {
 
 
  const BranchState* state = s.get(handle);
  // Must check has_logits, not just empty() - buffer is pre-allocated in create()
  if (!state || !state->has_logits) {
    return nullptr;
  }
 
  return state->logits_snapshot.data();
}
 
inline void set_logits(BranchHandle handle, std::span<const float> data, BranchStore& s) {
  BranchState* state = s.get(handle);
  if (!state) {
    throw std::runtime_error("set_logits: invalid handle");
  }
  if (state->n_vocab <= 0) {
    throw std::runtime_error("set_logits: invalid vocab size");
  }
  if (static_cast<int32_t>(data.size()) != state->n_vocab) {
    throw std::runtime_error(
        "set_logits: data length (" + std::to_string(data.size()) +
        ") does not match branch n_vocab (" + std::to_string(state->n_vocab) + ")");
  }
 
  std::memcpy(state->logits_snapshot.data(), data.data(),
              state->n_vocab * sizeof(float));
  state->has_logits = true;
  ++state->revision;
}
 
inline void get_logits_at(BranchHandle handle,
                          std::span<const int32_t> indices,
                          std::span<float> out,
                          BranchStore& s) {
  if (indices.size() != out.size()) {
    throw std::runtime_error("get_logits_at: indices/out size mismatch");
  }
  const BranchState* state = s.get(handle);
  if (!state) {
    throw std::runtime_error("get_logits_at: invalid handle");
  }
  if (!state->has_logits) {
    throw std::runtime_error("get_logits_at: no captured logits for handle");
  }
  const int32_t n_vocab = state->n_vocab;
  const float* logits = state->logits_snapshot.data();
  for (size_t i = 0; i < indices.size(); ++i) {
    const int32_t idx = indices[i];
    // Bounds check: indices arrive from JS Int32Array via N-API — must validate
    // against n_vocab to prevent native OOB reads driven by JS-controlled offsets.
    if (idx < 0 || idx >= n_vocab) {
      throw std::runtime_error("get_logits_at: index out of range");
    }
    out[i] = logits[idx];
  }
}
 
inline llama_token sample(BranchHandle handle, BranchStore& s) {
 
 
  BranchState* state = s.get(handle);
  llama_sampler* chain = state ? s.get_sampler_chain(state->sampler_chain) : nullptr;
  if (!state || !chain) {
    return -1;
  }
 
  // Must have logits captured before sampling
  if (!state->has_logits) {
    LLOYAL_LOG_DEBUG("[branch::sample] No logits captured - call prefill()/step() first");
    return -1;
  }
 
  llama_sampler* gram = s.get_grammar_sampler(state->grammar);
 
  // The candidate array is rebuilt from the branch's own logits snapshot for every pass, so a pass that
  // decided nothing (a refused pick) costs one more fill, never another decode.
  auto fill = [&]() -> llama_token_data_array {
    for (int i = 0; i < state->n_vocab; i++) {
      state->candidates_buffer[i] = llama_token_data{
          static_cast<llama_token>(i),
          state->logits_snapshot[i],
          0.0f};
    }
    return llama_token_data_array{
        state->candidates_buffer.data(),
        static_cast<size_t>(state->n_vocab),
        -1,
        false};
  };
 
  // One pass of the modifiers after the (optional) grammar mask: Logit Bias → Steer → Sampler Chain.
  auto shape = [&](llama_token_data_array& cur_p, bool mask_with_grammar) -> llama_token {
    if (mask_with_grammar && gram) {
      grammar::apply(gram, &cur_p);
    }
 
    // Logit bias — O(n_biases) via direct index (candidates are in token-ID order; grammar::apply
    // preserves order)
    if (!state->logit_bias.empty()) {
      for (const auto& bias : state->logit_bias) {
        if (bias.token >= 0 && bias.token < state->n_vocab) {
          cur_p.data[bias.token].logit += bias.bias;
        }
      }
    }
 
    if (state->steer_fn) {
      try {
        state->steer_fn(cur_p);
      } catch (const std::exception& e) {
        LLOYAL_LOG_DEBUG("[branch::sample] Steer exception: %s", e.what());
        // Continue sampling without steer on exception
      }
    }
 
    sampler::apply(chain, &cur_p);
    return cur_p.selected == -1 ? -1 : cur_p.data[cur_p.selected].id;
  };
 
  // Does the grammar admit this one token? One-element apply: the grammar decodes one piece instead of
  // the whole vocabulary, and apply() reads the parser state without advancing it.
  auto admits = [&](llama_token token) -> bool {
    llama_token_data single{token, 0.0f, 0.0f};
    llama_token_data_array one{&single, 1, -1, false};
    grammar::apply(gram, &one);
    return one.data[0].logit != -INFINITY;
  };
 
  llama_token_data_array cur_p = fill();
  const bool sample_first = gram && state->grammar_sample_first;
  llama_token token = shape(cur_p, /*mask_with_grammar=*/!sample_first);
 
  if (sample_first && token != -1 && !admits(token)) {
    // Refused: take the exact order for this token.
    cur_p = fill();
    token = shape(cur_p, /*mask_with_grammar=*/true);
  }
  ++state->revision;  // the chain ran: its PRNG may have advanced
 
  if (token == -1) {
    return -1;
  }
 
  // Capture filtered candidates for sampling metrics: the distribution the token was actually drawn from.
  state->last_token = token;
  state->last_candidates.clear();
  state->last_candidates.reserve(cur_p.size);
 
  for (size_t i = 0; i < cur_p.size; i++) {
    state->last_candidates.push_back(cur_p.data[i]);
  }
 
  return token;
}
 
inline void accept_token(
    BranchHandle handle,
    llama_token token,
    BranchStore& s) {
 
 
  BranchState* state = s.get(handle);
  if (!state) return;
  ++state->revision;
 
  // Accept in grammar (via anti-corruption layer)
  llama_sampler* gram = s.get_grammar_sampler(state->grammar);
  if (gram) {
    grammar::accept(gram, token);
  }
 
  // Accept in sampler chain for penalty tracking (via anti-corruption layer)
  llama_sampler* chain = s.get_sampler_chain(state->sampler_chain);
  if (chain) {
    sampler::accept(chain, token);
  }
 
  // Update model-level perplexity (from raw logits)
  // Guard on has_logits to avoid computing surprisal from zero-filled buffer
  if (state->metrics != 0 && state->has_logits) {
    float ms = metrics::model_surprisal(
        state->logits_snapshot.data(), state->n_vocab, token);
    s.add_model_surprisal(state->metrics, ms);
  }
 
  // Update sampling-level perplexity (from filtered candidates)
  if (state->metrics != 0 && !state->last_candidates.empty() &&
      token == state->last_token) {
    // Extract filtered logits and IDs from candidates
    std::vector<float> candidate_logits;
    std::vector<int32_t> candidate_ids;
    candidate_logits.reserve(state->last_candidates.size());
    candidate_ids.reserve(state->last_candidates.size());
 
    for (const auto& cand : state->last_candidates) {
      candidate_logits.push_back(cand.logit);
      candidate_ids.push_back(cand.id);
    }
 
    // Compute sampling-level surprisal
    float ss = metrics::sampling_surprisal(
        candidate_logits.data(),
        candidate_ids.data(),
        static_cast<int>(candidate_logits.size()),
        token
    );
    s.add_sampling_surprisal(state->metrics, ss);
  }
}
 
inline void apply_grammar(
    BranchHandle handle,
    float* logits,
    int n_vocab,
    BranchStore& s) {
 
 
  BranchState* state = s.get(handle);
  llama_sampler* gram = state ? s.get_grammar_sampler(state->grammar) : nullptr;
  if (!state || !gram) return;
 
  // Use pre-allocated candidates buffer if size matches, otherwise allocate
  std::vector<llama_token_data>* candidates_ptr;
  std::vector<llama_token_data> temp_buffer;
 
  if (n_vocab == state->n_vocab && !state->candidates_buffer.empty()) {
    candidates_ptr = &state->candidates_buffer;
  } else {
    temp_buffer.resize(n_vocab);
    candidates_ptr = &temp_buffer;
  }
 
  auto& candidates = *candidates_ptr;
  for (int i = 0; i < n_vocab; i++) {
    candidates[i] = llama_token_data{
        static_cast<llama_token>(i), logits[i], 0.0f};
  }
 
  llama_token_data_array cur_p = {
      candidates.data(),
      static_cast<size_t>(n_vocab),
      -1,
      false};
 
  grammar::apply(gram, &cur_p);
 
  // Copy masked logits back
  for (int i = 0; i < n_vocab; i++) {
    logits[i] = candidates[i].logit;
  }
}
 
inline std::vector<std::pair<llama_token, float>> get_legal_priors(
    BranchHandle handle,
    BranchStore& s) {
 
 
  BranchState* state = s.get(handle);
  if (!state || !state->has_logits) {
    return {};
  }
 
  // Reuse pre-allocated candidates buffer
  for (int i = 0; i < state->n_vocab; i++) {
    state->candidates_buffer[i] = llama_token_data{
        static_cast<llama_token>(i),
        state->logits_snapshot[i],
        0.0f};
  }
 
  llama_token_data_array cur_p = {
      state->candidates_buffer.data(),
      static_cast<size_t>(state->n_vocab),
      -1,
      false};
 
  // Apply grammar to mask illegal tokens
  llama_sampler* gram = s.get_grammar_sampler(state->grammar);
  if (gram) {
    grammar::apply(gram, &cur_p);
  }
 
  // Collect legal candidates (logit is finite after grammar masking)
  // Grammar masking sets illegal tokens to -INFINITY
  std::vector<std::pair<llama_token, float>> legal_priors;
  float max_logit = -std::numeric_limits<float>::infinity();
 
  for (size_t i = 0; i < cur_p.size; i++) {
    if (std::isfinite(cur_p.data[i].logit)) {  // Not masked
      legal_priors.emplace_back(cur_p.data[i].id, cur_p.data[i].logit);
      if (cur_p.data[i].logit > max_logit) {
        max_logit = cur_p.data[i].logit;
      }
    }
  }
 
  if (legal_priors.empty()) {
    return {};
  }
 
  // Compute softmax over legal moves only (numerically stable)
  float sum_exp = 0.0f;
  for (auto& [token, logit] : legal_priors) {
    float exp_val = std::exp(logit - max_logit);
    logit = exp_val;  // Temporarily store exp value
    sum_exp += exp_val;
  }
 
  // Normalize to probabilities
  for (auto& [token, prob] : legal_priors) {
    prob /= sum_exp;
  }
 
  return legal_priors;
}
 
inline float get_legal_logsumexp(BranchHandle handle, BranchStore& s) {
 
 
  BranchState* state = s.get(handle);
  if (!state || !state->has_logits) {
    return -std::numeric_limits<float>::infinity();
  }
 
  // Reuse pre-allocated candidates buffer
  for (int i = 0; i < state->n_vocab; i++) {
    state->candidates_buffer[i] = llama_token_data{
        static_cast<llama_token>(i),
        state->logits_snapshot[i],
        0.0f};
  }
 
  llama_token_data_array cur_p = {
      state->candidates_buffer.data(),
      static_cast<size_t>(state->n_vocab),
      -1,
      false};
 
  // Apply grammar to mask illegal tokens
  llama_sampler* gram = s.get_grammar_sampler(state->grammar);
  if (gram) {
    grammar::apply(gram, &cur_p);
  }
 
  // Numerically stable logsumexp over legal tokens (finite logits only)
  float max_logit = -std::numeric_limits<float>::infinity();
  for (size_t i = 0; i < cur_p.size; i++) {
    if (std::isfinite(cur_p.data[i].logit) && cur_p.data[i].logit > max_logit) {
      max_logit = cur_p.data[i].logit;
    }
  }
 
  if (!std::isfinite(max_logit)) {
    return -std::numeric_limits<float>::infinity();  // No legal tokens
  }
 
  float sum_exp = 0.0f;
  for (size_t i = 0; i < cur_p.size; i++) {
    if (std::isfinite(cur_p.data[i].logit)) {
      sum_exp += std::exp(cur_p.data[i].logit - max_logit);
    }
  }
 
  return max_logit + std::log(sum_exp);
}
 
inline bool is_token_legal(
    BranchHandle handle,
    llama_token token,
    BranchStore& s) {
 
 
  BranchState* state = s.get(handle);
  if (!state || token < 0 || token >= state->n_vocab) {
    return false;
  }
 
  // No grammar = all tokens legal
  llama_sampler* gram = s.get_grammar_sampler(state->grammar);
  if (!gram) {
    return true;
  }
 
  // Build 1-element candidate array (stack allocated, no heap)
  llama_token_data single_candidate = {
      token,
      state->has_logits ? state->logits_snapshot[token] : 0.0f,
      0.0f
  };
 
  llama_token_data_array cur_p = {
      &single_candidate,
      1,
      -1,
      false
  };
 
  // Apply grammar - will set logit to -INFINITY if illegal
  grammar::apply(gram, &cur_p);
 
  return std::isfinite(single_candidate.logit);
}
 
inline float get_token_prior_assume_legal(
    BranchHandle handle,
    llama_token token,
    float logsumexp,
    BranchStore& s) {
 
 
  BranchState* state = s.get(handle);
  if (!state || !state->has_logits || token < 0 || token >= state->n_vocab) {
    return 0.0f;
  }
 
  float logit = state->logits_snapshot[token];
  return std::exp(logit - logsumexp);
}
 
inline float get_token_prior(
    BranchHandle handle,
    llama_token token,
    float logsumexp,
    BranchStore& s) {
  if (!is_token_legal(handle, token, s)) {
    return 0.0f;
  }
  return get_token_prior_assume_legal(handle, token, logsumexp, s);
}
 
// ===== STATE ACCESSORS =====
 
 
inline llama_pos get_position(BranchHandle handle, BranchStore& s) {
 
  const BranchState* state = s.get(handle);
  return state ? state->position : -1;
}
 
inline llama_pos get_fork_head(BranchHandle handle, BranchStore& s) {
  const BranchState* state = s.get(handle);
  return state ? state->fork_head : 0;
}
 
inline float get_perplexity(BranchHandle handle, BranchStore& s) {
 
  const BranchState* state = s.get(handle);
  if (!state || state->metrics == 0) {
    return std::numeric_limits<float>::infinity();
  }
  return s.get_model_ppl(state->metrics);
}
 
inline float get_sampling_perplexity(BranchHandle handle, BranchStore& s) {
 
  const BranchState* state = s.get(handle);
  if (!state || state->metrics == 0) {
    return std::numeric_limits<float>::infinity();
  }
  return s.get_sampling_ppl(state->metrics);
}
 
inline float get_last_sampling_prior(BranchHandle handle, BranchStore& s) {
 
  const BranchState* state = s.get(handle);
 
  if (!state || state->last_candidates.empty() || state->last_token < 0) {
    return 0.0f;
  }
 
  // Extract candidates
  std::vector<float> candidate_logits;
  std::vector<int32_t> candidate_ids;
  candidate_logits.reserve(state->last_candidates.size());
  candidate_ids.reserve(state->last_candidates.size());
 
  for (const auto& cand : state->last_candidates) {
    candidate_logits.push_back(cand.logit);
    candidate_ids.push_back(cand.id);
  }
 
  // Compute surprisal from filtered distribution
  float surprisal = metrics::sampling_surprisal(
      candidate_logits.data(),
      candidate_ids.data(),
      static_cast<int>(candidate_logits.size()),
      state->last_token
  );
 
  // Convert to probability: P = exp(-surprisal)
  return std::exp(-surprisal);
}
 
inline int get_n_vocab(BranchHandle handle, BranchStore& s) {
 
  const BranchState* state = s.get(handle);
  return state ? state->n_vocab : 0;
}
 
 
// ===== RAII WRAPPER =====
 
class Branch {
public:
  Branch() : store_(nullptr), handle_(INVALID_HANDLE) {}
 
  Branch(BranchStore* store, BranchHandle handle)
      : store_(store), handle_(handle) {}
 
  ~Branch() {
    if (handle_ != INVALID_HANDLE && store_) {
      branch::pruneSubtree(handle_, *store_);
    }
  }
 
  Branch(Branch&& other) noexcept
      : store_(other.store_), handle_(other.handle_) {
    other.handle_ = INVALID_HANDLE;
  }
 
  Branch& operator=(Branch&& other) noexcept {
    if (this != &other) {
      if (handle_ != INVALID_HANDLE && store_) {
        branch::pruneSubtree(handle_, *store_);
      }
      store_ = other.store_;
      handle_ = other.handle_;
      other.handle_ = INVALID_HANDLE;
    }
    return *this;
  }
 
  Branch(const Branch&) = delete;
  Branch& operator=(const Branch&) = delete;
 
  template <SamplingParamsLike P>
  static Branch create(
      llama_context* ctx,
      const llama_model* model,
      BranchStore& store,
      llama_pos start_pos,
      const P& params,
      int n_batch = DEFAULT_N_BATCH,
      const char* grammar_str = nullptr,
      boundaries::BoundaryTracker* boundary_tracker = nullptr) {
    BranchHandle h = branch::create(ctx, model, store, start_pos, params, n_batch, grammar_str, boundary_tracker);
    return Branch(&store, h);
  }
 
  Branch fork() {
    BranchHandle h = branch::fork(handle_, *store_);
    return Branch(store_, h);
  }
 
  void prune() {
    branch::prune(handle_, *store_);
    handle_ = INVALID_HANDLE;
  }
 
  void pruneSubtree() {
    branch::pruneSubtree(handle_, *store_);
    handle_ = INVALID_HANDLE;
  }
 
  void force_snapshot_logits() {
    branch::force_snapshot_logits(handle_, *store_);
  }
 
  void prefill(const llama_token* tokens, size_t n) {
    branch::prefill(handle_, tokens, n, *store_);
  }
 
  void step(llama_token token) {
    branch::step(handle_, token, *store_);
  }
 
  const float* logits() const {
    return branch::get_logits(handle_, *store_);
  }
 
  llama_token sample() {
    return branch::sample(handle_, *store_);
  }
 
  void accept(llama_token token) {
    branch::accept_token(handle_, token, *store_);
  }
 
  bool is_eog(llama_token token) const {
    const BranchState* st = store_ ? store_->get(handle_) : nullptr;
    return st && st->model ? tokenizer::is_eog(st->model, token) : false;
  }
 
  template <SamplingParamsLike P>
  void setSamplerParams(const P& params) {
    branch::set_sampler_params(handle_, params, *store_);
  }
 
  void setGrammar(const char* grammar_str) {
    const BranchState* st = store_ ? store_->get(handle_) : nullptr;
    branch::set_grammar(handle_, st ? st->model : nullptr, grammar_str, *store_);
  }
 
  // ===== ACCESSORS =====
 
  llama_pos position() const { return branch::get_position(handle_, *store_); }
  llama_pos forkHead() const { return branch::get_fork_head(handle_, *store_); }
  float perplexity() const { return branch::get_perplexity(handle_, *store_); }
  int n_vocab() const { return branch::get_n_vocab(handle_, *store_); }
  bool valid() const { return handle_ != INVALID_HANDLE; }
  BranchHandle handle() const { return handle_; }
 
  // ===== TOPOLOGY =====
 
  BranchHandle parentHandle() const { return store_ ? store_->parent(handle_) : INVALID_HANDLE; }
  const std::vector<BranchHandle>& childHandles() const {
    static const std::vector<BranchHandle> empty;
    return store_ ? store_->children(handle_) : empty;
  }
  bool isLeaf() const { return store_ ? store_->isLeaf(handle_) : true; }
  bool isActive() const { return store_ ? store_->isActive(handle_) : false; }
 
private:
  BranchStore* store_;
  BranchHandle handle_;
};
 
}  // namespace lloyal::branch