quantcost 0.2.0__py3-none-any.whl

edgellm/hub.py

@@ -0,0 +1,125 @@
+"""Resolve a (model, precision) pair to a local ONNX file, downloading if needed.
+
+The default benchmark path deliberately consumes **pre-quantized ONNX already
+published on the Hugging Face Hub** rather than quantizing locally. That is what
+keeps the install light: a contributor needs ``onnxruntime`` and NumPy, not
+PyTorch, ``transformers``, ``optimum`` and a GPU. It also keeps results
+comparable, because every machine scores the *same bytes* of the same artifact
+instead of its own locally-produced quantization.
+
+The naming convention here is the one used across the ``onnx-community`` and
+``transformers.js`` model repos, so thousands of existing models work unchanged.
+"""
+
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass
+from pathlib import Path
+
+log = logging.getLogger(__name__)
+
+#: Canonical precision label -> filename stem inside the repo's ``onnx/`` folder.
+#: Keys are what the user types; values are the Hub convention.
+PRECISION_FILES: dict[str, str] = {
+    "fp32": "model",
+    "fp16": "model_fp16",
+    "int8": "model_int8",
+    "uint8": "model_uint8",
+    "q4": "model_q4",
+    "q4f16": "model_q4f16",
+    "bnb4": "model_bnb4",
+}
+
+#: What `--precisions` expands to when the user does not say. fp32 is the honest
+#: baseline every other number is relative to; int8 and q4 are the two choices a
+#: person actually weighs when shipping to a CPU.
+DEFAULT_PRECISIONS: tuple[str, ...] = ("fp32", "int8", "q4")
+
+#: Small, widely-mirrored instruct models that carry the full precision set.
+#: The default is the 135M: a complete three-precision sweep stays under a GB of
+#: download and finishes in minutes on a laptop CPU with no accelerator.
+DEFAULT_MODEL = "HuggingFaceTB/SmolLM2-135M-Instruct"
+
+
+class ModelResolutionError(RuntimeError):
+    """A model/precision combination is not available on the Hub."""
+
+
+@dataclass(frozen=True)
+class ResolvedModel:
+    """A downloaded ONNX artifact, ready to hand to ONNX Runtime."""
+
+    model_id: str
+    precision: str
+    onnx_path: Path
+    size_mb: float
+    tokenizer_path: Path
+
+    @property
+    def filename(self) -> str:
+        return PRECISION_FILES[self.precision] + ".onnx"
+
+
+def available_precisions(model_id: str, *, revision: str = "main") -> list[str]:
+    """Return the precision labels this repo actually publishes, in canonical order."""
+    from huggingface_hub import HfApi
+
+    files = set(HfApi().list_repo_files(model_id, revision=revision))
+    return [p for p, stem in PRECISION_FILES.items() if f"onnx/{stem}.onnx" in files]
+
+
+def resolve(
+    model_id: str,
+    precision: str,
+    *,
+    revision: str = "main",
+    cache_dir: str | None = None,
+) -> ResolvedModel:
+    """Download (or reuse from cache) one precision of ``model_id``.
+
+    Raises :class:`ModelResolutionError` with the list of precisions the repo does
+    publish, rather than letting a 404 surface, because "this model has no int4"
+    is an ordinary situation a contributor needs to act on.
+    """
+    from huggingface_hub import hf_hub_download
+    from huggingface_hub.errors import EntryNotFoundError
+
+    if precision not in PRECISION_FILES:
+        raise ModelResolutionError(
+            f"Unknown precision '{precision}'. Known: {', '.join(PRECISION_FILES)}"
+        )
+
+    stem = PRECISION_FILES[precision]
+    kwargs = {"repo_id": model_id, "revision": revision, "cache_dir": cache_dir}
+
+    try:
+        onnx_path = Path(hf_hub_download(filename=f"onnx/{stem}.onnx", **kwargs))
+    except EntryNotFoundError as exc:
+        have = available_precisions(model_id, revision=revision)
+        raise ModelResolutionError(
+            f"{model_id} does not publish '{precision}' "
+            f"(looked for onnx/{stem}.onnx). Available: {', '.join(have) or 'none'}"
+        ) from exc
+
+    total_bytes = onnx_path.stat().st_size
+
+    # Models over the 2 GB protobuf limit keep their weights in a sibling
+    # ``.onnx_data`` file. ONNX Runtime needs it next to the graph, and it is most
+    # of the real on-disk footprint, so it counts toward the reported size.
+    try:
+        data_path = Path(hf_hub_download(filename=f"onnx/{stem}.onnx_data", **kwargs))
+        total_bytes += data_path.stat().st_size
+        log.debug("external weights: %s", data_path)
+    except EntryNotFoundError:
+        pass
+
+    tokenizer_path = Path(hf_hub_download(filename="tokenizer.json", **kwargs))
+
+    return ResolvedModel(
+        model_id=model_id,
+        precision=precision,
+        onnx_path=onnx_path,
+        size_mb=round(total_bytes / (1024 * 1024), 2),
+        tokenizer_path=tokenizer_path,
+    )

