Hugging Face's logo

Spaces:
spitfire4794
/
test1111111
Sleeping

App Files Community
test1111111 / src /cism /server.py
spitfire4794's picture
spitfire4794
Serve SurjoLabs/Surjo-50m-SFT-Only (int8, 2T) on 7860
92dcc4e verified
Raw History Blame Contribute Delete
230 kB
"""Single-process, text-only OpenAI HTTP adapter for an already loaded engine.
Run the returned app with Uvicorn in one process (no multiple workers). Admission
is bounded by ``queue_size`` INCLUDING the active job. Tokenization, chat
formatting, prefill, and generation all run on one dedicated thread.
Serving decode (Stream S1 -- iteration-level continuous batching):
- Lone single-stream jobs run the dense ``engine.stream`` (``Session``) path,
bitwise-identical to before, so single-stream tokens/s are unregressed by
construction. The native paged-KV path (``create_paged_session``/block-16,
bitwise-identical math) is the default single-stream path
(``CISM_SERVE_PAGED=0`` opts back out to dense): cooled A/B measured
0.999x of dense single-stream (shared BlockPool, zero per-session
alloc; Stream S2) with ~26x prefill savings on shared prefixes via
copy-on-write fork. ``spec_k > 0``, Surjo,
and FWKV jobs always keep the dense ``engine.stream`` path.
- Stream S3 automatic prefix reuse: the paged single path above matches
each arrival against a per-engine radix prefix index (``paging.PrefixCache``)
and fork-shares the longest cached token prefix (suffix-only prefill).
Shared-prefix decode is token-identical to full prefill (native fork
re-checks token equality); stale anchors (S2 eviction / pool reset)
invalidate on first failed fork with exact full-prefill fallback.
Hit/miss counters via ``prefix_cache_stats``; ``CISM_SERVE_PREFIXCACHE=0``
disables. The dense ``BatchSession`` epoch path has no fork entry point,
so it keeps full-prefill rebuilds (documented v0 limit).
- Eligible co-queued greedy jobs (non-streaming, ``temperature == 0``,
``spec_k == 0``, no A2 opt-ins) on a dense engine with a steppable native
``BatchSession`` run in an iteration-level scheduler epoch: per-step ADD
(new arrivals join mid-decode via session rebuild over survivors'
``prompt + generated`` prefixes -- exact for greedy decoding), per-step
EVICT (finished/errored/cancelled leave immediately with one terminal
publish), oldest-first waiting queue with a starvation guard, and chunked
prefill (at most ``prefill_chunk_tokens`` new prompt tokens per rebuild;
env ``CISM_SERVE_PREFILL_CHUNK``). Sampled jobs (same exclusions,
``temperature > 0``) form same-parameter epochs: identical
(``temperature``, ``top_p``, ``top_k``) jobs share one native session
with each job's own serial seed per sequence (token-identical to
serial); sampled epochs form only on a drain (no mid-decode joins) and
never rotate (a rebuild would restart RNG streams), so they always run
to drain. Streaming, speculative, constrained/tool/logprob, Surjo/FWKV,
pool/batch, and test-double jobs always run serially (remaining v0
limits below).
- Kill switches: ``CISM_SERVE_PAGED=0`` (back to dense single),
``CISM_SERVE_MICROBATCH=0`` (disable batching entirely),
``CISM_SERVE_CONTINUOUS=0`` (fall back to the A1 FIFO micro-batch path),
``CISM_SERVE_PREFIXCACHE=0`` (disable S3 automatic prefix reuse;
``CISM_SERVE_PREFIXCACHE_ANCHORS`` bounds retained anchors, default 8).
Rotation tunables: ``CISM_SERVE_EPOCH_STEPS`` (default 1024),
``CISM_SERVE_MAX_WAIT_S`` (default 30), ``CISM_SERVE_JOIN_DEFER_TOKENS``
(default 2: skip a join rebuild when the live batch is that close to
draining).
Stream S1 v0 limits (documented, not faked):
- Batches are heterogeneous in prompt/limit/stop; greedy epochs are
seed-independent (mid-decode joins allowed), sampled epochs are
parameter-homogeneous form-on-drain groups (no joins, no rotation, so
per-seq RNG streams stay serial-identical). Per-job truncation keeps
``"length"``/``"stop"`` exact.
- No preemption with KV preservation: epoch rotation recomputes survivors'
prefixes (one prefill per rebuild); no swap-out, no priorities.
- No intra-prompt prefill chunking: one long prompt prefills atomically
inside a single native call (it is deferred while an epoch is live so it
never stalls decodes, then runs alone).
- No cross-engine fairness beyond FIFO drain order on a shared worker.
Stream A2 API surface (exact-math opt-in paths; never faked):
- Token logprobs (``logprobs`` on ``/v1/completions``, ``logprobs`` +
``top_logprobs`` on ``/v1/chat/completions``, non-streaming only): after
generation, each emitted token is rescored with the real engine
``logits()`` over its true prefix (the same native logits the decoder
used -- bitwise vs the reference in ``test_logits_and_cached_generation``),
with an exact float64 log-softmax. Cross-check: the mean negative
logprob equals ``engine.nll()`` over the same IDs within fp32 noise.
- JSON structured output (``response_format`` ``json_object``/``json_schema``
on chat, greedy only): a Python-side greedy loop calls the real engine
``logits()`` per step and takes the argmax over tokens whose full-sequence
decode keeps the output a viable JSON (resp. schema) prefix. Only a
documented schema subset is supported (object/array/string/integer/number/
boolean/null, properties, required, additionalProperties, items, enum,
const); anything else is a 400, never silently ignored.
- Tool calling (``tools`` + ``tool_choice`` on chat, non-streaming only):
the assistant's text is parsed for a real ``{"name", "arguments"}`` call
matching a declared tool, arguments are validated against the tool's
schema, the call is executed by a real local registry function (currently
``calc``: exact arithmetic, no imports/calls), and the result is fed back
as a follow-up turn. Bounded (at most ``_TOOL_MAX_CALLS`` executions), so
the loop always terminates. Tool names outside the registry are a 400:
the server cannot execute arbitrary client code and will not pretend to.
- Everything without an exact path stays 400/404 (embeddings, audio/image,
``n > 1``, penalties, ``logit_bias``, streaming+logprobs/tools/JSON,
batch+JSON/tools, pool engines+logprobs, unknown tools). Constrained, tool,
and logprob jobs always run serially -- they never join a micro-batch --
so the single-stream gate is unregressed by construction.
"""
from __future__ import annotations
import asyncio
from contextlib import asynccontextmanager, suppress
from dataclasses import dataclass, field
import ast
import hmac
import json
import math
import os
import queue
import statistics
import threading
import time
from typing import Annotated, Any, Iterator, Literal, Protocol
import uuid
from fastapi import APIRouter, Depends, FastAPI, Request
from fastapi.exceptions import RequestValidationError
from fastapi.responses import HTMLResponse, JSONResponse, StreamingResponse
from pydantic import BaseModel, ConfigDict, Field, model_validator
from starlette.exceptions import HTTPException
class _Engine(Protocol):
model_id: str
context_length: int
def encode(self, prompt: str | list[int]) -> list[int]: ...
def format_chat(self, messages: list[dict]) -> list[int]: ...
def stream(self, prompt: str | list[int], **kwargs: Any) -> Iterator[Any]: ...
class _StrictModel(BaseModel):
model_config = ConfigDict(extra="forbid", strict=True, allow_inf_nan=False)
class _StreamOptions(_StrictModel):
include_usage: bool = False
class _TextFormat(_StrictModel):
type: Literal["text"]
class _JsonObjectFormat(_StrictModel):
type: Literal["json_object"]
class _JsonSchemaBody(_StrictModel):
model_config = ConfigDict(extra="forbid", strict=True, allow_inf_nan=False,
populate_by_name=True)
name: _NonemptyString
schema_: dict[str, Any] = Field(alias="schema")
strict: bool | None = None
@property
def schema(self) -> dict[str, Any]:
return self.schema_
class _JsonSchemaFormat(_StrictModel):
type: Literal["json_schema"]
json_schema: _JsonSchemaBody
_ResponseFormat = _TextFormat | _JsonObjectFormat | _JsonSchemaFormat
class _FunctionDef(_StrictModel):
name: _NonemptyString
description: str | None = None
parameters: dict[str, Any] | None = None
strict: bool | None = None
class _ToolDef(_StrictModel):
type: Literal["function"]
function: _FunctionDef
class _NamedToolFunction(_StrictModel):
name: _NonemptyString
class _NamedToolChoice(_StrictModel):
type: Literal["function"]
function: _NamedToolFunction
_ToolChoice = Literal["auto", "none", "required"] | _NamedToolChoice
_NonemptyString = Annotated[str, Field(min_length=1)]
_Token = Annotated[int, Field(ge=0)]
class _GenerationRequest(_StrictModel):
model: _NonemptyString
max_tokens: int | None = Field(default=None, ge=1)
max_completion_tokens: int | None = Field(default=None, ge=1)
temperature: float = Field(default=0.0, ge=0.0, le=2.0)
top_p: float = Field(default=1.0, gt=0.0, le=1.0)
top_k: int = Field(default=0, ge=0)
seed: int = Field(default=0, ge=0)
stop: _NonemptyString | Annotated[list[_NonemptyString], Field(max_length=4)] | None = None
stream: bool = False
n: int = Field(default=1, ge=1, le=1)
stream_options: _StreamOptions | None = None
# CISM extension: greedy prompt-lookup speculation (0 = off).
spec_k: int = Field(default=0, ge=0, le=16)
@model_validator(mode="after")
def check_options(self) -> _GenerationRequest:
if self.max_tokens is not None and self.max_completion_tokens is not None:
raise ValueError("Specify only one token limit.")
if self.stream_options is not None and not self.stream:
raise ValueError("stream_options requires stream=true.")
return self
@property
def token_limit(self) -> int:
if self.max_completion_tokens is not None:
return self.max_completion_tokens
return self.max_tokens if self.max_tokens is not None else 128
class _CompletionRequest(_GenerationRequest):
prompt: str | Annotated[list[_Token], Field(min_length=1)]
# Exact-math token logprobs (top-N per emitted token, non-streaming only).
logprobs: int | None = Field(default=None, ge=0, le=20)
class _Message(_StrictModel):
role: Literal["system", "developer", "user", "assistant"]
content: str
class _ChatRequest(_GenerationRequest):
messages: Annotated[list[_Message], Field(min_length=1)]
response_format: _ResponseFormat | None = None
# Exact-math token logprobs (OpenAI chat shape, non-streaming only).
logprobs: bool = False
top_logprobs: int | None = Field(default=None, ge=0, le=20)
# Server-side tool loop (non-streaming only; registry tools only).
tools: Annotated[list[_ToolDef], Field(max_length=8)] | None = None
tool_choice: _ToolChoice | None = None
@model_validator(mode="after")
def check_a2_options(self) -> _ChatRequest:
if self.top_logprobs is not None and not self.logprobs:
raise ValueError("top_logprobs requires logprobs=true.")
if self.tool_choice is not None and not self.tools:
raise ValueError("tool_choice requires tools.")
return self
class _BatchRequest(_StrictModel):
"""Phase2 batched completions: list of prompts -> order-preserving results.
Single ``model`` for the whole batch (same-model same-quant fast path via
native ``BatchSession`` when available). ``prompts`` items mirror
``_CompletionRequest.prompt`` (text or token IDs). Sampling fields mirror
``_GenerationRequest`` except streaming (batch is non-streaming only);
``spec_k`` must be 0 because ``BatchSession`` has no speculation path.
"""
model: _NonemptyString
prompts: Annotated[list[str | Annotated[list[_Token], Field(min_length=1)]],
Field(min_length=1, max_length=64)]
max_tokens: int | None = Field(default=None, ge=1)
max_completion_tokens: int | None = Field(default=None, ge=1)
temperature: float = Field(default=0.0, ge=0.0, le=2.0)
top_p: float = Field(default=1.0, gt=0.0, le=1.0)
top_k: int = Field(default=0, ge=0)
seed: int = Field(default=0, ge=0)
stop: _NonemptyString | Annotated[list[_NonemptyString], Field(max_length=4)] | None = None
spec_k: int = Field(default=0, ge=0, le=16)
@model_validator(mode="after")
def check_options(self) -> _BatchRequest:
if self.max_tokens is not None and self.max_completion_tokens is not None:
raise ValueError("Specify only one token limit.")
return self
@property
def token_limit(self) -> int:
if self.max_completion_tokens is not None:
return self.max_completion_tokens
return self.max_tokens if self.max_tokens is not None else 128
class _LoadRequest(_StrictModel):
model: _NonemptyString
precision: Literal["fp32", "int8", "hybrid-int4", "hybrid-fp4"] = "int8"
act_precision: Literal["fp32", "int8"] | None = None
revision: str | None = None
threads: int = Field(default=2, ge=1, le=64)
trust_remote_code: bool = False
class _UnloadRequest(_StrictModel):
model: _NonemptyString
class _CacheRequest(_StrictModel):
model: _NonemptyString
decode_tokens: int = Field(default=64, ge=1, le=2048)
scrub_mb: int = Field(default=96, ge=1, le=1024)
class _TokenizeRequest(_StrictModel):
model: _NonemptyString
text: _NonemptyString
class _NllRequest(_StrictModel):
model: _NonemptyString
text: _NonemptyString
window: int = Field(default=128, ge=2, le=2048)
class _RuntimeRequest(_StrictModel):
model: _NonemptyString
lock_pages: bool | None = None
keep_warm: float | None = None
class _AutoBenchRequest(_StrictModel):
model: _NonemptyString
precision: Literal["fp32", "int8", "hybrid-int4", "hybrid-fp4"] = "int8"
threads: int = Field(default=2, ge=1, le=64)
max_tokens: int = Field(default=512, ge=1, le=2048)
thread_sweep: bool = False
spec_sweep: bool = False
bench_act_int8: bool = False
include_torch: bool = True
compile: bool = False
class _APIError(Exception):
def __init__(
self, status: int, message: str, code: str, param: str | None = None
) -> None:
super().__init__(message)
self.status = status
error_type = {
401: "authentication_error",
429: "rate_limit_error",
503: "server_error",
}.get(status, "invalid_request_error")
self.body = {"error": {
"message": message, "type": error_type, "param": param, "code": code,
}}
@dataclass
class _Chunk:
text: str
finish_reason: str | None
prompt_tokens: int
completion_tokens: int
spec_proposed: int = 0
spec_accepted: int = 0
# Full decoded token IDs behind this chunk (empty when the backend does
# not report them). Used only by opt-in exact-math paths (logprobs).
token_ids: list[int] = field(default_factory=list)
# Accumulated OpenAI-style tool calls for tool-loop jobs.
tool_calls: list[dict] = field(default_factory=list)
# Stream S3 automatic prefix-cache outcome for this job's prefill.
prefix_hit: bool = False
prefix_len: int = 0
@dataclass(eq=False)
class _Job:
body: _CompletionRequest | _ChatRequest
engine: _Engine | None = None
output: queue.Queue = field(default_factory=lambda: queue.Queue(maxsize=8))
cancel: threading.Event = field(default_factory=threading.Event)
finished: bool = False
consumed: bool = False
output_ready: asyncio.Event = field(default_factory=asyncio.Event)
notification_lock: threading.Lock = field(default_factory=threading.Lock)
notification_pending: bool = False
notified: bool = False
# Worker-side arrival time (monotonic) for the S1 starvation guard.
enqueued_at: float = 0.0
def notify(self, loop: asyncio.AbstractEventLoop) -> None:
with self.notification_lock:
if self.notification_pending or self.notified:
return
self.notification_pending = True
try:
loop.call_soon_threadsafe(self._wake)
except RuntimeError:
# Shutdown can close the loop before a native call returns.
self.notification_pending = False
self.cancel.set()
def _wake(self) -> None:
with self.notification_lock:
self.notification_pending = False
self.notified = True
self.output_ready.set()
@dataclass
class _ContState:
"""Per-job mutable state for one S1 continuous epoch (worker thread only)."""
job: _Job
engine: Any
prompt_ids: list[int]
stop: list[str] | None
eff: int # effective generation budget: min(token_limit, ctx - plen)
context: int
arrival: float
admittable: bool
# S1 epoch kind: "greedy" (mid-decode joins allowed -- seed-independent),
# "sampled" (form-only epochs -- per-seq RNG state cannot survive a
# rebuild), or "serial" (never batched).
kind: str = "serial"
gen: list[int] = field(default_factory=list)
@property
def plen(self) -> int:
return len(self.prompt_ids)
@property
def remaining(self) -> int:
return self.eff - len(self.gen)
class _Worker:
def __init__(self, engine: _Engine, capacity: int) -> None:
self.engine = engine
self.loop = asyncio.get_running_loop()
self.capacity = capacity
self.pending: queue.Queue[_Job] = queue.Queue(maxsize=capacity)
self.lock = threading.Lock()
self.jobs: set[_Job] = set()
self.stopping = threading.Event()
# A non-cooperative native call cannot be killed safely by Python.
self.thread = threading.Thread(target=self._run, name="cism-inference", daemon=True)
@property
def ready(self) -> bool:
return self.thread.is_alive() and not self.stopping.is_set()
def submit(self, body: _CompletionRequest | _ChatRequest,
engine: _Engine | None = None) -> _Job:
with self.lock:
if not self.ready:
raise _APIError(503, "Server is not ready.", "server_unavailable")
if len(self.jobs) >= self.capacity:
raise _APIError(429, "Generation queue is full.", "queue_full")
job = _Job(body, engine=engine)
job.enqueued_at = time.monotonic()
self.jobs.add(job)
self.pending.put_nowait(job)
return job
def release(self, job: _Job) -> None:
job.cancel.set()
job.notify(self.loop)
with self.lock:
job.consumed = True
if job.finished:
self.jobs.discard(job)
def stop(self) -> None:
with self.lock:
self.stopping.set()
for job in self.jobs:
job.cancel.set()
job.notify(self.loop)
self.jobs.clear()
def _put(self, job: _Job, value: _Chunk | _APIError) -> bool:
while not job.cancel.is_set() and not self.stopping.is_set():
try:
job.output.put(value, timeout=0.05)
job.notify(self.loop)
return True
except queue.Full:
pass
return False
def _run(self) -> None:
while True:
try:
first = self.pending.get(timeout=0.05)
except queue.Empty:
if self.stopping.is_set():
return
continue
takes = 1
try:
if first.cancel.is_set() or self.stopping.is_set():
self._finish(first)
elif self._use_continuous(first):
takes = self._run_continuous(first)
else:
batch = self._collect_batch(first)
takes = len(batch)
if len(batch) > 1:
self._run_micro_batch(batch)
else:
self._run_single(first)
finally:
for _ in range(takes):
self.pending.task_done()
def _use_continuous(self, job: _Job) -> bool:
"""True when ``job`` heads an S1 iteration-level epoch.
Requires both kill switches on (``CISM_SERVE_CONTINUOUS``,
``CISM_SERVE_MICROBATCH``), a batch-eligible body (non-streaming
greedy, or non-streaming sampled with valid params -- same
exclusions otherwise: no spec/A2 opt-ins), and a dense engine with
a steppable native ``BatchSession``. Everything else keeps the
exact A1 FIFO path, so streaming/speculative/Surjo/test-double
jobs are bit-for-bit unchanged. Non-admittable heads inside the
loop still run on the serial ``_run_single`` path.
"""
if not _serve_continuous_enabled() or not _serve_microbatch_enabled():
return False
if not (_micro_batch_eligible(job.body)
or _sampled_batch_eligible(job.body)):
return False
engine = job.engine or self.engine
if engine is None or _is_pool_or_batch_engine(engine):
return False
return _can_native_batch(engine)
def _collect_batch(self, first: _Job) -> list[_Job]:
"""Drain co-queued jobs without waiting (zero added latency).
A lone job returns alone and runs immediately on the single path, so
single-stream latency/tokens-s are unregressed by construction. Only
when jobs are already co-queued under concurrent load is a micro-batch
formed, capped at ``_MICRO_BATCH_MAX``.
"""
batch = [first]
if not _serve_microbatch_enabled():
return batch
while len(batch) < _MICRO_BATCH_MAX:
try:
batch.append(self.pending.get_nowait())
except queue.Empty:
break
return batch
def _finish(self, job: _Job) -> None:
with self.lock:
job.finished = True
if job.consumed:
self.jobs.discard(job)
def _run_single(self, job: _Job) -> None:
iterator = None
engine = job.engine or self.engine
try:
if job.cancel.is_set() or self.stopping.is_set():
return
body = job.body
try:
if isinstance(body, _ChatRequest):
prompt = engine.format_chat(
[message.model_dump() for message in body.messages]
)
tokens = engine.encode(prompt)
else:
tokens = engine.encode(body.prompt)
except (ValueError, TypeError):
raise _APIError(400, "Invalid prompt or chat messages.", "invalid_prompt") from None
if len(tokens) > engine.context_length:
raise _APIError(400, "Prompt exceeds the model context length.",
"context_length_exceeded")
if job.cancel.is_set() or self.stopping.is_set():
return
stop = [body.stop] if isinstance(body.stop, str) else body.stop
# Stream A2 opt-in jobs run serially on exact-math paths (never
# micro-batched, never paged): constrained JSON over engine.logits
# and multi-turn tool loops over engine.stream.
if isinstance(body, _ChatRequest) and _needs_tools(body):
self._run_tool_loop(job, engine, tokens, body, stop)
return
if _needs_constrained(body):
assert isinstance(body, _ChatRequest)
self._run_constrained_job(job, engine, tokens, body)
return
if body.spec_k == 0 and _serve_paged_enabled() and _can_paged(engine):
# Paged-KV single path: bitwise-identical to dense Session.
iterator = iter(_stream_paged(
engine, tokens, max_new_tokens=body.token_limit,
temperature=body.temperature, top_p=body.top_p,
top_k=body.top_k, seed=body.seed, stop=stop,
cancel=job.cancel, chunk_size=16,
))
else:
iterator = iter(engine.stream(
tokens, max_new_tokens=body.token_limit,
temperature=body.temperature, top_p=body.top_p,
top_k=body.top_k, seed=body.seed, stop=stop,
cancel=job.cancel, chunk_size=16,
spec_k=body.spec_k,
))
terminal = False
for chunk in iterator:
if not self._put(job, _Chunk(
chunk.text, chunk.finish_reason,
chunk.prompt_tokens, chunk.completion_tokens,
getattr(chunk, "spec_proposed", 0) or 0,
getattr(chunk, "spec_accepted", 0) or 0,
token_ids=list(getattr(chunk, "token_ids", None) or []),
prefix_hit=bool(getattr(chunk, "prefix_hit", False)),
prefix_len=int(getattr(chunk, "prefix_len", 0) or 0),
)):
break
if chunk.finish_reason is not None:
terminal = True
break
if not terminal and not job.cancel.is_set():
raise _APIError(503, "Generation did not complete.", "generation_failed")
except _APIError as error:
self._put(job, error)
except Exception:
self._put(job, _APIError(503, "Generation failed.", "generation_failed"))
finally:
if iterator is not None:
with suppress(Exception):
close = getattr(iterator, "close", None)
if close is not None:
close()
self._finish(job)
def _run_constrained_job(self, job: _Job, engine: Any,
tokens: list[int], body: _ChatRequest) -> None:
"""One greedy constrained-JSON turn over real engine logits."""
try:
if job.cancel.is_set() or self.stopping.is_set():
return
target = _constrained_target(body)
text, ids, finish, prompt_n, gen_n = _run_constrained_sync(
engine, tokens, token_limit=body.token_limit,
target=target, cancel=job.cancel)
self._put(job, _Chunk(text, finish, prompt_n, gen_n,
token_ids=list(ids)))
except _APIError as error:
self._put(job, error)
except Exception:
self._put(job, _APIError(503, "Generation failed.", "generation_failed"))
finally:
self._finish(job)
def _collect_turn(self, job: _Job, engine: Any, tokens: list[int],
body: _ChatRequest, stop: list[str] | None) -> tuple[str, list[int], str | None, int, int] | None:
"""Collect one dense ``engine.stream`` turn fully in memory.
Tool-loop turns (intermediate and final) never publish incrementally:
the caller publishes exactly one synthetic terminal chunk, so
``generate()`` observes the normal single-terminal protocol with the
accumulated ``tool_calls`` attached. Returns
(text, ids, finish, prompt_n, gen_n), or None when cancelled.
"""
iterator = None
try:
iterator = iter(engine.stream(
tokens, max_new_tokens=body.token_limit,
temperature=body.temperature, top_p=body.top_p,
top_k=body.top_k, seed=body.seed, stop=stop,
cancel=job.cancel, chunk_size=16, spec_k=body.spec_k,
))
parts: list[str] = []
ids: list[int] = []
last: Any = None
terminal = False
for chunk in iterator:
if job.cancel.is_set() or self.stopping.is_set():
return None
ids.extend(list(getattr(chunk, "token_ids", None) or []))
parts.append(chunk.text)
last = chunk
if chunk.finish_reason is not None:
terminal = True
break
if not terminal or last is None:
raise _APIError(503, "Generation did not complete.", "generation_failed")
return "".join(parts), ids, last.finish_reason, last.prompt_tokens, last.completion_tokens
finally:
if iterator is not None:
with suppress(Exception):
close = getattr(iterator, "close", None)
if close is not None:
close()
def _run_tool_loop(self, job: _Job, engine: Any, tokens: list[int],
body: _ChatRequest, stop: list[str] | None) -> None:
"""Bounded server-side tool loop with the real engine.
Each turn generates assistant text with ``engine.stream``; text that
parses as a ``{"name", "arguments"}`` call for a declared registry
tool is validated and executed for real, and the JSON result is fed
back as a follow-up user turn. At most ``_TOOL_MAX_CALLS`` executions,
so the loop always terminates. Every turn is collected in memory and
exactly one synthetic terminal chunk is published (with the
accumulated ``tool_calls`` attached), so ``generate()`` observes the
normal single-terminal protocol.
"""
try:
assert body.tools is not None
declared = {tool.function.name: tool.function for tool in body.tools}
choice = body.tool_choice
required_name = (choice.function.name
if isinstance(choice, _NamedToolChoice) else None)
require_call = choice == "required" or required_name is not None
messages = [message.model_dump() for message in body.messages]
tool_calls: list[dict] = []
final: tuple[str, list[int], str | None, int, int] | None = None
pending_call: tuple[str, dict] | None = None
current_tokens = list(tokens)
for turn in range(_TOOL_MAX_CALLS + 1):
if job.cancel.is_set() or self.stopping.is_set():
return
last_turn = turn == _TOOL_MAX_CALLS
turn_result = self._collect_turn(
job, engine, current_tokens, body, stop)
if turn_result is None:
return
text, _ids, finish, _pn, _gn = turn_result
parsed = _parse_tool_call(text)
if parsed is not None:
name, arguments = parsed
if name not in declared or (required_name is not None and name != required_name):
parsed = None
if parsed is None or last_turn:
if parsed is not None:
pending_call = parsed
final = turn_result
break
name, arguments = parsed
function = declared[name]
params_schema = (function.parameters if function.parameters is not None
else _TOOL_REGISTRY[name][1])
if _schema_support_error(params_schema, "tools.parameters") is not None:
result: dict[str, Any] = {
"error": "unsupported parameters schema for this tool"}
elif not _mini_validate(arguments, params_schema):
result = {"error": "arguments do not match the tool parameters schema"}
else:
try:
result = dict(_TOOL_REGISTRY[name][0](arguments))
except ValueError as exc:
result = {"error": str(exc)}
except Exception:
result = {"error": "tool execution failed"}
tool_calls.append({
"id": f"call_{len(tool_calls)}",
"type": "function",
"function": {"name": name, "arguments": json.dumps(
arguments, ensure_ascii=True, allow_nan=False,
separators=(",", ":"))},
})
messages = [*messages,
{"role": "assistant", "content": text},
{"role": "user", "content": (
f"Tool {name} returned: {json.dumps(result, ensure_ascii=True, allow_nan=False)}")}]
try:
current_tokens = engine.encode(
engine.format_chat(messages))
except (ValueError, TypeError):
raise _APIError(400, "Invalid prompt or chat messages.",
"invalid_prompt") from None
if len(current_tokens) > engine.context_length:
raise _APIError(400, "Prompt exceeds the model context length.",
"context_length_exceeded")
assert final is not None
text, ids, finish, prompt_n, gen_n = final
if pending_call is not None and require_call:
# Budget spent while the model still wants tools: report the
# pending call honestly instead of looping forever.
name, arguments = pending_call
tool_calls.append({
"id": f"call_{len(tool_calls)}",
"type": "function",
"function": {"name": name, "arguments": json.dumps(
arguments, ensure_ascii=True, allow_nan=False,
separators=(",", ":"))},
})
finish = "tool_calls"
self._put(job, _Chunk(text, finish, prompt_n, gen_n,
token_ids=list(ids),
tool_calls=tool_calls))
except _APIError as error:
self._put(job, error)
except Exception:
self._put(job, _APIError(503, "Generation failed.", "generation_failed"))
finally:
self._finish(job)
def _run_micro_batch(self, batch: list[_Job]) -> None:
"""FIFO continuous-batching step for co-queued greedy jobs (v0).
Policy: strict FIFO; the head job must be micro-batch eligible
(non-streaming, greedy, ``spec_k == 0``) on a dense engine with a
native single-weight-stream batch path. The eligible FIFO prefix with
the head's ``(stop, token_limit)`` key runs as one ``BatchSession``
micro-batch (chunked further by ``ContinuousSchedulerV0`` budgets when
configured on the worker); every other job runs serially in order.
Greedy decode is seed-independent, so coalescing is bitwise-identical
to serial execution. Per-job validation errors never poison the batch.
"""
try:
head = batch[0]
engine = head.engine or self.engine
key = _micro_batch_key(head.body) if _micro_batch_eligible(head.body) else None
group: list[_Job] = []
rest: list[_Job] = []
if (key is not None and not _is_pool_or_batch_engine(engine)
and _can_native_batch(engine)):
for job in batch:
if ((job.engine or self.engine) is engine
and _micro_batch_eligible(job.body)
and _micro_batch_key(job.body) == key):
group.append(job)
else:
rest.append(job)
else:
rest = list(batch)
if len(group) < 2:
rest = batch
group = []
for job in group:
if job.cancel.is_set() or self.stopping.is_set():
self._finish(job)
group = [job for job in group
if not job.cancel.is_set() and not self.stopping.is_set()]
if group:
self._run_batch_group(engine, group)
for job in rest:
if job.cancel.is_set() or self.stopping.is_set():
self._finish(job)
else:
self._run_single(job)
finally:
# Jobs finished inside _run_single/_run_batch_group are already
# marked; mark any leftovers (cancelled heads) here.
for job in batch:
with self.lock:
if not job.finished:
job.finished = True
if job.consumed:
self.jobs.discard(job)
def _run_batch_group(self, engine: Any, group: list[_Job]) -> None:
"""Run one compatible FIFO group as a single-weight-stream batch."""
body0 = group[0].body
stop = [body0.stop] if isinstance(body0.stop, str) else (list(body0.stop) if body0.stop else None)
token_limit = body0.token_limit
cfg = getattr(self, "scheduler_config", None) or {}
# Per-job encode/validation first: one bad prompt must not fail others.
valid: list[tuple[_Job, Any]] = []
for job in group:
body = job.body
try:
try:
if isinstance(body, _ChatRequest):
prompt = engine.format_chat(
[message.model_dump() for message in body.messages])
prompt_arg: Any = engine.encode(prompt)
else:
prompt_arg = body.prompt
probe = engine.encode(prompt_arg)
if len(probe) > engine.context_length:
raise _APIError(400, "Prompt exceeds the model context length.",
"context_length_exceeded")
if isinstance(prompt_arg, list) and len(prompt_arg) > engine.context_length:
raise _APIError(400, "Prompt exceeds the model context length.",
"context_length_exceeded")
except _APIError:
raise
except (ValueError, TypeError):
raise _APIError(400, "Invalid prompt or chat messages.", "invalid_prompt") from None
valid.append((job, prompt_arg))
except _APIError as error:
self._put(job, error)
self._finish(job)
except Exception:
self._put(job, _APIError(503, "Generation failed.", "generation_failed"))
self._finish(job)
if not valid:
return
try:
from .paging import ContinuousSchedulerV0
sched = ContinuousSchedulerV0(
num_blocks=int(cfg.get("scheduler_blocks", 512)), block_size=16,
max_batch_tokens=int(cfg.get("max_batch_tokens", 4096)),
prefill_chunk_tokens=int(cfg.get("prefill_chunk_tokens", 512)))
lens: list[int] = []
for job, prompt_arg in valid:
try:
lens.append(len(engine.encode(prompt_arg)))
except (ValueError, TypeError):
lens.append(token_limit)
for idx, length in enumerate(lens):
sched.add(idx, length, token_limit, stop=stop)
chunks = sched.plan_chunks()
except Exception:
chunks = [[{"req_id": i} for i in range(len(valid))]]
sched = None # type: ignore[assignment]
ordered: dict[int, tuple] = {}
failed: dict[int, _APIError] = {}
prompts = [prompt_arg for _, prompt_arg in valid]
try:
for chunk in chunks:
indices = [entry["req_id"] for entry in chunk]
sub = [prompts[i] for i in indices]
try:
rows = _run_batch_chunk_sync(
engine, sub, token_limit, 0.0, 1.0, 0, body0.seed, stop)
except _APIError as error:
for i in indices:
failed[i] = error
continue
except (ValueError, TypeError) as exc:
msg = str(exc)
error = _APIError(
400, "Prompt exceeds the model context length.",
"context_length_exceeded") if (
"context" in msg.lower() or "Prompt length" in msg) else _APIError(
400, "Invalid prompt or chat messages.", "invalid_prompt")
for i in indices:
failed[i] = error
continue
except Exception:
error = _APIError(503, "Generation failed.", "generation_failed")
for i in indices:
failed[i] = error
continue
if len(rows) != len(indices):
error = _APIError(503, "Generation did not complete.", "generation_failed")
for i in indices:
failed[i] = error
continue
for i, row in zip(indices, rows):
ordered[i] = row
if sched is not None:
for entry in chunk:
try:
sched.free(entry["req_id"])
except Exception:
pass
finally:
for pos, (job, _) in enumerate(valid):
try:
if pos in failed:
self._put(job, failed[pos])
else:
row = ordered.get(pos)
if row is None:
self._put(job, _APIError(503, "Generation did not complete.",
"generation_failed"))
else:
text, _ids, finish, prompt_n, gen_n = row
if finish is None:
self._put(job, _APIError(503, "Generation did not complete.",
"generation_failed"))
else:
self._put(job, _Chunk(text, finish, prompt_n, gen_n,
token_ids=list(_ids)))
finally:
self._finish(job)
def _run_continuous(self, first: _Job) -> int:
"""Iteration-level scheduler loop for S1 (vLLM-style, Python-side).
ADD: each iteration drains new arrivals into an
:class:`paging.IterationScheduler` waiting queue (oldest-first) and
admits group-matching jobs into the active batch -- including
mid-decode joins, which rebuild the native ``BatchSession`` with
survivors' ``prompt + generated`` prefixes (greedy continuation is
an exact function of the prefix, so joined decoding is
bitwise-identical to an uninterrupted run; per-job remaining
budgets truncate with ``"length"`` exactly as a lone run would).
EVICT: one shared-session ``BatchSession`` weight-stream pass per
iteration (``next_tokens(1)`` while the batch is shared or waiters
are present, so finished/errored/cancelled leave immediately;
``min(16, remaining)`` chunks for a lone idle job, mirroring
``Engine.stream`` -- this keeps single-stream tokens/s at parity by
avoiding a Python GIL round-trip per token when nobody else needs
the batch; join latency is bounded by one chunk). Per-job
remaining budgets truncate with ``"length"`` exactly as a lone run
would.
CHUNKED PREFILL: admissions per rebuild are capped at
``prefill_chunk_tokens`` new prompt tokens (env
``CISM_SERVE_PREFILL_CHUNK``); an oversized single prompt is
deferred while an epoch is live and runs alone once it drains
(intra-prompt chunking still needs native support -- documented
limit). Epochs rotate after ``CISM_SERVE_EPOCH_STEPS`` steps or
when a same-engine waiter exceeds ``CISM_SERVE_MAX_WAIT_S``.
Non-admissible jobs (streaming/sampled/spec/A2, other engines) keep
global FIFO order: when the batch is empty and the head is not
admittable it runs on the serial ``_run_single`` path. Returns the
number of ``pending`` gets consumed (``first`` + drains).
"""
from .paging import ContinuousSchedulerV0, IterationScheduler
cfg = getattr(self, "scheduler_config", None) or {}
chunk = _continuous_prefill_chunk(
int(cfg.get("prefill_chunk_tokens", 512)))
budget = int(cfg.get("max_batch_tokens", 4096))
sched = IterationScheduler(
max_batch=_MICRO_BATCH_MAX, prefill_chunk_tokens=chunk,
max_wait_s=_continuous_max_wait_s(), max_batch_tokens=budget)
max_wait = sched.max_wait_s
join_defer = _continuous_join_defer_tokens()
states: dict[int, _ContState] = {}
takes = 0
def arrive(job: _Job) -> None:
if job.cancel.is_set() or self.stopping.is_set():
self._finish(job)
return
engine = job.engine or self.engine
if engine is None:
self._put(job, _APIError(503, "Server is not ready.",
"server_unavailable"))
self._finish(job)
return
try:
prompt_ids, stop = _cont_encode(engine, job.body)
except _APIError as error:
self._put(job, error)
self._finish(job)
return
except Exception:
self._put(job, _APIError(503, "Generation failed.",
"generation_failed"))
self._finish(job)
return
limit = job.body.token_limit
context = int(engine.context_length)
eff = min(limit, context - len(prompt_ids))
if eff <= 0:
# Prompt fills the context: zero new tokens, terminal
# "length" with empty text (mirrors native max_new=0).
self._put(job, _Chunk("", "length", len(prompt_ids), 0,
token_ids=[]))
self._finish(job)
return
admittable = _micro_batch_eligible(job.body)
if admittable:
kind = "greedy"
group = (id(engine), "greedy")
elif _sampled_batch_eligible(job.body):
kind = "sampled"
admittable = True
group = (id(engine), "sampled", job.body.temperature,
job.body.top_p, job.body.top_k)
else:
kind = "serial"
group = (id(engine), "serial")
states[id(job)] = _ContState(
job, engine, prompt_ids, stop, eff, context,
arrival=(job.enqueued_at or time.monotonic()),
admittable=admittable, kind=kind)
sched.add(id(job), len(prompt_ids), group,
states[id(job)].arrival, admittable=admittable)
def fits(rid: int, _admitted: list) -> bool:
# Context-window compatibility: with maxR = max remaining over
# the would-be active set, every member's extended prompt must
# still fit (native requires prompt + max_new <= context), and
# the token budget must hold (lone head exempt).
active = [states[r] for r in sched.running_ids]
cand = states[rid]
# Sampled RNG exactness: a sampled job never joins a live epoch
# (a rebuild would restart its RNG stream), and nothing joins a
# live sampled epoch. Sampled epochs form only on a drain.
if active and (cand.kind == "sampled"
or sched.group_of(sched.running_ids[0])[1] == "sampled"):
return False
if cand.plen + cand.eff > cand.context:
return False
pool = active + [cand]
max_r = max(s.remaining for s in pool)
if any(s.plen + len(s.gen) + max_r > s.context for s in pool):
return False
if len(pool) > 1 and sum(s.plen + s.eff for s in pool) > budget:
return False
if active and join_defer > 0:
starved = (time.monotonic() - cand.arrival) > max_wait
if not starved and all(s.remaining <= join_defer
for s in active):
return False # drain is imminent; join would only
# force a pointless recompute-prefill
return True
arrive(first)
takes += 1
session = None
slot_rids: list[int] = []
fresh = False # first step after a (re)build stays granular
epoch_steps = 0
epoch_cap = _continuous_epoch_steps()
try:
while len(sched) > 0:
if self.stopping.is_set():
if session is not None:
try:
session.cancel()
except Exception:
pass
session = None
break
while True:
try:
job = self.pending.get_nowait()
except queue.Empty:
break
takes += 1
arrive(job)
now = time.monotonic()
for rid in sched.waiting_ids:
state = states.get(rid)
if state is not None and state.job.cancel.is_set():
states.pop(rid, None)
try:
sched.drop(rid, now)
except Exception:
pass
self._finish(state.job)
if not sched.running_ids:
session = None
slot_rids = []
head = sched.head_id()
if head is None:
break
state = states[head]
if not state.admittable:
states.pop(head, None)
try:
sched.drop(head, now)
except Exception:
pass
self._run_single(state.job)
continue
admitted, _ = sched.admit(now, fits=fits)
if not admitted:
# Head admittable but unfittable: fail it alone so
# one bad job never wedges the queue (per-job
# isolation, as in the A1 batch group).
states.pop(head, None)
try:
sched.drop(head, now)
except Exception:
pass
self._put(state.job, _APIError(
503, "Generation failed.", "generation_failed"))
self._finish(state.job)
continue
session, slot_rids = self._cont_build(states, sched)
fresh = True
epoch_steps = 0
continue
admitted, rotate = sched.admit(now, fits=fits)
if admitted:
session, slot_rids = self._cont_build(
states, sched, session=session)
fresh = True
epoch_steps = 0
if rotate or epoch_steps >= epoch_cap:
running_kind = (sched.group_of(sched.running_ids[0])[1]
if sched.running_ids else "greedy")
if running_kind == "sampled":
# Never rotate a live sampled epoch: rotation
# requeues + rebuilds, restarting per-seq RNG
# streams and diverging from serial. Reset the
# counter and fall through to stepping -- sampled
# epochs always run to drain; waiters (even
# starved same-group ones) join the next epoch.
# (A bare `continue` here would livelock: rotation
# changes no state, so the live epoch would never
# step again.)
epoch_steps = 0
else:
if session is not None:
try:
session.cancel()
except Exception:
pass
session = None
slot_rids = []
fresh = False
sched.requeue_running()
epoch_steps = 0
continue
# Step size: granular (1 token) while the batch is shared
# or waiters are present (prompt evict/publish + fast
# joins); chunked (16, mirroring Engine.stream) for a lone
# idle job, so a lone request follows the same native call
# pattern as the dense path (avoids a GIL round-trip per
# token at parity).
if len(sched.running_ids) > 1 or sched.waiting_ids or fresh:
count = 1
else:
count = min(16, states[sched.running_ids[0]].remaining)
if self._cont_step(states, sched, session, slot_rids, count):
session = None # all drained; loop re-admits
slot_rids = []
else:
epoch_steps += 1
fresh = False
finally:
for rid in list(states):
state = states.pop(rid)
try:
self._put(state.job, _APIError(
503, "Generation failed.", "generation_failed"))
except Exception:
pass
try:
self._finish(state.job)
except Exception:
pass
return takes
def _cont_build(self, states: dict[int, _ContState], sched: Any,
session: Any = None) -> tuple[Any, list[int]]:
"""(Re)build the epoch's shared native session for mid-decode joins.
Prompts are survivors' ``prompt + generated`` prefixes plus new
arrivals' prompts; ``max_new`` is the max remaining budget (per-job
truncation keeps ``"length"`` exact). Returns ``(session,
slot_rids)`` where ``slot_rids`` is the build-time member order
(native slot ``i`` forever maps to ``slot_rids[i]`` until the next
rebuild). Postcondition: a ``None`` session implies empty running
state (every member failed, so the caller simply re-admits).
"""
if session is not None:
try:
session.cancel()
except Exception:
pass
engine = states[sched.running_ids[0]].engine
def fail_running(error: _APIError) -> tuple[Any, list[int]]:
for rid in sched.running_ids:
state = states.pop(rid, None)
if state is None:
continue
self._put(state.job, error)
self._finish(state.job)
for rid in list(sched.running_ids):
try:
sched.finish(rid, time.monotonic())
except Exception:
pass
return None, []
members = [rid for rid in sched.running_ids
if states[rid].remaining > 0]
if not members:
return fail_running(_APIError(503, "Generation did not complete.",
"generation_failed"))
prompts = [states[r].prompt_ids + states[r].gen for r in members]
max_r = max(states[r].remaining for r in members)
body0 = states[members[0]].job.body
eos = _engine_eos(engine)
# Epoch-homogeneous sampling params (group key): greedy epochs run
# temperature=0 (seed-independent, join-safe); sampled epochs carry
# the group's (temperature, top_p, top_k) with each job's own serial
# seed per sequence, so batched sampled decode is token-identical
# to serial execution.
if states[members[0]].kind == "sampled":
temperature, top_p, top_k = (body0.temperature, body0.top_p,
body0.top_k)
seq_seeds = [states[r].job.body.seed for r in members]
else:
temperature, top_p, top_k = 0.0, 1.0, 0
seq_seeds = []
seed = states[members[0]].job.body.seed
lock = getattr(engine, "_generation_lock", None)
if lock is not None:
lock.acquire()
try:
try:
fresh = engine._native.create_batch_session(
prompts, max_new_tokens=max_r, temperature=temperature,
top_p=top_p, top_k=top_k, seed=seed, eos_token_ids=eos,
seq_seeds=seq_seeds)
except (ValueError, TypeError) as exc:
message = str(exc)
if "context" in message.lower() or "KV cache" in message:
error = _APIError(400, "Prompt exceeds the model context length.",
"context_length_exceeded")
else:
error = _APIError(400, "Invalid prompt or chat messages.",
"invalid_prompt")
return fail_running(error)
except Exception:
return fail_running(_APIError(503, "Generation failed.",
"generation_failed"))
return fresh, list(members)
finally:
if lock is not None:
lock.release()
def _cont_step(self, states: dict[int, _ContState], sched: Any,
session: Any, slot_rids: list[int], count: int) -> bool:
"""Run one decode pass (``count`` tokens) over the active batch.
Native slot ``i`` maps to ``slot_rids[i]`` (build-time order, stable
across evictions since finished seqs are ``cancel_seq``-masked, not
removed). Finished / errored / cancelled sequences leave
immediately with exactly one terminal publish. Returns True when no
active member remains (epoch drained).
"""
from .paging import ContinuousSchedulerV0
alive = [states[r] for r in slot_rids if r in states]
if not alive:
return True
engine = alive[0].engine
lock = getattr(engine, "_generation_lock", None)
def drop(rid: int, value: _Chunk | _APIError | None) -> None:
state = states.pop(rid, None)
try:
sched.finish(rid, time.monotonic())
except Exception:
pass
if state is not None:
if value is not None:
self._put(state.job, value)
self._finish(state.job)
if lock is not None:
lock.acquire()
try:
try:
out = session.next_tokens(max(1, count))
reasons = list(session.finish_reasons)
except Exception:
for rid in slot_rids:
drop(rid, _APIError(503, "Generation failed.",
"generation_failed"))
return True
finally:
if lock is not None:
lock.release()
if len(out) != len(slot_rids) or len(reasons) != len(slot_rids):
for rid in slot_rids:
drop(rid, _APIError(503, "Generation did not complete.",
"generation_failed"))
return True
for index, rid in enumerate(slot_rids):
state = states.get(rid)
if state is None:
continue # evicted earlier; slot masked via cancel_seq
toks, reason = out[index], reasons[index]
if state.job.cancel.is_set() or self.stopping.is_set():
try:
session.cancel_seq(index)
except Exception:
pass
drop(rid, None)
continue
if reason == "cancelled":
drop(rid, None)
continue
state.gen.extend(list(toks or []))
finish = reason if reason in ("stop", "length") else None
if finish is None and state.remaining <= 0:
# Per-job budget spent while the shared session runs on:
# terminal "length", exactly as a lone run would report.
finish = "length"
if finish is None:
continue
try:
text = engine.tokenizer.decode(
list(state.gen), skip_special_tokens=True,
clean_up_tokenization_spaces=False)
except Exception:
try:
session.cancel_seq(index)
except Exception:
pass
drop(rid, _APIError(503, "Generation failed.",
"generation_failed"))
continue
if state.stop:
text, stop_finish = ContinuousSchedulerV0.apply_stop(
text, state.stop)
finish = stop_finish or finish
try:
session.cancel_seq(index)
except Exception:
pass
drop(rid, _Chunk(text, finish, state.plen, len(state.gen),
token_ids=list(state.gen)))
return not sched.running_ids
async def receive(
self, job: _Job, request: Request, *, streaming: bool = False
) -> _Chunk | _APIError:
# Before ASGI 2.4 StreamingResponse owns the disconnect listener.
spec = request.scope.get("asgi", {}).get("spec_version", "2.0")
poll_disconnect = not streaming or tuple(map(int, spec.split("."))) >= (2, 4)
while True:
if self.stopping.is_set():
raise _APIError(503, "Server is shutting down.", "server_unavailable")
if job.cancel.is_set():
raise _APIError(503, "Request disconnected.", "request_cancelled")
try:
return job.output.get_nowait()
except queue.Empty:
if not self.thread.is_alive():
raise _APIError(503, "Inference worker is unavailable.", "server_unavailable")
# Re-arm only when drained. Keep any scheduled callback marked as
# pending until it runs, even if the consumer drains a new burst.
# Thus neither a fast producer nor a slow client can queue callbacks
# without bound. Rechecking after clear prevents a lost wakeup.
with job.notification_lock:
job.notified = False
job.output_ready.clear()
try:
return job.output.get_nowait()
except queue.Empty:
pass
try:
async with asyncio.timeout(0.05):
await job.output_ready.wait()
except TimeoutError:
if poll_disconnect and await request.is_disconnected():
job.cancel.set()
raise _APIError(503, "Request disconnected.", "request_cancelled")
class _JobResponse(StreamingResponse):
def __init__(self, content, worker: _Worker, job: _Job) -> None:
super().__init__(content, media_type="text/event-stream", headers={
"Cache-Control": "no-cache", "X-Accel-Buffering": "no",
})
self.worker = worker
self.job = job
async def __call__(self, scope, receive, send) -> None:
# A send failure can leave the generator suspended at yield, so its
# finally block alone cannot reliably cancel inference on disconnect.
try:
await super().__call__(scope, receive, send)
finally:
self.worker.release(self.job)
await self.body_iterator.aclose()
def _usage(chunk: _Chunk) -> dict[str, int]:
proposed, accepted = chunk.spec_proposed, chunk.spec_accepted
return {
"prompt_tokens": chunk.prompt_tokens,
"completion_tokens": chunk.completion_tokens,
"total_tokens": chunk.prompt_tokens + chunk.completion_tokens,
"spec_proposed": proposed,
"spec_accepted": accepted,
"spec_accept_rate": round(accepted / proposed, 4) if proposed else None,
}
async def _rescore_logprobs(engine: Any, body: _CompletionRequest | _ChatRequest,
gen_ids: list[int], text: str) -> dict[str, Any]:
"""Rescore emitted IDs with real engine logits (thread-offloaded).
The prompt is re-encoded on the tokenizer-only path (identical to the
worker's encoding); scoring itself runs in a worker thread since native
``logits()`` calls block. Raises _APIError on any failure -- logprobs
are never synthesized.
"""
top_n = _logprobs_top_n(body)
chat = isinstance(body, _ChatRequest)
if text and not gen_ids:
raise _APIError(503, "Engine did not report token IDs; logprobs unavailable.",
"generation_failed")
def _run() -> dict[str, Any]:
try:
if chat:
assert isinstance(body, _ChatRequest)
prompt_ids = engine.encode(engine.format_chat(
[message.model_dump() for message in body.messages]))
else:
assert isinstance(body, _CompletionRequest)
prompt_ids = engine.encode(body.prompt)
except (ValueError, TypeError):
raise _APIError(400, "Invalid prompt or chat messages.", "invalid_prompt") from None
if gen_ids and not all(isinstance(t, int) for t in gen_ids):
raise _APIError(503, "Engine reported invalid token IDs.", "generation_failed")
surfaces, chosen, tops = _score_token_logprobs(engine, prompt_ids, gen_ids, top_n)
if chat:
return {"content": [
{"token": token, "logprob": logprob,
"top_logprobs": [{"token": t, "logprob": lp} for t, lp in top]}
for token, logprob, top in zip(surfaces, chosen, tops)]}
return {"tokens": surfaces, "token_logprobs": chosen,
"top_logprobs": [{token: logprob for token, logprob in top}
for top in tops]}
return await asyncio.to_thread(_run)
# -- Stream A2: exact-math opt-in paths (logprobs / JSON / tools) -------------
#: Max tool executions per chat request; the loop always terminates.
_TOOL_MAX_CALLS = 4
#: Documented JSON-schema subset for constrained decoding and tool args.
_SCHEMA_TYPES = frozenset({"object", "array", "string", "integer", "number",
"boolean", "null"})
_SCHEMA_VALUE_KEYS = frozenset({"type", "properties", "required",
"additionalProperties", "items", "enum", "const"})
def _wants_logprobs(body: _CompletionRequest | _ChatRequest) -> bool:
if isinstance(body, _ChatRequest):
return bool(body.logprobs)
return body.logprobs is not None
def _logprobs_top_n(body: _CompletionRequest | _ChatRequest) -> int:
if isinstance(body, _ChatRequest):
return body.top_logprobs or 0
return body.logprobs or 0
def _needs_constrained(body: _CompletionRequest | _ChatRequest) -> bool:
fmt = getattr(body, "response_format", None)
return fmt is not None and getattr(fmt, "type", "text") != "text"
def _constrained_target(body: _CompletionRequest | _ChatRequest) -> Any:
"""None (plain text), True (any JSON), or a schema dict."""
fmt = getattr(body, "response_format", None)
if fmt is None or getattr(fmt, "type", "text") == "text":
return None
if getattr(fmt, "type", None) == "json_object":
return True
return fmt.json_schema.schema # type: ignore[union-attr]
def _needs_tools(body: _CompletionRequest | _ChatRequest) -> bool:
tools = getattr(body, "tools", None)
if not tools:
return False
return getattr(body, "tool_choice", None) != "none"
# -- minimal JSON-schema subset (validate + support check) --------------------
def _schema_support_error(schema: Any, where: str = "schema") -> str | None:
"""None when ``schema`` is inside the supported subset, else a reason.
Supported: type (object/array/string/integer/number/boolean/null, or a
list thereof), properties, required, additionalProperties (bool), items,
enum (JSON scalars), const. Everything else ($ref, allOf/anyOf/oneOf,
pattern, format, length/range keywords, ...) is rejected so the server
never silently ignores a constraint.
"""
if not isinstance(schema, dict):
return f"{where} must be an object"
for key in schema:
if key not in _SCHEMA_VALUE_KEYS:
return f"unsupported keyword {key!r} at {where}"
kind = schema.get("type", None)
if kind is not None:
kinds = kind if isinstance(kind, list) else [kind]
if not kinds or any(k not in _SCHEMA_TYPES for k in kinds):
return f"unsupported type at {where}"
if "properties" in schema:
props = schema["properties"]
if not isinstance(props, dict):
return f"properties must be an object at {where}"
for name, sub in props.items():
reason = _schema_support_error(sub, f"{where}.properties.{name}")
if reason is not None:
return reason
if "required" in schema:
required = schema["required"]
if (not isinstance(required, list) or not required
or any(not isinstance(name, str) for name in required)):
return f"required must be a nonempty string list at {where}"
if "additionalProperties" in schema and not isinstance(
schema["additionalProperties"], bool):
return f"additionalProperties must be a boolean at {where}"
if "items" in schema:
reason = _schema_support_error(schema["items"], f"{where}.items")
if reason is not None:
return reason
if "enum" in schema:
enum = schema["enum"]
if (not isinstance(enum, list) or not enum
or any(isinstance(v, (dict, list)) for v in enum)):
return f"enum must be a nonempty scalar list at {where}"
return None
def _mini_validate(value: Any, schema: Any) -> bool:
"""Validate ``value`` against the supported schema subset."""
if not isinstance(schema, dict):
return False
if "const" in schema and value != schema["const"]:
return False
if "enum" in schema and value not in schema["enum"]:
return False
kind = schema.get("type", None)
if kind is not None:
kinds = set(kind) if isinstance(kind, list) else {kind}
if not any(_mini_type_ok(value, k) for k in kinds):
return False
if len(kinds) == 1:
kind = next(iter(kinds))
else:
kind = None
if kind == "object" or (kind is None and "properties" in schema):
if not isinstance(value, dict):
return kind is None
props = schema.get("properties", {})
for name, sub in props.items():
if name in value and not _mini_validate(value[name], sub):
return False
for name in schema.get("required", []):
if name not in value:
return False
if schema.get("additionalProperties") is False:
if any(name not in props for name in value):
return False
return True
if kind == "array" or (kind is None and "items" in schema):
if not isinstance(value, list):
return kind is None
items = schema.get("items")
return items is None or all(_mini_validate(v, items) for v in value)
return True
def _mini_type_ok(value: Any, kind: str) -> bool:
if kind == "null":
return value is None
if kind == "boolean":
return isinstance(value, bool)
if kind == "integer":
return isinstance(value, int) and not isinstance(value, bool)
if kind == "number":
return isinstance(value, (int, float)) and not isinstance(value, bool)
if kind == "string":
return isinstance(value, str)
if kind == "array":
return isinstance(value, list)
if kind == "object":
return isinstance(value, dict)
return False
# -- JSON prefix viability (structural + top-level key allowlist) -------------
def _skip_ws(text: str, i: int) -> int:
while i < len(text) and text[i] in " \t\n\r":
i += 1
return i
def _parse_string(text: str, i: int) -> tuple[str, int]:
"""Parse text[i] == '\"'. Returns (ok|partial|bad, next index)."""
n = len(text)
j = i + 1
while j < n:
c = text[j]
if c == '"':
return "ok", j + 1
if c == "\\":
j += 1
if j >= n:
return "partial", j
e = text[j]
if e not in "\"\\/bfnrtu":
return "bad", j
if e == "u":
hex_digits = text[j + 1:j + 5]
if len(hex_digits) < 4:
rest = text[j + 1:]
if all(ch in "0123456789abcdefABCDEF" for ch in rest):
return "partial", n
return "bad", j
if any(ch not in "0123456789abcdefABCDEF" for ch in hex_digits):
return "bad", j
j += 4
j += 1
continue
if ord(c) < 0x20:
return "bad", j
j += 1
return "partial", n
def _parse_literal(text: str, i: int, word: str) -> tuple[str, int]:
rest = text[i:i + len(word)]
if rest == word:
return "ok", i + len(word)
if word.startswith(rest) and i + len(rest) >= len(text):
return "partial", len(text)
return "bad", i
def _parse_number(text: str, i: int) -> tuple[str, int]:
n = len(text)
j = i
if j < n and text[j] == "-":
j += 1
digits = 0
while j < n and text[j].isdigit():
j += 1
digits += 1
if digits == 0:
# Lone "-" (or empty) is a partial number, anything else is bad.
return ("partial", j) if j > i else ("bad", i)
if j < n and text[j] == ".":
j += 1
frac = 0
while j < n and text[j].isdigit():
j += 1
frac += 1
if frac == 0:
return "partial", j
if j < n and text[j] in "eE":
j += 1
if j < n and text[j] in "+-":
j += 1
exp = 0
while j < n and text[j].isdigit():
j += 1
exp += 1
if exp == 0:
return "partial", j
return "ok", j
def _parse_value(text: str, i: int) -> tuple[str, int]:
"""Returns (ok|partial|bad, next index) for one JSON value at i."""
n = len(text)
i = _skip_ws(text, i)
if i >= n:
return "partial", i
c = text[i]
if c == "{":
i += 1
i = _skip_ws(text, i)
if i >= n:
return "partial", i
if text[i] == "}":
return "ok", i + 1
while True:
i = _skip_ws(text, i)
if i >= n:
return "partial", i
if text[i] != '"':
return "bad", i
status, i = _parse_string(text, i)
if status != "ok":
return status, i
i = _skip_ws(text, i)
if i >= n or text[i] != ":":
return ("partial", i) if i >= n else ("bad", i)
i += 1
status, i = _parse_value(text, i)
if status != "ok":
return status, i
i = _skip_ws(text, i)
if i >= n:
return "partial", i
if text[i] == ",":
i += 1
continue
if text[i] == "}":
return "ok", i + 1
return "bad", i
if c == "[":
i += 1
i = _skip_ws(text, i)
if i >= n:
return "partial", i
if text[i] == "]":
return "ok", i + 1
while True:
status, i = _parse_value(text, i)
if status != "ok":
return status, i
i = _skip_ws(text, i)
if i >= n:
return "partial", i
if text[i] == ",":
i += 1
continue
if text[i] == "]":
return "ok", i + 1
return "bad", i
if c == '"':
return _parse_string(text, i)
if c == "t":
return _parse_literal(text, i, "true")
if c == "f":
return _parse_literal(text, i, "false")
if c == "n":
return _parse_literal(text, i, "null")
if c == "-" or c.isdigit():
return _parse_number(text, i)
return "bad", i
def _json_prefix_viable(text: str) -> bool:
"""True when ``text`` is a strict prefix of (or equal to) some valid JSON."""
status, i = _parse_value(text, 0)
if status == "bad":
return False
if status == "partial":
return True
return _skip_ws(text, i) >= len(text)
def _top_level_keys(text: str) -> list[str]:
"""Complete top-level object keys visible in a (possibly partial) prefix."""
keys: list[str] = []
i = _skip_ws(text, 0)
n = len(text)
if i >= n or text[i] != "{":
return keys
depth = 0
while i < n:
c = text[i]
if c in "{[":
depth += 1
i += 1
elif c in "}]":
depth -= 1
i += 1
if depth <= 0:
break
elif c == '"':
status, j = _parse_string(text, i)
if status != "ok":
break
if depth == 1:
k = _skip_ws(text, j)
if k < n and text[k] == ":":
try:
keys.append(json.loads(text[i:j]))
except ValueError:
pass
i = j
else:
i += 1
return keys
def _schema_prefix_ok(text: str, schema: Any) -> bool:
"""Top-level container/key guidance for the supported schema subset."""
if not isinstance(schema, dict):
return True
kind = schema.get("type", None)
expected: str | None = None
if kind == "object" or (kind is None and "properties" in schema):
expected = "object"
elif kind == "array" or (kind is None and "items" in schema):
expected = "array"
stripped = text.lstrip()
if not stripped:
return True
if expected == "object":
if not stripped.startswith("{"):
return False
if schema.get("additionalProperties") is False and isinstance(
schema.get("properties"), dict):
allowed = set(schema["properties"])
if any(key not in allowed for key in _top_level_keys(text)):
return False
elif expected == "array":
if not stripped.startswith("["):
return False
return True
def _prefix_viable(text: str, schema: Any | None) -> bool:
if not _json_prefix_viable(text):
return False
return schema is None or _schema_prefix_ok(text, schema)
def _complete_and_valid(text: str, target: Any) -> bool:
"""True only for a complete document that validates (schema or any JSON)."""
try:
value = json.loads(text)
except ValueError:
return False
if target is True or target is None:
return True
return _mini_validate(value, target)
# -- exact log-softmax rescoring over real engine logits ----------------------
def _log_softmax_rows(logits: Any) -> list[float]:
flat = list(logits)
if not flat:
raise ValueError("empty logits")
peak = max(float(v) for v in flat)
if not math.isfinite(peak):
raise ValueError("non-finite logits")
total = math.fsum(math.exp(float(v) - peak) for v in flat)
offset = peak + math.log(total)
return [float(v) - offset for v in flat]
def _argsort_desc(values: Any) -> list[int]:
try:
import numpy as _np # hard dependency (pyproject); exact float64 path
except ImportError: # pragma: no cover - portable fallback
_np = None # type: ignore[assignment]
if _np is not None:
arr = _np.asarray(list(values), dtype=_np.float64)
return [int(i) for i in _np.argsort(arr, kind="stable")[::-1]]
order = list(range(len(list(values))))
order.sort(key=lambda i: float(list(values)[i]), reverse=True)
return order
def _score_token_logprobs(engine: Any, prompt_ids: list[int],
gen_ids: list[int], top_n: int) -> tuple[list[str], list[float], list[list[tuple[str, float]]]]:
"""Exact per-token logprobs: real ``engine.logits`` + float64 log-softmax.
``prompt_ids`` is the exact encoded prompt; ``gen_ids`` the emitted IDs.
Returns (token surfaces, chosen logprobs, top-N lists) in order.
"""
logits_fn = getattr(engine, "logits", None)
if not callable(logits_fn):
raise _APIError(400, "This engine does not expose logits; logprobs are unsupported.",
"unsupported", "logprobs")
decode = engine.tokenizer.decode
surfaces: list[str] = []
chosen: list[float] = []
tops: list[list[tuple[str, float]]] = []
for pos, tid in enumerate(gen_ids):
try:
row = _log_softmax_rows(logits_fn(list(prompt_ids) + gen_ids[:pos]))
except _APIError:
raise
except (ValueError, TypeError, RuntimeError) as exc:
raise _APIError(400, f"Logprob rescoring failed: {exc}.", "scoring_failed")
if tid < 0 or tid >= len(row):
raise _APIError(503, "Engine reported an out-of-vocabulary token.", "generation_failed")
order = _argsort_desc(row)
keep = order[:max(top_n, 1)] if top_n else [tid]
# Decode surfaces lazily but exactly (single-token decode each).
table: dict[int, str] = {}
for cand in dict.fromkeys([tid, *keep]):
try:
table[cand] = decode([cand], skip_special_tokens=True,
clean_up_tokenization_spaces=False)
except Exception:
raise _APIError(503, "Engine tokenizer failed during rescoring.",
"generation_failed") from None
surfaces.append(table[tid])
chosen.append(row[tid])
tops.append([(table[c], row[c]) for c in keep[:top_n]] if top_n else [])
return surfaces, chosen, tops
# -- greedy constrained JSON decoding over real engine logits -----------------
def _run_constrained_sync(engine: Any, prompt_ids: list[int], *,
token_limit: int, target: Any,
cancel: threading.Event) -> tuple[str, list[int], str, int, int]:
"""Greedy argmax over schema-viable tokens using real engine logits.
Returns (text, token_ids, finish_reason, prompt_len, gen_len).
``finish_reason`` is ``"stop"`` only for a complete, validating
document; otherwise ``"length"`` (budget/context spent).
"""
schema = None if target is True else target
config = getattr(engine, "config", None)
vocab = config.get("vocab_size") if isinstance(config, dict) else None
if not isinstance(vocab, int) or vocab <= 0:
raise _APIError(400, "This engine does not expose a vocabulary; structured output is unsupported.",
"unsupported", "response_format")
logits_fn = getattr(engine, "logits", None)
if not callable(logits_fn):
raise _APIError(400, "This engine does not expose logits; structured output is unsupported.",
"unsupported", "response_format")
try:
context = int(engine.context_length)
except (TypeError, ValueError):
raise _APIError(400, "This engine does not expose a context length.",
"unsupported", "response_format") from None
decode = engine.tokenizer.decode
eos = set(_engine_eos(engine))
budget = min(token_limit, max(context - len(prompt_ids), 0))
prefix = list(prompt_ids)
generated: list[int] = []
text = ""
if _complete_and_valid(text, target):
return text, generated, "stop", len(prompt_ids), 0
finish = "length"
for _ in range(budget):
if cancel.is_set():
raise _APIError(503, "Request disconnected.", "request_cancelled")
if len(prefix) >= context:
break
try:
row = _log_softmax_rows(logits_fn(list(prefix)))
except _APIError:
raise
except (ValueError, TypeError, RuntimeError) as exc:
raise _APIError(400, f"Constrained decoding failed: {exc}.", "generation_failed")
if len(row) != vocab:
raise _APIError(503, "Engine logits do not match the model vocabulary.",
"generation_failed")
picked: tuple[int, str] | None = None
for tid in _argsort_desc(row):
try:
full = decode(generated + [tid], skip_special_tokens=True,
clean_up_tokenization_spaces=False)
except Exception:
continue
if tid in eos and not _complete_and_valid(full, target):
continue
if _prefix_viable(full, schema):
picked = (tid, full)
break
if picked is None:
raise _APIError(503, "Constrained decoding found no valid continuation.",
"generation_failed")
tid, full = picked
generated.append(tid)
prefix.append(tid)
text = full
if _complete_and_valid(text, target):
finish = "stop"
break
return text, generated, finish, len(prompt_ids), len(generated)
# -- tool calling: real parse + real local execution + bounded loop -----------
def _eval_arith(node: ast.AST, depth: int = 0) -> Any:
if depth > 64:
raise ValueError("expression too deeply nested")
if isinstance(node, ast.Expression):
return _eval_arith(node.body, depth + 1)
if isinstance(node, ast.Constant):
if isinstance(node.value, bool) or not isinstance(node.value, (int, float)):
raise ValueError("only plain numbers are supported")
if isinstance(node.value, float) and not math.isfinite(node.value):
raise ValueError("non-finite literal")
return node.value
if isinstance(node, ast.UnaryOp) and isinstance(node.op, (ast.UAdd, ast.USub)):
value = _eval_arith(node.operand, depth + 1)
return value if isinstance(node.op, ast.UAdd) else -value
if isinstance(node, ast.BinOp) and isinstance(
node.op, (ast.Add, ast.Sub, ast.Mult, ast.Div, ast.FloorDiv,
ast.Mod, ast.Pow)):
left = _eval_arith(node.left, depth + 1)
right = _eval_arith(node.right, depth + 1)
try:
if isinstance(node.op, ast.Add):
return left + right
if isinstance(node.op, ast.Sub):
return left - right
if isinstance(node.op, ast.Mult):
return left * right
if isinstance(node.op, ast.Div):
return left / right
if isinstance(node.op, ast.FloorDiv):
return left // right
if isinstance(node.op, ast.Mod):
return left % right
if isinstance(right, int) and right < 0 and isinstance(left, int):
raise ValueError("negative integer exponents are unsupported")
if isinstance(right, (int, float)) and abs(right) > 10**6:
raise ValueError("exponent too large")
return left ** right
except ZeroDivisionError:
raise ValueError("division by zero") from None
except OverflowError:
raise ValueError("arithmetic overflow") from None
raise ValueError("only + - * / // % **, parentheses, and numbers are supported")
def _tool_calc(arguments: Any) -> dict:
"""Exact local arithmetic: {"expression": "6*7"} -> {"result": 42}."""
if not isinstance(arguments, dict):
raise ValueError("arguments must be an object")
expression = arguments.get("expression", None)
if not isinstance(expression, str) or not expression.strip():
raise ValueError("expression must be a nonempty string")
if len(expression) > 4096:
raise ValueError("expression too long")
try:
tree = ast.parse(expression, mode="eval")
except SyntaxError as exc:
raise ValueError(f"invalid expression: {exc.msg}") from None
result = _eval_arith(tree)
if isinstance(result, float) and not math.isfinite(result):
raise ValueError("non-finite result")
return {"result": result}
_TOOL_REGISTRY: dict[str, tuple[Any, dict, str]] = {
"calc": (_tool_calc,
{"type": "object",
"properties": {"expression": {"type": "string"}},
"required": ["expression"], "additionalProperties": False},
"Evaluate a decimal arithmetic expression (+ - * / // % **)."),
}
def _parse_tool_call(text: str) -> tuple[str, dict] | None:
"""Parse a ``{"name", "arguments"}`` tool call from assistant text.
Accepts raw JSON or a single ```json fenced block; returns None when the
text is plain assistant content (not a tool call).
"""
stripped = text.strip()
if stripped.startswith("```"):
lines = stripped.split("\n")
try:
end = next(i for i, line in enumerate(lines[1:], 1)
if line.strip().startswith("```"))
except StopIteration:
return None
stripped = "\n".join(lines[1:end]).strip()
if stripped.lower().startswith("json"):
stripped = stripped[4:].strip()
candidate: Any = None
if stripped.startswith("{"):
try:
candidate = json.loads(stripped)
except ValueError:
candidate = None
if candidate is None:
start, end = stripped.find("{"), stripped.rfind("}")
if 0 <= start < end:
try:
candidate = json.loads(stripped[start:end + 1])
except ValueError:
candidate = None
if not isinstance(candidate, dict):
return None
name = candidate.get("name")
arguments = candidate.get("arguments", {})
if not isinstance(name, str) or not name:
return None
if isinstance(arguments, str):
try:
arguments = json.loads(arguments)
except ValueError:
return None
if not isinstance(arguments, dict):
return None
return name, arguments
def _validate_a2_request(body: _CompletionRequest | _ChatRequest, engine: Any) -> None:
"""Fail-fast 400s for Stream A2 options (before queue admission)."""
chat = isinstance(body, _ChatRequest)
if _wants_logprobs(body):
if body.stream:
raise _APIError(400, "logprobs are not supported with stream=true; use non-streaming.",
"invalid_parameter", "logprobs")
if _needs_constrained(body) or _needs_tools(body):
raise _APIError(400, "logprobs cannot be combined with response_format tools yet.",
"invalid_parameter", "logprobs")
if not callable(getattr(engine, "logits", None)):
raise _APIError(400, "This engine does not expose logits; logprobs are unsupported.",
"unsupported", "logprobs")
target = _constrained_target(body) if chat else None
if target is not None:
if not chat: # pragma: no cover - completions carry no response_format
raise _APIError(400, "response_format is only supported on chat completions.",
"invalid_parameter", "response_format")
assert isinstance(body, _ChatRequest)
if body.stream:
raise _APIError(400, "response_format json is not supported with stream=true; use non-streaming.",
"invalid_parameter", "response_format")
if body.temperature != 0 or body.top_p != 1.0 or body.top_k != 0 or body.spec_k != 0:
raise _APIError(400, "response_format json requires greedy decoding "
"(temperature=0, top_p=1, top_k=0, spec_k=0).",
"invalid_parameter", "response_format")
if body.stop:
raise _APIError(400, "response_format json cannot be combined with stop strings.",
"invalid_parameter", "response_format")
if _is_pool_or_batch_engine(engine):
raise _APIError(400, "Structured output is not supported for pool/batch engines.",
"unsupported", "response_format")
if target is not True:
reason = _schema_support_error(target)
if reason is not None:
raise _APIError(400, f"Unsupported response_format schema: {reason}.",
"unsupported_schema", "response_format")
if not callable(getattr(engine, "logits", None)):
raise _APIError(400, "This engine does not expose logits; structured output is unsupported.",
"unsupported", "response_format")
if chat and _needs_tools(body):
assert isinstance(body, _ChatRequest)
if body.stream:
raise _APIError(400, "tools are not supported with stream=true; use non-streaming.",
"invalid_parameter", "tools")
if _is_pool_or_batch_engine(engine):
raise _APIError(400, "Tool calling is not supported for pool/batch engines.",
"unsupported", "tools")
assert body.tools is not None
for tool in body.tools:
name = tool.function.name
if name not in _TOOL_REGISTRY:
raise _APIError(400, f"Unknown tool {name!r}: the server executes only its built-in "
"registry and will not fake external tools.",
"unknown_tool", "tools")
declared = tool.function.parameters
if declared is not None:
reason = _schema_support_error(declared)
if reason is not None:
raise _APIError(400, f"Unsupported tool parameters schema: {reason}.",
"unsupported_schema", "tools")
choice = body.tool_choice
if isinstance(choice, _NamedToolChoice):
wanted = choice.function.name
if all(tool.function.name != wanted for tool in body.tools):
raise _APIError(400, f"tool_choice names {wanted!r} which is not in tools.",
"invalid_parameter", "tool_choice")
# -- Stream A1: paged-KV single-stream + FIFO micro-batching -------------------
#: Dense-only paged decode; Surjo/FWKV keep their own session types.
_PAGED_EXCLUDED_MODEL_TYPES = frozenset({"surjo", "fwkv"})
#: Max co-queued jobs coalesced into one BatchSession micro-batch.
_MICRO_BATCH_MAX = 8
def _serve_paged_enabled() -> bool:
# Default-on since Stream S2 (shared BlockPool, 0.999x dense single-stream):
return os.environ.get("CISM_SERVE_PAGED", "1") != "0"
def _serve_prefixcache_enabled() -> bool:
"""Kill switch: ``CISM_SERVE_PREFIXCACHE=0`` disables automatic reuse.
Default-on (Stream S3): the paged single-stream path matches each
arrival against a per-engine radix prefix index and fork-shares the
longest cached prefix via S2 copy-on-write, prefilling only the
suffix. Misses and stale anchors fall back to full prefill (exact).
"""
return os.environ.get("CISM_SERVE_PREFIXCACHE", "1") != "0"
def _prefix_cache_max_anchors(default: int = 8) -> int:
"""Retained prefix-anchor sessions per engine (pool-pressure bound)."""
try:
value = int(os.environ.get("CISM_SERVE_PREFIXCACHE_ANCHORS", str(default)))
except ValueError:
return default
return value if value > 0 else default
#: Priority demotion for retained prefix anchors: under pool pressure S2's
#: ``evict_one`` frees the lowest-priority entry first, so idle anchors
#: (never decoded again, only forked from) are evicted before active work.
_PREFIX_ANCHOR_PRIORITY = -8
def _prefix_cache_for(engine: Any) -> Any | None:
"""Per-engine S3 radix index, or None when the path is ineligible.
The cache lives on the engine object itself (freed with it -- no
global registry to leak). Requires the S3 kill switch on, the paged
path eligible (``_can_paged``), and a native fork entry point.
"""
if not _serve_prefixcache_enabled() or not _can_paged(engine):
return None
native = getattr(engine, "_native", None)
if native is None or not callable(getattr(native, "create_paged_fork_session", None)):
return None
cache = getattr(engine, "_prefix_cache", None)
if cache is None:
try:
from .paging import PrefixCache
cache = PrefixCache(max_anchors=_prefix_cache_max_anchors())
engine._prefix_cache = cache
except (AttributeError, TypeError, ValueError):
return None
return cache
def prefix_cache_stats(engine: Any | None = None) -> dict[str, Any]:
"""Hit/miss counters for the S3 automatic prefix cache (ops/tests).
With an engine returns its radix-index stats (``enabled: False`` when
the kill switch or eligibility excludes it); with None reports whether
the switch is on. Never raises: stats must not break serving.
"""
try:
if engine is None:
return {"enabled": _serve_prefixcache_enabled()}
cache = getattr(engine, "_prefix_cache", None)
if cache is None:
return {"enabled": _serve_prefixcache_enabled()
and _can_paged(engine), "hits": 0, "misses": 0,
"fork_failures": 0, "inserts": 0, "evictions": 0,
"anchors": 0, "hit_rate": None}
stats = dict(cache.stats())
stats["enabled"] = _serve_prefixcache_enabled()
return stats
except Exception:
return {"enabled": False}
def _serve_microbatch_enabled() -> bool:
return os.environ.get("CISM_SERVE_MICROBATCH", "1") != "0"
def _serve_continuous_enabled() -> bool:
"""Kill switch: ``CISM_SERVE_CONTINUOUS=0`` falls back to A1 FIFO."""
return os.environ.get("CISM_SERVE_CONTINUOUS", "1") != "0"
def _continuous_prefill_chunk(default: int) -> int:
"""New prompt tokens admitted per epoch rebuild (chunked prefill)."""
try:
value = int(os.environ.get("CISM_SERVE_PREFILL_CHUNK", str(default)))
except ValueError:
return default
return value if value > 0 else default
def _continuous_epoch_steps() -> int:
"""Max decode steps before an epoch rotates (starvation backstop)."""
try:
value = int(os.environ.get("CISM_SERVE_EPOCH_STEPS", "1024"))
except ValueError:
return 1024
return value if value > 0 else 1024
def _continuous_max_wait_s() -> float:
"""Starvation guard: bound on an admittable waiter's queue wait."""
try:
value = float(os.environ.get("CISM_SERVE_MAX_WAIT_S", "30"))
except ValueError:
return 30.0
return value if value > 0 else 30.0
def _continuous_join_defer_tokens() -> int:
"""Skip a join rebuild when the live batch is this close to draining.
A newcomer arriving with <= N tokens left in every active member waits
for the drain instead of forcing a recompute-prefill that saves
nothing (starved waiters bypass the deferral via the max-wait guard).
"""
try:
value = int(os.environ.get("CISM_SERVE_JOIN_DEFER_TOKENS", "2"))
except ValueError:
return 2
return value if value >= 0 else 2
def _cont_encode(engine: Any, body: _CompletionRequest | _ChatRequest,
) -> tuple[list[int], list[str] | None]:
"""Validate + encode one epoch arrival (per-job isolation).
Raises ``_APIError`` (400 invalid/context) without touching any batch.
"""
try:
if isinstance(body, _ChatRequest):
prompt = engine.format_chat(
[message.model_dump() for message in body.messages])
prompt_ids = engine.encode(prompt)
else:
prompt_ids = engine.encode(body.prompt)
except (ValueError, TypeError):
raise _APIError(400, "Invalid prompt or chat messages.",
"invalid_prompt") from None
if len(prompt_ids) > engine.context_length:
raise _APIError(400, "Prompt exceeds the model context length.",
"context_length_exceeded")
stop = ([body.stop] if isinstance(body.stop, str)
else (list(body.stop) if body.stop else None))
return prompt_ids, stop
def _can_paged(engine: Any) -> bool:
"""True when ``engine`` can serve a job on the native paged-KV path.
Dense ``Engine`` only: needs ``_native.create_paged_session``,
``tokenizer``, ``config`` (model_type not surjo/fwkv), ``_eos``, and
``_generation_lock``. Pool/batch engines, test doubles, and Surjo/FWKV
return False and keep their existing paths bitwise identical.
"""
native = getattr(engine, "_native", None)
if native is None or not callable(getattr(native, "create_paged_session", None)):
return False
if getattr(engine, "tokenizer", None) is None:
return False
if getattr(engine, "_generation_lock", None) is None:
return False
config = getattr(engine, "config", None)
if isinstance(config, dict) and config.get("model_type") in _PAGED_EXCLUDED_MODEL_TYPES:
return False
return isinstance(getattr(engine, "_eos", None), list)
def _stream_paged(engine: Any, tokens: 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,
cancel: threading.Event | None = None,
chunk_size: int = 16) -> Iterator[_Chunk]:
"""Yield :class:`_Chunk` deltas over a native ``PagedSession``.
Exact-math mirror of ``Engine.stream`` for ``spec_k == 0`` (same
validation, same context clamp, same lock/cancel/monitor discipline, same
stop-string truncation with partial-prefix holdback, same incremental
decode check). Native ``test_paging_native.py`` proves paged blocked-SDPA
is bitwise-identical to the dense ``Session`` path, so greedy and sampled
outputs match ``engine.stream`` token-for-token; spec tallies report 0
(no speculation on the paged path -- callers route ``spec_k > 0`` to the
dense path).
"""
if type(chunk_size) is not int or not 1 <= chunk_size <= 1024:
raise ValueError("chunk_size must be an integer in [1, 1024]")
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 [])
tokens = list(tokens)
vocab = engine.config.get("vocab_size") if isinstance(engine.config, dict) else None
for token in tokens:
if type(token) is not int or (vocab is not None and not 0 <= token < vocab):
raise ValueError("Token IDs must be integers within the model vocabulary")
if len(tokens) > engine.context_length:
raise ValueError("Prompt length exceeds the model context length")
max_new_tokens = min(max_new_tokens, engine.context_length - len(tokens))
cancel = cancel if cancel is not None else threading.Event()
lock = engine._generation_lock
while not lock.acquire(timeout=0.05):
if cancel.is_set():
return
session = None
monitor_done = threading.Event()
monitor = None
monitor_started = False
released = False
def cleanup() -> None:
nonlocal released
if released:
return
try:
monitor_done.set()
if session is not None:
session.cancel()
if monitor_started:
monitor.join()
finally:
released = True
lock.release()
try:
if cancel.is_set():
return
# Stream S3: longest cached token prefix is fork-shared via S2
# copy-on-write (child prefills only the suffix). Sampling config
# is irrelevant to reuse (KV is sampling-independent); the fork
# re-checks token equality, so a hit is exact by construction.
# Stale anchors (S2 evict_one / pool reset) raise here -> invalidate
# the branch and fall back to full prefill (still exact).
cache = _prefix_cache_for(engine)
prefix_hit_len = 0
if cache is not None:
anchor, hit_len = cache.match(tokens)
if anchor is not None:
try:
session = engine._native.create_paged_fork_session(
tokens, max_new_tokens=max_new_tokens,
temperature=temperature, top_p=top_p, top_k=top_k,
seed=seed, eos_token_ids=list(engine._eos),
src=anchor, prefix_len=hit_len)
prefix_hit_len = int(hit_len)
cache.note_hit()
except (ValueError, RuntimeError):
cache.invalidate_session(anchor)
cache.note_fork_failed()
cache.note_miss()
session = None
else:
cache.note_miss()
if session is None:
session = engine._native.create_paged_session(
tokens, max_new_tokens=max_new_tokens,
temperature=temperature, top_p=top_p, top_k=top_k,
seed=seed, eos_token_ids=list(engine._eos))
def watch_cancel() -> None:
while not monitor_done.wait(0.01):
if cancel.is_set():
session.cancel()
return
monitor = threading.Thread(target=watch_cancel, name="cism-cancel", daemon=True)
monitor.start()
monitor_started = True
generated: list[int] = []
emitted = ""
first = True
while True:
if cancel.is_set():
session.cancel()
return
batch = session.next_tokens(1 if first or stops else chunk_size)
first = False
generated.extend(batch)
reason = session.finish_reason or None
if reason == "cancelled":
return
text = engine.tokenizer.decode(generated, skip_special_tokens=True,
clean_up_tokenization_spaces=False)
if reason is None:
text = text.rstrip("\ufffd")
matches = [text.find(s) for s in stops if s in text]
if matches:
text = text[:min(matches)]
reason = "stop"
session.cancel()
elif reason is None:
hold = 0
for sequence in stops:
for size in range(1, min(len(sequence), len(text) + 1)):
if text.endswith(sequence[:size]):
hold = max(hold, size)
if hold:
text = text[:-hold]
if not text.startswith(emitted):
raise ValueError("Tokenizer decoding rewrote emitted text; incremental streaming is unsupported")
delta = text[len(emitted):]
emitted = text
if reason is not None:
# Retain this finished session as an S3 fork anchor (pool KV
# stays resident; demoted priority so pressure evicts idle
# anchors before active work). Cancelled streams never anchor.
if (cache is not None and reason != "cancelled"
and not cancel.is_set()):
try:
session.set_priority(_PREFIX_ANCHOR_PRIORITY)
cache.insert(tokens, session)
except (ValueError, RuntimeError):
pass
cleanup()
yield _Chunk(delta, reason, len(tokens), len(generated),
token_ids=list(batch),
prefix_hit=bool(prefix_hit_len),
prefix_len=prefix_hit_len)
if reason is not None:
break
if not batch:
raise RuntimeError("Native decoder made no progress")
finally:
cleanup()
def _micro_batch_eligible(body: _CompletionRequest | _ChatRequest) -> bool:
"""Only non-streaming greedy jobs may join a micro-batch.
Greedy (``temperature == 0``) decode is seed-independent, so coalescing
jobs with different seeds/sampling trivia stays bitwise-identical to
serial execution. ``spec_k`` has no batch path; streaming jobs keep the
incremental serial path (v0 limit). Stream A2 opt-in jobs (logprobs
rescoring, constrained JSON, tool loops) always run serially so the
shared fast path is untouched by construction.
"""
return (not body.stream and body.spec_k == 0
and body.temperature == 0
and not _wants_logprobs(body)
and not _needs_constrained(body)
and not _needs_tools(body))
def _sampled_batch_eligible(body: _CompletionRequest | _ChatRequest) -> bool:
"""Only non-streaming sampled jobs may join a sampled micro-batch epoch.
Same serial-path exclusions as greedy (streaming, ``spec_k``, A2
opt-ins). ``temperature > 0`` (``== 0`` is the greedy path). Sampling
params must be in their valid ranges; per-sequence seeds ride along in
the native session, so coalescing same-param jobs stays token-identical
to serial execution. Sampled epochs never take mid-decode joins and
never rotate (a rebuild would restart per-seq RNG streams) -- enforced
in ``fits`` and the rotation gate, not here.
"""
try:
temperature = float(body.temperature)
top_p = float(body.top_p)
top_k = int(body.top_k)
except (TypeError, ValueError):
return False
return (not body.stream and body.spec_k == 0
and temperature > 0 and temperature == temperature
and 0 < top_p <= 1.0 and top_k >= 0
and not _wants_logprobs(body)
and not _needs_constrained(body)
and not _needs_tools(body))
def _micro_batch_key(body: _CompletionRequest | _ChatRequest) -> tuple:
stop = [body.stop] if isinstance(body.stop, str) else tuple(body.stop or ())
return (tuple(stop), body.token_limit)
def _is_pool_or_batch_engine(engine: Any) -> bool:
"""True for EnginePool/BatchEngine (generate_batch, no stream).
Plain Engine/FakeEngine expose ``stream``; pool/batch engines expose
``generate_batch`` without ``stream``. This keeps the single-stream
``_Worker`` path bitwise identical for plain engines.
"""
return (hasattr(engine, "generate_batch") and callable(getattr(engine, "generate_batch"))
and not (hasattr(engine, "stream") and callable(getattr(engine, "stream"))))
def _engine_eos(engine: Any) -> list[int]:
eos = getattr(engine, "_eos", None)
if isinstance(eos, list):
return list(eos)
config = getattr(engine, "config", None)
tokenizer = getattr(engine, "tokenizer", None)
raw = None
if isinstance(config, dict):
raw = config.get("eos_token_id", None)
if raw is None:
raw = getattr(tokenizer, "eos_token_id", None)
if raw is None:
return []
return raw if isinstance(raw, list) else [raw]
def _can_native_batch(engine: Any) -> bool:
"""Plain dense Engine with a native BatchSession path (no Surjo)."""
native = getattr(engine, "_native", None)
if native is None or not hasattr(native, "create_batch_session"):
return False
config = getattr(engine, "config", None)
if isinstance(config, dict) and config.get("model_type") == "surjo":
return False
return True
def _run_batch_chunk_sync(engine: Any, chunk_prompts: list, token_limit: int,
temperature: float, top_p: float, top_k: int,
seed: int, stop_list: list[str] | None):
"""Run one scheduler chunk synchronously, preserving chunk order.
Dispatch:
- BatchEngine/EnginePool (``_is_pool_or_batch_engine``): use its
``generate_batch`` (pool = N weight streams in parallel threads;
batch = single weight stream via native ``BatchSession``).
- Plain dense Engine with native batch support: single-weight-stream
``_native.create_batch_session`` under ``_generation_lock`` so
concurrent batches for *unlike* models (different engines/locks) run
in parallel while same-engine batches serialize.
- Otherwise (Surjo/no native batch): sequential ``engine.generate`` loop
(order-preserving fallback).
Returns a list of ``(text, token_ids, finish_reason, prompt_len, gen_len)``.
"""
stops = list(stop_list or [])
if _is_pool_or_batch_engine(engine):
try:
results = engine.generate_batch(
chunk_prompts, max_new_tokens=token_limit,
temperature=temperature, top_p=top_p, top_k=top_k,
seed=seed, **({"stop": stops} if stops else {}))
except TypeError:
# Older pool mocks without stop= support.
results = engine.generate_batch(
chunk_prompts, max_new_tokens=token_limit,
temperature=temperature, top_p=top_p, top_k=top_k, seed=seed)
if stops:
from .paging import ContinuousSchedulerV0 as _S
fixed = []
for r in results:
text, reason = _S.apply_stop(r.text, stops)
fixed.append(type(r)(text, r.token_ids,
reason or r.finish_reason,
r.prompt_tokens, r.completion_tokens,
getattr(r, "spec_proposed", 0),
getattr(r, "spec_accepted", 0)))
results = fixed
return [(r.text, list(r.token_ids), r.finish_reason,
r.prompt_tokens, r.completion_tokens) for r in results]
if _can_native_batch(engine):
# Encode via the engine (BOS fallback + vocab checks identical to
# the single-stream worker).
encoded = [engine.encode(p) for p in chunk_prompts]
context = int(engine.context_length)
effective = min([token_limit] + [context - len(t) for t in encoded])
eos = _engine_eos(engine)
lock = getattr(engine, "_generation_lock", None)
acquire = lock.acquire if lock is not None else None
# Blocking acquire is fine inside asyncio.to_thread.
if acquire is not None:
acquire()
try:
session = engine._native.create_batch_session(
encoded, max_new_tokens=effective, temperature=temperature,
top_p=top_p, top_k=top_k, seed=seed, eos_token_ids=eos)
batched = session.next_tokens(effective)
reasons = list(session.finish_reasons)
finally:
if lock is not None:
lock.release()
decode = engine.tokenizer.decode
out = []
for tokens, token_ids, reason in zip(encoded, batched, reasons):
text = 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"
out.append((text, list(token_ids), reason, len(tokens), len(token_ids)))
return out
# Fallback: sequential generate/stream (Surjo / custom / test doubles).
# Prefer Engine.generate when available; otherwise collect Engine.stream
# chunks (FakeEngine in tests only implements stream).
if hasattr(engine, "generate") and callable(getattr(engine, "generate")):
out = []
for prompt in chunk_prompts:
try:
r = engine.generate(prompt, max_new_tokens=token_limit,
temperature=temperature, top_p=top_p,
top_k=top_k, seed=seed,
stop=stops or None)
except TypeError:
r = engine.generate(prompt, max_new_tokens=token_limit,
temperature=temperature, top_p=top_p,
top_k=top_k, seed=seed)
if stops:
from .paging import ContinuousSchedulerV0 as _S2
text, reason = _S2.apply_stop(r.text, stops)
r = type(r)(text, r.token_ids, reason or r.finish_reason,
r.prompt_tokens, r.completion_tokens,
getattr(r, "spec_proposed", 0),
getattr(r, "spec_accepted", 0))
out.append((r.text, list(r.token_ids), r.finish_reason,
r.prompt_tokens, r.completion_tokens))
return out
# Stream-only backends (e.g. FakeEngine test double).
out = []
for prompt in chunk_prompts:
parts: list[str] = []
ids: list[int] = []
finish = None
prompt_n = 0
gen_n = 0
for chunk in engine.stream(prompt, max_new_tokens=token_limit,
temperature=temperature, top_p=top_p,
top_k=top_k, seed=seed,
stop=stops or None,
cancel=threading.Event(),
chunk_size=16, spec_k=0):
parts.append(chunk.text)
ids.extend(chunk.token_ids)
prompt_n = chunk.prompt_tokens
gen_n = chunk.completion_tokens
if chunk.finish_reason is not None:
finish = chunk.finish_reason
break
if finish is None:
# Mirror the FIFO worker: a missing terminal chunk is a failure.
raise RuntimeError("Generation did not complete")
out.append(("".join(parts), ids, finish, prompt_n, gen_n))
return out
def create_app(engine: _Engine | dict[str, _Engine], api_key: str | None = None,
queue_size: int = 8, *, allow_remote_code_load: bool = False,
load_lock_pages: bool = False, load_keep_warm: float = 0.0,
load_threads_default: int = 2,
load_local_files_only: bool = False,
load_act_precision: str = "fp32",
max_batch_tokens: int = 4096,
scheduler_blocks: int = 512,
prefill_chunk_tokens: int = 512) -> FastAPI:
"""Create an app around engine(s) owned by the caller.
A mapping of model ID to engine serves several loaded models; requests route
by ``model`` field through one shared FIFO worker, so only one generation
runs at a time regardless of how many models are loaded. Passing a single
engine keeps the original behavior.
Phase2 batching (``POST /v1/batch``)::
{"model": "<id>", "prompts": ["a", "b", [1, 2]], "max_tokens": 64}
returns order-preserving ``choices`` (``choices[i]`` is ``prompts[i]``).
Dispatch: same-model same-quant batches use one native ``BatchSession``
weight stream (``BatchEngine`` or plain dense ``Engine._native``); pool
engines use ``EnginePool.generate_batch`` (N weight streams); concurrent
batches for *unlike* models overlap in ``asyncio.to_thread`` (separate
locks), while same-engine batches serialize. Single-stream ``/v1/completions`` + ``/v1/chat/completions`` for plain
engines still go through the shared ``_Worker`` on the dense
``engine.stream`` path (unregressed by construction); the bitwise-identical
native paged-KV path (``create_paged_session``) is opt-in via
``CISM_SERVE_PAGED=1``. Eligible co-queued greedy jobs on dense engines
with a steppable native ``BatchSession`` run in an S1 iteration-level
epoch (per-step add/evict, oldest-first + starvation guard, chunked
prefill; ``CISM_SERVE_CONTINUOUS=0`` falls back to the A1 FIFO
micro-batch); streaming/sampled/speculative/A2 jobs always run
serially. Pool/batch engines serve single-stream via the same batch
machinery with ``B=1``. ``/v1/batch`` chunks FIFO waiting->running requests by
``max_batch_tokens`` (alias ``max_batched_tokens``, ~512 operating point)
plus ``prefill_chunk_tokens`` (128-512) via ``ContinuousSchedulerV0``
(block16, per-seq EOS/stop, BlockManager eviction, recompute-requeue);
see its docstring for v0 limits (no preempt-swap, no cross-seq sharing).
Runtime model management: ``POST /v1/models/load`` downloads and loads a
HF model on a background thread (one load at a time). The load applies
``load_lock_pages``/``load_keep_warm``/``load_threads_default`` to newly
loaded models. ``trust_remote_code=True`` in a load request is rejected
unless ``allow_remote_code_load`` is enabled server-side, because loading
executes downloaded tokenizer code with this process's permissions.
``POST /v1/models/unload`` releases idle models (409 if generating).
``None`` reads CISM_API_KEY once at creation; an empty string disables auth.
Health/readiness probes are public. ``queue_size`` bounds all admitted jobs,
including active inference and responses still being consumed. Cancellation
is cooperative; shutdown waits at most one second for the daemon worker.
"""
if type(queue_size) is not int or queue_size < 1:
raise ValueError("queue_size must be a positive integer")
if type(max_batch_tokens) is not int or max_batch_tokens < 1:
raise ValueError("max_batch_tokens must be a positive integer")
if type(scheduler_blocks) is not int or scheduler_blocks < 1:
raise ValueError("scheduler_blocks must be a positive integer")
if type(prefill_chunk_tokens) is not int or prefill_chunk_tokens < 1:
raise ValueError("prefill_chunk_tokens must be a positive integer")
if isinstance(engine, dict):
model_map = dict(engine)
if not model_map:
raise ValueError("At least one model is required")
else:
model_map = {engine.model_id: engine}
registry_lock = threading.Lock()
load_status: dict[str, str] = {}
load_in_flight = threading.Event()
bench_state: dict[str, Any] = {"status": "idle", "progress": "", "result": None, "error": None}
bench_in_flight = threading.Event()
if api_key is None:
api_key = os.environ.get("CISM_API_KEY")
if api_key is not None and not isinstance(api_key, str):
raise TypeError("api_key must be a string or None")
secret = api_key.encode("utf-8") if api_key else None
@asynccontextmanager
async def lifespan(app: FastAPI):
previous = app.state.worker
if previous is not None and previous.thread.is_alive():
raise RuntimeError("The previous inference worker is still running")
# Pool/Batch engines never use the FIFO worker's default engine (they
# bypass it via the batch machinery), so keep it None for them to
# avoid ever calling Pool.generate_batch as Engine.stream.
default_engine = next(iter(model_map.values())) if len(model_map) == 1 else None
if default_engine is not None and _is_pool_or_batch_engine(default_engine):
default_engine = None
worker = _Worker(default_engine, queue_size)
worker.scheduler_config = {
"scheduler_blocks": scheduler_blocks,
"max_batch_tokens": max_batch_tokens,
"prefill_chunk_tokens": prefill_chunk_tokens,
}
app.state.worker = worker
worker.thread.start()
try:
yield
finally:
worker.stop()
deadline = time.monotonic() + 1.0
while worker.thread.is_alive() and time.monotonic() < deadline:
await asyncio.sleep(0.01)
app = FastAPI(
title="CISM Text Generation API", lifespan=lifespan,
description="Text completions with exact-math logprobs, constrained JSON, and registry tool calls.",
docs_url=None, redoc_url=None, openapi_url=None,
)
app.state.worker = None
@app.exception_handler(_APIError)
async def api_error_handler(request: Request, error: _APIError):
headers = {"WWW-Authenticate": "Bearer"} if error.status == 401 else None
if error.status == 429:
headers = {"Retry-After": "1"}
return JSONResponse(error.body, status_code=error.status, headers=headers)
@app.exception_handler(RequestValidationError)
async def validation_error_handler(request: Request, error: RequestValidationError):
# Pydantic errors include raw inputs and exception context. Never expose them.
return JSONResponse(_APIError(
400, "Invalid request body or unsupported parameter. Only text generation is supported.",
"invalid_request",
).body, status_code=400)
@app.exception_handler(HTTPException)
async def http_error_handler(request: Request, error: HTTPException):
message = "Not found." if error.status_code == 404 else "HTTP request failed."
return JSONResponse(_APIError(error.status_code, message, "http_error").body,
status_code=error.status_code)
@app.exception_handler(Exception)
async def unexpected_error_handler(request: Request, error: Exception):
return JSONResponse(_APIError(503, "Server unavailable.", "server_unavailable").body,
status_code=503)
async def authorize(request: Request) -> None:
if secret is None:
return
scheme, _, credential = request.headers.get("authorization", "").partition(" ")
matches = hmac.compare_digest(credential.encode("utf-8"), secret)
if scheme.lower() != "bearer" or not matches:
raise _APIError(401, "Invalid API key.", "invalid_api_key")
router = APIRouter(prefix="/v1", dependencies=[Depends(authorize)])
@app.get("/health")
async def health():
return {"status": "ok"}
@app.get("/ready")
async def ready():
worker = app.state.worker
if worker is None or not worker.ready:
raise _APIError(503, "Server is not ready.", "server_unavailable")
return {"status": "ready"}
@app.get("/", response_class=HTMLResponse)
async def index():
# Static, dependency-free tester page. It never contains the API key;
# visitors paste their own credential, like any other API client.
# no-cache keeps model dropdowns fresh across server restarts.
return HTMLResponse(_WEBUI_HTML, headers={"Cache-Control": "no-cache"})
@router.get("/models")
async def models():
with registry_lock:
items = list(model_map.items())
data = []
for model_id, model in items:
entry = {"id": model_id, "object": "model", "created": 0, "owned_by": "cism",
"context_length": model.context_length}
# Optional engine introspection (real engines provide info).
info = getattr(model, "info", None)
if callable(info):
info = info()
if isinstance(info, dict):
for key in ("precision", "threads", "kernel", "architecture", "quantization",
"act_precision"):
if key in info:
entry[key] = info[key]
if isinstance(info.get("weight_bytes"), (int, float)):
entry["weight_mb"] = round(info["weight_bytes"] / (1024 * 1024), 2)
data.append(entry)
return {"object": "list", "data": data}
@router.get("/models/loading")
async def models_loading():
with registry_lock:
return {"loading": dict(load_status)}
def _apply_load_options(loaded_engine) -> None:
if load_lock_pages and loaded_engine.lock_pages() == 0:
load_status[loaded_engine.model_id] = "loaded (page locking failed)"
if load_keep_warm:
loaded_engine.start_keep_warm(load_keep_warm)
def _register_loaded_engine(loaded_engine) -> None:
with registry_lock:
model_map[loaded_engine.model_id] = loaded_engine
load_status.pop(loaded_engine.model_id, None)
_apply_load_options(loaded_engine)
def _load_via_loader(model: str, precision: str, threads: int,
revision: str | None = None, trust_remote_code: bool = False,
act_precision: str | None = None, reserved: bool = False):
"""Load through the single-slot loader; returns None if it is busy.
``reserved=True`` means the caller already holds the load slot (the
``/models/load`` endpoint reserves it under the registry lock before
starting its worker); without it the flag would read as someone
else's load and the worker would refuse its own reservation.
"""
with registry_lock:
existing = model_map.get(model)
existing_info = getattr(existing, "info", {}) if existing is not None else {}
if callable(existing_info):
existing_info = existing_info()
if not isinstance(existing_info, dict):
existing_info = {}
want_act = act_precision or load_act_precision
if (existing is not None and existing_info.get("precision") == precision and
existing_info.get("act_precision", "fp32") == want_act and
existing_info.get("threads", threads) == threads):
if reserved:
load_in_flight.clear() # the endpoint's reservation is done
return existing
if load_in_flight.is_set() and not reserved:
return None
load_in_flight.set()
load_status[model] = "loading"
try:
from cism.engine import Engine
loaded_engine = Engine.from_pretrained(
model, precision=precision, revision=revision, threads=threads,
trust_remote_code=trust_remote_code,
local_files_only=load_local_files_only,
act_precision=want_act)
except Exception as error:
with registry_lock:
load_status[model] = f"failed: {type(error).__name__}: {error}"
raise
finally:
# The loader releases its own slot on every exit path, whoever
# called it (the /models/load worker or the bench harness).
load_in_flight.clear()
with registry_lock:
if load_status.get(model) == "loading":
load_status.pop(model, None)
_register_loaded_engine(loaded_engine)
return loaded_engine
def _load_worker(request_body: _LoadRequest) -> None:
try:
_load_via_loader(request_body.model, request_body.precision,
request_body.threads, revision=request_body.revision,
trust_remote_code=request_body.trust_remote_code,
act_precision=request_body.act_precision,
reserved=True)
except Exception:
pass # failure is recorded in load_status for the pollers
finally:
load_in_flight.clear()
def _bench_worker(request_body: _AutoBenchRequest) -> None:
from . import autobench
def report_progress(message: str) -> None:
with registry_lock:
bench_state["progress"] = message
try:
result = autobench.run_benchmark(
request_body.model, precision=request_body.precision,
threads=request_body.threads, max_new_tokens=request_body.max_tokens,
local_files_only=load_local_files_only,
thread_sweep=request_body.thread_sweep,
spec_ks=(0, 2, 4, 8, 12, 16) if request_body.spec_sweep else (0,),
act_modes=("fp32", "int8") if request_body.bench_act_int8 else ("fp32",),
include_torch=request_body.include_torch,
include_compile=request_body.compile,
load_engine=lambda requested, threads=None, act_precision=None: _load_via_loader(
request_body.model, requested or request_body.precision,
threads or request_body.threads, act_precision=act_precision),
progress=report_progress)
with registry_lock:
bench_state.update(status="done", result=result, progress="", error=None)
except Exception as error:
with registry_lock:
bench_state.update(status="failed", error=f"{type(error).__name__}: {error}",
progress="", result=None)
finally:
bench_in_flight.clear()
@router.get("/bench/auto")
async def bench_auto_status():
with registry_lock:
return dict(bench_state)
@router.post("/bench/auto", status_code=202)
async def bench_auto(body: _AutoBenchRequest):
with registry_lock:
if bench_in_flight.is_set():
raise _APIError(409, "A benchmark is already running.", "bench_in_progress")
if body.model not in model_map and load_in_flight.is_set():
raise _APIError(409, "A model load is already in progress.", "load_in_progress")
bench_in_flight.set()
bench_state.update(status="running", progress="starting", result=None, error=None)
threading.Thread(target=_bench_worker, args=(body,), name="cism-bench", daemon=True).start()
return {"status": "running", "model": body.model, "precision": body.precision}
@router.post("/models/load", status_code=202)
async def models_load(body: _LoadRequest):
with registry_lock:
if body.model in model_map:
# Same weights + activation mode: no-op. A different precision
# (the shootout's load-bench cycle) must fall through and reload.
existing = model_map[body.model]
existing_info = getattr(existing, "info", {})
if callable(existing_info):
existing_info = existing_info()
if not isinstance(existing_info, dict):
existing_info = {}
want_act = body.act_precision or load_act_precision
if (existing_info.get("precision") == body.precision and
existing_info.get("act_precision", "fp32") == want_act):
return {"model": body.model, "status": "ready"}
if body.trust_remote_code and not allow_remote_code_load:
raise _APIError(403,
"trust_remote_code is disabled on this server; restart with "
"--allow-remote-code-load to permit downloaded tokenizer code.",
"remote_code_disabled")
if load_in_flight.is_set():
raise _APIError(429, "Another model load is already in progress.", "load_in_progress")
load_in_flight.set()
# Stale failures from earlier attempts would otherwise resurface in
# every later poll; the UI has already shown them once.
for failed_id in [m for m, s in load_status.items() if s.startswith("failed")]:
del load_status[failed_id]
load_status[body.model] = "loading"
threading.Thread(target=_load_worker, args=(body,), name="cism-loader", daemon=True).start()
return {"model": body.model, "status": "loading"}
@router.post("/models/unload")
async def models_unload(body: _UnloadRequest):
with registry_lock:
loaded_engine = model_map.pop(body.model, None)
if loaded_engine is None:
raise _APIError(404, "Requested model is not available.", "model_not_found", "model")
load_status.pop(body.model, None)
# Idle check: an active generation holds the engine's generation lock.
generation_lock = getattr(loaded_engine, "_generation_lock", None)
if generation_lock is not None and not generation_lock.acquire(blocking=False):
with registry_lock:
model_map[body.model] = loaded_engine # keep serving in-flight work
raise _APIError(409, "Model is generating; unload after it finishes.", "model_busy")
try:
info = loaded_engine.info if hasattr(loaded_engine, "info") else {}
if callable(info):
info = info()
if isinstance(info, dict) and info.get("pages_locked"):
loaded_engine.unlock_pages()
if getattr(loaded_engine, "keep_warm_active", False):
loaded_engine.stop_keep_warm()
except Exception:
with registry_lock:
model_map[body.model] = loaded_engine
raise _APIError(503, "Model unload failed; it remains available.", "unload_failed")
if generation_lock is not None:
generation_lock.release()
return {"model": body.model, "status": "unloaded"}
@router.post("/cache/report")
async def cache_report(body: _CacheRequest):
engine = model_map.get(body.model)
if engine is None:
raise _APIError(404, "Requested model is not available.", "model_not_found", "model")
def _run() -> dict:
import numpy as np
info = getattr(engine, "info", None)
if callable(info):
info = info()
weight_bytes = info.get("weight_bytes") if isinstance(info, dict) else None
if not isinstance(weight_bytes, (int, float)) or weight_bytes <= 0:
raise _APIError(400, "This engine does not expose weight bytes.",
"invalid_engine", "model")
try:
prompt = engine.encode("cache report probe")
def timed_scan(repeats: int, evictions: int) -> float:
# Unique 2x-L3 writes evict residency; the scan is timed alone.
scrub = np.empty(body.scrub_mb * 1024 * 1024, dtype=np.uint8)
samples = []
for _ in range(repeats):
for _ in range(evictions):
scrub[:] = 1
start = time.perf_counter()
engine.scan()
samples.append(time.perf_counter() - start)
return statistics.median(samples)
warm_seconds = timed_scan(20, evictions=0)
cold_seconds = timed_scan(5, evictions=2)
session = engine._native.create_session(
prompt, max_new_tokens=body.decode_tokens, eos_token_ids=[])
session.next_tokens(1)
start = time.perf_counter()
generated = session.next_tokens(body.decode_tokens - 1)
decode_seconds = time.perf_counter() - start
except _APIError:
raise
except (ValueError, RuntimeError) as exc:
raise _APIError(400, f"Cache probe failed: {exc}.", "cache_probe_failed")
tokens_per_second = len(generated) / decode_seconds if generated and decode_seconds > 0 else 0
effective_gbps = weight_bytes * tokens_per_second / 1e9
warm_gbps = weight_bytes / warm_seconds / 1e9 if warm_seconds > 0 else 0
cold_gbps = weight_bytes / cold_seconds / 1e9 if cold_seconds > 0 else 0
ratio = cold_seconds / warm_seconds if warm_seconds > 0 else 0
if ratio >= 2:
residency = "cache-resident when warm (cold scan is much slower)"
elif ratio >= 1.3:
residency = "partially cache-resident"
else:
residency = "no measurable residency benefit from this probe"
if not tokens_per_second:
bottleneck = "decode produced no tokens"
elif effective_gbps <= cold_gbps:
bottleneck = ("compute-bound: decode bandwidth is below even the cold DRAM scan; "
"the cache is not the bottleneck, the kernels are")
elif effective_gbps < warm_gbps * 0.8:
bottleneck = "DRAM-bandwidth-limited: weights are effectively streaming from RAM"
else:
bottleneck = "bandwidth-limited at cache-resident rates"
return {"model": engine.model_id, "weight_bytes": weight_bytes,
"warm_scan_gbps": round(warm_gbps, 2),
"cold_scan_gbps": round(cold_gbps, 2),
"warm_to_cold_ratio": round(ratio, 2),
"decode_tokens_per_second": round(tokens_per_second, 1),
"decode_effective_gbps": round(effective_gbps, 2),
"residency": residency, "decode_bottleneck": bottleneck}
return await asyncio.to_thread(_run)
async def _generate_single_via_batch(engine: Any, body: _CompletionRequest | _ChatRequest,
request: Request):
"""Serve one completions/chat request via the batch machinery (B=1).
Used only for pool/batch engines. Returns the same JSON/SSE shape as
the FIFO worker path (non-streaming choices + usage; streaming SSE
emits the full text as one delta for v0). Raises _APIError on
validation/context failures with sanitized messages.
"""
chat = isinstance(body, _ChatRequest)
stop = [body.stop] if isinstance(body.stop, str) else (list(body.stop) if body.stop else None)
# Resolve the prompt without holding any lock (tokenizer-only).
try:
if chat:
assert isinstance(body, _ChatRequest)
prompt_tokens = engine.format_chat(
[message.model_dump() for message in body.messages])
# format_chat may return tokens already; encode validates them.
prompt: Any = engine.encode(prompt_tokens)
else:
assert isinstance(body, _CompletionRequest)
prompt = body.prompt
# Early context check for string prompts (encode to measure).
if isinstance(prompt, str):
prompt_tokens_probe = engine.encode(prompt)
if len(prompt_tokens_probe) > engine.context_length:
raise _APIError(400, "Prompt exceeds the model context length.",
"context_length_exceeded")
prompt = body.prompt
except _APIError:
raise
except (ValueError, TypeError):
raise _APIError(400, "Invalid prompt or chat messages.", "invalid_prompt") from None
token_limit = body.token_limit
def _run_one():
try:
rows = _run_batch_chunk_sync(
engine, [prompt], token_limit, body.temperature,
body.top_p, body.top_k, body.seed, stop)
except (ValueError, TypeError) as exc:
msg = str(exc)
if "context" in msg.lower() or "Prompt length" in msg:
raise _APIError(400, "Prompt exceeds the model context length.",
"context_length_exceeded") from None
raise _APIError(400, "Invalid prompt or chat messages.", "invalid_prompt") from None
if len(rows) != 1:
raise _APIError(503, "Generation did not complete.", "generation_failed")
return rows[0]
try:
text, _ids, finish, prompt_n, gen_n = await asyncio.to_thread(_run_one)
except _APIError:
raise
except Exception:
raise _APIError(503, "Generation failed.", "generation_failed") from None
if finish is None:
raise _APIError(503, "Generation did not complete.", "generation_failed")
base = {
"id": ("chatcmpl-" if chat else "cmpl-") + uuid.uuid4().hex,
"object": "chat.completion" if (chat and not body.stream) else (
"chat.completion.chunk" if chat else "text_completion"),
"created": int(time.time()), "model": engine.model_id,
}
usage = {"prompt_tokens": prompt_n, "completion_tokens": gen_n,
"total_tokens": prompt_n + gen_n,
"spec_proposed": 0, "spec_accepted": 0, "spec_accept_rate": None}
if not body.stream:
choice: dict[str, Any] = {"index": 0, "finish_reason": finish, "logprobs": None}
if chat:
choice["message"] = {"role": "assistant", "content": text}
else:
choice["text"] = text
return {**base, "choices": [choice], "usage": usage}
# Streaming via batch path (v0): single delta + terminal + DONE.
include_usage = body.stream_options is not None and body.stream_options.include_usage
def event(choices: list[dict], usage_payload: dict | None = None,
completion_tokens: int | None = None) -> str:
payload: dict[str, Any] = {**base, "choices": choices}
if include_usage:
payload["usage"] = usage_payload
if completion_tokens is not None:
payload["cism_completion_tokens"] = completion_tokens
return "data: " + json.dumps(payload, ensure_ascii=True, allow_nan=False) + "\n\n"
async def events():
try:
if chat:
yield event([{"index": 0, "delta": {"role": "assistant", "content": ""},
"finish_reason": None}], None, gen_n)
if text:
content = {"delta": {"content": text}} if chat else {"text": text, "logprobs": None}
yield event([{"index": 0, **content, "finish_reason": None}], None, gen_n)
content = {"delta": {}} if chat else {"text": "", "logprobs": None}
yield event([{"index": 0, **content, "finish_reason": finish}], None, gen_n)
if include_usage:
yield event([], usage)
yield "data: [DONE]\n\n"
except Exception:
error = _APIError(503, "Generation failed.", "generation_failed")
yield "data: " + json.dumps(error.body) + "\n\n"
yield "data: [DONE]\n\n"
async def _gen_wrapper():
async for chunk in events():
yield chunk
# Plain StreamingResponse (no _Worker job to release).
return StreamingResponse(_gen_wrapper(), media_type="text/event-stream", headers={
"Cache-Control": "no-cache", "X-Accel-Buffering": "no",
})
async def generate(body: _CompletionRequest | _ChatRequest, request: Request):
engine = model_map.get(body.model)
if engine is None:
raise _APIError(404, "Requested model is not available.", "model_not_found", "model")
# SDK engines expose config; lightweight engine implementations need not.
vocab_size = getattr(engine, "config", {}).get("vocab_size")
if vocab_size is not None and body.top_k > vocab_size:
raise _APIError(400, "top_k exceeds the model vocabulary size.",
"invalid_parameter", "top_k")
known_length = len(body.prompt) if isinstance(body, _CompletionRequest) and isinstance(body.prompt, list) else 0
if known_length > engine.context_length:
raise _APIError(400, "Prompt exceeds the model context length.",
"context_length_exceeded")
if body.spec_k >= 2 and body.temperature != 0:
raise _APIError(400, "spec_k requires temperature=0 (greedy speculation).",
"invalid_parameter", "spec_k")
if not 0 <= body.seed < 2**64:
raise _APIError(400, "seed must be an unsigned 64-bit integer.",
"invalid_parameter", "seed")
# Stream A2: fail-fast 400s before queue admission (never faked).
_validate_a2_request(body, engine)
# Phase2: pool/batch engines bypass the shared FIFO worker and use the
# batch machinery with B=1 (single weight stream for BatchEngine /
# native path, N streams for EnginePool). Plain engines keep the exact
# FIFO _Worker path below (bitwise identical).
if _is_pool_or_batch_engine(engine):
if body.spec_k != 0:
raise _APIError(400, "spec_k is not supported for pool/batch engines.",
"invalid_parameter", "spec_k")
worker = app.state.worker
if worker is None or not worker.ready:
raise _APIError(503, "Server is not ready.", "server_unavailable")
with worker.lock:
if len(worker.jobs) >= worker.capacity:
raise _APIError(429, "Generation queue is full.", "queue_full")
return await _generate_single_via_batch(engine, body, request)
worker = app.state.worker
if worker is None:
raise _APIError(503, "Server is not ready.", "server_unavailable")
job = worker.submit(body, engine)
handed_off = False
chat = isinstance(body, _ChatRequest)
base = {
"id": ("chatcmpl-" if chat else "cmpl-") + uuid.uuid4().hex,
"object": "chat.completion.chunk" if chat else "text_completion",
"created": int(time.time()), "model": engine.model_id,
}
try:
first = await worker.receive(job, request)
if isinstance(first, _APIError):
raise first
if not body.stream:
parts = [first.text]
ids = list(first.token_ids)
tool_calls = list(first.tool_calls)
chunk = first
while chunk.finish_reason is None:
chunk = await worker.receive(job, request)
if isinstance(chunk, _APIError):
raise chunk
parts.append(chunk.text)
ids.extend(chunk.token_ids)
tool_calls.extend(chunk.tool_calls)
text = "".join(parts)
logprobs_payload: Any = None
if _wants_logprobs(body):
logprobs_payload = await _rescore_logprobs(engine, body, ids, text)
choice = {"index": 0, "finish_reason": chunk.finish_reason, "logprobs": logprobs_payload}
if chat:
assert isinstance(body, _ChatRequest)
message: dict[str, Any] = {"role": "assistant", "content": text}
if tool_calls:
message["tool_calls"] = tool_calls
choice["message"] = message
base["object"] = "chat.completion"
else:
choice["text"] = text
return {**base, "choices": [choice], "usage": _usage(chunk)}
include_usage = body.stream_options is not None and body.stream_options.include_usage
def event(choices: list[dict], usage: dict | None = None,
completion_tokens: int | None = None) -> str:
payload = {**base, "choices": choices}
if include_usage:
payload["usage"] = usage
if completion_tokens is not None:
# Nonstandard live progress field; OpenAI clients ignore it.
payload["cism_completion_tokens"] = completion_tokens
return "data: " + json.dumps(payload, ensure_ascii=True, allow_nan=False) + "\n\n"
async def events():
try:
if chat:
yield event([{"index": 0, "delta": {"role": "assistant", "content": ""},
"finish_reason": None}], None, first.completion_tokens)
chunk = first
while True:
if chunk.text:
content = {"delta": {"content": chunk.text}} if chat else {"text": chunk.text, "logprobs": None}
yield event([{"index": 0, **content, "finish_reason": None}],
None, chunk.completion_tokens)
if chunk.finish_reason is not None:
content = {"delta": {}} if chat else {"text": "", "logprobs": None}
yield event([{"index": 0, **content, "finish_reason": chunk.finish_reason}],
None, chunk.completion_tokens)
if include_usage:
yield event([], _usage(chunk))
break
chunk = await worker.receive(job, request, streaming=True)
if isinstance(chunk, _APIError):
raise chunk
yield "data: [DONE]\n\n"
except _APIError as error:
yield "data: " + json.dumps(error.body) + "\n\n"
yield "data: [DONE]\n\n"
except Exception:
error = _APIError(503, "Generation failed.", "generation_failed")
yield "data: " + json.dumps(error.body) + "\n\n"
yield "data: [DONE]\n\n"
response = _JobResponse(events(), worker, job)
handed_off = True
return response
finally:
if not handed_off:
worker.release(job)
@router.post("/tokenize")
async def tokenize(body: _TokenizeRequest):
"""Prompt accounting: token count, context fit, KV footprint."""
engine = model_map.get(body.model)
if engine is None:
raise _APIError(404, "Requested model is not available.", "model_not_found", "model")
try:
tokens = engine.encode(body.text)
except (ValueError, TypeError) as exc:
raise _APIError(400, f"Invalid prompt: {exc}.", "invalid_prompt")
context = getattr(engine, "context_length", None)
fits = len(tokens) <= context if isinstance(context, int) else None
kv = _kv_cache_info(engine)
if kv is not None:
kv = {**kv, "bytes_at_prompt": kv["bytes_per_token"] * len(tokens)}
return {"model": body.model, "tokens": tokens, "count": len(tokens),
"context_length": context, "fits_context": fits, "kv_cache": kv}
@router.post("/nll")
async def score(body: _NllRequest):
"""Standalone perplexity scoring over pasted text."""
engine = model_map.get(body.model)
if engine is None:
raise _APIError(404, "Requested model is not available.", "model_not_found", "model")
try:
result = await asyncio.to_thread(engine.nll, body.text, body.window)
except ValueError as exc:
raise _APIError(400, f"Scoring failed: {exc}.", "scoring_failed")
return {"model": body.model, **result}
@router.post("/models/runtime")
async def models_runtime(body: _RuntimeRequest):
"""Page locking and keep-warm control for a loaded model."""
engine = model_map.get(body.model)
if engine is None:
raise _APIError(404, "Requested model is not available.", "model_not_found", "model")
status: dict[str, Any] = {}
try:
if body.lock_pages is True:
if not hasattr(engine, "lock_pages"):
raise _APIError(400, "This engine does not support page locking.",
"unsupported", "lock_pages")
status["page_lock_bytes"] = engine.lock_pages()
elif body.lock_pages is False:
if not hasattr(engine, "unlock_pages"):
raise _APIError(400, "This engine does not support page locking.",
"unsupported", "lock_pages")
status["page_unlock_bytes"] = engine.unlock_pages()
if body.keep_warm is not None:
if not hasattr(engine, "start_keep_warm"):
raise _APIError(400, "This engine does not support keep-warm.",
"unsupported", "keep_warm")
if body.keep_warm == 0:
engine.stop_keep_warm()
status["keep_warm"] = "stopped"
else:
engine.start_keep_warm(body.keep_warm)
status["keep_warm"] = body.keep_warm
except _APIError:
raise
except ValueError as exc:
raise _APIError(400, f"Invalid runtime option: {exc}.", "invalid_parameter")
except RuntimeError as exc:
raise _APIError(409, f"Runtime change refused: {exc}.", "runtime_busy")
status["keep_warm_active"] = bool(getattr(engine, "keep_warm_active", False))
return {"model": body.model, **status}
@router.post("/batch")
async def batch(body: _BatchRequest, request: Request):
"""Phase2 batched completions (order-preserving, non-streaming).
Request ``{"model", "prompts": [...], ...sampling}`` returns
``{"object": "batch.completion", "choices": [{index, text,
finish_reason}], "usage": {...}}`` with ``choices[i]`` for
``prompts[i]``. Same-model batches run as one native ``BatchSession``
weight stream when available (``BatchEngine`` or dense ``Engine``);
pool engines fan out over ``EnginePool`` slots; concurrent batches for
unlike models overlap (separate locks), same-engine batches serialize.
Large batches are split FIFO waiting->running by ``max_batch_tokens``
plus ``prefill_chunk_tokens`` via ``ContinuousSchedulerV0`` (block16,
per-seq EOS/stop, BlockManager eviction, recompute-requeue: evicted
seqs restart from scratch). Limits: no preempt-swap yet, no
cross-seq KV sharing. ``spec_k`` must be 0 (no speculation in batch).
"""
engine = model_map.get(body.model)
if engine is None:
raise _APIError(404, "Requested model is not available.", "model_not_found", "model")
vocab_size = getattr(engine, "config", {}).get("vocab_size")
if vocab_size is not None and body.top_k > vocab_size:
raise _APIError(400, "top_k exceeds the model vocabulary size.",
"invalid_parameter", "top_k")
if body.spec_k != 0:
raise _APIError(400, "spec_k is not supported for batch requests.",
"invalid_parameter", "spec_k")
if not 0 <= body.seed < 2**64:
raise _APIError(400, "seed must be an unsigned 64-bit integer.",
"invalid_parameter", "seed")
# Early context check for token-list prompts with known lengths.
for prompt in body.prompts:
if isinstance(prompt, list) and len(prompt) > engine.context_length:
raise _APIError(400, "Prompt exceeds the model context length.",
"context_length_exceeded")
worker = app.state.worker
if worker is None or not worker.ready:
raise _APIError(503, "Server is not ready.", "server_unavailable")
with worker.lock:
queued = len(worker.jobs)
# Coarse admission: a batch counts as one job; oversized batches and a
# full FIFO both return 429 with Retry-After (same as single-stream).
if len(body.prompts) > worker.capacity:
raise _APIError(429, "Batch is larger than the generation queue.", "queue_full")
if queued >= worker.capacity:
raise _APIError(429, "Generation queue is full.", "queue_full")
if await request.is_disconnected():
raise _APIError(503, "Request disconnected.", "request_cancelled")
stop = [body.stop] if isinstance(body.stop, str) else (list(body.stop) if body.stop else None)
token_limit = body.token_limit
def _run_all():
from .paging import ContinuousSchedulerV0
# Measure prompt lens via engine.encode (validates + BOS fallback).
try:
encoded_lens: list[int] = []
for p in body.prompts:
encoded_lens.append(len(engine.encode(p)))
except (ValueError, TypeError):
raise _APIError(400, "Invalid prompt or chat messages.", "invalid_prompt")
for length in encoded_lens:
if length > engine.context_length:
raise _APIError(400, "Prompt exceeds the model context length.",
"context_length_exceeded")
# FIFO scheduler: waiting->running budget + prefill chunk +
# block16 eviction with recompute-requeue, trivial prefix-reuse
# accounting (no KV sharing, no preempt-swap in v0).
sched = ContinuousSchedulerV0(
num_blocks=scheduler_blocks, block_size=16,
max_batch_tokens=max_batch_tokens,
prefill_chunk_tokens=prefill_chunk_tokens)
for idx, length in enumerate(encoded_lens):
sched.add(idx, length, token_limit, stop=stop)
chunks = sched.plan_chunks()
num_chunks = len(chunks)
ordered: list | None = [None] * len(body.prompts)
prompt_list = list(body.prompts)
for chunk in chunks:
indices = [entry["req_id"] for entry in chunk]
sub_prompts = [prompt_list[i] for i in indices]
# Seed offset preserves per-seq determinism across chunks for
# the single-weight-stream path (seq i uses seed+i globally).
base_seed = (body.seed + min(indices)) % (2**64) if indices else body.seed
if _is_pool_or_batch_engine(engine):
# Pool uses one seed for all seqs (existing semantics);
# keep the request seed for every chunk.
chunk_seed = body.seed
else:
chunk_seed = base_seed
try:
rows = _run_batch_chunk_sync(
engine, sub_prompts, token_limit, body.temperature,
body.top_p, body.top_k, chunk_seed, stop)
except (ValueError, TypeError) as exc:
msg = str(exc)
if "context" in msg.lower() or "Prompt length" in msg:
raise _APIError(400, "Prompt exceeds the model context length.",
"context_length_exceeded") from None
raise _APIError(400, "Invalid prompt or chat messages.", "invalid_prompt") from None
if len(rows) != len(indices):
raise _APIError(503, "Generation did not complete.", "generation_failed")
for global_idx, row in zip(indices, rows):
text, _ids, finish, _pn, _gn = row
if finish is None:
raise _APIError(503, "Generation did not complete.", "generation_failed")
assert ordered is not None
ordered[global_idx] = (text, _ids, finish, _pn, _gn)
# Free the planner's per-chunk blocks (accounting only; the
# native sessions already released their state).
for entry in chunk:
try:
sched.free(entry["req_id"])
except Exception:
pass
assert ordered is not None and all(r is not None for r in ordered)
return ordered, num_chunks
try:
ordered, num_chunks = await asyncio.to_thread(_run_all)
except _APIError:
raise
except Exception:
raise _APIError(503, "Generation failed.", "generation_failed") from None
choices = []
prompt_tokens_total = 0
completion_tokens_total = 0
for index, row in enumerate(ordered):
assert row is not None
text, _ids, finish, prompt_n, gen_n = row
prompt_tokens_total += prompt_n
completion_tokens_total += gen_n
choices.append({"index": index, "text": text,
"finish_reason": finish, "logprobs": None})
return {"id": "batch-" + uuid.uuid4().hex, "object": "batch.completion",
"created": int(time.time()), "model": engine.model_id,
"choices": choices,
"usage": {"prompt_tokens": prompt_tokens_total,
"completion_tokens": completion_tokens_total,
"total_tokens": prompt_tokens_total + completion_tokens_total},
"scheduler": {"max_batch_tokens": max_batch_tokens,
"max_batched_tokens": max_batch_tokens,
"prefill_chunk_tokens": prefill_chunk_tokens,
"block_size": 16, "num_blocks": scheduler_blocks,
"chunks": num_chunks}}
@router.post("/completions")
async def completions(body: _CompletionRequest, request: Request):
return await generate(body, request)
@router.post("/chat/completions")
async def chat_completions(body: _ChatRequest, request: Request):
return await generate(body, request)
app.include_router(router)
return app
def _kv_cache_info(engine: Any) -> dict | None:
"""Per-token and full-context KV cache bytes, or None without config."""
config = getattr(engine, "config", None)
context = getattr(engine, "context_length", None)
if not isinstance(config, dict) or not isinstance(context, int):
return None
try:
per_token = (int(config["num_hidden_layers"]) * int(config["num_key_value_heads"]) *
int(config["head_dim"]) * 2 * 4)
except (KeyError, TypeError, ValueError):
return None
if per_token <= 0:
return None
return {"bytes_per_token": per_token,
"bytes_at_full_context": per_token * context}
_WEBUI_HTML = """<!doctype html>
<html lang="en">
<head>
<meta charset="utf-8">
<meta name="viewport" content="width=device-width, initial-scale=1">
<title>CISM tester</title>
<link rel="icon" href="data:,">
<meta name="description" content="CISM tester: generate, benchmark, and compare native CPU models.">
<title>CISM tester</title>
<link rel="icon" href="data:image/svg+xml,%3Csvg xmlns='http://www.w3.org/2000/svg' viewBox='0 0 32 32'%3E%3Crect width='32' height='32' rx='7' fill='%23171310'/%3E%3Ctext x='16' y='23' font-family='Consolas,monospace' font-size='19' font-weight='bold' fill='%23e0974a' text-anchor='middle'%3EC%3C/text%3E%3C/svg%3E">
<style>
:root {
color-scheme: dark;
--bg: #100d0b;
--panel: #171310;
--panel-2: #1f1815;
--line: #2e2521;
--text: #ece2d5;
--muted: #ab9d8b;
--accent: #e0974a;
--accent-dim: #b06f2e;
--good: #7fbf6a;
--bad: #ef6a5a;
--radius: 10px;
--radius-s: 6px;
}
* { box-sizing: border-box; }
html { scroll-behavior: smooth; }
body {
margin: 0 auto; padding: 20px 16px 48px; max-width: 1180px;
font: 14px/1.55 -apple-system, "Segoe UI", system-ui, Roboto, sans-serif;
color: var(--text);
background:
radial-gradient(1100px 480px at 15% -10%, #26190f 0%, transparent 60%),
radial-gradient(900px 420px at 95% 0%, #1c130c 0%, transparent 55%),
var(--bg);
}
.skip { position: absolute; left: -9999px; }
.skip:focus { left: 12px; top: 12px; background: var(--panel-2); padding: 6px 10px; z-index: 10; }
.topbar {
display: flex; flex-wrap: wrap; gap: 12px 20px; align-items: center;
justify-content: space-between; margin-bottom: 16px;
}
.brand { display: flex; gap: 12px; align-items: center; }
.mark {
width: 38px; height: 38px; border-radius: 9px; display: flex;
align-items: center; justify-content: center;
font: bold 21px Consolas, monospace; color: var(--accent);
background: var(--panel); border: 1px solid var(--line);
box-shadow: 0 4px 18px -8px #e0974a66;
}
.brand h1 { font-size: 19px; margin: 0; letter-spacing: -0.02em; }
.brand p { margin: 0; color: var(--muted); font-size: 12.5px; }
.top-actions { display: flex; flex-wrap: wrap; gap: 8px 12px; align-items: center; }
.pill {
font-size: 12.5px; color: var(--muted); border: 1px solid var(--line);
background: var(--panel); border-radius: 999px; padding: 4px 12px;
font-variant-numeric: tabular-nums; white-space: nowrap;
}
.pill[data-state="ok"] { color: var(--good); border-color: #3fb95055; }
.pill[data-state="bad"] { color: var(--bad); border-color: #f8514955; }
.pill.busy { animation: pulse 1.2s ease-in-out infinite; }
@keyframes pulse { 50% { opacity: 0.55; } }
.dash { display: grid; grid-template-columns: minmax(330px, 390px) 1fr; gap: 14px; }
.span { grid-column: 1 / -1; }
@media (max-width: 920px) { .dash { grid-template-columns: 1fr; } }
.panel {
background: var(--panel); border: 1px solid var(--line);
border-radius: var(--radius); padding: 16px 16px 18px; margin: 0;
box-shadow: 0 14px 34px -24px #000000cc;
}
.panel > h2 { font-size: 13px; margin: 0 0 12px; color: var(--muted);
font-weight: 600; letter-spacing: 0.04em; text-transform: uppercase; font-size: 12px; }
.panel > h2 .hint { text-transform: none; letter-spacing: 0; font-weight: 400; }
label { display: inline-block; margin: 4px 10px 4px 0; color: var(--muted); font-size: 13px; }
label.inline { margin: 0; }
input, textarea, select, button {
font: inherit; font-size: 13.5px; background: var(--panel-2); color: var(--text);
border: 1px solid var(--line); border-radius: var(--radius-s); padding: 5px 9px;
transition: border-color 0.2s, background 0.2s, transform 0.08s;
}
textarea { width: 100%; resize: vertical; line-height: 1.5; }
textarea#output, pre#output { font-family: ui-monospace, Consolas, monospace; }
input[type=number] { width: 84px; font-variant-numeric: tabular-nums; }
input[type=text], input[type=password] { width: 210px; }
select { max-width: 100%; }
input:focus-visible, textarea:focus-visible, select:focus-visible, button:focus-visible {
outline: 2px solid var(--accent); outline-offset: 1px;
}
button { cursor: pointer; }
button:hover:not(:disabled) { border-color: var(--accent); transform: translateY(-1px); }
button:active:not(:disabled) { transform: translateY(1px) scale(0.98); }
button.primary { background: var(--accent-dim); border-color: var(--accent-dim); color: #fff; font-weight: 600; }
button.primary:hover:not(:disabled) { background: #2f81f7; }
button.danger:hover:not(:disabled) { border-color: var(--bad); color: var(--bad); }
button:disabled { opacity: 0.5; cursor: default; }
.row { display: flex; flex-wrap: wrap; gap: 6px 14px; align-items: center; }
.row + .row { margin-top: 8px; }
.btnrow { display: flex; flex-wrap: wrap; gap: 8px; margin-top: 10px; }
.cards { display: grid; grid-template-columns: 1fr; gap: 10px; margin-top: 12px; }
.card {
background: var(--panel-2); border: 1px solid var(--line);
border-radius: var(--radius-s); padding: 10px 12px;
}
.card.active { border-color: var(--accent); }
.card-title { font-weight: 600; overflow-wrap: anywhere; margin-bottom: 6px; }
.badges { display: flex; flex-wrap: wrap; gap: 6px; margin-bottom: 8px; }
.badge {
font-size: 12px; color: var(--muted); border: 1px solid var(--line);
border-radius: 4px; padding: 1px 7px; font-variant-numeric: tabular-nums;
}
.badge.hot { color: var(--accent); border-color: #e0974a66; }
.card-actions { display: flex; gap: 8px; }
.empty { color: var(--muted); border: 1px dashed var(--line); border-radius: var(--radius-s);
padding: 14px; margin-top: 12px; font-size: 13px; }
#live { color: var(--accent); min-height: 20px; font-variant-numeric: tabular-nums; }
#output {
min-height: 120px; white-space: pre-wrap; word-break: break-word;
background: var(--panel-2); border: 1px solid var(--line); border-radius: var(--radius-s);
padding: 12px; margin: 8px 0 0;
}
#output:empty::before { content: "Run a generation to see output here."; color: var(--muted); }
#stats, #error { min-height: 20px; margin: 8px 0 0; font-variant-numeric: tabular-nums; }
#stats { color: var(--muted); }
#error { color: var(--bad); white-space: pre-wrap; }
#error:empty, #stats:empty { display: none; }
table { border-collapse: collapse; margin-top: 10px; width: 100%; }
th, td { padding: 5px 12px 5px 0; border-bottom: 1px solid var(--line);
font-variant-numeric: tabular-nums; }
th { color: var(--muted); font-weight: 600; font-size: 12px; text-transform: uppercase;
letter-spacing: 0.05em; text-align: left; }
td:last-child, th:last-child { padding-right: 0; }
.statusline { color: var(--muted); min-height: 18px; font-size: 13px; }
</style>
</head>
<body>
<a class="skip" href="#panel-request">Skip to generation</a>
<header class="topbar">
<div class="brand"><span class="mark" aria-hidden="true">C</span><div><h1>CISM tester</h1><p>native CPU inference console</p></div></div>
<div class="top-actions">
<span id="health" class="pill">health: unknown</span>
<label class="inline">API key <input id="key" type="password" autocomplete="off" placeholder="only if CISM_API_KEY is set"></label>
<label class="inline"><input id="remember" type="checkbox"> remember</label>
<button id="load">Refresh models</button>
</div>
</header>
<main class="dash">
<section class="panel" id="panel-models">
<h2>Models <span class="hint">fetch one, pin it down, compare it</span></h2>
<div class="row">
<label>Add model <input id="load-id" size="34" placeholder="HF repo id, e.g. SurjoLabs/Blaze"></label>
<label>precision
<select id="load-precision">
<option value="int8" selected>int8</option>
<option value="fp32">fp32</option>
<option value="hybrid-int4">hybrid-int4</option>
<option value="hybrid-fp4">hybrid-fp4</option>
</select>
</label>
<label>threads <input id="load-threads" type="number" min="1" max="64" value="2"></label>
<label>act-precision
<select id="load-act" title="activation precision for quantized matrices (empty = server default)">
<option value="" selected>server default</option>
<option value="fp32">fp32</option>
<option value="int8">int8 (experimental)</option>
</select>
</label>
<label>revision <input id="load-revision" size="12" placeholder="optional"></label>
<label><input id="load-trust" type="checkbox"> trust remote code</label>
<button id="load-model">Load</button>
<button id="unload-model">Unload selected</button>
</div>
<div id="load-status" class="statusline"></div>
<div class="row" style="margin-top:8px">
<button id="lock-toggle">Pin pages</button>
<label>keep-warm <input id="keep-warm-secs" type="number" min="0" value="0" title="seconds between sweeps; 0 with Stop turns it off"></label>
<button id="keep-warm-start">Start warm</button>
<button id="keep-warm-stop">Stop warm</button>
<span id="runtime-status" class="statusline"></span>
</div>
<div class="cards" id="model-cards"></div>
</section>
<section class="panel" id="panel-target">
<h2>Active model</h2>
<div class="row">
<label>Model <select id="model"></select></label>
</div>
<div class="row">
<span id="engine-info" style="color:#8b949e"></span>
</div>
</section>
<section class="panel" id="panel-request">
<h2>Request <span class="hint">say it plain, tune what matters</span></h2>
<div class="row">
<label><input type="radio" name="mode" value="completions" checked> completions</label>
<label><input type="radio" name="mode" value="chat"> chat</label>
<label><input id="stream" type="checkbox" checked> stream</label>
</div>
<label>System (chat only)
<textarea id="system" rows="2" placeholder="optional"></textarea>
</label>
<label>Prompt
<textarea id="prompt" rows="4">The future of artificial intelligence is</textarea>
</label>
<div class="row">
<button id="count-tokens">Count tokens</button>
<span id="token-info" class="statusline"></span>
</div>
<div class="row">
<label>max tokens <input id="max" type="number" min="1" value="64"></label>
<label>temperature <input id="temp" type="number" step="0.05" min="0" max="2" value="0"></label>
<label>top_p <input id="topp" type="number" step="0.01" min="0" max="1" value="1"></label>
<label>top_k <input id="topk" type="number" min="0" value="0"></label>
<label>seed <input id="seed" type="number" min="0" value="0"></label>
<label>spec K <input id="spec" type="number" min="0" max="16" value="0" title="greedy prompt-lookup speculation; 0=off, requires temperature 0"></label>
<label>stop <input id="stopwords" size="22" placeholder="comma-separated, optional"></label>
</div>
<div class="row">
<button id="go" class="primary">Generate</button>
<button id="stop" disabled>Stop</button>
<button id="bench">Bench (3x128)</button>
</div>
</section>
<section class="panel span" id="panel-measure">
<h2>Measure <span class="hint">numbers, not adjectives</span></h2>
<div class="row">
<button id="autobench">AutoBench</button>
<label>tokens <input id="bench-tokens" type="number" min="1" max="2048" value="512"></label>
<label title="sweep threads 1,2,4 on CISM and eager paths"><input id="bench-sweep" type="checkbox"> thread sweep</label>
<label title="time greedy spec K 0,2,4,8,12,16 per CISM variant"><input id="bench-spec" type="checkbox"> spec sweep</label>
<label title="include experimental int8 activations on the CISM side"><input id="bench-act8" type="checkbox"> act-int8</label>
<label title="uncheck to skip the slow torch baselines"><input id="bench-torch" type="checkbox" checked> torch</label>
<label title="compile baselines too (slow: minutes of inductor warmup)"><input id="bench-compile" type="checkbox"> compile</label>
<span id="bench-progress" style="color:#8b949e"></span>
</div>
<table id="bench-results" style="display:none"></table>
<div id="bench-extra"></div>
<h2 style="margin-top:16px">Precision shootout</h2>
<div class="row">
<label>Model <input id="shootout-id" size="34" placeholder="defaults to selected model"></label>
<label><input id="shoot-fp32" type="checkbox"> fp32</label>
<label><input id="shoot-int8" type="checkbox" checked> int8</label>
<label><input id="shoot-hybrid-int4" type="checkbox" checked> hybrid-int4</label>
<label><input id="shoot-hybrid-fp4" type="checkbox" checked> hybrid-fp4</label>
<label>tokens <input id="shootout-tokens" type="number" min="2" value="64"></label>
<label>runs <input id="shootout-runs" type="number" min="2" value="3" title="first run is warmup"></label>
<label>act <select id="shoot-act" title="activation precision for every shootout load">
<option value="" selected>server default</option>
<option value="fp32">fp32</option>
<option value="int8">int8 (experimental)</option>
</select></label>
<label>spec <input id="shoot-spec" type="number" min="0" max="16" value="0" title="prompt-lookup speculation per run; shootout always uses temperature 0"></label>
<button id="shootout">Run shootout</button>
<span id="shootout-status" style="color:#8b949e"></span>
</div>
<table id="shootout-results" style="display:none"></table>
<h2 style="margin-top:16px">Perplexity scoring</h2>
<div class="row">
<span style="color:#a89c8d">scores the prompt above on the selected model</span>
<label>window <input id="ppl-window" type="number" min="2" value="128"></label>
<button id="ppl-score">Score prompt</button>
<span id="ppl-status" class="statusline"></span>
</div>
<div id="ppl-results"></div>
<h2 style="margin-top:16px">Cache residency</h2>
<div class="row">
<span style="color:#a89c8d">uses the selected model</span>
<label>decode tokens <input id="cache-tokens" type="number" min="1" value="64"></label>
<label>scrub MiB <input id="cache-scrub" type="number" min="1" value="96" title="cache-eviction buffer, sized near 2x L3"></label>
<button id="cache-run">Run cache report</button>
<span id="cache-status" class="statusline"></span>
</div>
<div id="cache-results"></div>
</section>
<section class="panel span" id="panel-output">
<h2>Output</h2>
<div id="live" style="color:#58a6ff; min-height:18px"></div>
<pre id="output"></pre>
<div id="stats"></div>
<div id="error"></div>
<div class="btnrow">
<button id="copy">Copy output</button>
</div>
</section>
</main>
<script>
"use strict";
const $ = id => document.getElementById(id);
const output = $("output"), stats = $("stats"), errorBox = $("error"), live = $("live");
let controller = null;
const contexts = {};
const engineInfo = {};
function capMax() {
const limit = contexts[$("model").value];
if (!limit) return;
$("max").max = limit;
if (Number($("max").value) > limit) $("max").value = limit;
}
function showEngineInfo() {
const info = engineInfo[$("model").value] || {};
const parts = [];
if (info.precision) parts.push(info.precision);
if (info.architecture) parts.push(info.architecture);
if (info.threads) parts.push(info.threads + " thread" + (info.threads > 1 ? "s" : ""));
if (info.kernel) parts.push(info.kernel);
if (info.act_precision) parts.push("act " + info.act_precision);
if (info.quantization) parts.push(info.quantization);
if (info.weight_mb !== undefined) parts.push(info.weight_mb + " MB");
if (info.context_length) parts.push("ctx " + info.context_length);
$("engine-info").textContent = parts.join(" · ");
}
function withWhitespaceHint(text) {
if (output.textContent && !/\\S/.test(output.textContent))
return text + " | all-whitespace output: greedy loop on a space token; try temperature > 0";
return text;
}
function headers() {
const key = $("key").value.trim();
return key ? { "Authorization": "Bearer " + key } : {};
}
function showError(message) {
errorBox.textContent = typeof message === "string" ? message : JSON.stringify(message);
}
function mode() {
return document.querySelector('input[name=mode]:checked').value;
}
async function loadModels() {
errorBox.textContent = "";
try {
const response = await fetch("/v1/models", { headers: headers() });
const body = await response.json();
if (!response.ok) { showError(body.error && body.error.message || response.status); return; }
const select = $("model");
select.replaceChildren(...body.data.map(m => {
const option = document.createElement("option");
option.value = m.id;
// A select renders option text, not option value; both must be set.
option.textContent = m.id;
if (m.context_length) contexts[m.id] = m.context_length;
engineInfo[m.id] = m;
return option;
}));
if (!select.value && body.data.length) select.value = body.data[0].id;
capMax();
showEngineInfo();
refreshHealth();
renderModelCards(body.data);
} catch (err) {
showError(String(err));
}
}
function requestBody(streaming, maxOverride) {
const spec = Number($("spec").value) || 0;
if (spec && Number($("temp").value) !== 0) {
throw new Error("Speculation (spec K > 0) requires temperature 0");
}
const body = {
model: $("model").value.trim(),
max_tokens: maxOverride || Math.min(Number($("max").value) || 64,
contexts[$("model").value] || Infinity),
temperature: Number($("temp").value) || 0,
top_p: Number($("topp").value) || 1,
top_k: Number($("topk").value) || 0,
seed: Number($("seed").value) || 0,
spec_k: Number($("spec").value) || 0,
stream: streaming,
};
const stops = $("stopwords").value.split(",").map(s => s.trim()).filter(s => s);
if (stops.length === 1) body.stop = stops[0];
else if (stops.length > 1) body.stop = stops;
if (mode() === "chat") {
const messages = [];
const system = $("system").value.trim();
if (system) messages.push({ role: "system", content: system });
messages.push({ role: "user", content: $("prompt").value });
body.messages = messages;
} else {
body.prompt = $("prompt").value;
}
if (streaming) body.stream_options = { include_usage: true };
return body;
}
let loadWatch = null;
async function pollLoadStatus() {
// Returns "busy" while a load is in flight; "idle" once it settles.
try {
const response = await fetch("/v1/models/loading", { headers: headers() });
if (!response.ok) return "idle";
const body = await response.json();
const entries = Object.entries(body.loading || {});
const active = entries.filter(([, s]) => s === "loading");
const failed = entries.filter(([, s]) => s.startsWith("failed"));
if (failed.length) {
showError(failed.map(([id, s]) => id + ": " + s).join(" | "));
$("load-status").textContent = "";
}
if (active.length) {
$("load-status").textContent = active.map(([id, s]) => id + ": " + s).join(" | ");
return "busy";
}
// A watched load finished: the entry was removed server-side.
if (loadWatch) { $("load-status").textContent = loadWatch + ": loaded"; loadWatch = null; }
return "idle";
} catch (err) {
return "busy";
}
}
async function loadModel() {
const id = $("load-id").value.trim();
if (!id) { showError("Enter a model id to load"); return; }
errorBox.textContent = "";
$("load-status").textContent = id + ": starting…";
try {
const response = await fetch("/v1/models/load", {
method: "POST",
headers: Object.assign({ "Content-Type": "application/json" }, headers()),
body: JSON.stringify({
model: id,
precision: $("load-precision").value,
act_precision: $("load-act").value || null,
revision: $("load-revision").value.trim() || null,
threads: Number($("load-threads").value) || 2,
trust_remote_code: $("load-trust").checked,
}),
});
const body = await response.json();
if (!response.ok) {
$("load-status").textContent = "";
showError(body.error && body.error.message || "HTTP " + response.status);
return;
}
if (body.status === "ready") {
$("load-status").textContent = id + ": already loaded";
await loadModels();
return;
}
$("load-id").value = "";
loadWatch = id;
const poll = setInterval(async () => {
const state = await pollLoadStatus();
if (state !== "busy") {
clearInterval(poll);
if (state === "idle") await loadModels();
}
}, 1500);
} catch (err) {
showError(String(err));
}
}
function renderModelCards(models) {
const host = $("model-cards");
host.replaceChildren();
if (!models.length) {
const hint = document.createElement("div");
hint.className = "empty";
hint.textContent = "No models loaded. Enter an HF repo id above — or run a shootout, which loads each precision for you.";
host.appendChild(hint);
return;
}
for (const m of models) {
const card = document.createElement("article");
card.className = "card" + (m.id === $("model").value ? " active" : "");
const title = document.createElement("div");
title.className = "card-title";
title.textContent = m.id;
const badges = document.createElement("div");
badges.className = "badges";
const items = [];
if (m.precision) items.push([m.precision, false]);
if (m.act_precision) items.push(["act " + m.act_precision, m.act_precision !== "fp32"]);
if (m.weight_mb !== undefined) items.push([m.weight_mb + " MB", false]);
if (m.threads) items.push([m.threads + "t", false]);
if (m.kernel) items.push([m.kernel, false]);
for (const [text, hot] of items) {
const badge = document.createElement("span");
badge.className = "badge" + (hot ? " hot" : "");
badge.textContent = text;
badges.appendChild(badge);
}
const actions = document.createElement("div");
actions.className = "card-actions";
const use = document.createElement("button");
use.textContent = "Use";
use.addEventListener("click", () => {
$("model").value = m.id;
capMax();
showEngineInfo();
renderModelCards(models);
});
const drop = document.createElement("button");
drop.textContent = "Unload";
drop.className = "danger";
drop.addEventListener("click", () => { unloadModel(m.id); });
actions.append(use, drop);
card.append(title, badges, actions);
host.appendChild(card);
}
}
async function unloadModel(id) {
id = id || $("model").value;
if (!id) return;
errorBox.textContent = "";
try {
const response = await fetch("/v1/models/unload", {
method: "POST",
headers: Object.assign({ "Content-Type": "application/json" }, headers()),
body: JSON.stringify({ model: id }),
});
const body = await response.json();
if (!response.ok) { showError(body.error && body.error.message || "HTTP " + response.status); return; }
$("load-status").textContent = id + ": unloaded";
await loadModels();
} catch (err) {
showError(String(err));
}
}
// One generation; returns measurement or null on handled failure.
async function runGeneration(maxOverride) {
output.textContent = "";
stats.textContent = "";
live.textContent = "";
const started = performance.now();
let firstToken = 0, finish = null, usage = null;
const streaming = $("stream").checked;
const response = await fetch("/v1/" + (mode() === "chat" ? "chat/completions" : "completions"), {
method: "POST",
headers: Object.assign({ "Content-Type": "application/json" }, headers()),
body: JSON.stringify(requestBody(streaming, maxOverride)),
signal: controller.signal,
});
if (!response.ok) {
const body = await response.json().catch(() => null);
showError(body && body.error ? body.error.message : "HTTP " + response.status);
return null;
}
if (!streaming) {
const body = await response.json();
const choice = body.choices[0];
output.textContent = mode() === "chat" ? choice.message.content : choice.text;
const elapsed = (performance.now() - started) / 1000;
let summary = "finish=" + choice.finish_reason +
" prompt=" + body.usage.prompt_tokens + " completion=" + body.usage.completion_tokens +
" total=" + body.usage.total_tokens + " elapsed=" + elapsed.toFixed(3) + "s";
if (body.usage.spec_accept_rate !== undefined && body.usage.spec_accept_rate !== null)
summary += " accept=" + Math.round(body.usage.spec_accept_rate * 100) + "%";
stats.textContent = withWhitespaceHint(summary);
return { tokens: body.usage.completion_tokens, ttft: null, seconds: elapsed, decode: null };
}
const reader = response.body.getReader();
const decoder = new TextDecoder();
let buffer = "";
for (;;) {
const block = await reader.read();
if (block.done) break;
buffer += decoder.decode(block.value, { stream: true });
let index;
while ((index = buffer.indexOf("\\n\\n")) >= 0) {
const frame = buffer.slice(0, index);
buffer = buffer.slice(index + 2);
for (const line of frame.split("\\n")) {
if (!line.startsWith("data: ")) continue;
const payload = line.slice(6);
if (payload === "[DONE]") continue;
const event = JSON.parse(payload);
if (event.usage) usage = event.usage;
if (event.cism_completion_tokens !== undefined) {
const rate = firstToken
? " · " + (event.cism_completion_tokens / ((performance.now() - firstToken) / 1000)).toFixed(0) + " tok/s live"
: "";
live.textContent = event.cism_completion_tokens + " tokens" + rate;
}
const choice = event.choices && event.choices[0];
if (!choice) continue;
if (choice.finish_reason) finish = choice.finish_reason;
const text = mode() === "chat"
? (choice.delta && choice.delta.content) || ""
: choice.text || "";
if (text && !firstToken) firstToken = performance.now();
if (text) output.textContent += text;
}
}
}
const parts = [];
if (finish) parts.push("finish=" + finish);
if (usage) {
parts.push("prompt=" + usage.prompt_tokens, "completion=" + usage.completion_tokens,
"total=" + usage.total_tokens);
if (usage.spec_accept_rate !== undefined && usage.spec_accept_rate !== null)
parts.push("accept=" + Math.round(usage.spec_accept_rate * 100) + "%");
if (firstToken) parts.push("decode=" +
(usage.completion_tokens / ((performance.now() - firstToken) / 1000)).toFixed(1) + " tok/s");
}
parts.push("ttft=" + (firstToken ? ((firstToken - started) / 1000).toFixed(3) + "s" : "n/a"));
stats.textContent = withWhitespaceHint(parts.join(" "));
return { tokens: usage ? usage.completion_tokens : 0,
ttft: firstToken ? (firstToken - started) / 1000 : null,
seconds: (performance.now() - started) / 1000,
decode: firstToken && usage ? usage.completion_tokens / ((performance.now() - firstToken) / 1000) : null };
}
async function generate() {
errorBox.textContent = "";
$("go").disabled = true;
$("bench").disabled = true;
$("stop").disabled = false;
controller = new AbortController();
try {
await runGeneration(null);
} catch (err) {
if (err && err.name === "AbortError") stats.textContent += " (cancelled)";
else showError(String(err));
} finally {
$("go").disabled = false;
$("bench").disabled = false;
$("stop").disabled = true;
controller = null;
}
}
async function runAutoBench() {
const model = $("model").value.trim();
if (!model) { showError("Select a model first"); return; }
errorBox.textContent = "";
$("autobench").disabled = true;
$("bench-progress").textContent = "starting…";
$("bench-results").style.display = "none";
try {
const response = await fetch("/v1/bench/auto", {
method: "POST",
headers: Object.assign({ "Content-Type": "application/json" }, headers()),
// No precision selector: every AutoBench evaluates all four CISM
// variants (fp32, int8, hybrid-int4, hybrid-fp4) in one run.
body: JSON.stringify({ model: model,
threads: Number($("load-threads").value) || 2,
max_tokens: Number($("bench-tokens").value) || 512,
thread_sweep: $("bench-sweep").checked,
spec_sweep: $("bench-spec").checked,
bench_act_int8: $("bench-act8").checked,
include_torch: $("bench-torch").checked,
compile: $("bench-compile").checked }),
});
const body = await response.json();
if (!response.ok) {
$("bench-progress").textContent = "";
showError(body.error && body.error.message || "HTTP " + response.status);
// A 409 means a bench is already running server-side (perhaps started
// from another tab); adopt and display it instead of leaving the UI
// claiming idle.
if (response.status === 409) syncBenchState();
else $("autobench").disabled = false;
return;
}
syncBenchState();
} catch (err) {
$("autobench").disabled = false;
$("bench-progress").textContent = "";
showError(String(err));
}
}
let benchPoll = null;
function syncBenchState() {
// One shared poller for every tab/click; tracks the server-side bench.
if (benchPoll) return;
benchPoll = setInterval(async () => {
let state;
try { state = await (await fetch("/v1/bench/auto", { headers: headers() })).json(); }
catch { return; }
if (state.status === "running") {
$("autobench").disabled = true;
$("bench-progress").textContent = state.progress || "running…";
return;
}
clearInterval(benchPoll);
benchPoll = null;
$("autobench").disabled = false;
if (state.status === "done" && state.result) {
$("bench-progress").textContent = "done";
renderBench(state.result);
} else if (state.status === "failed") {
$("bench-progress").textContent = "";
showError(state.error || "benchmark failed");
}
}, 1000);
}
function benchCell(value) {
return value === undefined || value === null ? "—" : String(value);
}
function renderBench(result) {
if (!result) return;
const table = $("bench-results");
const rows = [];
const fmt = value => value === undefined || value === null ? "—" :
Number(value).toLocaleString(undefined, { maximumFractionDigits: 2 });
if (result.cism) {
let label = "CISM (" + result.precision + ")";
if (result.cism.weight_mb !== undefined) {
label += " · " + result.cism.weight_mb + " MB";
if (result.cism.fits_l3 === true) label += " · L3-resident";
else if (result.cism.fits_l3 === false) label += " · DRAM-streaming";
}
rows.push([label, fmt(result.cism.tokens_per_second),
fmt(result.cism.perplexity && result.cism.perplexity.perplexity)]);
}
if (result.cism_fp32) {
let label = "CISM fp32";
if (result.cism_fp32.weight_mb !== undefined) {
label += " · " + result.cism_fp32.weight_mb + " MB";
if (result.cism_fp32.fits_l3 === true) label += " · L3-resident";
else if (result.cism_fp32.fits_l3 === false) label += " · DRAM-streaming";
}
rows.push([label, fmt(result.cism_fp32.tokens_per_second),
fmt(result.cism_fp32.perplexity && result.cism_fp32.perplexity.perplexity)]);
}
if (result.cism_hybrid_int4) {
let label = "CISM hybrid-int4";
if (result.cism_hybrid_int4.weight_mb !== undefined) {
label += " · " + result.cism_hybrid_int4.weight_mb + " MB";
if (result.cism_hybrid_int4.fits_l3 === true) label += " · L3-resident";
else if (result.cism_hybrid_int4.fits_l3 === false) label += " · DRAM-streaming";
}
rows.push([label, fmt(result.cism_hybrid_int4.tokens_per_second),
fmt(result.cism_hybrid_int4.perplexity && result.cism_hybrid_int4.perplexity.perplexity)]);
}
if (result.cism_hybrid_fp4) {
let label = "CISM hybrid-fp4";
if (result.cism_hybrid_fp4.weight_mb !== undefined) {
label += " · " + result.cism_hybrid_fp4.weight_mb + " MB";
if (result.cism_hybrid_fp4.fits_l3 === true) label += " · L3-resident";
else if (result.cism_hybrid_fp4.fits_l3 === false) label += " · DRAM-streaming";
}
rows.push([label, fmt(result.cism_hybrid_fp4.tokens_per_second),
fmt(result.cism_hybrid_fp4.perplexity && result.cism_hybrid_fp4.perplexity.perplexity)]);
}
if (result.hf_eager) {
rows.push(["HF eager fp32 (generate)", fmt(result.hf_eager.tokens_per_second),
fmt(result.hf_eager.perplexity && result.hf_eager.perplexity.perplexity)]);
}
if (result.hf_eager_int8) {
rows.push(result.hf_eager_int8.failed
? ["HF eager int8 (failed)", result.hf_eager_int8.failed, "—"]
: ["HF eager int8 (dynamic)", fmt(result.hf_eager_int8.tokens_per_second),
fmt(result.hf_eager_int8.perplexity && result.hf_eager_int8.perplexity.perplexity)]);
}
if (result.hf_compile) {
rows.push(result.hf_compile.failed
? ["HF torch.compile fp32 (failed)", result.hf_compile.failed, "—"]
: ["HF torch.compile fp32", fmt(result.hf_compile.tokens_per_second), "—"]);
}
if (result.hf_compile_int8) {
rows.push(result.hf_compile_int8.failed
? ["HF torch.compile int8 (failed)", result.hf_compile_int8.failed, "—"]
: ["HF torch.compile int8", fmt(result.hf_compile_int8.tokens_per_second), "—"]);
}
if (result.hf_eager_int4) {
rows.push(result.hf_eager_int4.failed || result.hf_eager_int4.skipped
? ["HF eager int4 (torchao)", result.hf_eager_int4.failed || result.hf_eager_int4.skipped, "—"]
: ["HF eager int4 (torchao)", fmt(result.hf_eager_int4.tokens_per_second),
fmt(result.hf_eager_int4.perplexity && result.hf_eager_int4.perplexity.perplexity)]);
}
if (result.hf_compile_int4) {
rows.push(result.hf_compile_int4.failed
? ["HF torch.compile int4 (failed)", result.hf_compile_int4.failed, "—"]
: ["HF torch.compile int4", fmt(result.hf_compile_int4.tokens_per_second), "—"]);
}
if (result.torch_baselines && result.torch_baselines.skipped) {
rows.push(["HF baselines", result.torch_baselines.skipped, "—"]);
}
const speedups = [];
if (result.speedup_vs_eager) speedups.push(result.speedup_vs_eager + "x vs eager");
if (result.speedup_vs_compile) speedups.push(result.speedup_vs_compile + "x vs compile");
if (result.ppl_delta_percent !== null && result.ppl_delta_percent !== undefined) {
speedups.push((result.ppl_delta_percent >= 0 ? "+" : "") + result.ppl_delta_percent + "% PPL vs fp32");
}
if (result.hybrid_int4_speedup_vs_eager) speedups.push("int4: " + result.hybrid_int4_speedup_vs_eager + "x vs eager");
if (result.hybrid_int4_speedup_vs_compile) speedups.push("int4: " + result.hybrid_int4_speedup_vs_compile + "x vs compile");
if (result.hybrid_int4_ppl_delta_percent !== null && result.hybrid_int4_ppl_delta_percent !== undefined) {
speedups.push("int4: " + (result.hybrid_int4_ppl_delta_percent >= 0 ? "+" : "") +
result.hybrid_int4_ppl_delta_percent + "% PPL vs fp32");
}
if (result.hybrid_fp4_speedup_vs_eager) speedups.push("fp4: " + result.hybrid_fp4_speedup_vs_eager + "x vs eager");
if (result.hybrid_fp4_speedup_vs_compile) speedups.push("fp4: " + result.hybrid_fp4_speedup_vs_compile + "x vs compile");
if (result.hybrid_fp4_ppl_delta_percent !== null && result.hybrid_fp4_ppl_delta_percent !== undefined) {
speedups.push("fp4: " + (result.hybrid_fp4_ppl_delta_percent >= 0 ? "+" : "") +
result.hybrid_fp4_ppl_delta_percent + "% PPL vs fp32");
}
table.replaceChildren(...[
(() => {
const head = document.createElement("tr");
["method", "tok/s", "PPL"].forEach((title, index) => {
const th = document.createElement("th");
th.textContent = title;
th.style.textAlign = index ? "right" : "left";
head.appendChild(th);
});
return head;
})(),
...rows.map(cells => {
const tr = document.createElement("tr");
cells.forEach((cell, index) => {
const td = document.createElement("td");
td.textContent = benchCell(cell);
td.style.textAlign = index ? "right" : "left";
td.style.paddingRight = "12px";
tr.appendChild(td);
});
return tr;
}),
(() => {
const tr = document.createElement("tr");
const td = document.createElement("td");
td.colSpan = 3;
td.textContent = speedups.join(" | ");
td.style.paddingTop = "6px";
tr.appendChild(td);
return tr;
})(),
]);
table.style.display = "table";
renderBenchExtra(result);
}
async function bench() {
// Three fixed 128-token greedy runs; first is warmup, report the rest.
const previousMax = $("max").value, previousTemp = $("temp").value, previousSpec = $("spec").value;
$("max").value = 128; $("temp").value = 0; $("spec").value = 0;
errorBox.textContent = "";
$("go").disabled = true;
$("bench").disabled = true;
$("stop").disabled = false;
try {
const rates = [], ttfts = [];
for (let run = 0; run < 3; ++run) {
controller = new AbortController();
const result = await runGeneration(128);
if (!result) return;
if (run > 0) {
if (result.decode) rates.push(result.decode);
if (result.ttft !== null) ttfts.push(result.ttft);
}
}
const median = values => values.slice().sort((a, b) => a - b)[Math.floor(values.length / 2)];
const parts = ["bench (3x128, run 1 discarded):"];
if (rates.length) parts.push("median decode " + median(rates).toFixed(1) + " tok/s",
"best " + Math.max(...rates).toFixed(1));
if (ttfts.length) parts.push("median ttft " + median(ttfts).toFixed(3) + "s");
stats.textContent = parts.join(" ") + " | params restored";
} catch (err) {
if (err && err.name === "AbortError") stats.textContent += " (cancelled)";
else showError(String(err));
} finally {
$("max").value = previousMax; $("temp").value = previousTemp; $("spec").value = previousSpec;
capMax();
$("go").disabled = false;
$("bench").disabled = false;
$("stop").disabled = true;
controller = null;
}
}
$("remember").addEventListener("change", () => {
if ($("remember").checked) localStorage.setItem("cism-webui-key", $("key").value);
else localStorage.removeItem("cism-webui-key");
});
$("key").addEventListener("input", () => {
if ($("remember").checked) localStorage.setItem("cism-webui-key", $("key").value);
});
$("load").addEventListener("click", loadModels);
$("model").addEventListener("change", () => {
capMax();
showEngineInfo();
renderModelCards(Object.values(engineInfo));
});
$("load-model").addEventListener("click", loadModel);
$("unload-model").addEventListener("click", () => unloadModel());
$("bench").addEventListener("click", bench);
$("autobench").addEventListener("click", runAutoBench);
$("go").addEventListener("click", generate);
$("stop").addEventListener("click", () => { if (controller) controller.abort(); });
async function refreshHealth() {
const pill = $("health");
try {
const h = await (await fetch("/health")).json();
let text = "health: " + (h.status || "?");
let ok = h.status === "ok";
try {
const r = await (await fetch("/ready")).json();
text += " · ready: " + (r.status || "?");
ok = ok && r.status === "ready";
} catch { text += " · ready: n/a (worker busy or starting)"; }
pill.textContent = text;
pill.dataset.state = ok ? "ok" : "";
} catch {
pill.textContent = "health: unreachable";
pill.dataset.state = "bad";
}
}
$("copy").addEventListener("click", async () => {
errorBox.textContent = "";
try { await navigator.clipboard.writeText(output.textContent); }
catch (err) { showError(String(err)); }
});
async function waitForLoadIdle(id) {
for (;;) {
const response = await fetch("/v1/models/loading", { headers: headers() });
const body = await response.json();
const state = (body.loading || {})[id];
if (!state) return;
if (state.startsWith("failed")) throw new Error(id + ": " + state);
if (state !== "loading") return;
$("shootout-status").textContent = id + ": " + state + "…";
await new Promise(resolve => setTimeout(resolve, 1500));
}
}
async function shootoutTimedRun(maxTokens) {
// One non-streaming greedy run; returns tok/s or throws on HTTP error.
const started = performance.now();
const response = await fetch("/v1/completions", {
method: "POST",
headers: Object.assign({ "Content-Type": "application/json" }, headers()),
body: JSON.stringify({ model: $("shootout-id").value.trim(), prompt: $("prompt").value,
max_tokens: maxTokens, temperature: 0, stream: false,
spec_k: Number($("shoot-spec").value) || 0 }),
});
const body = await response.json().catch(() => null);
if (!response.ok) throw new Error(body && body.error ? body.error.message : "HTTP " + response.status);
return body.usage.completion_tokens / ((performance.now() - started) / 1000);
}
function renderShootout(rows) {
const table = $("shootout-results");
const head = document.createElement("tr");
["precision", "median tok/s"].forEach((title, index) => {
const th = document.createElement("th");
th.textContent = title;
th.style.textAlign = index ? "right" : "left";
head.appendChild(th);
});
const bodyRows = rows.map(cells => {
const tr = document.createElement("tr");
cells.forEach((cell, index) => {
const td = document.createElement("td");
td.textContent = cell;
td.style.textAlign = index ? "right" : "left";
td.style.paddingRight = "12px";
tr.appendChild(td);
});
return tr;
});
table.replaceChildren(head, ...bodyRows);
table.style.display = "table";
}
function renderBenchExtra(result) {
const host = $("bench-extra");
host.replaceChildren();
if (!result || !result.matrix || !result.matrix.length) return;
const fmt = value => value === undefined || value === null ? "—" :
Number(value).toLocaleString(undefined, { maximumFractionDigits: 1 });
const variants = [...new Set(result.matrix.map(r => r.precision + (r.act !== "fp32" ? " " + r.act : "")))];
const keyOf = r => r.precision + (r.act !== "fp32" ? " " + r.act : "");
const pivot = (title, rows, preformatted) => {
const table = document.createElement("table");
const head = document.createElement("tr");
["", ...variants].forEach((text, index) => {
const th = document.createElement("th");
th.textContent = text;
th.style.textAlign = index ? "right" : "left";
head.appendChild(th);
});
table.appendChild(head);
for (const row of rows) {
const tr = document.createElement("tr");
const label = document.createElement("td");
label.textContent = row.label;
tr.appendChild(label);
for (const v of variants) {
const td = document.createElement("td");
td.textContent = preformatted ? row.cells[v] : fmt(row.cells[v]);
td.style.textAlign = "right";
td.style.paddingRight = "12px";
tr.appendChild(td);
}
table.appendChild(tr);
}
const caption = document.createElement("div");
caption.textContent = title;
caption.style.color = "#a89c8d";
caption.style.margin = "12px 0 2px";
host.append(caption, table);
};
const threads = [...new Set(result.matrix.map(r => r.threads))].sort((a, b) => a - b);
if (threads.length > 1 || result.spec_ks.length > 1 || result.act_modes.length > 1) {
pivot("tok/s by threads (spec 0, fp32 act)",
threads.map(t => ({ label: t + " threads",
cells: Object.fromEntries(variants.map(v => [v, (result.matrix.find(
r => keyOf(r) === v && r.threads === t && r.spec_k === 0) || {}).tokens_per_second])) })));
const specRows = [...new Set(result.matrix.map(r => r.spec_k))].sort((a, b) => a - b);
if (specRows.length > 1) {
const specThreads = [...new Set(result.matrix.map(r => r.threads))].sort((a, b) => a - b);
for (const th of specThreads) {
const specCell = (v, k) => {
const row = result.matrix.find(
r => keyOf(r) === v && r.spec_k === k && r.threads === th) || {};
let text = fmt(row.tokens_per_second);
if (row.accept_rate !== undefined && row.accept_rate !== null)
text += " (" + Math.round(row.accept_rate * 100) + "% acc)";
return text;
};
pivot("tok/s by spec_k (" + th + " threads, fp32 act)",
specRows.map(k => ({ label: "spec " + k,
cells: Object.fromEntries(variants.map(v => [v, specCell(v, k)])) })), true);
}
}
if (result.best) {
const lines = Object.entries(result.best).map(([setup, info]) =>
setup + ": " + fmt(info.tokens_per_second) + " tok/s @" + info.threads + "t spec " + info.spec_k);
const note = document.createElement("div");
note.textContent = "best: " + lines.join(" · ");
note.style.color = "#a89c8d";
note.style.marginTop = "10px";
host.appendChild(note);
}
}
}
async function runShootout() {
const id = $("shootout-id").value.trim() || $("model").value.trim();
if (!id) { showError("Enter a shootout model id"); return; }
$("shootout-id").value = id;
const precisions = ["fp32", "int8", "hybrid-int4", "hybrid-fp4"].filter(p => $("shoot-" + p).checked);
if (!precisions.length) { showError("Tick at least one precision"); return; }
const maxTokens = Math.min(Number($("shootout-tokens").value) || 64, contexts[id] || Infinity);
const runs = Math.max(2, Number($("shootout-runs").value) || 3);
errorBox.textContent = "";
$("shootout").disabled = true;
$("shootout-results").style.display = "none";
const rows = [];
try {
for (const precision of precisions) {
$("shootout-status").textContent = id + " " + precision + ": loading…";
const response = await fetch("/v1/models/load", {
method: "POST",
headers: Object.assign({ "Content-Type": "application/json" }, headers()),
body: JSON.stringify({ model: id, precision: precision,
threads: Number($("load-threads").value) || 2,
act_precision: $("shoot-act").value || null }),
});
const body = await response.json().catch(() => null);
if (!response.ok) throw new Error(body && body.error ? body.error.message : "HTTP " + response.status);
await waitForLoadIdle(id);
await loadModels();
$("shootout-status").textContent = id + " " + precision + ": benching…";
const rates = [];
for (let run = 0; run < runs; ++run) {
const rate = await shootoutTimedRun(maxTokens);
if (run > 0) rates.push(rate);
}
rates.sort((a, b) => a - b);
const info = engineInfo[id] || {};
let label = precision;
if (info.weight_mb !== undefined) label += " · " + info.weight_mb + " MB";
rows.push([label, rates[Math.floor(rates.length / 2)].toFixed(1)]);
renderShootout(rows);
}
$("shootout-status").textContent = "done";
} catch (err) {
showError(String((err && err.message) || err));
$("shootout-status").textContent = "";
} finally {
$("shootout").disabled = false;
}
}
$("shootout").addEventListener("click", runShootout);
async function runCacheReport() {
const id = $("model").value.trim();
if (!id) { showError("Select a model first"); return; }
errorBox.textContent = "";
$("cache-run").disabled = true;
$("cache-status").textContent = "scanning…";
try {
const response = await fetch("/v1/cache/report", {
method: "POST",
headers: Object.assign({ "Content-Type": "application/json" }, headers()),
body: JSON.stringify({ model: id,
decode_tokens: Number($("cache-tokens").value) || 64,
scrub_mb: Number($("cache-scrub").value) || 96 }),
});
const body = await response.json().catch(() => null);
if (!response.ok) throw new Error(body && body.error ? body.error.message : "HTTP " + response.status);
$("cache-status").textContent = "done";
renderCache(body);
} catch (err) {
showError(String((err && err.message) || err));
$("cache-status").textContent = "";
} finally {
$("cache-run").disabled = false;
}
}
function renderCache(result) {
const host = $("cache-results");
host.replaceChildren();
const rows = [
["decode", (result.decode_tokens_per_second || "—") + " tok/s"],
["effective", (result.decode_effective_gbps || "—") + " GB/s"],
["warm scan", (result.warm_scan_gbps || "—") + " GB/s"],
["cold scan", (result.cold_scan_gbps || "—") + " GB/s"],
["warm/cold", result.warm_to_cold_ratio || "—"],
["residency", result.residency || "—"],
["bottleneck", result.decode_bottleneck || "—"],
];
for (const [name, value] of rows) {
const div = document.createElement("div");
const key = document.createElement("span");
key.textContent = name + ": ";
key.style.color = "#a89c8d";
const val = document.createElement("span");
val.textContent = String(value);
div.append(name, val);
host.appendChild(div);
}
}
$("cache-run").addEventListener("click", runCacheReport);
function fmtMB(bytes) {
if (bytes === undefined || bytes === null) return "—";
return bytes >= 1048576 ? (bytes / 1048576).toFixed(1) + " MB" : Math.round(bytes / 1024) + " KB";
}
async function countTokens() {
const id = $("model").value.trim();
if (!id) { showError("Select a model first"); return; }
$("token-info").textContent = "counting…";
try {
const response = await fetch("/v1/tokenize", {
method: "POST",
headers: Object.assign({ "Content-Type": "application/json" }, headers()),
body: JSON.stringify({ model: id, text: $("prompt").value }),
});
const body = await response.json().catch(() => null);
if (!response.ok) throw new Error(body && body.error ? body.error.message : "HTTP " + response.status);
const parts = [body.count + " tokens"];
if (body.context_length) {
parts.push(body.fits_context === false ? "exceeds context " + body.context_length : "of context " + body.context_length);
}
if (body.kv_cache) {
parts.push("KV at prompt " + fmtMB(body.kv_cache.bytes_at_prompt) +
" · full context " + fmtMB(body.kv_cache.bytes_at_full_context));
}
$("token-info").textContent = parts.join(" · ");
} catch (err) {
showError(String((err && err.message) || err));
$("token-info").textContent = "";
}
}
async function scorePrompt() {
const id = $("model").value.trim();
if (!id) { showError("Select a model first"); return; }
errorBox.textContent = "";
$("ppl-status").textContent = "scoring…";
$("ppl-results").replaceChildren();
try {
const response = await fetch("/v1/nll", {
method: "POST",
headers: Object.assign({ "Content-Type": "application/json" }, headers()),
body: JSON.stringify({ model: id, text: $("prompt").value,
window: Number($("ppl-window").value) || 128 }),
});
const body = await response.json().catch(() => null);
if (!response.ok) throw new Error(body && body.error ? body.error.message : "HTTP " + response.status);
$("ppl-status").textContent = "done";
const host = $("ppl-results");
const rows = [
["perplexity", String(Math.round(body.perplexity * 1000) / 1000)],
["mean NLL", String(Math.round(body.nll * 10000) / 10000)],
["scored tokens", String(body.scored_tokens)],
];
for (const [name, value] of rows) {
const div = document.createElement("div");
const key = document.createElement("span");
key.textContent = name + ": ";
key.style.color = "#a89c8d";
const val = document.createElement("span");
val.textContent = value;
div.append(name, val);
host.appendChild(div);
}
} catch (err) {
showError(String((err && err.message) || err));
$("ppl-status").textContent = "";
}
}
const pinnedPages = {};
async function togglePin() {
const id = $("model").value.trim();
if (!id) { showError("Select a model first"); return; }
errorBox.textContent = "";
const pin = !pinnedPages[id];
$("runtime-status").textContent = pin ? "pinning…" : "unpinning…";
try {
const response = await fetch("/v1/models/runtime", {
method: "POST",
headers: Object.assign({ "Content-Type": "application/json" }, headers()),
body: JSON.stringify({ model: id, lock_pages: pin }),
});
const body = await response.json().catch(() => null);
if (!response.ok) throw new Error(body && body.error ? body.error.message : "HTTP " + response.status);
pinnedPages[id] = pin;
$("lock-toggle").textContent = pin ? "Unpin pages" : "Pin pages";
const bytes = pin ? body.page_lock_bytes : body.page_unlock_bytes;
$("runtime-status").textContent = id + (pin ? " pinned " : " unpinned ") + fmtMB(bytes);
} catch (err) {
showError(String((err && err.message) || err));
$("runtime-status").textContent = "";
}
}
async function keepWarm(start) {
const id = $("model").value.trim();
if (!id) { showError("Select a model first"); return; }
errorBox.textContent = "";
$("runtime-status").textContent = start ? "starting…" : "stopping…";
try {
const response = await fetch("/v1/models/runtime", {
method: "POST",
headers: Object.assign({ "Content-Type": "application/json" }, headers()),
body: JSON.stringify({ model: id, keep_warm: start ? (Number($("keep-warm-secs").value) || 2) : 0 }),
});
const body = await response.json().catch(() => null);
if (!response.ok) throw new Error(body && body.error ? body.error.message : "HTTP " + response.status);
$("runtime-status").textContent = start ? ("warm every " + body.keep_warm + "s") : "warm stopped";
} catch (err) {
showError(String((err && err.message) || err));
$("runtime-status").textContent = "";
}
}
$("count-tokens").addEventListener("click", countTokens);
$("ppl-score").addEventListener("click", scorePrompt);
$("lock-toggle").addEventListener("click", togglePin);
$("keep-warm-start").addEventListener("click", () => keepWarm(true));
$("keep-warm-stop").addEventListener("click", () => keepWarm(false));
$("key").value = localStorage.getItem("cism-webui-key") || "";
if ($("key").value) $("remember").checked = true;
loadModels();
syncBenchState();
</script>
</body>
</html>
"""