kolibri1: n-gram spec-decode on the engine path — 14.1 tok/s bench, byte-identical chains

.agents/issues/MODEL-TEXT-kolibri-1-kolibri1-for-causal-lm/ISSUE-LOCAL-01M4P22PX6BPRQAY1CK20FZQF6.md

@@ -7,7 +7,7 @@ GitHub: -
 Mirror: PENDING
 Availability: FULL
 Created: 2026-10-11
-Updated: 2026-10-11
+Updated: 2026-10-12
 Closed: 2026-10-11
 
 ## Problem
@@ -17,3 +17,104 @@ The CPU kolibri decode bench measures 5.76 tok/s (t32+32GiB cache+NEON) against
 ## Resolution
 
 PROBE COMPLETE 2026-10-11 (branch row/kolibri-spec-probe, commit pending; report docs/bench-evidence/kolibri1-spec-probe-20261011.md). Findings: (1) PORTABILITY: the tree's n-gram spec-decode machinery is generic end to end — scheduler spec slots (scheduler.cpp:247-251,700-741,948-950,1226-1265), engine draft install/rollback (core.cpp:101-127), runner verify (runner.cpp:3546+ sample_tokens_with_rejection, greedy accept-iff-argmax per rejection_sampler.h:13-51), proposer (runner.cpp:4517-4570 -> ngram_proposer.cpp), and GPUModelRunner is the production runner on CPU (model_loader.h:897). Only the MTP/DFlash/DSpark DRAFTERS and the GDN/mamba spec-state widening are model-specific; n-gram needs neither. ONE blocker: LoadedEngine::MakeKVCacheMaybeSpec (model_loader.cpp:1756-1763) builds MakeQwen3_5KVCacheSpec for ANY spec config, ignoring the model's make_kv_cache; for kolibri it injects a MambaSpec+fa_draft group and the runner refuses at runner.cpp:1074 (pinned by test_kolibri1_ngram_spec, green). (2) MEASURED ACCEPTANCE (exact simulation with the production proposer over the W3 golden chains, pinned by test_kolibri1_ngram_accept, 90/90 green): k=5 window 1-4 -> 103 verify passes vs 248 plain steps (2.408x), 150/297 drafts accepted (50.5%); k=8 -> 2.583x; upstream default window 5/5 is worse (1.687x). Bench chain (128 distinct random prompt, 64-token 2-cycle chain captured from the bench run): k=5 -> 12 passes vs 63 (5.25x), 52/53 accepted (98.1%). (3) PROTOTYPE (contained, bench-level per the brief's fallback since the engine path is blocked by the one seam): test_kolibri1_spec_probe drives the production proposer + ONE production ForwardKolibri1Forward over the k+1 verify sequence on the real checkpoint — bench prompt: spec chain byte-identical to plain, 53/53 accepted (100%), decode 14.51 tok/s vs 4.92 plain (2.95x, CLEARS the 12 tok/s target); W3 goldens: 8/8 byte-identical, 146/270 accepted (54.1%), decode 8.99 vs 6.84 tok/s (1.31x, short-prompt diluted). Green 17/17. (4) Engine reach pinned: spec-OFF engine generates on CPU; spec-ON throws the precise refusal (test_kolibri1_ngram_spec, 6/6 green). (5) INTEGRATION PLAN (not done — serving wiring deliberately out of scope): dispatch MakeKVCacheMaybeSpec's widened branch on use_eagle() or a recurrent group in the model's own KV config (kolibri has none, so its make_kv_cache is already spec-correct for n-gram); then flip the pinned blocker test into the engine-level e2e gate with a row spec; exercise the --speculative-config flag once via the C ABI; tune k/window on serving-shaped traffic; batch>1 acceptance owed. Candidate (b) tiny external draft model NOT recommended. Baseline reproduced: bench 5.156 tok/s this run (task baseline 5.76), anchor 109726 PASS.
+
+INTEGRATION LANDED 2026-10-12 (branch row/kolibri-spec-engine). The probe's
+one blocker is fixed and the n-gram proposer now runs in the kolibri ENGINE
+path behind the existing opt-in (EngineParams::speculative_config — the same
+object the C ABI v6 field and the server's --speculative-config populate).
+
+THE SEAM FIX (the only product change): LoadedEngine::MakeKVCacheMaybeSpec
+(src/vllm/entrypoints/model_loader.cpp) now builds the MODEL's own KV spec
+(ModelRegistry::MakeKVCache) and takes the widened MakeQwen3_5KVCacheSpec
+branch only when spec->use_eagle() OR the model's own groups carry a
+recurrent/mamba group. A draft-free n-gram config on a model with no recurrent
+group keeps the model's own spec: the drafts are ordinary scheduled tokens, so
+the target's groups are the whole cache need and NO draft-proposal group is
+added (the fa_draft group is consumed only by the MTP draft head's own paged
+KV, runner.cpp:1865-1898, a forward n-gram does not have — and that allocation
+is skipped outright on kolibri's multi-cache topology). The default (no spec
+config) path is byte-identical: exactly ModelRegistry::MakeKVCache. qwen3.5's
+ngram path keeps the widened spec (it HAS GDN layers), so the landed 27B ngram
+gate is untouched. RED-FIRST: with the fix reverted, the new seam unit fails
+(REQUIRE( 3 == 2 ) — the injected fa+gdn+fa_draft set) and the e2e fails at
+construction with the pinned refusal "vt: runner: MambaSpec state shapes must
+be positive at runner.cpp:1074" (the spec-OFF arm runs all 9 prompts first,
+0 drafts). GREEN with the fix: unit 21/21; e2e 83/83.
+
+NO NEW ENGINE WIRING WAS NEEDED: the scheduler spec slots, the engine's draft
+install/rollback, the runner's verify (sample_tokens_with_rejection) and the
+draft-free proposer (propose_drafts_ngram) were already generic; the e2e gate
+proves the whole loop executed (53/53 drafts accepted on the bench prompt).
+tests/vllm/models/test_kolibri1_ngram_spec is flipped from the pinned blocker
+into (1) a checkpoint-free seam unit (kolibri under an ngram spec config keeps
+its own 2 groups — no MambaSpec, no fa_draft, byte-identical to spec-off —
+while qwen3.5 still widens to fa+gdn+fa_draft with the conv row at
+(K-1)+k taps and num_speculative_blocks=k) and (2) the engine e2e gate.
+tests/capi/test_capi.cpp gains the ABI reach pin, in TWO halves (2026-10-12
+review fix — the first shape pinned only the PARSE: a missing checkpoint fails
+with VLLM_ERR_MODEL_LOAD whether or not the parsed config is forwarded, so it
+could not see a dropped assignment). The PARSE pin: a valid ngram
+speculative_config is ACCEPTED at the ABI layer (MODEL_LOAD on a missing
+checkpoint, not INVALID_ARGUMENT), and an invalid ngram window is a CALLER
+error (VLLM_ERR_INVALID_ARGUMENT, never MODEL_LOAD). The FORWARDING pin: a
+drafter-chain document the ABI layer accepts is REFUSED BY NAME by the
+engine's up-front ResolveSpecConfig, before any model I/O — so against a
+missing checkpoint the load fails with the spec-validation refusal, not the
+missing-directory error. Mutation-verified both ways: reverting the whole
+parse+assignment block fails the parse pin; dropping only the assignment into
+EngineParams fails the forwarding pin (the reviewer's exact mutation).
+
+MEASURED THROUGH THE ENGINE (LoadedEngine, t32 + 32 GiB dequant cache, under
+flock, VT_ASYNC_SCHED=0 so BOTH arms run the synchronous scheduler — ngram
+forces sync via async_scheduling_compatible()==use_eagle(); with the default
+async spec-OFF arm the goldens A/B is confounded by pipelined input prep and
+is reported separately):
+- EXACTNESS: spec-ON == spec-OFF BYTE-IDENTICAL on 9/9 prompts (the
+  deterministic bench prompt 128-in/64-out + the 8 W3 golden prompts). The
+  spec-OFF bench chain ends at the pinned anchor 109726.
+- ACCEPTANCE (engine counters spec_drafts_accepted/proposed): bench prompt
+  53/53 — the pinned number (exact sim over the torch chain: 52/53; the probe
+  harness over this tree's bf16 chain: 53/53). Goldens aggregate 146/270 —
+  EXACTLY the probe harness's 146/270 (sim over torch chains: 150/297). The
+  engine reproduces the harness's acceptance counts exactly, per prompt too.
+- ENGINE tok/s (CORRECTED DENOMINATOR, 2026-10-12 review fix): decode tok/s
+  divides the tokens generated INSIDE the decode interval — the full output
+  count less the tokens already present when the first output arrived — by
+  the interval from that first output to finish; wall tok/s keeps the full
+  count over the whole request. (The first report divided the FULL count by
+  that interval, billing the first step's tokens with no generation time;
+  the corrected decode figures read modestly lower — the operator's
+  overcounted verification measured 14.06 bench / 8.60 goldens decode.)
+  Re-measured 2026-10-12 on this box (t32 + 32 GiB cache, under flock,
+  VT_ASYNC_SCHED=0): bench prompt decode 4.53 -> 13.78 tok/s (3.04x; wall 3.04
+  -> 5.89, 1.93x) — CLEARS the 12 tok/s target through the engine (harness
+  reference: 4.92 -> 14.51, 2.95x). Goldens aggregate decode 6.27 -> 8.13
+  tok/s (1.30x; wall 5.46 -> 6.91, 1.27x; harness reference: 6.84 -> 8.99,
+  1.31x; short-prompt diluted, same shape). Default-scheduler run for the
+  record: bench decode 4.80 -> 15.72 (wall 3.18 -> 6.34), goldens decode 6.35
+  -> 8.98 (wall 5.39 -> 7.60) — the async spec-OFF arm's pipelined input prep
+  is the confound; not the A/B of record. The denominator fix is
+  REPORTING-ONLY: acceptance (53/53 bench, 146/270 goldens) and byte-identity
+  (9/9) are unchanged.
+- W3 VERDICT: unchanged with spec off — test_kolibri1_w3 900/900 PASS (the
+  default path is byte-identical; the KV-spec unit pins it too) and the decode
+  bench anchor 109726 PASS. With spec ON: byte-identical to spec-OFF (9/9), so
+  the torch-golden relationship is identical by transitivity — 4/8 prompts
+  token-identical end to end with the same per-prompt shared prefixes the
+  probe measured (141/145 argmax positions, 4 adjudicated near-tie flips
+  unchanged). No hard flip.
+- REGRESSION SURFACE: test_kolibri1_ngram_accept 90/90 (pinned sim stays
+  green); test_kolibri1_spec_probe (the harness) green; seam-adjacent suites
+  test_loaded_engine_dense 157, test_model_registry 1229,
+  test_deepseek_v4_scaffold 684, test_kv_pool_min_fit 16, test_runner 1918,
+  test_kv_cache_fp8_wiring 487 — all green (4491 assertions).
+
+REMAINING DELTA (owed, none blocking): (a) the engine's reject-heavy path — on
+the most reject-heavy golden (10/84 acceptance) the verify T=1+kr forward
+makes spec-ON ~1.8x SLOWER wall than plain on that prompt (10.20 s vs 6.81 s);
+the aggregate still wins 1.20x. A per-step cost trace of the reject path
+(scheduler rollback + draft re-splice + rejected-slot re-schedule) is the
+named follow-up. (b) Batch>1 acceptance (probe plan step 6). (c) Real-traffic
+(non-degenerate-chain) acceptance (probe §9 caveat: the measured rates are an
+upper bound). (d) Serving-shaped k/window tuning — config, not code
+(k=5, window 1-4 measured on these chains; upstream default 5/5 is worse).

.agents/specs/kolibri-1-cpu.md

@@ -367,6 +367,44 @@ the lane ON — the default flip is blocked on (1) the idle-host A/B and (2) a
 near-tie adjudication of that anchor (gate-semantics change, own spec).
 Evidence: docs/bench-evidence/kolibri1-neon-paged-attn-20261008.md.
 
+N-GRAM SPEC DECODE LANDED ON THE ENGINE PATH (2026-10-12, branch
+`row/kolibri-spec-engine`, issue ISSUE-LOCAL-01M4P22PX6BPRQAY1CK20FZQF6): the
+spec-probe's one seam is fixed — `LoadedEngine::MakeKVCacheMaybeSpec`
+(src/vllm/entrypoints/model_loader.cpp) builds the model's own KV spec and
+takes the widened qwen3.5 branch only when `spec->use_eagle()` or the model's
+own groups carry a recurrent/mamba group. A draft-free n-gram config on
+kolibri keeps kolibri's own two groups (full + sliding attention): the drafts
+are ordinary scheduled tokens, so the target's groups are the whole cache need
+and no draft-proposal group is added (the `fa_draft` group serves only the MTP
+draft head's paged KV, a forward n-gram does not have). The default (no spec
+config) path is byte-identical. No new engine wiring was needed:
+`EngineParams::speculative_config` (the C ABI v6 field / the server's
+--speculative-config) already fed the generic scheduler slots, engine draft
+pull, runner verify and n-gram proposer. `test_kolibri1_ngram_spec` is flipped
+from the pinned blocker into a checkpoint-free seam unit (kolibri keeps its
+own groups under an ngram spec; qwen3.5 still widens) plus the engine e2e
+gate: spec-ON == spec-OFF BYTE-IDENTICAL on 9/9 prompts (bench prompt +
+8 W3 goldens), acceptance exactly 53/53 on the bench prompt and 146/270 on
+the goldens (the probe harness's numbers), and engine-level decode tok/s
+4.53 -> 13.78 on the bench prompt (3.04x, CLEARS the 12 tok/s target;
+harness reference 4.92 -> 14.51) and 6.27 -> 8.13 on the goldens (1.30x),
+measured with the synchronous scheduler on BOTH arms (ngram forces sync;
+VT_ASYNC_SCHED=0, t32 + 32 GiB cache, under the lock). Decode tok/s uses the
+CORRECTED denominator (2026-10-12 review fix): the tokens generated inside
+the decode interval — the full count less the tokens already present when
+the first output arrived — over that interval; wall tok/s keeps the full
+count (bench wall 3.04 -> 5.89; goldens wall 5.46 -> 6.91). The first report
+divided the full count by the decode interval, overcounting decode by the
+first step's tokens (the operator's overcounted verification: 14.06 bench /
+8.60 goldens decode); the correction is reporting-only — acceptance
+(53/53, 146/270) and byte-identity (9/9) are unchanged. W3 unchanged: 900/900
+with spec off; with spec ON the chains are byte-identical to spec-OFF, so the
+torch-golden relationship (4/8 prompts end-to-end identical, 4 adjudicated
+near-tie flips) is unchanged. Owed: the engine's reject-heavy path costs more
+than the harness's on the lowest-acceptance golden (verify T=1+kr for ~0.5
+emitted tokens/pass; aggregate still wins); batch>1 acceptance; real-traffic
+acceptance; serving-shaped k/window tuning (config, not code).
+
 ## R7 resolution — the tokenizer engine accepts the Kolibri-1 split regex
 
 ### Scope — what is actually in the file

include/vllm.h

@@ -481,8 +481,10 @@ typedef struct vllm_model_params {
    * the emitted tokens under greedy sampling, so it is a throughput knob.
    * NULL or "" => speculation DISABLED, the byte-identical default engine.
    * A malformed document or unsupported method fails vllm_engine_load with
-   * VLLM_ERR_INVALID_ARGUMENT. Only MTP is supported today, on the Qwen3.5/3.6
-   * checkpoints that ship an mtp.* head (safetensors only). Borrowed for the
+   * VLLM_ERR_INVALID_ARGUMENT. "mtp" needs a Qwen3.5/3.6 checkpoint that
+   * ships an mtp.* head (safetensors only); "ngram" is draft-free and runs on
+   * any target (measured on the kolibri1 CPU checkpoint: the speculative
+   * chain is byte-identical to speculative-off). Borrowed for the
    * vllm_engine_load call only. */
   const char* speculative_config;
   /* ── Automatic prefix caching (ABI v7) ─────────────────────────────────────

include/vllm/entrypoints/model_loader.h

@@ -748,10 +748,15 @@ class LoadedEngine {
   static std::pair<int, int> ResolveSchedulerBlockSizes(
       const vllm::v1::KVCacheConfig& kv_cfg, int block_size,
       bool enable_prefix_caching);
-  // SPEC-MTP I5d: build the KV-cache spec, widened for speculation when a spec
-  // config is set (the extra GDN k+1 state slots + widened conv row + the
-  // `fa_draft` full-attn group, MakeQwen3_5KVCacheSpec num_spec>0). With no spec
-  // config it is exactly ModelRegistry::MakeKVCache (byte-identical).
+  // SPEC-MTP I5d: build the model's own KV-cache spec (ModelRegistry::MakeKVCache)
+  // and, when a spec config is set AND the method or the model needs spec state
+  // (an eagle method, or the model's own groups carry a recurrent/mamba group),
+  // rebuild it widened (the extra GDN k+1 state slots + widened conv row + the
+  // `fa_draft` full-attn group, MakeQwen3_5KVCacheSpec num_spec>0). A draft-free
+  // method (ngram) on a model with no recurrent group keeps the model's own spec:
+  // the drafts are ordinary scheduled tokens, so the target's groups are the
+  // whole cache need. With no spec config this is exactly
+  // ModelRegistry::MakeKVCache (byte-identical).
  public:
   //
   // PUBLIC because it is the only entry to the KV geometry the engine actually

src/vllm/entrypoints/model_loader.cpp

@@ -1614,10 +1614,14 @@ std::optional<vllm::SpeculativeConfig> LoadedEngine::ResolveSpecConfig(
   // SPEC-NGRAM (ROAD-V1-D3): the draft-FREE n-gram proposer. num_speculative_tokens
   // (k) is REQUIRED (speculative.py:1224-1234); the prompt_lookup window defaults to
   // 5/5. There is NO draft checkpoint to load and NO n_predict-module constraint —
-  // the drafts come from matching the sequence's own suffix. The GDN spec verify
-  // machinery (widened conv + k+1 state slots) is still needed and is reused via
-  // MakeQwen3_5KVCacheSpec(num_spec>0) exactly like MTP; the fa_draft draft-KV group
-  // it allocates is simply unused (ngram has no draft-model forward).
+  // the drafts come from matching the sequence's own suffix. On a model WITH a
+  // recurrent (GDN) group the spec verify machinery (widened conv + k+1 state
+  // slots) is still needed and MakeKVCacheMaybeSpec reuses it via
+  // MakeQwen3_5KVCacheSpec(num_spec>0) exactly like MTP; the fa_draft draft-KV
+  // group it allocates is simply unused (ngram has no draft-model forward). On a
+  // model WITHOUT one (kolibri) the model's own KV spec is already spec-correct:
+  // the drafts are ordinary scheduled tokens and the cache need is the target's
+  // groups alone.
   if (cli.method == "ngram") {
     if (!cli.num_speculative_tokens.has_value()) {
       throw std::invalid_argument(
@@ -1753,15 +1757,31 @@ vllm::v1::KVCacheConfig LoadedEngine::MakeKVCacheMaybeSpec(
                  block_size, resolved_bs);
     block_size = resolved_bs;
   }
-  vllm::v1::KVCacheConfig kv;
-  if (spec.has_value()) {
+  // SPEC-NGRAM / kolibri (ISSUE-LOCAL-01M4P22PX6BPRQAY1CK20FZQF6): build the
+  // MODEL's own KV spec first, always. The widened qwen3.5 spec below exists
+  // for exactly two needs: (a) an eagle method (mtp/dflash/dspark), whose
+  // draft-model forward writes its own K/V into the extra `fa_draft`
+  // full-attn group, and (b) a model with a recurrent (GDN/mamba) group,
+  // whose conv row widens to (K-1)+k taps and whose SSM state needs the k+1
+  // snapshot slots a k+1 verify step rewinds. A draft-FREE method (ngram) on
+  // a model with NO recurrent group — kolibri: sliding/full attention only —
+  // needs neither: the drafts are ordinary scheduled tokens verified by the
+  // target's own forward, so the target's own groups are the whole cache need
+  // and there is no draft-proposal group to add. Building
+  // MakeQwen3_5KVCacheSpec for such a model instead injected a MambaSpec
+  // group with empty shapes (the checkpoint carries no linear_* keys) plus an
+  // unused fa_draft group, and the runner refused at construction
+  // (runner.cpp:1074 "MambaSpec state shapes must be positive"). The default
+  // (no spec config) path is untouched: exactly ModelRegistry::MakeKVCache,
+  // byte-identical.
+  vllm::v1::KVCacheConfig kv =
+      ModelRegistry::MakeKVCache(model, config, block_size, num_blocks);
+  if (spec.has_value() && (spec->use_eagle() || kv.has_mamba_layers())) {
     // Speculation is Qwen3.5/3.6-only at this pin (both gate checkpoints); build
     // the widened spec KV directly (extra GDN k+1 state slots + widened conv row
     // + the `fa_draft` full-attn group). MakeQwen3_5KVCacheSpec(num_spec>0).
     kv = vllm::MakeQwen3_5KVCacheSpec(config, block_size, num_blocks,
                                       spec->ResolvedNumSpeculativeTokens());
-  } else {
-    kv = ModelRegistry::MakeKVCache(model, config, block_size, num_blocks);
   }
   // FIX-KV-GROUP-LAYER-COUNT (#1963, #1966). THE single funnel: both branches
   // above return through here, so one call reaches every architecture and both