edgellm/leaderboard.py

@@ -0,0 +1,210 @@
+"""Aggregate every submitted card into a leaderboard.
+
+The interesting question is not "which machine is fastest" — that is just a list
+of who owns the newest CPU. It is **"does quantization actually pay off, and where"**.
+So the leaderboard is organised around the per-precision speedup distribution
+across machines, with the raw per-machine rows underneath.
+
+Cards flagged as not comparable by :mod:`edgellm.validate` are listed but kept out
+of the aggregate statistics, so a run with a different token budget cannot quietly
+move the headline numbers.
+"""
+
+from __future__ import annotations
+
+import json
+import statistics
+from pathlib import Path
+from typing import Any
+
+from edgellm.validate import validate_card
+
+BASELINE = "fp32"
+
+
+def _plural(count: int, noun: str) -> str:
+    return f"{count} {noun}" if count == 1 else f"{count} {noun}s"
+
+
+def _speedups(rows: list[dict]) -> dict[str, float]:
+    by_precision = {r["precision"]: r for r in rows}
+    base = by_precision.get(BASELINE)
+    if not base or not base.get("tokens_per_second"):
+        return {}
+    return {
+        r["precision"]: r["tokens_per_second"] / base["tokens_per_second"]
+        for r in rows
+        if r["precision"] != BASELINE and r.get("tokens_per_second")
+    }
+
+
+def _ppl_deltas(rows: list[dict]) -> dict[str, float]:
+    by_precision = {r["precision"]: r for r in rows}
+    base = by_precision.get(BASELINE)
+    if not base or not base.get("perplexity"):
+        return {}
+    return {
+        r["precision"]: (r["perplexity"] / base["perplexity"] - 1) * 100
+        for r in rows
+        if r["precision"] != BASELINE and r.get("perplexity")
+    }
+
+
+def build_leaderboard(card_paths: list[Path]) -> dict[str, Any]:
+    """Load, validate and aggregate cards into a JSON-serializable leaderboard."""
+    entries: list[dict[str, Any]] = []
+    rejected: list[dict[str, str]] = []
+
+    for path in sorted(card_paths):
+        report = validate_card(path)
+        if not report.ok:
+            rejected.append({"file": path.name, "reason": "; ".join(report.errors)})
+            continue
+
+        card = json.loads(path.read_text())
+        machine = card["machine"]
+        entries.append(
+            {
+                "file": path.name,
+                "comparable": report.comparable,
+                "model_id": card["model_id"],
+                "created_utc": card["created_utc"],
+                "submitted_by": card.get("submitted_by", ""),
+                "notes": card.get("notes", ""),
+                "cpu": machine["cpu"],
+                "arch": machine["arch"],
+                "os": machine["os"],
+                "physical_cores": machine["physical_cores"],
+                "ram_gb": machine["ram_gb"],
+                "provider": machine["provider"],
+                "threads": machine["intra_op_threads"],
+                "thread_policy": machine.get("thread_policy", "unknown"),
+                "onnxruntime": machine["onnxruntime"],
+                "rows": card["rows"],
+                "speedups": _speedups(card["rows"]),
+                "ppl_deltas": _ppl_deltas(card["rows"]),
+            }
+        )
+
+    aggregate = _aggregate(e for e in entries if e["comparable"])
+    return {
+        "entries": sorted(entries, key=lambda e: (e["model_id"], e["cpu"])),
+        "aggregate": aggregate,
+        "rejected": rejected,
+        "machine_count": len({e["cpu"] for e in entries}),
+        "model_count": len({e["model_id"] for e in entries}),
+    }
+
+
+def _aggregate(entries) -> dict[str, Any]:
+    """Per-(model, precision) speedup and quality-cost distribution across machines."""
+    buckets: dict[tuple[str, str], dict[str, list[float]]] = {}
+    for entry in entries:
+        for precision, speedup in entry["speedups"].items():
+            bucket = buckets.setdefault((entry["model_id"], precision), {"speedup": [], "ppl": []})
+            bucket["speedup"].append(speedup)
+        for precision, delta in entry["ppl_deltas"].items():
+            bucket = buckets.setdefault((entry["model_id"], precision), {"speedup": [], "ppl": []})
+            bucket["ppl"].append(delta)
+
+    out = []
+    for (model_id, precision), values in sorted(buckets.items()):
+        speedups = values["speedup"]
+        ppls = values["ppl"]
+        out.append(
+            {
+                "model_id": model_id,
+                "precision": precision,
+                "machines": len(speedups),
+                "speedup_median": round(statistics.median(speedups), 3) if speedups else None,
+                "speedup_min": round(min(speedups), 3) if speedups else None,
+                "speedup_max": round(max(speedups), 3) if speedups else None,
+                "slower_than_baseline": sum(1 for s in speedups if s < 0.95),
+                "ppl_delta_median_pct": round(statistics.median(ppls), 2) if ppls else None,
+            }
+        )
+    return {"by_model_precision": out}
+
+
+def render_leaderboard_markdown(board: dict[str, Any]) -> str:
+    """The committed LEADERBOARD.md."""
+    lines = [
+        "# Leaderboard",
+        "",
+        "<!-- Generated by `quantcost leaderboard`. Do not edit by hand:",
+        "     every number here is derived from a card in results/community/. -->",
+        "",
+        f"**{_plural(board['machine_count'], 'machine')} · "
+        f"{_plural(board['model_count'], 'model')} · "
+        f"{_plural(len(board['entries']), 'card')} submitted**",
+        "",
+        "## Does quantization pay off?",
+        "",
+        "Speedup is throughput relative to that machine's own fp32 baseline, so the",
+        "column compares quantization choices rather than hardware budgets.",
+        "",
+    ]
+
+    aggregate = board["aggregate"]["by_model_precision"]
+    if aggregate:
+        lines += [
+            "| Model | Precision | Machines | Median speedup | Range | "
+            "Slower than fp32 | Median PPL change |",
+            "| --- | --- | --- | --- | --- | --- | --- |",
+        ]
+        for row in aggregate:
+            span = (
+                f"{row['speedup_min']:.2f}x – {row['speedup_max']:.2f}x"
+                if row["speedup_min"] is not None
+                else "—"
+            )
+            median = f"{row['speedup_median']:.2f}x" if row["speedup_median"] else "—"
+            ppl = (
+                f"{row['ppl_delta_median_pct']:+.1f}%"
+                if row["ppl_delta_median_pct"] is not None
+                else "—"
+            )
+            slower = f"{row['slower_than_baseline']}/{row['machines']}" if row["machines"] else "—"
+            lines.append(
+                f"| `{row['model_id']}` | {row['precision']} | {row['machines']} | "
+                f"{median} | {span} | {slower} | {ppl} |"
+            )
+    else:
+        lines.append("_No comparable cards yet._")
+
+    lines += ["", "## Per-machine results", ""]
+    for entry in board["entries"]:
+        flag = "" if entry["comparable"] else " _(not ranked — non-default settings)_"
+        credit = f" — submitted by {entry['submitted_by']}" if entry["submitted_by"] else ""
+        lines += [
+            f"### {entry['cpu']} · {entry['os']} ({entry['arch']}){flag}",
+            "",
+            f"`{entry['model_id']}` · {entry['physical_cores']} physical cores · "
+            f"{entry['ram_gb']:.0f} GB RAM · "
+            f"{entry['threads']} threads ({entry['thread_policy']}) · "
+            f"onnxruntime {entry['onnxruntime']} · {entry['provider']}{credit}",
+            "",
+            "| Precision | Size (MB) | tok/s | Speedup | Peak RAM (MB) | Perplexity |",
+            "| --- | --- | --- | --- | --- | --- |",
+        ]
+        for row in entry["rows"]:
+            speed = entry["speedups"].get(row["precision"])
+            speed_text = (
+                "baseline" if row["precision"] == BASELINE else (f"{speed:.2f}x" if speed else "—")
+            )
+            ppl = f"{row['perplexity']:.2f}" if row.get("perplexity") else "—"
+            lines.append(
+                f"| {row['precision']} | {row['size_mb']:.0f} | {row['tokens_per_second']:.2f} | "
+                f"{speed_text} | {row['peak_ram_mb']:.0f} | {ppl} |"
+            )
+        if entry["notes"]:
+            lines += ["", f"> {entry['notes']}"]
+        lines.append("")
+
+    if board["rejected"]:
+        lines += ["## Rejected cards", ""]
+        for item in board["rejected"]:
+            lines.append(f"- `{item['file']}`: {item['reason']}")
+        lines.append("")
+
+    return "\n".join(lines)

