MoonBit TreeSpec

A MoonBit library for speculative decoding and reproducible algorithm experiments. It implements sampled single-path verification, deterministic candidate-tree residual verification, context-sensitive tree construction, and matched autoregressive comparisons.

Scope

The library works on token IDs and logits. A caller supplies pure draft/target model callbacks. Included fixtures run offline without model downloads. Both sequential and callback-based batch adapters are available. Batch adapters receive full token contexts and validate row order, result length and vocabulary shape; tensor kernels, tokenizers, trained weights and real KV caches remain outside this build. Reported logical batches are not measured hardware speedups.

Run

Install the MoonBit toolchain, open this module directory, then run:

moon test --deny-warn
moon run cmd/main -- demo
moon run cmd/main -- workload
moon run cmd/main -- tree-plan
moon run cmd/main -- experiment
moon run cmd/main -- batch-experiment
moon run cmd/main -- benchmark
moon run cmd/main -- preflight

experiment performs context-sensitive tree decoding and an autoregressive baseline using the same target callback, prompt, and 32-token output budget. It reports actual target/draft queries separately from hypothetical batches. The demo uses seed 20260905; library users configure seed, output length, width, depth, and node budget with TreeExperimentConfig::new.

batch-experiment exercises the same decoding contract through batch callbacks. Its request counters are actual provider invocations; whether that becomes lower latency depends on the embedding backend’s parallelism.

benchmark runs a matched baseline/adaptive comparison and emits a stable report for a benchmark record. Replace its synthetic callbacks with a pure batch callback to a real model backend; TreeSpec reports provider-request work, not wall-clock latency or throughput.

Before calling a real provider, use diagnose_batch_provider(contexts, vocab, provider). It checks the batch contract twice for row count, vocabulary width, finite logits and deterministic repeated results; preflight runs a small offline example of that contract.

Library example

In a consumer package, import clhhhhhh/moonbit-treespec as @treespec. After this module is published, install it with moon add clhhhhhh/moonbit-treespec.

fn main {
  let config = match @treespec.TreeExperimentConfig::new(2, 3, 14, 32, 12345) {
    Ok(value) => value
    Err(_) => abort("invalid configuration")
  }
  let model = fn(context : Array[Int]) -> Result[Array[Double], String] {
    Ok(if context.length() % 2 == 0 { [0.0, 2.0] } else { [2.0, 0.0] })
  }
  match @treespec.compare_tree_decoding([0], model, model, config) {
    Ok(result) => println(result.render())
    Err(_) => abort("decoding failed")
  }
}

This consumer example is compiled and exercised by cmd/library_example.

Algorithms and engineering

• Stable softmax, temperature/top-k/top-p/min-p, token masks and stop matchers.
• Sampled draft proposals with independent random inputs and residual correction.
• Fixed-width/entropy-adaptive trees queried at each branch’s own prefix, including one-call-per-level draft batching.
• End-to-end entropy-adaptive batch decoding: per-parent uncertainty chooses width, while accepted-path rate adjusts the next proposal depth.
• Sequential residual sibling verification, rejected-branch fallback, correct query positions, and vocabulary/topology validation.
• Flat and level-by-level target batch callbacks with duplicate sibling contexts removed before inference.
• Replay fixtures, logical cache-slot transactions and budgeted sessions.
• Independent session snapshots and cancellation of pending reservations.
• Matched experiments with real callback counters and explicit simulation limits.
• Reproducible matched benchmark records for external batch-model callbacks.

See algorithm and proof sketch and architecture. Breaking changes in this correction are listed in CHANGELOG.md.

Verification

moon fmt --check
moon check --deny-warn
moon test --deny-warn
moon build
moon run cmd/library_example

CI is configured for format, check, tests, build, and example execution. An online green run must be verified after pushing.

License and origin

Apache-2.0; see LICENSE. The code is a MoonBit implementation informed by the speculative-sampling literature, including Leviathan et al. (2023). The deterministic candidate-tree residual procedure is explained in our algorithm notes. No upstream implementation or model weights are bundled.