llm

AsyncVirtualLM

Bases: AsyncLM

Async language model using vLLM v1 with live engine requests.

Concurrent calls are auto-batched. A one-token extension of an idle context appends to that context's engine request, and identical contexts in one batch share one row.

Source code in genlm/backend/llm/vllm.py

class AsyncVirtualLM(AsyncLM):  # pragma: no cover
    """Async language model using vLLM v1 with live engine requests.

    Concurrent calls are auto-batched. A one-token extension of an idle
    context appends to that context's engine request, and identical
    contexts in one batch share one row.
    """

    def __init__(self, llm_engine):
        """Initialize an `AsyncVirtualLM` instance.

        Args:
            llm_engine (LLM): The vLLM engine instance.
        """
        self.llm_engine = llm_engine
        self.tokenizer = llm_engine.get_tokenizer()
        self._params = SamplingParams(
            n=1, max_tokens=1 << 20, ignore_eos=True, detokenize=False
        )

        # name -> LoRARequest. vLLM caches adapter weights by int id, so ids
        # are never reused. Loop-side only; the engine thread never reads this map.
        self._lora_requests = {}
        self._next_lora_id = 1

        # Only the engine thread reads or writes the state below.
        self._requests = {}  # rid -> (tuple(ids), lora_name), recency order
        self._by_content = {}  # (tuple(ids), lora_name) -> rid
        self._pending = {}  # rid -> [(future, loop)] awaiting a row or an exception
        self._served = 0  # rows resolved; the engine thread's progress signal
        self._next_rid = 0
        self._max_requests = 1 << 30  # tightened by from_name

        self._work = queue.SimpleQueue()
        self._engine_thread = None  # started by from_name once the engine is wired
        self._sched = None
        self._core = (
            None  # in-process engine-core client, stepped by the engine thread
        )
        self._block_hasher = None

        super().__init__(tokenizer=self.tokenizer)

    @classmethod
    def from_name(cls, model_name, engine_opts=None, **kwargs):
        """Create a `AsyncVirtualLM` instance from a model name.

        Args:
            model_name (str): Name of the model to load.
            engine_opts (dict): Additional options to pass to the `LLM` engine.
            **kwargs: Additional arguments passed to `AsyncVirtualLM` constructor.

        Returns:
            (AsyncVirtualLM): An `AsyncVirtualLM` instance.
        """
        if not HAS_VLLM:
            raise ImportError(  # pragma: no cover
                "vLLM not available. Install vLLM or use AsyncTransformer instead."
            )

        engine_opts = {
            "enable_prefix_caching": True,
            "disable_log_stats": True,
            "gpu_memory_utilization": 0.9,
            "async_scheduling": True,
            "scheduler_cls": BackendScheduler,
            **(engine_opts or {}),
        }

        global _PENDING_CAPTURE
        capture = {"armed": False, "resolve": None}
        _PENDING_CAPTURE = capture
        try:
            llm = LLM(model=model_name, tokenizer=model_name, **engine_opts)
        finally:
            _PENDING_CAPTURE = None
        if not capture["armed"]:
            raise RuntimeError(
                "capture sampler was not installed: vLLM's "
                "ModelState.custom_sampler hook has moved"
            )

        inst = cls(llm, **kwargs)
        capture["resolve"] = _weak_resolve(inst)
        inst._capture = capture

        # The engine thread drives the engine core directly: LLMEngine's output
        # processor would reject requests it never registered.
        inst._core = llm.llm_engine.engine_core
        engine_core = inst._core.engine_core
        sched = engine_core.scheduler
        if not isinstance(sched, BackendScheduler):
            raise RuntimeError(
                f"engine scheduler is {type(sched).__name__}, not BackendScheduler"
            )
        inst._sched = sched
        inst._block_hasher = engine_core.request_block_hasher
        inst._max_requests = max(
            8, llm.llm_engine.vllm_config.scheduler_config.max_num_seqs - 8
        )
        inst._engine_thread = threading.Thread(
            target=_engine_loop,
            args=(inst._work, weakref.ref(inst)),
            name="genlm-engine",
            daemon=True,
        )
        inst._engine_thread.start()
        return inst

    @property
    def underlying_model(self):
        """Access the underlying model for advanced use cases."""
        engine_core = self.llm_engine.llm_engine.engine_core.engine_core
        return engine_core.model_executor.driver_worker.worker.model_runner.model

    # -- LoRA -------------------------------------------------------------

    def add_new_lora(self, lora_path, lora_name="lora_1"):
        """Register a LoRA adapter under ``lora_name``.

        Re-registering an existing name rebinds it to the weights at
        ``lora_path``; forwards still pending under the old weights fail.
        Forwards select the adapter per call via ``lora_name=``.

        Args:
            lora_path (str): Path to the adapter weights directory or identifier in HuggingFace's model hub.
            lora_name (str): Name to assign to the loaded adapter.
        """
        if lora_name in self._lora_requests:
            del self._lora_requests[lora_name]
            self._work.put(("abort_adapter", lora_name, None, None))
        lid = self._next_lora_id
        self._next_lora_id += 1
        self._lora_requests[lora_name] = LoRARequest(lora_name, lid, lora_path)

    def remove_lora(self, lora_name):
        """Unregister ``lora_name`` and evict its weights. Forwards still pending
        under it fail.

        Args:
            lora_name (str): Name of the adapter to remove.
        """
        req = self._lora_requests.pop(lora_name)
        self._work.put(("remove_lora", (lora_name, req.lora_int_id), None, None))

    def _abort_adapter(self, lora_name):
        """(engine thread) Abort every request under ``lora_name``."""
        self._abort(
            [rid for rid, (_, name) in self._requests.items() if name == lora_name],
            RuntimeError(
                f"adapter {lora_name!r} was rebound or removed mid-forward"
            ),
        )

    # -- the batch --------------------------------------------------------

    async def next_token_logprobs(self, token_ids, lora_name=None):
        """Request log probabilities of next token asynchronously with auto-batching.

        Args:
            token_ids (list[int]): A list of token IDs, representing a prompt to the language model.
            lora_name (str, optional): Name of the LoRA adapter to use. Defaults to None (the base model).

        Returns:
            result (torch.Tensor): Normalized log probability tensor.
        """
        (future,) = await self._enqueue([token_ids], lora_name)
        return (await future).clone()

    async def batch_next_token_logprobs(self, token_ids_list, lora_name=None):
        """Batch request log probabilities for multiple token sequences asynchronously.

        Concurrent callers, batched or single, are dispatched together as
        one batch.

        Args:
            token_ids_list (list[list[int]]): A list of token ID lists.
            lora_name (str, optional): Name of the LoRA adapter to use. Defaults to None (the base model).

        Returns:
            (torch.Tensor): A ``[N, vocab]`` tensor of normalized log probabilities.
        """
        futures = await self._enqueue(token_ids_list, lora_name)
        return torch.stack(await asyncio.gather(*futures))

    async def _enqueue(self, token_ids_list, lora_name):
        """Queue requests whose futures share captured rows; public returns copy once."""
        self._check_alive()
        if lora_name is not None and lora_name not in self._lora_requests:
            raise ValueError(UNKNOWN_ADAPTER.format(lora_name))
        if any(not token_ids for token_ids in token_ids_list):
            raise ValueError("token_ids must not be empty")

        loop = asyncio.get_running_loop()
        futures = [loop.create_future() for _ in token_ids_list]
        batch = await self._join_batch(
            [
                ((tuple(token_ids), lora_name), loop, future)
                for token_ids, future in zip(token_ids_list, futures)
            ]
        )
        if batch is not None:
            queries = self._bind_adapters(batch)
            if queries:
                self._work.put(("queries", queries, None, None))
        return futures

    def _check_alive(self):
        if self._engine_thread is None:
            raise RuntimeError(
                "engine thread not running: this model was cleaned up"
                if getattr(self, "_engine_cleaned", False)
                else "engine thread not running; construct via from_name()"
            )

    def _bind_adapters(self, batch):
        """Resolve each query's adapter as the batch is dispatched.

        Nothing may yield between this and the ``_work`` put, or a rebind
        could land between them.
        """
        bound = []
        for key, loop, future in batch:
            lora_name = key[1]
            if lora_name is None:
                bound.append((key, None, loop, future))
            elif lora_name in self._lora_requests:
                bound.append((key, self._lora_requests[lora_name], loop, future))
            else:  # removed while this query sat in the batch
                self._resolve(
                    loop, future, exc=ValueError(UNKNOWN_ADAPTER.format(lora_name))
                )
        return bound

    # -- the engine thread --------------------------------------------------

    def _engine_step(self, item):
        """(engine thread) Handle one work item, then step the engine until no
        row is pending. A failure goes to every pending future. Returns whether
        the thread should stop."""
        stopping = item is _STOP
        try:
            if not stopping:
                self._handle(item)
            stalled = 0
            steps = 0
            while not stopping and self._pending:
                served = self._served
                self._core.get_output()
                steps += 1
                while True:
                    try:
                        nxt = self._work.get_nowait()
                    except queue.Empty:
                        break
                    if nxt is _STOP:
                        stopping = True
                        break
                    self._handle(nxt)
                    stalled = 0
                if stopping:
                    break
                if self._served == served:
                    stalled += 1
                    if stalled > 4096:
                        raise RuntimeError(
                            "engine made no progress on pending rows: "
                            f"{list(self._pending)}"
                        )
                else:
                    stalled = 0
        except BaseException as exc:
            self._fail_pending(exc)
        if stopping:
            self._fail_pending(RuntimeError("backend was shut down"))
        return stopping

    def _handle(self, item):
        """(engine thread) Execute one work item. A barrier item resolves its
        own future with the result or the failure; a query batch leaves its
        futures in ``_pending`` for ``_resolve`` or ``_fail_pending``."""
        kind, arg, future, loop = item
        if kind == "queries":
            self._reconcile(arg)
            return
        try:
            if kind == "abort_adapter":
                self._abort_adapter(arg)
            elif kind == "release":
                self._evict_idle(len(self._requests))
            elif kind == "remove_lora":
                lora_name, lora_int_id = arg
                self._abort_adapter(lora_name)
                # One step flushes the finished ids through the runner
                # before the weights they used disappear.
                self._core.get_output()
                self.llm_engine.llm_engine.remove_lora(lora_int_id)
        except BaseException as exc:
            if future is not None:
                self._resolve(loop, future, exc=exc)
            raise
        if future is not None:
            self._resolve(loop, future, None)

    def _reconcile(self, batch):
        """(engine thread) Reconcile one batch against the request table, each
        context before its extensions. A failure fails every query in the batch."""
        try:
            self._reconcile_inner(batch)
        except BaseException as exc:
            for _, _, loop, future in batch:
                self._resolve(loop, future, exc=exc)
            raise

    def _reconcile_inner(self, batch):
        self._evict_under_pressure()
        grouped = {}
        for key, lora_request, loop, future in batch:
            entry = grouped.setdefault(key, (lora_request, []))
            entry[1].append((future, loop))
        for key in sorted(grouped, key=lambda k: len(k[0])):
            lora_request, waiters = grouped[key]
            ids, lora_name = key
            rid = None
            # A request is extendable only while no row is pending on it: one
            # request cannot serve both a context and its extension.
            cand = self._by_content.get((ids[:-1], lora_name))
            if cand is not None and cand not in self._pending:
                request = self._sched.requests.get(cand)
                if request is None:
                    self._forget(cand)  # engine dropped it (e.g. preempt races)
                else:
                    rid = cand
                    self._sched.feed_token(request, ids[-1])
                    self._rekey(rid, key)
            if rid is None:
                rid = self._new_rid()
                self._sched.add_request(
                    Request(
                        request_id=rid,
                        prompt_token_ids=list(ids),
                        sampling_params=self._params,
                        pooling_params=None,
                        lora_request=lora_request,
                        block_hasher=self._block_hasher,
                    )
                )
                self._admit(rid, key)
            self._pending[rid] = waiters

        self._evict_over_cap()

    def _new_rid(self):
        self._next_rid += 1
        return f"{_REQ_PREFIX}{self._next_rid}"

    # -- delivery -----------------------------------------------------------

    def _receive_rows(self, rows):
        """(engine thread) Capture callback: resolve the forward's pending
        futures, one callback per event loop."""
        by_loop = {}
        for req_id, row in rows.items():
            waiters = self._pending.pop(req_id, None)
            if not waiters:
                continue
            self._served += 1
            for future, loop in waiters:
                by_loop.setdefault(loop, []).append((future, row))
        for loop, items in by_loop.items():

            def _set(items=items):
                for future, result in items:
                    if not future.done():
                        future.set_result(result)

            try:
                loop.call_soon_threadsafe(_set)
            except RuntimeError:
                pass  # the loop closed; the rows have no reader

    def _fail_pending(self, exc):
        """(engine thread) Abort every request with a pending row."""
        self._abort(list(self._pending), exc)

    def _fail_waiters(self, waiters, exc):
        for future, loop in waiters:
            self._resolve(loop, future, exc=exc)

    @staticmethod
    def _resolve(loop, future, result=None, exc=None):
        def _set():
            if not future.done():
                if exc is not None:
                    future.set_exception(exc)
                else:
                    future.set_result(result)

        try:
            loop.call_soon_threadsafe(_set)
        except RuntimeError:
            pass  # the loop closed; the result has no reader

    async def _submit(self, kind, arg=None):
        self._check_alive()
        loop = asyncio.get_running_loop()
        future = loop.create_future()
        self._work.put((kind, arg, future, loop))
        return await future

    # -- the request table ------------------------------------------------

    def _abort(self, rids, exc):
        """(engine thread) Fail the futures pending on ``rids`` and drop the
        requests from the table and the engine.

        Never raises: it runs in the engine thread's failure handler, where an
        exception would kill the thread.
        """
        rids = list(rids)
        for rid in rids:
            waiters = self._pending.pop(rid, None)
            if waiters:
                self._fail_waiters(waiters, exc)
            self._forget(rid)
        if rids:
            try:
                self._sched.finish_requests(rids, RequestStatus.FINISHED_ABORTED)
            except BaseException:  # pragma: no cover
                logging.getLogger(__name__).exception(
                    "could not finish aborted requests %s", rids
                )

    def _admit(self, rid, key):
        """Register a live request at ``key``. A newcomer takes the content
        index; a displaced incumbent stays in the table as an idle row."""
        self._requests[rid] = key
        self._by_content[key] = rid

    def _rekey(self, rid, key):
        """Re-key an extended request (re-insertion keeps recency order)."""
        self._forget(rid)
        self._admit(rid, key)

    def _forget(self, rid):
        """Drop ``rid`` from the table and, if it still holds it, the index."""
        key = self._requests.pop(rid, None)
        if key is not None and self._by_content.get(key) == rid:
            del self._by_content[key]

    def _evict_idle(self, n):
        """(engine thread) Finish up to ``n`` oldest requests with no pending row."""
        victims = []
        for rid in self._requests:
            if len(victims) >= n:
                break
            if rid not in self._pending:
                victims.append(rid)
        for rid in victims:
            self._forget(rid)
        if victims:
            self._sched.finish_requests(victims, RequestStatus.FINISHED_ABORTED)

    def _evict_over_cap(self):
        over = len(self._requests) - self._max_requests
        if over > 0:
            self._evict_idle(over)

    def _evict_under_pressure(self):
        """(engine thread) Evict the oldest idle requests when KV-cache usage reaches 90%."""
        if self._sched.kv_cache_manager.usage >= 0.9:
            self._evict_idle(max(1, len(self._requests) // 8))

    async def release_all(self):
        """Evict every idle request (end of an inference run)."""
        await self._submit("release")

    # -- sync paths -----------------------------------------------------

    def next_token_logprobs_sync(self, token_ids, lora_name=None):
        """Request log probabilities of next token synchronously.

        Does not support auto-batching. For batched sync calls, use
        ``batch_next_token_logprobs_sync`` instead.

        Args:
            token_ids (list[int]): A list of token IDs, representing a prompt to the language model.
            lora_name (str, optional): Name of the LoRA adapter to use. Defaults to None (the base model).

        Returns:
            (torch.Tensor): Normalized log probability tensor.
        """
        return asyncio.run(self.next_token_logprobs(token_ids, lora_name=lora_name))

    def batch_next_token_logprobs_sync(self, token_ids_list, lora_name=None):
        """
        Request log probabilities of next tokens in a batch synchronously.

        Args:
            token_ids_list (list[list[int]]): A list of token ID lists, each representing a prompt to the language model.
            lora_name (str, optional): Name of the LoRA adapter to use. Defaults to None (the base model).

        Returns:
            (torch.Tensor): A tensor of normalized log probability tensors, one for each prompt in the input list.
        """
        return asyncio.run(
            self.batch_next_token_logprobs(token_ids_list, lora_name=lora_name)
        )

    def cleanup(self):
        """Explicitly clean up GPU resources. Call this when done with the model."""
        self._cleanup_engine()

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.cleanup()
        return False

    async def __aenter__(self):
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        self.cleanup()
        return False

    def __del__(self):
        """Clean up resources on deletion."""
        self._cleanup_engine()

    def _cleanup_engine(self):
        """Clean up the vLLM engine and associated resources.

        This is invoked from both :meth:`cleanup` (explicit, during normal
        program flow) and :meth:`__del__` (implicit, possibly at
        interpreter shutdown). The narrow exception classes below cover
        the races and idempotency issues we know about:

        * ``ImportError`` / ``AttributeError`` arise when ``__del__`` runs
          after ``sys.meta_path`` is already torn down during interpreter
          shutdown.
        * ``AssertionError`` is raised by vLLM's
          ``destroy_distributed_environment`` if it's called twice.
        * ``RuntimeError`` can surface from CUDA when the driver is
          already being torn down.

        Anything else is re-raised so real bugs are not swallowed.
        """
        if getattr(self, "_engine_cleaned", False):
            return
        self._engine_cleaned = True
        thread = getattr(self, "_engine_thread", None)
        if thread is not None and thread.is_alive():
            self._work.put(_STOP)
            thread.join(timeout=10)
        self._engine_thread = None
        # Only release what the engine thread can no longer reach: a thread
        # that outlived its join would fault on the handles this drops.
        released = thread is None or not thread.is_alive()
        try:
            import gc

            capture = getattr(self, "_capture", None)
            if capture is not None:
                capture["resolve"] = _drop_rows
            if released:
                self._sched = self._core = self._block_hasher = None
                self.llm_engine = None
            # vLLM's internals are cyclic; only a collection frees them.
            gc.collect()
            if torch.cuda.is_available():
                torch.cuda.empty_cache()
                torch.cuda.synchronize()
            destroy_model_parallel()
            destroy_distributed_environment()
        except (ImportError, AttributeError, AssertionError, RuntimeError):
            pass

__init__(llm_engine)

Initialize an AsyncVirtualLM instance.

Parameters:

Name	Type	Description	Default
llm_engine	LLM	The vLLM engine instance.	required

Source code in genlm/backend/llm/vllm.py

def __init__(self, llm_engine):
    """Initialize an `AsyncVirtualLM` instance.

    Args:
        llm_engine (LLM): The vLLM engine instance.
    """
    self.llm_engine = llm_engine
    self.tokenizer = llm_engine.get_tokenizer()
    self._params = SamplingParams(
        n=1, max_tokens=1 << 20, ignore_eos=True, detokenize=False
    )

    # name -> LoRARequest. vLLM caches adapter weights by int id, so ids
    # are never reused. Loop-side only; the engine thread never reads this map.
    self._lora_requests = {}
    self._next_lora_id = 1

    # Only the engine thread reads or writes the state below.
    self._requests = {}  # rid -> (tuple(ids), lora_name), recency order
    self._by_content = {}  # (tuple(ids), lora_name) -> rid
    self._pending = {}  # rid -> [(future, loop)] awaiting a row or an exception
    self._served = 0  # rows resolved; the engine thread's progress signal
    self._next_rid = 0
    self._max_requests = 1 << 30  # tightened by from_name

    self._work = queue.SimpleQueue()
    self._engine_thread = None  # started by from_name once the engine is wired
    self._sched = None
    self._core = (
        None  # in-process engine-core client, stepped by the engine thread
    )
    self._block_hasher = None

    super().__init__(tokenizer=self.tokenizer)

from_name(model_name, engine_opts=None, **kwargs) classmethod

Create a AsyncVirtualLM instance from a model name.

Parameters:

Name	Type	Description	Default
model_name	str	Name of the model to load.	required
engine_opts	dict	Additional options to pass to the LLM engine.	None
**kwargs		Additional arguments passed to AsyncVirtualLM constructor.	{}

Returns:

Type	Description
AsyncVirtualLM	An AsyncVirtualLM instance.

Source code in genlm/backend/llm/vllm.py

@classmethod
def from_name(cls, model_name, engine_opts=None, **kwargs):
    """Create a `AsyncVirtualLM` instance from a model name.

    Args:
        model_name (str): Name of the model to load.
        engine_opts (dict): Additional options to pass to the `LLM` engine.
        **kwargs: Additional arguments passed to `AsyncVirtualLM` constructor.

    Returns:
        (AsyncVirtualLM): An `AsyncVirtualLM` instance.
    """
    if not HAS_VLLM:
        raise ImportError(  # pragma: no cover
            "vLLM not available. Install vLLM or use AsyncTransformer instead."
        )

    engine_opts = {
        "enable_prefix_caching": True,
        "disable_log_stats": True,
        "gpu_memory_utilization": 0.9,
        "async_scheduling": True,
        "scheduler_cls": BackendScheduler,
        **(engine_opts or {}),
    }

    global _PENDING_CAPTURE
    capture = {"armed": False, "resolve": None}
    _PENDING_CAPTURE = capture
    try:
        llm = LLM(model=model_name, tokenizer=model_name, **engine_opts)
    finally:
        _PENDING_CAPTURE = None
    if not capture["armed"]:
        raise RuntimeError(
            "capture sampler was not installed: vLLM's "
            "ModelState.custom_sampler hook has moved"
        )

    inst = cls(llm, **kwargs)
    capture["resolve"] = _weak_resolve(inst)
    inst._capture = capture

    # The engine thread drives the engine core directly: LLMEngine's output
    # processor would reject requests it never registered.
    inst._core = llm.llm_engine.engine_core
    engine_core = inst._core.engine_core
    sched = engine_core.scheduler
    if not isinstance(sched, BackendScheduler):
        raise RuntimeError(
            f"engine scheduler is {type(sched).__name__}, not BackendScheduler"
        )
    inst._sched = sched
    inst._block_hasher = engine_core.request_block_hasher
    inst._max_requests = max(
        8, llm.llm_engine.vllm_config.scheduler_config.max_num_seqs - 8
    )
    inst._engine_thread = threading.Thread(
        target=_engine_loop,
        args=(inst._work, weakref.ref(inst)),
        name="genlm-engine",
        daemon=True,
    )
    inst._engine_thread.start()
    return inst

underlying_model property

Access the underlying model for advanced use cases.

add_new_lora(lora_path, lora_name='lora_1')

Register a LoRA adapter under lora_name.

Re-registering an existing name rebinds it to the weights at lora_path; forwards still pending under the old weights fail. Forwards select the adapter per call via lora_name=.

Parameters:

Name	Type	Description	Default
lora_path	str	Path to the adapter weights directory or identifier in HuggingFace's model hub.	required
lora_name	str	Name to assign to the loaded adapter.	'lora_1'

Source code in genlm/backend/llm/vllm.py

def add_new_lora(self, lora_path, lora_name="lora_1"):
    """Register a LoRA adapter under ``lora_name``.

    Re-registering an existing name rebinds it to the weights at
    ``lora_path``; forwards still pending under the old weights fail.
    Forwards select the adapter per call via ``lora_name=``.

    Args:
        lora_path (str): Path to the adapter weights directory or identifier in HuggingFace's model hub.
        lora_name (str): Name to assign to the loaded adapter.
    """
    if lora_name in self._lora_requests:
        del self._lora_requests[lora_name]
        self._work.put(("abort_adapter", lora_name, None, None))
    lid = self._next_lora_id
    self._next_lora_id += 1
    self._lora_requests[lora_name] = LoRARequest(lora_name, lid, lora_path)

remove_lora(lora_name)

Unregister lora_name and evict its weights. Forwards still pending under it fail.

Parameters:

Name	Type	Description	Default
lora_name	str	Name of the adapter to remove.	required

Source code in genlm/backend/llm/vllm.py

def remove_lora(self, lora_name):
    """Unregister ``lora_name`` and evict its weights. Forwards still pending
    under it fail.

    Args:
        lora_name (str): Name of the adapter to remove.
    """
    req = self._lora_requests.pop(lora_name)
    self._work.put(("remove_lora", (lora_name, req.lora_int_id), None, None))

next_token_logprobs(token_ids, lora_name=None) async

Request log probabilities of next token asynchronously with auto-batching.

Parameters:

Name	Type	Description	Default
token_ids	list[int]	A list of token IDs, representing a prompt to the language model.	required
lora_name	str	Name of the LoRA adapter to use. Defaults to None (the base model).	None

Returns:

Name	Type	Description
result	Tensor	Normalized log probability tensor.

Source code in genlm/backend/llm/vllm.py

async def next_token_logprobs(self, token_ids, lora_name=None):
    """Request log probabilities of next token asynchronously with auto-batching.

    Args:
        token_ids (list[int]): A list of token IDs, representing a prompt to the language model.
        lora_name (str, optional): Name of the LoRA adapter to use. Defaults to None (the base model).

    Returns:
        result (torch.Tensor): Normalized log probability tensor.
    """
    (future,) = await self._enqueue([token_ids], lora_name)
    return (await future).clone()

batch_next_token_logprobs(token_ids_list, lora_name=None) async

Batch request log probabilities for multiple token sequences asynchronously.

Concurrent callers, batched or single, are dispatched together as one batch.

Parameters:

Name	Type	Description	Default
token_ids_list	list[list[int]]	A list of token ID lists.	required
lora_name	str	Name of the LoRA adapter to use. Defaults to None (the base model).	None

Returns:

Type	Description
Tensor	A [N, vocab] tensor of normalized log probabilities.

Source code in genlm/backend/llm/vllm.py

async def batch_next_token_logprobs(self, token_ids_list, lora_name=None):
    """Batch request log probabilities for multiple token sequences asynchronously.

    Concurrent callers, batched or single, are dispatched together as
    one batch.

    Args:
        token_ids_list (list[list[int]]): A list of token ID lists.
        lora_name (str, optional): Name of the LoRA adapter to use. Defaults to None (the base model).

    Returns:
        (torch.Tensor): A ``[N, vocab]`` tensor of normalized log probabilities.
    """
    futures = await self._enqueue(token_ids_list, lora_name)
    return torch.stack(await asyncio.gather(*futures))

release_all() async

Evict every idle request (end of an inference run).

Source code in genlm/backend/llm/vllm.py

async def release_all(self):
    """Evict every idle request (end of an inference run)."""
    await self._submit("release")

next_token_logprobs_sync(token_ids, lora_name=None)

Request log probabilities of next token synchronously.

Does not support auto-batching. For batched sync calls, use batch_next_token_logprobs_sync instead.

Parameters:

Name	Type	Description	Default
token_ids	list[int]	A list of token IDs, representing a prompt to the language model.	required
lora_name	str	Name of the LoRA adapter to use. Defaults to None (the base model).	None

Returns:

Type	Description
Tensor	Normalized log probability tensor.

Source code in genlm/backend/llm/vllm.py

def next_token_logprobs_sync(self, token_ids, lora_name=None):
    """Request log probabilities of next token synchronously.

    Does not support auto-batching. For batched sync calls, use
    ``batch_next_token_logprobs_sync`` instead.

    Args:
        token_ids (list[int]): A list of token IDs, representing a prompt to the language model.
        lora_name (str, optional): Name of the LoRA adapter to use. Defaults to None (the base model).

    Returns:
        (torch.Tensor): Normalized log probability tensor.
    """
    return asyncio.run(self.next_token_logprobs(token_ids, lora_name=lora_name))

batch_next_token_logprobs_sync(token_ids_list, lora_name=None)

Request log probabilities of next tokens in a batch synchronously.

Parameters:

Name	Type	Description	Default
token_ids_list	list[list[int]]	A list of token ID lists, each representing a prompt to the language model.	required
lora_name	str	Name of the LoRA adapter to use. Defaults to None (the base model).	None

Returns:

Type	Description
Tensor	A tensor of normalized log probability tensors, one for each prompt in the input list.

Source code in genlm/backend/llm/vllm.py

def batch_next_token_logprobs_sync(self, token_ids_list, lora_name=None):
    """
    Request log probabilities of next tokens in a batch synchronously.

    Args:
        token_ids_list (list[list[int]]): A list of token ID lists, each representing a prompt to the language model.
        lora_name (str, optional): Name of the LoRA adapter to use. Defaults to None (the base model).

    Returns:
        (torch.Tensor): A tensor of normalized log probability tensors, one for each prompt in the input list.
    """
    return asyncio.run(
        self.batch_next_token_logprobs(token_ids_list, lora_name=lora_name)
    )

cleanup()

Explicitly clean up GPU resources. Call this when done with the model.

Source code in genlm/backend/llm/vllm.py

def cleanup(self):
    """Explicitly clean up GPU resources. Call this when done with the model."""
    self._cleanup_engine()

__del__()

Clean up resources on deletion.

Source code in genlm/backend/llm/vllm.py

def __del__(self):
    """Clean up resources on deletion."""
    self._cleanup_engine()

AsyncTransformer

Bases: AsyncLM

Asynchronous wrapper around a HuggingFace causal language model with caching support.

This class provides an asynchronous interface to HuggingFace language models with automatic batching and caching (output and KV) for improved efficiency.

Source code in genlm/backend/llm/hf.py

class AsyncTransformer(AsyncLM):
    """Asynchronous wrapper around a HuggingFace causal language model with caching support.

    This class provides an asynchronous interface to HuggingFace language models with automatic batching
    and caching (output and KV) for improved efficiency.
    """

    @classmethod
    def from_name(cls, model_id, bitsandbytes_opts=None, hf_opts=None, **kwargs):
        """Create an AsyncTransformer instance from a pretrained HuggingFace model.

        Args:
            model_id (str): Model identifier in HuggingFace's model hub.
            bitsandbytes_opts (dict, optional): Additional configuration options for bitsandbytes quantization.
                Defaults to None.
            hf_opts (dict, optional): Additional configuration options for loading the HuggingFace model.
                Defaults to None.
            **kwargs: Additional arguments passed to the `AsyncTransformer` constructor

        Returns:
            (AsyncTransformer): An initialized `AsyncTransformer` instance.
        """
        if bitsandbytes_opts:
            bnb_config = BitsAndBytesConfig(**bitsandbytes_opts)
        else:
            bnb_config = None

        _hf_opts = {
            "device_map": "auto",
            "torch_dtype": "auto",
        }
        if hf_opts:
            _hf_opts.update(hf_opts)

        tok = AutoTokenizer.from_pretrained(model_id)
        mod = AutoModelForCausalLM.from_pretrained(
            model_id, quantization_config=bnb_config, **_hf_opts
        )

        return cls(mod, tok, **kwargs)

    @torch.no_grad()
    def __init__(self, hf_model, hf_tokenizer, batch_size=20, timeout=0.02):
        """Initialize an AsyncTransformer instance.

        Args:
            hf_model: A HuggingFace CausalLM model instance.
            hf_tokenizer: A HuggingFace Tokenizer.
            batch_size (int, optional): Maximum queries to process in one batch during auto-batching.
                Defaults to 20.
            timeout (float, optional): Seconds to wait since last query before processing current batch.
                Defaults to 0.02.
        """
        self.model = hf_model
        self.tokenizer = hf_tokenizer
        self.device = hf_model.device
        # One cache trie per LoRA adapter (None = base); tries must never mix.
        self._caches = defaultdict(TokenTrie)

        # Queries to be batched. Each query is a sequence of tokens,
        # and a Future to be called when the query is resolved.
        self.queries = []
        self.batch_size = batch_size
        self.timeout = timeout
        self.timer = None

        self.model.eval()

        super().__init__(tokenizer=self.tokenizer)

    def clear_cache(self):
        """Clear the cache of log probabilities and key/value pairs."""
        self._caches = defaultdict(TokenTrie)

    def clear_kv_cache(self):
        """Clear any key and value vectors from the cache."""
        for trie in self._caches.values():
            trie.clear_kv_cache()

    def reset_async_queries(self):
        """Clear any pending language model queries from the queue. Use this method when an exception prevented an inference algorithm from executing
        to completion."""
        self.queries = []

    @torch.no_grad()
    def cache_kv(self, prompt_tokens, lora_name=None):
        """Cache the key and value vectors for a prompt. Future queries that have this prompt as a prefix will only run the LLM on new tokens.

        Args:
            prompt_tokens (list[int]): token ids for the prompt to cache.
            lora_name (str, optional): Name of the LoRA adapter to use. Defaults to None (the base model).
        """
        self._activate(lora_name)
        result = self.model(torch.tensor([prompt_tokens]).to(self.device))
        node = self._caches[lora_name].extend_cache(
            0, prompt_tokens, result.logits[0], 0
        )
        node.past_key_values = result.past_key_values

    def _has_adapter(self, lora_name):
        return (
            getattr(self.model, "_hf_peft_config_loaded", False)
            and lora_name in self.model.peft_config
        )

    def add_new_lora(self, lora_path, lora_name="lora_1"):
        """Load a LoRA adapter into the base model.

        Loading under an existing name replaces that adapter. Select it per call with `lora_name`.

        Args:
            lora_path (str): Path to the adapter weights directory or identifier in HuggingFace's model hub.
            lora_name (str): Name to assign to the loaded adapter.
        """
        if self._has_adapter(lora_name):
            self.remove_lora(lora_name)
        self.model.load_adapter(lora_path, lora_name)

    def remove_lora(self, lora_name):
        """Remove a previously loaded LoRA adapter and its cache."""
        self.model.delete_adapter(lora_name)
        self._caches.pop(lora_name, None)

    def _activate(self, lora_name):
        """Activate `lora_name` on the model, or disable adapters if it is None.

        Adapters outside the transformers peft mixin are left untouched under None.

        Raises:
            ValueError: If `lora_name` has not been loaded.
        """
        if lora_name is None:
            if getattr(self.model, "_hf_peft_config_loaded", False):
                self.model.disable_adapters()
        else:
            if not self._has_adapter(lora_name):
                raise ValueError(UNKNOWN_ADAPTER.format(lora_name))
            self.model.set_adapter(lora_name)
            self.model.enable_adapters()

    @torch.no_grad()
    def batch_evaluate_queries(self):
        """
        Process a batch of queued language model queries, one forward per LoRA adapter.

        This method is called internally when the `batch_size` has been met or the `timeout` has expired.
        """

        queries, self.queries = self.queries, []
        if len(queries) == 0:
            return

        by_lora = defaultdict(list)
        for query in queries:
            by_lora[query.lora_name].append(query)
        for lora_name, group in by_lora.items():
            # An exception escaping here would leave every awaiting caller hanging.
            try:
                self._activate(lora_name)
                self._evaluate_queries(group)
            except Exception as exc:
                for query in group:
                    if not query.future.done():
                        query.future.set_exception(exc)

    @torch.no_grad()
    def _evaluate_queries(self, queries):
        """Run one padded forward over `queries` under the currently active adapter."""
        query_groups = defaultdict(list)
        for query in queries:
            key = tuple(query.prompt)  # XXX: cache based on past_len too?
            query_groups[key].append(query)

        # Use one representative query from each group
        unique_queries = [group[0] for group in query_groups.values()]

        past_example = next(
            (q.past for q in unique_queries if q.past is not None), None
        )
        max_past_length = max(q.past_len for q in unique_queries)
        max_query_length = max(len(q.prompt) for q in unique_queries)

        padding_token_id = (
            self.tokenizer.pad_token_id
            if self.tokenizer.pad_token_id is not None
            else 0
        )

        input_ids = torch.tensor(
            [
                q.prompt_padded(padding_token_id, max_query_length)
                for q in unique_queries
            ]
        ).to(self.device)
        attn_masks = torch.tensor(
            [
                q.attention_mask(max_past_length, max_query_length)
                for q in unique_queries
            ]
        ).to(self.device)
        posn_ids = torch.tensor(
            [q.position_ids(max_past_length, max_query_length) for q in unique_queries]
        ).to(self.device)
        if past_example is not None:
            example = past_example.layers[0].keys
            pasts = [
                [
                    torch.cat(
                        (
                            *(
                                q.past_padded(
                                    layer,
                                    side,
                                    max_past_length,
                                    example.dtype,
                                    self.device,
                                    example.shape,
                                )
                                for q in unique_queries
                            ),
                        ),
                        dim=0,
                    )
                    for side in ("keys", "values")
                ]
                for layer in range(len(past_example))
            ]
        else:
            pasts = None

        pasts = DynamicCache(None if pasts is None else [tuple(kv) for kv in pasts])

        results = self.model(
            input_ids,
            attention_mask=attn_masks,
            position_ids=posn_ids,
            past_key_values=pasts,
            use_cache=True,
        )

        assert len(results.logits) == len(unique_queries)

        for i, q in enumerate(unique_queries):
            result = results.logits[i]
            for dup_query in query_groups[tuple(q.prompt)]:
                dup_query.future.set_result(result)

    @torch.no_grad()
    def add_query(self, query, future, past, lora_name=None):
        """Add a query to be evaluated in the next batch.

        This method is called internally when a `next_token_logprobs` request is made.

        Args:
            query (list[int]): Token IDs representing the query prompt
            future (asyncio.Future): Future to store the result in
            past (DynamicCache|None): Past key/value states from previous evaluation,
                or None if this is a new query
            lora_name (str, optional): Name of the LoRA adapter to use. Defaults to None (the base model).
        """
        self.queries.append(Query(query, future, past, lora_name))

        if self.timer:
            self.timer.cancel()
            self.timer = None
        if len(self.queries) >= self.batch_size:
            self.batch_evaluate_queries()
        else:
            self.timer = asyncio.get_running_loop().call_later(
                self.timeout, lambda: self.batch_evaluate_queries()
            )

    def walk_cache(self, token_ids, lora_name=None):
        """Walk the cache tree to find the deepest node matching a sequence of tokens.

        Args:
            token_ids (list[int]): Sequence of token IDs to follow in the cache tree
            lora_name (str, optional): Which adapter's cache tree to walk. Defaults to None (the base model).

        Returns:
            tuple:
                - CacheNode: The deepest node in the cache tree that matches the token sequence
                - int: Number of tokens matched from the start of token_ids
                - DynamicCache|None: Past key/value states from the deepest cached node,
                    or None if no cached states were found
                - int: Base index indicating where the past states start in token_ids
        """
        # Walk while tokens can be found
        node = self._caches[lora_name]
        next_token_index = 0

        past = None
        base = 0
        while next_token_index < len(token_ids):
            if node.past_key_values is not None:
                past = node.past_key_values
                base = next_token_index
            if node.has_token(token_ids[next_token_index]):
                node = node.get_token(token_ids[next_token_index])
                next_token_index += 1
            else:
                break

        return node, next_token_index, past, base

    @torch.no_grad()
    async def next_token_logprobs(self, token_ids, lora_name=None):
        """Request log probabilities of next token. This version is asynchronous because it automatically batches concurrent requests; use with `await`.

        Args:
            token_ids (list[int]): a list of token ids, representing a prompt to the language model.
            lora_name (str, optional): Name of the LoRA adapter to use. Defaults to None (the base model).

        Returns:
            logprobs (torch.Tensor): a tensor of with the language model's log (normalized) probabilities for the next token following the prompt.
        """
        if not token_ids:
            raise ValueError("Token ids must not be empty")

        node, next_token_index, past, base = self.walk_cache(token_ids, lora_name)

        # If we processed all tokens, then we're done.
        if next_token_index == len(token_ids):
            return node.logprobs

        # Create a future with the prompt
        future = asyncio.get_running_loop().create_future()
        self.add_query(token_ids[base:], future, past, lora_name)
        logits = await future

        # Create new nodes
        node = node.extend_cache(next_token_index, token_ids, logits, base)

        return node.logprobs

    @torch.no_grad()
    def next_token_logprobs_sync(self, token_ids, lora_name=None):
        """Request log probabilities of next token. Not asynchronous, and does not support auto-batching.

        Args:
            token_ids (list[int]): a list of token ids, representing a prompt to the language model.
            lora_name (str, optional): Name of the LoRA adapter to use. Defaults to None (the base model).

        Returns:
            logprobs (torch.Tensor): a tensor with the language model's log (normalized) probabilities for the next token following the prompt.
        """
        if not token_ids:
            raise ValueError("Token ids must not be empty")

        # Walk while tokens can be found
        node, next_token_index, past, base = self.walk_cache(token_ids, lora_name)

        if next_token_index == len(token_ids):
            return node.logprobs

        self._activate(lora_name)
        logits = self.model(
            torch.tensor([token_ids[base:]]).to(self.device),
            past_key_values=node.past_key_values,
            use_cache=node.past_key_values is not None,
        ).logits[0]

        node = node.extend_cache(next_token_index, token_ids, logits, base)

        return node.logprobs

    def next_token_logprobs_uncached(self, token_ids, lora_name=None):
        """Request log probabilities of next token. No KV or output caching, and does not support auto-batching.

        Args:
            token_ids (list[int]): a list of token ids, representing a prompt to the language model.
            lora_name (str, optional): Name of the LoRA adapter to use. Defaults to None (the base model).

        Returns:
            logprobs (torch.Tensor): a tensor with the language model's log (normalized) probabilities for the next token following the prompt.
        """
        if not token_ids:
            raise ValueError("Token ids must not be empty")

        with torch.no_grad():
            self._activate(lora_name)
            logits = self.model(
                torch.tensor([token_ids]).to(self.device),
                past_key_values=None,
                use_cache=False,
            ).logits[0]
            return torch.log_softmax(logits[-1], dim=0, dtype=torch.float32)

from_name(model_id, bitsandbytes_opts=None, hf_opts=None, **kwargs) classmethod

Create an AsyncTransformer instance from a pretrained HuggingFace model.

Parameters:

Name	Type	Description	Default
model_id	str	Model identifier in HuggingFace's model hub.	required
bitsandbytes_opts	dict	Additional configuration options for bitsandbytes quantization. Defaults to None.	None
hf_opts	dict	Additional configuration options for loading the HuggingFace model. Defaults to None.	None
**kwargs		Additional arguments passed to the AsyncTransformer constructor	{}

Returns:

Type	Description
AsyncTransformer	An initialized AsyncTransformer instance.

Source code in genlm/backend/llm/hf.py

@classmethod
def from_name(cls, model_id, bitsandbytes_opts=None, hf_opts=None, **kwargs):
    """Create an AsyncTransformer instance from a pretrained HuggingFace model.

    Args:
        model_id (str): Model identifier in HuggingFace's model hub.
        bitsandbytes_opts (dict, optional): Additional configuration options for bitsandbytes quantization.
            Defaults to None.
        hf_opts (dict, optional): Additional configuration options for loading the HuggingFace model.
            Defaults to None.
        **kwargs: Additional arguments passed to the `AsyncTransformer` constructor

    Returns:
        (AsyncTransformer): An initialized `AsyncTransformer` instance.
    """
    if bitsandbytes_opts:
        bnb_config = BitsAndBytesConfig(**bitsandbytes_opts)
    else:
        bnb_config = None

    _hf_opts = {
        "device_map": "auto",
        "torch_dtype": "auto",
    }
    if hf_opts:
        _hf_opts.update(hf_opts)

    tok = AutoTokenizer.from_pretrained(model_id)
    mod = AutoModelForCausalLM.from_pretrained(
        model_id, quantization_config=bnb_config, **_hf_opts
    )

    return cls(mod, tok, **kwargs)

__init__(hf_model, hf_tokenizer, batch_size=20, timeout=0.02)

Initialize an AsyncTransformer instance.

Parameters:

Name	Type	Description	Default
hf_model		A HuggingFace CausalLM model instance.	required
hf_tokenizer		A HuggingFace Tokenizer.	required
batch_size	int	Maximum queries to process in one batch during auto-batching. Defaults to 20.	20
timeout	float	Seconds to wait since last query before processing current batch. Defaults to 0.02.	0.02

Source code in genlm/backend/llm/hf.py

@torch.no_grad()
def __init__(self, hf_model, hf_tokenizer, batch_size=20, timeout=0.02):
    """Initialize an AsyncTransformer instance.

    Args:
        hf_model: A HuggingFace CausalLM model instance.
        hf_tokenizer: A HuggingFace Tokenizer.
        batch_size (int, optional): Maximum queries to process in one batch during auto-batching.
            Defaults to 20.
        timeout (float, optional): Seconds to wait since last query before processing current batch.
            Defaults to 0.02.
    """
    self.model = hf_model
    self.tokenizer = hf_tokenizer
    self.device = hf_model.device
    # One cache trie per LoRA adapter (None = base); tries must never mix.
    self._caches = defaultdict(TokenTrie)

    # Queries to be batched. Each query is a sequence of tokens,
    # and a Future to be called when the query is resolved.
    self.queries = []
    self.batch_size = batch_size
    self.timeout = timeout
    self.timer = None

    self.model.eval()

    super().__init__(tokenizer=self.tokenizer)

clear_cache()

Clear the cache of log probabilities and key/value pairs.

Source code in genlm/backend/llm/hf.py

def clear_cache(self):
    """Clear the cache of log probabilities and key/value pairs."""
    self._caches = defaultdict(TokenTrie)

clear_kv_cache()

Clear any key and value vectors from the cache.

Source code in genlm/backend/llm/hf.py

def clear_kv_cache(self):
    """Clear any key and value vectors from the cache."""
    for trie in self._caches.values():
        trie.clear_kv_cache()

reset_async_queries()

Clear any pending language model queries from the queue. Use this method when an exception prevented an inference algorithm from executing to completion.

Source code in genlm/backend/llm/hf.py

def reset_async_queries(self):
    """Clear any pending language model queries from the queue. Use this method when an exception prevented an inference algorithm from executing
    to completion."""
    self.queries = []

cache_kv(prompt_tokens, lora_name=None)

Cache the key and value vectors for a prompt. Future queries that have this prompt as a prefix will only run the LLM on new tokens.

Parameters:

Name	Type	Description	Default
prompt_tokens	list[int]	token ids for the prompt to cache.	required
lora_name	str	Name of the LoRA adapter to use. Defaults to None (the base model).	None

Source code in genlm/backend/llm/hf.py

@torch.no_grad()
def cache_kv(self, prompt_tokens, lora_name=None):
    """Cache the key and value vectors for a prompt. Future queries that have this prompt as a prefix will only run the LLM on new tokens.

    Args:
        prompt_tokens (list[int]): token ids for the prompt to cache.
        lora_name (str, optional): Name of the LoRA adapter to use. Defaults to None (the base model).
    """
    self._activate(lora_name)
    result = self.model(torch.tensor([prompt_tokens]).to(self.device))
    node = self._caches[lora_name].extend_cache(
        0, prompt_tokens, result.logits[0], 0
    )
    node.past_key_values = result.past_key_values

add_new_lora(lora_path, lora_name='lora_1')

Load a LoRA adapter into the base model.

Loading under an existing name replaces that adapter. Select it per call with lora_name.

Parameters:

Name	Type	Description	Default
lora_path	str	Path to the adapter weights directory or identifier in HuggingFace's model hub.	required
lora_name	str	Name to assign to the loaded adapter.	'lora_1'

Source code in genlm/backend/llm/hf.py

def add_new_lora(self, lora_path, lora_name="lora_1"):
    """Load a LoRA adapter into the base model.

    Loading under an existing name replaces that adapter. Select it per call with `lora_name`.

    Args:
        lora_path (str): Path to the adapter weights directory or identifier in HuggingFace's model hub.
        lora_name (str): Name to assign to the loaded adapter.
    """
    if self._has_adapter(lora_name):
        self.remove_lora(lora_name)
    self.model.load_adapter(lora_path, lora_name)

remove_lora(lora_name)

Remove a previously loaded LoRA adapter and its cache.

Source code in genlm/backend/llm/hf.py

def remove_lora(self, lora_name):
    """Remove a previously loaded LoRA adapter and its cache."""
    self.model.delete_adapter(lora_name)
    self._caches.pop(lora_name, None)

batch_evaluate_queries()

Process a batch of queued language model queries, one forward per LoRA adapter.

This method is called internally when the batch_size has been met or the timeout has expired.

Source code in genlm/backend/llm/hf.py

@torch.no_grad()
def batch_evaluate_queries(self):
    """
    Process a batch of queued language model queries, one forward per LoRA adapter.

    This method is called internally when the `batch_size` has been met or the `timeout` has expired.
    """

    queries, self.queries = self.queries, []
    if len(queries) == 0:
        return

    by_lora = defaultdict(list)
    for query in queries:
        by_lora[query.lora_name].append(query)
    for lora_name, group in by_lora.items():
        # An exception escaping here would leave every awaiting caller hanging.
        try:
            self._activate(lora_name)
            self._evaluate_queries(group)
        except Exception as exc:
            for query in group:
                if not query.future.done():
                    query.future.set_exception(exc)

add_query(query, future, past, lora_name=None)

Add a query to be evaluated in the next batch.

This method is called internally when a next_token_logprobs request is made.

Parameters:

Name	Type	Description	Default
query	list[int]	Token IDs representing the query prompt	required
future	Future	Future to store the result in	required
past	DynamicCache | None	Past key/value states from previous evaluation, or None if this is a new query	required
lora_name	str	Name of the LoRA adapter to use. Defaults to None (the base model).	None

Source code in genlm/backend/llm/hf.py

@torch.no_grad()
def add_query(self, query, future, past, lora_name=None):
    """Add a query to be evaluated in the next batch.

    This method is called internally when a `next_token_logprobs` request is made.

    Args:
        query (list[int]): Token IDs representing the query prompt
        future (asyncio.Future): Future to store the result in
        past (DynamicCache|None): Past key/value states from previous evaluation,
            or None if this is a new query
        lora_name (str, optional): Name of the LoRA adapter to use. Defaults to None (the base model).
    """
    self.queries.append(Query(query, future, past, lora_name))

    if self.timer:
        self.timer.cancel()
        self.timer = None
    if len(self.queries) >= self.batch_size:
        self.batch_evaluate_queries()
    else:
        self.timer = asyncio.get_running_loop().call_later(
            self.timeout, lambda: self.batch_evaluate_queries()
        )

walk_cache(token_ids, lora_name=None)

Walk the cache tree to find the deepest node matching a sequence of tokens.

Parameters:

Name	Type	Description	Default
token_ids	list[int]	Sequence of token IDs to follow in the cache tree	required
lora_name	str	Which adapter's cache tree to walk. Defaults to None (the base model).	None

Returns:

Name	Type	Description
tuple		CacheNode: The deepest node in the cache tree that matches the token sequence int: Number of tokens matched from the start of token_ids DynamicCache|None: Past key/value states from the deepest cached node, or None if no cached states were found int: Base index indicating where the past states start in token_ids

Source code in genlm/backend/llm/hf.py

def walk_cache(self, token_ids, lora_name=None):
    """Walk the cache tree to find the deepest node matching a sequence of tokens.

    Args:
        token_ids (list[int]): Sequence of token IDs to follow in the cache tree
        lora_name (str, optional): Which adapter's cache tree to walk. Defaults to None (the base model).

    Returns:
        tuple:
            - CacheNode: The deepest node in the cache tree that matches the token sequence
            - int: Number of tokens matched from the start of token_ids
            - DynamicCache|None: Past key/value states from the deepest cached node,
                or None if no cached states were found
            - int: Base index indicating where the past states start in token_ids
    """
    # Walk while tokens can be found
    node = self._caches[lora_name]
    next_token_index = 0

    past = None
    base = 0
    while next_token_index < len(token_ids):
        if node.past_key_values is not None:
            past = node.past_key_values
            base = next_token_index
        if node.has_token(token_ids[next_token_index]):
            node = node.get_token(token_ids[next_token_index])
            next_token_index += 1
        else:
            break

    return node, next_token_index, past, base

next_token_logprobs(token_ids, lora_name=None) async

Request log probabilities of next token. This version is asynchronous because it automatically batches concurrent requests; use with await.

Parameters:

Name	Type	Description	Default
token_ids	list[int]	a list of token ids, representing a prompt to the language model.	required
lora_name	str	Name of the LoRA adapter to use. Defaults to None (the base model).	None

Returns:

Name	Type	Description
logprobs	Tensor	a tensor of with the language model's log (normalized) probabilities for the next token following the prompt.

Source code in genlm/backend/llm/hf.py

@torch.no_grad()
async def next_token_logprobs(self, token_ids, lora_name=None):
    """Request log probabilities of next token. This version is asynchronous because it automatically batches concurrent requests; use with `await`.

    Args:
        token_ids (list[int]): a list of token ids, representing a prompt to the language model.
        lora_name (str, optional): Name of the LoRA adapter to use. Defaults to None (the base model).

    Returns:
        logprobs (torch.Tensor): a tensor of with the language model's log (normalized) probabilities for the next token following the prompt.
    """
    if not token_ids:
        raise ValueError("Token ids must not be empty")

    node, next_token_index, past, base = self.walk_cache(token_ids, lora_name)

    # If we processed all tokens, then we're done.
    if next_token_index == len(token_ids):
        return node.logprobs

    # Create a future with the prompt
    future = asyncio.get_running_loop().create_future()
    self.add_query(token_ids[base:], future, past, lora_name)
    logits = await future

    # Create new nodes
    node = node.extend_cache(next_token_index, token_ids, logits, base)

    return node.logprobs

next_token_logprobs_sync(token_ids, lora_name=None)

Request log probabilities of next token. Not asynchronous, and does not support auto-batching.

Parameters:

Name	Type	Description	Default
token_ids	list[int]	a list of token ids, representing a prompt to the language model.	required
lora_name	str	Name of the LoRA adapter to use. Defaults to None (the base model).	None

Returns:

Name	Type	Description
logprobs	Tensor	a tensor with the language model's log (normalized) probabilities for the next token following the prompt.

Source code in genlm/backend/llm/hf.py

@torch.no_grad()
def next_token_logprobs_sync(self, token_ids, lora_name=None):
    """Request log probabilities of next token. Not asynchronous, and does not support auto-batching.

    Args:
        token_ids (list[int]): a list of token ids, representing a prompt to the language model.
        lora_name (str, optional): Name of the LoRA adapter to use. Defaults to None (the base model).

    Returns:
        logprobs (torch.Tensor): a tensor with the language model's log (normalized) probabilities for the next token following the prompt.
    """
    if not token_ids:
        raise ValueError("Token ids must not be empty")

    # Walk while tokens can be found
    node, next_token_index, past, base = self.walk_cache(token_ids, lora_name)

    if next_token_index == len(token_ids):
        return node.logprobs

    self._activate(lora_name)
    logits = self.model(
        torch.tensor([token_ids[base:]]).to(self.device),
        past_key_values=node.past_key_values,
        use_cache=node.past_key_values is not None,
    ).logits[0]

    node = node.extend_cache(next_token_index, token_ids, logits, base)

    return node.logprobs

next_token_logprobs_uncached(token_ids, lora_name=None)

Request log probabilities of next token. No KV or output caching, and does not support auto-batching.

Parameters:

Name	Type	Description	Default
token_ids	list[int]	a list of token ids, representing a prompt to the language model.	required
lora_name	str	Name of the LoRA adapter to use. Defaults to None (the base model).	None

Returns:

Name	Type	Description
logprobs	Tensor	a tensor with the language model's log (normalized) probabilities for the next token following the prompt.

Source code in genlm/backend/llm/hf.py

def next_token_logprobs_uncached(self, token_ids, lora_name=None):
    """Request log probabilities of next token. No KV or output caching, and does not support auto-batching.

    Args:
        token_ids (list[int]): a list of token ids, representing a prompt to the language model.
        lora_name (str, optional): Name of the LoRA adapter to use. Defaults to None (the base model).

    Returns:
        logprobs (torch.Tensor): a tensor with the language model's log (normalized) probabilities for the next token following the prompt.
    """
    if not token_ids:
        raise ValueError("Token ids must not be empty")

    with torch.no_grad():
        self._activate(lora_name)
        logits = self.model(
            torch.tensor([token_ids]).to(self.device),
            past_key_values=None,
            use_cache=False,
        ).logits[0]
        return torch.log_softmax(logits[-1], dim=0, dtype=torch.float32)

AsyncLM

Bases: ABC

Abstract base class for asynchronous language models.

This class provides an interface for language models that can generate token probabilities asynchronously. It handles tokenization and vocabulary management.

Parameters:

Name	Type	Description	Default
tokenizer		A Hugging Face tokenizer instance compatible with the language model	required

Source code in genlm/backend/llm/base.py

class AsyncLM(ABC):
    """Abstract base class for asynchronous language models.

    This class provides an interface for language models that can generate token probabilities
    asynchronously. It handles tokenization and vocabulary management.

    Args:
        tokenizer: A Hugging Face tokenizer instance compatible with the language model
    """

    def __init__(self, tokenizer):
        self.tokenizer = tokenizer
        self.byte_vocab, self.str_vocab = decode_vocab(self.tokenizer)
        self._batch_queue = []
        self._batch_armed = False

    async def _join_batch(self, entries, *, timeout=0.0):
        """Join the batch of concurrent queries on this event loop.

        The first caller leads: it holds the batch open until an event-loop pass adds
        no new query, then receives the drained batch; every other caller receives
        `None`. The leader must evaluate the batch in its own coroutine, never a
        background task. If the leader dies first, every other queued future fails
        with `BatchAbandoned`.

        Args:
            entries (list): Queries to add to the batch, each a tuple ending in its future.
            timeout (float, optional): Seconds to wait once for late callers.
                Defaults to 0.0.

        Returns:
            (list | None): The drained batch for the leader, `None` for every other caller.
        """
        self._batch_queue.extend(entries)
        if self._batch_armed:
            return None
        self._batch_armed = True
        try:
            waited = not timeout
            while True:
                n = len(self._batch_queue)
                await asyncio.sleep(0)
                if len(self._batch_queue) > n:
                    continue
                if waited:
                    break
                waited = True
                await asyncio.sleep(timeout)
            batch, self._batch_queue = self._batch_queue, []
            return batch
        except BaseException as exc:
            batch, self._batch_queue = self._batch_queue, []
            # Skip this caller's own entries: it is unwinding, so an exception set
            # on them would go unretrieved.
            mine = {id(e) for e in entries}
            fail_futures([e for e in batch if id(e) not in mine], batch_abandoned(exc))
            raise
        finally:
            self._batch_armed = False

    @abstractmethod
    async def next_token_logprobs(self, token_ids, lora_name=None):
        """Request log probabilities of next token asynchronously.

        Args:
            token_ids (list[int]): A list of token IDs representing the prompt.
            lora_name (str, optional): Name of the LoRA adapter to use. Defaults to None (the base model).

        Returns:
            (torch.Tensor): Normalized log probability tensor.
        """
        pass

    @abstractmethod
    def next_token_logprobs_sync(self, token_ids, lora_name=None):
        """Request log probabilities of next token synchronously.

        Args:
            token_ids (list[int]): A list of token IDs representing the prompt.
            lora_name (str, optional): Name of the LoRA adapter to use. Defaults to None (the base model).

        Returns:
            (torch.Tensor): Normalized log probability tensor.
        """
        pass

    async def batch_next_token_logprobs(self, token_ids_list, lora_name=None):
        """Batch request log probabilities for multiple token sequences asynchronously.

        Args:
            token_ids_list (list[list[int]]): A list of token ID lists.
            lora_name (str, optional): Name of the LoRA adapter to use. Defaults to None (the base model).

        Returns:
            (torch.Tensor): A tensor of log probability tensors.
        """
        logprobs = await asyncio.gather(
            *[
                self.next_token_logprobs(token_ids, lora_name=lora_name)
                for token_ids in token_ids_list
            ]
        )

        return torch.stack(logprobs)

    def batch_next_token_logprobs_sync(self, token_ids_list, lora_name=None):
        """Batch request log probabilities for multiple token sequences synchronously.

        Args:
            token_ids_list (list[list[int]]): A list of token ID lists.
            lora_name (str, optional): Name of the LoRA adapter to use. Defaults to None (the base model).

        Returns:
            (torch.Tensor): A tensor of log probability tensors.
        """
        return torch.stack(
            [
                self.next_token_logprobs_sync(token_ids, lora_name=lora_name)
                for token_ids in token_ids_list
            ]
        )

    def add_new_lora(self, lora_path, lora_name):
        """Load a LoRA adapter into the base model.

        Loading under an existing name replaces that adapter. Select it per call with `lora_name`.

        Args:
            lora_path (str): Path to the adapter weights directory or identifier in HuggingFace's model hub.
            lora_name (str): Name to assign to the loaded adapter.

        """
        raise NotImplementedError(
            "add_new_lora must be implemented by subclasses"
        )  # pragma: no cover

    def remove_lora(self, lora_name):
        """Remove a previously loaded LoRA adapter.

        Args:
            lora_name (str): Name of the adapter to remove.

        """
        raise NotImplementedError(
            "remove_lora must be implemented by subclasses"
        )  # pragma: no cover

    def set_lora(self, lora_path, lora_name):
        """Unsupported: adapter selection is per-request, via ``lora_name=``."""
        raise RuntimeError(
            "set_lora() was removed: there is no active-adapter global anymore. "
            "Pass lora_name= per call (next_token_logprobs/sample/...)."
        )

    def clear_lora(self):
        """Unsupported: ``lora_name=None`` (the default) is the base model."""
        raise RuntimeError(
            "clear_lora() was removed: omit lora_name (None = base model)."
        )

    def clear_cache(self):
        """Clear any caches used by the language model. No-op in base class."""
        pass  # pragma: no cover

    async def sample(
        self,
        prompt_token_ids,
        max_tokens,
        eos_token_ids,
        temperature=1.0,
        seed=None,
        lora_name=None,
    ):
        """Sample from the language model.

        Args:
            prompt_token_ids (list[int]): The token IDs of the prompt.
            eos_token_ids (list[int]): The token IDs of the end-of-sequence tokens.
            temperature (float, optional): The temperature to use to rescale the logits. Defaults to 1.0.
            max_tokens (int): The maximum number of tokens to generate.
            seed (int, optional): The seed for the random number generator. Defaults to None.
            lora_name (str, optional): Name of the LoRA adapter to use. Defaults to None (the base model).

        Returns:
            (list[int]): The sampled token IDs.
        """
        if seed is not None:
            generator = torch.Generator()
            generator.manual_seed(seed)
        else:
            generator = None

        generated_token_ids = []
        for _ in range(max_tokens):
            logprobs = await self.next_token_logprobs(
                prompt_token_ids + generated_token_ids, lora_name=lora_name
            )
            probs = torch.softmax(logprobs / temperature, dim=-1)
            next_token_id = torch.multinomial(
                probs.cpu() if seed is not None else probs,
                num_samples=1,
                generator=generator,
            ).item()
            if next_token_id in eos_token_ids:
                break
            generated_token_ids.append(next_token_id)

        return generated_token_ids

    async def batch_sample(
        self,
        prompt_token_ids_list,
        max_tokens,
        eos_token_ids,
        temperature=1.0,
        seed=None,
        lora_name=None,
    ):
        """Batch sample from the language model.

        Args:
            prompt_token_ids_list (list[list[int]]): The token IDs of the prompts.
            max_tokens (int): The maximum number of tokens to generate.
            eos_token_ids (list[int]): The token IDs of the end-of-sequence token.
            temperature (float): The temperature to use for the logits.
            seed (int, optional): The seed for the random number generator. Defaults to None.
            lora_name (str, optional): Name of the LoRA adapter to use. Defaults to None (the base model).

        Returns:
            (list[list[int]]): The sampled token IDs.
        """
        return await asyncio.gather(
            *[
                self.sample(
                    prompt_token_ids=prompt_token_ids,
                    max_tokens=max_tokens,
                    eos_token_ids=eos_token_ids,
                    temperature=temperature,
                    seed=seed,
                    lora_name=lora_name,
                )
                for prompt_token_ids in prompt_token_ids_list
            ]
        )

next_token_logprobs(token_ids, lora_name=None) abstractmethod async

Request log probabilities of next token asynchronously.

Parameters:

Name	Type	Description	Default
token_ids	list[int]	A list of token IDs representing the prompt.	required
lora_name	str	Name of the LoRA adapter to use. Defaults to None (the base model).	None

Returns:

Type	Description
Tensor	Normalized log probability tensor.

Source code in genlm/backend/llm/base.py

@abstractmethod
async def next_token_logprobs(self, token_ids, lora_name=None):
    """Request log probabilities of next token asynchronously.

    Args:
        token_ids (list[int]): A list of token IDs representing the prompt.
        lora_name (str, optional): Name of the LoRA adapter to use. Defaults to None (the base model).

    Returns:
        (torch.Tensor): Normalized log probability tensor.
    """
    pass

next_token_logprobs_sync(token_ids, lora_name=None) abstractmethod

Request log probabilities of next token synchronously.

Parameters:

Name	Type	Description	Default
token_ids	list[int]	A list of token IDs representing the prompt.	required
lora_name	str	Name of the LoRA adapter to use. Defaults to None (the base model).	None

Returns:

Type	Description
Tensor	Normalized log probability tensor.

Source code in genlm/backend/llm/base.py

@abstractmethod
def next_token_logprobs_sync(self, token_ids, lora_name=None):
    """Request log probabilities of next token synchronously.

    Args:
        token_ids (list[int]): A list of token IDs representing the prompt.
        lora_name (str, optional): Name of the LoRA adapter to use. Defaults to None (the base model).

    Returns:
        (torch.Tensor): Normalized log probability tensor.
    """
    pass

batch_next_token_logprobs(token_ids_list, lora_name=None) async

Batch request log probabilities for multiple token sequences asynchronously.

Parameters:

Name	Type	Description	Default
token_ids_list	list[list[int]]	A list of token ID lists.	required
lora_name	str	Name of the LoRA adapter to use. Defaults to None (the base model).	None

Returns:

Type	Description
Tensor	A tensor of log probability tensors.

Source code in genlm/backend/llm/base.py

async def batch_next_token_logprobs(self, token_ids_list, lora_name=None):
    """Batch request log probabilities for multiple token sequences asynchronously.

    Args:
        token_ids_list (list[list[int]]): A list of token ID lists.
        lora_name (str, optional): Name of the LoRA adapter to use. Defaults to None (the base model).

    Returns:
        (torch.Tensor): A tensor of log probability tensors.
    """
    logprobs = await asyncio.gather(
        *[
            self.next_token_logprobs(token_ids, lora_name=lora_name)
            for token_ids in token_ids_list
        ]
    )

    return torch.stack(logprobs)

batch_next_token_logprobs_sync(token_ids_list, lora_name=None)

Batch request log probabilities for multiple token sequences synchronously.

Parameters:

Name	Type	Description	Default
token_ids_list	list[list[int]]	A list of token ID lists.	required
lora_name	str	Name of the LoRA adapter to use. Defaults to None (the base model).	None

Returns:

Type	Description
Tensor	A tensor of log probability tensors.

Source code in genlm/backend/llm/base.py

def batch_next_token_logprobs_sync(self, token_ids_list, lora_name=None):
    """Batch request log probabilities for multiple token sequences synchronously.

    Args:
        token_ids_list (list[list[int]]): A list of token ID lists.
        lora_name (str, optional): Name of the LoRA adapter to use. Defaults to None (the base model).

    Returns:
        (torch.Tensor): A tensor of log probability tensors.
    """
    return torch.stack(
        [
            self.next_token_logprobs_sync(token_ids, lora_name=lora_name)
            for token_ids in token_ids_list
        ]
    )

add_new_lora(lora_path, lora_name)

Load a LoRA adapter into the base model.

Loading under an existing name replaces that adapter. Select it per call with lora_name.

Parameters:

Name	Type	Description	Default
lora_path	str	Path to the adapter weights directory or identifier in HuggingFace's model hub.	required
lora_name	str	Name to assign to the loaded adapter.	required

Source code in genlm/backend/llm/base.py

def add_new_lora(self, lora_path, lora_name):
    """Load a LoRA adapter into the base model.

    Loading under an existing name replaces that adapter. Select it per call with `lora_name`.

    Args:
        lora_path (str): Path to the adapter weights directory or identifier in HuggingFace's model hub.
        lora_name (str): Name to assign to the loaded adapter.

    """
    raise NotImplementedError(
        "add_new_lora must be implemented by subclasses"
    )  # pragma: no cover

remove_lora(lora_name)

Remove a previously loaded LoRA adapter.

Parameters:

Name	Type	Description	Default
lora_name	str	Name of the adapter to remove.	required

Source code in genlm/backend/llm/base.py

def remove_lora(self, lora_name):
    """Remove a previously loaded LoRA adapter.

    Args:
        lora_name (str): Name of the adapter to remove.

    """
    raise NotImplementedError(
        "remove_lora must be implemented by subclasses"
    )  # pragma: no cover

set_lora(lora_path, lora_name)

Unsupported: adapter selection is per-request, via lora_name=.

Source code in genlm/backend/llm/base.py

def set_lora(self, lora_path, lora_name):
    """Unsupported: adapter selection is per-request, via ``lora_name=``."""
    raise RuntimeError(
        "set_lora() was removed: there is no active-adapter global anymore. "
        "Pass lora_name= per call (next_token_logprobs/sample/...)."
    )

clear_lora()

Unsupported: lora_name=None (the default) is the base model.

Source code in genlm/backend/llm/base.py

def clear_lora(self):
    """Unsupported: ``lora_name=None`` (the default) is the base model."""
    raise RuntimeError(
        "clear_lora() was removed: omit lora_name (None = base model)."
    )

clear_cache()

Clear any caches used by the language model. No-op in base class.

Source code in genlm/backend/llm/base.py

def clear_cache(self):
    """Clear any caches used by the language model. No-op in base class."""
    pass  # pragma: no cover

sample(prompt_token_ids, max_tokens, eos_token_ids, temperature=1.0, seed=None, lora_name=None) async

Sample from the language model.

Parameters:

Name	Type	Description	Default
prompt_token_ids	list[int]	The token IDs of the prompt.	required
eos_token_ids	list[int]	The token IDs of the end-of-sequence tokens.	required
temperature	float	The temperature to use to rescale the logits. Defaults to 1.0.	1.0
max_tokens	int	The maximum number of tokens to generate.	required
seed	int	The seed for the random number generator. Defaults to None.	None
lora_name	str	Name of the LoRA adapter to use. Defaults to None (the base model).	None

Returns:

Type	Description
list[int]	The sampled token IDs.

Source code in genlm/backend/llm/base.py

async def sample(
    self,
    prompt_token_ids,
    max_tokens,
    eos_token_ids,
    temperature=1.0,
    seed=None,
    lora_name=None,
):
    """Sample from the language model.

    Args:
        prompt_token_ids (list[int]): The token IDs of the prompt.
        eos_token_ids (list[int]): The token IDs of the end-of-sequence tokens.
        temperature (float, optional): The temperature to use to rescale the logits. Defaults to 1.0.
        max_tokens (int): The maximum number of tokens to generate.
        seed (int, optional): The seed for the random number generator. Defaults to None.
        lora_name (str, optional): Name of the LoRA adapter to use. Defaults to None (the base model).

    Returns:
        (list[int]): The sampled token IDs.
    """
    if seed is not None:
        generator = torch.Generator()
        generator.manual_seed(seed)
    else:
        generator = None

    generated_token_ids = []
    for _ in range(max_tokens):
        logprobs = await self.next_token_logprobs(
            prompt_token_ids + generated_token_ids, lora_name=lora_name
        )
        probs = torch.softmax(logprobs / temperature, dim=-1)
        next_token_id = torch.multinomial(
            probs.cpu() if seed is not None else probs,
            num_samples=1,
            generator=generator,
        ).item()
        if next_token_id in eos_token_ids:
            break
        generated_token_ids.append(next_token_id)

    return generated_token_ids

batch_sample(prompt_token_ids_list, max_tokens, eos_token_ids, temperature=1.0, seed=None, lora_name=None) async

Batch sample from the language model.

Parameters:

Name	Type	Description	Default
prompt_token_ids_list	list[list[int]]	The token IDs of the prompts.	required
max_tokens	int	The maximum number of tokens to generate.	required
eos_token_ids	list[int]	The token IDs of the end-of-sequence token.	required
temperature	float	The temperature to use for the logits.	1.0
seed	int	The seed for the random number generator. Defaults to None.	None
lora_name	str	Name of the LoRA adapter to use. Defaults to None (the base model).	None

Returns:

Type	Description
list[list[int]]	The sampled token IDs.

Source code in genlm/backend/llm/base.py

async def batch_sample(
    self,
    prompt_token_ids_list,
    max_tokens,
    eos_token_ids,
    temperature=1.0,
    seed=None,
    lora_name=None,
):
    """Batch sample from the language model.

    Args:
        prompt_token_ids_list (list[list[int]]): The token IDs of the prompts.
        max_tokens (int): The maximum number of tokens to generate.
        eos_token_ids (list[int]): The token IDs of the end-of-sequence token.
        temperature (float): The temperature to use for the logits.
        seed (int, optional): The seed for the random number generator. Defaults to None.
        lora_name (str, optional): Name of the LoRA adapter to use. Defaults to None (the base model).

    Returns:
        (list[list[int]]): The sampled token IDs.
    """
    return await asyncio.gather(
        *[
            self.sample(
                prompt_token_ids=prompt_token_ids,
                max_tokens=max_tokens,
                eos_token_ids=eos_token_ids,
                temperature=temperature,
                seed=seed,
                lora_name=lora_name,
            )
            for prompt_token_ids in prompt_token_ids_list
        ]
    )

MockAsyncLM

Bases: AsyncLM

Mock implementation of AsyncLM used for testing.

Source code in genlm/backend/llm/base.py

class MockAsyncLM(AsyncLM):
    """Mock implementation of AsyncLM used for testing."""

    def __init__(self, tokenizer):
        """Initialize a `MockAsyncLM` instance.

        Args:
            tokenizer: Hugging Face tokenizer instance
        """
        super().__init__(tokenizer)
        self._rng = np.random.RandomState(42)

    @classmethod
    def from_name(cls, model_name, **kwargs):
        """Create a MockAsyncLM instance over the vocabulary of the model's tokenizer.

        Args:
            model_name (str): Name of pretrained model to load tokenizer from
            **kwargs: Additional arguments passed to `MockAsyncLM` constructor

        Returns:
            (MockAsyncLM): `MockAsyncLM` instance initialized with tokenizer from `model_name`
        """
        from transformers import AutoTokenizer

        return cls(AutoTokenizer.from_pretrained(model_name), **kwargs)

    async def next_token_logprobs(self, token_ids, lora_name=None):
        """Get next token log probabilities asynchronously.

        Args:
            token_ids (list[int]): Input token IDs.
            lora_name (str, optional): Must be None; the mock has no adapters.

        Returns:
            (torch.Tensor): Normalized log probability tensor.
        """
        if lora_name is not None:
            raise ValueError(UNKNOWN_ADAPTER.format(lora_name))
        return self._get_logprobs(token_ids)

    def next_token_logprobs_sync(self, token_ids, lora_name=None):
        """Get next token log probabilities synchronously.

        Args:
            token_ids (list[int]): Input token IDs.
            lora_name (str, optional): Must be None; the mock has no adapters.

        Returns:
            (torch.Tensor): Normalized log probability tensor.
        """
        if lora_name is not None:
            raise ValueError(UNKNOWN_ADAPTER.format(lora_name))
        return self._get_logprobs(token_ids)

    def _get_logprobs(self, token_ids):
        """Generate random but deterministic log probabilities for given tokens.

        Uses token_ids to seed the random generator, ensuring same inputs produce same outputs.

        Args:
            token_ids (list[int]): Input token IDs.

        Returns:
            (torch.Tensor): Normalized log probability tensor.
        """
        seed = sum([(i + 1) * t for i, t in enumerate(token_ids)])
        self._rng.seed(seed)
        logits = torch.from_numpy(
            self._rng.rand(len(self.byte_vocab)).astype(np.float32)
        )
        return torch.log_softmax(logits, dim=-1)

__init__(tokenizer)

Initialize a MockAsyncLM instance.

Parameters:

Name	Type	Description	Default
tokenizer		Hugging Face tokenizer instance	required

Source code in genlm/backend/llm/base.py

def __init__(self, tokenizer):
    """Initialize a `MockAsyncLM` instance.

    Args:
        tokenizer: Hugging Face tokenizer instance
    """
    super().__init__(tokenizer)
    self._rng = np.random.RandomState(42)

from_name(model_name, **kwargs) classmethod

Create a MockAsyncLM instance over the vocabulary of the model's tokenizer.

Parameters:

Name	Type	Description	Default
model_name	str	Name of pretrained model to load tokenizer from	required
**kwargs		Additional arguments passed to MockAsyncLM constructor	{}

Returns:

Type	Description
MockAsyncLM	MockAsyncLM instance initialized with tokenizer from model_name

Source code in genlm/backend/llm/base.py

@classmethod
def from_name(cls, model_name, **kwargs):
    """Create a MockAsyncLM instance over the vocabulary of the model's tokenizer.

    Args:
        model_name (str): Name of pretrained model to load tokenizer from
        **kwargs: Additional arguments passed to `MockAsyncLM` constructor

    Returns:
        (MockAsyncLM): `MockAsyncLM` instance initialized with tokenizer from `model_name`
    """
    from transformers import AutoTokenizer

    return cls(AutoTokenizer.from_pretrained(model_name), **kwargs)

next_token_logprobs(token_ids, lora_name=None) async

Get next token log probabilities asynchronously.

Parameters:

Name	Type	Description	Default
token_ids	list[int]	Input token IDs.	required
lora_name	str	Must be None; the mock has no adapters.	None

Returns:

Type	Description
Tensor	Normalized log probability tensor.

Source code in genlm/backend/llm/base.py

async def next_token_logprobs(self, token_ids, lora_name=None):
    """Get next token log probabilities asynchronously.

    Args:
        token_ids (list[int]): Input token IDs.
        lora_name (str, optional): Must be None; the mock has no adapters.

    Returns:
        (torch.Tensor): Normalized log probability tensor.
    """
    if lora_name is not None:
        raise ValueError(UNKNOWN_ADAPTER.format(lora_name))
    return self._get_logprobs(token_ids)

next_token_logprobs_sync(token_ids, lora_name=None)

Get next token log probabilities synchronously.

Parameters:

Name	Type	Description	Default
token_ids	list[int]	Input token IDs.	required
lora_name	str	Must be None; the mock has no adapters.	None

Returns:

Type	Description
Tensor	Normalized log probability tensor.

Source code in genlm/backend/llm/base.py

def next_token_logprobs_sync(self, token_ids, lora_name=None):
    """Get next token log probabilities synchronously.

    Args:
        token_ids (list[int]): Input token IDs.
        lora_name (str, optional): Must be None; the mock has no adapters.

    Returns:
        (torch.Tensor): Normalized log probability tensor.
    """
    if lora_name is not None:
        raise ValueError(UNKNOWN_ADAPTER.format(lora_name))
    return self._get_logprobs(token_ids)

AsyncMlxLM

Bases: AsyncLM

Asynchronous MLX-based language model wrapper.

This class provides an async interface to MLX language models with automatic batching, caching, and KV cache management. It extends AsyncLM to provide efficient batched inference with prefix caching.

The model automatically batches concurrent requests, continues the live KV cache when a batch extends the previous one, and optionally caches computed log probabilities for reuse.

Source code in genlm/backend/llm/mlx.py

class AsyncMlxLM(AsyncLM):
    """Asynchronous MLX-based language model wrapper.

    This class provides an async interface to MLX language models with
    automatic batching, caching, and KV cache management. It extends
    AsyncLM to provide efficient batched inference with prefix caching.

    The model automatically batches concurrent requests, continues the live
    KV cache when a batch extends the previous one, and optionally caches
    computed log probabilities for reuse.
    """

    def __init__(
        self,
        mlx_lm_model,
        tokenizer,
        timeout=0.0,
        prefill_step_size=2048,
        cache_size=0,
        cache_opts=None,
    ):
        """Initialize an `AsyncMlxLM` instance.

        Args:
            mlx_lm_model: The MLX language model instance.
            tokenizer: The tokenizer for encoding/decoding text.
            timeout (float, optional): Seconds to wait once for late concurrent
                requests before processing a batch. Defaults to 0.
            prefill_step_size (int, optional): Number of tokens to process
                per step during prompt prefilling.
            cache_size (int, optional): Maximum size of the output cache. If 0,
                caching is disabled. Defaults to 0.
            cache_opts (dict, optional): Additional options to pass to the
                [`OutputCache`][genlm.backend.cache.OutputCache] constructor.
                Defaults to None (no extra options).
        """
        self.mlx_lm_model = mlx_lm_model
        self.timeout = timeout
        self.prefill_step_size = prefill_step_size
        self._adapters = _Adapters(mlx_lm_model)
        # KV rows never cross adapters, so each adapter keeps its own pool.
        self._slots = defaultdict(
            partial(_SlotPool, mlx_lm_model, prefill_step_size)
        )
        self.cache = (
            OutputCache(maxsize=cache_size, **(cache_opts or {}))
            if cache_size > 0
            else None
        )
        # mlx_lm.load returns a TokenizerWrapper; AsyncLM needs the HuggingFace
        # tokenizer it wraps.
        super().__init__(tokenizer=getattr(tokenizer, "_tokenizer", tokenizer))

    @classmethod
    def from_name(cls, model_name, **kwargs):
        """Create an `AsyncMlxLM` instance from a model name.

        Args:
            model_name (str): Name of the model to load. Can be a Hugging Face
                model identifier or local path.
            **kwargs: Additional arguments passed to `AsyncMlxLM` constructor,
                such as `timeout`, `prefill_step_size`, `cache_size`.

        Returns:
            AsyncMlxLM: An `AsyncMlxLM` instance with the loaded model and tokenizer.
        """
        model, tokenizer = mlx_lm.load(model_name)
        return cls(model, tokenizer, **kwargs)

    def add_new_lora(self, lora_path, lora_name="lora_1"):
        """Register the adapter at ``lora_path`` under ``lora_name``.

        Re-registering a name rebinds it. Every adapter on a model must share the
        layout the first one wrapped it with.

        Args:
            lora_path (str): Directory holding ``adapter_config.json`` and
                ``adapters.safetensors``.
            lora_name (str, optional): Name to select the adapter by.
        """
        self._adapters.add(lora_name, lora_path)
        self.clear_cache()

    def remove_lora(self, lora_name):
        """Unregister ``lora_name``.

        Args:
            lora_name (str): Name of the adapter to remove.
        """
        self._adapters.remove(lora_name)
        self.clear_cache()

    def clear_cache(self):
        """Clear the output cache, the live KV rows, and the MLX device cache."""
        if self.cache is not None:
            self.cache.clear()
        self._slots.clear()
        mx.clear_cache()

    def _forward(self, prompts, lora_name):
        """Next-token log-probs under one adapter, ``[len(prompts), vocab]``."""
        with _wired(self.mlx_lm_model):
            self._adapters.select(lora_name)
            logits = self._slots[lora_name].logits(prompts).astype(mx.float32)
            logprobs = logits - mx.logsumexp(logits, axis=-1, keepdims=True)
            mx.eval(logprobs)
        return _to_torch(logprobs)

    def _resolve(self, keys):
        """``{key: logprobs}`` for distinct ``(context, lora_name)`` keys.

        Uncached keys are forwarded one batch per adapter, then memoized.
        """
        out, todo = {}, defaultdict(list)
        for key in keys:
            if self.cache is not None and key in self.cache:
                out[key] = self.cache[key]
            else:
                todo[key[1]].append(key)
        for lora_name, batch in todo.items():
            logprobs = self._forward([list(k[0]) for k in batch], lora_name)
            for row, key in enumerate(batch):
                out[key] = logprobs[row]
                if self.cache is not None:
                    self.cache[key] = logprobs[row]
        return out

    def _batch_evaluate(self, queries):
        """Resolve a batch of queries, deduplicating equal keys.

        Every future gets its row or a failure; a failed forward would
        otherwise leave the co-callers awaiting.
        """
        if not queries:
            return
        futures = defaultdict(list)
        for key, future in queries:
            futures[key].append(future)
        try:
            logprobs = self._resolve(list(futures))
        except Exception as exc:
            fail_futures(queries, exc)
            raise
        except BaseException as exc:
            fail_futures(queries, batch_abandoned(exc))
            raise
        for key, waiting in futures.items():
            for future in waiting:
                future.set_result(logprobs[key])

    async def next_token_logprobs(self, token_ids, lora_name=None):
        """Request log probabilities of next token. This version is asynchronous because it automatically batches concurrent requests; use with `await`.

        Args:
            token_ids (list[int]): a list of token ids, representing a prompt to the language model.
            lora_name (str, optional): Name of the LoRA adapter to use. Defaults to None (the base model).

        Returns:
            logprobs (torch.Tensor): a tensor of with the language model's log (normalized) probabilities for the next token following the prompt.
        """
        if not token_ids:
            raise ValueError("Token ids must not be empty")
        key = (tuple(token_ids), lora_name)
        if self.cache is not None and key in self.cache:
            return self.cache[key]
        future = asyncio.get_running_loop().create_future()
        batch = await self._join_batch([(key, future)], timeout=self.timeout)
        if batch is not None:
            self._batch_evaluate(batch)
        return await future

    def next_token_logprobs_sync(self, token_ids, lora_name=None):
        """Request log probabilities of next token synchronously.

        Args:
            token_ids (list[int]): A list of token IDs, representing a prompt to the language model.
            lora_name (str, optional): Name of the LoRA adapter to use. Defaults to None (the base model).

        Returns:
            (torch.Tensor): Normalized log probability tensor.
        """
        if not token_ids:
            raise ValueError("Token ids must not be empty")
        key = (tuple(token_ids), lora_name)
        return self._resolve([key])[key]

    def batch_next_token_logprobs_sync(self, token_ids_list, lora_name=None):
        """Next-token log-probs for each sequence, in one batched forward.

        Args:
            token_ids_list (list[list[int]]): Prompts' token ids.
            lora_name (str, optional): Name of the LoRA adapter to use. Defaults to None (the base model).

        Returns:
            (torch.Tensor): ``[len(token_ids_list), vocab]`` log-probabilities.
        """
        if any(not ids for ids in token_ids_list):
            raise ValueError("Token ids must not be empty")
        keys = [(tuple(ids), lora_name) for ids in token_ids_list]
        logprobs = self._resolve(list(dict.fromkeys(keys)))
        return torch.stack([logprobs[k] for k in keys])

    async def sample(
        self,
        prompt_token_ids,
        max_tokens,
        eos_token_ids,
        temperature=1.0,
        seed=None,
        lora_name=None,
    ):
        """Sample from the language model.

        Args:
            prompt_token_ids (list[int]): The token IDs of the prompt to
                start generation from.
            max_tokens (int): The maximum number of tokens to generate.
            eos_token_ids (list[int]): The token IDs that signal
                end-of-sequence. Generation stops when one of these is
                sampled.
            temperature (float, optional): The temperature to use for
                sampling. Higher values make the distribution more uniform,
                lower values make it more peaked. Defaults to 1.0.
            seed (int, optional): The seed for the random number generator.
                If provided, sets the random seed before sampling.
                Defaults to None.
            lora_name (str, optional): Name of the LoRA adapter to use. Defaults to None (the base model).

        Returns:
            (list[int]): The sampled token IDs.
        """
        self._adapters.select(lora_name)
        if seed is not None:
            mx.random.seed(seed)

        generated = []
        for token, _ in generate_step(
            mx.array(prompt_token_ids),
            self.mlx_lm_model,
            max_tokens=max_tokens,
            sampler=make_sampler(temp=temperature),
        ):
            if token in eos_token_ids:
                break
            generated.append(token)
        return generated

__init__(mlx_lm_model, tokenizer, timeout=0.0, prefill_step_size=2048, cache_size=0, cache_opts=None)

Initialize an AsyncMlxLM instance.

Parameters:

Name	Type	Description	Default
mlx_lm_model		The MLX language model instance.	required
tokenizer		The tokenizer for encoding/decoding text.	required
timeout	float	Seconds to wait once for late concurrent requests before processing a batch. Defaults to 0.	0.0
prefill_step_size	int	Number of tokens to process per step during prompt prefilling.	2048
cache_size	int	Maximum size of the output cache. If 0, caching is disabled. Defaults to 0.	0
cache_opts	dict	Additional options to pass to the OutputCache constructor. Defaults to None (no extra options).	None

Source code in genlm/backend/llm/mlx.py

def __init__(
    self,
    mlx_lm_model,
    tokenizer,
    timeout=0.0,
    prefill_step_size=2048,
    cache_size=0,
    cache_opts=None,
):
    """Initialize an `AsyncMlxLM` instance.

    Args:
        mlx_lm_model: The MLX language model instance.
        tokenizer: The tokenizer for encoding/decoding text.
        timeout (float, optional): Seconds to wait once for late concurrent
            requests before processing a batch. Defaults to 0.
        prefill_step_size (int, optional): Number of tokens to process
            per step during prompt prefilling.
        cache_size (int, optional): Maximum size of the output cache. If 0,
            caching is disabled. Defaults to 0.
        cache_opts (dict, optional): Additional options to pass to the
            [`OutputCache`][genlm.backend.cache.OutputCache] constructor.
            Defaults to None (no extra options).
    """
    self.mlx_lm_model = mlx_lm_model
    self.timeout = timeout
    self.prefill_step_size = prefill_step_size
    self._adapters = _Adapters(mlx_lm_model)
    # KV rows never cross adapters, so each adapter keeps its own pool.
    self._slots = defaultdict(
        partial(_SlotPool, mlx_lm_model, prefill_step_size)
    )
    self.cache = (
        OutputCache(maxsize=cache_size, **(cache_opts or {}))
        if cache_size > 0
        else None
    )
    # mlx_lm.load returns a TokenizerWrapper; AsyncLM needs the HuggingFace
    # tokenizer it wraps.
    super().__init__(tokenizer=getattr(tokenizer, "_tokenizer", tokenizer))

from_name(model_name, **kwargs) classmethod

Create an AsyncMlxLM instance from a model name.

Parameters:

Name	Type	Description	Default
model_name	str	Name of the model to load. Can be a Hugging Face model identifier or local path.	required
**kwargs		Additional arguments passed to AsyncMlxLM constructor, such as timeout, prefill_step_size, cache_size.	{}

Returns:

Name	Type	Description
AsyncMlxLM		An AsyncMlxLM instance with the loaded model and tokenizer.

Source code in genlm/backend/llm/mlx.py

@classmethod
def from_name(cls, model_name, **kwargs):
    """Create an `AsyncMlxLM` instance from a model name.

    Args:
        model_name (str): Name of the model to load. Can be a Hugging Face
            model identifier or local path.
        **kwargs: Additional arguments passed to `AsyncMlxLM` constructor,
            such as `timeout`, `prefill_step_size`, `cache_size`.

    Returns:
        AsyncMlxLM: An `AsyncMlxLM` instance with the loaded model and tokenizer.
    """
    model, tokenizer = mlx_lm.load(model_name)
    return cls(model, tokenizer, **kwargs)

add_new_lora(lora_path, lora_name='lora_1')

Register the adapter at lora_path under lora_name.

Re-registering a name rebinds it. Every adapter on a model must share the layout the first one wrapped it with.

Parameters:

Name	Type	Description	Default
lora_path	str	Directory holding adapter_config.json and adapters.safetensors.	required
lora_name	str	Name to select the adapter by.	'lora_1'

Source code in genlm/backend/llm/mlx.py

def add_new_lora(self, lora_path, lora_name="lora_1"):
    """Register the adapter at ``lora_path`` under ``lora_name``.

    Re-registering a name rebinds it. Every adapter on a model must share the
    layout the first one wrapped it with.

    Args:
        lora_path (str): Directory holding ``adapter_config.json`` and
            ``adapters.safetensors``.
        lora_name (str, optional): Name to select the adapter by.
    """
    self._adapters.add(lora_name, lora_path)
    self.clear_cache()

remove_lora(lora_name)

Unregister lora_name.

Parameters:

Name	Type	Description	Default
lora_name	str	Name of the adapter to remove.	required

Source code in genlm/backend/llm/mlx.py

def remove_lora(self, lora_name):
    """Unregister ``lora_name``.

    Args:
        lora_name (str): Name of the adapter to remove.
    """
    self._adapters.remove(lora_name)
    self.clear_cache()

clear_cache()

Clear the output cache, the live KV rows, and the MLX device cache.

Source code in genlm/backend/llm/mlx.py

def clear_cache(self):
    """Clear the output cache, the live KV rows, and the MLX device cache."""
    if self.cache is not None:
        self.cache.clear()
    self._slots.clear()
    mx.clear_cache()

next_token_logprobs(token_ids, lora_name=None) async

Request log probabilities of next token. This version is asynchronous because it automatically batches concurrent requests; use with await.

Parameters:

Name	Type	Description	Default
token_ids	list[int]	a list of token ids, representing a prompt to the language model.	required
lora_name	str	Name of the LoRA adapter to use. Defaults to None (the base model).	None

Returns:

Name	Type	Description
logprobs	Tensor	a tensor of with the language model's log (normalized) probabilities for the next token following the prompt.

Source code in genlm/backend/llm/mlx.py

async def next_token_logprobs(self, token_ids, lora_name=None):
    """Request log probabilities of next token. This version is asynchronous because it automatically batches concurrent requests; use with `await`.

    Args:
        token_ids (list[int]): a list of token ids, representing a prompt to the language model.
        lora_name (str, optional): Name of the LoRA adapter to use. Defaults to None (the base model).

    Returns:
        logprobs (torch.Tensor): a tensor of with the language model's log (normalized) probabilities for the next token following the prompt.
    """
    if not token_ids:
        raise ValueError("Token ids must not be empty")
    key = (tuple(token_ids), lora_name)
    if self.cache is not None and key in self.cache:
        return self.cache[key]
    future = asyncio.get_running_loop().create_future()
    batch = await self._join_batch([(key, future)], timeout=self.timeout)
    if batch is not None:
        self._batch_evaluate(batch)
    return await future

next_token_logprobs_sync(token_ids, lora_name=None)

Request log probabilities of next token synchronously.

Parameters:

Name	Type	Description	Default
token_ids	list[int]	A list of token IDs, representing a prompt to the language model.	required
lora_name	str	Name of the LoRA adapter to use. Defaults to None (the base model).	None

Returns:

Type	Description
Tensor	Normalized log probability tensor.

Source code in genlm/backend/llm/mlx.py

def next_token_logprobs_sync(self, token_ids, lora_name=None):
    """Request log probabilities of next token synchronously.

    Args:
        token_ids (list[int]): A list of token IDs, representing a prompt to the language model.
        lora_name (str, optional): Name of the LoRA adapter to use. Defaults to None (the base model).

    Returns:
        (torch.Tensor): Normalized log probability tensor.
    """
    if not token_ids:
        raise ValueError("Token ids must not be empty")
    key = (tuple(token_ids), lora_name)
    return self._resolve([key])[key]

batch_next_token_logprobs_sync(token_ids_list, lora_name=None)

Next-token log-probs for each sequence, in one batched forward.

Parameters:

Name	Type	Description	Default
token_ids_list	list[list[int]]	Prompts' token ids.	required
lora_name	str	Name of the LoRA adapter to use. Defaults to None (the base model).	None

Returns:

Type	Description
Tensor	[len(token_ids_list), vocab] log-probabilities.

Source code in genlm/backend/llm/mlx.py

def batch_next_token_logprobs_sync(self, token_ids_list, lora_name=None):
    """Next-token log-probs for each sequence, in one batched forward.

    Args:
        token_ids_list (list[list[int]]): Prompts' token ids.
        lora_name (str, optional): Name of the LoRA adapter to use. Defaults to None (the base model).

    Returns:
        (torch.Tensor): ``[len(token_ids_list), vocab]`` log-probabilities.
    """
    if any(not ids for ids in token_ids_list):
        raise ValueError("Token ids must not be empty")
    keys = [(tuple(ids), lora_name) for ids in token_ids_list]
    logprobs = self._resolve(list(dict.fromkeys(keys)))
    return torch.stack([logprobs[k] for k in keys])

sample(prompt_token_ids, max_tokens, eos_token_ids, temperature=1.0, seed=None, lora_name=None) async

Sample from the language model.

Parameters:

Name	Type	Description	Default
prompt_token_ids	list[int]	The token IDs of the prompt to start generation from.	required
max_tokens	int	The maximum number of tokens to generate.	required
eos_token_ids	list[int]	The token IDs that signal end-of-sequence. Generation stops when one of these is sampled.	required
temperature	float	The temperature to use for sampling. Higher values make the distribution more uniform, lower values make it more peaked. Defaults to 1.0.	1.0
seed	int	The seed for the random number generator. If provided, sets the random seed before sampling. Defaults to None.	None
lora_name	str	Name of the LoRA adapter to use. Defaults to None (the base model).	None

Returns:

Type	Description
list[int]	The sampled token IDs.

Source code in genlm/backend/llm/mlx.py

async def sample(
    self,
    prompt_token_ids,
    max_tokens,
    eos_token_ids,
    temperature=1.0,
    seed=None,
    lora_name=None,
):
    """Sample from the language model.

    Args:
        prompt_token_ids (list[int]): The token IDs of the prompt to
            start generation from.
        max_tokens (int): The maximum number of tokens to generate.
        eos_token_ids (list[int]): The token IDs that signal
            end-of-sequence. Generation stops when one of these is
            sampled.
        temperature (float, optional): The temperature to use for
            sampling. Higher values make the distribution more uniform,
            lower values make it more peaked. Defaults to 1.0.
        seed (int, optional): The seed for the random number generator.
            If provided, sets the random seed before sampling.
            Defaults to None.
        lora_name (str, optional): Name of the LoRA adapter to use. Defaults to None (the base model).

    Returns:
        (list[int]): The sampled token IDs.
    """
    self._adapters.select(lora_name)
    if seed is not None:
        mx.random.seed(seed)

    generated = []
    for token, _ in generate_step(
        mx.array(prompt_token_ids),
        self.mlx_lm_model,
        max_tokens=max_tokens,
        sampler=make_sampler(temp=temperature),
    ):
        if token in eos_token_ids:
            break
        generated.append(token)
    return generated

AsyncSGLTransformer

Bases: AsyncLM

Asynchronous wrapper around a SGLang inference engine.

This class provides an asynchronous interface to SGLang inference engine with automatic batching and caching. It extends AsyncLM to provide efficient batched inference.

The model automatically batches concurrent requests and uses a cache to store computed log probabilities for reuse.

Source code in genlm/backend/llm/sgl.py

class AsyncSGLTransformer(AsyncLM):
    """Asynchronous wrapper around a SGLang inference engine.

    This class provides an asynchronous interface to SGLang inference engine with
    automatic batching and caching. It extends AsyncLM to provide efficient
    batched inference.

    The model automatically batches concurrent requests and uses a cache to store
    computed log probabilities for reuse.
    """

    def __init__(self, sgl_model, cache_size=0, cache_opts=None):
        """Initialize an `AsyncSGLTransformer` instance.

        Args:
            sgl_model: The SGLang inference engine instance.
            cache_size (int, optional): Maximum number of log probabilities to keep in memory.
            cache_opts (dict, optional): Additional configuration options for the cache.
        """
        self.model = sgl_model
        self.tokenizer = sgl_model.tokenizer

        cache_opts = {} if cache_opts is None else cache_opts
        self.cache = (
            OutputCache(maxsize=cache_size, **cache_opts)
            if cache_size > 0
            else None
        )

        self._queue: Optional[asyncio.Queue] = None
        self._task: Optional[asyncio.Task] = None

        self._pending: Dict[Tuple[int, ...], List[asyncio.Future]] = {}
        self._inflight: Dict[Tuple[int, ...], Request] = {}

        self._rid_to_token_ids: Dict[str, Tuple[int, ...]] = {}

        super().__init__(tokenizer=self.tokenizer)

    @classmethod
    def from_name(cls, model_id, engine_opts=None, gpu_id=0, **kwargs):
        """Create an `AsyncSGLTransformer` instance from a model name.

        Args:
            model_id (str): The name of the model to load.
            engine_opts (dict, optional): Additional configuration options for the SGLang inference engine.
            gpu_id (int, optional): The GPU ID to use for the inference engine.
            **kwargs: Additional arguments passed to the `AsyncSGLTransformer` constructor.

        Returns:
            (AsyncSGLTransformer): An initialized `AsyncSGLTransformer` instance.
        """
        _engine_opts = {
            "sampling_backend": "pytorch",
            "skip_tokenizer_init": False,
            "model_path": model_id,
            "grammar_backend": "none",
            "allow_auto_truncate": False,
            # `_batch_evaluate` drives the normal event loop; the overlap loop
            # returns no results to it.
            "disable_overlap_schedule": True,
            "mem_fraction_static": 0.9,  # default value is 0.9
        }
        if engine_opts:
            _engine_opts.update(engine_opts)
        server_args = ServerArgs(**_engine_opts)
        port_args = PortArgs.init_new(server_args)
        # Scheduler reads its config from the process-wide runtime context; publish it first.
        publish(server_args, role="scheduler")
        # By keyword: the rank parameters' positions vary across sglang versions.
        mod = Scheduler(
            server_args,
            port_args,
            gpu_id,
            tp_rank=0,
            moe_ep_rank=0,
            pp_rank=0,
            attn_cp_rank=0,
            moe_dp_rank=0,
            dp_rank=0,
        )
        mod.result_queue = deque()
        # Normally set by `run_event_loop`, which this driver bypasses. A schedule
        # stream aliasing the forward stream erases scheduler/forward overlap.
        mod.schedule_stream = mod.device_module.Stream(priority=0)
        redraws = 0
        while (
            getattr(mod.schedule_stream, "cuda_stream", None)
            == getattr(mod.forward_stream, "cuda_stream", object())
            and redraws < 64
        ):
            mod.schedule_stream = mod.device_module.Stream(priority=0)
            redraws += 1
        mod._war_barrier_enabled = is_cuda() or envs.SGLANG_ENABLE_WAR_BARRIER.get()
        return cls(mod, **kwargs)

    def clear_cache(self):
        """Clear the logprobs output cache."""
        if self.cache:
            self.cache.clear()

    def clear_kv_cache(self):
        """Clear the SGLang cache."""
        return self.model.flush_cache()

    def reset_async_queries(self):
        """Clear any pending language model queries from the queue. Use this method when an exception prevented an inference algorithm from executing
        to completion."""

        for waiters in self._pending.values():
            for fut in waiters:
                fut.cancel()

        self._pending.clear()
        self._inflight.clear()
        self._rid_to_token_ids.clear()

        if self._queue:
            while True:
                try:
                    _, fut = self._queue.get_nowait()
                    fut.cancel()
                except asyncio.QueueEmpty:
                    break

        if self._task and not self._task.done():
            self._task.cancel()
        self._task = None
        self._queue = None

    def _start(self):
        """Start the background loop if it is not already running."""
        if not self._task or self._task.done():
            self._queue = asyncio.Queue()
            self._task = asyncio.create_task(self._background_loop())

    def _queue_request(self, token_ids: Tuple[int]):
        """Queue a request to the SGLang inference engine.

        Args:
            token_ids (tuple[int]): The token IDs of the request.

        Returns:
            (asyncio.Future): A future that will be set with the result of the request.
        """
        if not self._task or self._task.done():
            self._start()
        fut = asyncio.get_running_loop().create_future()
        self._queue.put_nowait((token_ids, fut))
        return fut

    async def next_token_logprobs(self, token_ids: List[int], lora_name=None):
        """Request log probabilities of next token. This version is asynchronous because it automatically batches concurrent requests; use with `await`.

        Args:
            token_ids (list[int]): a list of token ids, representing a prompt to the language model.
            lora_name (str, optional): Must be ``None``; this backend serves no adapters.

        Returns:
            logprobs (torch.Tensor): a tensor of with the language model's log (normalized) probabilities for the next token following the prompt.
        """
        if lora_name is not None:
            raise ValueError(UNKNOWN_ADAPTER.format(lora_name))
        if not token_ids:
            raise ValueError("Token ids must not be empty")

        key = tuple(token_ids)

        if self.cache is not None and key in self.cache:
            return self.cache[key]

        out = await self._queue_request(key)

        if self.cache is not None:
            self.cache[key] = out

        return out

    def next_token_logprobs_sync(self, token_ids: List[int], lora_name=None):
        """Request log probabilities of next token synchronously.

        Args:
            token_ids (list[int]): A list of token IDs, representing a prompt to the language model.
            lora_name (str, optional): Must be ``None``; this backend serves no adapters.

        Returns:
            (torch.Tensor): Normalized log probability tensor.
        """
        return self.batch_next_token_logprobs_sync(
            [token_ids], lora_name=lora_name
        )[0]

    def batch_next_token_logprobs_sync(
        self, token_ids_list: List[List[int]], lora_name=None
    ):
        """Request log probabilities of next tokens in a batch synchronously.

        Args:
            token_ids_list (list[list[int]]): A list of token ID lists, each representing a prompt.
            lora_name (str, optional): Must be ``None``; this backend serves no adapters.

        Returns:
            (torch.Tensor): A tensor of normalized log probability tensors.
        """
        if lora_name is not None:
            raise ValueError(UNKNOWN_ADAPTER.format(lora_name))
        results = {}
        to_compute = []

        for token_ids in token_ids_list:
            if not token_ids:
                raise ValueError("Token ids must not be empty")
            key = tuple(token_ids)
            if self.cache is not None and key in self.cache:
                results[key] = self.cache[key]
            elif key not in results:
                to_compute.append(key)
                results[key] = None

        if to_compute:
            requests = []
            for key in to_compute:
                req = _make_request(key)
                self._rid_to_token_ids[req.rid] = key
                requests.append(req)

            for key, logprobs in self._batch_evaluate(requests):
                results[key] = logprobs
                if self.cache is not None:
                    self.cache[key] = logprobs

        return torch.stack([results[tuple(t)] for t in token_ids_list])

    def _register(self, token_ids: Tuple[int], future: asyncio.Future):
        """Register a request with the SGLang inference engine.

        Args:
            token_ids (Tuple[int]): The token IDs of the request.
            future (asyncio.Future): A future that will be set with the result of the request.

        Returns:
            (Request | None): The Request object that was registered, or None if the request future was cancelled.
        """
        if future.cancelled():
            return None

        key = tuple(token_ids)

        self._pending.setdefault(key, []).append(future)

        if key in self._inflight:
            return None

        req = _make_request(token_ids)
        self._rid_to_token_ids[req.rid] = key
        self._inflight[key] = req
        return req

    async def _drain_queue(self) -> List[Request]:
        """Wait for at least one item, then drain all available items from the queue."""
        assert self._queue is not None

        requests = []

        # Wait for at least one item
        token_ids, future = await self._queue.get()
        req = self._register(token_ids, future)
        if req is not None:
            requests.append(req)

        while True:
            try:
                token_ids, future = self._queue.get_nowait()
                req = self._register(token_ids, future)
                if req is not None:
                    requests.append(req)
            except asyncio.QueueEmpty:
                break

        return requests

    def _batch_evaluate(self, requests: List[Request]):
        """Evaluate a batch of requests and return the token IDs and log probabilities."""
        if not requests:
            return  # pragma: no cover

        sched = self.model
        sched.process_input_requests(requests)

        # One drain of sglang's normal event loop; `running_batch` and `last_batch`
        # must carry across calls.
        with sched.device_module.StreamContext(sched.schedule_stream):
            while True:
                plan = sched.get_next_batch_to_run(
                    running_batch=sched.running_batch,
                    last_batch=sched.last_batch,
                )
                sched.running_batch = plan.running_batch
                batch = plan.batch_to_run
                sched.cur_batch_for_debug = batch
                if batch is None:
                    break
                with torch.inference_mode():
                    batch_result = sched.run_batch(batch)
                    sched.process_batch_result(batch, batch_result)
                    logprobs = torch.log_softmax(
                        batch_result.logits_output.next_token_logits, dim=-1
                    ).to("cpu")

                    for i, req in enumerate(batch.reqs):
                        if req.finished():
                            token_ids = self._rid_to_token_ids.pop(req.rid, None)
                            if token_ids is None:
                                continue  # pragma: no cover
                            yield token_ids, logprobs[i]
                sched.last_batch = batch

    async def _background_loop(self):
        """Background task that processes queued requests from the queue."""
        assert self._queue is not None
        try:
            while True:
                requests = await self._drain_queue()
                for token_ids, logprobs in self._batch_evaluate(requests):
                    waiters = self._pending.pop(token_ids, [])
                    self._inflight.pop(token_ids, None)
                    for f in waiters:
                        f.set_result(logprobs)

        except asyncio.CancelledError:
            raise
        except Exception as e:
            for waiters in self._pending.values():
                for f in waiters:
                    f.set_exception(e)
            self._pending.clear()
            self._inflight.clear()
            self._rid_to_token_ids.clear()
            raise

    def _cleanup_engine(self):
        """Clean up the SGLang inference engine and distributed environment."""
        if getattr(self, "model", None) is None:
            return  # pragma: no cover
        try:
            self.reset_async_queries()
            destroy_model_parallel()
            destroy_distributed_environment()
        except Exception:  # pragma: no cover
            pass  # pragma: no cover

    def __del__(self):  # pragma: no cover
        """Clean up the SGLang inference engine when the instance is deleted."""
        self._cleanup_engine()

__init__(sgl_model, cache_size=0, cache_opts=None)

Initialize an AsyncSGLTransformer instance.

Parameters:

Name	Type	Description	Default
sgl_model		The SGLang inference engine instance.	required
cache_size	int	Maximum number of log probabilities to keep in memory.	0
cache_opts	dict	Additional configuration options for the cache.	None

Source code in genlm/backend/llm/sgl.py

def __init__(self, sgl_model, cache_size=0, cache_opts=None):
    """Initialize an `AsyncSGLTransformer` instance.

    Args:
        sgl_model: The SGLang inference engine instance.
        cache_size (int, optional): Maximum number of log probabilities to keep in memory.
        cache_opts (dict, optional): Additional configuration options for the cache.
    """
    self.model = sgl_model
    self.tokenizer = sgl_model.tokenizer

    cache_opts = {} if cache_opts is None else cache_opts
    self.cache = (
        OutputCache(maxsize=cache_size, **cache_opts)
        if cache_size > 0
        else None
    )

    self._queue: Optional[asyncio.Queue] = None
    self._task: Optional[asyncio.Task] = None

    self._pending: Dict[Tuple[int, ...], List[asyncio.Future]] = {}
    self._inflight: Dict[Tuple[int, ...], Request] = {}

    self._rid_to_token_ids: Dict[str, Tuple[int, ...]] = {}

    super().__init__(tokenizer=self.tokenizer)

from_name(model_id, engine_opts=None, gpu_id=0, **kwargs) classmethod

Create an AsyncSGLTransformer instance from a model name.

Parameters:

Name	Type	Description	Default
model_id	str	The name of the model to load.	required
engine_opts	dict	Additional configuration options for the SGLang inference engine.	None
gpu_id	int	The GPU ID to use for the inference engine.	0
**kwargs		Additional arguments passed to the AsyncSGLTransformer constructor.	{}

Returns:

Type	Description
AsyncSGLTransformer	An initialized AsyncSGLTransformer instance.

Source code in genlm/backend/llm/sgl.py

@classmethod
def from_name(cls, model_id, engine_opts=None, gpu_id=0, **kwargs):
    """Create an `AsyncSGLTransformer` instance from a model name.

    Args:
        model_id (str): The name of the model to load.
        engine_opts (dict, optional): Additional configuration options for the SGLang inference engine.
        gpu_id (int, optional): The GPU ID to use for the inference engine.
        **kwargs: Additional arguments passed to the `AsyncSGLTransformer` constructor.

    Returns:
        (AsyncSGLTransformer): An initialized `AsyncSGLTransformer` instance.
    """
    _engine_opts = {
        "sampling_backend": "pytorch",
        "skip_tokenizer_init": False,
        "model_path": model_id,
        "grammar_backend": "none",
        "allow_auto_truncate": False,
        # `_batch_evaluate` drives the normal event loop; the overlap loop
        # returns no results to it.
        "disable_overlap_schedule": True,
        "mem_fraction_static": 0.9,  # default value is 0.9
    }
    if engine_opts:
        _engine_opts.update(engine_opts)
    server_args = ServerArgs(**_engine_opts)
    port_args = PortArgs.init_new(server_args)
    # Scheduler reads its config from the process-wide runtime context; publish it first.
    publish(server_args, role="scheduler")
    # By keyword: the rank parameters' positions vary across sglang versions.
    mod = Scheduler(
        server_args,
        port_args,
        gpu_id,
        tp_rank=0,
        moe_ep_rank=0,
        pp_rank=0,
        attn_cp_rank=0,
        moe_dp_rank=0,
        dp_rank=0,
    )
    mod.result_queue = deque()
    # Normally set by `run_event_loop`, which this driver bypasses. A schedule
    # stream aliasing the forward stream erases scheduler/forward overlap.
    mod.schedule_stream = mod.device_module.Stream(priority=0)
    redraws = 0
    while (
        getattr(mod.schedule_stream, "cuda_stream", None)
        == getattr(mod.forward_stream, "cuda_stream", object())
        and redraws < 64
    ):
        mod.schedule_stream = mod.device_module.Stream(priority=0)
        redraws += 1
    mod._war_barrier_enabled = is_cuda() or envs.SGLANG_ENABLE_WAR_BARRIER.get()
    return cls(mod, **kwargs)

clear_cache()

Clear the logprobs output cache.

Source code in genlm/backend/llm/sgl.py

def clear_cache(self):
    """Clear the logprobs output cache."""
    if self.cache:
        self.cache.clear()

clear_kv_cache()

Clear the SGLang cache.

Source code in genlm/backend/llm/sgl.py

def clear_kv_cache(self):
    """Clear the SGLang cache."""
    return self.model.flush_cache()

reset_async_queries()

Clear any pending language model queries from the queue. Use this method when an exception prevented an inference algorithm from executing to completion.

Source code in genlm/backend/llm/sgl.py

def reset_async_queries(self):
    """Clear any pending language model queries from the queue. Use this method when an exception prevented an inference algorithm from executing
    to completion."""

    for waiters in self._pending.values():
        for fut in waiters:
            fut.cancel()

    self._pending.clear()
    self._inflight.clear()
    self._rid_to_token_ids.clear()

    if self._queue:
        while True:
            try:
                _, fut = self._queue.get_nowait()
                fut.cancel()
            except asyncio.QueueEmpty:
                break

    if self._task and not self._task.done():
        self._task.cancel()
    self._task = None
    self._queue = None

next_token_logprobs(token_ids, lora_name=None) async

Request log probabilities of next token. This version is asynchronous because it automatically batches concurrent requests; use with await.

Parameters:

Name	Type	Description	Default
token_ids	list[int]	a list of token ids, representing a prompt to the language model.	required
lora_name	str	Must be None; this backend serves no adapters.	None

Returns:

Name	Type	Description
logprobs	Tensor	a tensor of with the language model's log (normalized) probabilities for the next token following the prompt.

Source code in genlm/backend/llm/sgl.py

async def next_token_logprobs(self, token_ids: List[int], lora_name=None):
    """Request log probabilities of next token. This version is asynchronous because it automatically batches concurrent requests; use with `await`.

    Args:
        token_ids (list[int]): a list of token ids, representing a prompt to the language model.
        lora_name (str, optional): Must be ``None``; this backend serves no adapters.

    Returns:
        logprobs (torch.Tensor): a tensor of with the language model's log (normalized) probabilities for the next token following the prompt.
    """
    if lora_name is not None:
        raise ValueError(UNKNOWN_ADAPTER.format(lora_name))
    if not token_ids:
        raise ValueError("Token ids must not be empty")

    key = tuple(token_ids)

    if self.cache is not None and key in self.cache:
        return self.cache[key]

    out = await self._queue_request(key)

    if self.cache is not None:
        self.cache[key] = out

    return out

next_token_logprobs_sync(token_ids, lora_name=None)

Request log probabilities of next token synchronously.

Parameters:

Name	Type	Description	Default
token_ids	list[int]	A list of token IDs, representing a prompt to the language model.	required
lora_name	str	Must be None; this backend serves no adapters.	None

Returns:

Type	Description
Tensor	Normalized log probability tensor.

Source code in genlm/backend/llm/sgl.py

def next_token_logprobs_sync(self, token_ids: List[int], lora_name=None):
    """Request log probabilities of next token synchronously.

    Args:
        token_ids (list[int]): A list of token IDs, representing a prompt to the language model.
        lora_name (str, optional): Must be ``None``; this backend serves no adapters.

    Returns:
        (torch.Tensor): Normalized log probability tensor.
    """
    return self.batch_next_token_logprobs_sync(
        [token_ids], lora_name=lora_name
    )[0]

batch_next_token_logprobs_sync(token_ids_list, lora_name=None)

Request log probabilities of next tokens in a batch synchronously.

Parameters:

Name	Type	Description	Default
token_ids_list	list[list[int]]	A list of token ID lists, each representing a prompt.	required
lora_name	str	Must be None; this backend serves no adapters.	None

Returns:

Type	Description
Tensor	A tensor of normalized log probability tensors.

Source code in genlm/backend/llm/sgl.py

def batch_next_token_logprobs_sync(
    self, token_ids_list: List[List[int]], lora_name=None
):
    """Request log probabilities of next tokens in a batch synchronously.

    Args:
        token_ids_list (list[list[int]]): A list of token ID lists, each representing a prompt.
        lora_name (str, optional): Must be ``None``; this backend serves no adapters.

    Returns:
        (torch.Tensor): A tensor of normalized log probability tensors.
    """
    if lora_name is not None:
        raise ValueError(UNKNOWN_ADAPTER.format(lora_name))
    results = {}
    to_compute = []

    for token_ids in token_ids_list:
        if not token_ids:
            raise ValueError("Token ids must not be empty")
        key = tuple(token_ids)
        if self.cache is not None and key in self.cache:
            results[key] = self.cache[key]
        elif key not in results:
            to_compute.append(key)
            results[key] = None

    if to_compute:
        requests = []
        for key in to_compute:
            req = _make_request(key)
            self._rid_to_token_ids[req.rid] = key
            requests.append(req)

        for key, logprobs in self._batch_evaluate(requests):
            results[key] = logprobs
            if self.cache is not None:
                self.cache[key] = logprobs

    return torch.stack([results[tuple(t)] for t in token_ids_list])

__del__()

Clean up the SGLang inference engine when the instance is deleted.

Source code in genlm/backend/llm/sgl.py

def __del__(self):  # pragma: no cover
    """Clean up the SGLang inference engine when the instance is deleted."""
    self._cleanup_engine()

load_model_by_name(name, backend=None, llm_opts=None)

Load a language model by name.

Parameters:

Name	Type	Description	Default
name	str	Hugging Face model name (e.g. "openai-community/gpt2", "meta-llama/Llama-3.2-1B-Instruct")	required
backend	str	Backend to use for inference. Can be "vllm", "hf", "mlx", "sgl", or "mock". If None, defaults to "vllm" if CUDA is available, otherwise "hf".	None
llm_opts	dict	Additional options to pass to the backend constructor. See AsyncVirtualLM and AsyncTransformer documentation for details.	None

Returns:

Type	Description
AsyncLM	An asynchronous language model.

Raises:

Type	Description
ValueError	If an invalid backend is specified.

Source code in genlm/backend/llm/__init__.py

def load_model_by_name(name, backend=None, llm_opts=None):
    """Load a language model by name.

    Args:
        name (str): Hugging Face model name (e.g. "openai-community/gpt2", "meta-llama/Llama-3.2-1B-Instruct")
        backend (str, optional): Backend to use for inference. Can be "vllm", "hf", "mlx", "sgl", or "mock".
            If None, defaults to "vllm" if CUDA is available, otherwise "hf".
        llm_opts (dict, optional): Additional options to pass to the backend constructor.
            See AsyncVirtualLM and AsyncTransformer documentation for details.

    Returns:
        (AsyncLM): An asynchronous language model.

    Raises:
        (ValueError): If an invalid backend is specified.
    """
    if backend is None:
        backend = "vllm" if torch.cuda.is_available() else "hf"

    if llm_opts is None:
        llm_opts = {}

    if backend == "vllm":
        return AsyncVirtualLM.from_name(name, **llm_opts)
    elif backend == "hf":
        return AsyncTransformer.from_name(name, **llm_opts)
    elif backend == "mock":
        return MockAsyncLM.from_name(name, **llm_opts)
    elif backend == "mlx":
        return AsyncMlxLM.from_name(name, **llm_opts)
    elif backend == "sgl":
        return AsyncSGLTransformer.from_name(name, **llm_opts)
    else:
        raise ValueError(f"Invalid backend: {backend}")