Non-neural speculative draft benchmark for Qwen math traces

README.md

Non-Neural Drafting Baseline for Qwen Math Traces

This gist contains a small reproducibility bundle for measuring a deterministic draft program on qwen35 math reasoning traces.

The measurement is an offline speculative-decoding analogue: at each response position, a causal program sees only the prompt plus already accepted response prefix, drafts 16 tokens, and counts how many leading tokens match the held-out trace continuation.

No model calls, gradients, or future trace text are used at eval time.

Headline Result

On the held-out eval split of vikhyatk/qwen-math-traces:

• 16 draft tokens per position
• 2.05 accepted tokens per position
• 2.27 oracle-best accepted tokens per position
• 1.14 accepted tokens per decode-step call
• 53.8% every-position hit rate
• 2.5% full 16-token draft acceptance

These numbers come from deterministic copy, n-gram, and backoff rules over the train split, with source scoring calibrated on a separate calibration split.

Files:

• minimal_spec_decode.py: standalone gist-sized benchmark script.
• results.json: headline metrics and histogram from the best stable run.
• accepted_histogram.svg: generated chart.

Related artifacts:

• HF dataset: https://huggingface.co/datasets/vikhyatk/qwen-math-traces
• Source workspace includes dataset prep and chart-generation scripts.

Run

uv run python minimal_spec_decode.py

The default script uses the public dataset and Qwen tokenizer. For a quick smoke test:

uv run python minimal_spec_decode.py \
  --train-limit 200 \
  --calibration-limit 32 \
  --eval-limit 16 \
  --max-completion-tokens 128 \
  --max-index-positions 100000 \
  --backoff-max-index-positions 100000

Dataset

vikhyatk/qwen-math-traces is a deterministic subset of moondream/synth-math-reasoning-v2 filtered to:

• model == "qwen35"
• setting == "think"

The public dataset contains 7,500 rows split as:

• train: 6,988 rows
• calibration: 256 rows
• eval: 256 rows

Full Local Command

The results.json metrics were produced with the larger research script using:

uv run python spec_math_draft.py \
  --max-examples 7500 \
  --eval-size 256 \
  --max-completion-tokens 512 \
  --policy-calibration-size 256 \
  --backoff-context-lengths 4,6,8,12,16,24,32,48 \
  --backoff-min-context-len 4 \
  --max-index-positions 4000000 \
  --backoff-max-index-positions 4000000

accepted_histogram.svg

minimal_spec_decode.py

#!/usr/bin/env python3
"""Minimal non-neural speculative-decoding analogue for qwen math traces.

This script is intentionally small enough to paste into a gist. It loads a
public trace dataset, builds causal deterministic draft rules on the train
split, calibrates source scores on a calibration split, then evaluates on a
held-out eval split.

Run:

    uv run python artifacts/minimal_spec_decode.py

The metric: at each response position, predict 16 tokens from the prompt plus
already accepted prefix, then count how many leading tokens match the true
continuation.
"""

from __future__ import annotations

import argparse
import json
import math
import re
from collections import Counter, defaultdict
from dataclasses import dataclass
from statistics import mean
from typing import Iterable


@dataclass(frozen=True)
class Example:
    prompt: str
    completion: str
    prompt_tokens: tuple[int, ...]
    completion_tokens: tuple[int, ...]


@dataclass(frozen=True)
class Candidate:
    tokens: tuple[int, ...]
    source: str
    score: float


def encode(tokenizer, text: str) -> tuple[int, ...]:
    return tuple(tokenizer.encode(text, add_special_tokens=False)) if text else ()


def decode(tokenizer, tokens: tuple[int, ...]) -> str:
    return tokenizer.decode(tokens, skip_special_tokens=False) if tokens else ""


def accepted_prefix(prediction: tuple[int, ...], actual: tuple[int, ...], limit: int) -> int:
    accepted = 0
    for left, right in zip(prediction[:limit], actual[:limit]):
        if left != right:
            break
        accepted += 1
    return accepted


def parse_ints(value: str) -> list[int]:
    return sorted({int(part) for part in value.split(",") if part.strip()})


