diff --git a/README.md b/README.md index 0810eb9..f72c4e9 100644 --- a/README.md +++ b/README.md @@ -32,7 +32,7 @@ The engine is a single C file (`c/glm.c`, ~2,400 lines) plus small headers. No B - **MLA attention** (q/kv-LoRA, interleaved partial RoPE) with **compressed KV-cache**: 576 floats/token instead of 32,768 (57× smaller — GLM-5.2 has 64 heads and no GQA). - **DeepSeek-V3-style sigmoid router** (noaux_tc, routed_scaling_factor), shared expert, first-3-dense layers. - **Native MTP speculative decoding** — GLM-5.2's own multi-token-prediction head (layer 78) drafts tokens that the main model verifies in one batched forward. **The head must be int8** (the converter does this by default): at int4 draft acceptance collapses to 0–4% and speculation never engages; at int8 it's 39–59% acceptance, **2.2–2.8 tokens/forward** (community-measured, [#8](https://github.com/JustVugg/colibri/issues/8)). Lossless *in exact arithmetic* — but **not byte-identical to non-speculative greedy in practice** ([#100](https://github.com/JustVugg/colibri/issues/100)). This isn't MTP-specific: colibrì's quantized integer kernels are shape-dependent, so any batched (S>1) or GPU forward rounds slightly differently from the single-token path, and int4 GLM-5.2 sits close enough to argmax ties that such a rounding change can flip a token. MTP, the CUDA expert tier, and batched prefill are three different ways to trip the same sensitivity (community-confirmed in #100: swapping only the kernel family forks greedy output on 3/5 prompts, with **zero speculation**). Every emitted token is still the argmax of a *valid* forward — the continuation stays correct — it just isn't the same stream. For byte-exact reproducibility: `DRAFT=0` (no speculation), plus `IDOT=0 COLI_CUDA=0` if you also want kernel-family/GPU independence. Under sampling, rejection sampling keeps the distribution correct. Honest caveat from the same measurement: on a **cold** cache each verified draft routes to extra experts (~660 → ~1100 expert-loads/token), so speculation can be a net *time* loss until the cache/pin warms up. -- **Grammar-forced speculative drafts** (`GRAMMAR=file.gbnf`, [#48](https://github.com/JustVugg/colibri/issues/48)) — on constrained-output workloads (JSON/NDJSON, function calling, structured extraction) the grammar itself is a third draft source: wherever it admits exactly **one** legal byte (braces, quotes, key names, enum bodies), that forced span is tokenized and injected as pre-accepted drafts with ~1.0 acceptance — no draft head, no lookup table, and it engages even with the int4 MTP head from [#8](https://github.com/JustVugg/colibri/issues/8). It never constrains sampling: forced spans are verified in the same batch-union forward as any draft, so a wrong or out-of-sync grammar cannot change the output — worst case is rejected drafts, and an adaptive guard turns the source off below 50% acceptance. Byte-level GBNF subset (literals, char classes, `| ( ) ? * +`, comments); `GRAMMAR_DRAFT=n` caps the forced span per forward (default 24). Composes with `DRAFT`/MTP, which fill the free-text gaps between forced spans. +- **Grammar-forced speculative drafts** (`GRAMMAR=file.gbnf`, [#48](https://github.com/JustVugg/colibri/issues/48)) — on constrained-output workloads (JSON/NDJSON, function calling, structured extraction) the grammar itself is a third draft source: wherever it admits exactly **one** legal byte (braces, quotes, key names, enum bodies), that forced span is tokenized and injected as pre-accepted drafts with ~1.0 acceptance — no draft head, no lookup table, and it engages even with the int4 MTP head from [#8](https://github.com/JustVugg/colibri/issues/8). It never constrains sampling: forced spans are verified in the same batch-union forward as any draft, so a wrong or out-of-sync grammar cannot change the output — worst case is rejected drafts, and an adaptive guard turns the source off below 50% acceptance. Byte-level GBNF subset (literals, char classes, `| ( ) ? * +`, comments); `GRAMMAR_DRAFT=n` caps the forced span per forward (default 24). Composes with `DRAFT`/MTP, which fill the free-text gaps between forced spans. Full reference — mechanism, measured A/Bs, when it pays, prior art: [docs/grammar-draft.md](docs/grammar-draft.md). - **True sampling** — temperature + nucleus, defaults tuned for int4 reality (0.7 / 0.90; the official 1.0 / 0.95 samples quantization noise from the tail). - **Integer-dot kernels** (Q8_0-style int8 activations, AVX2 `maddubs`): int8 matmuls 1.4–2.5× faster (119 GFLOP/s measured), int4 1.8× in batch — routing decided per shape by measurement (int4 single-row stays f32: it measured slower). - **MLA weight absorption** (DeepSeek trick) for decode: no per-token k/v reconstruction — the query absorbs `kv_b`, context is projected after attention. Validated exact: TF 32/32 and generation 20/20 with absorption forced everywhere. diff --git a/docs/grammar-draft.md b/docs/grammar-draft.md new file mode 100644 index 0000000..707fe5c --- /dev/null +++ b/docs/grammar-draft.md @@ -0,0 +1,85 @@ +# Grammar-forced speculative drafts + +*The canonical reference for the `GRAMMAR=` draft source (method F). History: idea in +[#48](https://github.com/JustVugg/colibri/issues/48), implementation in +[#70](https://github.com/JustVugg/colibri/pull/70), consolidated write-up with A/B +measurements and corrections in [#146](https://github.com/JustVugg/colibri/issues/146).* + +## The mechanism in one paragraph + +On constrained-output workloads (JSON/NDJSON, function calling, structured extraction) +the grammar itself is a draft source. A byte-level GBNF subset is compiled into a +set-of-stacks PDA; wherever the grammar admits **exactly one** legal next byte +(braces, quotes, key names, enum bodies, fixed separators), the engine walks that +forced span, tokenizes it, and injects it as pre-accepted draft tokens into the +**same batch-union verify forward** used by MTP and n-gram drafts. It never +constrains sampling: forced spans are *verified by the target model*, so a wrong, +stale, or desynced grammar cannot change the output — worst case is rejected drafts, +and an adaptive guard disables the source below 50% acceptance. It composes with +`DRAFT`/MTP, which fill the free-text gaps between forced spans. + +## Why it pays disproportionately in this engine + +In a disk-streaming MoE, the marginal cost of a decode step is **unique expert bytes +off disk**, not FLOPs. The batch-union forward loads each unique expert once for +*all* positions in the verify batch, and structural tokens route heavily into experts +the neighboring free-text positions already pulled in. So every accepted forced span +converts almost directly into **disk reads avoided per emitted token**. On a dense +CPU engine the same trick saves only compute; here it saves the scarcest resource. +Fewer forwards also means less LRU churn, which compounds with cache hit rate. + +## Usage + +```bash +GRAMMAR=fit.gbnf TEMP=0 PROMPT="..." ./glm 64 4 8 +``` + +- `GRAMMAR=` — arm the draft source with a grammar. Byte-level GBNF + subset: literals (escapes `\" \\ \n \r \t \xHH`), byte classes `[...]`/`[^...]`, + rule refs, groups, postfix `? * +`, `|`, `#` comments. Root rule must be `root`. +- `GRAMMAR_DRAFT=n` — cap the forced span per forward (default 24, max 48). +- Arming is lazy (the walker starts at the first byte the root admits, skipping + preambles) and desync-tolerant (a non-conforming byte kills the walker for the + current span; it re-arms at the next opportunity). +- The engine reports at end of run: `grammar: NN% acceptance (a/b forced drafts)`. + +## What to expect (measured, current `main`, M3 Max — details in #146) + +| workload shape | tok/forward | end-to-end | +|---|---|---| +| conforming compact NDJSON (PR #70 era) | 1.60 | large | +| current-main NDJSON with sloppy spacing / long free-text fields | 1.21–1.22, 87% acceptance | ~+5% | +| prose-dominated windows (preambles, long `reason` strings) | ~1.0 | ~nil | + +**The win tracks structural-span density.** Free-text content forces nothing; +keys, separators, quotes and enum bodies force everything. Two practical levers: + +1. **Prompt for compact output** (short enum-like fields, no markdown fences). +2. **Whitespace-tolerant grammars**: a compact-only grammar desyncs at the first + stray space and forfeits every span after it. Emitting an optional-whitespace + rule at separators costs those single bytes (two legal bytes → not forced) but + keeps the walker alive, so every multi-byte span after them still drafts. + Because drafts are verified, tolerance is strictly acceptance-positive. + +## Guarantees and limits + +- **Lossless by construction**: greedy output is byte-identical with and without + a grammar (verified in the #70 A/B); under sampling the rejection step preserves + the sampling distribution. The tokenization boundary of a forced byte span is + not guaranteed to coincide with the model's — verification absorbs the + difference (worst case the last draft of a span is rejected). +- Adaptive shut-off below 50% acceptance, so a mismatched grammar degrades to + baseline speed, never below it for long. +- Benchmarking note (learned the hard way, see #146): hold the PIN hot-store and + page-cache warmth constant across A/B rungs and report the expert hit-rate + column next to any tok/s claim. + +## Prior art + +The umbrella idea is not colibri's: **jump-forward decoding** (SGLang; +[XGrammar](https://arxiv.org/abs/2411.15100)) skips grammar-deterministic tokens by +*constraining output*, and Outlines' coalescence does the analogous FSM move. What +this implementation adds: forced spans as a **draft source verified in the target's +own forward** (lossless even under a wrong grammar, composes with MTP/n-gram in one +union batch), deployed where the win is denominated in expert I/O rather than +forward passes.