"""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