edgellm/models.py

@@ -0,0 +1,91 @@
+"""Model + tokenizer loading.
+
+:class:`ModelLoader` wraps Hugging Face ``transformers`` so the rest of the code
+never touches ``AutoModelForCausalLM`` directly and device/dtype resolution lives
+in exactly one place.
+"""
+
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass
+
+import torch
+from transformers import (
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    PreTrainedModel,
+    PreTrainedTokenizerBase,
+)
+
+from edgellm.config import ModelConfig
+
+logger = logging.getLogger(__name__)
+
+_DTYPES: dict[str, torch.dtype] = {
+    "float32": torch.float32,
+    "float16": torch.float16,
+    "bfloat16": torch.bfloat16,
+}
+
+
+@dataclass
+class LoadedModel:
+    """A loaded model bundled with its tokenizer and resolved placement."""
+
+    model: PreTrainedModel
+    tokenizer: PreTrainedTokenizerBase
+    device: str
+    dtype: torch.dtype
+
+
+class ModelLoader:
+    """Load a causal-LM + tokenizer from a :class:`ModelConfig`."""
+
+    def __init__(self, config: ModelConfig) -> None:
+        self.config = config
+
+    def resolve_device(self) -> str:
+        """Turn ``device: auto`` into a concrete device string for this machine."""
+        requested = self.config.device
+        if requested != "auto":
+            return requested
+        if torch.cuda.is_available():
+            return "cuda"
+        if torch.backends.mps.is_available():
+            return "mps"
+        return "cpu"
+
+    def resolve_dtype(self) -> torch.dtype:
+        try:
+            return _DTYPES[self.config.dtype]
+        except KeyError as exc:
+            raise ValueError(
+                f"Unsupported dtype '{self.config.dtype}'. Choose one of {list(_DTYPES)}."
+            ) from exc
+
+    def load(self) -> LoadedModel:
+        """Load the model and tokenizer and move them onto the resolved device."""
+        device = self.resolve_device()
+        dtype = self.resolve_dtype()
+        logger.info("Loading %s (dtype=%s, device=%s)", self.config.id, self.config.dtype, device)
+
+        tokenizer = AutoTokenizer.from_pretrained(
+            self.config.id,
+            revision=self.config.revision,
+            trust_remote_code=self.config.trust_remote_code,
+        )
+        if tokenizer.pad_token is None and tokenizer.eos_token is not None:
+            # Small decoder-only models often ship without a pad token.
+            tokenizer.pad_token = tokenizer.eos_token
+
+        model = AutoModelForCausalLM.from_pretrained(
+            self.config.id,
+            revision=self.config.revision,
+            dtype=dtype,
+            trust_remote_code=self.config.trust_remote_code,
+        )
+        model.to(device)
+        model.eval()
+
+        return LoadedModel(model=model, tokenizer=tokenizer, device=device, dtype=dtype)

