Blame for src/cism/batch.py

"""True batched decode on one weight stream (dense Llama/Qwen3 first).

Wraps a single native ``Model`` with :class:`BatchEngine.generate_batch`,
which runs ``B`` sequences through one shared weight stream via the native
``BatchSession`` (``Model.create_batch_session``):

- per-sequence prompt / position / generated / rng / eos / finish state,
- concatenated ``X:[sumT, hidden]`` per-layer projections via a single
  ``gemm`` per projection (weights stream once for all ``B`` sequences),
- per-sequence attention and sampling loops.

For ``sumT == 1`` the native kernels run in the same order as the
single-sequence ``Session`` path, giving bitwise-identical results.
Surjo hybrid batching is explicitly out of scope (``from_weights`` rejects
``model_type == "surjo"``).
"""

from __future__ import annotations

import math
import threading
from typing import Any

from .engine import GenerationResult


class BatchEngine:
    """One native dense model serving true batched generations.

    Unlike :class:`cism.pool.EnginePool` (``N`` engines, ``N`` weight
    streams), this holds a single native ``Model``: each
    :meth:`generate_batch` call builds one native ``BatchSession`` whose
    decode steps stream each weight matrix once for all ``B`` sequences.
    Calls are serialized with an internal lock; the native session already
    owns its mutex, but the shared ``WorkerPool`` barrier requires one
    active generation at a time for deterministic results.
    """

    def __init__(self, native: Any, tokenizer: Any, config: dict, model_id: str,
                 *, revision: str | None = None) -> None:
        self._native = native
        self.tokenizer = tokenizer
        self.config = dict(config)
        self.model_id = model_id
        self.revision = revision
        self.context_length = config["max_position_embeddings"]
        self._lock = threading.Lock()
        eos = config.get("eos_token_id", getattr(tokenizer, "eos_token_id", None))
        self._eos = [] if eos is None else eos if isinstance(eos, list) else [eos]
        for token in self._eos:
            self._validate_token(token)

    @classmethod
    def from_weights(cls, config: dict, weights: dict, tokenizer: Any, *,
                     model_id: str = "local", precision: str = "fp32",
                     revision: str | None = None, threads: int = 1,
                     act_precision: str = "fp32") -> BatchEngine:
        from .loader import validate_config
        try:
            from . import _native
        except ImportError as exc:
            raise RuntimeError("CISM native extension is missing. Build with: python -m pip install .") from exc
        if isinstance(config, dict) and config.get("model_type") == "surjo":
            raise ValueError("batch decode is dense-only in this version (Surjo later)")
        normalized = validate_config(config)
        return cls(_native.Model(normalized, weights, precision, threads, act_precision),
                   tokenizer, normalized, model_id, revision=revision)

    @property
    def info(self) -> dict:
        return {**self._native.info, "model": self.model_id, "revision": self.revision,
                "context_length": self.context_length, "architecture": self.config["model_type"],
                "capabilities": ["text", "batching"],
                "batching": "single-weight-stream BatchSession (dense only)"}

    def _validate_token(self, token: int) -> None:
        if type(token) is not int or not 0 <= token < self.config["vocab_size"]:
            raise ValueError("Token IDs must be integers within the model vocabulary")

    def encode(self, prompt: str | list[int]) -> list[int]:
        if isinstance(prompt, str):
            tokens = list(self.tokenizer.encode(prompt, add_special_tokens=True))
        elif isinstance(prompt, list):
            tokens = prompt.copy()
        else:
            raise TypeError("prompt must be text or a list of token IDs")
        if not tokens:
            bos = self.config.get("bos_token_id", getattr(self.tokenizer, "bos_token_id", None))
            if bos is None:
                raise ValueError("Empty prompt requires a BOS token")
            tokens = [bos]
        for token in tokens:
            self._validate_token(token)
        return tokens

    def format_chat(self, messages: list[dict]) -> list[int]:
        if not getattr(self.tokenizer, "chat_template", None):
            raise ValueError("This tokenizer has no chat template; use batch prompts")
        rendered = self.tokenizer.apply_chat_template(
            messages, tokenize=True, add_generation_prompt=True)
        # Bare ID list for most tokenizers; a BatchEncoding
        # ({"input_ids", "attention_mask"}) when the tokenizer's __call__
        # returns one.
        if not isinstance(rendered, list):
            rendered = rendered["input_ids"]
        return self.encode(list(rendered))

    def generate_batch(self, prompts: list[str | list[int]], *,
                       max_new_tokens: int = 128,
                       temperature: float = 0.0, top_p: float = 1.0, top_k: int = 0,
                       seed: int = 0,
                       stop: list[str] | None = None) -> list[GenerationResult]:
        """Generate ``B`` prompts in one weight-stream batch, preserving order.

        Each prompt is encoded (with BOS fallback for empty prompts), then all
        ``B`` sequences decode together: prefill concatenates every prompt
        into one ``[sumT, hidden]`` activation batch (one ``gemm`` per
        projection), and each decode step concatenates one token per
        unfinished sequence the same way. Sampling and attention stay
        per-sequence, so ``B=1`` is bitwise-identical to serial decoding and
        ``batch[i]`` matches a lone ``B=1`` run of ``prompts[i]`` exactly
        (same ``seed`` seeds sequence ``i`` as ``seed + i``; sequence 0 always
        uses ``seed``).

        ``stop`` is an optional list of nonempty strings applied per sequence
        after decoding (first occurrence wins, ``finish_reason="stop"``,
        token IDs retained), mirroring ``Engine.stream`` truncation for
        completed batch text. ``None`` (default) disables stop handling and
        preserves the exact pre-stop behavior.
        """
        if not isinstance(prompts, (list, tuple)):
            raise TypeError("prompts must be a list or tuple")
        prompt_list = list(prompts)
        if len(prompt_list) == 0:
            return []
        if type(max_new_tokens) is not int or max_new_tokens < 0:
            raise ValueError("max_new_tokens must be a nonnegative integer")
        if type(seed) is not int or not 0 <= seed < 2**64:
            raise ValueError("seed must be an unsigned 64-bit integer")
        if type(top_k) is not int or top_k < 0:
            raise ValueError("top_k must be a nonnegative integer")
        for name, value in (("temperature", temperature), ("top_p", top_p)):
            if type(value) not in (float, int) or not math.isfinite(value):
                raise ValueError(f"{name} must be finite")
        if temperature < 0 or not 0 < top_p <= 1:
            raise ValueError("temperature must be nonnegative and top_p in (0, 1]")
        if stop is not None and (not isinstance(stop, list) or any(not isinstance(s, str) or not s for s in stop)):
            raise ValueError("stop must be a list of nonempty strings")
        stops = list(stop or [])
        encoded = [self.encode(p) for p in prompt_list]
        for tokens in encoded:
            if len(tokens) > self.context_length:
                raise ValueError("Prompt length exceeds the model context length")
        # Shared native max_new must fit every sequence; clamp to the smallest
        # remaining window (prompts of equal length -- the common benchmark
        # shape -- are unaffected).
        effective = min([max_new_tokens] + [self.context_length - len(t) for t in encoded])
        with self._lock:
            session = self._native.create_batch_session(
                encoded, max_new_tokens=effective,
                temperature=temperature, top_p=top_p, top_k=top_k,
                seed=seed, eos_token_ids=self._eos)
            batched = session.next_tokens(effective)
            reasons = session.finish_reasons
            # generated_tokens property returns per-sequence counts.
            _generated = session.generated_tokens
        results: list[GenerationResult] = []
        for tokens, token_ids, reason in zip(encoded, batched, reasons):
            text = self.tokenizer.decode(list(token_ids), skip_special_tokens=True,
                                         clean_up_tokenization_spaces=False)
            if stops:
                hits = [text.find(s) for s in stops if s in text]
                if hits:
                    text = text[:min(hits)]
                    reason = "stop"
            results.append(GenerationResult(text, list(token_ids), reason,
                                            len(tokens), len(token_ids)))
        return results