class ContinuationIndex:
    def __init__(
        self,
        context_lengths: list[int],
        draft_len: int,
        max_positions: int,
        max_continuations_per_context: int,
    ) -> None:
        self.context_lengths = sorted(set(context_lengths))
        self.draft_len = draft_len
        self.max_positions = max_positions
        self.max_continuations_per_context = max_continuations_per_context
        self.tables: dict[int, dict[tuple[int, ...], Counter[tuple[int, ...]]]] = {
            length: defaultdict(Counter) for length in self.context_lengths
        }
        self.positions_added = 0

    def add(self, tokens: tuple[int, ...]) -> None:
        last = len(tokens) - self.draft_len
        for pos in range(max(0, last)):
            if self.max_positions and self.positions_added >= self.max_positions:
                return
            continuation = tuple(tokens[pos : pos + self.draft_len])
            for length in self.context_lengths:
                if pos >= length:
                    self.tables[length][tuple(tokens[pos - length : pos])][continuation] += 1
            self.positions_added += 1

    def prune(self) -> None:
        limit = self.max_continuations_per_context
        for table in self.tables.values():
            for key, counter in list(table.items()):
                if len(counter) > limit:
                    table[key] = Counter(dict(counter.most_common(limit)))

    def predict(self, context: tuple[int, ...]) -> Candidate | None:
        for length in sorted(self.context_lengths, reverse=True):
            if len(context) < length:
                continue
            counter = self.tables[length].get(tuple(context[-length:]))
            if not counter:
                continue
            continuation, count = counter.most_common(1)[0]
            return Candidate(continuation, f"ngram:{length}", length * 8.0 + math.log1p(count))
        return None


class BackoffIndex:
    def __init__(
        self,
        context_lengths: list[int],
        min_context_len: int,
        draft_len: int,
        max_positions: int,
    ) -> None:
        self.context_lengths = sorted({x for x in context_lengths if x > 0})
        self.min_context_len = min_context_len
        self.draft_len = draft_len
        self.max_positions = max_positions
        self.tables: dict[int, dict[tuple[int, ...], Counter[int]]] = {
            length: defaultdict(Counter) for length in self.context_lengths
        }
        self.positions_added = 0

    def add(self, tokens: tuple[int, ...]) -> None:
        for pos, token in enumerate(tokens):
            if self.max_positions and self.positions_added >= self.max_positions:
                return
            for length in self.context_lengths:
                if pos >= length:
                    self.tables[length][tuple(tokens[pos - length : pos])][token] += 1
            self.positions_added += 1

    def predict_one(self, context: tuple[int, ...]) -> tuple[int, int, float] | None:
        for length in sorted(self.context_lengths, reverse=True):
            if length < self.min_context_len or len(context) < length:
                continue
            counter = self.tables[length].get(tuple(context[-length:]))
            if not counter:
                continue
            token, count = counter.most_common(1)[0]
            return token, length, count / counter.total()
        return None

    def predict(self, context: tuple[int, ...]) -> Candidate | None:
        generated: list[int] = []
        lengths: list[int] = []
        confidences: list[float] = []
        working = tuple(context)
        for _ in range(self.draft_len):
            pred = self.predict_one(working)
            if pred is None:
                break
            token, length, confidence = pred
            generated.append(token)
            lengths.append(length)
            confidences.append(confidence)
            working = working + (token,)
        if not generated:
            return None
        score = mean(lengths) * 5.5 + mean(confidences) * 8.0
        return Candidate(tuple(generated), f"backoff:{lengths[0]}", score)


class PrefixCopyState:
    def __init__(self, full_tokens: tuple[int, ...], context_lengths: list[int], draft_len: int):
        self.full_tokens = full_tokens
        self.context_lengths = sorted({x for x in context_lengths if x > 0})
        self.draft_len = draft_len
        self.tables: dict[int, dict[tuple[int, ...], tuple[int, ...]]] = {
            length: {} for length in self.context_lengths
        }
        self.next_start = {length: 0 for length in self.context_lengths}

    def advance(self, known_len: int) -> None:
        for length in self.context_lengths:
            table = self.tables[length]
            limit = known_len - length - self.draft_len
            start = self.next_start[length]
            while start <= limit:
                table[tuple(self.full_tokens[start : start + length])] = tuple(
                    self.full_tokens[start + length : start + length + self.draft_len]
                )
                start += 1
            self.next_start[length] = start

    def predict(self, known_len: int) -> Candidate | None:
        self.advance(known_len)
        for length in sorted(self.context_lengths, reverse=True):
            if known_len >= length:
                continuation = self.tables[length].get(tuple(self.full_tokens[known_len - length : known_len]))
                if continuation:
                    return Candidate(continuation, f"copy:{length}", length * 9.0 + 0.5)
        return None