edgellm/ort_lite.py

@@ -0,0 +1,236 @@
+"""A torch-free ONNX Runtime decoder: greedy generation with a KV cache, in NumPy.
+
+Why this exists instead of ``optimum.onnxruntime``: Optimum is the right tool when
+you are exporting and quantizing models, but it pulls in PyTorch and
+``transformers``, which together are a multi-gigabyte install and the single
+biggest reason a stranger abandons a benchmark run. The decode loop for a
+causal LM is about a hundred lines of NumPy, so the default path implements it
+directly and depends only on ``onnxruntime``, ``numpy`` and ``tokenizers``.
+
+The graph contract (verified against the ``onnx-community`` / ``transformers.js``
+exports, which is what the Hub publishes for thousands of models):
+
+    inputs   input_ids                  int64  [batch, sequence]
+             attention_mask             int64  [batch, total_sequence]
+             position_ids               int64  [batch, sequence]
+             past_key_values.{i}.key    T      [batch, kv_heads, past_sequence, head_dim]
+             past_key_values.{i}.value  T      ... same
+    outputs  logits                     float  [batch, sequence, vocab]
+             present.{i}.key            T      [batch, kv_heads, total_sequence, head_dim]
+             present.{i}.value          T      ... same
+
+Layer count, KV-head count, head dim and the cache dtype are all read off the
+session rather than from ``config.json``, so the runner is model-agnostic and
+works for fp16 graphs (whose cache is float16) without special-casing.
+"""
+
+from __future__ import annotations
+
+import logging
+import re
+import time
+
+import numpy as np
+
+from edgellm.base import GenerationResult, InferenceRunner
+from edgellm.config import GenerationConfig
+
+log = logging.getLogger(__name__)
+
+_ORT_TO_NUMPY = {
+    "tensor(float)": np.float32,
+    "tensor(float16)": np.float16,
+    "tensor(bfloat16)": np.float32,  # ORT has no bf16 numpy view; feed fp32 zeros.
+    "tensor(int64)": np.int64,
+    "tensor(int32)": np.int32,
+}
+
+_PAST_RE = re.compile(r"^past_key_values\.(\d+)\.(key|value)$")
+
+
+class OnnxLiteRunner(InferenceRunner):
+    """Greedy ONNX Runtime generation with a KV cache, implemented in NumPy."""
+
+    def __init__(
+        self,
+        onnx_path: str,
+        tokenizer_path: str,
+        *,
+        provider: str = "CPUExecutionProvider",
+        name: str = "ort-lite",
+        intra_op_threads: int | None = None,
+    ) -> None:
+        import onnxruntime as ort
+        from tokenizers import Tokenizer
+
+        available = ort.get_available_providers()
+        if provider not in available:
+            raise RuntimeError(
+                f"Execution provider '{provider}' is not available in this ONNX "
+                f"Runtime build (have: {', '.join(available)})."
+            )
+
+        # Thread count is the single largest lever on CPU throughput, so it is
+        # always pinned to a concrete number and recorded in the result card.
+        # Leaving it at ORT's default would store 0 ("decide for me"), which
+        # makes two cards look identical while having run very differently.
+        # The *fast* cores only, not every physical core. On a heterogeneous CPU
+        # one thread landing on an efficiency core gates the whole parallel
+        # region; see edgellm.card.performance_cores for the measurements.
+        if intra_op_threads is None:
+            from edgellm.card import performance_cores
+
+            intra_op_threads, self.thread_policy = performance_cores()
+        else:
+            self.thread_policy = "explicit"
+
+        opts = ort.SessionOptions()
+        opts.intra_op_num_threads = intra_op_threads
+        opts.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
+
+        self.name = name
+        self.provider = provider
+        self.session = ort.InferenceSession(onnx_path, opts, providers=[provider])
+        self.tokenizer = Tokenizer.from_file(tokenizer_path)
+        self.intra_op_threads = opts.intra_op_num_threads
+
+        self._inspect_graph()
+
+    # ------------------------------------------------------------------ setup
+
+    def _inspect_graph(self) -> None:
+        """Read layer count, KV geometry and cache dtype off the graph itself."""
+        inputs = {i.name: i for i in self.session.get_inputs()}
+        self.input_names = set(inputs)
+
+        layers = sorted(
+            int(m.group(1))
+            for name in inputs
+            if (m := _PAST_RE.match(name)) and m.group(2) == "key"
+        )
+        if not layers:
+            raise ValueError(
+                "This ONNX graph exposes no 'past_key_values.*' inputs, so it is not a "
+                "KV-cached decoder export. Use a model from onnx-community or any repo "
+                "exported with Optimum's `--task text-generation-with-past`."
+            )
+        if layers != list(range(len(layers))):
+            raise ValueError(f"Non-contiguous KV cache layer indices: {layers}")
+        self.num_layers = len(layers)
+
+        spec = inputs["past_key_values.0.key"]
+        shape = spec.shape
+        if len(shape) != 4 or not isinstance(shape[1], int) or not isinstance(shape[3], int):
+            raise ValueError(
+                f"Unexpected KV cache shape {shape}; expected "
+                "[batch, kv_heads, past_sequence, head_dim] with static head dims."
+            )
+        self.kv_heads, self.head_dim = int(shape[1]), int(shape[3])
+        self.cache_dtype = _ORT_TO_NUMPY.get(spec.type, np.float32)
+
+        self.vocab_size = None
+        for out in self.session.get_outputs():
+            if out.name == "logits" and isinstance(out.shape[-1], int):
+                self.vocab_size = int(out.shape[-1])
+
+        self.uses_position_ids = "position_ids" in self.input_names
+
+        log.debug(
+            "graph: %d layers, %d kv heads, head_dim %d, cache %s, position_ids %s",
+            self.num_layers,
+            self.kv_heads,
+            self.head_dim,
+            np.dtype(self.cache_dtype).name,
+            self.uses_position_ids,
+        )
+
+    def _empty_cache(self, batch: int = 1) -> dict[str, np.ndarray]:
+        """Zero-length KV cache: the graph's signal that this is the prefill step."""
+        empty = np.zeros((batch, self.kv_heads, 0, self.head_dim), dtype=self.cache_dtype)
+        feed = {}
+        for i in range(self.num_layers):
+            feed[f"past_key_values.{i}.key"] = empty
+            feed[f"past_key_values.{i}.value"] = empty
+        return feed
+
+    # ------------------------------------------------------------- inference
+
+    def _forward(
+        self,
+        input_ids: np.ndarray,
+        past: dict[str, np.ndarray],
+        past_len: int,
+    ) -> tuple[np.ndarray, dict[str, np.ndarray]]:
+        """One step (prefill or decode). Returns logits and the next cache."""
+        batch, seq = input_ids.shape
+        total = past_len + seq
+
+        feed: dict[str, np.ndarray] = {
+            "input_ids": input_ids,
+            "attention_mask": np.ones((batch, total), dtype=np.int64),
+            **past,
+        }
+        if self.uses_position_ids:
+            feed["position_ids"] = np.arange(past_len, total, dtype=np.int64)[None, :].repeat(
+                batch, axis=0
+            )
+
+        outputs = self.session.run(None, feed)
+        names = [o.name for o in self.session.get_outputs()]
+        by_name = dict(zip(names, outputs, strict=True))
+
+        logits = by_name["logits"]
+        next_past = {
+            f"past_key_values.{i}.{kind}": by_name[f"present.{i}.{kind}"]
+            for i in range(self.num_layers)
+            for kind in ("key", "value")
+        }
+        return logits, next_past
+
+    def forward_logits(self, input_ids: np.ndarray) -> np.ndarray:
+        """Single full forward pass with no cache — what perplexity scoring needs."""
+        logits, _ = self._forward(input_ids, self._empty_cache(input_ids.shape[0]), 0)
+        return logits
+
+    def generate(self, prompt: str, generation: GenerationConfig) -> GenerationResult:
+        """Greedy decode ``max_new_tokens`` tokens, timing the whole loop.
+
+        Decoding is always greedy here regardless of ``generation.temperature``:
+        a benchmark needs a fixed token count and a deterministic path, and
+        sampling would make throughput depend on the RNG. ``min_new_tokens`` is
+        implied — generation never stops early on EOS, so every backend and
+        precision is measured over exactly the same amount of work.
+        """
+        ids = self.tokenizer.encode(prompt).ids
+        input_ids = np.asarray([ids], dtype=np.int64)
+        prompt_tokens = len(ids)
+        want = generation.max_new_tokens
+
+        start = time.perf_counter()
+
+        logits, past = self._forward(input_ids, self._empty_cache(), 0)
+        past_len = prompt_tokens
+        next_id = int(np.argmax(logits[0, -1]))
+        generated = [next_id]
+
+        for _ in range(want - 1):
+            logits, past = self._forward(np.asarray([[next_id]], dtype=np.int64), past, past_len)
+            past_len += 1
+            next_id = int(np.argmax(logits[0, -1]))
+            generated.append(next_id)
+
+        latency_s = time.perf_counter() - start
+
+        text = self.tokenizer.decode(generated, skip_special_tokens=True)
+        return GenerationResult(
+            backend=self.name,
+            prompt=prompt,
+            text=text,
+            prompt_tokens=prompt_tokens,
+            generated_tokens=len(generated),
+            latency_s=latency_s,
+            tokens_per_second=len(generated) / latency_s if latency_s > 0 else 0.0,
+        )
+
+    def close(self) -> None:
+        self.session = None