|
| | BranchStore (size_t initial_capacity=16) |
| | Construct a branch store with initial slot capacity.
|
| |
| | ~BranchStore () |
| | Destructor — frees CPU resources.
|
| |
| Allocation | allocate () |
| | Allocate a branch slot + KV lease atomically.
|
| |
| void | release (BranchHandle handle) |
| | Release a branch slot + evict its KV lease.
|
| |
| void | init_tenancy (llama_context *ctx) |
| | Initialize KV tenancy after context creation.
|
| |
| void | drain () |
| | Explicit teardown — evict all leases while context is alive.
|
| |
| void | retainOnly (BranchHandle winner) |
| | Keep only the winner — nuclear KV + CPU cleanup.
|
| |
| void | fast_forward (BranchHandle parent, BranchHandle child) |
| | Fast-forward parent to child: the child's state becomes the parent's.
|
| |
| size_t | available () const |
| | Number of vacant seq_ids available for acquisition.
|
| |
| KvPressure | kv_pressure () const |
| | KV cache pressure snapshot — O(1), no tree walking.
|
| |
| void | add_cells_used (uint32_t n) |
| | Increment cells_used counter (for standalone prefill/step outside BranchStore methods)
|
| |
| BranchHandle | parent (BranchHandle h) const |
| | Get a branch's parent handle.
|
| |
| llama_pos | fork_head (BranchHandle h) const |
| | Get a branch's fork head (parent position at fork time)
|
| |
| const std::vector< BranchHandle > & | children (BranchHandle h) const |
| | Get a branch's child handles.
|
| |
| bool | isLeaf (BranchHandle h) const |
| | Test whether a branch is a leaf (no children)
|
| |
| bool | isActive (BranchHandle h) const |
| | Test whether a branch holds a KV lease.
|
| |
| BranchState * | get (BranchHandle handle) |
| | Look up branch state by handle.
|
| |
| const BranchState * | get (BranchHandle handle) const |
| | Look up branch state by handle.
|
| |
| template<SamplingParamsLike P> |
| SamplerChainHandle | create_sampler (const P ¶ms, int32_t n_vocab) |
| | Create a sampler chain and register it.
|
| |
| SamplerChainHandle | clone_sampler (SamplerChainHandle h) |
| | Clone a sampler chain (for fork)
|
| |
| void | free_sampler (SamplerChainHandle h) |
| | Free a sampler chain.
|
| |
| llama_sampler * | get_sampler_chain (SamplerChainHandle h) const |
| | Dereference a sampler chain handle (non-owning)
|
| |
| bool | sampler_has_dist (SamplerChainHandle h) const |
| | Check if a sampler chain ends with dist (stochastic) or greedy.
|
| |
| GrammarHandle | create_grammar (const llama_model *model, const char *grammar_str, const char *root="root") |
| | Create a grammar sampler and register it.
|
| |
| 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") |
| | Create a lazy grammar (unconstrained until trigger fires)
|
| |
| GrammarHandle | clone_grammar (GrammarHandle h) |
| | Clone a grammar (for fork)
|
| |
| void | free_grammar (GrammarHandle h) |
| | Free a grammar.
|
| |
| llama_sampler * | get_grammar_sampler (GrammarHandle h) const |
| | Dereference a grammar handle (non-owning)
|
| |
| MetricsHandle | create_metrics () |
| | Create a metrics tracker and register it.
|
| |
| MetricsHandle | clone_metrics (MetricsHandle h) |
| | Clone a metrics tracker (for fork)
|
| |
| void | free_metrics (MetricsHandle h) |
| | Free a metrics tracker.
|
| |
| void | add_model_surprisal (MetricsHandle h, float surprisal) |
| | Add model-level surprisal to a metrics tracker.
|
| |
| void | add_sampling_surprisal (MetricsHandle h, float surprisal) |
| | Add sampling-level surprisal to a metrics tracker.
|
| |
| float | get_model_ppl (MetricsHandle h) const |
| | Get model-level perplexity from a metrics tracker.
|
| |
| float | get_sampling_ppl (MetricsHandle h) const |
| | Get sampling-level perplexity from a metrics tracker.
|
| |
| void | decode_each (std::span< const DecodeEachItem > items) |
| | Decode one token per branch in a single GPU dispatch.
|
| |
| void | decode_scatter (std::span< const DecodeScatterItem > items) |
| | Decode variable token counts per branch with auto-chunking.
|
| |
| 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) |
| | Decode embedding rows into ONE branch's KV (multimodal ingress)
|
| |
| DecodeSegmentsResult | decode_segments (BranchHandle handle, decode::SegmentSource &source, const std::function< void(const decode::SegmentProgress &)> &on_progress={}) |
| | Prefill a heterogeneous segment sequence into ONE branch.
|
| |
| void | merge_logits (BranchHandle dst_handle, std::span< const BranchHandle > expert_handles, float alpha) |
| | Merge experts' logits_snapshot into dst's logits_snapshot.
|
| |
Handle table and batched decode orchestrator for branch management.
Provides two concerns:
Slot management — A pool of BranchState slots addressed by opaque handles with generation counters for ABA prevention. Slot 0 is permanently reserved (handle 0 = INVALID_HANDLE). Auto-grows by doubling up to 65535 slots. Methods: allocate(), release(), get().
Batched decode — Orchestrates multi-branch GPU dispatches that amortize llama_decode() overhead across N branches. Each method validates handles, builds the appropriate decode primitive's input, dispatches, captures logits into per-branch snapshots, and advances positions atomically. Methods: decode_each(), decode_scatter().
Batched decode methods vs free-function decode:
- Warning
- n_seq_max constraint: Each live branch consumes one KV cache sequence ID (llama_seq_id) managed by kv::tenancy. Call init_tenancy(ctx) after context creation to set the ceiling. The hard limit is
llama_n_seq_max(ctx) (typically 256). allocate() acquires both a slot and a lease atomically; release()/drain() return both resources symmetrically.
-
Thread safety: External synchronization required. The store holds no lock: callers must not touch a branch, or the shared context, while native work on it is in flight.
- See also
- branch::create() to initialize a branch in this store
-
branch::fork() to clone a branch into a new sequence
-
branch::prune() / branch::destroy() for teardown
Definition at line 490 of file branch.hpp.
| void lloyal::branch::BranchStore::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 |
|
) |
| |
|
inline |
Decode embedding rows into ONE branch's KV (multimodal ingress)
The branch-level wrapper over decode::embd(): dispatches the rows into the branch's sequence and does the bookkeeping the token paths do for tokens — with the one multimodal difference: cells grow by n_tokens (every row is a KV cell) while position advances by n_pos (max(nx,ny) under M-RoPE; == n_tokens for plain-position models). The gap is tracked as embedding-row slack so release()/retainOnly() recover the true cell count.
Rows are an interior prefix: no logits are captured unless want_logits (the rows-terminal prefill), in which case the final row's logits land in logits_snapshot like a token prefill's would.
Always its own dispatch: a legacy llama_batch carries rows or tokens, not both. A multimodal prefill goes through decode_segments, which places text and rows in order through the rail, with the same books.
- Parameters
-
| handle | Branch to decode into (must be valid + leased) |
| rows | Embedding rows, n_tokens x n_embd_inp floats |
| n_tokens | Number of rows (cells added) |
| n_embd_inp | Row width — llama_model_n_embd_inp(model) |
| n_pos | Position advance (max(nx,ny) under M-RoPE) |
| pos | Section-major positions, n_tokens * n_pos_per_embd |
| n_pos_per_embd | 1 (plain) or 4 (M-RoPE) |
| non_causal | Bracket the decode in non-causal attention |
| want_logits | Capture the final row's logits (rows-terminal) |
- Exceptions
-
| std::runtime_error | if the handle is invalid or decode fails |
Handling a failed prefill
This is NOT atomic. Rows chunk by n_batch — an image larger than the batch is the normal case, not an edge — so a failure on a later chunk leaves earlier ones committed while the counters here never move.
There is no rollback to reach for. Rewinding the committed rows means a partial-range seq_rm, which cannot restore a recurrent carrier (see kv::remove_range's warning). Attempting it rolls attention back, leaves the carrier folded, and looks like it worked — worse than not trying.
In order of preference:
- PREVENT. Exhausting the KV is the likeliest cause and the only one knowable in advance. A source reports its cost before anything decodes — see decode::SegmentSource::cells() — so admit or refuse the prefill against your own budget instead of discovering it midway.
- CONTAIN. If it throws anyway the branch is poisoned: prune it. Its OWN accounting is still consistent for that prune, because neither position nor cells_used_ advanced — release() subtracts exactly what the branch legitimately owned, and whole-sequence eviction reclaims the orphaned rows. Siblings are untouched.
- RECOVER. Replay the content onto a fresh branch. A branch stores its position and never what was decoded into it, so the content comes from you — the same reason replay is the portable correction at all: it re-runs the folds rather than trying to undo them.
- See also
- kv::remove_range() for why rewinding is not an option
-
decode::embd() for the underlying primitive
-
decode_scatter() for the token rail
- Examples
- /home/runner/work/liblloyal/liblloyal/include/lloyal/branch.hpp.
Definition at line 1340 of file branch.hpp.
Prefill a heterogeneous segment sequence into ONE branch.
The branch's side of decode::segments, the rail that places the stream: it starts at the branch's position, lands in the branch's sequence, chunks rows by the branch's n_batch, and moves the books (position, cells, slack, logits) as each segment lands, through the same helpers the token and embedding rails use. Positions come from the branch's own state and never leave it.
Only the FINAL segment's logits survive: every llama_decode resets the output buffer, so an interior segment's snapshot is immediately dead. Every TEXT segment computes logits and is captured, the interior captures simply overwritten. EMBD segments request output explicitly, so only a trailing one does. Either way the branch ends holding the last segment's.
Segments are requested and dispatched strictly in order — see SegmentSource's in-order contract; a source may reuse its row buffer between at() calls. Every segment must be material: an empty one is rejected, because terminality is positional and a trailing empty segment would leave the real last one without logits.
- Warning
- NOT atomic. Once the first segment is dispatched the branch has been mutated, so a later throw — from the source or from a decode — leaves it partially advanced. The branch is poisoned at that point: prune it and rebuild rather than continuing from it. A DecodeError from any segment after the first carries
partial = true for that reason, whichever of its calls failed — the prefill is one operation to the caller.
- Parameters
-
| handle | Branch to prefill (must be valid + leased) |
| source | Yields the segments; owns production, not placement |
| on_progress | Told, as each segment of an input lands, how much of that input is in the KV and of how much (SegmentProgress). The prompt's own text reports nothing. Never called after a failure. |
- Returns
- Cells added and position advance (they differ under M-RoPE)
- Exceptions
-
| std::runtime_error | if the handle is invalid or a decode fails |
- See also
- decode::segments() for the rail
- Examples
- /home/runner/work/liblloyal/liblloyal/include/lloyal/branch.hpp.
Definition at line 1430 of file branch.hpp.
Fast-forward parent to child: the child's state becomes the parent's.
git's merge --ff-only. With fork() and prune() it gives any branch a checkpoint it can keep or discard while the rest of the tree runs on: fork, work on the fork, then fast-forward the parent to it or prune it.
The parent keeps its identity (handle, parent, children, fork_head, n_batch, steer_fn) and takes everything that follows the child's token trajectory: its sequence and KV, position, image slack, logits snapshot, sampler chain and parameters, grammar and grammar order, last sample, metrics and logit bias. A grammar holds its rules and its parser position in one sampler, and the chain holds its parameters with the penalty window and PRNG, so they move whole: configuration set on the child after the fork comes with it. The child's steer_fn is dropped with the child.
No cells move. The parent's former sequence is evicted whole, which is safe on recurrent layers; siblings still sharing its cells keep them. The parent's former sampler, grammar and metrics are freed with the child's slot, and cells_used is unchanged: the child's cells were counted when they were decoded, and the parent's are all shared with the child.
No native work may be in flight on the context while it runs, on these branches or any other: evicting the parent's former sequence is a KV operation on the shared context (external synchronization, as for every BranchStore operation).
- Exceptions
-
| std::runtime_error,changing | nothing, unless the child was forked from this parent, the parent has neither decoded nor changed since, the child has no children and is not poisoned, and both hold a lease |
- Examples
- /home/runner/work/liblloyal/liblloyal/include/lloyal/branch.hpp.
Definition at line 697 of file branch.hpp.
| void lloyal::branch::BranchStore::merge_logits |
( |
BranchHandle |
dst_handle, |
|
|
std::span< const BranchHandle > |
expert_handles, |
|
|
float |
alpha |
|
) |
| |
|
inline |
Merge experts' logits_snapshot into dst's logits_snapshot.
Accumulates additively (in log-probability space): dst.logits_snapshot[t] += alpha * sum_i(experts[i].logits_snapshot[t])
Mutates dst in place. After this call, sample/produce on dst draws from the combined distribution via dst's existing sampler chain.
Backend-agnostic: operates purely on cached logits buffers, no KV ops, no GPU dispatch. Works identically for transformer and recurrent backends.
Inspired by DExperts-style contrastive decoding (Liu et al., 2021 — https://arxiv.org/abs/2105.03023). Here the experts are different KV states of the SAME model rather than separate models: their live logits_snapshot at each step reflects their own KV history, shaping dst's next-token choice. A negative alpha makes an anti-expert.
- Parameters
-
| dst_handle | Destination branch (must have captured logits) |
| expert_handles | Span of expert branches (all must have captured logits) |
| alpha | Expert weight (additive in log space) |
- Exceptions
-
| std::runtime_error | if any handle invalid, logits not captured, or n_vocab mismatch across branches; a refusal changes nothing |
Definition at line 1504 of file branch.hpp.
1.9.8