class TextCursor:
    def __init__(self, tokenizer, full_tokens: tuple[int, ...], start_len: int) -> None:
        self.tokenizer = tokenizer
        self.full_tokens = full_tokens
        self.last_known = start_len
        self.text = ""
        self.current_line = ""

    def advance(self, known_len: int) -> None:
        if known_len <= self.last_known:
            return
        piece = decode(self.tokenizer, self.full_tokens[self.last_known : known_len])
        self.text += piece
        self.current_line = (self.current_line + piece).split("\n")[-1]
        self.last_known = known_len


def segment_after_delimiters(line: str) -> str:
    for delimiter in (":", "=", "\"", "'", "(", "[", "{", "*"):
        if delimiter in line:
            line = line.rsplit(delimiter, 1)[1]
    return line.lstrip()


def prompt_span_candidates(prompt: str, max_span_chars: int) -> list[str]:
    spans: set[str] = set()
    starts = [m.start() for m in re.finditer(r"(?<!\w)(?:[$]?\d|[A-Za-z])", prompt)]
    boundaries = {len(prompt)}
    boundaries.update(m.end() for m in re.finditer(r"[,.;:!?)](?:\s|$)", prompt))
    for start in starts:
        for end in sorted(boundaries):
            if end <= start:
                continue
            if end - start > max_span_chars:
                break
            span = prompt[start:end].strip()
            if len(span) >= 8:
                spans.add(span)
        word_match = re.match(r"\S+(?:\s+\S+){0,8}", prompt[start:])
        if word_match and len(word_match.group(0)) >= 8:
            spans.add(word_match.group(0)[:max_span_chars].strip())
    return sorted(spans, key=lambda span: (-len(span), span))


class PromptAutocompleteState:
    def __init__(self, tokenizer, example: Example, draft_len: int) -> None:
        self.tokenizer = tokenizer
        self.draft_len = draft_len
        self.cursor = TextCursor(tokenizer, example.prompt_tokens + example.completion_tokens, len(example.prompt_tokens))
        self.prefix_map: dict[str, str] = {}
        for span in prompt_span_candidates(example.prompt, 240):
            for prefix_len in range(3, min(len(span), 240)):
                self.prefix_map.setdefault(span[:prefix_len].lower(), span)

    def predict(self, known_len: int) -> Candidate | None:
        self.cursor.advance(known_len)
        segment = segment_after_delimiters(self.cursor.current_line)
        if len(segment) < 3:
            return None
        span = self.prefix_map.get(segment.lower())
        if not span or not span.lower().startswith(segment.lower()):
            return None
        tokens = encode(self.tokenizer, span[len(segment) :])[: self.draft_len]
        if not tokens:
            return None
        return Candidate(tokens, "prompt-autocomplete", 52.0 + min(len(segment), 56) * 0.55)


class SourceAveragePolicy:
    def __init__(self, accepted_by_source: dict[str, list[int]], prior_strength: float) -> None:
        all_values = [value for values in accepted_by_source.values() for value in values]
        self.global_avg = mean(all_values) if all_values else 0.0
        self.prior_strength = prior_strength
        self.source_avg: dict[str, float] = {}
        for source, values in accepted_by_source.items():
            self.source_avg[source] = (sum(values) + self.global_avg * prior_strength) / (
                len(values) + prior_strength
            )

    def score(self, candidate: Candidate) -> float:
        return self.source_avg.get(candidate.source, self.global_avg) * 1000.0 + candidate.score * 0.0001


def build_examples(rows, tokenizer, max_completion_tokens: int) -> list[Example]:
    examples: list[Example] = []
    for row in rows:
        prompt = row["problem"]
        if row.get("problem_source"):
            prompt = f"{prompt}\n\n{row['problem_source']}"
        completion = row["raw_reasoning"].strip()
        prompt_tokens = encode(tokenizer, prompt)
        completion_tokens = encode(tokenizer, completion)
        if max_completion_tokens:
            completion_tokens = completion_tokens[:max_completion_tokens]
        if completion_tokens:
            examples.append(Example(prompt, completion, prompt_tokens, completion_tokens))
    return examples


def maybe_limit(rows, limit: int):
    return rows.select(range(min(limit, len(rows)))) if limit > 0 else rows


