From 8f3979d937c2d575d559979566c14d64dbf7da9f Mon Sep 17 00:00:00 2001 From: Jacob Fu <141651335+FuJacob@users.noreply.github.com> Date: Fri, 17 Jul 2026 23:59:55 -0700 Subject: [PATCH] Simplify inference to one sequence --- .github/copilot-instructions.md | 55 +- .gitignore | 1 + Package.swift | 13 - README.md | 210 +++--- Sources/CotabbyInferenceBench/baseline_a.cpp | 207 ------ Sources/CotabbyInferenceBench/baseline_b.cpp | 195 ----- .../CotabbyInferenceBench/bench_common.cpp | 73 -- Sources/CotabbyInferenceBench/bench_common.h | 80 -- Sources/CotabbyInferenceBench/main.cpp | 129 ---- .../CotabbyInferenceEngine.cpp | 684 +++--------------- .../include/CotabbyInferenceEngine.h | 91 +-- .../LlamaMiddlewareTests.swift | 613 ++++------------ 12 files changed, 405 insertions(+), 1946 deletions(-) delete mode 100644 Sources/CotabbyInferenceBench/baseline_a.cpp delete mode 100644 Sources/CotabbyInferenceBench/baseline_b.cpp delete mode 100644 Sources/CotabbyInferenceBench/bench_common.cpp delete mode 100644 Sources/CotabbyInferenceBench/bench_common.h delete mode 100644 Sources/CotabbyInferenceBench/main.cpp diff --git a/.github/copilot-instructions.md b/.github/copilot-instructions.md index 194a1d9..ac4cdb8 100644 --- a/.github/copilot-instructions.md +++ b/.github/copilot-instructions.md @@ -1,30 +1,41 @@ - +# CotabbyInference Coding Instructions -# RocketRide — AI Pipeline Builder +CotabbyInference is Cotabby's product-specific C++ wrapper around llama.cpp. Preserve its narrow +single-sequence boundary unless a shipping Cotabby requirement explicitly changes it. -Use RocketRide when building AI pipelines, document processing, RAG systems, or data integration. +## Ownership -## Documentation +- `CotabbyInferenceEngine` owns one model, one context, and at most one active sequence. +- Cotabby's Swift `LlamaRuntimeCore` owns generation serialization, token budgets, prompt/cache + compatibility, streaming, and lifecycle admission. +- The native engine owns sampler state, seed/pending-token handoff, KV mutation, token masks, + cancellation observation, and native resource release. +- Do not add editor, Accessibility, UI, or hosted-service concepts to this package. -Full docs: `.rocketride/docs/` +## Native Safety -**Read the relevant doc(s) before generating any RocketRide code.** +- Keep the public header free of llama.cpp and ggml types. +- Treat `SampleResult.piece` as borrowed memory that expires when sequence storage changes. +- Serialize every llama context mutation through `decode_mutex`. +- Keep cancellation atomic and nonblocking; an active `llama_decode` is not preemptible. +- Never destroy a sequence or unload the model while another non-cancellation operation is using it. +- Preserve the changing external sequence ID even though the internal llama sequence slot is fixed + at zero; it protects replacement sequences from late cancellation. -| File | Read when... | -| --------------------------------- | ----------------------------------------------------------------- | -| ROCKETRIDE_README.md | Starting any RocketRide work — overview + mandatory setup steps | -| ROCKETRIDE_QUICKSTART.md | Writing first pipeline — complete working examples (Python & TS) | -| ROCKETRIDE_PIPELINE_RULES.md | Defining pipelines — structure, lane wiring, config rules | -| ROCKETRIDE_COMPONENT_REFERENCE.md | Choosing/configuring components — all providers and config fields | -| ROCKETRIDE_COMMON_MISTAKES.md | Before finalizing — known pitfalls to avoid | -| ROCKETRIDE_python_API.md | Python SDK — client methods, types, patterns | -| ROCKETRIDE_typescript_API.md | TypeScript SDK — client methods, types, patterns | -| ROCKETRIDE_OBSERVABILITY.md | Consuming runtime logs, lifecycle events, and pipeline traces | +## Change Strategy -## Before Writing ANY RocketRide Code +- Add only APIs used by the current Cotabby integration or backed by a concrete consumer plan. +- Prefer removing obsolete experimental surfaces over retaining speculative generality. +- Treat the llama.cpp binary URL/checksum as a native compatibility boundary. +- Update `README.md`, package tests, and local `.internal/` notes when architecture changes. +- Keep `.internal/` gitignored; it is local study material, not PR content. -1. Read `.rocketride/docs/ROCKETRIDE_README.md` for mandatory setup requirements -2. Read the relevant API doc (Python or TypeScript) for your language -3. Read `.rocketride/docs/ROCKETRIDE_PIPELINE_RULES.md` + `.rocketride/docs/ROCKETRIDE_COMPONENT_REFERENCE.md` -4. Read `.rocketride/docs/ROCKETRIDE_COMMON_MISTAKES.md` before finalizing - +## Validation + +Run `swift test` for compilation and no-model contracts. When a GGUF is available, also run: + +~~~bash +COTABBY_TEST_MODEL_PATH=/absolute/path/model.gguf swift test +~~~ + +For public API changes, build the current Cotabby app against the local package before merging. diff --git a/.gitignore b/.gitignore index 7bfabf4..d2c8b95 100644 --- a/.gitignore +++ b/.gitignore @@ -8,3 +8,4 @@ DerivedData/ .netrc .rocketride/ .env +.internal/ diff --git a/Package.swift b/Package.swift index 2f786df..1827dbc 100644 --- a/Package.swift +++ b/Package.swift @@ -33,18 +33,5 @@ let package = Package( .interoperabilityMode(.Cxx), ] ), - // Phase 0 spike: standalone executable that compares aggregate decode - // tok/s between the current "separate llama_context per sequence" - // architecture and a prospective shared-context-with-batched-decode - // architecture. Used to decide whether the Phase 1 refactor is worth - // doing on M-series + Metal. Not linked into CotabbyInference itself. - .executableTarget( - name: "CotabbyInferenceBench", - dependencies: ["llama-cpp"], - path: "Sources/CotabbyInferenceBench", - cxxSettings: [ - .unsafeFlags(["-std=c++17"]), - ] - ), ] ) diff --git a/README.md b/README.md index bd10941..1b2aca3 100644 --- a/README.md +++ b/README.md @@ -1,81 +1,93 @@ # CotabbyInference -A C++ middleware layer for running LLM inference on-device, built on top of [llama.cpp](https://github.com/ggml-org/llama.cpp). Designed as the inference backend for [Cotabby](https://github.com/fujacob/cotabby), a macOS AI assistant. +CotabbyInference is Cotabby's narrow C++ boundary around +[llama.cpp](https://github.com/ggml-org/llama.cpp). It owns the mapped GGUF model, one llama +context, one autocomplete sequence, sampler state, KV-cache mutation, cancellation, and the +token-level signals consumed by the macOS app. -[![MIT License](https://img.shields.io/badge/license-MIT-blue.svg)](LICENSE) +The package deliberately does not own prompts, request identity, streaming order, normalization, +editor focus, overlays, or insertion. Those remain product responsibilities in Cotabby. -## Why CotabbyInference? +## Current Architecture -llama.cpp is powerful but low-level. CotabbyInference sits on top of it and provides: +~~~text +Cotabby LlamaRuntimeCore + -> CotabbyInferenceEngine + -> one mapped llama_model + -> one llama_context / KV allocation + -> zero or one SequenceState using llama seq_id 0 + -> one sampler chain +~~~ -- **Concurrent sequences** -- run up to 4 independent inference streams at once (e.g. autocomplete + summarization), each with its own context, sampler, and sampling config. No shared decode mutex, no contention. -- **Per-sequence isolation** -- every sequence owns its own `llama_context` and sampler chain. One sequence can be cancelled, trimmed, or destroyed without affecting the others. -- **Thread-safe cancellation** -- cancel any running sequence from any thread via an atomic flag. The next decode or sample call returns immediately with a `cancelled` status. -- **KV cache control** -- trim the KV cache per-sequence to reuse prompt prefixes without re-decoding. Useful for autocomplete where the user keeps typing. -- **Clean Swift interop** -- the public header uses PIMPL to hide all llama.cpp internals. Swift consumers import a single `CotabbyInference` module with no transitive C++ dependencies. The engine class is move-only for `~Copyable` compatibility in Swift 6.2. -- **Zero-config GPU** -- pass `-1` for GPU layers and the engine offloads everything it can to Metal. No manual layer counting. +Cotabby serializes generation and prefill through its runtime lock, so the middleware does not +reserve unused secondary sequence capacity or run a batching worker. Prompt decode, feedback decode, +KV trim, and sequence destruction use one native context mutex. Cancellation is the intentional +cross-thread operation and uses a one-way atomic flag. -## Requirements +The public sequence ID changes whenever a sequence is recreated even though llama's internal slot +is always zero. That prevents a late cancellation from accidentally targeting a replacement +sequence. -- macOS 14+ -- Swift 6.2+ -- Xcode 26+ +## Responsibilities -## Installation - -Add CotabbyInference to your `Package.swift`: - -```swift -dependencies: [ - .package(url: "https://github.com/FuJacob/cotabbyinference.git", from: "0.2.0"), -], -targets: [ - .target( - name: "YourTarget", - dependencies: [ - .product(name: "CotabbyInference", package: "cotabbyinference"), - ], - swiftSettings: [ - .interoperabilityMode(.Cxx), - ] - ), -] -``` - -## Usage - -```swift -import CotabbyInference +- Load and unload one memory-mapped GGUF model. +- Allocate one context with exactly the configured token window. +- Tokenize Cotabby's base-model continuation prompts. +- Build the sampler chain and precompute invalid-token, line-break, and word-boundary masks. +- Decode a prompt and capture its first seed token while the final logits row is live. +- Return one sampled UTF-8 piece at a time. +- Expose sampled EOS, raw-argmax EOG intent, cancellation, and optional log-probability. +- Trim the sequence KV cache for verified prefix reuse. +- Release sampler, context, model, and llama backend resources in order. -var engine = CotabbyInferenceEngine() +## Swift Package -// Load a GGUF model (-1 for all GPU layers, 2048 context, 512 batch) -let status = engine.loadModel("/path/to/model.gguf", -1, 2048, 512) +The package contains: -// Tokenize -let prompt = "The quick brown fox" -let tokens = engine.tokenize(prompt, Int32(prompt.utf8.count)) +- `llama-cpp`: checksum-pinned binary build `b9310`; +- `CotabbyInferenceEngine`: the C++ library imported by Cotabby through Swift C++ interop; +- `CotabbyInferenceTests`: no-model contract tests and optional GGUF-backed integration tests. + +Cotabby currently consumes the package's `main` branch through `project.yml` and records an exact +revision in `Package.resolved`. + +## Basic Usage + +~~~swift +import CotabbyInference + +var engine = CotabbyInferenceEngine() +guard engine.loadModel("/path/to/model.gguf", -1, 2048, 512) == .ok else { + fatalError("Model load failed") +} +defer { engine.unloadModel() } -// Create a sequence with sampling parameters let config = SamplingConfig( - max_prediction_tokens: 64, - temperature: 0.7, - top_k: 40, - top_p: 0.95, - min_p: 0.05, - repetition_penalty: 1.1, - seed: 0 + temperature: 0.1, + top_k: 20, + top_p: 0.7, + min_p: 0.08, + repetition_penalty: 1.05, + seed: 0x00C0_FFEE, + single_line: true ) -let seqId = engine.createSequence(config) -// Decode prompt into KV cache -var tokenArray = Array(tokens) -engine.decodePrompt(seqId, &tokenArray, Int32(tokenArray.count), 0) +let sequenceID = engine.createSequence(config) +guard sequenceID >= 0 else { + fatalError("A sequence is already active or the model is unavailable") +} +defer { engine.destroySequence(sequenceID) } + +let prompt = "The quick brown fox" +var tokens = Array(engine.tokenize(prompt, Int32(prompt.utf8.count))) +guard engine.decodePrompt(sequenceID, &tokens, Int32(tokens.count), 0) == .ok else { + fatalError("Prompt decode failed") +} -// Sample tokens -while true { - let result = engine.sampleNext(seqId) - if result.is_eos { break } +// The caller owns the generation budget. +for _ in 0 ..< 8 { + let result = engine.sampleNext(sequenceID) + if result.is_eos || result.was_cancelled { break } if let piece = result.piece, result.piece_length > 0 { let text = String( @@ -88,48 +100,60 @@ while true { print(text, terminator: "") } } +~~~ -// Cleanup -engine.destroySequence(seqId) -engine.unloadModel() -``` +`SampleResult.piece` is borrowed sequence storage. Copy it before another sampling call or sequence +destruction. -### Running multiple sequences concurrently +## Generation Semantics -Each sequence is fully independent -- different sampling configs, different prompts, different lifetimes: +`decodePrompt` decodes the prompt and immediately samples one seed token. The first `sampleNext` +returns that saved seed without another decode. Each later call feedback-decodes the previously +returned token into KV and samples the next token. -```swift -// Autocomplete: low temperature, short output -let autocompleteConfig = SamplingConfig( - max_prediction_tokens: 8, temperature: 0.1, - top_k: 20, top_p: 0.7, min_p: 0.08, - repetition_penalty: 1.05, seed: 42 -) -let seqA = engine.createSequence(autocompleteConfig) +Cotabby controls the maximum token count in Swift. The engine controls token selection and reports: -// Summary: higher temperature, longer output -let summaryConfig = SamplingConfig( - max_prediction_tokens: 256, temperature: 0.5, - top_k: 40, top_p: 0.95, min_p: 0.05, - repetition_penalty: 1.4, seed: 0 -) -let seqB = engine.createSequence(summaryConfig) +- `is_eos`: the sampled token is an end-of-generation token; +- `was_cancelled`: the native sequence cancellation flag was observed; +- `argmax_is_eog`: the raw model distribution most strongly wanted to stop even if stochastic + sampling selected visible text; +- `logprob`: the selected token's raw-model log-probability when enabled. + +## KV Reuse and Cancellation + +`trimKV` removes a suffix from fixed llama sequence slot zero and invalidates any saved seed or +pending feedback token. Cotabby independently validates request continuity, UTF-8 prefix, token +prefix, and sampling compatibility before calling it. + +`cancelSequence` is thread-safe and nonblocking. Prompt decode checks cancellation between chunks; +sample generation checks before work and after feedback decode. An active llama decode is not +preempted mid-call. Cotabby destroys a natively cancelled sequence because the flag is intentionally +one-way. -// Both run against the same loaded model, with separate contexts. -// Cancel one without affecting the other: -engine.cancelSequence(seqA) -``` +## Testing -## Architecture +Run compile-time and no-model contracts: -CotabbyInference gives each sequence its own `llama_context` and sampler chain. This means: +~~~bash +swift test +~~~ -- No shared decode mutex -- sequences never block each other -- Clean cancellation via per-sequence atomic flags -- Independent KV caches that can be trimmed separately -- Up to 4 concurrent sequences (the memory overhead per context is ~2-4 MB for small models) +Run the full native path with a local GGUF: -The public C++ API uses PIMPL to keep all llama.cpp headers out of the public interface. Swift consumers link against the `CotabbyInference` module only -- no need to deal with `llama.h`, `ggml.h`, or any transitive C dependencies. +~~~bash +COTABBY_TEST_MODEL_PATH=/absolute/path/model.gguf swift test +~~~ + +The model-backed suite covers single-sequence admission/replacement, prompt decode, sampling, KV +trim, cancellation, mid-word continuation, optional log-probability, scaffolding-token masking, and +argmax-EOG behavior. CI does not currently provide a GGUF, so these tests skip there unless the +environment variable is configured. + +## Requirements + +- macOS 14+ +- Swift 6.2+ +- Xcode 26+ ## License diff --git a/Sources/CotabbyInferenceBench/baseline_a.cpp b/Sources/CotabbyInferenceBench/baseline_a.cpp deleted file mode 100644 index 88ad2d1..0000000 --- a/Sources/CotabbyInferenceBench/baseline_a.cpp +++ /dev/null @@ -1,207 +0,0 @@ -#include "bench_common.h" - -#include -#include -#include -#include - -#include - -namespace { - -llama_sampler* make_greedy_sampler() { - auto params = llama_sampler_chain_default_params(); - llama_sampler* chain = llama_sampler_chain_init(params); - if (!chain) return nullptr; - llama_sampler_chain_add(chain, llama_sampler_init_greedy()); - return chain; -} - -llama_context* create_isolated_context( - llama_model* model, - const BenchConfig& cfg -) { - auto p = llama_context_default_params(); - p.n_ctx = static_cast(cfg.ctx_size); - p.n_batch = static_cast(cfg.batch_size); - p.n_ubatch = static_cast(cfg.batch_size); - p.n_seq_max = 1; - int t = static_cast( - std::max(1u, std::thread::hardware_concurrency()) - ); - p.n_threads = t; - p.n_threads_batch = t; - p.offload_kqv = true; - return llama_init_from_model(model, p); -} - -// Walks the prompt in chunks of `batch_cap`, setting logits=1 only on the -// final token of the final chunk. That leaves a single logits row at batch -// index 0 of the last decode, which is where the first sample reads from. -bool decode_prompt( - llama_context* ctx, - const std::vector& tokens, - int batch_cap -) { - llama_batch batch = llama_batch_init(batch_cap, 0, 1); - int cursor = 0; - int n = static_cast(tokens.size()); - while (cursor < n) { - int chunk_end = std::min(cursor + batch_cap, n); - int chunk_size = chunk_end - cursor; - batch.n_tokens = chunk_size; - for (int i = 0; i < chunk_size; ++i) { - int idx = cursor + i; - batch.token[i] = tokens[idx]; - batch.pos[i] = idx; - batch.n_seq_id[i] = 1; - if (batch.seq_id && batch.seq_id[i]) batch.seq_id[i][0] = 0; - bool is_last = (chunk_end == n && i == chunk_size - 1); - batch.logits[i] = is_last ? 1 : 0; - } - if (llama_decode(ctx, batch) != 0) { - llama_batch_free(batch); - return false; - } - cursor = chunk_end; - } - llama_batch_free(batch); - return true; -} - -// Steady-state sampling loop, run once per sequence in its own thread. Takes -// the seed token from warmup as input. Each iteration feedback-decodes the -// previous sample then takes a new one. Mirrors CotabbyInferenceEngine's -// sampleNext pattern minus the per-call batch allocation. -int run_sampling_loop( - llama_context* ctx, - llama_sampler* sampler, - llama_token seed_token, - int start_position, - int sample_count -) { - if (sample_count <= 0) return 0; - - llama_token next = seed_token; - int position = start_position; - int sampled = 0; - - llama_batch batch = llama_batch_init(1, 0, 1); - for (int i = 0; i < sample_count; ++i) { - batch.n_tokens = 1; - batch.token[0] = next; - batch.pos[0] = position; - batch.n_seq_id[0] = 1; - if (batch.seq_id && batch.seq_id[0]) batch.seq_id[0][0] = 0; - batch.logits[0] = 1; - if (llama_decode(ctx, batch) != 0) { - llama_batch_free(batch); - return sampled; - } - position++; - next = llama_sampler_sample(sampler, ctx, -1); - llama_sampler_accept(sampler, next); - sampled++; - } - llama_batch_free(batch); - return sampled; -} - -} // namespace - -BenchResult run_baseline_a_two_contexts( - llama_model* model, - const BenchConfig& cfg -) { - BenchResult r; - r.scenario = "a_two_contexts"; - r.num_sequences = cfg.num_sequences; - r.prompt_tokens = cfg.prompt_tokens; - r.sample_tokens = cfg.sample_tokens; - - const llama_vocab* vocab = llama_model_get_vocab(model); - auto prompt = make_synthetic_prompt(vocab, cfg.prompt_tokens); - - std::vector ctxs(cfg.num_sequences, nullptr); - std::vector samplers(cfg.num_sequences, nullptr); - - auto cleanup = [&]() { - for (auto* s : samplers) if (s) llama_sampler_free(s); - for (auto* c : ctxs) if (c) llama_free(c); - }; - - if (prompt.empty()) { - r.error = "Failed to tokenize synthetic prompt"; - cleanup(); - return r; - } - - // Allocate all resources upfront. Any failure here aborts before any - // timing happens. - for (int i = 0; i < cfg.num_sequences; ++i) { - ctxs[i] = create_isolated_context(model, cfg); - if (!ctxs[i]) { - r.error = "Failed to create context"; - cleanup(); - return r; - } - samplers[i] = make_greedy_sampler(); - if (!samplers[i]) { - r.error = "Failed to create sampler"; - cleanup(); - return r; - } - } - - // Warmup: prompt decode + first sample per context. Both untimed so the - // timed section measures only steady-state decode+sample work, which is - // what the user actually feels in the real app (prompt is mostly cached - // across requests via KV reuse). Pulling the seed sample out also makes - // the count symmetric with baseline B's timed section. - std::vector seed_tokens(cfg.num_sequences, 0); - for (int i = 0; i < cfg.num_sequences; ++i) { - if (!decode_prompt(ctxs[i], prompt, cfg.batch_size)) { - r.error = "Prompt decode failed"; - cleanup(); - return r; - } - // -1 reads from the last logits row, which is where the prompt's - // final-token logits live (we set logits=1 only on that position - // during prompt decode). - seed_tokens[i] = llama_sampler_sample(samplers[i], ctxs[i], -1); - llama_sampler_accept(samplers[i], seed_tokens[i]); - } - - // Timed: one thread per sequence, each runs (sample_tokens - 1) feedback - // decode + sample iterations against its own context. - Timer timer; - timer.start(); - std::vector threads; - std::atomic total{0}; - threads.reserve(cfg.num_sequences); - for (int i = 0; i < cfg.num_sequences; ++i) { - threads.emplace_back([&, i]() { - int n = run_sampling_loop( - ctxs[i], - samplers[i], - seed_tokens[i], - cfg.prompt_tokens, - cfg.sample_tokens - 1 - ); - total.fetch_add(n); - }); - } - for (auto& t : threads) t.join(); - r.elapsed_seconds = timer.elapsed_seconds(); - r.total_tokens_sampled = total.load(); - - if (r.elapsed_seconds > 0.0) { - r.aggregate_tokens_per_second = - r.total_tokens_sampled / r.elapsed_seconds; - r.per_sequence_tokens_per_second = - r.aggregate_tokens_per_second / cfg.num_sequences; - } - - cleanup(); - return r; -} diff --git a/Sources/CotabbyInferenceBench/baseline_b.cpp b/Sources/CotabbyInferenceBench/baseline_b.cpp deleted file mode 100644 index 43f25b9..0000000 --- a/Sources/CotabbyInferenceBench/baseline_b.cpp +++ /dev/null @@ -1,195 +0,0 @@ -#include "bench_common.h" - -#include -#include -#include - -#include - -namespace { - -llama_sampler* make_greedy_sampler() { - auto params = llama_sampler_chain_default_params(); - llama_sampler* chain = llama_sampler_chain_init(params); - if (!chain) return nullptr; - llama_sampler_chain_add(chain, llama_sampler_init_greedy()); - return chain; -} - -// Single shared context with room for `num_sequences` distinct seq_ids. The -// KV cache size is multiplied by `num_sequences` so each sequence still gets -// its configured `ctx_size` of slots — Phase 0 also wants to confirm whether -// n_ctx is shared or per-sequence in the b9310 build. -llama_context* create_shared_context( - llama_model* model, - const BenchConfig& cfg -) { - auto p = llama_context_default_params(); - p.n_ctx = static_cast(cfg.ctx_size * cfg.num_sequences); - p.n_batch = static_cast(cfg.batch_size); - p.n_ubatch = static_cast(cfg.batch_size); - p.n_seq_max = static_cast(cfg.num_sequences); - int t = static_cast( - std::max(1u, std::thread::hardware_concurrency()) - ); - p.n_threads = t; - p.n_threads_batch = t; - p.offload_kqv = true; - return llama_init_from_model(model, p); -} - -// Decodes one sequence's prompt into the shared context, tagged with `seq_id`. -// Only the final token of the final chunk has logits=1, so the next sampler -// call reads from batch index 0 of that decode. -bool decode_prompt_for_seq( - llama_context* ctx, - llama_seq_id seq_id, - const std::vector& tokens, - int batch_cap -) { - llama_batch batch = llama_batch_init(batch_cap, 0, 1); - int cursor = 0; - int n = static_cast(tokens.size()); - while (cursor < n) { - int chunk_end = std::min(cursor + batch_cap, n); - int chunk_size = chunk_end - cursor; - batch.n_tokens = chunk_size; - for (int i = 0; i < chunk_size; ++i) { - int idx = cursor + i; - batch.token[i] = tokens[idx]; - batch.pos[i] = idx; - batch.n_seq_id[i] = 1; - if (batch.seq_id && batch.seq_id[i]) batch.seq_id[i][0] = seq_id; - bool is_last = (chunk_end == n && i == chunk_size - 1); - batch.logits[i] = is_last ? 1 : 0; - } - if (llama_decode(ctx, batch) != 0) { - llama_batch_free(batch); - return false; - } - cursor = chunk_end; - } - llama_batch_free(batch); - return true; -} - -} // namespace - -BenchResult run_baseline_b_shared_context( - llama_model* model, - const BenchConfig& cfg -) { - BenchResult r; - r.scenario = "b_shared_context"; - r.num_sequences = cfg.num_sequences; - r.prompt_tokens = cfg.prompt_tokens; - r.sample_tokens = cfg.sample_tokens; - - const llama_vocab* vocab = llama_model_get_vocab(model); - auto prompt = make_synthetic_prompt(vocab, cfg.prompt_tokens); - - llama_context* ctx = nullptr; - std::vector samplers(cfg.num_sequences, nullptr); - - auto cleanup = [&]() { - for (auto* s : samplers) if (s) llama_sampler_free(s); - if (ctx) llama_free(ctx); - }; - - if (prompt.empty()) { - r.error = "Failed to tokenize synthetic prompt"; - cleanup(); - return r; - } - - ctx = create_shared_context(model, cfg); - if (!ctx) { - r.error = "Failed to create shared context"; - cleanup(); - return r; - } - - for (int s = 0; s < cfg.num_sequences; ++s) { - samplers[s] = make_greedy_sampler(); - if (!samplers[s]) { - r.error = "Failed to create sampler"; - cleanup(); - return r; - } - } - - std::vector last_tokens(cfg.num_sequences, 0); - std::vector positions(cfg.num_sequences, cfg.prompt_tokens); - - // Warmup: decode each prompt and immediately sample one token while that - // sequence's prompt logits are still resident. The next prompt decode for - // sequence s+1 overwrites the logits buffer, so we must sample-before- - // advance. The sampled token becomes the seed input to the timed loop. - for (int s = 0; s < cfg.num_sequences; ++s) { - if (!decode_prompt_for_seq( - ctx, static_cast(s), - prompt, cfg.batch_size)) { - r.error = "Prompt decode failed"; - cleanup(); - return r; - } - // -1 reads from the last logits row, which is where this sequence's - // prompt-final-token logits live until the next decode overwrites - // them. - last_tokens[s] = llama_sampler_sample(samplers[s], ctx, -1); - llama_sampler_accept(samplers[s], last_tokens[s]); - } - - // Timed: build one batch carrying num_sequences tokens (different seq_ids), - // decode all of them in a single llama_decode call, then sample one new - // token per sequence from its respective logit row. Each iteration - // produces one new token per sequence. - llama_batch batch = llama_batch_init(cfg.num_sequences + 4, 0, 1); - Timer timer; - timer.start(); - bool decode_failed = false; - - for (int step = 1; step < cfg.sample_tokens; ++step) { - batch.n_tokens = cfg.num_sequences; - for (int s = 0; s < cfg.num_sequences; ++s) { - batch.token[s] = last_tokens[s]; - batch.pos[s] = positions[s]; - batch.n_seq_id[s] = 1; - if (batch.seq_id && batch.seq_id[s]) { - batch.seq_id[s][0] = static_cast(s); - } - batch.logits[s] = 1; - positions[s]++; - } - if (llama_decode(ctx, batch) != 0) { - decode_failed = true; - break; - } - for (int s = 0; s < cfg.num_sequences; ++s) { - last_tokens[s] = llama_sampler_sample(samplers[s], ctx, s); - llama_sampler_accept(samplers[s], last_tokens[s]); - } - } - - r.elapsed_seconds = timer.elapsed_seconds(); - llama_batch_free(batch); - - if (decode_failed) { - r.error = "llama_decode failed mid-loop"; - cleanup(); - return r; - } - - // Exclude the warmup seed sample from total_tokens_sampled so the - // numerator is comparable to baseline A's timed-only count. - r.total_tokens_sampled = (cfg.sample_tokens - 1) * cfg.num_sequences; - if (r.elapsed_seconds > 0.0) { - r.aggregate_tokens_per_second = - r.total_tokens_sampled / r.elapsed_seconds; - r.per_sequence_tokens_per_second = - r.aggregate_tokens_per_second / cfg.num_sequences; - } - - cleanup(); - return r; -} diff --git a/Sources/CotabbyInferenceBench/bench_common.cpp b/Sources/CotabbyInferenceBench/bench_common.cpp deleted file mode 100644 index dbc9c5e..0000000 --- a/Sources/CotabbyInferenceBench/bench_common.cpp +++ /dev/null @@ -1,73 +0,0 @@ -#include "bench_common.h" - -#include -#include - -#include - -std::vector make_synthetic_prompt( - const llama_vocab* vocab, - int target_tokens -) { - if (!vocab || target_tokens <= 0) return {}; - - // Real text so the tokenizer produces a sensible distribution. The seed is - // short; we pad until tokenization yields at least `target_tokens` tokens. - static const char seed[] = - "The quick brown fox jumps over the lazy dog. "; - std::string text; - while (static_cast(text.size()) < target_tokens * 6) { - text += seed; - } - - bool add_bos = llama_vocab_get_add_bos(vocab); - int capacity = target_tokens * 4 + 16; - std::vector tokens(capacity); - int n = llama_tokenize( - vocab, - text.c_str(), - static_cast(text.size()), - tokens.data(), - static_cast(capacity), - add_bos, - false - ); - if (n <= 0) return {}; - - tokens.resize(n); - if (static_cast(tokens.size()) > target_tokens) { - tokens.resize(target_tokens); - } - return tokens; -} - -void print_result(const BenchResult& r) { - if (!r.error.empty()) { - std::printf( - "{\"scenario\":\"%s\",\"error\":\"%s\"}\n", - r.scenario.c_str(), - r.error.c_str() - ); - return; - } - std::printf( - "{" - "\"scenario\":\"%s\"," - "\"num_sequences\":%d," - "\"prompt_tokens\":%d," - "\"sample_tokens\":%d," - "\"elapsed_seconds\":%.4f," - "\"total_tokens_sampled\":%d," - "\"aggregate_tokens_per_second\":%.2f," - "\"per_sequence_tokens_per_second\":%.2f" - "}\n", - r.scenario.c_str(), - r.num_sequences, - r.prompt_tokens, - r.sample_tokens, - r.elapsed_seconds, - r.total_tokens_sampled, - r.aggregate_tokens_per_second, - r.per_sequence_tokens_per_second - ); -} diff --git a/Sources/CotabbyInferenceBench/bench_common.h b/Sources/CotabbyInferenceBench/bench_common.h deleted file mode 100644 index de1c7bf..0000000 --- a/Sources/CotabbyInferenceBench/bench_common.h +++ /dev/null @@ -1,80 +0,0 @@ -#pragma once - -#include -#include -#include -#include - -// Phase 0 spike: compare aggregate decode tok/s between the current "separate -// llama_context per sequence" architecture and a prospective "single shared -// llama_context with n_seq_max > 1, batched llama_decode" architecture. -// -// The goal is a yes/no decision: on the M-series + Metal hardware our users -// run, does batched decode actually move the needle (>= 1.4x aggregate tok/s) -// or is the GPU command queue serializing everything regardless? -// -// Two baselines, identical warmup, identical token count. Both report total -// samples produced in the timed section and elapsed wall-clock so the caller -// can compute tok/s without us baking in a definition. - -struct BenchConfig { - std::string model_path; - std::string scenario; - int num_sequences = 2; - int prompt_tokens = 256; - int sample_tokens = 200; - int gpu_layers = -1; - int batch_size = 512; - int ctx_size = 2048; - bool verbose = false; -}; - -struct BenchResult { - std::string scenario; - int num_sequences = 0; - int prompt_tokens = 0; - int sample_tokens = 0; - double elapsed_seconds = 0.0; - int total_tokens_sampled = 0; - double aggregate_tokens_per_second = 0.0; - double per_sequence_tokens_per_second = 0.0; - std::string error; -}; - -class Timer { -public: - void start() { t0_ = std::chrono::steady_clock::now(); } - double elapsed_seconds() const { - auto dt = std::chrono::steady_clock::now() - t0_; - return std::chrono::duration(dt).count(); - } -private: - std::chrono::steady_clock::time_point t0_; -}; - -// Pads a short English seed string to a target token count using the supplied -// vocab. Returning real tokens (not random IDs) keeps the sampler in a regime -// the model was trained for, so the decode cost matches what real usage hits. -std::vector make_synthetic_prompt( - const struct llama_vocab* vocab, - int target_tokens -); - -// Emits a single-line JSON record to stdout. One result per invocation; the -// caller is expected to run the binary multiple times and compare lines. -void print_result(const BenchResult& r); - -// Baseline A: N separate llama_context instances, each decoded from its own -// thread. This is what `CotabbyInferenceEngine` does today. -BenchResult run_baseline_a_two_contexts( - struct llama_model* model, - const BenchConfig& cfg -); - -// Baseline B: one shared llama_context with n_seq_max = N, batched -// llama_decode calls carrying tokens for all sequences. This is the candidate -// architecture for Phase 1. -BenchResult run_baseline_b_shared_context( - struct llama_model* model, - const BenchConfig& cfg -); diff --git a/Sources/CotabbyInferenceBench/main.cpp b/Sources/CotabbyInferenceBench/main.cpp deleted file mode 100644 index 5cceef6..0000000 --- a/Sources/CotabbyInferenceBench/main.cpp +++ /dev/null @@ -1,129 +0,0 @@ -#include "bench_common.h" - -#include -#include -#include -#include - -#include - -namespace { - -void silenced_log(ggml_log_level, const char*, void*) {} - -bool parse_args(int argc, char** argv, BenchConfig& cfg) { - for (int i = 1; i < argc; ++i) { - std::string a = argv[i]; - auto next = [&]() -> const char* { - if (i + 1 >= argc) return nullptr; - return argv[++i]; - }; - if (a == "--model") { - const char* v = next(); - if (v) cfg.model_path = v; - } else if (a == "--scenario") { - const char* v = next(); - if (v) cfg.scenario = v; - } else if (a == "--num-sequences") { - const char* v = next(); - if (v) cfg.num_sequences = std::atoi(v); - } else if (a == "--prompt-tokens") { - const char* v = next(); - if (v) cfg.prompt_tokens = std::atoi(v); - } else if (a == "--sample-tokens") { - const char* v = next(); - if (v) cfg.sample_tokens = std::atoi(v); - } else if (a == "--gpu-layers") { - const char* v = next(); - if (v) cfg.gpu_layers = std::atoi(v); - } else if (a == "--batch-size") { - const char* v = next(); - if (v) cfg.batch_size = std::atoi(v); - } else if (a == "--ctx-size") { - const char* v = next(); - if (v) cfg.ctx_size = std::atoi(v); - } else if (a == "--verbose") { - cfg.verbose = true; - } else if (a == "--help" || a == "-h") { - return false; - } - } - return !cfg.model_path.empty() && !cfg.scenario.empty(); -} - -void usage() { - std::fprintf(stderr, - "CotabbyInferenceBench — Phase 0 spike for batched-decode evaluation\n" - "\n" - "Usage:\n" - " CotabbyInferenceBench --model PATH --scenario SCENARIO [options]\n" - "\n" - "Scenarios:\n" - " a_two_contexts N separate llama_context instances, one thread each\n" - " (current CotabbyInferenceEngine architecture)\n" - " b_shared_context One shared llama_context with n_seq_max=N, batched\n" - " llama_decode (candidate Phase 1 architecture)\n" - "\n" - "Options:\n" - " --num-sequences N Number of concurrent sequences (default 2)\n" - " --prompt-tokens N Synthetic prompt length per sequence (default 256)\n" - " --sample-tokens N Tokens to generate per sequence (default 200)\n" - " --gpu-layers N llama n_gpu_layers; -1 means all (default -1)\n" - " --batch-size N Decode batch size (default 512)\n" - " --ctx-size N Per-sequence KV slots (default 2048)\n" - " --verbose Don't silence llama's internal logging\n" - "\n" - "Output: one line of JSON to stdout with elapsed_seconds,\n" - " total_tokens_sampled, aggregate_tokens_per_second.\n" - "\n" - "Both scenarios exclude prompt decode and the seed sample from the\n" - "timed section, so the numerator counts (sample_tokens - 1) *\n" - "num_sequences steady-state samples and is directly comparable.\n" - ); -} - -} // namespace - -int main(int argc, char** argv) { - BenchConfig cfg; - if (!parse_args(argc, argv, cfg)) { - usage(); - return 1; - } - - if (!cfg.verbose) { - llama_log_set(silenced_log, nullptr); - } - llama_backend_init(); - - auto model_params = llama_model_default_params(); - model_params.n_gpu_layers = cfg.gpu_layers; - model_params.use_mmap = true; - model_params.use_mlock = false; - - llama_model* model = llama_model_load_from_file( - cfg.model_path.c_str(), model_params); - if (!model) { - std::fprintf(stderr, "Failed to load model: %s\n", cfg.model_path.c_str()); - llama_backend_free(); - return 2; - } - - BenchResult r; - if (cfg.scenario == "a_two_contexts") { - r = run_baseline_a_two_contexts(model, cfg); - } else if (cfg.scenario == "b_shared_context") { - r = run_baseline_b_shared_context(model, cfg); - } else { - std::fprintf(stderr, "Unknown scenario: %s\n", cfg.scenario.c_str()); - llama_model_free(model); - llama_backend_free(); - return 3; - } - - print_result(r); - - llama_model_free(model); - llama_backend_free(); - return r.error.empty() ? 0 : 4; -} diff --git a/Sources/CotabbyInferenceEngine/CotabbyInferenceEngine.cpp b/Sources/CotabbyInferenceEngine/CotabbyInferenceEngine.cpp index 3bcf259..5fc867f 100644 --- a/Sources/CotabbyInferenceEngine/CotabbyInferenceEngine.cpp +++ b/Sources/CotabbyInferenceEngine/CotabbyInferenceEngine.cpp @@ -2,17 +2,13 @@ #include #include -#include #include -#include -#include -#include +#include #include #include #include #include #include -#include #include #include @@ -84,13 +80,11 @@ static int resolveDecodeThreadCount() { } // --------------------------------------------------------------------------- -// Per-sequence state +// Sequence state // -// Phase 1 architecture: all sequences share a single `llama_context` allocated -// in `Impl`. Each sequence owns its own sampler chain, KV-cache position -// counter, cancellation flag, and detokenization buffer. The `seq_id` is the -// internal `llama_seq_id` slot (0..MAX_SEQUENCES-1) used to tag this -// sequence's tokens in the shared KV cache. +// Cotabby owns one autocomplete stream, so the engine owns at most one live sequence. The public +// ID still changes each time a sequence is created; that prevents a late cancellation from +// targeting a replacement sequence that reuses llama's fixed internal slot 0. // // `seed_token` / `has_seed_token` carries the first sample produced by // `decodePrompt`. We sample it right after the prompt's final decode while @@ -103,9 +97,8 @@ static int resolveDecodeThreadCount() { // --------------------------------------------------------------------------- struct SequenceState { - llama_seq_id seq_id = -1; + int32_t external_id = -1; llama_sampler* sampler = nullptr; - SamplingConfig config{}; int kv_position_count = 0; std::atomic cancelled{false}; std::string last_piece; @@ -135,50 +128,6 @@ struct SequenceState { if (sampler) { llama_sampler_free(sampler); } } - SequenceState() = default; - SequenceState(SequenceState&& o) noexcept - : seq_id(o.seq_id), - sampler(o.sampler), - config(o.config), - kv_position_count(o.kv_position_count), - cancelled(o.cancelled.load()), - last_piece(std::move(o.last_piece)), - seed_token(o.seed_token), - has_seed_token(o.has_seed_token), - pending_input_token(o.pending_input_token), - has_pending_input(o.has_pending_input), - force_word_continuation(o.force_word_continuation), - compute_logprob(o.compute_logprob), - seed_logprob(o.seed_logprob), - seed_argmax_is_eog(o.seed_argmax_is_eog) { - o.sampler = nullptr; - } - SequenceState& operator=(SequenceState&&) = delete; - SequenceState(const SequenceState&) = delete; - SequenceState& operator=(const SequenceState&) = delete; -}; - -// --------------------------------------------------------------------------- -// Pending decode + sample request handed to the decoder thread. -// -// Holds raw pointers into a SequenceState entry. SequenceState entries live in -// a node-based unordered_map (stable addresses across inserts), and the -// public contract forbids destroying a sequence with sampleNext in flight, so -// the pointers stay valid for the request's lifetime. -// --------------------------------------------------------------------------- - -struct PendingRequest { - llama_seq_id seq_id = -1; - llama_token token = 0; - int position = 0; - llama_sampler* sampler = nullptr; - std::atomic* cancelled_ptr = nullptr; - std::string* piece_buffer = nullptr; - SampleResult* result_out = nullptr; - // Snapshot of the sequence's compute_logprob flag at staging time, so the decoder thread - // never has to re-resolve the sequence entry. - bool compute_logprob = true; - std::promise done; }; // --------------------------------------------------------------------------- @@ -186,29 +135,7 @@ struct PendingRequest { // --------------------------------------------------------------------------- struct CotabbyInferenceEngine::Impl { - // Concurrent sequence slots. The shared context's KV allocation scales linearly with this - // (n_ctx = context_window_tokens * MAX_SEQUENCES), so every unused slot is resident RAM — - // hundreds of MB at a 2048-token window on multi-GB models. The app holds at most ONE live - // sequence (it destroys the old autocomplete sequence before building a fresh one); 2 keeps - // one spare slot for a concurrent secondary consumer (e.g. an eval or test driving two - // sequences side by side) without paying for two more that nothing has ever used. - static constexpr int MAX_SEQUENCES = 2; - - // Microseconds the decoder thread waits after the first request arrives - // before flushing. This is the knob that lets multi-sequence callers pile - // up tokens for a batched `llama_decode`. Too short → no batching; too - // long → single-sequence callers feel extra latency per token. 200µs was - // chosen as a starting point and should be tuned via the bench. - static constexpr int BATCH_WINDOW_MICROS = 200; - - // Lock-free mirror of `sequences.size()` for the decoder thread's flush - // decision. The decoder holds `decode_mutex` at that point, and taking - // `sequences_mutex` inside it would create a lock-order hazard with - // `destroySequence` (which nests decode_mutex inside sequences scope); - // an atomic mirror avoids nesting entirely. Updated under - // `sequences_mutex` at every map mutation, so it can lag a concurrent - // create/destroy by at most one flush decision. - std::atomic live_sequence_count{0}; + static constexpr llama_seq_id SEQUENCE_ID = 0; llama_model* model = nullptr; const llama_vocab* vocab = nullptr; @@ -227,55 +154,26 @@ struct CotabbyInferenceEngine::Impl { std::vector nonprintable_bias; std::vector linebreak_bias; std::vector starts_new_word; - // How many of the nonprintable entries came from the scaffolding-piece rule rather than - // the control/unknown/unused attributes. Surfaced via getMaskedScaffoldingTokenCount so - // tests and diagnostics can confirm the rule's reach on a given vocabulary. - int scaffolding_masked_count = 0; - - // Public-facing sequence map (external int32_t IDs → state) and the - // internal `llama_seq_id` slot allocator. - mutable std::mutex sequences_mutex; - std::unordered_map sequences; + + // One product sequence with a monotonically changing external identity. The mutex protects + // create/destroy and lookup; callers still must not destroy the sequence while another method + // is using the returned state pointer. + mutable std::mutex sequence_mutex; + std::unique_ptr sequence; int32_t next_external_id = 1; - bool seq_slot_in_use[MAX_SEQUENCES] = {false}; - // Decoder thread. Owns all `llama_decode` calls on `shared_ctx` after - // model load, including both sample-step batches (built from - // `sampleNext` requests) and prompt-decode chunks (forwarded from - // `decodePrompt` via the same queue path). + // Serializes every llama context mutation. Cotabby's Swift wrapper already orders generation, + // prefill, trim, and reset, while cancellation only touches the sequence's atomic flag. std::mutex decode_mutex; - std::condition_variable request_cv; - std::vector pending; - std::thread decoder_thread; - bool decoder_should_stop = false; - bool decoder_running = false; - - int allocateSeqSlot() { - for (int i = 0; i < MAX_SEQUENCES; ++i) { - if (!seq_slot_in_use[i]) { - seq_slot_in_use[i] = true; - return i; - } - } - return -1; - } - - void releaseSeqSlot(int slot) { - if (slot >= 0 && slot < MAX_SEQUENCES) { - seq_slot_in_use[slot] = false; - } - } SequenceState* findSequence(int32_t id) { - std::lock_guard lock(sequences_mutex); - auto it = sequences.find(id); - return it != sequences.end() ? &it->second : nullptr; + std::lock_guard lock(sequence_mutex); + return sequence && sequence->external_id == id ? sequence.get() : nullptr; } const SequenceState* findSequence(int32_t id) const { - std::lock_guard lock(sequences_mutex); - auto it = sequences.find(id); - return it != sequences.end() ? &it->second : nullptr; + std::lock_guard lock(sequence_mutex); + return sequence && sequence->external_id == id ? sequence.get() : nullptr; } llama_sampler* buildSampler(const SamplingConfig& cfg) const { @@ -286,7 +184,7 @@ struct CotabbyInferenceEngine::Impl { // Quality mask: control/unknown/unused tokens can never be sampled as visible text, and // for single-line fields line-break tokens are masked too. Placed first so the -inf bias // is absolute regardless of the temperature/top-k stages that follow. EOG is intentionally - // left sampleable so the stop check in processBatch/sampleNext still fires. + // left sampleable so the stop check in sampleNext still fires. std::vector mask = nonprintable_bias; if (cfg.single_line && !linebreak_bias.empty()) { mask.insert(mask.end(), linebreak_bias.begin(), linebreak_bias.end()); @@ -348,7 +246,6 @@ struct CotabbyInferenceEngine::Impl { nonprintable_bias.clear(); linebreak_bias.clear(); starts_new_word.clear(); - scaffolding_masked_count = 0; if (!vocab) return; const int32_t n = llama_vocab_n_tokens(vocab); @@ -378,7 +275,6 @@ struct CotabbyInferenceEngine::Impl { if (special_written > 0 && isScaffoldingMarkerPiece(piece, special_written)) { masked = true; - ++scaffolding_masked_count; } } if (masked) { @@ -460,156 +356,9 @@ struct CotabbyInferenceEngine::Impl { return llama_vocab_is_eog(vocab, argmax) || argmax == llama_vocab_eos(vocab); } - void destroyAllSequences() { - std::lock_guard lock(sequences_mutex); - for (auto& [id, seq] : sequences) { - releaseSeqSlot(seq.seq_id); - } - sequences.clear(); - live_sequence_count.store(0, std::memory_order_release); - } - - void startDecoderThread() { - if (decoder_running) return; - decoder_should_stop = false; - decoder_running = true; - decoder_thread = std::thread([this]() { decoderRun(); }); - } - - void stopDecoderThread() { - if (!decoder_running) return; - { - std::lock_guard lock(decode_mutex); - decoder_should_stop = true; - } - request_cv.notify_all(); - if (decoder_thread.joinable()) { - decoder_thread.join(); - } - decoder_running = false; - } - - // Decoder loop: collect pending sample-step requests, batch them into one - // llama_decode, sample each sequence's next token using its own sampler, - // and resolve every request's promise so the caller threads can return. - void decoderRun() { - while (true) { - std::vector batch; - { - std::unique_lock lock(decode_mutex); - request_cv.wait(lock, [&] { - return decoder_should_stop || !pending.empty(); - }); - if (decoder_should_stop && pending.empty()) return; - - // Brief flush window: when only one request has arrived, - // wait a short period for siblings to pile in so the next - // llama_decode can batch them together. Multi-sequence - // workloads naturally fall into lockstep here because each - // sequence resubmits as soon as its previous sample returns. - // With at most one live sequence there is no sibling to wait - // for, so the window is pure added latency on every sampled - // token (the autocomplete app holds exactly one sequence); - // skip it and flush immediately. - if (live_sequence_count.load(std::memory_order_acquire) > 1) { - request_cv.wait_for( - lock, - std::chrono::microseconds(BATCH_WINDOW_MICROS), - [&] { return decoder_should_stop; } - ); - } - - if (pending.empty()) { - if (decoder_should_stop) return; - continue; - } - - batch = std::move(pending); - pending.clear(); - - // Process while still holding decode_mutex so prompt-decode - // calls in `decodePrompt` cannot race with the sample-step - // llama_decode below. llama_decode + sampling for a small - // batch is ~10ms; callers blocked on staging will simply - // queue for the next cycle. - processBatch(batch); - } - } - } - - void processBatch(std::vector& reqs) { - if (reqs.empty() || !shared_ctx) return; - - llama_batch batch = llama_batch_init( - static_cast(reqs.size() + 4), 0, 1 - ); - batch.n_tokens = static_cast(reqs.size()); - for (int i = 0; i < static_cast(reqs.size()); ++i) { - batch.token[i] = reqs[i].token; - batch.pos[i] = static_cast(reqs[i].position); - batch.n_seq_id[i] = 1; - if (batch.seq_id && batch.seq_id[i]) { - batch.seq_id[i][0] = reqs[i].seq_id; - } - batch.logits[i] = 1; - } - - int status = llama_decode(shared_ctx, batch); - - for (int i = 0; i < static_cast(reqs.size()); ++i) { - PendingRequest& req = reqs[i]; - SampleResult r{}; - r.token = 0; - r.piece = nullptr; - r.piece_length = 0; - r.is_eos = false; - r.was_cancelled = false; - - if (status != 0) { - r.is_eos = true; - } else if (req.cancelled_ptr && - req.cancelled_ptr->load(std::memory_order_acquire)) { - r.was_cancelled = true; - } else { - llama_token next = llama_sampler_sample( - req.sampler, shared_ctx, i - ); - // Computed on the raw logits row after sampling: llama_sampler_sample applies - // the chain to a copied candidate array, so row i is still the model's - // unbiased distribution here. - r.argmax_is_eog = argmaxIsEOG(i); - if (next == llama_vocab_eos(vocab) || - llama_vocab_is_eog(vocab, next)) { - r.token = next; - r.is_eos = true; - } else { - llama_sampler_accept(req.sampler, next); - - std::string& piece = *req.piece_buffer; - piece.resize(64); - while (true) { - int written = llama_token_to_piece( - vocab, next, piece.data(), - static_cast(piece.size()), 0, false - ); - if (written >= 0) { piece.resize(written); break; } - piece.resize(static_cast(-written) + 1); - } - - r.token = next; - r.piece = piece.c_str(); - r.piece_length = static_cast(piece.size()); - // Two O(vocab) passes per token — skip entirely when the caller's confidence - // gate is off and the value would be discarded. - r.logprob = req.compute_logprob ? computeLogprob(i, next) : 0.0f; - } - } - - *req.result_out = r; - req.done.set_value(); - } - - llama_batch_free(batch); + void destroySequenceState() { + std::lock_guard lock(sequence_mutex); + sequence.reset(); } }; @@ -680,18 +429,13 @@ EngineStatus CotabbyInferenceEngine::loadModel(const char* path, int gpu_layers, // burns extra package power for nothing when layers are Metal-offloaded anyway. impl_->thread_count = resolveDecodeThreadCount(); - // Shared context sized to hold MAX_SEQUENCES sequences each with up to - // `context_window_tokens` KV slots. llama.cpp's `n_ctx` is the total slot - // budget across all sequences in a context, so we multiply to give each - // sequence the configured window without sequences stealing slots from - // each other. + // Cotabby has one autocomplete sequence, so the configured context window is the complete KV + // budget. Reserving a second unused llama sequence used to double this allocation. auto ctx_params = llama_context_default_params(); - ctx_params.n_ctx = static_cast( - context_window_tokens * Impl::MAX_SEQUENCES - ); + ctx_params.n_ctx = static_cast(context_window_tokens); ctx_params.n_batch = static_cast(batch_size); ctx_params.n_ubatch = static_cast(batch_size); - ctx_params.n_seq_max = static_cast(Impl::MAX_SEQUENCES); + ctx_params.n_seq_max = 1; ctx_params.n_threads = static_cast(impl_->thread_count); ctx_params.n_threads_batch = static_cast(impl_->thread_count); ctx_params.offload_kqv = true; @@ -708,15 +452,13 @@ EngineStatus CotabbyInferenceEngine::loadModel(const char* path, int gpu_layers, // createSequence read these, and the hot path then needs no per-token tokenizer work. impl_->buildTokenMasks(); - impl_->startDecoderThread(); return EngineStatus::ok; } void CotabbyInferenceEngine::unloadModel() { if (!impl_) return; - impl_->stopDecoderThread(); - impl_->destroyAllSequences(); + impl_->destroySequenceState(); if (impl_->shared_ctx) { llama_free(impl_->shared_ctx); @@ -736,10 +478,6 @@ void CotabbyInferenceEngine::unloadModel() { } } -bool CotabbyInferenceEngine::isModelLoaded() const { - return impl_ && impl_->model != nullptr && impl_->shared_ctx != nullptr; -} - // --------------------------------------------------------------------------- // Sequence lifecycle // --------------------------------------------------------------------------- @@ -747,62 +485,36 @@ bool CotabbyInferenceEngine::isModelLoaded() const { int32_t CotabbyInferenceEngine::createSequence(SamplingConfig config) { if (!impl_->model || !impl_->shared_ctx) return -1; - std::lock_guard lock(impl_->sequences_mutex); - int slot = impl_->allocateSeqSlot(); - if (slot < 0) return -1; + std::lock_guard lock(impl_->sequence_mutex); + if (impl_->sequence) return -1; llama_sampler* sampler = impl_->buildSampler(config); - if (!sampler) { - impl_->releaseSeqSlot(slot); - return -1; - } - - SequenceState state; - state.seq_id = static_cast(slot); - state.sampler = sampler; - state.config = config; + if (!sampler) return -1; int32_t id = impl_->next_external_id++; - impl_->sequences.emplace(id, std::move(state)); - impl_->live_sequence_count.store( - static_cast(impl_->sequences.size()), std::memory_order_release); + auto state = std::make_unique(); + state->external_id = id; + state->sampler = sampler; + impl_->sequence = std::move(state); return id; } void CotabbyInferenceEngine::destroySequence(int32_t sequence_id) { if (!impl_) return; - // Look up the internal slot once. Caller's contract is to not destroy a - // sequence with sampleNext in flight, so the entry is stable for the - // duration of this call. - llama_seq_id slot_to_wipe = -1; - { - std::lock_guard seq_lock(impl_->sequences_mutex); - auto it = impl_->sequences.find(sequence_id); - if (it == impl_->sequences.end()) return; - slot_to_wipe = it->second.seq_id; - } + std::lock_guard sequence_lock(impl_->sequence_mutex); + if (!impl_->sequence || impl_->sequence->external_id != sequence_id) return; - // Wipe this sequence's KV slots in the shared context before releasing - // the slot, otherwise stale positions linger and reusing the slot later - // would mix old tokens with new ones. Hold decode_mutex so the wipe does - // not race with the decoder thread's llama_decode call. - if (impl_->shared_ctx && slot_to_wipe >= 0) { + // Wipe slot 0 before a replacement sequence can reuse it. The caller must not destroy a + // sequence while decode/sample is using it; decode_mutex protects the llama context mutation. + if (impl_->shared_ctx) { std::lock_guard decode_lock(impl_->decode_mutex); llama_memory_t memory = llama_get_memory(impl_->shared_ctx); if (memory) { - llama_memory_seq_rm(memory, slot_to_wipe, 0, -1); + llama_memory_seq_rm(memory, Impl::SEQUENCE_ID, 0, -1); } } - - std::lock_guard seq_lock(impl_->sequences_mutex); - auto it = impl_->sequences.find(sequence_id); - if (it != impl_->sequences.end()) { - impl_->releaseSeqSlot(it->second.seq_id); - impl_->sequences.erase(it); - impl_->live_sequence_count.store( - static_cast(impl_->sequences.size()), std::memory_order_release); - } + impl_->sequence.reset(); } // --------------------------------------------------------------------------- @@ -811,19 +523,11 @@ void CotabbyInferenceEngine::destroySequence(int32_t sequence_id) { std::vector CotabbyInferenceEngine::tokenize(const char* text, int text_length) const { - // Preserve the historical contract: add BOS per model metadata, treat any - // special-token text as plaintext (parse_special = false). - bool add_bos = impl_->vocab ? llama_vocab_get_add_bos(impl_->vocab) : false; - return tokenizeWithOptions(text, text_length, add_bos, false); -} - -std::vector CotabbyInferenceEngine::tokenizeWithOptions( - const char* text, int text_length, - bool add_special, bool parse_special) const { if (!impl_->vocab || !text || text_length <= 0) { return {}; } + const bool add_bos = llama_vocab_get_add_bos(impl_->vocab); int capacity = text_length + 8; while (true) { @@ -834,8 +538,8 @@ std::vector CotabbyInferenceEngine::tokenizeWithOptions( static_cast(text_length), tokens.data(), static_cast(capacity), - add_special, - parse_special + add_bos, + false ); if (n > 0) { @@ -849,80 +553,11 @@ std::vector CotabbyInferenceEngine::tokenizeWithOptions( } } -bool CotabbyInferenceEngine::hasChatTemplate() const { - if (!impl_->model) { - return false; - } - return llama_model_chat_template(impl_->model, /*name=*/nullptr) != nullptr; -} - -int CotabbyInferenceEngine::applyChatTemplate( - const char* system_text, - const char* user_text, - bool add_assistant, - char* buffer, - int buffer_size) const { - if (!impl_->model || !system_text || !user_text || - !buffer || buffer_size <= 0) { - return 0; - } - - const char* tmpl = llama_model_chat_template(impl_->model, /*name=*/nullptr); - if (!tmpl) { - return 0; - } - - // Borrowed `const char*` from the caller; valid for this call's duration. - llama_chat_message chat[2] = { - { "system", system_text }, - { "user", user_text } - }; - - int32_t n = llama_chat_apply_template( - tmpl, - chat, - 2, - add_assistant, - buffer, - static_cast(buffer_size) - ); - - // Contract of llama_chat_apply_template: returns the total byte length of - // the formatted prompt; negative means the template is unsupported by - // llama.cpp's predefined list. A positive value larger than the buffer - // means the output did not fit and the caller must retry with a bigger - // buffer. Map all three onto this function's documented C-ABI contract. - if (n < 0) { - return 0; // genuine render failure → caller falls back to raw - } - if (n > buffer_size) { - return -n; // too small → -(required size); caller resizes and retries - } - return n; // success: n bytes written (n <= buffer_size) -} - -int CotabbyInferenceEngine::detokenize(int32_t token, char* buffer, - int buffer_size) const { - if (!impl_->vocab || !buffer || buffer_size <= 0) return 0; - - int written = llama_token_to_piece( - impl_->vocab, - token, - buffer, - static_cast(buffer_size), - 0, - false - ); - - return written; -} - // --------------------------------------------------------------------------- // Prompt decoding // -// Prompt decode runs synchronously on the calling thread, but takes -// `decode_mutex` so it serializes with the decoder thread's sample-step -// llama_decode calls. After the prompt's final decode succeeds, we +// Prompt decode runs synchronously on the calling thread under `decode_mutex`. After the prompt's +// final decode succeeds, we // immediately sample one "seed" token using this sequence's sampler while // the prompt's logits are still live in the shared context. That seed is // handed back via the very next `sampleNext` call without any further @@ -970,7 +605,7 @@ EngineStatus CotabbyInferenceEngine::decodePrompt(int32_t sequence_id, batch.pos[i] = static_cast(start_position + token_index); batch.n_seq_id[i] = 1; if (batch.seq_id && batch.seq_id[i]) { - batch.seq_id[i][0] = seq->seq_id; + batch.seq_id[i][0] = Impl::SEQUENCE_ID; } bool is_last = (chunk_end == end && i == chunk_size - 1); batch.logits[i] = is_last ? 1 : 0; @@ -995,10 +630,8 @@ EngineStatus CotabbyInferenceEngine::decodePrompt(int32_t sequence_id, seq->force_word_continuation = false; } - // Seed sample: take one token from the prompt's logits row right now, - // before any other sequence's decode can overwrite the shared logits - // buffer. The seed will be returned by the next sampleNext call as-is - // and feedback-decoded by the call after that. + // Seed sample: take one token from the prompt's final logits row. The seed will be returned by + // the next sampleNext call as-is and feedback-decoded by the call after that. llama_token seed = llama_sampler_sample(seq->sampler, impl_->shared_ctx, -1); llama_sampler_accept(seq->sampler, seed); seq->seed_token = seed; @@ -1013,9 +646,8 @@ EngineStatus CotabbyInferenceEngine::decodePrompt(int32_t sequence_id, // --------------------------------------------------------------------------- // Sampling // -// First call after decodePrompt: return the seed token directly (no decode -// queued). Subsequent calls: queue the previously sampled token for feedback -// decode via the decoder thread, then return its sampled result. +// First call after decodePrompt returns the seed token directly. Subsequent calls synchronously +// feedback-decode the previously sampled token, then sample and return the next one. // --------------------------------------------------------------------------- SampleResult CotabbyInferenceEngine::sampleNext(int32_t sequence_id) { @@ -1067,8 +699,8 @@ SampleResult CotabbyInferenceEngine::sampleNext(int32_t sequence_id) { seq->last_piece.resize(static_cast(-written) + 1); } - // The seed has not yet been added to KV. Queue it as the next - // feedback-decode input so the call after this one has fresh logits. + // The seed has not yet been added to KV. Remember it as the next feedback-decode input so + // the call after this one can produce fresh logits. seq->pending_input_token = next; seq->has_pending_input = true; @@ -1086,127 +718,67 @@ SampleResult CotabbyInferenceEngine::sampleNext(int32_t sequence_id) { return result; } - // Steady-state path: queue the previously sampled token (or the seed, - // if this is the call right after the seed was delivered) for feedback - // decode via the decoder thread, which batches it with any other - // sequences that happen to be sampling at the same time. - PendingRequest req; - req.seq_id = seq->seq_id; - req.token = seq->pending_input_token; - req.position = seq->kv_position_count; - req.sampler = seq->sampler; - req.cancelled_ptr = &seq->cancelled; - req.piece_buffer = &seq->last_piece; - req.result_out = &result; - req.compute_logprob = seq->compute_logprob; - auto done_future = req.done.get_future(); - - { - std::lock_guard lock(impl_->decode_mutex); - impl_->pending.push_back(std::move(req)); - impl_->request_cv.notify_one(); - } - - done_future.wait(); - - if (result.is_eos || result.was_cancelled) { - return result; - } - - // Feedback decode advanced KV by one position; record the just-sampled - // token as input for the next call. - seq->kv_position_count++; - seq->pending_input_token = result.token; - seq->has_pending_input = true; - return result; -} - -// --------------------------------------------------------------------------- -// Constrained generation primitives -// -// These let a Swift caller run the select-then-commit loop manually instead of -// using sampleNext: read the logits row, classify/choose a token under its own -// constraints, then commit the choice with acceptToken. The vocab queries are -// pure reads on the loaded vocab; getNextTokenLogits copies the live logits row -// (still resident from the last decode); acceptToken mirrors the KV-advancing -// decode used elsewhere so the shared context produces fresh logits afterward. -// --------------------------------------------------------------------------- - -int CotabbyInferenceEngine::getVocabSize() const { - // No vocab means no model loaded; report 0 so callers don't size buffers off a stale value. - if (!impl_ || !impl_->vocab) return 0; - return llama_vocab_n_tokens(impl_->vocab); -} - -bool CotabbyInferenceEngine::isEndOfGenerationToken(int32_t token) const { - if (!impl_ || !impl_->vocab) return false; - return llama_vocab_is_eog(impl_->vocab, token); -} - -int32_t CotabbyInferenceEngine::endOfSequenceToken() const { - // -1 is not a valid token id, so it doubles as the "no model" sentinel. - if (!impl_ || !impl_->vocab) return -1; - return llama_vocab_eos(impl_->vocab); -} - -int CotabbyInferenceEngine::getNextTokenLogits(int32_t sequence_id, - float* out, int out_capacity) const { - if (!impl_ || !impl_->vocab || !impl_->shared_ctx || !out) return 0; - - const SequenceState* seq = impl_->findSequence(sequence_id); - if (!seq) return 0; - - // Refuse to write past the caller's buffer; the full vocab-size row is all-or-nothing. - const int32_t n = llama_vocab_n_tokens(impl_->vocab); - if (n <= 0 || out_capacity < n) return 0; - - // -1 is the most recent decode's logits row, which is what the caller wants to inspect - // before choosing the next token. Null means no live logits (e.g. nothing decoded yet). - const float* logits = llama_get_logits_ith(impl_->shared_ctx, -1); - if (!logits) return 0; - - std::memcpy(out, logits, static_cast(n) * sizeof(float)); - return n; -} - -EngineStatus CotabbyInferenceEngine::acceptToken(int32_t sequence_id, int32_t token) { - if (!impl_ || !impl_->model || !impl_->shared_ctx) return EngineStatus::not_loaded; - - SequenceState* seq = impl_->findSequence(sequence_id); - if (!seq) return EngineStatus::error; - - if (seq->cancelled.load(std::memory_order_acquire)) { - return EngineStatus::cancelled; - } - - // Feed the chosen token to the sampler so repetition/penalty state matches what sampleNext - // would have produced; the caller selected the token externally but the sampler must still see it. - llama_sampler_accept(seq->sampler, token); - - // Serialize with the decoder thread so this manual decode never races an in-flight batch. + // Cotabby has no sibling sequence to batch with, so feedback decode happens directly on the + // calling worker under the same context lock used by prompt decode and KV mutation. std::lock_guard lock(impl_->decode_mutex); - - // Single-token feedback decode, mirroring the KV-advancing step in sampleNext/processBatch: - // pos = current KV position, request logits so a fresh row is ready for getNextTokenLogits. llama_batch batch = llama_batch_init(1, 0, 1); batch.n_tokens = 1; - batch.token[0] = token; + batch.token[0] = seq->pending_input_token; batch.pos[0] = static_cast(seq->kv_position_count); batch.n_seq_id[0] = 1; if (batch.seq_id && batch.seq_id[0]) { - batch.seq_id[0][0] = seq->seq_id; + batch.seq_id[0][0] = Impl::SEQUENCE_ID; } batch.logits[0] = 1; - int status = llama_decode(impl_->shared_ctx, batch); + const int status = llama_decode(impl_->shared_ctx, batch); + if (status != 0) { + result.is_eos = true; + } else if (seq->cancelled.load(std::memory_order_acquire)) { + result.was_cancelled = true; + } else { + const llama_token next = llama_sampler_sample(seq->sampler, impl_->shared_ctx, 0); + result.argmax_is_eog = impl_->argmaxIsEOG(0); + result.token = next; + + if (next == llama_vocab_eos(impl_->vocab) || + llama_vocab_is_eog(impl_->vocab, next)) { + result.is_eos = true; + } else { + llama_sampler_accept(seq->sampler, next); + seq->last_piece.resize(64); + while (true) { + const int written = llama_token_to_piece( + impl_->vocab, + next, + seq->last_piece.data(), + static_cast(seq->last_piece.size()), + 0, + false + ); + if (written >= 0) { + seq->last_piece.resize(written); + break; + } + seq->last_piece.resize(static_cast(-written) + 1); + } + result.piece = seq->last_piece.c_str(); + result.piece_length = static_cast(seq->last_piece.size()); + result.logprob = seq->compute_logprob ? impl_->computeLogprob(0, next) : 0.0f; + } + } llama_batch_free(batch); - if (status != 0) { - return EngineStatus::error; + if (result.is_eos || result.was_cancelled) { + return result; } + // Feedback decode advanced KV by one position; record the just-sampled + // token as input for the next call. seq->kv_position_count++; - return EngineStatus::ok; + seq->pending_input_token = result.token; + seq->has_pending_input = true; + return result; } // --------------------------------------------------------------------------- @@ -1221,13 +793,12 @@ bool CotabbyInferenceEngine::trimKV(int32_t sequence_id, int keep_positions) { llama_memory_t memory = llama_get_memory(impl_->shared_ctx); if (!memory) return false; - // Serialize with the decoder thread; we don't want to remove KV slots - // mid-batch. + // Serialize with prompt and feedback decode; never remove KV while llama is mutating it. std::lock_guard lock(impl_->decode_mutex); bool ok = llama_memory_seq_rm( memory, - seq->seq_id, + Impl::SEQUENCE_ID, static_cast(keep_positions), -1 ); @@ -1243,11 +814,6 @@ bool CotabbyInferenceEngine::trimKV(int32_t sequence_id, int keep_positions) { return ok; } -int CotabbyInferenceEngine::getKVPositionCount(int32_t sequence_id) const { - const SequenceState* seq = impl_->findSequence(sequence_id); - return seq ? seq->kv_position_count : 0; -} - void CotabbyInferenceEngine::setForceWordContinuation(int32_t sequence_id, bool enabled) { if (!impl_) return; SequenceState* seq = impl_->findSequence(sequence_id); @@ -1264,46 +830,6 @@ void CotabbyInferenceEngine::setComputeLogprob(int32_t sequence_id, bool enabled } } -// --------------------------------------------------------------------------- -// KV state snapshot / restore -// -// Thin wrappers over llama's single-sequence state copy. They serialize with the decoder thread -// via decode_mutex so a snapshot or restore never races an in-flight llama_decode. Callers pair a -// snapshot with the KV position count they observed and pass it back on restore, so this engine's -// own position bookkeeping stays consistent with the restored KV cache. -// --------------------------------------------------------------------------- - -size_t CotabbyInferenceEngine::snapshotSize(int32_t sequence_id) const { - if (!impl_ || !impl_->shared_ctx) return 0; - const SequenceState* seq = impl_->findSequence(sequence_id); - if (!seq) return 0; - return llama_state_seq_get_size(impl_->shared_ctx, seq->seq_id); -} - -size_t CotabbyInferenceEngine::snapshotSequence(int32_t sequence_id, uint8_t* dst, size_t capacity) { - if (!impl_ || !impl_->shared_ctx || !dst) return 0; - SequenceState* seq = impl_->findSequence(sequence_id); - if (!seq) return 0; - std::lock_guard lock(impl_->decode_mutex); - return llama_state_seq_get_data(impl_->shared_ctx, dst, capacity, seq->seq_id); -} - -bool CotabbyInferenceEngine::restoreSequence(int32_t sequence_id, const uint8_t* src, - size_t size, int position_count) { - if (!impl_ || !impl_->shared_ctx || !src) return false; - SequenceState* seq = impl_->findSequence(sequence_id); - if (!seq) return false; - std::lock_guard lock(impl_->decode_mutex); - const size_t read = llama_state_seq_set_data(impl_->shared_ctx, src, size, seq->seq_id); - if (read == 0) return false; - seq->kv_position_count = position_count; - // The restored blob invalidates any seed/pending token captured before; force the caller to - // re-prime via decodePrompt before the next sampleNext. - seq->has_seed_token = false; - seq->has_pending_input = false; - return true; -} - // --------------------------------------------------------------------------- // Cancellation // --------------------------------------------------------------------------- @@ -1334,7 +860,3 @@ int CotabbyInferenceEngine::getThreadCount() const { int CotabbyInferenceEngine::getGPULayerCount() const { return impl_->gpu_layer_count; } - -int CotabbyInferenceEngine::getMaskedScaffoldingTokenCount() const { - return impl_->scaffolding_masked_count; -} diff --git a/Sources/CotabbyInferenceEngine/include/CotabbyInferenceEngine.h b/Sources/CotabbyInferenceEngine/include/CotabbyInferenceEngine.h index 8da29de..383bcaf 100644 --- a/Sources/CotabbyInferenceEngine/include/CotabbyInferenceEngine.h +++ b/Sources/CotabbyInferenceEngine/include/CotabbyInferenceEngine.h @@ -1,11 +1,9 @@ #pragma once -#include #include #include #include struct SamplingConfig { - int max_prediction_tokens; float temperature; int top_k; float top_p; @@ -57,52 +55,15 @@ class CotabbyInferenceEngine { EngineStatus loadModel(const char* path, int gpu_layers, int context_window_tokens, int batch_size); void unloadModel(); - bool isModelLoaded() const; // Sequence lifecycle + // The engine owns at most one sequence. `createSequence` returns -1 while another sequence is + // alive; destroy the current sequence before creating its replacement. int32_t createSequence(SamplingConfig config); void destroySequence(int32_t sequence_id); // Tokenization (thread-safe, read-only on vocab) std::vector tokenize(const char* text, int text_length) const; - // Like `tokenize`, but the caller controls BOS/EOS injection and whether - // special/control tokens in the text (e.g. chat-template markers like - // <|im_start|>) are recognized as their token IDs instead of plain text. - // The plain `tokenize` keeps `parse_special = false` for backward - // compatibility; the chat-template path needs `true` so rendered markers - // tokenize correctly. - std::vector tokenizeWithOptions(const char* text, int text_length, - bool add_special, - bool parse_special) const; - int detokenize(int32_t token, char* buffer, int buffer_size) const; - - // Chat templates - // - // `hasChatTemplate` reports whether the loaded model ships a chat template - // in its GGUF metadata. Instruct models (Qwen, Gemma, Llama) do; raw base - // models do not. Callers use this to decide between the structured - // chat-template prompt path and the legacy raw-continuation path so a - // user-supplied base model keeps working. - bool hasChatTemplate() const; - // Renders a system + user turn through the model's built-in chat template - // into `buffer`. `add_assistant` appends the assistant-turn opening marker - // so the model continues as the assistant. Returns: - // > 0 : number of bytes written (<= buffer_size) — the formatted prompt. - // < 0 : -(required buffer size); the buffer was too small, retry at that size. - // = 0 : no model, no template, or render failure — caller falls back to raw. - // - // Autocomplete needs exactly one system turn (rules + context) and one user - // turn (the text to continue), so the signature takes those two directly - // rather than a message array. This buffer-based C ABI mirrors `detokenize` - // and deliberately avoids std::string / struct parameter and return types, - // so it bridges cleanly into the Swift objcxx interop mode the app target - // uses (where a std::string return does not bridge). - int applyChatTemplate(const char* system_text, - const char* user_text, - bool add_assistant, - char* buffer, - int buffer_size) const; - // Prompt decoding EngineStatus decodePrompt(int32_t sequence_id, const int32_t* tokens, int token_count, @@ -111,39 +72,8 @@ class CotabbyInferenceEngine { // Sampling SampleResult sampleNext(int32_t sequence_id); - // Constrained generation primitives - // - // These decouple token *selection* from the engine: a Swift caller can read the - // raw next-token logits, classify candidate tokens, pick one under its own - // constraints (grammar, vocabulary subset, etc.), and then commit it via - // `acceptToken` to advance the sequence. This is the manual alternative to - // `sampleNext`, which both selects and commits in one step. - - // Total vocabulary size, i.e. the number of float logits a row from - // `getNextTokenLogits` contains. 0 when no model is loaded. - int getVocabSize() const; - // Whether `token` is an end-of-generation marker (EOS or any other model - // stop token). Callers use this to terminate manual generation loops. - bool isEndOfGenerationToken(int32_t token) const; - // The model's end-of-sequence token id, or -1 when no model is loaded. - int32_t endOfSequenceToken() const; - // Copies the next-token logits row for `sequence_id` (the distribution - // produced by the most recent decode) into `out`. `out_capacity` must be at - // least `getVocabSize()`; exactly that many floats are written. Returns the - // number of floats written, or 0 on any error (null buffer, unknown - // sequence, too-small capacity, or no live logits). - int getNextTokenLogits(int32_t sequence_id, float* out, int out_capacity) const; - // Commits `token` as the chosen next token for `sequence_id`: feeds it to the - // sequence's sampler (so penalty/repetition state stays consistent) and - // feedback-decodes it to advance the KV cache by one position, leaving fresh - // logits at the new position for the next `getNextTokenLogits` call. This is - // the commit half of the manual select-then-commit loop. Guards not_loaded / - // cancelled and returns `error` if the decode fails. - EngineStatus acceptToken(int32_t sequence_id, int32_t token); - // KV cache management bool trimKV(int32_t sequence_id, int keep_positions); - int getKVPositionCount(int32_t sequence_id) const; // Constrains the FIRST token of the next generation on `sequence_id` to continue the current // word: tokens whose decoded text begins with whitespace are masked for that one token, then @@ -153,19 +83,9 @@ class CotabbyInferenceEngine { // Controls whether `SampleResult.logprob` is computed for this sequence. Defaults to true // (the historical behavior). The log-probability costs two O(vocab-size) passes per generated // token, so callers whose confidence gating is disabled should pass false to skip it; results - // then report `logprob == 0`. A method rather than a `SamplingConfig` field so existing - // memberwise `SamplingConfig` initializers in Swift keep compiling unchanged. + // then report `logprob == 0`. void setComputeLogprob(int32_t sequence_id, bool enabled); - // Single-sequence KV state snapshot/restore. `snapshotSize` reports the buffer size needed, - // `snapshotSequence` copies the sequence's KV state into `dst` (returns bytes written, 0 on - // failure), and `restoreSequence` loads a previously captured blob back and sets the KV - // position to `position_count` (returns false on failure). Blobs are model- and context- - // specific; never restore one across a model reload. - size_t snapshotSize(int32_t sequence_id) const; - size_t snapshotSequence(int32_t sequence_id, uint8_t* dst, size_t capacity); - bool restoreSequence(int32_t sequence_id, const uint8_t* src, size_t size, int position_count); - // Cancellation (thread-safe, non-blocking) void cancelSequence(int32_t sequence_id); @@ -174,11 +94,6 @@ class CotabbyInferenceEngine { int getBatchSize() const; int getThreadCount() const; int getGPULayerCount() const; - // Number of vocabulary tokens that were hard-masked at model load because their rendered - // piece is chat/instruct/FIM scaffolding that the GGUF did not flag as a control token. - // 0 is the common (and healthy) case: well-formed GGUFs flag these as control already, - // and control tokens are masked by the base rule rather than counted here. - int getMaskedScaffoldingTokenCount() const; private: struct Impl; diff --git a/Tests/CotabbyInferenceTests/LlamaMiddlewareTests.swift b/Tests/CotabbyInferenceTests/LlamaMiddlewareTests.swift index 53dfaa2..b7741ba 100644 --- a/Tests/CotabbyInferenceTests/LlamaMiddlewareTests.swift +++ b/Tests/CotabbyInferenceTests/LlamaMiddlewareTests.swift @@ -1,605 +1,288 @@ -import XCTest import CotabbyInference +import XCTest final class LlamaMiddlewareTests: XCTestCase { - - func testEngineStartsUnloaded() { - let engine = CotabbyInferenceEngine() - XCTAssertFalse(engine.isModelLoaded()) - } - func testUnloadWhenNothingLoadedIsIdempotent() { - var engine = CotabbyInferenceEngine() // var: unloadModel mutates + var engine = CotabbyInferenceEngine() engine.unloadModel() engine.unloadModel() - XCTAssertFalse(engine.isModelLoaded()) } func testLoadModelWithBadPathReturnsError() { var engine = CotabbyInferenceEngine() - let status = engine.loadModel("/nonexistent/path.gguf", -1, 2048, 512) - XCTAssertEqual(status, EngineStatus.error) - XCTAssertFalse(engine.isModelLoaded()) + XCTAssertEqual( + engine.loadModel("/nonexistent/path.gguf", -1, 2048, 512), + EngineStatus.error + ) } - func testCreateSequenceWithoutModelReturnsMinus1() { + func testCreateSequenceWithoutModelReturnsMinusOne() { var engine = CotabbyInferenceEngine() - let config = SamplingConfig( - max_prediction_tokens: 8, - temperature: 0.1, - top_k: 20, - top_p: 0.7, - min_p: 0.08, - repetition_penalty: 1.05, - seed: 0, - single_line: false - ) - let seqId = engine.createSequence(config) - XCTAssertEqual(seqId, -1) + XCTAssertEqual(engine.createSequence(Self.samplingConfig()), -1) } - func testDestroySequenceWithInvalidIdDoesNotCrash() { + func testInvalidSequenceOperationsDoNotCrash() { var engine = CotabbyInferenceEngine() engine.destroySequence(999) engine.destroySequence(-1) - } - - func testCancelSequenceWithInvalidIdDoesNotCrash() { - var engine = CotabbyInferenceEngine() engine.cancelSequence(999) + engine.setForceWordContinuation(999, true) + engine.setComputeLogprob(999, false) } func testTokenizeWithoutModelReturnsEmpty() { let engine = CotabbyInferenceEngine() let text = "hello" - let tokens = engine.tokenize(text, Int32(text.utf8.count)) - XCTAssertTrue(tokens.isEmpty) - } - - func testTokenizeWithOptionsWithoutModelReturnsEmpty() { - let engine = CotabbyInferenceEngine() - let text = "hello" - let tokens = engine.tokenizeWithOptions( - text, Int32(text.utf8.count), false, true - ) - XCTAssertTrue(tokens.isEmpty) - } - - func testHasChatTemplateWithoutModelIsFalse() { - let engine = CotabbyInferenceEngine() - XCTAssertFalse(engine.hasChatTemplate()) - } - - func testApplyChatTemplateWithoutModelReturnsZero() { - let engine = CotabbyInferenceEngine() - var buffer = [CChar](repeating: 0, count: 256) - let written = engine.applyChatTemplate( - "You complete text.", "The quick brown", true, &buffer, Int32(buffer.count) - ) - // No model loaded → 0 (caller falls back to the raw path). - XCTAssertEqual(written, 0) + XCTAssertTrue(engine.tokenize(text, Int32(text.utf8.count)).isEmpty) } func testDiagnosticsDefaultToZero() { let engine = CotabbyInferenceEngine() XCTAssertEqual(engine.getContextWindowTokens(), 0) XCTAssertEqual(engine.getBatchSize(), 0) + XCTAssertEqual(engine.getThreadCount(), 0) XCTAssertEqual(engine.getGPULayerCount(), 0) } func testDecodePromptWithoutModelReturnsNotLoaded() { var engine = CotabbyInferenceEngine() var tokens: [Int32] = [1, 2, 3] - let status = engine.decodePrompt(1, &tokens, Int32(tokens.count), 0) - XCTAssertEqual(status, EngineStatus.not_loaded) + XCTAssertEqual( + engine.decodePrompt(1, &tokens, Int32(tokens.count), 0), + EngineStatus.not_loaded + ) } - func testEndToEndWithModel() throws { - guard let modelPath = ProcessInfo.processInfo.environment["COTABBY_TEST_MODEL_PATH"], - FileManager.default.fileExists(atPath: modelPath) else { - throw XCTSkip("Set COTABBY_TEST_MODEL_PATH to a .gguf file to run this test") - } - + func testEndToEndSingleSequenceLifecycle() throws { + let modelPath = try Self.modelPath() var engine = CotabbyInferenceEngine() + XCTAssertEqual(engine.loadModel(modelPath, -1, 2048, 512), EngineStatus.ok) + defer { engine.unloadModel() } - // Load - let loadStatus = engine.loadModel(modelPath, -1, 2048, 512) - XCTAssertEqual(loadStatus, EngineStatus.ok) - XCTAssertTrue(engine.isModelLoaded()) XCTAssertEqual(engine.getContextWindowTokens(), 2048) XCTAssertEqual(engine.getBatchSize(), 512) XCTAssertGreaterThan(engine.getThreadCount(), 0) - // Idempotent re-load - let reloadStatus = engine.loadModel(modelPath, -1, 2048, 512) - XCTAssertEqual(reloadStatus, EngineStatus.ok) + // Repeating an identical load remains an idempotent no-op. + XCTAssertEqual(engine.loadModel(modelPath, -1, 2048, 512), EngineStatus.ok) - // Tokenize let prompt = "The quick brown fox" - let tokens = engine.tokenize(prompt, Int32(prompt.utf8.count)) + var tokens = Array(engine.tokenize(prompt, Int32(prompt.utf8.count))) XCTAssertFalse(tokens.isEmpty) - // Chat-template path: instruct models ship a template; if present, - // rendering a simple conversation must produce a non-empty prompt that - // tokenizes (with parse_special) to a non-empty token list. - if engine.hasChatTemplate() { - // Render system + user through the model's template into a caller buffer. - var buffer = [CChar](repeating: 0, count: 4096) - let written = engine.applyChatTemplate( - "You complete text.", "The quick brown", true, &buffer, Int32(buffer.count) - ) - XCTAssertGreaterThan(written, 0, "Model reports a template but rendering produced no bytes") - - let rendered = buffer.prefix(Int(written)).withUnsafeBufferPointer { ptr in - String( - bytes: UnsafeRawBufferPointer(ptr), - encoding: .utf8 - ) - } - let renderedSwift = try XCTUnwrap(rendered, "Rendered template was not valid UTF-8") - XCTAssertFalse(renderedSwift.isEmpty) - - let templated = engine.tokenizeWithOptions( - renderedSwift, Int32(renderedSwift.utf8.count), false, true - ) - XCTAssertFalse(templated.isEmpty) - } - - // Detokenize a content token (the prompt's last token). Index 0 can be BOS, a control - // token that renders to zero bytes with special=false, so we avoid it here. - var buf = [CChar](repeating: 0, count: 64) - let written = engine.detokenize(tokens[tokens.count - 1], &buf, Int32(buf.count)) - XCTAssertGreaterThan(written, 0) - - // Create autocomplete sequence - let autoConfig = SamplingConfig( - max_prediction_tokens: 8, - temperature: 0.1, - top_k: 20, - top_p: 0.7, - min_p: 0.08, - repetition_penalty: 1.05, - seed: 42, - single_line: false + let sequence = engine.createSequence(Self.samplingConfig()) + XCTAssertGreaterThan(sequence, 0) + XCTAssertEqual( + engine.createSequence(Self.samplingConfig(seed: 99)), + -1, + "The engine must reject a second live sequence" ) - let seqA = engine.createSequence(autoConfig) - XCTAssertGreaterThan(seqA, 0) - // Decode prompt - var tokenArray = Array(tokens) - let decodeStatus = engine.decodePrompt( - seqA, &tokenArray, Int32(tokenArray.count), 0 + XCTAssertEqual( + engine.decodePrompt(sequence, &tokens, Int32(tokens.count), 0), + EngineStatus.ok ) - XCTAssertEqual(decodeStatus, EngineStatus.ok) - XCTAssertEqual(engine.getKVPositionCount(seqA), Int32(tokenArray.count)) - // Sample a few tokens var generated = "" - for _ in 0..<4 { - let result = engine.sampleNext(seqA) + for _ in 0 ..< 4 { + let result = engine.sampleNext(sequence) if result.is_eos { break } XCTAssertFalse(result.was_cancelled) - if let piece = result.piece, result.piece_length > 0 { - generated += String( - bytes: UnsafeBufferPointer( - start: UnsafeRawPointer(piece) - .assumingMemoryBound(to: UInt8.self), - count: Int(result.piece_length) - ), - encoding: .utf8 - ) ?? "" - } + generated += Self.string(from: result) } XCTAssertFalse(generated.isEmpty, "Expected at least one generated token") - // Trim KV back to prompt (remove sampled tokens) - let trimOk = engine.trimKV(seqA, Int32(tokenArray.count)) - XCTAssertTrue(trimOk) - XCTAssertEqual(engine.getKVPositionCount(seqA), Int32(tokenArray.count)) - - // Create a second concurrent sequence (summary config) - let summaryConfig = SamplingConfig( - max_prediction_tokens: 60, - temperature: 0.5, - top_k: 40, - top_p: 0.95, - min_p: 0.05, - repetition_penalty: 1.4, - seed: 0, - single_line: false - ) - let seqB = engine.createSequence(summaryConfig) - XCTAssertGreaterThan(seqB, 0) - XCTAssertNotEqual(seqA, seqB) - - // Both sequences exist simultaneously - XCTAssertGreaterThan(engine.getKVPositionCount(seqA), 0) - XCTAssertEqual(engine.getKVPositionCount(seqB), 0) + // Hybrid/recurrent and SWA model caches can reject partial KV removal. Cotabby treats + // that as a cache-reuse miss and rebuilds the sequence, so lifecycle coverage must not + // require a model-specific optimization to succeed. + _ = engine.trimKV(sequence, Int32(tokens.count)) - // Destroy both - engine.destroySequence(seqB) - engine.destroySequence(seqA) + engine.destroySequence(sequence) + let replacement = engine.createSequence(Self.samplingConfig(seed: 100)) + XCTAssertGreaterThan(replacement, 0) + XCTAssertNotEqual(replacement, sequence) + engine.destroySequence(replacement) - // Double-destroy is safe - engine.destroySequence(seqA) - - // Unload - engine.unloadModel() - XCTAssertFalse(engine.isModelLoaded()) + // Stale and repeated destruction must not affect a later sequence identity. + engine.destroySequence(sequence) } - // Multi-sequence sequential test. Drives the new shared-context decoder - // thread through two interleaved sampling loops to verify the seed-token - // / feedback-decode handoff produces valid tokens for both sequences - // when their sampleNext calls alternate. - func testInterleavedMultiSequenceSampling() throws { - guard let modelPath = ProcessInfo.processInfo.environment["COTABBY_TEST_MODEL_PATH"] else { - try XCTSkipIf(true, "Set COTABBY_TEST_MODEL_PATH to a .gguf file to run this test") - return - } - - var engine = CotabbyInferenceEngine() - XCTAssertEqual(engine.loadModel(modelPath, -1, 1024, 256), EngineStatus.ok) - defer { engine.unloadModel() } - - let prompt = "The quick brown fox jumps over the lazy dog." - let tokens = engine.tokenize(prompt, Int32(prompt.utf8.count)) - XCTAssertGreaterThan(tokens.size(), 0) - - let configA = SamplingConfig( - max_prediction_tokens: 16, temperature: 0, - top_k: 0, top_p: 0, min_p: 0, - repetition_penalty: 0, seed: 1, - single_line: false - ) - let configB = SamplingConfig( - max_prediction_tokens: 16, temperature: 0, - top_k: 0, top_p: 0, min_p: 0, - repetition_penalty: 0, seed: 2, - single_line: false - ) - - let seqA = engine.createSequence(configA) - let seqB = engine.createSequence(configB) - XCTAssertGreaterThan(seqA, 0) - XCTAssertGreaterThan(seqB, 0) - - var tokenArr = Array(tokens) - XCTAssertEqual( - engine.decodePrompt(seqA, &tokenArr, Int32(tokenArr.count), 0), - EngineStatus.ok - ) - XCTAssertEqual( - engine.decodePrompt(seqB, &tokenArr, Int32(tokenArr.count), 0), - EngineStatus.ok - ) - - // Alternate sampleNext between the two sequences. With greedy - // sampling and identical prompts, the first sampled tokens for both - // sequences should be identical (different samplers reading the - // same logits row at separate decodePrompt times). - var sampledA: [Int32] = [] - var sampledB: [Int32] = [] - for _ in 0..<8 { - let rA = engine.sampleNext(seqA) - let rB = engine.sampleNext(seqB) - XCTAssertFalse(rA.was_cancelled) - XCTAssertFalse(rB.was_cancelled) - if rA.is_eos || rB.is_eos { break } - sampledA.append(rA.token) - sampledB.append(rB.token) - } - XCTAssertEqual(sampledA.count, sampledB.count) - XCTAssertGreaterThan(sampledA.count, 0) - XCTAssertEqual(sampledA, sampledB, - "Greedy sampling with identical prompts should match across sequences") - - engine.destroySequence(seqA) - engine.destroySequence(seqB) - } - - // Cancellation regression: setting cancelled on a sequence mid-loop must - // cause subsequent sampleNext calls to return was_cancelled=true so - // callers can break out without waiting for the full prediction budget. func testCancellationStopsSamplingPromptly() throws { - guard let modelPath = ProcessInfo.processInfo.environment["COTABBY_TEST_MODEL_PATH"] else { - try XCTSkipIf(true, "Set COTABBY_TEST_MODEL_PATH to a .gguf file to run this test") - return - } - + let modelPath = try Self.modelPath() var engine = CotabbyInferenceEngine() XCTAssertEqual(engine.loadModel(modelPath, -1, 1024, 256), EngineStatus.ok) defer { engine.unloadModel() } + let sequence = engine.createSequence(Self.samplingConfig(temperature: 0)) let prompt = "Hello" - let tokens = engine.tokenize(prompt, Int32(prompt.utf8.count)) - let config = SamplingConfig( - max_prediction_tokens: 32, temperature: 0, - top_k: 0, top_p: 0, min_p: 0, - repetition_penalty: 0, seed: 0, - single_line: false - ) - - let seq = engine.createSequence(config) - var tokenArr = Array(tokens) + var tokens = Array(engine.tokenize(prompt, Int32(prompt.utf8.count))) XCTAssertEqual( - engine.decodePrompt(seq, &tokenArr, Int32(tokenArr.count), 0), + engine.decodePrompt(sequence, &tokens, Int32(tokens.count), 0), EngineStatus.ok ) - // Sample a couple of tokens first. - for _ in 0..<2 { - let r = engine.sampleNext(seq) - XCTAssertFalse(r.was_cancelled) - } - - // Cancel. The next sampleNext should return was_cancelled=true - // without doing any further model work. - engine.cancelSequence(seq) - let cancelled = engine.sampleNext(seq) - XCTAssertTrue(cancelled.was_cancelled, - "sampleNext after cancelSequence must return was_cancelled=true") - - engine.destroySequence(seq) - } - - func testSetForceWordContinuationWithoutModelDoesNotCrash() { - var engine = CotabbyInferenceEngine() - engine.setForceWordContinuation(999, true) - engine.setForceWordContinuation(-1, false) - XCTAssertFalse(engine.isModelLoaded()) - } - - func testSnapshotSizeWithoutModelIsZero() { - let engine = CotabbyInferenceEngine() - XCTAssertEqual(engine.snapshotSize(1), 0) + _ = engine.sampleNext(sequence) + engine.cancelSequence(sequence) + XCTAssertTrue(engine.sampleNext(sequence).was_cancelled) + engine.destroySequence(sequence) } - // With the first-token word-continuation constraint set, the seed token must not start a new - // word, i.e. its decoded text must not begin with whitespace. func testForceWordContinuationConstrainsFirstToken() throws { - guard let modelPath = ProcessInfo.processInfo.environment["COTABBY_TEST_MODEL_PATH"] else { - try XCTSkipIf(true, "Set COTABBY_TEST_MODEL_PATH to a .gguf file to run this test") - return - } + let modelPath = try Self.modelPath() var engine = CotabbyInferenceEngine() XCTAssertEqual(engine.loadModel(modelPath, -1, 1024, 256), EngineStatus.ok) defer { engine.unloadModel() } - let config = SamplingConfig( - max_prediction_tokens: 8, temperature: 0, - top_k: 0, top_p: 0, min_p: 0, - repetition_penalty: 0, seed: 0, - single_line: false - ) - let seq = engine.createSequence(config) - XCTAssertGreaterThan(seq, 0) - - // Prompt ends mid-word ("writ"); the forced continuation must finish the word. + let sequence = engine.createSequence(Self.samplingConfig(temperature: 0)) let prompt = "I am writ" var tokens = Array(engine.tokenize(prompt, Int32(prompt.utf8.count))) - XCTAssertGreaterThan(tokens.count, 0) - - engine.setForceWordContinuation(seq, true) - XCTAssertEqual(engine.decodePrompt(seq, &tokens, Int32(tokens.count), 0), EngineStatus.ok) - - let result = engine.sampleNext(seq) - if !result.is_eos, let piece = result.piece, result.piece_length > 0 { - let text = String( - bytes: UnsafeBufferPointer( - start: UnsafeRawPointer(piece).assumingMemoryBound(to: UInt8.self), - count: Int(result.piece_length) - ), - encoding: .utf8 - ) ?? "" - if let firstChar = text.first { - XCTAssertFalse( - firstChar.isWhitespace, - "Forced word continuation must not begin the first token with whitespace" - ) - } - } - engine.destroySequence(seq) - } - - // Snapshotting a sequence then restoring it must return the engine's KV position bookkeeping - // to the captured value. - func testSnapshotRestorePreservesPosition() throws { - guard let modelPath = ProcessInfo.processInfo.environment["COTABBY_TEST_MODEL_PATH"] else { - try XCTSkipIf(true, "Set COTABBY_TEST_MODEL_PATH to a .gguf file to run this test") - return - } - var engine = CotabbyInferenceEngine() - XCTAssertEqual(engine.loadModel(modelPath, -1, 1024, 256), EngineStatus.ok) - defer { engine.unloadModel() } - - let config = SamplingConfig( - max_prediction_tokens: 8, temperature: 0, - top_k: 0, top_p: 0, min_p: 0, - repetition_penalty: 0, seed: 0, - single_line: false + engine.setForceWordContinuation(sequence, true) + XCTAssertEqual( + engine.decodePrompt(sequence, &tokens, Int32(tokens.count), 0), + EngineStatus.ok ) - let seq = engine.createSequence(config) - let prompt = "The quick brown fox" - var tokens = Array(engine.tokenize(prompt, Int32(prompt.utf8.count))) - XCTAssertEqual(engine.decodePrompt(seq, &tokens, Int32(tokens.count), 0), EngineStatus.ok) - let position = engine.getKVPositionCount(seq) - XCTAssertGreaterThan(position, 0) - - let size = engine.snapshotSize(seq) - XCTAssertGreaterThan(size, 0) - var buffer = [UInt8](repeating: 0, count: Int(size)) - let written = engine.snapshotSequence(seq, &buffer, size) - XCTAssertGreaterThan(written, 0) - - // Advance past the snapshot point, then restore back to it. - _ = engine.sampleNext(seq) - XCTAssertTrue(engine.restoreSequence(seq, buffer, written, position)) - XCTAssertEqual(engine.getKVPositionCount(seq), position) - - engine.destroySequence(seq) + let result = engine.sampleNext(sequence) + if !result.is_eos, let first = Self.string(from: result).first { + XCTAssertFalse(first.isWhitespace) + } + engine.destroySequence(sequence) } - // A sampled (non-EOS) token must carry a finite log-probability that is <= 0, so the app can use - // it as a confidence signal. func testSampleNextReportsFiniteLogprob() throws { - guard let modelPath = ProcessInfo.processInfo.environment["COTABBY_TEST_MODEL_PATH"] else { - try XCTSkipIf(true, "Set COTABBY_TEST_MODEL_PATH to a .gguf file to run this test") - return - } + let modelPath = try Self.modelPath() var engine = CotabbyInferenceEngine() XCTAssertEqual(engine.loadModel(modelPath, -1, 1024, 256), EngineStatus.ok) defer { engine.unloadModel() } - let config = SamplingConfig( - max_prediction_tokens: 4, temperature: 0, - top_k: 0, top_p: 0, min_p: 0, - repetition_penalty: 0, seed: 0, - single_line: false - ) - let seq = engine.createSequence(config) + let sequence = engine.createSequence(Self.samplingConfig(temperature: 0)) let prompt = "The quick brown fox" var tokens = Array(engine.tokenize(prompt, Int32(prompt.utf8.count))) - XCTAssertEqual(engine.decodePrompt(seq, &tokens, Int32(tokens.count), 0), EngineStatus.ok) + XCTAssertEqual( + engine.decodePrompt(sequence, &tokens, Int32(tokens.count), 0), + EngineStatus.ok + ) - let result = engine.sampleNext(seq) + let result = engine.sampleNext(sequence) if !result.is_eos { - XCTAssertTrue(result.logprob.isFinite, "logprob must be finite") - XCTAssertLessThanOrEqual(result.logprob, 0.0001, "a log-probability must be <= 0") + XCTAssertTrue(result.logprob.isFinite) + XCTAssertLessThanOrEqual(result.logprob, 0.0001) } - engine.destroySequence(seq) - } - - func testSetComputeLogprobWithInvalidIdDoesNotCrash() { - var engine = CotabbyInferenceEngine() - engine.setComputeLogprob(999, false) - engine.setComputeLogprob(-1, true) + engine.destroySequence(sequence) } - // Opting out of log-probabilities must zero `logprob` on both the seed token (first sampleNext) - // and the steady-state decoder path, while leaving the sampled tokens themselves untouched. - func testSetComputeLogprobFalseZeroesLogprob() throws { - guard let modelPath = ProcessInfo.processInfo.environment["COTABBY_TEST_MODEL_PATH"] else { - try XCTSkipIf(true, "Set COTABBY_TEST_MODEL_PATH to a .gguf file to run this test") - return - } + func testDisablingLogprobSkipsSeedAndSteadyStateWork() throws { + let modelPath = try Self.modelPath() var engine = CotabbyInferenceEngine() XCTAssertEqual(engine.loadModel(modelPath, -1, 1024, 256), EngineStatus.ok) defer { engine.unloadModel() } - let config = SamplingConfig( - max_prediction_tokens: 4, temperature: 0, - top_k: 0, top_p: 0, min_p: 0, - repetition_penalty: 0, seed: 0, - single_line: false - ) - let seq = engine.createSequence(config) - engine.setComputeLogprob(seq, false) - + let sequence = engine.createSequence(Self.samplingConfig(temperature: 0)) + engine.setComputeLogprob(sequence, false) let prompt = "The quick brown fox" var tokens = Array(engine.tokenize(prompt, Int32(prompt.utf8.count))) - XCTAssertEqual(engine.decodePrompt(seq, &tokens, Int32(tokens.count), 0), EngineStatus.ok) + XCTAssertEqual( + engine.decodePrompt(sequence, &tokens, Int32(tokens.count), 0), + EngineStatus.ok + ) - // Seed token (computed at decodePrompt) and two steady-state tokens. - for _ in 0..<3 { - let result = engine.sampleNext(seq) + for _ in 0 ..< 3 { + let result = engine.sampleNext(sequence) if result.is_eos || result.was_cancelled { break } - XCTAssertEqual(result.logprob, 0.0, "logprob must be exactly 0 when computation is disabled") + XCTAssertEqual(result.logprob, 0) } - engine.destroySequence(seq) - } - - func testMaskedScaffoldingCountDefaultsToZero() { - let engine = CotabbyInferenceEngine() - XCTAssertEqual(engine.getMaskedScaffoldingTokenCount(), 0) + engine.destroySequence(sequence) } - // Single-token chat/template markers must never surface as sampled text. Most vocabularies - // flag them as control tokens (masked by the base rule); the scaffolding rule covers ones - // that ship unflagged. Either way the observable contract is the same: no sampled piece is - // ever a complete marker string, even at high temperature with a template-bait prompt. func testSamplingNeverEmitsScaffoldingMarkerPieces() throws { - guard let modelPath = ProcessInfo.processInfo.environment["COTABBY_TEST_MODEL_PATH"] else { - try XCTSkipIf(true, "Set COTABBY_TEST_MODEL_PATH to a .gguf file to run this test") - return - } + let modelPath = try Self.modelPath() var engine = CotabbyInferenceEngine() XCTAssertEqual(engine.loadModel(modelPath, -1, 1024, 256), EngineStatus.ok) defer { engine.unloadModel() } - XCTAssertGreaterThanOrEqual(engine.getMaskedScaffoldingTokenCount(), 0) let markers: Set = [ "<|im_start|>", "<|im_end|>", "<|user|>", "<|assistant|>", "<|system|>", "<|start_header_id|>", "<|end_header_id|>", "<|eot_id|>", "<|end|>", "<|endoftext|>", "", "", "[INST]", "[/INST]" ] - let config = SamplingConfig( - max_prediction_tokens: 64, temperature: 1.8, - top_k: 0, top_p: 0, min_p: 0, - repetition_penalty: 1.0, seed: 7, - single_line: false - ) - let seq = engine.createSequence(config) + let sequence = engine.createSequence(Self.samplingConfig(temperature: 1.8, seed: 7)) let prompt = "<|im_start|>user\nWrite a reply<|im_end|>\n<|im_start|>assistant\n" var tokens = Array(engine.tokenize(prompt, Int32(prompt.utf8.count))) - XCTAssertEqual(engine.decodePrompt(seq, &tokens, Int32(tokens.count), 0), EngineStatus.ok) + XCTAssertEqual( + engine.decodePrompt(sequence, &tokens, Int32(tokens.count), 0), + EngineStatus.ok + ) - for _ in 0..<64 { - let result = engine.sampleNext(seq) + for _ in 0 ..< 64 { + let result = engine.sampleNext(sequence) if result.is_eos || result.was_cancelled { break } - guard let piecePointer = result.piece else { continue } - let piece = String( - bytes: UnsafeRawBufferPointer(start: piecePointer, count: Int(result.piece_length)), - encoding: .utf8 - ) ?? "" - XCTAssertFalse( - markers.contains(piece), - "Sampled a complete scaffolding marker piece: \(piece)" - ) + XCTAssertFalse(markers.contains(Self.string(from: result))) } - engine.destroySequence(seq) + engine.destroySequence(sequence) } - // Under greedy sampling with penalties disabled, the sampled token IS the raw argmax, so - // `argmax_is_eog` must agree with `isEndOfGenerationToken(sampled token)` on every step. - // This pins the flag's semantics without needing the model to reach a natural EOS. - func testArgmaxIsEOGMatchesGreedyChoice() throws { - guard let modelPath = ProcessInfo.processInfo.environment["COTABBY_TEST_MODEL_PATH"] else { - try XCTSkipIf(true, "Set COTABBY_TEST_MODEL_PATH to a .gguf file to run this test") - return - } + func testArgmaxIsEOGMatchesGreedySample() throws { + let modelPath = try Self.modelPath() var engine = CotabbyInferenceEngine() XCTAssertEqual(engine.loadModel(modelPath, -1, 1024, 256), EngineStatus.ok) defer { engine.unloadModel() } - let config = SamplingConfig( - max_prediction_tokens: 24, temperature: 0, - top_k: 0, top_p: 0, min_p: 0, - repetition_penalty: 1.0, seed: 0, - single_line: false + let sequence = engine.createSequence( + Self.samplingConfig(temperature: 0, repetitionPenalty: 1) ) - let seq = engine.createSequence(config) let prompt = "The capital of France is" var tokens = Array(engine.tokenize(prompt, Int32(prompt.utf8.count))) - XCTAssertEqual(engine.decodePrompt(seq, &tokens, Int32(tokens.count), 0), EngineStatus.ok) + XCTAssertEqual( + engine.decodePrompt(sequence, &tokens, Int32(tokens.count), 0), + EngineStatus.ok + ) var steps = 0 - for _ in 0..<24 { - let result = engine.sampleNext(seq) + for _ in 0 ..< 24 { + let result = engine.sampleNext(sequence) if result.was_cancelled { break } - // Greedy + no penalties: sampled token == raw argmax. A masked control token can in - // principle displace the raw argmax from the greedy pick, but EOG tokens are never - // masked, so the EOG verdicts still agree. - XCTAssertEqual( - result.argmax_is_eog, - engine.isEndOfGenerationToken(result.token), - "argmax_is_eog disagreed with the greedy token's EOG status at step \(steps)" - ) + XCTAssertEqual(result.argmax_is_eog, result.is_eos) steps += 1 if result.is_eos { break } } - XCTAssertGreaterThan(steps, 0, "Expected at least one sampled step") - engine.destroySequence(seq) + XCTAssertGreaterThan(steps, 0) + engine.destroySequence(sequence) + } +} + +private extension LlamaMiddlewareTests { + static func modelPath() throws -> String { + guard let path = ProcessInfo.processInfo.environment["COTABBY_TEST_MODEL_PATH"], + FileManager.default.fileExists(atPath: path) else { + throw XCTSkip("Set COTABBY_TEST_MODEL_PATH to a .gguf file to run model-backed tests") + } + return path + } + + static func samplingConfig( + temperature: Float = 0.1, + repetitionPenalty: Float = 1.05, + seed: UInt32 = 42 + ) -> SamplingConfig { + SamplingConfig( + temperature: temperature, + top_k: 20, + top_p: 0.7, + min_p: 0.08, + repetition_penalty: repetitionPenalty, + seed: seed, + single_line: false + ) + } + + static func string(from result: SampleResult) -> String { + guard let piece = result.piece, result.piece_length > 0 else { return "" } + return String( + bytes: UnsafeBufferPointer( + start: UnsafeRawPointer(piece).assumingMemoryBound(to: UInt8.self), + count: Int(result.piece_length) + ), + encoding: .utf8 + ) ?? "" } }