Spaces:
Sleeping
Sleeping
| """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) | |
| 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) | |
| 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 | |