def build_indexes(args, train_examples: list[Example]) -> tuple[ContinuationIndex, BackoffIndex]:
    exact = ContinuationIndex(
        parse_ints(args.context_lengths),
        args.draft_len,
        args.max_index_positions,
        args.max_continuations_per_context,
    )
    backoff = BackoffIndex(
        parse_ints(args.backoff_context_lengths),
        args.backoff_min_context_len,
        args.draft_len,
        args.backoff_max_index_positions,
    )
    for ex in train_examples:
        tokens = ex.prompt_tokens + ex.completion_tokens
        exact.add(tokens)
        backoff.add(tokens)
    exact.prune()
    return exact, backoff


def candidates_for(example: Example, tokenizer, exact: ContinuationIndex, backoff: BackoffIndex):
    full = example.prompt_tokens + example.completion_tokens
    prompt_len = len(example.prompt_tokens)
    copy_state = PrefixCopyState(full, parse_ints("4,6,8,12,16,24,32,48,64,96,128"), exact.draft_len)
    prompt_state = PromptAutocompleteState(tokenizer, example, exact.draft_len)

    def at(position: int) -> list[Candidate]:
        known_len = prompt_len + position
        context = full[:known_len]
        out: list[Candidate] = []
        for candidate in (
            copy_state.predict(known_len),
            prompt_state.predict(known_len),
            backoff.predict(context),
            exact.predict(context),
        ):
            if candidate is not None:
                out.append(candidate)
        return out

    return at


def collect_calibration(examples, tokenizer, exact, backoff, draft_len: int, stride: int, max_positions: int):
    accepted_by_source: dict[str, list[int]] = defaultdict(list)
    for ex in examples:
        at = candidates_for(ex, tokenizer, exact, backoff)
        positions = 0
        for pos in range(0, len(ex.completion_tokens), stride):
            if max_positions and positions >= max_positions:
                break
            actual = ex.completion_tokens[pos : pos + draft_len]
            for candidate in at(pos):
                accepted_by_source[candidate.source].append(accepted_prefix(candidate.tokens, actual, draft_len))
            positions += 1
    return accepted_by_source


def summarize(values: list[int], draft_len: int) -> dict:
    hist = dict(sorted(Counter(values).items()))
    nonzero = [x for x in values if x > 0]
    return {
        "positions": len(values),
        "avg_accepted": mean(values) if values else 0.0,
        "hit_rate": len(nonzero) / len(values) if values else 0.0,
        "full_draft_rate": sum(x >= draft_len for x in values) / len(values) if values else 0.0,
        "avg_on_hit": mean(nonzero) if nonzero else 0.0,
        "histogram": hist,
    }


def evaluate(examples, tokenizer, exact, backoff, policy, draft_len: int) -> tuple[dict, dict]:
    every_values: list[int] = []
    decode_values: list[int] = []
    for ex in examples:
        at = candidates_for(ex, tokenizer, exact, backoff)
        for pos in range(len(ex.completion_tokens)):
            actual = ex.completion_tokens[pos : pos + draft_len]
            choices = at(pos)
            if not choices:
                every_values.append(0)
            else:
                best = max(choices, key=policy.score)
                every_values.append(accepted_prefix(best.tokens, actual, draft_len))

        pos = 0
        while pos < len(ex.completion_tokens):
            actual = ex.completion_tokens[pos : pos + draft_len]
            choices = at(pos)
            if not choices:
                decode_values.append(0)
                pos += 1
                continue
            best = max(choices, key=policy.score)
            accepted = accepted_prefix(best.tokens, actual, draft_len)
            decode_values.append(accepted)
            pos += accepted if accepted == draft_len else accepted + 1
    return summarize(every_values, draft_len), summarize(decode_values, draft_len)


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--dataset", default="vikhyatk/qwen-math-traces")
    parser.add_argument("--tokenizer", default="Qwen/Qwen2.5-7B-Instruct")
    parser.add_argument("--draft-len", type=int, default=16)
    parser.add_argument("--max-completion-tokens", type=int, default=512)
    parser.add_argument("--context-lengths", default="1,2,3,4,6,8,12,16,24,32,48,64")
    parser.add_argument("--backoff-context-lengths", default="4,6,8,12,16,24,32,48")
    parser.add_argument("--backoff-min-context-len", type=int, default=4)
    parser.add_argument("--max-index-positions", type=int, default=4_000_000)
    parser.add_argument("--backoff-max-index-positions", type=int, default=4_000_000)
    parser.add_argument("--max-continuations-per-context", type=int, default=8)
    parser.add_argument("--policy-prior-strength", type=float, default=200.0)
    parser.add_argument("--train-limit", type=int, default=0)
    parser.add_argument("--calibration-limit", type=int, default=0)
    parser.add_argument("--eval-limit", type=int, default=0)
    args = parser.parse_args()

    from datasets import load_dataset
    from transformers import AutoTokenizer

    tokenizer = AutoTokenizer.from_pretrained(args.tokenizer, token=True)
    ds = load_dataset(args.dataset)
    train = build_examples(maybe_limit(ds["train"], args.train_limit), tokenizer, args.max_completion_tokens)
    calibration = build_examples(
        maybe_limit(ds["calibration"], args.calibration_limit),
        tokenizer,
        args.max_completion_tokens,
    )
    eval_examples = build_examples(maybe_limit(ds["eval"], args.eval_limit), tokenizer, args.max_completion_tokens)

    exact, backoff = build_indexes(args, train)
    accepted_by_source = collect_calibration(
        calibration, tokenizer, exact, backoff, args.draft_len, stride=2, max_positions=256
    )
    policy = SourceAveragePolicy(accepted_by_source, args.policy_prior_strength)
    every, decode_summary = evaluate(eval_examples, tokenizer, exact, backoff, policy, args.draft_len)

    print(
        json.dumps(
            {
                "dataset": args.dataset,
                "train_examples": len(train),
                "calibration_examples": len(calibration),
                "eval_examples": len(eval_examples),
                "draft_len": args.draft_len,
                "exact_positions_added": exact.positions_added,
                "backoff_positions_added": backoff.positions_added,
                "every_position": every,
                "decode_step": decode_summary,
            },
            indent=2,
            sort_keys=True,
        )
    )


if __name__ == "__main__":
    main()

results.json

{
  "benchmark": {
    "date": "2026-07-07",
    "dataset_recipe": "qwen35 think subset, 7500 rows, seed 0, 256 eval rows",
    "draft_len": 16,
    "command": "uv run python spec_math_draft.py --max-examples 7500 --eval-size 256 --max-completion-tokens 512 --policy-calibration-size 256 --backoff-context-lengths 4,6,8,12,16,24,32,48 --backoff-min-context-len 4 --max-index-positions 4000000 --backoff-max-index-positions 4000000",
    "rules": [
      "prefix-copy",
      "prompt-autocomplete",
      "backoff-ngram",
      "global-ngram"
    ],
    "policy": "source-average",
    "selected": {
      "positions": 129919,
      "avg_accepted": 2.046,
      "hit_rate": 0.538,
      "full_draft_rate": 0.025,
      "avg_on_hit": 3.802,
      "histogram": {
        "0": 60003,
        "1": 24610,
        "2": 13996,
        "3": 8527,
        "4": 5096,
        "5": 3481,
        "6": 2520,
        "7": 1888,
        "8": 1461,
        "9": 1148,
        "10": 946,
        "11": 794,
        "12": 668,
        "13": 582,
        "14": 508,
        "15": 439,
        "16": 3252
      }
    },
    "oracle_best": {
      "positions": 129919,
      "avg_accepted": 2.270,
      "hit_rate": 0.577,
      "full_draft_rate": 0.028,
      "avg_on_hit": 3.937,
      "histogram": {
        "0": 55009,
        "1": 24588,
        "2": 15117,
        "3": 9418,
        "4": 5823,
        "5": 3942,
        "6": 2862,
        "7": 2131,
        "8": 1660,
        "9": 1302,
        "10": 1062,
        "11": 878,
        "12": 747,
        "13": 641,
        "14": 565,
        "15": 486,
        "16": 3688
      }
    },
    "decode_step": {
      "calls": 61002,
      "avg_accepted": 1.139,
      "hit_rate": 0.410,
      "full_draft_rate": 0.007,
      "avg_on_hit": 2.776
    }
  },
  "scaling_curve": [
    {
      "label": "2k default",
      "rows": 2000,
      "avg_accepted": 1.699
    },
    {
      "label": "5k broad",
      "rows": 5000,
      "avg_accepted": 1.873
    },
    {
      "label": "5k larger caps",
      "rows": 5000,
      "avg_accepted": 1.963
    },
    {
      "label": "7.5k larger caps",
      "rows": 7500,
      "avg_accepted": 2.046
    }
  ]
}